Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu; Tri Dao

Artigos · arXiv 2312.00752

Mamba: Modelagem de Sequências em Tempo Linear com Espaços de Estados Seletivos

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao

Publicado em 2023-12-01cs.LGcs.AI arXiv PDF

Ler com tradução e IA

Resumo em três linhas

🎯 Transformers escalam quadraticamente com o comprimento da sequência e exigem grandes KV caches durante a inferência autorregressiva, enquanto modelos subquadráticos anteriores, como SSMs LTI, têm dificuldades em tarefas discretas de raciocínio de conteúdo, como modelagem de linguagem.
🔧 Torna os parâmetros do SSM dependentes da entrada (seletivos) para filtrar ou propagar informações dinamicamente, combinando isso a um algoritmo de scan associativo paralelo fundido e recomputação ciente de hardware para evitar materializar o estado oculto expandido na memória HBM da GPU.
📈 O Mamba-3B iguala ou supera Transformers com o dobro do seu tamanho em tarefas downstream zero-shot (por exemplo, média de 63,3% vs. 59,1% do Pythia-2.8B e 61,7% do Pythia-6.9B), oferecendo throughput de inferência até 5× maior.

Resumo

Os modelos de fundação, que hoje sustentam a maioria das aplicações mais promissoras em aprendizado profundo, são quase universalmente baseados na arquitetura Transformer e em seu módulo central de atenção. Diversas arquiteturas de tempo subquadrático — como atenção linear, convoluções com gating e modelos recorrentes, e modelos de espaço de estados estruturados (SSMs) — foram desenvolvidas para lidar com a ineficiência computacional dos Transformers em sequências longas; no entanto, elas não alcançaram o mesmo desempenho da atenção em modalidades essenciais, como a linguagem. Identificamos que uma fraqueza fundamental desses modelos é a incapacidade de realizar raciocínio baseado em conteúdo e propomos várias melhorias. Primeiro, simplesmente permitir que os parâmetros do SSM sejam funções da entrada aborda sua fraqueza em modalidades discretas, possibilitando que o modelo propague ou esqueça informações seletivamente ao longo da dimensão do comprimento da sequência, dependendo do token atual. Segundo, embora essa alteração impeça o uso de convoluções eficientes, projetamos um algoritmo paralelo ciente do hardware em modo recorrente. Integramos esses SSMs seletivos em uma arquitetura de rede neural fim a fim simplificada, sem blocos de atenção ou mesmo de MLP (Mamba). O Mamba apresenta inferência rápida (throughput 5 $\times$ maior que o dos Transformers) e escalabilidade linear no comprimento da sequência, e seu desempenho melhora em dados reais com sequências de até um milhão de elementos. Como um backbone genérico de modelagem de sequências, o Mamba atinge desempenho de ponta em várias modalidades, como linguagem, áudio e genômica. Na modelagem de linguagem, nosso modelo Mamba-3B supera Transformers do mesmo tamanho e iguala o desempenho de Transformers com o dobro de seu tamanho, tanto no pré-treinamento quanto na avaliação downstream.

Resumo original (inglês)

Foundation models, now powering most of the exciting applications in deep learning, are almost universally based on the Transformer architecture and its core attention module. Many subquadratic-time architectures such as linear attention, gated convolution and recurrent models, and structured state space models (SSMs) have been developed to address Transformers' computational inefficiency on long sequences, but they have not performed as well as attention on important modalities such as language. We identify that a key weakness of such models is their inability to perform content-based reasoning, and make several improvements. First, simply letting the SSM parameters be functions of the input addresses their weakness with discrete modalities, allowing the model to selectively propagate or forget information along the sequence length dimension depending on the current token. Second, even though this change prevents the use of efficient convolutions, we design a hardware-aware parallel algorithm in recurrent mode. We integrate these selective SSMs into a simplified end-to-end neural network architecture without attention or even MLP blocks (Mamba). Mamba enjoys fast inference (5 $\times$ higher throughput than Transformers) and linear scaling in sequence length, and its performance improves on real data up to million-length sequences. As a general sequence model backbone, Mamba achieves state-of-the-art performance across several modalities such as language, audio, and genomics. On language modeling, our Mamba-3B model outperforms Transformers of the same size and matches Transformers twice its size, both in pretraining and downstream evaluation.

Guia de leitura

Problema e motivação

Os modelos de fundação baseiam-se predominantemente na arquitetura Transformer e no mecanismo de autoatenção. Embora a atenção densa direcione informações de forma eficaz ao longo de uma janela de contexto sequencial, ela impõe duas limitações fundamentais: custo computacional quadrático $O (L^{2})$ em relação ao comprimento da sequência $L$ e a necessidade de armazenar chaves e valores passados (o KV cache), o que resulta em inferência com tempo linear por passo e gargalos de memória.

