Por que estas notas

Tenho estudado SSM e Mamba para minha dissertação de mestrado, e ando tomando notas. A ideia central com este texto é sedimentar tais leituras.

E por que estudar SSM e Mamba? Por eficiência. O mecanismo de atenção¹ processa uma sequência de comprimento \(L\) comparando cada posição com todas as demais, o que custa \(\mathcal{O}(L^2)\) em tempo e memória e torna proibitivo o processamento de contextos muito longos. Modelos de espaço de estados (SSM, de state space models) oferecem uma alternativa com custo \(\mathcal{O}(L)\), mas historicamente pagavam esse barateamento com perda de capacidade de raciocínio dependente de conteúdo. Os três artigos a seguir contam a história de como essa lacuna foi sendo fechada:

1. Mamba² introduz a seleção — deixar os parâmetros do SSM dependerem da entrada — e um algoritmo de varredura ciente do hardware que mantém o custo linear.
2. Mamba-2³ revela que SSMs e atenção são duas faces da mesma operação com matrizes estruturadas (a dualidade de espaços de estados estruturados, ou SSD), e usa essa equivalência para construir um algoritmo 2 a 8 vezes mais rápido.
3. Mamba-3⁴ refina a recorrência com três princípios extraídos da teoria de SSMs: discretização trapezoidal, estados complexos e uma formulação de múltiplas entradas e saídas (MIMO).

Preliminares: o modelo de espaço de estados

Um SSM, em sua forma contínua, é um sistema linear que mapeia um sinal de entrada \(x(t) \in \mathbb{R}\) em uma saída \(y(t) \in \mathbb{R}\) por meio de um estado latente \(h(t) \in \mathbb{R}^{N}\):

