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

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

December 1, 2023

Autores: Albert Gu, Tri Dao

cs.AI

Resumo

Modelos de base, que agora impulsionam a maioria das aplicações empolgantes em aprendizado profundo, são quase universalmente baseados na arquitetura Transformer e em seu módulo central de atenção. Muitas arquiteturas de tempo subquadrático, como atenção linear, convolução com portas e modelos recorrentes, e modelos de espaço de estados estruturados (SSMs), foram desenvolvidas para abordar a ineficiência computacional dos Transformers em sequências longas, mas elas não performaram tão bem quanto a atenção em modalidades importantes, como linguagem. Identificamos que uma fraqueza chave desses modelos é sua incapacidade de realizar raciocínio baseado em conteúdo, e fazemos várias melhorias. Primeiro, simplesmente permitir que os parâmetros do SSM sejam funções da entrada aborda sua fraqueza com modalidades discretas, permitindo 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, mesmo que essa mudança impeça o uso de convoluções eficientes, projetamos um algoritmo paralelo consciente do hardware no modo recorrente. Integramos esses SSMs seletivos em uma arquitetura de rede neural simplificada de ponta a ponta sem atenção ou mesmo blocos MLP (Mamba). O Mamba desfruta de inferência rápida (5 vezes maior taxa de transferência que Transformers) e escalonamento linear no comprimento da sequência, e seu desempenho melhora em dados reais até sequências de milhões de comprimento. Como um modelo de sequência geral, o Mamba alcança 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 Transformers duas vezes maiores, tanto no pré-treinamento quanto na avaliação downstream.

English

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 (5times 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.