¿Qué es un modelo Mamba?

Los modelos Mamba son quizás la primera arquitectura de deep learning que rivaliza con la eficacia de los modelos 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 hacer coincidir transformadores de tamaño equivalente en importantes evaluaciones de referencia de LLM, a la vez que a menudo es significativamente más eficiente en términos de latencia y requisitos de memoria.

La arquitectura Mamba fue presentada por primera vez por Tri Dao y Albert Gu en el documento de 2023, "Mamba: Linear-Time Sequence Modeling with Selective State Spaces". Un año después, continuaron el artículo original de Mamba con otro artículo que exploraba las conexiones entre los SSM y los transformadores y presentaba una versión refinada y significativamente más rápida de la arquitectura Mamba, a la que llamaron Mamba-2.

Aunque los transformadores han seguido siendo el modo dominante de LLM en los dos años posteriores al lanzamiento del artículo 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, como 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 funcionar 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, basándose en alguna entrada. Conceptual y matemáticamente, están relacionados con las redes neuronales recurrentes (RNN) que dominaban el procesamiento del 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 de Markov ocultos (HMM).

Como su nombre indica, los SSM hacen predicciones sobre el siguiente estado en un sistema dinámico modelando el espacio de estado: una representación matemática de todas las variables de estado que describen el estado de un sistema y el rango de posibilidades para cada una de esas variables en conjunto con entre 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 dos ecuaciones:

• La ecuación de estado,  $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 donde los SSM se usan convencionalmente, como la teoría de control, a menudo se asume que estas matrices son fijas: representan la dinámica de un sistema establecido, y el SSM se usa para encontrar las entradas x que conducen a resultados deseables y. En las concepciones más modernas de los SSM, esas matrices son en sí mismas parámetros que deben optimizarse a través del machine learning. En los modelos de deep learning, esas matrices se representan mediante los pesos 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 influye la entrada, como el siguiente token en una secuencia de texto, en cada variable de estado.

En el modelado del 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) y cómo la entrada influye directamente en la salida (mediada por la matriz D). Como la matriz D es esencialmente externa a la modelización de h(t) en sí, a menudo se omite en los diagramas y discusiones sobre los SSM en favor de centrarse en las matrices principales A, B y C.

En un Mamba LLM, 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 deep learning) 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 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) utilizan el método simple de retenedor de orden cero (ZOH).

La discretización de un SSM permite utilizarlo como un RNN para tareas de secuencia a secuencia. Los parámetros y ecuaciones de un SSM discretizado suelen reescribirse para distinguirlos de sus equivalentes en tiempo continuo, utilizando la notación de subíndices que se suele emplear para las RNN. En esta notación, h_t representa el espacio de estado actualizado que generará el modelo y h_t-1 representa el estado anterior, es decir, el espacio de estado actual.

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

El modelado de datos de texto utilizando SSM discretos estándar no es práctico debido a una serie de deficiencias que comparten con las RNN. Dos de esas deficiencias se abordaron con la introducción de modelos de secuencias espaciales de estados estructurados (o "modelos S4") por parte de Albert Gu et al en 2021: la ineficacia de su entrenamiento y su incapacidad para modelar secuencias largas.

Aunque el éxito de los modelos S4, y sus numerosas variantes, como los SSM diagonales (DSS), los S4 diagonales (S4D) y los modelos H3, allanó directamente el camino para lo que se convirtió en Mamba.

La ventaja de que los SSM discretizados sean equivalentes a una instancia específica de una RNN es que las RNN son extremadamente rápidas en la inferencia. Sin embargo, la desventaja es que las RNN son extremadamente lentos de entrenar.

Por suerte, los SSM discretizados tienen una propiedad importante que las distingue de otras RNN: modelan exclusivamente dependencias lineales. En otras palabras, solo utilizan operaciones simples y sencillas de multiplicación y suma. Como demuestra el documento S4, estas recurrencias lineales simples, repetidas e interdependientes pueden desenrollarse en un kernel de convolución unidimensional, que mapea directamente la entrada x a la salida y en un solo paso. Esto puede calcularse de forma muy eficiente utilizando la transformada rápida de Fourier.

El único "problema" 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 sí durante el entrenamiento. Por lo tanto, un SSM estructurado disfruta 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 a la hora de modelar dependencias de larga distancia. En otras palabras, no se les da bien entender la relación entre pasos de una secuencia que están muy separados, como las palabras al principio y al final de un párrafo, lo que hace que no se les dé bien modelar secuencias largas en general.

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 el comportamiento de las matrices A y B estructurando sus valores iniciales utilizando 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 "Cómo entrenar a su HiPPO"¹.

