O que é um modelo Mamba?

Os modelos Mamba são talvez a primeira arquitetura de deep learning a rivalizar com a eficácia dos modelos transformadores na tarefa pela qual os transformadores originalmente ganharam reputação: modelagem de linguagem. Principalmente a arquitetura Mamba demonstrou a capacidade de combinar transformadores de tamanho equivalente em avaliações de benchmark de LLM, sendo muitas vezes consideravelmente mais eficiente em termos de requisitos de latência e memória.

A arquitetura Mamba foi introduzida pela primeira vez por Tri Dao e Albert Gu no artigo de 2023, “Mamba: Linear-Time Sequence Modeling with Selective State Spaces”. Um ano depois, seguiram o artigo original do Mamba com outro artigo que explorou as conexões entre SSMs e transformadores e apresentou uma versão refinada e consideravelmente mais rápida da arquitetura Mamba que eles chamaram de Mamba-2.

Embora os transformadores tenham permanecido como o modo dominante de LLM nos dois anos seguintes ao lançamento do documento original do Mamba, a arquitetura foi incorporada a um número cada vez maior de modelos de código aberto. Alguns, como o Codestral Mamba da Mistral IA, são modelos Mamba puros. Muitos outros, incluindo a série Jamba da AI2I e o IBM Granite 4.0, são modelos híbridos que incorporam camadas de atenção (transformador) e camadas SSM (Mamba). Além de seus benefícios baseados no desempenho, a proliferação de modelos baseados no Mamba promete democratizar o acesso à IA em virtude da execução sem problemas em um hardware comparativamente barato.

Os SSMs foram originalmente projetados para prever o próximo estado de uma sequência contínua, como um sinal elétrico, um padrão climático ou a trajetória de um objeto em movimento a partir de alguma entrada. Conceitual e matematicamente, estão relacionados às redes neurais recorrentes (RNNs) que dominaram o processamento de linguagem natural (NLP) antes da introdução dos transformadores em 2017, bem como a outros algoritmos de aprendizado de máquina, incluindo redes neurais convolucionais (CNNs) e modelos de Markov ocultos (HMMs).

Como seu nome sugere, os SSMs fazem previsões sobre o próximo estado em um sistema dinâmico modelando o espaço de estado: uma representação matemática de todas as variáveis de estado que descrevem o estado de um sistema e o intervalo de possibilidades de cada uma dessas variáveis combinadas entre si.

Um SSM pega uma sequência de input x(t) e a mapeia para uma representação de estado latente h(t) — análoga ao estado oculto de uma RNN — para prever uma sequência de saída y(t). No núcleo de qualquer SSM estão 2 equações:

• A equação de estado,  $h^{\prime }\left(t\right)=A*h\left(t\right)+B*x\left(t\right)$ 
• A equação de resultado,  $y\left(t\right)=C*h\left(t\right)+D*x\left(t\right)$ 

Os principais parâmetros do modelo são A, B, C e D, que normalmente assumem a forma de uma matriz de pesos. Nos campos em que os SSMs são convencionalmente usados, como a teoria do controle, essas matrizes são frequentemente consideradas fixas: representam a dinâmica de um sistema estabelecido e o SSM é utilizado para encontrar as entradas x que levam a saídas desejáveis y. Em concepções mais modernas de SSMs, essas matrizes são parâmetros a serem otimizados por meio do aprendizado de máquina. Em modelos de deep learning, essas matrizes são representadas pelos pesos que podem ser aprendidos de redes neurais.

A equação de estado descreve como o estado se altera. Os valores da matriz A determinam como cada variável de estado evolui com o tempo se deixada sem interferência. Os valores da matriz B determinam como o input—como o próximo token em uma sequência de texto—influencia cada variável de estado.

Na modelagem de linguagem, o estado atual representa o contexto de uma sequência de texto, atualizado após cada token. Sua função é equivalente à do cache KV em um modelo de transformador.

