Princeton y CMU amplían los límites de la IA con el modelo de secuencia Mamba

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 los SSM selectivos

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 y referencias

Una discusión: Mecanismo de selección

B Trabajo relacionado

C Mecánica de los SSM selectivos

Algoritmo D basado en hardware para SSM selectivos

E Detalles experimentales y resultados adicionales

Abstracto

Los modelos básicos, que ahora impulsan la mayoría de las aplicaciones interesantes en el aprendizaje profundo, se basan casi universalmente en la arquitectura Transformer y su módulo de atención central. Se han desarrollado muchas arquitecturas de tiempo subcuadrático, como la atención lineal, los modelos de convolución controlada y recurrentes, y los modelos de espacio de estado estructurado (SSM) para abordar la ineficiencia computacional de los Transformers en secuencias largas, pero no han tenido el mismo rendimiento que la atención en modalidades importantes como el lenguaje. Identificamos que una debilidad clave de dichos modelos es su incapacidad para realizar razonamiento basado en contenido, y realizamos varias mejoras. Primero, simplemente dejar que los parámetros SSM sean funciones de la entrada aborda su debilidad con modalidades discretas, lo que permite que el modelo propague u olvide información de manera selectiva a lo largo de la dimensión de longitud de la secuencia según el token actual. Segundo, 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 de extremo a extremo simplificada sin atención o incluso bloques MLP (Mamba). Mamba ofrece una inferencia rápida (un rendimiento 5 veces superior al de Transformers) y un escalamiento lineal en la longitud de la secuencia, y su rendimiento mejora con respecto a los datos reales hasta secuencias de un millón de longitudes. Como eje central de un modelo de secuencia general, Mamba logra un rendimiento de vanguardia en varias modalidades, como lenguaje, audio y genómica. En cuanto al modelado del lenguaje, nuestro modelo Mamba-3B supera a Transformers del mismo tamaño y se equipara a Transformers del doble de tamaño, tanto en el preentrenamiento como en la evaluación posterior.

1 Introducción

Los modelos de base (FM), o modelos de gran tamaño entrenados previamente con datos masivos y luego adaptados para tareas posteriores, han surgido como un paradigma eficaz en el aprendizaje automático moderno. La columna vertebral de estos FM son a menudo modelos de secuencia, que operan con secuencias arbitrarias de entradas 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). Si bien este concepto es agnóstico a una elección particular de arquitectura de 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 trae 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 estos inconvenientes (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 surgido 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), con inspiración de los modelos de espacio de estados clásicos (Kalman 1960). Esta clase de modelos se puede calcular de manera muy eficiente como una recurrencia o una convolución, con escalamiento lineal o casi lineal en la longitud de la secuencia. Además, tienen mecanismos de principios para modelar dependencias de largo alcance (Gu, Dao, et al. 2020) en ciertas modalidades de datos, y han dominado puntos de referencia como el Long Range Arena (Tay, Dehghani, Abnar, et al. 2021). Muchos tipos 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 efectivos 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 de manera eficiente de una manera dependiente de la entrada (es decir, centrarse en entradas particulares 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 parametrizando los parámetros del SSM en función de la entrada. Esto permite que el modelo filtre la información irrelevante y recuerde la información relevante indefinidamente.

Algoritmo que tiene en cuenta el hardware. Este cambio simple plantea un desafío 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. Superamos esto con un algoritmo que tiene en cuenta el hardware que calcula el modelo de forma recurrente con 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 3 veces más rápido en GPU A100).

Arquitectura . Simplificamos las arquitecturas de modelos de secuencias profundas anteriores combinando el diseño de arquitecturas SSM anteriores (Dao, Fu, Saab, et al. 2023) con el bloque MLP de Transformadores en un solo bloque, lo que conduce a 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 adecuados como columna vertebral de los modelos de base general que operan en secuencias. (i) Alta calidad: la selectividad brinda un sólido desempeño 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 desarrollar el modelo de manera autorregresiva durante la inferencia solo requiere un tiempo constante por paso ya que no requiere un caché de elementos previos. (iii) Contexto largo: la calidad y la eficiencia juntas producen mejoras en el desempeño en datos reales hasta una longitud de secuencia de 1M.

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 claves para los modelos de lenguaje de gran tamaño, Mamba no solo las resuelve fácilmente, sino que puede extrapolar soluciones a un tiempo indefinido (>1 millón de tokens).

• Audio y genómica. Mamba supera a los modelos de última generación anteriores, 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 (por ejemplo, reduciendo la FID en un conjunto de datos de generación de voz desafiante en más de la mitad). En ambos entornos, su rendimiento mejora con un contexto más largo, hasta secuencias de un millón de longitudes.

• Modelado del lenguaje. Mamba es el primer modelo de secuencia de tiempo lineal que realmente logra un rendimiento de calidad de Transformer, tanto en la perplejidad previa al entrenamiento como en las evaluaciones posteriores. Con leyes de escalamiento de hasta parámetros 1B, demostramos que Mamba supera el rendimiento de una amplia gama de líneas base, incluidas recetas de entrenamiento de Transformer modernas muy sólidas basadas en LLaMa (Touvron et al. 2023). Nuestro modelo de lenguaje Mamba tiene un rendimiento de generación 5 veces mayor en comparación con Transformers de tamaño similar, y la calidad de Mamba-3B coincide con la de Transformers del doble de su tamaño (por ejemplo, 4 puntos más de promedio en razonamiento de sentido común en comparación con Pythia-3B e incluso supera a Pythia-7B).