El documento S4 señala que "la simple modificación de un SSM de una matriz aleatoria A a [la matriz HiPPO] mejoró su rendimiento en el punto de referencia secuencial MNIST de 60 % a 98 %," resolviendo de forma eficaz el problema de memoria a largo plazo de los SSM. Las variaciones posteriores de los SSM estructurados, como DSS, S5 y Mamba, utilizan diferentes esquemas de inicialización (a menudo más sencillos) para A y B que, no obstante, conservan los principios básicos de HiPPO: implementar una estructura diagonal que imponga actualizaciones estables y cierto grado de independencia entre cada valor de la matriz.

En el núcleo de la arquitectura Mamba hay dos innovaciones. El primero es el modelo de espacio de estados selectivo, que proporciona a Mamba una capacidad crucial que antes solo poseían los modelos de transformadores: la capacidad de centrarse selectivamente o ignorar partes específicas del historial de entradas pasadas en función de su relevancia actual. El otro es el escaneo paralelo compatible con el hardware, un algoritmo que optimiza la forma en que una unidad de procesamiento gráfico (GPU) gestiona 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 asignar entrada a salida utilizando todo el historial de entrada.Esto es aceptable o incluso deseable para algunas tareas de modelado de secuencias, pero supone una desventaja importante para la mayoría de las tareas avanzadas de modelado lingüístico.

Para corregir esta incapacidad de omitir o enfatizar dinámicamente partes específicas de su historial de entrada, Dao y Gu propusieron una nueva clase de modelos de espacio de estados con un "escaneo selectivo". 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". A su arquitectura basada en S6 la apodaron "Mamba" porque, entre otras razones, todas esas S suenan como el silbido de una serpiente.

Mamba puede entenderse mejor como una arquitectura de redes neuronales que contiene el modelo espacial de estados selectivos en su esencia. Para hacer una analogía simple, Mamba es para los SSM selectivos como el modelo de 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 temporal lineal (LTI). Para dotar a los SSM de la capacidad de priorizar o despriorizar selectivamente la información específica del pasado en función del contexto actual, Dao y Gu reconfiguraron su SSM de forma que los valores de los parámetros clave del modelo fueran diferentes para distintas entradas.

Más concretamente, los SSM selectivos hacen que el tamaño de 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 el embedding vectorial de x_t a través de tres capas de proyección lineal paralelas, en otras palabras, capas de redes neuronales (o capas MLP). Esto es equivalente a cómo las cabezas de consulta, clave y valor en paralelo generan los respectivos vectores Q, K y V de una entrada en un modelo de transformador.

Multiplicando el embedding vectorial de x_t por los términos de peso y sesgo en esa red de proyección lineal se obtienen los valores resultantes de ∆ _t, B_t y C_t. Los términos de peso y sesgo de las propias capas de proyección lineal se aprenden durante el preentrenamiento del modelo en conjuntos de datos masivos de muestras de texto y, a continuación, (opcionalmente) se refinan mediante un fine-tuning posterior.

• El valor de ∆_t determina la magnitud de la influencia de x_t en la memoria del modelo del contexto que ha visto hasta ahora: en otras palabras, en cuánto se actualizará del estado oculto h_t-1 a h_t. Un tamaño de paso mayor ∆_t produce mayores cambios y acelera el decaimiento (en otras palabras, el "olvido") de la información más antigua contenida en el 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 x_t es un token para la palabra "ayer", B_t podría ajustarse de forma que actualice el estado para reflejar que el contexto resultante probablemente pertenece al pasado.
  
  
• Los cambios en la matriz C_t determinan cómo se traduce esta información contextual en influencia sobre la salida del modelo y_t. Continuando el ejemplo en el que x_k es un token para “ayer”, C_k podría verse influenciado de una manera que haga que los verbos que posteriormente genere el modelo se conjuguen en tiempo pasado.

En particular, no se realizan tales ajustes basados en entrada 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 la gestionan 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 núcleo de transición ya no es constante: el quid del mecanismo de selectividad es que la transición de h_t-1 a h_t ahora depende del contexto.

En lugar de ello, Mamba utiliza una solución alternativa inteligente para lograr beneficios de paralelización similares. Dado que el SSM utiliza únicamente multiplicación y suma, sus cálculos están sujetos a la propiedad asociativa habitual de las matemáticas: pueden agruparse de distintas maneras sin modificar el resultado final. Esto permite dividir los numerosos cálculos secuenciales en fragmentos pequeños e independientes que pueden procesarse en paralelo mediante una GPU a través de un escaneo de suma de prefijos paralelo.

