El mundo de la IA tiene un nuevo favorito, y no es un Transformer

Tabla de enlaces

Resumen y 1. Introducción

2 modelos de espacio de estados

3 Modelos de espacio de estados selectivos y 3.1 Motivación: la selección como medio de compresión

3.2 Mejora de los SSM con selección

3.3 Implementación eficiente de mecanismos de estabilización selectiva

3.4 Una arquitectura SSM simplificada

3.5 Propiedades de los mecanismos de selección

3.6 Detalles adicionales del modelo

4 Evaluación empírica y 4.1 Tareas sintéticas

4.2 Modelado del lenguaje

4.3 Modelado de ADN

4.4 Modelado y generación de audio

4.5 Puntos de referencia de velocidad y memoria

4.6 Ablaciones de modelos

5 Discusión

6 Conclusión, agradecimientos y referencias

Una discusión: Mecanismo de selección

B Trabajo relacionado y B.1 S4 Variantes y derivados

B.2 Arquitecturas SSM

B.3 Relación con las RNN

B.4 Atención lineal y B.5 Modelos de contexto largo

C Mecánica de los SSM selectivos

Algoritmo D basado en hardware para SSM selectivos

E Detalles experimentales y resultados adicionales y E.1 Tareas sintéticas

E.2 Modelado del lenguaje

E.3 Modelado de ADN

E.4 Detalles de audio

E.5 Indicador de eficiencia

Abstracto

Los modelos fundamentales, que actualmente impulsan la mayoría de las aplicaciones más interesantes del aprendizaje profundo, se basan casi universalmente en la arquitectura Transformer y su módulo de atención. Se han desarrollado numerosas arquitecturas de tiempo subcuadrático, como la atención lineal, los modelos de convolución controlada y recurrentes, y los modelos de espacio de estados estructurados (SSM), para abordar la ineficiencia computacional de los Transformers en secuencias largas, pero su rendimiento no ha sido tan bueno como el de la atención en modalidades importantes como el lenguaje. Identificamos que una debilidad clave de estos modelos es su incapacidad para realizar razonamiento basado en contenido e implementamos varias mejoras. En primer lugar, simplemente permitiendo que los parámetros del SSM sean funciones de la entrada, se soluciona su debilidad con modalidades discretas, lo que permite que el modelo propague u omita información selectivamente a lo largo de la longitud de la secuencia, dependiendo del token actual. En segundo lugar, aunque este cambio impide el uso de convoluciones eficientes, diseñamos un algoritmo paralelo con reconocimiento de hardware en modo recurrente. Integramos estos SSM selectivos en una arquitectura simplificada de red neuronal de extremo a extremo sin atención ni bloques MLP (Mamba). Mamba ofrece una inferencia rápida (cinco veces mayor rendimiento que Transformers) y un escalado lineal en la longitud de secuencia, y su rendimiento mejora con datos reales hasta secuencias de un millón de longitud. Como eje central de un modelo de secuencia general, Mamba alcanza un rendimiento de vanguardia en diversas modalidades, como lenguaje, audio y genómica. En el modelado del lenguaje, nuestro modelo Mamba-3B supera a Transformers del mismo tamaño y se iguala a Transformers del doble de tamaño, tanto en el preentrenamiento como en la evaluación posterior.

1 Introducción

