¿Qué es un modelo Mamba?

Los modelos Mamba son quizás la primera arquitectura de aprendizaje profundo que compite con la eficacia de los modelos de transformadores en la tarea por la que los transformadores ganaron originalmente su fama: el modelado del lenguaje. En particular, la arquitectura Mamba ha demostrado la capacidad de igualar a transformadores de tamaño equivalente en destacadas evaluaciones de punto de referencia de LLM, que suele ser significativamente más eficiente en términos de latencia y requisitos de memoria.

La arquitectura Mamba fue introducida por primera vez por Tri Dao y Albert Gu en el artículo de 2023 "Mamba: Linear-Time Sequence Modeling with Selective State Spaces". Un año más tarde, tras el artículo original sobre Mamba, publicaron otro en el que exploraban las conexiones entre los SSM y los transformadores, y presentaban una versión refinada y mucho más rápida de la arquitectura Mamba, a la que bautizaron como Mamba-2.

Aunque los transformadores han seguido siendo el modo dominante de LLM en los 2 años posteriores al lanzamiento del documento original de Mamba, la arquitectura se ha incorporado a un número creciente de modelos de código abierto. Algunos, como Codestral Mamba de Mistral AI, son modelos puros de Mamba. Muchos más, incluida la serie Jamba de AI2I e IBM Granite 4.0, son modelos híbridos que incorporan capas de atención (transformador) y capas SSM (Mamba). Además de sus beneficios basados en el rendimiento, la proliferación de modelos basados en Mamba promete democratizar el acceso a la IA en virtud de su funcionamiento sin problemas en hardware comparativamente económico.

Los SSM se diseñaron originalmente para predecir el siguiente estado de una secuencia continua, como una señal eléctrica, un patrón meteorológico o la trayectoria de un objeto en movimiento, en función de alguna entrada. Desde el punto de vista conceptual y matemático, están relacionados con las redes neuronales recurrentes (RNN) que dominaban el procesamiento de lenguaje natural (PLN) antes de la introducción de los transformadores en 2017, así como con otros algoritmos de machine learning, incluidas las redes neuronales convolucionales (CNN) y los modelos ocultos de Markov (HMM).

Como su nombre indica, los SSM hacen predicciones sobre el siguiente estado en un sistema dinámico modelando el espacio de estados: una representación matemática de todas las variables de estados que describen el estado de un sistema y el rango de posibilidades para cada una de esas variables en conjunto con las demás.

Un SSM toma una secuencia de entrada x(t) y la asigna a una representación de estado latente h(t), análoga al estado oculto de un RNN, para predecir una secuencia de salida y(t). En el núcleo de cualquier SSM hay 2 ecuaciones:

• La ecuación de estados,  $h^{\prime }\left(t\right)=A*h\left(t\right)+B*x\left(t\right)$ 
• La ecuación de salida,  $y\left(t\right)=C*h\left(t\right)+D*x\left(t\right)$ 

Los parámetros clave del modelo son A, B, C y D, que suelen adoptar la forma de una matriz de ponderaciones. En los campos en los que se utilizan convencionalmente los SSM, como la teoría de control, a menudo se supone que estas matrices son fijas: representan la dinámica de un sistema establecido, y el SSM se utiliza para encontrar las entradas x que llevan a las salidas deseables y. En concepciones más modernas de los SSM, esas matrices son en sí mismas parámetros que deben optimizarse mediante machine learning. En los modelos de aprendizaje profundo, esas matrices están representadas por la ponderaciones aprendibles de una red neuronal.

La ecuación de estado describe cómo cambia el estado. Los valores de la matriz A determinan cómo evoluciona cada variable de estado con el tiempo si se deja a sí misma. Los valores de la matriz B determinan cómo la entrada, como el siguiente token en una secuencia de texto, influye en cada variable de estado.

En el modelado de lenguaje, el estado actual representa el contexto de una secuencia de texto, actualizado después de cada token. Su función es equivalente a la de la caché KV en un modelo de transformador.