\[ \begin{aligned} h'(t) &= \mathbf{A},h(t) + \mathbf{B},x(t) \\ y(t) &= \mathbf{C},h(t) \end{aligned} \]

onde \(\mathbf{A} \in \mathbb{R}^{N \times N}\) governa a dinâmica do estado, \(\mathbf{B} \in \mathbb{R}^{N \times 1}\) injeta a entrada e \(\mathbf{C} \in \mathbb{R}^{1 \times N}\) projeta o estado na saída. O tamanho \(N\) do estado é a "memória" do modelo: tudo o que a sequência viu até o instante \(t\) precisa caber em \(h(t)\).

Para operar sobre sequências discretas, o sistema é discretizado com um passo \(\Delta\). A discretização por retenção de ordem zero (ZOH, de zero-order hold), que assume a entrada constante dentro de cada intervalo \(\Delta\), produz

\[ \begin{aligned} \bar{\mathbf{A}} &= \exp(\Delta \mathbf{A}) \\ \bar{\mathbf{B}} &= (\Delta \mathbf{A})^{-1}\bigl(\exp(\Delta \mathbf{A}) - \mathbf{I}\bigr)\cdot \Delta \mathbf{B} \end{aligned} \]

de modo que a recorrência discreta se torna

\[ \begin{aligned} h_t &= \bar{\mathbf{A}},h_{t-1} + \bar{\mathbf{B}},x_t \\ y_t &= \mathbf{C},h_t \end{aligned} \]

Há, aqui, uma dualidade que sustenta toda a literatura de SSMs. Se \(\bar{\mathbf{A}}, \bar{\mathbf{B}}, \mathbf{C}\) são fixos — isto é, se o sistema é invariante no tempo (LTI, de linear time-invariant) —, então desenrolar a recorrência mostra que \(y\) é uma convolução de \(x\) com um núcleo fixo:

\[ y_t = \sum_{k=0}^{t} \mathbf{C},\bar{\mathbf{A}}^{,k},\bar{\mathbf{B}};x_{t-k} \quad\Longleftrightarrow\quad y = \bar{\mathbf{K}} \ast x, \qquad \bar{\mathbf{K}} = (\mathbf{C}\bar{\mathbf{B}}, \mathbf{C}\bar{\mathbf{A}}\bar{\mathbf{B}}, \dots, \mathbf{C}\bar{\mathbf{A}}^{L-1}\bar{\mathbf{B}}) \]

Essa equivalência é a chave do desempenho dos modelos S4⁵: pode-se treinar em modo convolucional, paralelizando ao longo de toda a sequência com a transformada de Fourier, e inferir em modo recorrente, gerando um token por vez com estado de tamanho constante. O preço, porém, é que o núcleo \(\bar{\mathbf{K}}\) é o mesmo para toda a sequência: o modelo não consegue tratar um token de forma diferente de outro com base em seu conteúdo. É exatamente essa limitação que o Mamba ataca.

Mamba: a seleção como mecanismo

A limitação dos SSMs invariantes no tempo

O artigo de Gu e Dao² parte de uma observação sobre tarefas sintéticas simples, como cópia seletiva e indução de padrões. Resolver essas tarefas exige ignorar tokens irrelevantes e lembrar os relevantes — ou seja, um comportamento dependente de conteúdo. Um SSM LTI não pode fazer isso: como \(\bar{\mathbf{A}}, \bar{\mathbf{B}}, \mathbf{C}\) não dependem da entrada, o modelo comprime o contexto de maneira uniforme, sem decidir o que reter. A atenção, em contraste, resolve essas tarefas trivialmente, porque compara conteúdos — mas ao custo quadrático.

A pergunta do artigo é, então: é possível dar a um SSM a capacidade de seleção dependente de conteúdo sem renunciar ao custo linear?

O mecanismo de seleção (S6)

A resposta é deixar os parâmetros variarem com a entrada. No SSM seletivo — batizado de S6, por ser um S4 com mecanismo de seleção —, os parâmetros deixam de ser constantes e passam a ser funções do token \(x_t\):

\[ \begin{aligned} \mathbf{B}_t &= \text{Linear}_B(x_t) \\ \mathbf{C}_t &= \text{Linear}_C(x_t) \\ \Delta_t &= \text{softplus}\bigl(\text{Linear}_\Delta(x_t) + \text{bias}\bigr) \end{aligned} \]

A matriz \(\mathbf{A}\) permanece fixa, mas, por ser discretizada por \(\Delta_t\), a transição efetiva \(\bar{\mathbf{A}}_t = \exp(\Delta_t \mathbf{A})\) também passa a depender da entrada. Na prática, \(\mathbf{A}\) é mantida diagonal e parametrizada como \(\mathbf{A} = -\exp(\mathbf{A}_{\log})\), o que garante autovalores reais negativos e, portanto, dinâmica estável e contrativa. A recorrência torna-se variante no tempo:

\[ \begin{aligned} h_t &= \bar{\mathbf{A}}_t,h_{t-1} + \bar{\mathbf{B}}_t,x_t \\ y_t &= \mathbf{C}_t,h_t \end{aligned} \]

O papel de \(\Delta_t\) merece destaque, porque é o coração da seleção. Pode-se lê-lo como uma porta (gate): um \(\Delta_t\) grande faz \(\bar{\mathbf{A}}_t \to 0\) e \(\bar{\mathbf{B}}_t\) dominar, de modo que o estado é "reiniciado" e foca o token atual; um \(\Delta_t\) pequeno faz \(\bar{\mathbf{A}}_t \to \mathbf{I}\) e \(\bar{\mathbf{B}}_t \to 0\), de modo que o token atual é ignorado e o estado é mantido. O modelo aprende, token a token, o quanto prestar atenção a cada entrada — e é nesse sentido que a seleção generaliza os mecanismos de gating de RNNs clássicas.

A varredura paralela ciente do hardware

A seleção tem um custo: ao tornar os parâmetros dependentes da entrada, o modelo perde a propriedade LTI e, com ela, a equivalência com a convolução. Não há mais um núcleo fixo \(\bar{\mathbf{K}}\) a aplicar via Fourier; resta a recorrência, que é sequencial por natureza.

A solução do artigo é calcular a recorrência como uma varredura associativa paralela (parallel scan). A recorrência \(h_t = \bar{\mathbf{A}}_t h_{t-1} + \bar{\mathbf{B}}_t x_t\) é uma operação associativa sobre os pares \((\bar{\mathbf{A}}_t, \bar{\mathbf{B}}_t x_t)\), e varreduras associativas admitem algoritmos paralelos em \(\mathcal{O}(\log L)\) passos⁶. A isso, o artigo soma três decisões de engenharia ciente do hardware:

• Fusão de núcleos (kernel fusion): os parâmetros discretizados são materializados na memória rápida (SRAM) da GPU, e não na memória principal (HBM), evitando o tráfego que dominaria o tempo de execução.
• Varredura sem materializar o estado expandido: o estado de dimensão \(N\) é mantido na SRAM e só a saída é escrita de volta.
• Recomputação: os estados intermediários não são guardados para o passo de retropropagação, mas recomputados, trocando memória por operações — o que é vantajoso em GPUs modernas.

O resultado é um SSM seletivo que escala linearmente em \(L\), gera tokens com estado de tamanho constante e atinge, segundo o artigo, vazão de inferência cerca de cinco vezes maior que a de um Transformer comparável.

O bloco Mamba

A arquitetura final é deliberadamente simples. Em vez de alternar blocos de atenção e blocos de perceptron multicamadas (MLP), como faz o Transformer, o Mamba empilha um único bloco homogêneo. Esse bloco combina, em um caminho, uma projeção que expande a dimensão, uma convolução curta em profundidade (que captura contexto local antes da varredura), uma ativação SiLU e o SSM seletivo; em paralelo, um segundo caminho fornece um gate multiplicativo (no espírito do Gated Linear Unit). A saída dos dois caminhos é combinada e projetada de volta à dimensão original. Sem atenção e sem MLP, todo o processamento da sequência recai sobre o SSM seletivo — o que dá à arquitetura sua economia e sua elegância.

  flowchart TB
    x["Entrada"] --> inproj["Projeção de entrada"]
    inproj -->|"ramo x"| conv["Convolução curta em profundidade"]
    conv --> act["SiLU"]
    act --> ssm["SSM seletivo (S6)"]
    inproj -->|"ramo z (gate)"| gate["SiLU"]
    ssm --> mul["⊙ multiplicação"]
    gate --> mul
    mul --> outproj["Projeção de saída"]
    outproj --> y["Saída"]

Mamba-2: a dualidade de espaços de estados estruturados

O segundo artigo³, de Dao e Gu, é mais teórico e, em certo sentido, mais ambicioso: ele explica por que o Mamba funciona, situando-o num arcabouço que também contém a atenção. A tese, anunciada já no título — "Transformers are SSMs" —, é que as duas arquiteturas são fatorizações diferentes do mesmo objeto matemático.

SSMs como matrizes semisseparáveis

O ponto de partida é escrever o SSM não como recorrência, mas como uma única transformação matricial. Desenrolando a recorrência (com \(\mathbf{A}_t\) variante no tempo), cada saída é

\[ y_j = \sum_{i \le j} \mathbf{C}_j^\top \bigl(\bar{\mathbf{A}}_j \bar{\mathbf{A}}_{j-1}\cdots \bar{\mathbf{A}}_{i+1}\bigr)\bar{\mathbf{B}}_i,x_i \]

ou seja, \(y = \mathbf{M}x\), onde \(\mathbf{M}\) é a matriz triangular inferior de entradas

\[ \mathbf{M}_{ji} = \mathbf{C}_j^\top \bar{\mathbf{A}}_{j:i},\bar{\mathbf{B}}_i, \qquad \bar{\mathbf{A}}_{j:i} = \prod_{k=i+1}^{j}\bar{\mathbf{A}}_k \]

Matrizes dessa forma são chamadas semisseparáveis (mais precisamente, \(N\)-semisseparáveis sequenciais): qualquer submatriz inteiramente contida no triângulo inferior tem posto no máximo \(N\). Essa estrutura é o que permite calcular \(y = \mathbf{M}x\) em tempo linear, apesar de \(\mathbf{M}\) ser uma matriz \(L \times L\).

A forma dual: atenção mascarada

A decisão arquitetural decisiva do Mamba-2 é restringir \(\mathbf{A}_t\) a um escalar vezes a identidade, \(\bar{\mathbf{A}}_t = a_t \mathbf{I}\). Com isso, o produto de transições vira um produto de escalares, e a matriz \(\mathbf{M}\) fatora-se de modo notável:

\[ \mathbf{M}_{ji} = \Bigl(\prod_{k=i+1}^{j} a_k\Bigr),\mathbf{C}_j^\top \mathbf{B}_i = \mathbf{L}_{ji},(\mathbf{C}_j^\top \mathbf{B}_i) \]

onde \(\mathbf{L}_{ji} = \prod_{k=i+1}^{j} a_k\) é uma matriz triangular inferior de produtos cumulativos. Em notação compacta,

\[ \mathbf{M} = \mathbf{L} \circ (\mathbf{C},\mathbf{B}^\top) \]

onde \(\circ\) é o produto de Hadamard. Aqui está a dualidade: \(\mathbf{C}\mathbf{B}^\top\) desempenha o papel de \(\mathbf{Q}\mathbf{K}^\top\) na atenção, e \(\mathbf{L}\) é uma máscara causal — só que, em vez de uma máscara de zeros e uns, é uma máscara de decaimento, cujas entradas são produtos cumulativos de \(a_k\). Calcular \(y = (\mathbf{L} \circ \mathbf{C}\mathbf{B}^\top)x\) diretamente é precisamente uma forma de atenção linear mascarada, com custo quadrático em \(L\). Calcular a mesma coisa pela recorrência é o SSM, com custo linear. São a mesma função, computada por dois caminhos — a dualidade de espaços de estados estruturados (SSD, de structured state space duality).

O algoritmo SSD

A equivalência não é só conceitual: ela rende um algoritmo melhor. A ideia é particionar a sequência em blocos e tratar cada parte da matriz \(\mathbf{M}\) com o método mais conveniente:

• Os blocos diagonais — interações dentro de um mesmo bloco curto — são computados pela forma dual quadrática, que é pequena e aproveita bem as unidades de multiplicação de matrizes (tensor cores) das GPUs.
• Os blocos fora da diagonal — o fluxo de informação entre blocos — são computados pela forma recorrente linear, propagando um estado resumido de bloco a bloco.

Esse algoritmo "por blocos" (chunkwise) combina o melhor dos dois mundos: a paralelização eficiente da atenção dentro de blocos e a escala linear do SSM entre blocos. Na prática, é de 2 a 8 vezes mais rápido que a varredura seletiva do Mamba original e permite estados substancialmente maiores, porque o gargalo de memória da varredura é aliviado. Mantém-se o custo \(\mathcal{O}(L)\), com aproveitamento muito maior do hardware.

  flowchart TB
    seq["Sequência de comprimento L"] --> split["Particiona em blocos"]
    split --> diag["Blocos diagonais
(dentro do bloco):
forma dual quadrática, tipo atenção"]
    split --> off["Blocos fora da diagonal
(entre blocos):
forma recorrente linear, estado resumido"]
    diag --> comb["Combina"]
    off --> comb
    comb --> out["Saída y = M x"]

Mudanças arquiteturais

A restrição de \(\mathbf{A}\) a escalar abre espaço para tratar o SSM como uma forma de atenção multi-cabeça. O Mamba-2 introduz, então, cabeças múltiplas e produz \(\mathbf{B}, \mathbf{C}, \Delta\) em paralelo a partir da entrada (em vez de sequencialmente, como no Mamba-1), o que melhora a paralelização e a interação com o paralelismo de tensores em treinos de grande escala. O resultado é um bloco mais simples e mais escalável, sem perda de qualidade em modelagem de linguagem.

Mamba-3: refinando os princípios de espaço de estados

O terceiro artigo⁴, de Lahoti e colaboradores, mantém o fio recorrente da família Mamba, mas não é apenas um ajuste pontual do bloco anterior. Ele reorganiza a arquitetura em torno de uma receita mais próxima de Llama — alternando blocos Mamba-3 e blocos SwiGLU — e aperfeiçoa o mecanismo de SSM em três frentes, cada uma derivada de um princípio da teoria de espaços de estados. O motivador é a inferência: como o custo de gerar tokens com Transformers é alto, vale a pena extrair mais expressividade de cada unidade de estado de um SSM sem aumentar, ou aumentando muito pouco, a latência de decodificação.

Recorrência trapezoidal

A primeira melhoria é trocar a discretização. O Mamba-3 chama a regra usada por Mamba-1 e Mamba-2 de exponential-Euler: depois de resolver analiticamente a parte exponencial da transição, a integral de entrada é aproximada por uma regra de Euler no endpoint direito. Isso recupera a atualização

\[ h_t = \exp(\Delta_t A_t)h_{t-1} + \Delta_t B_t x_t. \]

O Mamba-3 substitui essa aproximação por uma discretização exponencial-trapezoidal generalizada. A ideia continua sendo causal: a atualização não olha para \(x_{t+1}\). Ela mistura o endpoint anterior e o endpoint atual do intervalo,

\[ \begin{aligned} h_t &= \exp(\Delta_t A_t)h_{t-1} {}+ (1-\lambda_t)\Delta_t\exp(\Delta_t A_t)B_{t-1}x_{t-1} {}+ \lambda_t\Delta_t B_t x_t \\ &= \alpha_t h_{t-1} + \beta_t B_{t-1}x_{t-1} + \gamma_t B_t x_t, \end{aligned} \]

com

\[ \alpha_t = \exp(\Delta_t A_t), \qquad \beta_t = (1-\lambda_t)\Delta_t\exp(\Delta_t A_t), \qquad \gamma_t = \lambda_t\Delta_t. \]

Quando \(\lambda_t = 1\), a regra volta ao caso exponential-Euler. Quando \(\lambda_t = 1/2\), aproxima a regra trapezoidal clássica. No modelo final, porém, \(\lambda_t\) é aprendido como uma porta dependente do token, tipicamente \(\lambda_t = \sigma(u_t)\). Por isso, o ganho não deve ser lido como literalmente "sem parâmetros": há uma parametrização adicional pequena para a porta. O ponto central é outro: a recorrência passa a conter uma convolução causal de largura dois sobre o fluxo de entrada do estado, \(B_t x_t\), sem reintroduzir a convolução curta externa usada em Mamba-2.

Estados complexos e rastreamento de estado

A segunda melhoria ataca uma limitação conhecida dos SSMs com transições reais. Quando a transição efetiva tem autovalores reais não negativos, a dinâmica do estado é essencialmente de escala: cada componente cresce ou decai sem rotação. Isso é útil para memória associativa, mas ruim para rastreamento de estado (state tracking) — tarefas como contar paridade, acompanhar contadores modulares ou rastrear estados que oscilam, em que o modelo precisa de dinâmica periódica, e não de mero decaimento.

O Mamba-3 introduz estados de valor complexo. Autovalores complexos correspondem a dinâmica rotacional: o estado não apenas decai, mas gira, codificando naturalmente amplitude e fase. Na implementação, essa dinâmica pode ser escrita como um SSM real com blocos de rotação \(2 \times 2\), isto é, como uma forma de RoPE dependente dos dados aplicada às projeções \(\mathbf{B}\) e \(\mathbf{C}\). Isso restaura a capacidade de representar padrões periódicos e oscilatórios e, com ela, a competência em tarefas de rastreamento de estado em que os modelos reais falhavam.

Formulação MIMO

A terceira melhoria diz respeito à capacidade por unidade de estado. A forma SSD do Mamba-2, ao restringir \(\mathbf{A}\) a escalar, torna cada atualização efetivamente de uma entrada e uma saída (SISO, de single-input single-output): a interação \(\mathbf{B}_i\) com \(\mathbf{C}_j\) é de posto um por cabeça. O Mamba-3 generaliza para uma formulação de múltiplas entradas e múltiplas saídas (MIMO), em que o estado processa vários canais simultaneamente e as projeções \(\mathbf{B}\) e \(\mathbf{C}\) produzem interações de posto maior.

O efeito prático é aumentar a expressividade sem aumentar o tamanho do estado — e, portanto, sem aumentar muito a latência de decodificação, que é governada em grande parte pelo estado a ser lido e escrito a cada token. Como a decodificação dos SSMs costuma ser limitada por movimentação de memória, a formulação MIMO aumenta a intensidade aritmética: faz mais multiplicações úteis sobre praticamente o mesmo tráfego de estado. O artigo reporta que o Mamba-3 atinge perplexidade comparável à do Mamba-2 usando metade do tamanho de estado, o que expressa esse ganho de densidade informacional.

Síntese

As três melhorias são complementares e, combinadas, deslocam a fronteira em três eixos: recuperação de informação (retrieval), rastreamento de estado e modelagem de linguagem. Em escala de 1,5 bilhão de parâmetros, o artigo reporta ganho de 1,8 ponto percentual de acurácia sobre modelos comparáveis, preservando a eficiência de inferência que é a razão de ser de toda a família. A mensagem é que ainda havia expressividade a extrair do arcabouço de SSMs — bastava revisitar a discretização, o domínio dos autovalores e a estrutura das projeções.

Lidos em sequência, os três artigos formam um arco coerente. O Mamba identificou que a invariância no tempo era a causa da fraqueza dos SSMs em tarefas dependentes de conteúdo e a removeu com o mecanismo de seleção, pagando o preço com um algoritmo de varredura ciente do hardware. O Mamba-2 explicou esse sucesso de um ângulo mais alto, mostrando que SSMs e atenção são fatorizações do mesmo cálculo com matrizes semisseparáveis, e converteu essa dualidade em um algoritmo bem mais rápido. O Mamba-3 voltou aos princípios — discretização, autovalores, posto das projeções — para espremer mais expressividade de cada unidade de estado, sem comprometer a inferência.

O fio condutor é a tensão permanente entre dois custos: o quadrático da atenção, que dá expressividade plena à troca de informação entre posições, e o linear dos SSMs, que comprime tudo num estado de tamanho fixo. A família Mamba é, em larga medida, uma sequência de respostas cada vez mais refinadas à pergunta de quanto da expressividade da atenção se consegue recuperar dentro de um orçamento linear — e cada artigo recupera um pouco mais.

Referências