o problema
O Transformer é eficaz porque a attention roteia informação densamente dentro da janela: cada token vê todos os outros, sem compressão nenhuma. Essa mesma propriedade é o defeito. Nada fora da janela existe para o modelo, o treino escala quadrático no comprimento da sequência e a inferência autoregressiva precisa guardar o contexto inteiro no KV cache. Uma pilha enorme de variantes subquadráticas de attention apareceu para consertar isso, e nenhuma delas se mostrou empiricamente competitiva em escala.
A alternativa mais promissora eram os structured state space models — S4 e derivados. Eles têm estado finito, treinam em tempo linear, geram em tempo constante por passo e dominaram benchmarks de dependência longa. Mas todos eles são LTI: os parâmetros são fixos ao longo do tempo. Isso é uma escolha de engenharia, não de modelagem — só com dinâmica constante dá para calcular o modelo como convolução global e escapar de materializar o estado expandido. O custo aparece em dado discreto e denso de informação, como texto e DNA, onde os SSMs ficavam claramente atrás de attention. O paper nomeia a falha: sem dinâmica dependente do input, o modelo não sabe filtrar. Ele resolve a tarefa de cópia com espaçamento fixo, porque basta contar tempo, e quebra na cópia seletiva, onde o espaçamento é aleatório e é preciso olhar o conteúdo.
a ideia
Tratar modelagem de sequência como compressão de contexto. Attention não comprime e por isso é cara. Recorrência comprime e por isso é barata, mas o resultado depende de quão bem aquele estado finito guardou o que importava. Se o gargalo é a qualidade da compressão, a solução não é estado maior: é o modelo poder decidir, token a token, o que entra no estado e o que é descartado. Os autores chamam isso de seleção, e insistem que não é a mesma coisa que gating arquitetural — multiplicação elementwise que não interage ao longo do eixo da sequência não muda o espaçamento entre tokens e não resolve a cópia seletiva.
como funciona
A mudança é curta. No S4, Δ, A, B e C são parâmetros; no S6, B e C viram projeções lineares do input e Δ vira softplus de um parâmetro somado a uma projeção do input, projetado a uma dimensão pequena e broadcast para os D canais. A ganha só invariância: ele afeta o modelo apenas via Ā = exp(ΔA), então selecionar Δ já basta. Os tensores ganham uma dimensão L e o modelo deixa de ser invariante no tempo — o que elimina a convolução e obriga a computar a recorrência.
Aí entra a parte de sistemas. A recorrência ingênua precisaria materializar em HBM o estado de tamanho (B,L,D,N), N vezes maior que entrada e saída. O kernel fundido lê (Δ, A, B, C) da HBM para a SRAM, discretiza lá dentro, roda um parallel scan associativo lá dentro e escreve de volta só a saída (B,L,D). Os estados intermediários não são salvos: são recomputados no backward. Resultado prático: o scan é 20 a 40 vezes mais rápido que uma implementação padrão em PyTorch, fica à frente do FlashAttention-2 a partir de 2K de sequência, e a camada consome mais ou menos a mesma memória de ativação que uma implementação otimizada de Transformer.
O bloco também simplifica. Em vez de intercalar bloco tipo H3 com bloco MLP, o Mamba funde os dois num bloco só, homogêneo, com fator de expansão E=2 e SiLU, sem attention e sem MLP separado.
o que isso custou
O próprio paper admite três coisas. A primeira é que seleção tem viés indutivo, e ele nem sempre ajuda: em waveform de áudio, trocar S4 por S6 piora — sinal amostrado uniformemente gosta de LTI. Esse foi o único experimento do paper que usou parametrização complexa em vez de real. A segunda é escala: a avaliação para em 1,3B para as scaling laws e 2,8B no downstream, abaixo dos modelos abertos fortes e abaixo dos 7B em que RWKV e RetNet foram medidos, e os autores dizem que escalar SSM pode trazer problemas de engenharia não discutidos ali. A terceira é ecossistema: fine-tuning, in-context learning, quantização, RLHF — tudo isso é conhecido em Transformer e é pergunta aberta em SSM.
Vale também ler a tabela de ablação com atenção. Aumentar a dimensão de estado N só melhora quando B e C são seletivos; com B e C constantes, ir de N=1 a N=16 quase não move a perplexidade.
onde isso aparece hoje
Os números que sustentaram a adoção: no Pile, Mamba-2,8B faz 6,22 de perplexidade contra 6,73 do Pythia-2,8B, e a média de raciocínio de senso comum, 63,3, passa a do Pythia-6,9B. Em induction heads, treinado em 256 tokens, ele generaliza até 2²⁰ — nenhum outro método passa do dobro. Em DNA, a qualidade melhora até 1M de contexto enquanto o HyenaDNA piora.
Código e checkpoints saíram abertos junto com o paper, o que fez do Mamba a referência da linha de SSM seletivo, e os próprios autores voltaram ao tema em 2024 com o Mamba-2. O paper já ablava a arquitetura híbrida: intercalar bloco Mamba com MHA sai levemente melhor que o Mamba puro — uma pista, no próprio texto, do caminho que as arquiteturas híbridas seguiriam.