La ecuación de salida describe cómo el estado actual influye en la salida (mediada por la matriz C), así como cómo la entrada influye directamente en la salida (mediada por la matriz D). Debido a que la matriz D es esencialmente externa al modelado de h(t),a menudo se omite de los diagramas y debates de SSM a favor de centrarse en las matrices centrales A, B y C.

En un LLM de Mamba, la ecuación de salida se utiliza para generar el siguiente token.

Los SSM tradicionales están diseñados para modelar entradas continuas, pero las secuencias de texto (y la mayoría de las demás modalidades de datos procesadas por los modelos modernos de aprendizaje profundo) son entradas discretas. El uso de SSM para modelar una secuencia discreta requiere un medio para representar sus pasos de tiempo distintos y específicos como parte de una señal continua.

Conceptualmente, la discretización equivale a muestrear el valor de una función continua en momentos específicos. Esto implica la introducción de un nuevo parámetro (el tamaño del paso, escrito como ∆) que determina durante cuánto tiempo se muestrea o “mantiene” ese valor en cada paso de tiempo discreto t. Los ajustes de ∆ son similares a los cambios en cualidades como la resolución de los datos (para datos de series de tiempo) o la velocidad de fotogramas (para datos de video). Existen múltiples métodos de “discretización”, pero la mayoría de las variantes modernas de SSM (incluido Mamba) utiliza el método simple de retención de orden cero (ZOH).

Discretizar un SSM permite usarlo como una RNN para tareas de secuencia a secuencia. Los parámetros y las ecuaciones de un SSM discretizado suelen reescribirse para distinguirlos de sus equivalentes de tiempo continuo mediante la notación de subíndice usada habitualmente para las RNN. En esta notación, h_t representa el espacio de estados actualizado que generará el modelo y h_t-1 representa el estado anterior, es decir, el espacio de estados actual.

 $h_t=\overline{A}{h}_{t-1}+\overline{B}{x}_t$ 
 $y_t=\overline{C}{h}_t$