A equação de saída descreve como o estado atual influencia a saída (conforme mediado pela matriz C), bem como a entrada influencia a saída diretamente (conforme mediado pela matriz D). Como a matriz D é essencialmente externa à modelagem de h(t) em si, ela é frequentemente omitida nos diagramas e nas discussões de SSMs em favor do foco nas matrizes centrais A, B e C.

Em um Mamba LLM, a equação de saída é usada para gerar o próximo token.

Os SSMs tradicionais são projetados para modelar inputs contínuos, mas sequências de texto (e a maioria das outras modalidades de dados processadas por modelos modernos de deep learning) são inputs discretos. A utilização de SSMs para modelar uma sequência discreta exige um meio para representar suas etapas de tempo distintas e específicas como parte de um sinal contínuo.

Conceitualmente, a discretização equivale à amostragem do valor de uma função contínua em momentos específicos. Isso implica na introdução de um novo parâmetro — o tamanho da etapa , descrito como ∆— que determina por quanto tempo esse valor é amostrado ou "mantido" em cada etapa de tempo discreta t. Os ajustes em ∆ são semelhantes a alterações em qualidades, como a resolução dos dados (para dados de séries temporais) ou a taxa de quadros (para dados de vídeo). Há vários métodos de "discretização", mas a maioria das variantes modernas de SSM (incluindo o Mamba) utiliza o método simples de retenção de ordem zero (ZOH).

A discretização de um SSM permite que ele seja usado como uma RNN para tarefas de sequência para sequência. Os parâmetros e as equações de um SSM discretizado são geralmente reescritos para distingui-los de seus equivalentes de tempo contínuo, usando a notação de subscrito normalmente empregada para RNNs. Nessa notação, h_t representa o espaço de estado atualizado que o modelo gerará e h_t-1 representa o estado anterior a ele, ou seja, o espaço de estado atual.

 $h_t=\overline{A}{h}_{t-1}+\overline{B}{x}_t$ 
 $y_t=\overline{C}{h}_t$

A modelagem de dados de texto com SSMs discretos padrão é impraticável devido a uma série de deficiências que compartilham com as RNNs. Duas dessas deficiências foram abordadas com a introdução de modelos de sequência de espaço de estado estruturado (ou "modelos S4") por Albert Gu et al em 2021: a ineficiência de seu treinamento e sua incapacidade de modelar sequências longas.

Embora o sucesso dos modelos S4, e seus muitos derivados, como os modelos diagonais SSMs (DSS), S4 diagonais (S4D) e H3, tenha aberto o caminho para o que se tornou o Mamba.

O benefício de SSMs discretizados serem equivalentes a uma instância específica de uma RNN é que as RNNs são extremamente rápidas na inferência. No entanto, a desvantagem é que as RNNs são extremamente lentas para treinar.

Felizmente, os SSMs discretizados têm uma propriedade importante que os distingue de outras RNNs: modelam exclusivamente dependências lineares. Em outras palavras, utilizam somente operações de multiplicação e adição simples e diretas. Como o artigo sobre S4 demonstra, essas recorrências lineares simples, repetidas e interdependentes podem ser aplicadas em um kernel de convolução unidimensional, que vincula diretamente a entrada x à saída y em uma única etapa:. Isso pode ser computado com muita eficiência utilizando a transformada rápida de Fourier.

O único problema é que isso só é possível quando cada etapa de toda a sequência de entrada é conhecida. Isso não é possível durante a inferência, mas esse é o caso durante o treinamento. Portanto um SSM estruturado desfruta do melhor dos dois mundos: durante o treinamento, ele pode ser operado de forma muito eficiente como uma CNN; durante a inferência, ela pode ser operada de forma muito eficiente como uma RNN.

Como a maioria das RNNs, as SSMs padrão são inerentemente fracas na modelagem de dependências de longa distância. Em outras palavras, eles não são bons em entender a relação entre etapas em uma sequência que estão muito distantes, como palavras no início e no fim de um parágrafo – o que os torna fracos na modelagem de sequências longas.

