Mamba: Modelado de Secuencias en Tiempo Lineal con Espacios de Estado Selectivos

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

December 1, 2023

Autores: Albert Gu, Tri Dao

cs.AI

Resumen

Los modelos fundacionales, que ahora impulsan la mayoría de las aplicaciones más destacadas en el aprendizaje profundo, se basan casi universalmente en la arquitectura Transformer y su módulo central de atención. Se han desarrollado muchas arquitecturas de tiempo subcuadrático, como la atención lineal, las convoluciones con compuertas, los modelos recurrentes y los modelos de espacio de estados estructurados (SSM, por sus siglas en inglés), para abordar la ineficiencia computacional de los Transformers en secuencias largas. Sin embargo, no han tenido un rendimiento tan bueno como la atención en modalidades importantes como el lenguaje. Identificamos que una debilidad clave de estos modelos es su incapacidad para realizar razonamientos basados en contenido, y proponemos varias mejoras. En primer lugar, simplemente permitir que los parámetros del SSM sean funciones de la entrada aborda su debilidad con modalidades discretas, lo que permite al modelo propagar o descartar información selectivamente a lo largo de la dimensión de longitud de la secuencia según el token actual. En segundo lugar, aunque este cambio impide el uso de convoluciones eficientes, diseñamos un algoritmo paralelo consciente del hardware en modo recurrente. Integramos estos SSM selectivos en una arquitectura de red neuronal simplificada de extremo a extremo, sin atención ni bloques MLP (Mamba). Mamba disfruta de una inferencia rápida (5 veces mayor rendimiento que los Transformers) y una escalabilidad lineal en la longitud de la secuencia, y su rendimiento mejora en datos reales hasta secuencias de millones de elementos. Como columna vertebral general para modelos de secuencias, Mamba logra un rendimiento de vanguardia en varias modalidades, como lenguaje, audio y genómica. En modelado de lenguaje, nuestro modelo Mamba-3B supera a Transformers del mismo tamaño y se iguala a Transformers del doble de su tamaño, tanto en preentrenamiento como en evaluación en tareas posteriores.

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.