State Space Models (SSM)

Un modelo de espacio de estados (SSM) asigna una secuencia de entrada a una secuencia de salida mediante un vector de estado oculto que actualiza en cada paso, por lo que su memoria mantiene el mismo tamaño por mucho que crezca la secuencia. La idea procede de la teoría del control, donde un sistema lineal se expresa mediante una actualización de estado h' = Ah + Bx y una salida y = Ch + Dx, y donde el filtro de Kalman estima esos estados a partir de mediciones ruidosas desde 1960. Los SSM de aprendizaje profundo aprenden A, B y C a partir de los datos y apilan muchas de estas capas. Desde S4, en 2021, y Mamba, en 2023, compiten con los transformers en secuencias largas de texto, audio y datos de series temporales, con un cómputo que crece linealmente con la longitud de la secuencia.

Cómo funcionan los SSM#

Una capa SSM neuronal convierte las ecuaciones continuas en una recurrencia discreta con un tamaño de paso Δ aprendido, mediante la regla de retención de orden cero Ā = exp(ΔA). En cada paso calcula h_t = Ā h_(t-1) + B̄ x_t y y_t = C h_t. Cuando A, B y C se mantienen constantes en todos los pasos temporales, al desplegar la recurrencia se observa que la salida es una convolución de la entrada con el núcleo (CB̄, CĀB̄, CĀ²B̄, ...). Por tanto, la misma capa puede ejecutarse como una convolución durante el entrenamiento, procesando toda la secuencia en paralelo en una GPU, y como una recurrencia durante la generación, con una actualización de estado de tamaño constante por token. En PyTorch, ambas representaciones producen las mismas salidas:

import torch

torch.manual_seed(0)
L, N = 64, 16  # longitud de la secuencia, tamaño del estado
x = torch.randn(L)

# SSM continuo diagonal, discretizado con un tamaño de paso delta (retención de orden cero)
A = -(torch.rand(N) + 0.5)
B, C, delta = torch.randn(N), torch.randn(N), 0.1
A_bar = torch.exp(delta * A)
B_bar = (A_bar - 1) / A * B

# Representación recurrente: una actualización de estado por paso, como en la generación
h, y_rec = torch.zeros(N), []
for t in range(L):
    h = A_bar * h + B_bar * x[t]
    y_rec.append(C @ h)
y_rec = torch.stack(y_rec)

# Representación convolucional: precalcula el núcleo K_k = C . (A_bar^k * B_bar), como en el entrenamiento
K = torch.stack([C @ (A_bar**k * B_bar) for k in range(L)])
y_conv = torch.stack([K[: t + 1].flip(0) @ x[: t + 1] for t in range(L)])

print("Same outputs:", torch.allclose(y_rec, y_conv, atol=1e-5))

Imprime Same outputs: True. La recurrencia conserva solo 16 números de estado en cada paso, mientras que el núcleo precalculado permite al entrenamiento calcular las 64 salidas en paralelo.

De S4 a Mamba#

Los primeros SSM profundos eran demasiado lentos y consumían demasiada memoria para entrenarlos a gran escala. S4, de Albert Gu, Karan Goel y Christopher Ré, dio a la matriz de estado una estructura que permite calcular el núcleo de convolución a bajo coste e incorporó una inicialización diseñada para recordar historiales largos. Alcanzó el estado del arte en todas las tareas del benchmark Long Range Arena, se convirtió en el primer modelo en resolver su tarea Path-X de secuencias de 16K pasos y generaba 60 veces más rápido que un transformer.

Un núcleo fijo trata todas las entradas del mismo modo, así que no puede decidir omitir el contenido de relleno y conservar un dato clave. Mamba, de Albert Gu y Tri Dao, hace que Δ, B y C dependan del token actual, lo que permite al modelo elegir qué guardar en su estado y qué olvidar. Estos parámetros selectivos rompen la representación convolucional, por lo que Mamba calcula la recurrencia mediante un escaneo paralelo adaptado al hardware. Mamba alcanzó cinco veces el rendimiento de inferencia de los transformers, y Mamba-3B superó a transformers de su mismo tamaño e igualó a otros del doble de tamaño. Cada bloque Mamba envuelve el SSM selectivo con una convolución corta y una compuerta SiLU:

Mamba-2 mostró que los SSM selectivos y una forma de atención lineal son dos representaciones del mismo cálculo, un marco que sus autores llaman dualidad de espacio de estados, y lo utilizó para hacer que la capa principal fuera entre 2 y 8 veces más rápida.

SSM frente a arquitecturas relacionadas#

• Transformers: La autoatención compara cada token con todos los demás, así que puede copiar exactamente cualquier token anterior, con un coste que crece cuadráticamente y una caché KV que crece con la ventana de contexto. Un SSM comprime el historial en un estado fijo, lo que resulta más barato pero con pérdida de información. El artículo de 2024 Repeat After Me demostró que los transformers copian y recuperan mejor información de entradas largas que los SSM de estado fijo.
• Redes neuronales recurrentes (RNNs): Una RNN o una LSTM aplica una actualización no lineal, así que el entrenamiento debe recorrer la secuencia token a token. Las actualizaciones de los SSM son lineales en el estado, lo que permite usar una convolución o un escaneo paralelo.
• Modelos ocultos de Markov (HMMs): Un HMM tiene un conjunto finito de estados discretos con probabilidades de transición. El modelo lineal de espacio de estados gaussiano que sustenta el filtro de Kalman es su equivalente continuo, y los SSM profundos conservan ese vector de estado continuo.

Dónde se usan los SSM#

La mayoría de los usos en producción se dan en modelos de lenguaje grandes (LLM) híbridos que conservan unas pocas capas de atención para la recuperación exacta. Jamba de AI21 intercala capas transformer y Mamba con capas de mezcla de expertos (MoE), cabe en una sola GPU de 80 GB y ofrece buenos resultados con contextos de hasta 256K tokens; por su parte, los modelos híbridos Granite 4.0-H de IBM combinan Mamba-2 y capas transformer en una proporción de 9:1. Vision Mamba aplica escaneos selectivos a parches de imagen, y los SSM también modelan secuencias de audio sin procesar y de ADN, demasiado largas para la atención completa. La forma clásica se ejecuta en el seguimiento de Ultralytics YOLO26: el rastreador predeterminado TrackTrack, al igual que ByteTrack y BoT-SORT, predice la siguiente caja de cada objeto con un filtro de Kalman cuyo estado de 8 dimensiones contiene el centro y el tamaño de la caja, así como sus velocidades. El estado se actualiza con una matriz de transición de velocidad constante y una matriz de observación lineal que lee las cuatro coordenadas de la caja.