Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu; Tri Dao

Artículos · arXiv 2312.00752

Mamba: Modelado de secuencias en tiempo lineal con espacios de estados selectivos

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao

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

Leer con traducción e IA

Resumen en tres líneas

🎯 Los Transformers escalan de forma cuadrática con la longitud de la secuencia y requieren grandes cachés KV durante la inferencia autorregresiva, mientras que los modelos subcuadráticos previos, como los modelos de espacio de estados LTI, tienen dificultades en tareas discretas de razonamiento basado en contenido como el modelado del lenguaje.
🔧 Hace que los parámetros del SSM dependan de la entrada (selectivos) para filtrar o propagar información dinámicamente, y combina esto con un algoritmo de escaneo asociativo paralelo fusionado y de recomputación consciente del hardware para evitar materializar el estado oculto expandido en la HBM de la GPU.
📈 Mamba-3B iguala o supera a Transformers con el doble de su tamaño en tareas zero-shot downstream (p. ej., un promedio de 63.3% frente al 59.1% de Pythia-2.8B y el 61.7% de Pythia-6.9B), al tiempo que ofrece un throughput de inferencia hasta 5× mayor.

Resumen

Los modelos fundacionales, que ahora impulsan la mayoría de las aplicaciones más destacadas en el aprendizaje profundo, se basan de forma casi universal en la arquitectura Transformer y en 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 y modelos recurrentes, y los modelos de espacio de estados estructurados (SSM), para abordar la ineficiencia computacional de los Transformers en secuencias largas; sin embargo, no han alcanzado un rendimiento equiparable al de la atención en modalidades cruciales como el lenguaje. Identificamos que una debilidad clave de tales modelos es su incapacidad para realizar razonamiento basado en el contenido, y aportamos varias mejoras. En primer lugar, permitir que los parámetros del SSM sean funciones de la entrada resuelve directamente su debilidad con modalidades discretas, lo que permite al modelo propagar u olvidar información selectivamente a lo largo de la dimensión de la secuencia en función del token actual. En segundo lugar, a pesar de que este cambio impide el uso de convoluciones eficientes, diseñamos un algoritmo paralelo optimizado para el 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 ofrece una inferencia rápida (un rendimiento o throughput 5 $\times$ superior al de los Transformers) y un escalado lineal respecto a la longitud de la secuencia, y su rendimiento mejora en datos reales con secuencias de hasta un millón de elementos. Como arquitectura troncal de secuencia general, Mamba logra un rendimiento de vanguardia en diversas modalidades como lenguaje, audio y genómica. En el modelado del lenguaje, nuestro modelo Mamba-3B supera a los Transformers del mismo tamaño e iguala a Transformers con el doble de parámetros, tanto en el preentrenamiento como en la evaluación downstream.

Resumen 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.

Guía de lectura

Problema y motivación

Los modelos fundacionales dependen de forma abrumadora de la arquitectura Transformer y de la autoatención. Aunque la atención densa enruta la información de manera eficaz a lo largo de la ventana de contexto de una secuencia, conlleva dos limitaciones fundamentales: un costo computacional cuadrático $O (L^{2})$ con respecto a la longitud de la secuencia $L$ , y el requisito de almacenar claves y valores pasados (la caché KV), lo que conduce a una inferencia en tiempo lineal por paso y a cuellos de botella de memoria.

Los enfoques subcuadráticos previos —como la atención lineal, las arquitecturas recurrentes y los modelos de espacio de estados estructurados (SSM) como S4— escalan de forma lineal o pseudolineal en el tiempo. Sin embargo, históricamente los SSM estructurados no han logrado igualar la calidad de los Transformers en modalidades discretas y densas en información, como el texto. Los autores identifican la causa raíz: los SSM previos son lineales e invariantes en el tiempo (LTI). Dado que su dinámica de transición no cambia según el contenido de la entrada, no pueden comprimir el contexto de forma selectiva, filtrar el ruido irrelevante ni modular la retención ante un espaciado variable.

Ideas clave