Para resolver isso, Gu e seus co-autores (um dos quais foi Tri Dao) usaram uma técnica chamada HiPPO — abreviação de Operadores de Projeção Polinomial de Alta Ordem — para definir a forma como as matrizes A e B se comportam estruturando seus valores iniciais utilizando uma fórmula derivada de polinômios ortogonais. Isso contrasta com a prática padrão de aprendizado de máquina, na qual os pesos do modelo são inicializados aleatoriamente no início do treinamento do modelo. Para S4, Dao e Gu propuseram esquemas de inicialização derivados dos polinômios de Legendre. Eles exploraram fórmulas adicionais em um artigo de acompanhamento, intitulado “Como treinar seu HIPPO". ¹

O artigo S4 observa que “a simples alteração de um SSM de uma matriz aleatória A para [a matriz HiPPO] melhorou seu desempenho no benchmark MNIST sequencial de 60% para 98%”, resolvendo efetivamente o problema de memória de longo prazo dos SSMs. Variações posteriores de SSMs estruturados, como DSS, S5 e Mamba, utilizam esquemas de inicialização diferentes (geralmente mais simples) para A e B que, no entanto, mantêm os princípios básicos do HiPPO: implementar uma estrutura diagonal que impõe atualizações estáveis e algum grau de independência entre cada valor na matriz.

No núcleo da arquitetura Mamba há duas inovações. O primeiro é o modelo de espaço de estado seletivo, que disponibiliza ao Mamba um recurso crucial que antes possuía apenas os modelos de transformadores: a capacidade de focar seletivamente ou ignorar partes específicas do histórico de inputs passados com base em sua relevância atual. A outra é a varredura paralela com reconhecimento de hardware, um algoritmo que otimiza a maneira como uma unidade de processamento gráfico (GPU) lida com os cálculos do modelo em sua hierarquia de memória para maximizar a velocidade e a eficiência computacional.

Em transformadores, essa capacidade é disponibilizada pelo mecanismo de atenção que ajusta os pesos de atenção que enfatizam ou deixam de enfatizar a influência de cada token anterior com base em sua relevância para o token de input atual. Os SSMs comuns são explicitamente projetados para mapear a entrada para a saída utilizando todo o histórico de entrada. Isso é aceitável ou até desejável para algumas tarefas de modelagem de sequência, mas é uma desvantagem considerável para a maioria das tarefas avançadas de modelagem de linguagem.

Para remediar essa incapacidade de omitir ou enfatizar dinamicamente partes específicas de seu histórico de input, Dao e Gu propuseram uma nova classe de modelos de espaço de estado com uma "verificação seletiva". No artigo de Mamba, os autores observam que “às vezes abreviam SSMs seletivos como modelos S6, porque são modelos S4 com um mecanismo de seleção e calculados com uma varredura”. Eles apelidaram sua arquitetura baseada no S6 de “Mamba” porque, entre outras razões, todos esses S soam como o silvo de uma cobra.

O Mamba pode ser melhor compreendido como uma arquitetura de rede neurais que contém o modelo de espaço de estado seletivo em seu núcleo. Para uma analogia simples, o Mamba está para os SSMs seletivos assim como o modelo transformador está para o mecanismo de atenção.

Um SSM tradicional tem dinâmica fixa: as regras que regem como o estado oculto evolui de uma etapa para a próxima — os parâmetros do modelo — são os mesmos para cada entrada e em cada etapa na sequência. Essa propriedade é conhecida como invariância linear no tempo (LTI). Para fornecer aos SSMs a capacidade de priorizar ou retirar a prioridade seletivamente de informações passadas específicas com base no contexto presente, Dao e Gu reconfiguraram seus SSM de tal forma que os valores dos parâmetros do modelo chave serão diferentes para diferentes inputs.

Mais especificamente, SSMs seletivos tornam o tamanho do _passo∆ t e as matrizes B_t e C_t funções diretas do token de entrada atual x_t. Isso é alcançado inicialmente passando a embedding de vetor de x_t por meio de três camadas de projeção linear paralelas — em outras palavras, camadas de redes neurais padrão (ou camadas de MLP). Isso é equivalente a como a consulta paralela, a chave e as características do valor geram os vetores Q, K e V respectivos de uma entrada em um modelo de transformador.

