Mamba: Redefiniendo el modelado de secuencias y superando la arquitectura de los Transformadores

En este artículo sobre Mamba, exploraremos cómo este innovador modelo de espacio de estado (SSM) revoluciona el modelado de secuencias. Desarrollado por Albert Gu y Tri Dao, Mamba se distingue por su eficiencia en el procesamiento de secuencias complejas en campos como el procesamiento de lenguaje, la genómica y el análisis de audio. Su modelado de secuencias lineal con espacios de estado selectivos garantiza un rendimiento excepcional en estas diversas modalidades.

Nos adentraremos en la capacidad de Mamba para superar los desafíos computacionales que enfrentan los Transformadores tradicionales, especialmente con secuencias largas. Su enfoque selectivo en los modelos de espacio de estado permite una inferencia más rápida y una escalabilidad lineal con la longitud de la secuencia, lo que mejora significativamente el rendimiento.

La singularidad de Mamba radica en su capacidad de procesamiento rápido, su capa SSM selectiva y su diseño amigable con el hardware inspirado en FlashAttention. Estas características permiten que Mamba supere a muchos modelos existentes, incluidos aquellos basados en el enfoque de los transformadores, lo que lo convierte en un avance notable en el aprendizaje automático.

Transformadores vs Mamba

Los Transformadores, como GPT-4, han establecido estándares en el procesamiento de lenguaje natural. Sin embargo, su eficiencia disminuye con secuencias más largas. Es aquí donde Mamba se adelanta, con su capacidad para procesar secuencias largas de manera más eficiente y su arquitectura única que simplifica todo el proceso.

Los Transformadores son hábiles para manejar secuencias de datos, como texto para modelos de lenguaje. A diferencia de los modelos anteriores que procesaban los datos secuencialmente, los Transformadores procesan las secuencias enteras simultáneamente, lo que les permite capturar relaciones complejas dentro de los datos.

Utilizan un mecanismo de atención, que permite al modelo centrarse en diferentes partes de la secuencia al hacer predicciones.

Esta atención se calcula utilizando tres conjuntos de pesos: consultas, claves y valores, derivados de los datos de entrada. Cada elemento de una secuencia se compara con cada otro elemento, proporcionando un peso que indica la importancia o “atención” que cada elemento debe recibir al predecir el siguiente elemento de la secuencia.

Los Transformadores mantienen dos bloques principales: el codificador, que procesa los datos de entrada, y el decodificador, que genera la salida. El codificador consta de múltiples capas, cada una con dos subcapas: un mecanismo de autoatención de múltiples cabezas y una red neuronal feed-forward completamente conectada simple y position-wise. La normalización y las conexiones residuales se utilizan en cada subcapa para ayudar en el entrenamiento de redes profundas.

El decodificador también tiene capas con dos subcapas similares al codificador, pero agrega una tercera subcapa que realiza una atención de múltiples cabezas sobre la salida del codificador. La naturaleza secuencial del decodificador garantiza que las predicciones para una posición solo puedan considerar posiciones anteriores, preservando la propiedad auto-regresiva.

En contraste con los Transformadores, el modelo Mamba adopta un enfoque diferente. Mientras que los Transformadores abordan el problema de las secuencias largas utilizando mecanismos de atención más complejos, Mamba utiliza espacios de estado selectivos, lo que proporciona una forma más computacionalmente eficiente de procesar secuencias.

La singularidad de Mamba radica en su capacidad de procesamiento rápido, su capa SSM selectiva y su diseño amigable con el hardware inspirado en FlashAttention. Estas características permiten que Mamba supere a muchos modelos existentes, incluidos aquellos basados en el enfoque de los transformadores, lo que lo convierte en un avance notable en el aprendizaje automático.

Mamba

Lo que hace que Mamba sea verdaderamente único es su alejamiento de los bloques de atención y MLP tradicionales. Esta simplificación conduce a un modelo más ligero y rápido que se escala linealmente con la longitud de la secuencia, una hazaña sin precedentes en sus predecesores.

Las características clave de Mamba incluyen:

