Mamba: Redefiniendo el modelado de secuencias y superando la arquitectura Transformers

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

Profundizaremos en la capacidad de Mamba para superar los desafíos computacionales que enfrentan los Transformers tradicionales, especialmente con secuencias largas. Su enfoque selectivo en los modelos de espacio de estados permite una inferencia más rápida y una escalabilidad lineal con la longitud de la secuencia, mejorando 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 los basados en el enfoque transformer, lo que lo convierte en un avance notable en el aprendizaje automático.

Transformers vs Mamba Los Transformers, como GPT-4, han establecido puntos de referencia en el procesamiento del lenguaje natural. Sin embargo, su eficiencia disminuye con secuencias más largas. Aquí es 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 Transformers son hábiles para manejar secuencias de datos, como texto para modelos de idioma. A diferencia de modelos anteriores que procesaban datos secuencialmente, los Transformers procesan secuencias completas simultáneamente, lo que les permite capturar relaciones complejas dentro de los datos.

Utilizan un mecanismo de atención, que permite al modelo enfocarse 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 todos los demás elementos, proporcionando un peso que indica la importancia o “atención” que debe recibir cada elemento al predecir el siguiente elemento de la secuencia.

Los Transformers 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 totalmente conectada simple de avance de posición. Se utilizan normalización y conexiones residuales 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 multicabezal sobre la salida del codificador. La naturaleza secuencial del decodificador garantiza que las predicciones para una posición sólo puedan considerar posiciones anteriores, preservando la propiedad autorregresiva.

A diferencia de los Transformers, el modelo Mamba adopta un enfoque diferente. Mientras que los Transformers tratan el problema de las secuencias largas utilizando mecanismos de atención más complejos, Mamba utiliza espacios de estados selectivos, proporcionando un enfoque más computacional para superar las limitaciones anteriores.

Aquí hay una descripción general de alto nivel de cómo funciona un transformer:

Procesamiento de entrada: Los Transformers primero codifican los datos de entrada en un formato que el modelo pueda entender, a menudo utilizando incrustaciones que también incorporan la posición de cada elemento en la secuencia.

Mecanismo de atención: En su núcleo, el mecanismo de atención calcula una puntuación que representa cuánto enfoque poner en otras partes de la secuencia de entrada al comprender un elemento actual.

Arquitectura codificador-decodificador: El modelo transformer se compone de un codificador para procesar la entrada y un decodificador para generar la salida. Cada uno consta de múltiples capas que refinan la comprensión del modelo sobre la entrada.

Atención multicabezal: Dentro del codificador y el decodificador, la atención multicabezal permite al modelo atender simultáneamente a diferentes partes de la secuencia desde diferentes espacios de representación, mejorando su capacidad para aprender de diversos contextos.

Redes neuronales de avance de posición: Después de la atención, una red neuronal simple procesa la salida de cada posición por separado e idénticamente. Esto se combina con la entrada a través de una conexión residual y va seguido de una normalización de capa.

Generación de salida: El decodificador luego predice una secuencia de salida, influenciada por el contexto del codificador y lo que ha generado hasta ahora.

La capacidad del transformer para manejar secuencias en paralelo y su robusto mecanismo de atención lo hacen poderoso para tareas como traducción y generación de texto.

En contraste, el modelo Mamba opera de manera diferente utilizando espacios de estados selectivos para procesar secuencias. Este enfoque aborda la ineficiencia computacional en los Transformers cuando se trata con secuencias extensas. El diseño de Mamba permite una inferencia más rápida y escala linealmente con la longitud de la secuencia, estableciendo un nuevo paradigma para el modelado de secuencias que podría ser más eficiente, especialmente a medida que las secuencias se vuelven cada vez más extensas.

MAMBA

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

Las características clave de Mamba incluyen:

SSM selectivos: Estos permiten a Mamba filtrar información irrelevante y concentrarse en datos relevantes, mejorando su manejo de secuencias. Esta selectividad es crucial para un razonamiento eficiente basado en contenido.

Algoritmo optimizado para hardware: Mamba utiliza un algoritmo paralelo optimizado para hardware moderno, especialmente GPU. Este diseño permite un cálculo más rápido y reduce los requisitos de memoria en comparación con los modelos tradicionales.

Arquitectura simplificada: Al integrar SSM selectivos y eliminar 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, incluidos el lenguaje, el audio y la genómica, destacando tanto en el preentrenamiento como en las tareas específicas del dominio.

Por ejemplo, en el modelado del lenguaje, Mamba iguala o supera el rendimiento de modelos Transformer más grandes.

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