Los modelos de base (MF), o modelos de gran tamaño preentrenados con datos masivos y adaptados para tareas posteriores, se han convertido en un paradigma eficaz en el aprendizaje automático moderno. La columna vertebral de estos MF son a menudo modelos de secuencia, que operan con secuencias arbitrarias de entradas provenientes de una amplia variedad de dominios, como el lenguaje, las imágenes, el habla, el audio, las series temporales y la genómica (Brown et al., 2020; Dosovitskiy et al., 2020; Ismail Fawaz et al., 2019; Oord et al., 2016; Poli et al., 2023; Sutskever, Vinyals y Quoc V. Le, 2014). Aunque este concepto es agnóstico a una elección particular de arquitectura del modelo, los FM modernos se basan predominantemente en un solo tipo de modelo de secuencia: el Transformador (Vaswani et al. 2017) y su capa de atención central (Bahdanau, Cho y Bengio 2015). La eficacia de la autoatención se atribuye a su capacidad para enrutar información densamente dentro de una ventana de contexto, lo que le permite modelar datos complejos. Sin embargo, esta propiedad conlleva desventajas fundamentales: una incapacidad para modelar nada fuera de una ventana finita y una escala cuadrática con respecto a la longitud de la ventana. Ha aparecido una enorme cantidad de investigación sobre variantes más eficientes de la atención para superar estas desventajas (Tay, Dehghani, Bahri, et al. 2022), pero a menudo a expensas de las mismas propiedades que la hacen efectiva. Hasta el momento, ninguna de estas variantes ha demostrado ser empíricamente efectiva a escala en todos los dominios. Recientemente, los modelos de secuencias de espacio de estados estructurados (SSM) (Gu, Goel y Ré, 2022; Gu, Johnson, Goel et al., 2021) han emergido como una clase prometedora de arquitecturas para el modelado de secuencias. Estos modelos pueden interpretarse como una combinación de redes neuronales recurrentes (RNN) y redes neuronales convolucionales (CNN), inspirados en los modelos clásicos de espacio de estados (Kalman, 1960). Esta clase de modelos puede calcularse de forma muy eficiente, ya sea como recurrencia o convolución, con escalamiento lineal o casi lineal en la longitud de la secuencia. Además, cuentan con mecanismos fundamentales para modelar dependencias de largo alcance (Gu, Dao et al., 2020) en ciertas modalidades de datos, y han dominado los benchmarks como el Long Range Arena (Tay, Dehghani, Abnar et al., 2021). Muchas variantes de SSM (Gu, Goel y Ré, 2022; Gu, Gupta et al., 2022; Gupta, Gu y Berant, 2022; Y. Li et al., 2023; Ma et al., 2023; Orvieto et al., 2023; Smith, Warrington y Linderman, 2023) han tenido éxito en dominios que involucran datos de señales continuas, como audio y visión (Goel et al., 2022; Nguyen, Goel et al., 2022; Saon, Gupta y Cui, 2023). Sin embargo, han sido menos eficaces en el modelado de datos discretos y con alta densidad de información, como el texto.

Proponemos una nueva clase de modelos de espacio de estados selectivos que mejoran el trabajo previo en varios ejes para lograr el poder de modelado de los transformadores mientras se escala linealmente en la longitud de la secuencia.

Mecanismo de Selección. En primer lugar, identificamos una limitación clave de los modelos anteriores: la capacidad de seleccionar datos eficientemente según la entrada (es decir, centrarse en entradas específicas o ignorarlas). Basándonos en la intuición basada en tareas sintéticas importantes, como la copia selectiva y las cabezas de inducción, diseñamos un mecanismo de selección simple mediante la parametrización de los parámetros del SSM en función de la entrada. Esto permite al modelo filtrar información irrelevante y recordar la información relevante indefinidamente.

Algoritmo basado en hardware . Este simple cambio supone un reto técnico para el cálculo del modelo; de hecho, todos los modelos SSM anteriores deben ser invariantes en el tiempo y en la entrada para ser computacionalmente eficientes. Esto se soluciona con un algoritmo basado en hardware que calcula el modelo recurrentemente mediante un escaneo en lugar de convolución, pero no materializa el estado expandido para evitar el acceso de E/S entre diferentes niveles de la jerarquía de memoria de la GPU. La implementación resultante es más rápida que los métodos anteriores, tanto en teoría (escalando linealmente en longitud de secuencia, en comparación con el pseudo-lineal para todos los SSM basados en convolución) como en hardware moderno (hasta tres veces más rápido en GPU A100).

Arquitectura. Simplificamos las arquitecturas de modelos de secuencias profundas previas combinando el diseño de arquitecturas SSM previas (Dao, Fu, Saab, et al., 2023) con el bloque MLP de Transformadores en un solo bloque, lo que resulta en un diseño de arquitectura simple y homogéneo (Mamba) que incorpora espacios de estados selectivos. Los SSM selectivos, y por extensión la arquitectura Mamba, son modelos completamente recurrentes con propiedades clave que los hacen idóneos como base de modelos de base general que operan en secuencias. (i) Alta calidad: la selectividad ofrece un rendimiento sólido en modalidades densas como el lenguaje y la genómica. (ii) Entrenamiento e inferencia rápidos: el cálculo y la memoria escalan linealmente en la longitud de la secuencia durante el entrenamiento, y el desarrollo autorregresivo del modelo durante la inferencia requiere solo un tiempo constante por paso, ya que no requiere una caché de elementos previos. (iii) Contexto extenso: la calidad y la eficiencia, en conjunto, producen mejoras de rendimiento en datos reales hasta una longitud de secuencia de 1 M.

Validamos empíricamente el potencial de Mamba como columna vertebral de FM de secuencia general, tanto en la calidad previa al entrenamiento como en el desempeño de tareas específicas del dominio, en varios tipos de modalidades y entornos:

• Sintéticos. En tareas sintéticas importantes, como la copia y la inducción de cabezas, que se han propuesto como clave para grandes modelos lingüísticos, Mamba no solo las resuelve con facilidad, sino que también puede extrapolar soluciones con una longitud indefinida (más de 1 millón de tokens).