Abordagens subquadráticas anteriores — como atenção linear, arquiteturas recorrentes e modelos de espaço de estados estruturados (SSMs), como o S4 — escalam de maneira linear ou pseudolinear no tempo. No entanto, os SSMs estruturados historicamente não conseguiram atingir a qualidade dos Transformers em modalidades discretas e densas em informação, como texto. Os autores identificam a causa raiz: os SSMs anteriores são lineares e invariantes no tempo (LTI). Como suas dinâmicas de transição não se alteram em função do conteúdo da entrada, eles não conseguem comprimir o contexto de forma seletiva, filtrar ruídos irrelevantes ou modular a retenção de informações através de espaçamentos variáveis.

Ideias centrais

Mecanismo de Seleção como Compressão de Contexto: A modelagem de sequências requer a compressão do histórico em um estado. SSMs LTI anteriores usam parâmetros de transição fixos, forçando-os a tratar todos os timesteps uniformemente. Ao parametrizar as matrizes do SSM com base no token de entrada atual, o modelo adquire a percepção de conteúdo necessária para filtrar ou memorizar informações seletivamente.
Superando o Gargalo da Não Convolução: Tornar os parâmetros dependentes da entrada quebra a equivalência matemática entre recorrências lineares e convoluções globais. Sem convoluções, uma recorrência ingênua exigiria materializar um estado latente expandido de tamanho $(B, L, D, N)$ na memória de alta largura de banda (HBM) da GPU. O Mamba resolve isso fundindo discretização, scan associativo paralelo e projeção de saída na rápida memória SRAM on-chip da GPU, mantendo o tráfego de HBM em níveis mínimos.
Arquitetura Unificada: O Mamba substitui blocos intercalados de Multi-Head Attention (MHA) e MLP por um único bloco homogêneo que combina o design com portas (gating) do H3 com o MLP com gating no estilo SwiGLU.

Método