Las tareas de copia estándar son simples para modelos lineales. La copia selectiva y las cabezas de inducción requieren una memoria dinámica y consciente del contenido para LLMs.

Los modelos de espacio de estado estructurado (S4) han surgido recientemente como una prometedora clase de modelos de secuencia, que abarcan rasgos de RNN, CNN y modelos de espacio de estado clásicos. Los modelos S4 derivan inspiración de sistemas continuos, específicamente un tipo de sistema que asigna funciones o secuencias unidimensionales 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 básico de espacio de estado estructurado. Toma una secuencia x y produce una salida y utilizando parámetros aprendidos A, B, C y un parámetro de retardo Δ. La transformación implica discretizar los parámetros (convertir funciones continuas en discretas) y aplicar la operación SSM, que es invariante en el tiempo, lo que significa que no cambia en 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, permitiendo que los modelos S4 mantengan una conexión con los sistemas de tiempo continuo. Esto dota a los modelos con propiedades adicionales, como la invariancia de resolución, y garantiza una normalización adecuada, mejorando la estabilidad y el rendimiento del modelo. La discretización también establece paralelismos con los mecanismos de control encontrados en RNN, que son fundamentales para gestionar el flujo de información a través de la red.

Invarianza lineal en el tiempo (LTI)

Una característica central de los modelos S4 es su invarianza lineal en el tiempo. Esta propiedad implica que la dinámica del modelo se mantiene consistente a lo largo del tiempo, con los parámetros fijos para todos los pasos de tiempo. 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, que plantea desafíos para modelar datos que requieren una dinámica adaptativa. El documento de investigación reciente presenta un enfoque que supera estas limitaciones al introducir parámetros dependientes del 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” cubre ampliamente cualquier proceso recurrente que involucre un estado latente y se ha 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 se han optimizado para un cálculo eficiente mientras conservan la capacidad de modelar secuencias complejas.

Los modelos S4 se pueden integrar en arquitecturas de redes neuronales de extremo a extremo, funcionando como transformaciones de secuencia independientes. Se pueden ver como análogos a las capas de convolución en 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

Los modelos 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 enfocarse o filtrar entradas selectivamente proporcionan un medio más efectivo para mantener este estado comprimido, lo que lleva a modelos de secuencia más eficientes y potentes. Esta selectividad es vital para que los modelos controlen de manera adaptativa 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 del lenguaje y más allá.

Los SSM selectivos mejoran los SSM convencionales al permitir que sus parámetros dependan de la entrada, lo que introduce un grado de adaptabilidad previamente inalcanzable con modelos invariantes en el tiempo. Esto resulta en SSM dependientes del tiempo que ya no pueden usar convoluciones para un 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 entrada a los parámetros B y C, y un parámetro de retardo Δ. Esto permite al modelo enfocarse selectivamente en ciertas partes de la secuencia de entrada x. Se discretizan los parámetros teniendo en cuenta la selección y se aplica la operación SSM de manera variable en el tiempo utilizando una operación de escaneo, que procesa los elementos secuencialmente, ajustando el enfoque dinámicamente a lo largo del tiempo.

Puntos destacados del rendimiento de Mamba

Mamba es la mejor 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, como se demuestra en el modelado de ADN y audio, superando a 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 dichos modelos en términos de eficiencia y escalabilidad.

Introducción a 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 del repositorio de Mamba. Si surgen problemas de compatibilidad con las versiones de PyTorch, usar la bandera -no-build-isolation con pip puede ayudar. Estos modelos, entrenados en conjuntos de datos extensos como Pile y SlimPajama, están diseñados para satisfacer diversas necesidades computacionales y puntos de referencia de rendimiento.

Mamba ofrece diferentes niveles de interfaces, desde la capa SSM selectiva hasta el bloque Mamba y las estructuras completas del modelo de lenguaje. El bloque Mamba, que es el módulo principal de la arquitectura, utiliza una capa causal Conv1d y se puede integrar 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, resaltando la simplicidad y flexibilidad del modelo.

Los modelos Mamba preentrenados 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 satisfacer diversos requisitos computacionales y de rendimiento, cumpliendo con los estándares dimensionales de GPT-3. Los usuarios pueden esperar un alto rendimiento y precisión de estos modelos, lo que hace de Mamba una opción competitiva para diversas aplicaciones, que incluyen pero no se limitan al modelado del lenguaje.

El impacto de Mamba

Mamba representa un salto hacia adelante en el modelado de secuencias, ofreciendo una poderosa alternativa a las arquitecturas Transformer para procesar 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 la base de código de Mamba y sus modelos preentrenados la convierten en una herramienta accesible y robusta para investigadores y desarrolladores en el campo de la IA y el aprendizaje profundo.