1. SSM selectivos: Estos permiten que Mamba filtre información irrelevante y se centre en los datos relevantes, mejorando su manejo de secuencias. Esta selectividad es crucial para el razonamiento basado en contenido eficiente.
2. Algoritmo consciente del hardware: Mamba utiliza un algoritmo paralelo que está optimizado para el hardware moderno, especialmente las GPU. Este diseño permite cálculos más rápidos y reduce los requisitos de memoria en comparación con los modelos tradicionales.
3. Arquitectura simplificada: Al integrar SSM selectivos y eliminar los bloques de atención y MLP, Mamba ofrece una estructura más simple y homogénea. Esto conduce a una mejor escalabilidad y rendimiento.

Mamba ha demostrado un rendimiento superior en varios dominios, incluyendo lenguaje, audio y genómica, destacándose tanto en tareas de preentrenamiento como en tareas específicas de dominio.

El código y los modelos preentrenados de Mamba están disponibles abiertamente para su uso en la comunidad en GitHub.

Tareas de copia estándar son simples para los modelos lineales. La copia selectiva y las cabezas de inducción requieren memoria dinámica y consciente del contenido para los LLM.

Los modelos de espacio de estado estructurado (S4) han surgido recientemente como una clase prometedora de modelos de secuencia, que abarca rasgos de las RNN, las CNN y los modelos de espacio de estado clásicos. Los modelos S4 derivan su inspiración de sistemas continuos, específicamente de un tipo de sistema que mapea funciones unidimensionales o secuencias a través de un estado latente implícito. En el contexto del aprendizaje profundo, representan una innovación significativa, proporcionando una nueva metodología para diseñar modelos de secuencia que sean eficientes y altamente adaptables.

La dinámica de los modelos S4

SSM (S4) Este es el modelo de espacio de estado estructurado básico. Toma una secuencia x y produce una salida y utilizando parámetros aprendidos A, B, C y un parámetro de retraso Δ. La transformación implica discretizar los parámetros (convirtiendo funciones continuas en discretas) y aplicar la operación SSM, que es invariante en el tiempo, lo que significa que no cambia a lo largo de los diferentes pasos de tiempo.

La importancia de la discretización

La discretización es un proceso clave que transforma los parámetros continuos en discretos a través de fórmulas fijas, lo que permite a los modelos S4 mantener una conexión con los sistemas de tiempo continuo. Esto dota a los modelos de propiedades adicionales, como la invarianza de resolución, y garantiza una normalización adecuada, lo que mejora la estabilidad y el rendimiento del modelo. La discretización también traza paralelos con los mecanismos de puerta encontrados en las RNN, que son críticos para gestionar el flujo de información a través de la red.

Invarianza lineal en el tiempo (LTI)

Una característica fundamental de los modelos S4 es su invarianza lineal en el tiempo. Esta propiedad implica que la dinámica del modelo permanece consistente en el tiempo, con los parámetros fijos para todos los pasos de tiempo. La LTI es una piedra angular de la recurrencia y las convoluciones, ofreciendo un marco simplificado pero poderoso para construir modelos de secuencia.

Superar limitaciones fundamentales

El marco S4 ha estado tradicionalmente limitado por su naturaleza LTI, lo que plantea desafíos en la modelización de datos que requieren dinámicas adaptativas. El reciente artículo de investigación presenta un enfoque que supera estas limitaciones al introducir parámetros que varían con el tiempo, eliminando así la restricción de LTI. Esto permite que los modelos S4 manejen un conjunto más diverso de secuencias y tareas, expandiendo significativamente su aplicabilidad.

El término “modelo de espacio de estado” abarca ampliamente cualquier proceso recurrente que involucre un estado latente y ha sido utilizado para describir varios conceptos en múltiples disciplinas. En el contexto del aprendizaje profundo, los modelos S4, o SSM estructurados, se refieren a una clase específica de modelos que han sido optimizados para cálculos eficientes mientras retienen la capacidad de modelar secuencias complejas.

Los modelos S4 pueden integrarse en arquitecturas de redes neuronales de extremo a extremo, funcionando como transformaciones de secuencia independientes. Pueden considerarse análogos a las capas de convolución en las CNN, proporcionando la columna vertebral para el modelado de secuencias en una variedad de arquitecturas de redes neuronales.

SSM vs SSM + Selección

Motivación para la selectividad en el modelado de secuencias

SSM estructurados