El mecanismo de selección como compresión de contexto: El modelado de secuencias requiere comprimir el historial en un estado. Los SSM LTI previos utilizan parámetros de transición fijos, lo que los obliga a tratar todos los pasos temporales de manera uniforme. Al parametrizar las matrices del SSM en función del token de entrada actual, el modelo adquiere la conciencia del contenido necesaria para filtrar o recordar información de manera selectiva.
Superación del cuello de botella no convolutivo: Hacer que los parámetros dependan de la entrada rompe la equivalencia matemática entre las recurrencias lineales y las convoluciones globales. Sin convoluciones, una recurrencia ingenua requeriría materializar un estado latente expandido de tamaño $(B, L, D, N)$ en la memoria de alto ancho de banda (HBM) de la GPU. Mamba resuelve esto fusionando la discretización, el escaneo asociativo paralelo y la proyección de salida dentro de la rápida memoria SRAM en chip de la GPU, manteniendo al mínimo el tráfico con la HBM.
Arquitectura unificada: Mamba reemplaza los bloques intercalados de atención multicabezal (MHA) y MLP con un único bloque homogéneo que combina el diseño con compuertas de H3 con el MLP con compuertas de estilo SwiGLU.

Método

Los modelos continuos de espacio de estados asignan una entrada $x (t) \in R$ a $y (t) \in R$ a través de un estado latente implícito $h (t) \in R^{N}$ : $h^{'} (t) = A h (t) + B x (t), y (t) = C h (t)$

Utilizando la discretización por retención de orden cero (ZOH) con un tamaño de paso $Δ$ , el sistema continuo se transforma en parámetros de recurrencia 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}$

El mecanismo de selección (S6)

En lugar de mantener $(Δ, B, C)$ constantes a lo largo del tiempo, Mamba hace que dependan de la 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}))$ , con $τ_{Δ} = softplus$ y $s_{Δ} (x) = Broadcast_{D} (Linear_{1} (x))$

Esto generaliza el mecanismo de compuertas de las RNN clásicas: cuando $N = 1, A = - 1, B = 1$ , la recurrencia se simplifica directamente a una actualización con compuertas $h_{t} = (1 - g_{t}) h_{t - 1} + g_{t} x_{t}$ , donde $g_{t} = σ (Linear (x_{t}))$ .

Expansión de estados consciente del hardware

Para mantener un alto rendimiento de procesamiento sin materializar el estado $(B, L, D, N)$ en HBM:

Los parámetros $(Δ, A, B, C)$ se cargan desde la HBM directamente a la SRAM en chip.
La discretización en $(\overline{A}, \overline{B})$ y un escaneo asociativo paralelo eficiente en trabajo se calculan íntegramente dentro de la SRAM.
El estado latente se contrae con $C$ a la forma $(B, L, D)$ antes de escribirse de nuevo en la HBM.
Los estados intermedios se recalculan en la SRAM durante la pasada hacia atrás (backward pass) en lugar de guardarse en la HBM.

Resultados

Extrapolación sintética: En la tarea de Induction Heads, Mamba entrenado con una longitud de secuencia de 256 extrapola perfectamente (100% de precisión) hasta una longitud de $2^{20} = 1, 048, 576$ ( $4000 \times$ la longitud de la secuencia), mientras que las líneas base de atención fallan o sufren errores por falta de memoria (OOM) más allá de $2^{14}$ . En Selective Copying, Mamba alcanza un 99.8% de precisión en comparación con el 57.0% de H3 y el 18.3% de S4.
Modelado del lenguaje: Evaluado bajo el escalado de Chinchilla con hasta 1.3B de parámetros en the Pile, Mamba iguala o supera a Transformer++ (rotary embeddings, SwiGLU, RMSNorm). En tareas downstream zero-shot (p. ej., LAMBADA, PIQA, HellaSwag, WinoGrande, ARC), Mamba-2.8B alcanza una precisión media del 63.3%, superando a Pythia-2.8B (59.1%), Pythia-6.9B (61.7%) e igualando o superando a modelos de 6B–7B parámetros.
Genómica y audio: En la clasificación de ADN de Great Apes con una longitud de secuencia de 1M ( $2^{20}$ ), Mamba logra una precisión del 71.67% (1.4M de parámetros) y del 81.31% (7M de parámetros) frente al 54.87% de HyenaDNA. En la generación de voz con SC09, Mamba reduce el FID de 1.99 (SaShiMi) a 0.94 (6.1M de parámetros) y a 0.67 (24.3M de parámetros).
Eficiencia: El kernel de escaneo fusionado se ejecuta de 20 a 40 veces más rápido que las implementaciones de escaneo estándar de PyTorch y es hasta 7 veces más rápido que FlashAttention-2 con un contexto de 32K. El throughput de inferencia es hasta 5 veces mayor que el de Transformers comparables.

Limitaciones y preguntas abiertas