Modelos de espaço de estados contínuos mapeiam uma entrada $x (t) \in R$ para $y (t) \in R$ por meio de um estado latente implícito $h (t) \in R^{N}$ : $h^{'} (t) = A h (t) + B x (t), y (t) = C h (t)$

Utilizando a discretização por retenção de ordem zero (ZOH) com tamanho de passo $Δ$ , o sistema contínuo se transforma em parâmetros de recorrência discretos: $\overline{A} = exp (Δ A), \overline{B} = (Δ A)^{- 1} (exp (Δ A) - I) \cdot Δ B$ $h_{t} = \overline{A} h_{t - 1} + \overline{B} x_{t}, y_{t} = C h_{t}$

O Mecanismo de Seleção (S6)

Em vez de manter $(Δ, B, C)$ constantes ao longo do tempo, o Mamba os torna dependentes da entrada:

$B_{t} = s_{B} (x_{t}) = Linear_{N} (x_{t})$
$C_{t} = s_{C} (x_{t}) = Linear_{N} (x_{t})$
$Δ_{t} = τ_{Δ} (Parameter + s_{Δ} (x_{t}))$ , com $τ_{Δ} = softplus$ e $s_{Δ} (x) = Broadcast_{D} (Linear_{1} (x))$

Isso generaliza o gating clássico de RNNs: quando $N = 1, A = - 1, B = 1$ , a recorrência se simplifica diretamente para uma atualização com porta $h_{t} = (1 - g_{t}) h_{t - 1} + g_{t} x_{t}$ , onde $g_{t} = σ (Linear (x_{t}))$ .

Expansão de Estado Ciente de Hardware

Para manter alto throughput sem materializar o estado $(B, L, D, N)$ na HBM:

Os parâmetros $(Δ, A, B, C)$ são carregados da HBM diretamente para a SRAM on-chip.
A discretização em $(\overline{A}, \overline{B})$ e um scan associativo paralelo com eficiência de trabalho são computados inteiramente dentro da SRAM.
O estado latente é contraído com $C$ para o formato $(B, L, D)$ antes de ser gravado de volta na HBM.
Os estados intermediários são recomputados na SRAM durante o passo para trás (backward pass) em vez de serem salvos na HBM.

Resultados

Extrapolação Sintética: Na tarefa de Induction Heads, o Mamba treinado com comprimento de sequência 256 extrapola perfeitamente (100% de acurácia) até o comprimento $2^{20} = 1, 048, 576$ ( $4000 \times$ o comprimento da sequência), enquanto as baselines de atenção falham ou sofrem OOM após $2^{14}$ . No Selective Copying, o Mamba atinge 99,8% de acurácia em comparação com 57,0% do H3 e 18,3% do S4.
Modelagem de Linguagem: Avaliado sob a lei de escala de Chinchilla com até 1,3B parâmetros no The Pile, o Mamba iguala ou supera o Transformer++ (embeddings rotativos, SwiGLU, RMSNorm). Em tarefas downstream zero-shot (ex.: LAMBADA, PIQA, HellaSwag, WinoGrande, ARC), o Mamba-2.8B atinge uma acurácia média de 63,3%, superando o Pythia-2.8B (59,1%), o Pythia-6.9B (61,7%) e igualando ou superando modelos de 6B–7B parâmetros.
Genômica e Áudio: Na classificação de DNA de grandes primatas com comprimento de sequência de 1M ( $2^{20}$ ), o Mamba atinge acurácia de 71,67% (1,4M de parâmetros) e 81,31% (7M de parâmetros) em comparação com 54,87% do HyenaDNA. Na geração de fala com SC09, o Mamba reduz o FID de 1,99 (SaShiMi) para 0,94 (6,1M de parâmetros) e 0,67 (24,3M de parâmetros).
Eficiência: O kernel fundido de scan executa de 20 a 40× mais rápido do que implementações padrão de scan em PyTorch e funciona até 7× mais rápido do que o FlashAttention-2 em contextos de 32K. O throughput de inferência é até 5× maior do que o de Transformers comparáveis.

Limitações e questões em aberto

Desempenho em Sinais Contínuos: Embora a seleção dependente da entrada melhore o raciocínio discreto (texto, DNA), ablações no pré-treinamento com formas de onda de áudio bruto mostram que modelos puramente lineares e invariantes no tempo (LTI, como o S4 padrão) apresentam desempenho superior em dados contínuos amostrados uniformemente.
Escala do Modelo: Os modelos avaliados limitam-se a cerca de 3B parâmetros e 300B tokens de pré-treinamento; o desempenho empírico em escalas de 7B–70B+ parâmetros e além ainda requer investigações adicionais.
Ecossistema Downstream: Os Transformers possuem um ecossistema consolidado de ferramentas para engenharia de prompts, manipulação in-context, quantização de KV-cache e ajuste fino com adaptadores. Entender como os SSMs se adaptam a esses modos de interação permanece como um campo de exploração ativo.

Termos-chave

Structured State Space Models (SSMs): Uma classe de modelos de sequência profundos derivados de sistemas dinâmicos lineares contínuos que mapeiam entradas unidimensionais por meio de um estado latente implícito.
Linear Time Invariance (LTI): Uma propriedade matemática na qual os parâmetros de transição do sistema permanecem estáticos em todos os timesteps, viabilizando representações convolucionais globais rápidas, mas restringindo o raciocínio dependente de conteúdo.
Selection Mechanism: A parametrização da dinâmica de transição do SSM (Δ, B, C) como funções do token de entrada, permitindo ao modelo propagar, redefinir ou ignorar dinamicamente o contexto ao longo do comprimento da sequência.
S6: Abreviação para modelos de espaço de estados estruturados seletivos computados por meio de um scan paralelo associativo em vez de convolução.
Parallel Associative Scan: Um algoritmo eficiente de cálculo de prefixos que avalia a recorrência linear paralelamente ao longo de comprimentos de sequência em hardwares aceleradores.
Hardware-Aware State Expansion: Uma estratégia de fusão de kernels e recomputação em GPU que realiza a discretização e o scan na rápida memória SRAM on-chip, evitando o tráfego de ida e volta de grandes estados intermediários para a memória de alta largura de banda (HBM).
Induction Heads: Um mecanismo de recuperação associativa sintética que avalia se um modelo é capaz de reproduzir um token observado anteriormente após seu token de contexto predecessor.
Mamba Block: Um bloco homogêneo e simplificado para modelos de sequência que integra projeções de entrada, convolução 1D, SSMs seletivos e gating multiplicativo, sem camadas separadas de atenção ou MLP.

Perguntas

Por que os SSMs seletivos não podem ser treinados usando convoluções FFT rápidas?

Convoluções globais rápidas via FFT exigem que a dinâmica do sistema (matrizes A, B e C) seja constante ao longo do tempo (invariância linear no tempo). Como o Mamba torna B, C e Δ funções do token de entrada a cada timestep, o kernel de convolução varia em cada posição, rompendo a equivalência convolucional e exigindo, em seu lugar, scans paralelos recorrentes.

Como o Mamba evita a explosão de memória da expansão do estado recorrente durante o treinamento?

Uma recorrência ingênua com dimensão de expansão de estado N exigiria materializar um estado intermediário de formato (B, L, D, N) na HBM da GPU. O Mamba utiliza fusão de kernels para carregar apenas as entradas de baixa dimensão na SRAM, computa a discretização e o scan paralelo dentro da rápida SRAM, grava de volta na HBM apenas a saída contraída de formato (B, L, D) e recomputa as ativações intermediárias durante a retropropagação.

Qual é a relação entre o parâmetro Δ do Mamba e o gating recorrente?

O parâmetro de tamanho de passo Δ atua como um mecanismo de gating generalizado. Quando discretizado sob retenção de ordem zero (ZOH) com N=1 e A=-1, a recorrência seletiva se reduz matematicamente a uma porta de RNN clássica: um Δ grande faz o modelo focar na entrada atual e redefinir o estado, enquanto um Δ pequeno preserva o histórico de estados e ignora a entrada.

O Mamba utiliza KV cache durante a geração?

Não. Como o Mamba opera como uma recorrência de espaço de estados autorregressiva durante a geração, ele mantém apenas um estado oculto de tamanho fixo (dimensão D × N por camada) em vez de armazenar em cache chaves e valores para todos os tokens passados, resultando em inferência com tempo constante O(1) por passo.

A explicação desta página foi escrita por IA a partir do artigo e pode conter erros. Consulte o artigo para os detalhes.

Ler com tradução e IA