El artículo argumenta que un aspecto fundamental del modelado de secuencias es la compresión del contexto en un estado manejable. Los modelos que pueden centrarse selectivamente o filtrar las entradas proporcionan un medio más efectivo para mantener este estado comprimido, lo que conduce a modelos de secuencia más eficientes y poderosos. Esta selectividad es vital para que los modelos puedan controlar adaptativamente cómo fluye la información a lo largo de la dimensión de la secuencia, una capacidad esencial para manejar tareas complejas en el modelado de lenguaje y más allá.

Los SSM selectivos mejoran los SSM convencionales al permitir que sus parámetros sean dependientes de la entrada, lo que introduce un grado de adaptabilidad que anteriormente no era alcanzable con modelos de tiempo invariante. Esto resulta en SSM que varían con el tiempo, que ya no pueden utilizar convoluciones para el cálculo eficiente, sino que dependen de un mecanismo de recurrencia lineal, una desviación significativa de los modelos tradicionales.

SSM + Selección (S6) Esta variante incluye un mecanismo de selección, agregando dependencia de la entrada a los parámetros B y C, y un parámetro de retraso Δ. Esto permite que el modelo se centre selectivamente en ciertas partes de la secuencia de entrada x. Los parámetros se discretizan teniendo en cuenta la selección, y la operación SSM se aplica de manera que varía con el tiempo, utilizando una operación de exploración, que procesa los elementos secuencialmente, ajustando el enfoque dinámicamente con el tiempo.

Aspectos destacados del rendimiento de Mamba

Mamba es el mejor de su clase en cada resultado de evaluación

En términos de rendimiento, Mamba sobresale tanto en velocidad de inferencia como en precisión. Su diseño permite un mejor uso de contextos más largos, lo que se demuestra tanto en el modelado de ADN como en el modelado de audio, superando a los modelos anteriores en tareas complejas que requieren dependencias de largo alcance. Su versatilidad también se destaca en evaluaciones de disparo cero en múltiples tareas, estableciendo un nuevo estándar para este tipo de modelos en términos de eficiencia y escalabilidad.

Comenzando con Mamba

Para aquellos interesados en aprovechar Mamba, los requisitos técnicos incluyen un sistema operativo Linux, una GPU NVIDIA, PyTorch 1.12+ y CUDA 11.6+. La instalación implica simples comandos pip para instalar los paquetes necesarios desde el repositorio de Mamba. Si surgen problemas de compatibilidad con versiones de PyTorch, usar la bandera –no-build-isolation con pip puede ayudar. Estos modelos, entrenados en extensos conjuntos de datos como el Pile y el conjunto de datos SlimPajama, están diseñados para cumplir con diversas necesidades computacionales y estándares de rendimiento.

Mamba ofrece diferentes niveles de interfaces, desde la capa SSM selectiva hasta el bloque Mamba y las estructuras de modelo de lenguaje completas. El bloque Mamba, que es el módulo principal de la arquitectura, utiliza una capa Conv1d causal y puede integrarse fácilmente en diseños de redes neuronales. El ejemplo de uso proporcionado en Python demuestra la instanciación de un modelo Mamba y el procesamiento de datos a través de él, destacando la simplicidad y la flexibilidad del sistema.

Los modelos preentrenados de Mamba están disponibles en Hugging Face, con tamaños que van desde 130M hasta 2.8B parámetros, entrenados en el extenso conjunto de datos Pile y el conjunto de datos SlimPajama. Estos modelos están diseñados para cumplir con diversas necesidades computacionales y de rendimiento, ajustándose a los estándares dimensionales de GPT-3. Los usuarios pueden esperar un alto rendimiento y precisión de estos modelos, lo que convierte a Mamba en una opción competitiva para diversas aplicaciones, incluyendo, pero no limitado a, el modelado de lenguaje.

El impacto de Mamba

Mamba representa un salto adelante en el modelado de secuencias, ofreciendo una poderosa alternativa a las arquitecturas de Transformadores para el procesamiento de datos densos en información. Su diseño se alinea con las demandas del hardware moderno, optimizando tanto el uso de memoria como las capacidades de procesamiento paralelo. La disponibilidad de código abierto de Mamba y sus modelos preentrenados lo convierte en una herramienta accesible y robusta para investigadores y desarrolladores en el campo de la IA y el aprendizaje profundo.