• Audio y genómica. Mamba supera a modelos de vanguardia previos, como SaShiMi, Hyena y Transformers, en el modelado de formas de onda de audio y secuencias de ADN, tanto en calidad previa al entrenamiento como en métricas posteriores (p. ej., reduciendo la FID en más de la mitad en un conjunto de datos de generación de voz complejo). En ambos entornos, su rendimiento mejora con contextos más largos, hasta secuencias de un millón de longitudes.

Modelado del lenguaje. Mamba es el primer modelo de secuencia de tiempo lineal que realmente alcanza el rendimiento de un Transformer, tanto en la perplejidad previa al entrenamiento como en las evaluaciones posteriores. Con leyes de escalado de hasta 1B de parámetros, demostramos que Mamba supera el rendimiento de una amplia gama de bases, incluyendo recetas de entrenamiento de Transformer modernas y muy robustas basadas en LLaMa (Touvron et al., 2023). Nuestro modelo de lenguaje Mamba tiene un rendimiento de generación quintuplicado en comparación con Transformers de tamaño similar, y la calidad de Mamba-3B es igual a la de Transformers del doble de tamaño (por ejemplo, un promedio de 4 puntos superior en razonamiento de sentido común en comparación con Pythia-3B, e incluso superior a Pythia-7B).

El código del modelo y los puntos de control entrenados previamente son de código abierto en https://github.com/state-spaces/mamba.

2 modelos de espacio de estados

Modelos generales de espacio de estados. Cabe destacar que el término "modelo de espacio de estados" tiene un significado muy amplio, que simplemente representa la noción de cualquier proceso recurrente con un estado latente. Se ha utilizado para referirse a diversos conceptos en diferentes disciplinas, incluyendo los procesos de decisión de Markov (MDP) (aprendizaje por refuerzo [Hafner et al., 2020]), el modelado causal dinámico (DCM) (neurociencia computacional [Friston, Harrison y Penny, 2003]), los filtros de Kalman (controles [Kalman, 1960]), los modelos ocultos de Markov (HMM) y los sistemas dinámicos lineales (LDS) (aprendizaje automático), así como los modelos recurrentes (y, en ocasiones, convolucionales) en general (aprendizaje profundo).

A lo largo de este artículo, utilizamos el término «SSM» para referirnos exclusivamente a la clase de SSM estructurados o modelos S4 (Gu, Goel y Ré, 2022; Gu, Gupta et al., 2022; Gupta, Gu y Berant, 2022; Hasani et al., 2023; Ma et al., 2023; Smith, Warrington y Linderman, 2023) y los usamos indistintamente. Para mayor comodidad, también podemos incluir derivados de dichos modelos, como aquellos que se centran en la recurrencia lineal o la convolución global (Y. Li et al., 2023; Orvieto et al., 2023; Poli et al., 2023), y aclarar los matices cuando sea necesario.

Arquitecturas SSM. Las SSM son transformaciones de secuencia independientes que pueden incorporarse en arquitecturas de redes neuronales de extremo a extremo. (A veces también las llamamos SSNN, que son a las capas SSM lo que las CNN son a las capas de convolución lineal). Analizamos algunas de las arquitecturas SSM más conocidas, muchas de las cuales también servirán como base principal.

• La atención lineal (Katharopoulos et al. 2020) es una aproximación de la autoatención que implica una recurrencia que puede verse como un SSM lineal degenerado.

• H3 (Dao, Fu, Saab, et al., 2023) generalizó esta recurrencia para usar S4; puede considerarse una arquitectura con un SSM intercalado entre dos conexiones con compuertas (Figura 3). H3 también inserta una convolución local estándar, que definen como un SSM de desplazamiento, antes de la capa principal del SSM.

• Hyena (Poli et al. 2023) utiliza la misma arquitectura que H3 pero reemplaza la capa S4 con una convolución global parametrizada por MLP (Romero et al. 2021).

• RetNet (Y. Sun et al. 2023) agrega una puerta adicional a la arquitectura y utiliza un SSM más simple, lo que permite una ruta de cálculo paralelizable alternativa, utilizando una variante de atención de múltiples cabezas (MHA) en lugar de convoluciones.

• RWKV (B. Peng et al., 2023) es una RNN reciente diseñada para el modelado del lenguaje basada en otra aproximación de atención lineal (Transformador sin atención [S. Zhai et al., 2021]). Su principal mecanismo «WKV» implica recurrencias LTI y puede considerarse como la relación de dos SSM.

Otros SSM y arquitecturas estrechamente relacionados se analizan con más detalle en un trabajo ampliado (Apéndice B). Destacamos en particular S5 (Smith, Warrington y Linderman, 2023), QRNN (Bradbury et al., 2016) y SRU (Lei et al., 2017), que consideramos los métodos más estrechamente relacionados con nuestro SSM selectivo principal.