Multiplicar a incorporação do vetor x_t pelos termos de peso e tendência nessa rede de projeção linear produz os valores resultantes de  ∆ _t, B_t e C_t. Os termos de peso e viés das próprias camadas de projeção linear são aprendidos durante o pré-treinamento do modelo em conjuntos de dados maciços de amostras de texto e, em seguida, (opcionalmente) refinados por meio de ajuste fino subsequente.

• O valor de ∆_t determina a magnitude da influência de x_t na memória do modelo do contexto que ele viu até agora: em outras palavras, em quanto de uma atualização haverá do estado oculto h_t-1 para h_t. Um tamanho de etapa maior ∆_t resultados em mudanças maiores e acelera a decadência - em outras palavras, o "esquecimento" de informações mais antigas contidas no estado. Por outro lado, um tamanho de etapa menor resulta em uma atualização menor. Em um tamanho de etapa pequeno o suficiente, a entrada atual não terá nenhum impacto no estado oculto.
  
  
• As alterações na matriz B_k determinam como o token de input atual atualiza o estado oculto. Por exemplo, se x_t for um token para a palavra "ontem", B_t poderá ser ajustado de forma a atualizar o estado para refletir que o contexto subsequente provavelmente pertence ao passado.
  
  
• Mudanças na matriz C_t determinam como essas informações contextuais se traduzem em influência sobre a produção do modelo yt. Continuando o exemplo em que x_k é um token para "ontem", C_k pode ser influenciado de forma que faça com que todos os verbos que forem posteriormente produzidos pelo modelo sejam conjugados no passado.

Principalmente nenhum desses ajustes baseados em entrada é feito na matriz A. Sua função permanece a mesma dos modelos S4: memorizar com eficiência todo o histórico de entradas anteriores. O papel de determinar quais partes desse histórico utilizar em um determinado momento é realizada pelas matrizes B e C.

Mas, uma vez que o modelo não é mais invariante no tempo, ele não pode mais usar o atalho de convolução durante o treinamento porque o kernel de transição não é mais constante: o ponto crucial do mecanismo de seletividade é que a transição de h_t-1 para h_t agora depende do contexto.

Em vez disso, o Mamba usa uma solução alternativa inteligente para obter benefícios de paralelização semelhantes. Como o SSM usa apenas multiplicação e adição, seus cálculos estão sujeitos à familiar propriedade associativa da matemática: eles podem ser agrupados de diferentes maneiras sem alterar o resultado final. Isso permite que os muitos cálculos sequenciais sejam divididos em pequenos blocos independentes que podem ser processados em paralelo por uma GPU por meio de uma varredura de soma de prefixos paralela.

Além disso, os resultados são combinados de maneira hierárquica específica que faz uso eficiente e otimizado dos diferentes tipos de memória de hardware em uma GPU, usando princípios semelhantes às técnicas FlashAttention, também desenvolvidas pela Tri Dao - agora onipresentes nos LLMs modernos.  

Dentro da arquitetura Mamba, o modelo S6 serve como um módulo do "bloco Mamba" maior, da mesma forma que o mecanismo de atenção serve como um módulo dentro do "bloco de atenção" maior. Ele combina o módulo S6 com uma arquitetura de rede neural restrita. Os modelos Mamba normalmente compreendem vários blocos Mamba (ou seja, uma série de camadas Mamba consecutivas em uma rede neural) antes da camada de saída que faz a previsão da saída final do modelo.

Antes de entrar no bloco Mamba, uma cópia do input é enviada diretamente para a ponta como uma conexão residual.O objetivo do funcionamento interno do bloco Mamba é não apenas determinar quais partes do contexto maior são relevantes para aquele input, mas determinar o quanto essa informação contextual deve modificar o significado original do input.

Dentro do bloco Mamba, o vetor de entrada original é processado da seguinte forma:

