English · Español
03 — Modelos de espacio de estados: RWKV, S4, Mamba y el selective scan¶
🇪🇸 La atención escala \(O(T^2)\) con la longitud de la secuencia. Los modelos de espacio de estados (SSM) la reemplazan por una recurrencia: cada token actualiza un estado oculto, sin atención al pasado. Tiempo lineal, memoria constante por capa. El "secret sauce" de Mamba sobre los S4 anteriores: los pesos de la recurrencia dependen del input (selectividad).
El coste de la attention es el elefante en cualquier sala de contexto largo. El cómputo es \(O(T^2 \cdot d_{\text{model}})\) para una longitud de contexto \(T\). La KV cache crece linealmente en \(T\). Para \(T = 100k\), estás en territorio serio de memoria y tiempo; para \(T = 1M\), es imposible sin hardware especializado.
La familia de modelos de espacio de estados (SSM) — incluyendo RWKV, S4, S5 y Mamba — esquiva la attention por completo. La representación de cada token se computa como una actualización recurrente de estado, no como un lookup suave sobre el pasado. Tiempo lineal. Memoria constante por capa. Sin KV cache.
Para la Fase 36 nos centramos en Mamba — la variante competitiva con transformers — y en el truco que la hace funcionar: el selective scan.
El punto de partida: recurrencia lineal¶
Una recurrencia lineal simple sobre una secuencia de entrada \(x_1, \ldots, x_T\):
con \(A \in \mathbb{R}^{N \times N}, B \in \mathbb{R}^{N \times d_{\text{in}}}, C \in \mathbb{R}^{d_{\text{out}} \times N}\), \(h_t \in \mathbb{R}^N\) el estado oculto.
Memoria: \(h_t\) es el único estado necesario para \(h_{t+1}\) — constante \(O(N)\). Tiempo: \(O(T)\) para la secuencia completa.
Comparar: un LSTM estándar es una recurrencia no lineal; una capa de attention no tiene recurrencia (cada interacción par de tokens es directa).
Entonces, ¿por qué no son los SSM la elección obvia? Dos razones:
- Sin no linealidad en la actualización de estado, la capacidad es limitada. Los SSM lineales vainilla rinden peor que los LSTM y que la attention.
- El parameter sharing entre timesteps significa que el modelo es independiente del input en la dinámica de la recurrencia. Las puertas de un LSTM dependen del input; aquí \(A, B, C\) son fijos.
S4 (Gu et al., 2022) abordó (1) con matrices \(A\) estructuradas (HiPPO) que tienen propiedades agradables en el dominio de la frecuencia. S4 fue competitivo en benchmarks de largo alcance, pero aún no en modelado de lenguaje.
Mamba (Gu y Dao, 2023) abordó (2) con selectividad.
El mecanismo selectivo de Mamba¶
Mamba parametriza \(B\), \(C\) y un tamaño de paso de discretización \(\Delta\) como dependientes del input:
Ahora la dinámica de la recurrencia cambia en cada token, condicionada al input. El modelo puede:
- "Abrir la puerta" para un token relevante (gran \(\Delta\), gran \(B\)): el estado oculto absorbe el nuevo input.
- "Cerrar la puerta" para un token irrelevante (pequeño \(\Delta\), pequeño \(B\)): el estado oculto se mantiene aproximadamente igual.
- Ajustar la proyección de salida \(C\) por token.
Esto es Mamba. El "selective" en su nombre es exactamente esta parametrización dependiente del input.
Discretización¶
Mamba define un sistema lineal en tiempo continuo \(\dot{h}(t) = A h(t) + B x(t)\) y lo discretiza por zero-order hold:
\(\Delta\) controla la escala temporal: un \(\Delta\) mayor significa que cada timestep avanza el sistema más en tiempo continuo, dando más peso al nuevo input.
En la práctica \(A\) se parametriza como una matriz diagonal, por lo que \(\exp(\Delta A)\) es elemento a elemento — un exponencial matricial tratable.
Parallel scan¶
La recurrencia \(h_t = \bar{A}_t h_{t-1} + \bar{B}_t x_t\) parece secuencial, pero es una operación asociativa: existe un algoritmo de parallel scan que computa todos los \(h_t\) con profundidad \(O(\log T)\) en GPU. La implementación oficial de Mamba tiene un kernel CUDA personalizado para esto; mamba-minimal usa una versión más lenta en Python/JAX.
Esto es lo que hace a Mamba rápido en entrenamiento a pesar de la recurrencia: el parallel scan explota el paralelismo de la GPU de una forma que las RNN no pueden.
Coste de inferencia¶
En tiempo de decode, Mamba es genuinamente constante por token:
- Actualizar \(h_t = \bar{A}_t h_{t-1} + \bar{B}_t x_t\): \(O(N \cdot d_{\text{in}})\) — independiente de la longitud pasada.
- Salida \(y_t = C_t h_t\): \(O(N \cdot d_{\text{out}})\).
Total: \(O(N \cdot (d_{\text{in}} + d_{\text{out}}))\), sin dependencia de la longitud pasada de la secuencia.
Memoria por capa en inferencia: \(h_t \in \mathbb{R}^N\). Sin KV cache creciente.
Para los modelos liberados de Mamba (~3B parámetros, \(N = 64\)): la inferencia por token es \(\sim 200\)k FLOPs — comparable a un transformer pequeño con contexto modesto. La ganancia aparece en contexto largo. Un decode de Mamba a 100k tokens tiene la misma velocidad que un decode de Mamba a 100 tokens. Un decode de transformer a 100k tokens es 1000× más lento que un decode a 100 tokens.
Modelos híbridos (Jamba)¶
Mamba puro tiene sus propios límites: empíricamente, lo hace peor que los transformers en tareas que requieren recall exacto (lookup in-context de un token específico). La hipótesis: la recurrencia comprime información en el estado; recuperar exactamente un token pasado específico es más difícil que el lookup directo de la attention.
Jamba (AI21, 2024) intercala bloques de Mamba con bloques de attention. La mayoría de las capas son Mamba (baratas, rápidas); unas pocas son attention (capacidad). Lo mejor de ambos: coste total cercano a lineal, recall exacto cuando hace falta.
Mencionarlo; no derivarlo. El principio de hibridación es lo que hay que llevarse.
RWKV — un primo¶
RWKV ("Receptance Weighted Key Value") es otra familia SSM — mismo objetivo (recurrencia en tiempo lineal con capacidad similar a la attention), mecánica distinta. Usa un patrón "time-mix" + "channel-mix" sin attention explícita. Open-source y dirigido por la comunidad; competitivo con transformers de tamaño medio; iteración continua.
Para la Fase 36: mencionar que RWKV existe, apuntar a su web. No derivar.
¿Ayudaría Mamba al tutor de gramática?¶
El test:
- Cuello de botella de Mamba: "el tiempo cuadrático de la attention y el crecimiento del KV cache en contexto muy largo".
- ¿El tutor de gramática lo tiene? Contexto máximo: 32 tokens. Coste de attention: \(32^2 \cdot d_{\text{model}}\). Con \(d_{\text{model}} = 64\), eso son 131k FLOPs — despreciable. KV cache para 32 tokens a esta escala: \(\sim 16\) KB en total. Sin cuello de botella.
- ¿Qué costaría Mamba? Reemplazar la attention por selective scan en
src/minimodel/transformer.py: reescritura sustancial. Nueva superficie de tuning (\(N\), escalado de \(\Delta\), parametrización de la selectividad). Nuevos modos de fallo (colapso de estado, inestabilidad). - Veredicto: Nunca. Las tareas del tutor de gramática implican acuerdo sintáctico de corto alcance (sujeto-verbo-tiempo), exactamente donde la propiedad de lookup directo de la attention es más útil que recurrencia-y-recall.
Punto sutil sobre Mamba y la gramática¶
Las tareas de gramática premian específicamente la attention. Una corrección de conjugación verbal necesita atender a la persona y número del sujeto — un lookup preciso y directo de un token pasado específico. Mamba puede aprender esto a través de su estado, pero lo hace de forma menos nítida que el mecanismo de puntería de la attention. Para el tutor de gramática, la attention es positivamente la elección correcta — no solo "bien para mantener", sino realmente mejor que las alternativas. Es un resultado más bonito que "no ayuda"; es "el diseño original era correcto y ahora podemos defenderlo".
Cuándo ayudaría Mamba (el contrafáctico)¶
Un agente de corrección gramatical con contexto de 1M tokens (proofreading de un libro, parsing de novelas enteras para consistencia de tiempos): la KV cache de la attention se vuelve prohibitiva. La recurrencia de memoria constante de Mamba se vuelve entonces atractiva. Pero también aceptarías cierta caída de calidad en queries de recall exacto — un tradeoff que requiere un diseño híbrido (estilo Jamba).
Fuera de alcance, pero Borja debería poder esbozar esta conversación al final de la fase.
Lo que esta fase NO cubre¶
- Implementar Mamba. Solo lectura vía
mamba-minimal. La codebase completa de Mamba tiene kernels CUDA; no está en nuestro terreno. - Derivación de S4 y HiPPO. Mencionado; la derivación está en el paper original. La Fase 36 se detiene en "S4 fue el predecesor; Mamba añadió selectividad".
- RWKV con profundidad matemática. Solo mención.
- Entrenar un Mamba sobre el corpus de gramática. Llevaría semanas de CPU y produciría un resultado claramente peor. La matemática en §"¿Ayudaría Mamba?" es la respuesta; no hace falta gastar el cómputo.
- Arquitecturas híbridas attention+Mamba. Solo vocabulario (Jamba).
Siguiente: theory/04-speculative-and-reasoning.md.