Mamba: o modelo de espaço de estados que desafia os Transformers

Mamba é um modelo de espaço de estados (State Space Model) que desafia os Transformers

• A IA atualmente domina o mundo, e os Transformers estão no centro disso
• Mamba pertence a uma classe alternativa de modelos chamada modelos de espaço de estados (SSM)
• Mamba oferece desempenho e escalabilidade semelhantes aos dos Transformers, mas consegue operar com sequências longas
• Um ponto de destaque é que o Mamba elimina o "gargalo quadrático" do "mecanismo de atenção", permitindo contextos longos
• Mamba roda até 5 vezes mais rápido que os Transformers

Problemas dos Transformers - atenção sozinha pode não ser suficiente

• Nos Transformers, cada token pode consultar todos os tokens anteriores, o que faz o modelo ficar mais lento à medida que o contexto cresce
• O armazenamento desse cache KV também exige complexidade espacial O(n)
• Existem técnicas para aliviar os gargalos dos Transformers atuais, mas uma solução fundamental exige uma abordagem diferente

Backbone de modelos fundacionais

• Dois componentes importantes de um bom backbone de arquitetura de ML
  • Comunicação (Communication) entre tokens
  • Computação (Computation) dentro dos tokens
• Um bloco Transformer é composto por Attention e MLPs
• O Mamba usa SSMs inspirados na teoria de controle para comunicação e mantém projeções no estilo MLP para computação

Motivação do Mamba - voltando ao Temple Run

• Estado (state) se refere às variáveis necessárias para determinar o comportamento futuro de um sistema
• O estado é uma compressão de tudo o que é preciso saber sobre o passado, convertendo-se em um processo de decisão de Markov

Discretização - vivendo em um mundo quantizado

• O processo de converter equações diferenciais em tempo contínuo em equações de diferenças em tempo discreto é chamado de discretização (discretisation)
• O Mamba usa discretização zero-order hold (ZOH)

Entendendo as matrizes do SSM

• As matrizes A, B, C e D são responsáveis por transição de estado, mapeamento de novas entradas para o estado, mapeamento do estado para a saída do SSM e passagem de novas entradas diretamente para a saída

Eficiência vs. efetividade: Attention is Focus, Selectivity is Prioritisation (atenção é foco, seletividade é priorização)

• Os Transformers são muito efetivos, mas não são eficientes
• A arquitetura Mamba oferece uma solução que empurra a fronteira de Pareto entre eficiência e efetividade

Mecanismo de seletividade

• Seletividade (Selectivity) permite converter cada token em um estado de acordo com sua necessidade
• O Mamba torna as matrizes A, B e C funções de x, tornando-as dependentes do contexto em vez de estáticas

Problemas da seletividade

• Ao aplicar o mecanismo de seletividade, o cálculo pode ficar mais lento do que em SSMs não seletivos
• Com otimização de hardware, o Mamba pode rodar mais rápido do que Transformers de tamanho semelhante

Machine learning e economia política - qual deve ser o tamanho do estado?

• O trade-off entre eficiência e efetividade em modelos de sequência é caracterizado por quão bem o estado comprime a informação
• A representação do estado é importante, e comprimir o estado de forma seletiva e dinâmica é o ponto central

Fluxo de informação em Transformers vs. Mamba

• Os Transformers aprendem por meio dos dados de treinamento e dos dados de contexto
• No Mamba, os dados de treinamento e de contexto ficam acessíveis de forma comprimida/filtrada

Substituição de estado como um novo paradigma de prompting

• Com modelos como o Mamba, é possível compartilhar bibliotecas de estados gerados a partir de dados especializados
• Os estados podem aplicar aprendizado de contexto infinito em tempo de inferência sem backprop

Mamba e interpretabilidade mecanística

• A interpretabilidade do Mamba se concentra em entender como a informação se move entre tokens

O que vem a seguir para Mamba e SSMs

• Modelos como o Mamba provavelmente terão ótimo desempenho em cenários que exigem contextos muito longos e memória de longo prazo

Agentes e segurança em IA

• Modelos de linguagem são inerentemente seguros, mas a possibilidade de raciocínio sequencial de longo prazo reacende a importância da segurança em IA baseada em agentes

A melhor colaboração entre Transformers e Mamba

• Vale a pena combinar o contexto longo do Mamba com a alta resolução dos Transformers para sequências curtas

Opinião do GN⁺

• O Mamba resolve gargalos dos Transformers e apresenta uma alternativa eficaz para processar sequências longas
• Essa tecnologia pode ser especialmente útil em áreas nas quais sequências longas de dados são importantes, como medicina, genética e processamento de linguagem natural
• Mais pesquisas são necessárias para verificar se o mecanismo de seletividade do Mamba é realmente eficaz
• A seletividade do Mamba pode ajudar a encontrar um equilíbrio entre a alta precisão oferecida pelos Transformers e a eficiência