• Primeiro, a entrada é passada por uma camada linear que tem o dobro da largura do próprio vetor de entrada, projetando-o em um espaço de dimensão superior. Por exemplo, se o modelo originalmente representa cada token de entrada x como embedding de 512 dimensões, multiplicar x pelos pesos da camada de projeção linear o expande para um vetor de 1024 dimensões.
  
  
• Em seguida, o vetor expandido é dividido em dois. Uma metade (que chamaremos de x _proj) é alimentada no caminho que passa pelo SSM e a outra metade (que chamaremos de z _proj) é alimentada em um caminho separado que passa por um mecanismo de gating. Para maior clareza, a etapa de expansão anterior é geralmente representada como sendo realizada por duas camadas lineares paralelas.
  
  
• Antes de x_proj alcançar o SSM, ele é alimentado em uma camada de convolução unidimensional. Essa camada de convolução extrai padrões locais (como dependências entre tokens vizinhos, como pares simples de verbo-sujeito). Isso permite que o SSM "foque" na compreensão contextual das dependências globais de longo alcance.
  
  
• A saída da camada de convolução serve como entrada para uma função de ativação não linear. A introdução da não linearidade é uma marca registrada de todas as redes neurais, permitindo que elas capturem padrões mais complexos. O artigo Mamba usa Sigmoid Linear Unit (SiLU). Chamaremos o vetor resultante de x_act.
  
  
• Enquanto isso, no caminho separado do mecanismo de gating, o z_proj também serve de entrada para uma função de ativação não linear, produzindo z_act.
  
  
• No caminho SSM, o_ato xé alimentado em três camadas de projeção linear paralelas que geram os respectivos valores para ∆_x, B_x e C_x, respectivamente.
  
  
• O SSM usa esses parâmetros dependentes de entrada (e as matrizes A e D) para calcular a atualização do espaço de estado e a saída y do SSM.
  
  
• O vetor de saída y do SSM agora é multiplicado elemento a elemento pelo vetor de saída do caminho do gating zact. Essencialmente, cada elemento em z_act atua como um botão de volume em um console de mixagem de áudio: se um determinado elemento de z_act estiver próximo de zero, a multiplicação com a parte correspondente de y produzirá um valor mais próximo de zero, e sua influência será reduzida. Por outro lado, se um determinado elemento de zact for grande, a multiplicação com y amplificará a influência de suas informações contextuais.
  
  
• O vetor resultante é projetado de volta ao seu tamanho original. Pode ser entendido como um vetor de atualizações contextuais ponderadas (ou não atualizações) para cada um dos elementos do vetor de entrada original.
  
  
• Por fim, esse vetor de atualizações é adicionado à cópia do vetor de entrada original que foi enviado diretamente para o final do bloco como uma conexão residual.
  
  
• O vetor de entrada original agora foi atualizado para refletir a compreensão contextual entregue pelo SSM seletivo. Agora pode ser enviado para a próxima camada do Mamba ou, nas camadas finais do modelo, servir como input para uma função softmax que produz a probabilidade respectiva de que o vetor totalmente atualizado corresponda a cada palavra no vocabulário do modelo.

Um ano após o artigo original do Mamba, Dao e Gu deram continuidade com “Transformadores são SSMs: modelos generalizados e algoritmos eficientes por meio da dualidade do espaço de estados estruturados”. Este artigo de acompanhamento ofereceu três contribuições principais:

• Uma exploração das conexões teóricas entre Mamba e transformadores e um vocabulário compartilhado entre as duas arquiteturas
• Uma série de esclarecimentos e explorações de diferentes escolhas de design para modelos Mamba
• Uma arquitetura modificada, Maba-2, informada e aprimorada por essas explorações de design

O algoritmo Mamba-2 é significativamente mais rápido e fácil de implementar do que o Mamba original: os autores forneceram uma base de código de "SSD mínimo" que implementa o SSM seletivo em cerca de 25 linhas de código.² Essa eficiência permite que o Mamba-2 use dimensões de estado oculto muito maiores sem tornar o modelo mais lento, permitindo modelos maiores, mais poderosos e mais expressos construídos com a arquitetura. Nos testes, os modelos Mamba-2 definitivamente corresponderam ou superaram os modelos Mamba e transformadores de tamanho correspondente em uma série de tarefas subsequentes.