Rendimiento en señales continuas: Aunque la selección dependiente de la entrada mejora el razonamiento discreto (texto, ADN), las ablaciones en el preentrenamiento de formas de onda de audio sin procesar muestran que los modelos puramente lineales e invariantes en el tiempo (LTI, como el S4 estándar) funcionan mejor en datos continuos y muestreados uniformemente.
Escala del modelo: Los modelos evaluados tienen un tope de alrededor de 3B parámetros y 300B tokens de preentrenamiento; el rendimiento empírico a escalas de 7B–70B+ parámetros y más allá requiere mayor investigación.
Ecosistema downstream: Los Transformers cuentan con un amplio conjunto de herramientas para ingeniería de prompts, manipulación en contexto, cuantización de caché KV y ajuste fino con adaptadores. Comprender cómo se adaptan los SSM a través de estos modos de interacción sigue siendo un campo de exploración activo.

Términos clave

Structured State Space Models (SSMs): Clase de modelos de secuencias profundos derivados de sistemas dinámicos lineales continuos que transforman entradas 1D a través de un estado latente implícito.
Linear Time Invariance (LTI): Propiedad matemática en la que los parámetros de transición del sistema permanecen estáticos a lo largo de todos los pasos temporales, lo que permite representaciones convolucionales globales rápidas pero limita el razonamiento dependiente del contenido.
Selection Mechanism: Parametrización de la dinámica de transición de los SSM (Δ, B, C) como funciones del token de entrada, lo que permite al modelo propagar, reiniciar o ignorar contexto dinámicamente a lo largo de la secuencia.
S6: Abreviatura de modelos de espacio de estados estructurados selectivos calculados mediante un escaneo paralelo asociativo en lugar de una convolución.
Parallel Associative Scan: Algoritmo de cálculo de prefijos eficiente en trabajo que evalúa la recurrencia lineal a lo largo de secuencias en paralelo sobre hardware acelerador.
Hardware-Aware State Expansion: Estrategia de fusión de kernels de GPU y recomputación que realiza la discretización y el escaneo en la rápida memoria SRAM en chip para evitar el trasvase de estados intermedios grandes hacia la memoria de alto ancho de banda (HBM).
Induction Heads: Mecanismo sintético de recuerdo asociativo que evalúa si un modelo puede reproducir un token observado previamente tras la aparición de su token de contexto predecesor.
Mamba Block: Bloque simplificado y homogéneo de modelado de secuencias que integra proyecciones de entrada, convolución 1D, SSM selectivos y compuertas multiplicativas sin requerir capas separadas de atención o MLP.

Preguntas

¿Por qué los SSM selectivos no pueden entrenarse utilizando convoluciones rápidas por FFT?

Las convoluciones globales rápidas mediante FFT requieren que la dinámica del sistema (las matrices A, B y C) sea constante a lo largo del tiempo (invarianza lineal en el tiempo). Dado que Mamba hace que B, C y Δ sean funciones del token de entrada en cada paso temporal, el kernel de convolución varía en cada posición, lo que rompe la equivalencia convolucional y exige utilizar escaneos paralelos recurrentes en su lugar.

¿Cómo evita Mamba la explosión de memoria provocada por la expansión del estado recurrente durante el entrenamiento?

Una recurrencia ingenua con una dimensión de expansión de estado N requeriría materializar un estado intermedio de tamaño (B, L, D, N) en la HBM de la GPU. Mamba utiliza la fusión de kernels para cargar únicamente entradas de baja dimensión en la SRAM, calcula la discretización y el escaneo paralelo dentro de la rápida SRAM, escribe de vuelta en la HBM únicamente la salida contraída (B, L, D) y recalcula las activaciones intermedias durante la retropropagación.

¿Cuál es la relación entre el parámetro Δ de Mamba y el mecanismo de compuertas recurrente?

El parámetro de tamaño de paso Δ actúa como un mecanismo de compuerta generalizado. Cuando se discretiza bajo retención de orden cero (ZOH) con N=1 y A=-1, la recurrencia selectiva se reduce matemáticamente a una compuerta clásica de RNN: un valor de Δ grande provoca que el modelo se enfoque en la entrada actual y reinicie el estado, mientras que un Δ pequeño preserva el estado histórico e ignora la entrada.

¿Utiliza Mamba una caché KV durante la generación?

No. Dado que Mamba opera como una recurrencia de espacio de estados autorregresiva durante la generación, solo mantiene un estado oculto de tamaño fijo (dimensión D × N por capa) en lugar de almacenar claves y valores en caché para todos los tokens pasados, lo que da como resultado una inferencia en tiempo constante O(1) por paso.

La explicación de esta página la escribió una IA a partir del artículo y puede contener errores. Consulte el artículo para los detalles.

Leer con traducción e IA