Modelar datos de texto utilizando SSM discretos estándar no es práctico debido a una serie de deficiencias que comparten con los RNN. Dos de esas deficiencias fueron abordadas por la introducción de modelos de secuencia espacial de estado estructurado (o “modelos S4") por Albert Gu et al en 2021: la ineficiencia de su entrenamiento y su incapacidad para modelar secuencias largas.

Aunque el éxito de los modelos S4, y sus muchos derivados, como los modelos diagonales SSM (DSS), diagonales S4 (S4D) y H3, sentó directamente las bases para lo que se convirtió en Mamba.

El beneficio de que los SSM discretizados sean el equivalente de una instancia específica de un RNN es que los RNN son extremadamente rápidos en la inferencia. Sin embargo, la desventaja es que las RNN son extremadamente lentas de capacitar.

Afortunadamente, los SSM discretizados tienen una propiedad importante que los distingue de otras RNN: modelan exclusivamente dependencias lineales. En otras palabras, solo utilizan operaciones simples y directas de multiplicación y suma. Como demuestra el documento de S4, estas recurrencias lineales simples, repetidas e interdependientes se pueden desenrollar en un kernel de convolución unidimensional, que asigna directamente la entrada x a la salida y en un solo paso. Esto se puede calcular de manera muy eficiente mediante la transformada rápida de Fourier.

La única “trampa” es que esto solo es posible cuando se conoce cada paso de toda la secuencia de entrada. Esto no es posible durante la inferencia, pero es el caso durante el entrenamiento. Por lo tanto, un SSM estructurado se beneficia de lo mejor de ambos mundos: durante el entrenamiento, puede funcionar de manera muy eficiente como una CNN; durante la inferencia, puede funcionar de manera muy eficiente como una RNN.

Como la mayoría de las RNN, los SSM estándar son inherentemente débiles para modelar dependencias de larga distancia. En otras palabras, no sirven para comprender la relación entre los pasos de una secuencia que están muy separados, como las palabras al principio y al final de un párrafo, lo que los hace débiles para modelar secuencias largas en conjunto.

Para resolver esto, Gu y sus coautores (uno de los cuales fue Tri Dao) utilizaron una técnica llamada HiPPO,abreviatura de High-order Polynomial Projection Operators, para definir la manera en que se comportan las matrices A y B estructurando sus valores iniciales mediante una fórmula derivada de polinomios ortogonales. Esto contrasta con la práctica estándar de machine learning, en la que las ponderaciones del modelo se inicializan aleatoriamente al inicio del entrenamiento del modelo. Para S4, Dao y Gu propusieron esquemas de inicialización derivados de polinomios de Legendre. Exploraron fórmulas adicionales en un documento de seguimiento, titulado "How to Train Your HiPPO".¹

El documento de S4 señala que "simplemente modificar un SSM de una matriz aleatoria A a [la matriz HiPPO] mejoró su rendimiento en el punto de referencia secuencial MNIST del 60 % al 98%", lo cual resolvió eficazmente el problema de memoria a largo plazo de los SSM. Las variaciones posteriores de los SSM estructurados, como DSS, S5 y Mamba, emplean diferentes esquemas de inicialización (a menudo más simples) para A y B que, sin embargo, conservan los principios básicos de HiPPO: implementar una estructura diagonal que impone actualizaciones estables y cierto grado de independencia entre cada valor de la matriz.

En el núcleo de la arquitectura de Mamba hay dos innovaciones. El primero es el modelo de espacio de estado selectivo, que proporciona a Mamba una capacidad crucial que antes solo poseían los modelos transformadores: la capacidad de centrarse selectivamente o ignorar partes específicas del historial de entrada pasado en función de su relevancia actual. El otro es el escaneo paralelo con reconocimiento de hardware , un algoritmo que optimiza la forma en que una unidad de procesamiento de gráficos (GPU) maneja los cálculos del modelo en su jerarquía de memoria para maximizar la velocidad y la eficiencia computacional.

En los transformadores, esta capacidad la proporciona el mecanismo de atención que ajusta los pesos de atención que enfatizan o restan énfasis a la influencia de cada token anterior en función de su relevancia para el token de entrada actual. Los SSM ordinarios están diseñados explícitamente para mapear la entrada a la salida utilizando todo el historial de entrada.Esto es aceptable o incluso deseable para algunas tareas de modelado de secuencias, pero es una desventaja significativa para la mayoría de las tareas avanzadas de modelado de lenguaje.

Para remediar esta incapacidad de omitir o enfatizar dinámicamente partes específicas de su historia de entrada, Dao y Gu propusieron una nueva clase de modelos de espacio de estados con una "exploración selectiva". En el artículo de Mamba, los autores comentan que "a veces abrevian los SSM selectivos como modelos S6, porque son modelos S4 con un mecanismo de selección y se calculan con un escaneo". Apodaron "Mamba" a su arquitectura basada en S6 porque, entre otras razones, todas esas S suenan como el silbido de una serpiente.

Mamba se puede entender mejor como una arquitectura de Neural Networks que contiene el modelo de espacio de estado selectivo en su núcleo. Para una analogía simple, Mamba es para los SSM selectivos lo que el modelo transformador es para el mecanismo de atención.

Un SSM tradicional tiene una dinámica fija: las reglas que rigen cómo evoluciona el estado oculto de un paso al siguiente (los parámetros del modelo) son las mismas para cada entrada y en cada paso de la secuencia. Esta propiedad se conoce como invariancia lineal en el tiempo (LTI). Para proporcionar a los SSM la capacidad de priorizar o despriorizar selectivamente información pasada específica en función del contexto actual, Dao y Gu reconfiguraron su SSM de modo que los valores de los parámetros clave del modelo sean diferentes para diferentes entradas.

Más específicamente, los SSM selectivos hacen que el tamaño del paso ∆_t y las matrices B_t y C_t sean funciones directas del token de entrada actual x_t. Esto se logra pasando primero la incorporación de x_t a través de tres capas de proyección lineal paralelas, en otras palabras, capas de red neuronal feedforward estándar (o capas MLP). Esto es equivalente a cómo las cabezas de consulta paralela, clave y valor generan los respectivos vectores Q, K y V de una entrada en un modelo de transformador.

Al multiplicar la incrustación vectorial de xt por los términos de_{peso y sesgo en esa red de proyecciónlineal, se obtienen los valores resultantes de ∆ t}, B_t y C_t. Los términos de ponderación y sesgo de las propias capas de proyección lineal se aprenden durante el entrenamiento previo del modelo en conjuntos de datos masivos de muestras de texto y luego (opcionalmente) se refinan mediante ajustes posteriores.

• El valor de ∆_t determina la magnitud de la influencia de xt en la memoria del modelo del contexto que ha visto hasta ahora: en otras palabras, en cuánta actualización habrá del estado oculto h_t-1 a _ht. Un tamaño de paso mayor ∆_t da como resultado mayores cambios y acelera el decaimiento (en otras palabras, el "olvido") de la información más antigua contenida dentro del estado. Por el contrario, un tamaño de paso más pequeño da como resultado una actualización más pequeña. Con un tamaño de paso lo suficientemente pequeño, la entrada actual no tendrá ningún impacto en el estado oculto.
  
  
• Los cambios en la matriz B_k determinan cómo el token de entrada actual actualiza el estado oculto. Por ejemplo, si xt es un token para la palabra “yesterday”, B_t podría ajustarse de manera que actualice el estado para reflejar que el contexto resultante probablemente pertenezca al pasado.
  
  
• Los cambios en la matriz C_t determinan cómo se traduce esta información contextual en influencia sobre los resultados y_t del modelo. Siguiendo el ejemplo en el que xk_k es un token para “ayer”, Ck_k podría ver influido de tal forma que cualquier verbo que el modelo produzca posteriormente se conjugue en pasado.

En particular, no se realizan tales ajustes basados en entradas en la matriz A. Su función sigue siendo la misma que en los modelos S4: memorizar de manera eficiente todo el historial de entradas pasadas. La función de determinar qué partes de ese historial utilizar en un momento dado está a cargo de las matrices B y C.

Pero una vez que el modelo ya no es invariante en el tiempo, ya no puede usar el atajo de convolución durante el entrenamiento porque el kernel de transición ya no es constante: el quid del mecanismo de selectividad es que la transición de h_t-1 a ht ahora depende en el contexto.

En su lugar, Mamba utiliza una solución inteligente para lograr beneficios de paralelización similares. Debido a que el SSM usa solo multiplicación y suma, sus cálculos están sujetos a la propiedad asociativa familiar de las matemáticas: se pueden agrupar de diferentes maneras sin cambiar el resultado final. Esto permite que los muchos cálculos secuenciales se dividan en pequeños fragmentos independientes que una GPU puede procesar en paralelo a través de un escaneo de suma de prefijos paralelos.

Además, los resultados se combinan de una manera jerárquica específica que hace un uso óptimamente eficiente de los diferentes tipos de memoria de hardware en una GPU, utilizando principios similares a las técnicas FlashAttention, que también fueron desarrolladas por Tri Dao, que ahora son omnipresentes en los LLM modernos. .  

Dentro de la arquitectura Mamba, el modelo S6 sirve como un módulo del "bloque Mamba" más grande, de manera similar a cómo el mecanismo de atención sirve como un módulo dentro del "bloque de atención" más grande. Combina el módulo S6 con una arquitectura de red neuronal cerrada. Los modelos Mamba suelen constar de varios bloques Mamba, es decir, una serie de capas Mamba consecutivas en una red neuronal, antes de la capa de resultados que realiza los resultados finales del modelo.

Antes de entrar en el bloque Mamba, se envía una copia de la entrada directamente al final como conexión residual. El propósito del funcionamiento interno del bloque Mamba es no solo determinar qué partes del contexto mayor son relevantes para esa entrada, sino determinar cuánto esa información contextual debe modificar el significado original de la entrada.

Dentro del bloque Mamba, el vector de entrada original se procesa del siguiente modo:

• Primero, la entrada se pasa a través de una capa lineal que es dos veces más ancha que el vector de entrada en sí, proyectándolo a un espacio de mayor dimensión. Por ejemplo, si el modelo representa originalmente cada token de entrada x como incorporación vectorial de 512 dimensiones, multiplicar x por los pesos de la capa de proyección lineal lo expande en un vector de 1024 dimensiones.
  
  
• A continuación, el vector expandido se divide en dos. Una mitad (que llamaremos x_proj) se introduce en la ruta que atraviesa el SSM, y la otra mitad (que llamaremos z_proj) se introduce en una ruta separada que se ejecuta a través de un mecanismo de puerta. Para mayor claridad, el paso de expansión anterior generalmente se representa como realizado por 2 capas lineales paralelas.
  
  
• Antes de que x_proj alcance el SSM, se alimenta a una capa de convolución unidimensional. Esta capa de convolución extrae patrones locales (como dependencias entre tokens vecinos, como simples pares verbo-sujeto). Esto permite que el SSM se “centre” en la comprensión contextual de las dependencias globales de largo alcance.
  
  
• La salida de la capa de convolución sirve como entrada para una función de activación no lineal. La introducción de la no linealidad es un sello distintivo de todas las redes neuronales, lo que les permite capturar patrones más complejos. El artículo de Mamba utiliza la unidad lineal sigmoide (SiLU). Llamaremos al vector resultante x_act.
  
  
• Mientras tanto, en la ruta del mecanismo de activaciónindependiente, z_proj también se ingresa a una función de activación no lineal, lo que produce z_act.
  
  
• En el trayecto del MUS, x_act se introduce en tres capas de proyección lineal paralelas que generan los valores respectivos para ∆_x, B_x y C_x, respectivamente.
  
  
• El SSM utiliza estos parámetros dependientes de la entrada (y las matrices A y D) para calcular la actualización del espacio de estado y la salida y del SSM.
  
  
• El vector de salida y ahora se multiplica por elementos por el vector de salida z_act de la ruta de activación. Esencialmente, cada elemento en z_act actúa como una perilla de volumen en una mesa de mezclas de audio: si un elemento dado de z_act está cerca de cero, la multiplicación con la parte correspondiente de y arrojará un valor más cercano a cero y su influencia disminuirá. Por el contrario, si un elemento dado de zact es grande, la multiplicación con y amplificará la influencia de su información contextual.
  
  
• El vector resultante se proyecta de nuevo a su tamaño original. Puede entenderse como un vector de actualizaciones contextuales ponderadas (o no actualizaciones) de cada uno de los elementos del vector de entrada original.
  
  
• Finalmente, ese vector de actualizaciones se agrega a la copia del vector de entrada original que se envió directamente al final del bloque como una conexión residual.
  
  
• El vector de entrada original ahora se ha actualizado para reflejar la comprensión contextual proporcionada por el SSM selectivo. Ahora se puede enviar a la siguiente capa de Mamba o, en las capas finales del modelo, servir como entrada a una función softmax que resulta la probabilidad respectiva de que el vector completamente actualizado corresponda a cada palabra del vocabulario del modelo.

Un año después del artículo original de Mamba, Dao y Gu lo siguieron con "Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality". Este documento de seguimiento ofreció tres contribuciones principales:

• Una exploración de las conexiones teóricas entre Mamba y los transformadores y un vocabulario compartido entre las dos arquitecturas.
• Una serie de aclaraciones y exploraciones de diferentes opciones de diseño para los modelos Mamba
• Una arquitectura modificada, Mamba-2, informada y mejorada por esas exploraciones de diseño.

El algoritmo Mamba-2 es significativamente más rápido y fácil de implementar que el Mamba original: los autores proporcionaron una base de código "minimal SSD" que implementa el SSM selectivo en aproximadamente 25 líneas de código.² Esta eficiencia permite a Mamba-2 utilizar dimensiones de estado oculto mucho más grandes sin ralentizar el modelo, lo que permite modelos más grandes, potentes y expresivos creados con la arquitectura. En las pruebas, los modelos Mamba-2 definitivamente igualaron o superaron a los modelos Mamba y transformadores de tamaño correspondiente en una serie de tareas posteriores.

Como se indica en las introducciones del artículo, el "objetivo principal de Dao y Gu [era] desarrollar un rico cuerpo de conexiones teóricas entre los SSM estructurados y las variantes de atención". Esto arrojó un nuevo marco conceptual que une a ambos, al que denominaron “dualidad espacial de estado” (SSD).³ Al hacerlo, abrieron la puerta para que Mamba se beneficiara de varios años de exploración y optimización de la arquitectura transformadora.

Un beneficio notable fue el desarrollo de un equivalente Mamba de atención multicabezal (MHA), en el que un bloque Mamba se puede dividir en múltiples "cabezas Mamba" similares a las múltiples "cabezas de atención" en los transformadores. Una variante de este enfoque, que consideraron análoga a la atención de consultas agrupadas , permite aún más eficiencia a través del paralelismo tensorial en las GPU.

En el bloque Mamba-2 —al que llaman el bloque paralelo Mamba (en contraposición al bloque “secuencial” original de Mamba”) —los parámetros dependientes de la entrada (en paralelo), B y C se generan en paralelo en la capa de proyección inicial. B y C, específicamente, se derivan simplemente copiando porciones de x_proj, en lugar de multiplicar x_proj a través de capas lineales dedicadas. Además de simplificar y reducir los parámetros totales del modelo, este paralelismo permite un entrenamiento a gran escala significativamente más eficiente.⁴

Tanto Mamba como los transformadores tienen sus respectivos puntos fuertes, pero los modelos basados en Mamba son generalmente superiores en todo lo relacionado con el uso de memoria y la velocidad: según el documento de Mamba, Mamba ofrece un rendimiento 5 veces mayor que los transformadores equivalentes.

Los transformadores son increíblemente precisos y versátiles, pero también increíblemente exigentes en cuanto a recursos computacionales. Durante el preentrenamiento (y el refinamiento), los requisitos de memoria de la autoatención se escalan cuadráticamente con la longitud de la secuencia: si duplica la longitud del contexto de una secuencia, el mecanismo de atención utiliza el cuádruple de recursos. Este “cuello de botella cuadrático” limita cada vez más la velocidad y la disponibilidad de memoria a medida que crece la ventana de contexto. Durante la inferencia, sus necesidades de memoria escalan linealmente.

Durante el entrenamiento, el uso de memoria de un modelo Mamba se escala solo linealmente durante el entrenamiento. Más importante aún, su uso de memoria durante la inferencia es constante: independientemente de cuántos tokens el modelo ha visto, el SSM mantiene una representación de tamaño fijo de su historial de entrada. Esto permite una longitud de contexto teóricamente ilimitada, limitada solo por limitaciones de hardware.

Dicho esto, el método más intensivo en memoria y computacionalmente redundante de los transformadores tiene sus propias ventajas. Por ejemplo, la investigación ha demostrado que los transformadores siguen superando tanto a Mamba como a Mamba-2 en tareas que requieren aprendizaje en contexto (como las instrucciones de pocas tomas), la copia o el razonamiento de contexto largo.

Afortunadamente, las fortalezas respectivas de los transformadores y Mamba no son mutuamente excluyentes. El documento de Mamba-2 sugiere que un modelo híbrido podría superar tanto a los transformadores puros como a los SSM, una noción validada formalmente por la investigación de NVIDIA más adelante en 2024.⁵ En términos generales, los modelos híbridos parecen combinar los beneficios de eficiencia de Mamba con los matices y rendimiento de aprendizaje en contexto, proporcionado por el mecanismo de atención más intensivo en recursos de los transformadores.

Para explorar esto más a fondo, IBM Research colaboró con Dao y Gu, junto con Minjia Zhang de la Universidad de Illinois en Urbana-Champaign (UIUC), en Bamba y Bamba V2. Bamba, a su vez, ha fundamentado muchos de los elementos arquitectónicos de IBM Granite 4.0.

La investigación de modelos híbridos sigue siendo un área de investigación activa, particularmente dentro de la comunidad de código abierto.