Como afirma a introdução do artigo, o "principal objetivo de Dao e Gu [foi] desenvolver um rico corpo de conexões teóricas entre SSMs estruturados e variantes de atenção". Isso resultou em um novo framework conceitual unindo os dois, que eles chamaram de "dualidade de espaço de estado" (SSD).³ Ao fazer isso, abriram as portas para a Mamba se beneficiar de vários anos de exploração e otimização da arquitetura de transformação.

Um benefício notável foi o desenvolvimento de um equivalente do Mamba da atenção de várias cabeças (MHA), no qual um bloco Mamba pode ser dividido em várias "cabeças Mamba", semelhante às várias "cabeças de atenção" dos transformadores. Uma variante dessa abordagem, que eles consideraram análoga à atenção de consulta agrupada, permite ainda mais eficiência por meio do paralelismo de tensor em GPUs.

No bloco Mamba-2, que eles chamam de bloco Mamba paralelo (em oposição ao bloco Mamba “sequencial” original), os parâmetros dependentes de entrada B e C são gerados em paralelo na camada de projeção inicial. B e C, especificamente, são derivados simplesmente copiando partes de x_proj, em vez da multiplicação de x_proj por meio de camadas lineares dedicadas. Além de simplificar e reduzir os parâmetros totais do modelo, esse paralelismo permite um treinamento em grande escala consideravelmente mais eficiente.⁴

Tanto o Mamba quanto os transformadores têm seus próprios pontos fortes, mas os modelos baseados em Mamba são geralmente superiores em todas as questões relacionadas ao uso de memória e velocidade: de acordo com o artigo do Mamba, o Mamba oferece uma taxa de transferência cinco vezes maior do que os transformadores equivalentes.

Os transformadores são incrivelmente precisos e versáteis, mas também incrivelmente exigentes em recursos computacionais. Durante o pré-treinamento (e ajuste fino), os requisitos de memória da autoatenção escalam quadraticamente com o comprimento da sequência: se você dobrar o comprimento do contexto de uma sequência, o mecanismo de atenção utilizará o quádruplo dos recursos. Esse "gargalo quadrático" limita cada vez mais a velocidade e a disponibilidade de memória à medida que a janela de contexto cresce. Durante a inferência, suas necessidades de memória são dimensionadas de forma linear.

Durante o treinamento, o uso de memória de um modelo Mamba escala apenas de forma linear durante o treinamento. Mais importante ainda, seu uso de memória durante a inferência é constante: independentemente de quantos tokens o modelo tenha visto, o SSM mantém uma representação de tamanho fixo de seu histórico de entrada. Isso permite um comprimento de contexto teoricamente ilimitado, restrito apenas por limitações de hardware.

Dito isso, o método com uso mais intenso de memória e computacionalmente redundante dos transformadores tem suas próprias vantagens. Por exemplo, a pesquisa demonstrou que os transformadores ainda superam o Mamba e o Mamba-2 em tarefas que exigem aprendizado no contexto (como prompt few-shot), cópia ou raciocínio de contexto longo.

Felizmente as respectivas forças dos transformadores e do Mamba não são mutuamente exclusivas. O artigo Mamba-2 sugere que um modelo híbrido poderia superar tanto os transformadores puros quanto os SSMs, uma noção formalmente validada pela pesquisa da NVIDIA no final de 2024.⁵ Em termos gerais, os modelos híbridos parecem combinar os benefícios de eficiência do Mamba com as nuances e o contexto desempenho de aprendizado proporcionados pelo mecanismo de atenção mais intensivo em recursos dos transformadores.

Para explorar isso ainda mais, a IBM Research trabalhou com Dao e Gu e também com Minjia Zhang, da University of Illinois em Urbana-Champaign (UIUC), no Bamba e Bamba V2. Por sua vez, o Bamba disponibilizou informações para muitos dos elementos arquitetônicos do IBM Granite 4.0.

A pesquisa de modelos híbridos continua sendo uma área de pesquisa ativa, especialmente na comunidade de código aberto.