Además, los resultados se combinan de una manera jerárquica específica que hace un uso óptimo y 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 redes neuronales, antes de la salida que realiza la predicción de salida final 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 general son relevantes para esa entrada, sino determinar en qué medida esa información contextual debe modificar el significado original de la entrada.

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

• Primero, la entrada pasa a través de una capa lineal que es dos veces más ancha que el vector de entrada en sí, proyectándola a un espacio de mayor dimensión. Por ejemplo, si el modelo representa originalmente cada token de entrada x como un embedding de vectores de 512 dimensiones, al multiplicar x por los pesos de la capa de proyección lineal se expande a un vector de 1024 dimensiones.
  
  
• A continuación, el vector expandido se divide en dos. La mitad (que llamaremos x_proj) conduce a la ruta que pasa por el SSM y la otra mitad (que llamaremos z_proj) pasa a una ruta diferente que pasa por un mecanismo de compuerta. Para mayor claridad, el paso de expansión anterior se suele representar como si lo realizaran dos capas lineales paralelas.
  
  
• Antes de que x_proj llegue al SSM, se introduce en una capa de convolución unidimensional. Esta capa de convolución extrae patrones locales (como dependencias entre tokens vecinos, como simples emparejamientos verbo-sujeto). Esto permite al SSM "centrarse" en la comprensión contextual de las dependencias globales a largo plazo.
  
  
• La salida de la capa de convolución sirve como la entrada a 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 documento Mamba utiliza la unidad lineal sigmoidea (SiLU). Llamaremos al vector resultante x_act.
  
  
• Mientras tanto, en la ruta del mecanismo de activación independiente, z_proj también se introduce en una función de activación no lineal, lo que da como resultado z_act.
  
  
• En la ruta SSM, x_act se introduce en tres capas de proyección lineal paralelas que generan los valores respectivos de_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 del SSM y.
  
  
• El vector de salida del SSM y ahora se multiplica elemento por elemento por el vector de salida de la ruta de acceso z_act. Básicamente, 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 dará 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 por y amplificará la influencia de su información contextual.
  
  
• El vector resultante se proyecta de nuevo a su tamaño original. Se puede entender como un vector de actualizaciones contextuales ponderadas (o no actualizaciones) para cada uno de los elementos del vector de entrada original.
  
  
• Por último, ese vector de actualizaciones se añade a la copia del vector de entrada original que se envió directamente al final del bloque como 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 Mamba o, en las capas finales del modelo, servir como entrada para una función softmax que genera 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 aportó tres contribuciones importantes:

• Una exploración de las conexiones teóricas entre Mamba y los transformadores junto a 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 "SSD mínima" que implementa el SSM selectivo en unas 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 igualaron o superaron definitivamente a los modelos Mamba y transformadores del tamaño correspondiente en una serie de tareas posteriores.

Como se afirma en la introducción 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 produjo un nuevo marco conceptual que unía los dos, al que llamaron "dualidad del espacio de estado" (SSD)³. Al hacerlo, abrieron la puerta para que Mamba pudiera obtener el beneficio de varios años de exploración y optimización de la arquitectura del transformador.

Un beneficio notable fue el desarrollo de un equivalente Mamba de la atención multicabezal (MHA), en el que un bloque Mamba se puede dividir en varias "cabezas Mamba" similares a las múltiples "cabezales de atención" de 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 bloque Mamba paralelo (a diferencia del bloque Mamba "secuencial" original), los parámetros dependientes de la entrada ∆, B y C se generan en paralelo en la capa de proyección inicial. B y C, en concreto, se derivan simplemente copiando partes 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 propios puntos fuertes, pero los modelos basados en Mamba son generalmente superiores en todos los aspectos relacionados con el uso de la memoria y la velocidad: según el documento de Mamba, Mamba ofrece un rendimiento cinco 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 fine-tuning), 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 solo se escala 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 únicamente 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 few-shot prompting), la copia o el razonamiento de contexto largo.

Afortunadamente, los puntos fuertes respectivos de los transformadores y Mamba no son mutuamente excluyentes. El artículo Mamba-2 sugiere que un modelo híbrido podría superar a los transformadores puros o a los SSM, una idea validada formalmente por un estudio de NVIDIA a finales de 2024⁵. En términos generales, los modelos híbridos parecen combinar los beneficios de eficiencia de Mamba con el rendimiento de aprendizaje matizado y contextual que proporciona el mecanismo de atención de los transformadores, que consume más recursos.

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 informado 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, especialmente dentro de la comunidad de código abierto.