Mamba y las SSMs: el retorno del estado persistente
Rubén Rodríguez Abril
Mamba es un modelo revolucionario que, a diferencia de los Transformers, mantiene un estado interno oculto y persistente que evoluciona con cada dato nuevo. Esta ‘memoria’ le permite recordar entre conversaciones. Y procesar secuencias largas (como genomas o audios) de forma más eficiente y rápida, superando a los Transformers en tareas complejas sin perder precisión.
La aparición de los modelos de lenguaje basados en transformers fue verdaderamente revolucionario en el procesamiento del lenguaje natural. Las matrices de atención permitieron capturar relaciones de naturaleza lingüística entre tokens de un texto, incluso cuando se encontraban muy alejados entre sí, superando las limitaciones de las RNNs, propensas a “olvidar” la información leída muchos pasos atrás.
Sin embargo, los transformers, a diferencia de las RNNs, carecen de estado interno que pueda memorizar información más allá de una conversación: su memoria, la caché KV, se restringe al contexto de la interacción en curso. Se trata de un modelo autorregresivo, que en cada paso se alimenta de la misma ventana de contexto, donde reside la secuencia en formación. No pueden, por lo tanto, retener memoria entre conversaciones, a menos que se les adjunte una mecanismo de almacenamiento auxiliar, como sucede en el caso de MemGPT.
SSMs – Modelos de espacios de estados
Para superar esta carencia y crear chatbots dotados de “mente” (estado interno), desde 2021 fueron introducidas en el ámbito del procesamiento del lenguaje natural nuevas arquitecturas, como las SSM (State Space Models) y su variedad más exitosa hasta ahora: Mamba.
En una SSM el estado interno h y la salida y se calculan mediante estas dos ecuaciones procedentes de la teoría de control de sistemas y procesamiento de señales:
ht = Aht-1 + Bxt
yt = Cht + Dxt
donde ht y h t-1 representan el estado interno en los tiempos t y t-1 mientras que xt y yt conforman la entrada y la salida, respectivamente. A, B, C y D son matrices aprendidas, cada una de las cuales tiene las siguientes características:
-A → matriz de estado (state matrix): regula la evolución del estado interno a lo largo del tiempo
-B → matriz de entrada (input matrix): define cómo la entrada influye en el estado interno de la máquina.
-C → matriz de salida (output matrix): define cómo construir la salida.
-D → matriz de transmisión directa (direct transmission matrix): es una suerte de “conexión directa” entre la entrada y la salida.
Figura 1. A la izquierda, se describe el flujo de información en una SSM. La señal (tanto de entrada como de salida) es continua. Existen bucles de retroalimentación, similares a los presentes en una RNN, que son controlados por la matriz A. En esta última, los elementos situados encima de la diagonal son nulos, de tal manera que el pasado puede influir sobre el futuro, pero no al revés. A la derecha del todo, se representa una adaptación de las ecuaciones a representaciones discretas (como es el caso del Deep Learning). Fuente: S4.
Ambas ecuaciones son lineales, y en un modelo de aprendizaje profundo deben de complementarse con una capa no lineal (p.e. función GeLU).
Entre las SSMs más conocidas hasta este momento podemos citar S4 (ideal para modelar secuencias ultra-largas en genómica y audio), H3 (con su mecanismo de memoria comprimida para documentos extensos), Linear RNNs (optimizadas para procesamiento en tiempo real en edge computing) e Hyena (que combina convoluciones con SSMs para síntesis de voz), y Mamba, que será el objeto de la siguiente sección.
Mamba
Aunque las SSMs son capaces de mantener un estado interno, tienen el inconveniente de que sus matrices (equivalentes a los pesos sinápticos de Deep Learning) tienen coeficientes fijos y constantes. Esto contrasta con lo que sucede en los transformers, cuyas matrices de atención son calculadas a partir de las claves (K) y las consultas (Q).
La naturaleza dinámica de estas matrices de atención permite que el modelo, al calcular cada token de salida, centre su atención en un subconjunto concreto de los tokens de entrada. Así es posible construir nuevos vectores/tokens que representen sintagmas, oraciones o relaciones entre diferentes partes del texto. Una unidad de atención puede entenderse, en cierto modo, una suerte de círculo de feedback: La información que entra en el mecanismo es utilizada en primer lugar para calculada la matriz de atención. Y tras ello se hace pasar de nuevo a través de dicha matriz para producir la salida.
Mamba, introducido en diciembre de 2023, pretende tomar lo mejor de ambos mundos.
Sus ecuaciones de estado son similares a las de las SSM descritas con anterioridad, aunque con la siguiente particularidad: Las matrices B y C no son fijas, sino que son funciones de los datos entrada. Para calcularlas introduce un mecanismo de selección, muy parecido al de la atención en transformers, que selecciona los datos de entrada más relevantes en cada paso. Y los utiliza para parametrizar dichas matrices.
Las nueva unidades, reciben el nombre de SSMs selectivas, y sus ecuaciones de estado y de salida (en su versión discretizada, la más difundida en NLP) son las siguientes:
ht = A̅ht-1 + B̅xt
yt = Ctht
A̅ = exp(ΔtA)
B̅ = A-1(exp(ΔtA) – I)Bt
donde Δ es el factor de olvido descrito en el comentario de la ilustración 2. En las ecuaciones, las matrices variables Bt y Ct (dependientes de los datos de entrada) están coloreadas en verde, las matriz fija A, en rojo.
Figura 2. Funcionamiento de una unidad SSM selectiva. En naranja se representa evolución a lo largo del tiempo del estado interno de Mamba, controlado por la matriz A. Δ es un factor de “olvido” y determina cuánto del estado interno anterior se preserva en el estado interno actual (Δt bajo, estado muy inercial, recuerda mucho; Δt alto, estado reactivo, olvida rápido). Modula la matriz A. Las matrices Bt y Ct, que cambian en cada paso, parametrizan la actualiación del estado interno a partir de la entrada, así como la construcción de la salida a partir del propio estado interno. D no aparece y es sustituida por una conexión residual externa. El estado interno actual se almacena en la caché (memoria SRAM) de la GPU. Los datos de entrada y de salida, así como las matrices son almacenados en la memoria externa HBM, más lenta, y se cargan en la memoria SRAM cuando se requieren para el cálculo. Fuente: Mamba.
Dentro de un LLM, la unidad SSM selectiva a menudo juega un papel muy similar al de las unidades de atención en los transformers. Es estrictamente lineal. Y por ello, en sus diferentes implementaciones, viene seguida de un perceptrón multicapas dotado una función de activación no lineal (por lo general, SiLU o ReLU), una transformación lineal y una conexión residual. Una implementación común suele ser la siguiente:
zt = SiLU(yt) (función no lineal)
ot = Woutzt (transformación lineal)
o’t = ot + xt (conexión residual)
Figura 3. A la derecha se describe un bloque de Mamba. En verde, proyecciones lineales que modifican la dimensionalidad de los datos. Los círculos naranjados representan no linealidades. Fuente: Mamba.
Estado interno puramente lineal
A diferencia de lo que sucede en las redes neuronales recurrentes, en los SSMs el estado interno evoluciona de forma estrictamente lineal. Ello tiene la ventaja de que se sus sucesivos valores pueden ser computados anticipadamente mediante el procedimiento de escaneo o scan, lo cual permite que se pueda usar el paralelismo para acelerar notablemente la inferencia. En Mamba, la no linealidad se coloca fuera del núcleo, después de la salida del bloque, y no afecta al estado interno.
La inexistencia de no linealidades impone limitaciones en la capacidad de razonamiento dentro del espacio latente, a las que ya se hizo referencia en otros artículos de esta serie. Según el Teorema de la Aproximación Lineal, una única capa neuronal dotada una función de activación no lineal y con suficiente número de unidades es capaz de simular cualquier función recursiva (computable). En cambio, la ausencia de elementos no lineales provoca que el sistema sólo sea capaz de detectar patrones lineales (p.e. discriminantes) en el espacio latente.
Evaluación
Mamba fue evaluada en múltiples tareas, obteniendo resultados destacados:
–Tareas sintéticas, como las tareas de copia y de inducción, y que miden la capacidad de memorización y generalización del modelo.
–Procesamiento de audio y modelización de secuencias de ADN: Mamba mejoró el rendimiento de los modelos más avanzados del estado de la técnica como ShaShiMi, Hyena o los transformers.
–Procesamiento de lenguaje natural: Mamba es el primer modelo secuencial en tiempo lineal cuyo rendimiento iguala, o incluso supera, al de los transformers. Con modelos escalados hasta 1.000 millones de parámetros, Mamba supera a una amplia cantidad de modelos fundacionales, particularmente a aquellos entrenados con las configuraciones optimizadas de Llama. Tiene una velocidad 5 veces mayor a la de transformers de tamaño similar y la versión Mamba-3B es superior en calidad a transformers del doble de tamaño.
SERIES
