Phase 10 - Lesson 21

Jamba — SSM-Transformer Híbrido

Los modelos de espacio de estados (SSM) y los transformers quieren cosas diferentes. Los transformers compran calidad mediante atención a un costo cuadrático. Los SSM compran inferência en tiempo lineal y memoria constante a través de una recurrencia, pero se quedan atrás en calidad. Jamba de AI21 (marzo de 2024) y Jamba 1.5 (agosto de 2024) los colocan en el mismo modelo: 1 capa Transformer por cada 7 capas Mamba, MoE en cada dos bloques y una ventana de contexto de 256k que cabe en una sola GPU de 80GB. Mamba-3 (ICLR 2026) ajusta la parte de SSM con espacios de estados de valores complejos y proyecciones MIMO. Esta lección lee ambas arquitecturas de extremo a extremo y explica por qué la receta híbrida ha sobrevivido a tres años de escalado cuando los intentos de contexto largo de SSM puro y Transformer puro no lo han hecho.

Type: Learn Languages: Python (stdlib, calculadora de mezcla de capas) Prerequisites: Phase 10 · 14 (arquitecturas de modelos abiertos), Phase 10 · 17 (atención esparsa nativa) Time: ~60 minutos

Objetivos de Aprendizaje

• Explicar las tres primitivas en un bloque Jamba (capas Transformer, capas Mamba, MoE) y la receta de intercalación 1:7:par.
• Describir cómo es la recurrencia de un SSM a alto nivel y por qué permite inferencia con memoria constante.
• Calcular la huella de caché KV de un modelo Jamba en un contexto de 256k y comparar con lo que necesitaría un modelo Transformer puro.
• Nombrar las tres innovaciones de Mamba-3 (discretización exponencial-trapezoidal, actualización de estado con valores complejos, MIMO) y el problema que cada una busca resolver.

El Problema

La atención es cuadrática con respecto a la longitud de la secuencia. Los modelos de espacio de estados son lineales. Esa diferencia se acumula: a 256k tokens, un mapa de atención de un Transformer tiene 65B de entradas por cabezal; el estado recurrente de un SSM tiene un tamaño fijo independientemente de la longitud de la secuencia.

Los modelos SSM puros (Mamba, Mamba-2) igualan la perplejidad de los Transformers a escalas pequeñas, pero se quedan atrás en tareas de seguimiento de estado (state-tracking) y fallan en algunas categorías de recuperación en contexto (in-context retrieval). La intuición: los SSM comprimen el historial en un estado fijo y, cuando el historial es largo, la información se filtra. La atención recuerda todo exactamente, pero paga un costo cuadrático.

La solución óbvia: usar ambos. Colocar capas Transformer donde la recuperación exacta importa. Usar capas SSM en otros lugares. Ajustar la proporción. Jamba es el primer modelo de nivel de producción que ofrece esta receta híbrida a escala (52B en total, 12B activos, contexto de 256k, una sola GPU de 80GB). Jamba 1.5 extiende la familia a 398B en total / 94B activos. Mamba-3 (ICLR 2026) es la mejor línea base de SSM puro actual en torno a la cual se pueden reconstruir los híbridos.

Esta lección lee los tres artículos y produce el modelo mental para "elegir la proporción correcta".

El Concepto

An SSM in one page

Un modelo de espacio de estados procesa una secuencia x_1, ..., x_N a través de un estado de tamaño fijo h:

h_t = A h_{t-1} + B x_t
y_t = C h_t

En cada paso, el estado evoluciona mediante una dinámica lineal A, toma la entrada B x_t y emite la salida C h_t. A, B, C se pueden aprender. Note la propiedad crítica: calcular y_t requiere solo de h_{t-1} y x_t, no de ningún x anterior. La memoria es constante. La inferencia es O(1) por token.

El truco para la calidad del modelado es la estructura de A. S4 (Gu 2021) utilizó una matriz altamente estructurada que podía evaluarse de manera eficiente como una convolución larga durante el entrenamiento. Mamba (Gu, Dao 2023) reemplazó los parámetros fijos A, B, C por parámetros dependientes de los datos (la parte "selectiva"). Mamba-2 (2024) simplificó aún más la estructura. Mamba-3 (2026) vuelve a agregar complejidad en lugares específicos.

La propiedad clave: para un LLM decodificador, una capa SSM es un reemplazo directo para una capa de atención, con un estado de tamaño fijo por capa en lugar de una caché KV creciente.

El bloque Jamba

Un bloque Jamba intercala capas según dos números:

• l: la proporción de atención a Mamba. Jamba usa l = 8, lo que significa 1 capa Transformer por cada 7 capas Mamba (7 Mamba + 1 Atención = 8 capas por grupo).
• e: la frecuencia de MoE. Jamba usa e = 2, lo que significa que una de cada dos capas aplica MoE.

La secuencia de capas dentro de un bloque:

M  M  M  M  M  M  M  A    (7 Mamba + 1 Attention)
|  M  |  M  |  M  |  M    (where | marks MoE applied)

Cada bloque Jamba consta de 8 capas. Con una profundidad de 4 bloques (32 capas en total), se obtienen 28 capas Mamba y 4 capas de Atención. 16 de ellas utilizan MoE.

Por qué la proporción 1:7

AI21 realizó ablaciones: ¿qué proporción de atención a Mamba ofrece la mejor perplejidad por parámetro Y recuperación en contexto en sus evaluaciones de contexto largo?

• Demasiada atención (1:1): la calidad aumenta pero la memoria y la velocidad se degradan.
• Muy poca atención (1:15): la memoria es excelente pero la recuperación en contexto falla.
• Punto óptimo: 1:7 o 1:8.

La intuición: las capas Transformer manejan la recuperación exacta y el seguimiento de estado. Las capas Mamba manejan la mayor parte del procesamiento a bajo costo.

Codificación posicional

Las capas Mamba son, por sí mismas, conscientes de la posición (a través de la recurrencia). Las capas de atención en los híbridos originales basados en Mamba no utilizaban RoPE; las capas SSM proporcionaban información de posición. Jamba 1.5 añade RoPE a las capas de atención para una generalización de contexto más largo, un refinamento a posteriori basado en la evaluación empírica de contexto largo.

El presupuesto de memoria

Para una configuración Jamba-1 (32 capas: 28 Mamba + 4 Atención, hidden 4096, 32 cabezales de atención):

• Caché KV (solo capas de atención): 2 * 4 * 32 * 128 * 256k * 2 = 8.4 GB en BF16 de 256k. Solo contribuyen las 4 capas de atención.
• Estado SSM: 28 * hidden * state_size por prefijo de token, pero este es un tamaño fijo por capa, no escala con la longitud de la secuencia. El estado típico de Mamba es 16 por característica (feature), hidden 4096: 28 * 4096 * 16 * 2 = 3.7 MB en total.

Compárelo con un Transformer puro de 32 capas, mismo hidden, MHA completo con 32 cabezales: 2 * 32 * 32 * 128 * 256k * 2 = 128 GB en BF16 de 256k. Una reducción de 8x en la caché KV. Incluso en comparación con la línea base GQA(8) que utilizan la mayoría de los modelos de 2024 (2 * 32 * 8 * 128 * 256k * 2 = 32 GB), el híbrido 1:7 de Jamba con 16 GB sigue siendo 2x más pequeño.

Esto es a lo que AI21 se refiere con "contexto de 256k en una sola GPU de 80GB". La caché KV de un Transformer puro con MHA completo no cabría; incluso una línea base GQA no deja espacio para pesos y activaciones; la de Jamba sí.

Mamba-3: la línea base de SSM puro en 2026

Mamba-3 (ICLR 2026, arXiv:2603.15569) introduce tres innovaciones en el lado de SSM puro:

1. Discretización exponencial-trapezoidal. Reemplaza la discretización del método de Euler en Mamba-2 con una recurrencia más expresiva. Se aplica una operación tipo convolución sobre la entrada del estado dentro de la recurrencia central, en lugar de una convolución externa sobre x_t.

2. Actualización de estado con valores complejos. Los Mamba anteriores redujeron la matriz de estado de compleja (S4) a diagonal real (Mamba) y luego a identidad escalada (Mamba-2). Mamba-3 vuelve a agregar valores complejos, lo que equivale a un embedding rotatorio (rotary embedding) dependiente de los datos en el estado. Esto restaura las capacidades de seguimiento de estado que costaron las simplificaciones anteriores de valores reales.

3. Proyecciones multi-entrada multi-salida (MIMO). En lugar de proyecciones escalares por característica, utiliza proyecciones con valores de matriz. Mejora el poder de modelado y la utilización del hardware en tiempo de inferencia sin aumentar la latência de decodificación.

Con 1.5B de parámetros, Mamba-3 mejora la precisión media downstream en 0.6 puntos sobre Gated DeltaNet; la variante MIMO añade 1.2 más para una ganancia total de 1.8 puntos. Con el mismo tamaño de estado, Mamba-3 iguala a Mamba-2 con la mitad del estado.

Mamba-3 aún no se distribuye en un híbrido de producción a escala, pero es el candidato obvio para el lado de SSM del próximo modelo de la clase Jamba.

Cuándo recurrir a un híbrido

Los híbridos ganan cuando:

• El contexto es lo suficientemente largo como para que la caché KV de un Transformer puro se vuelva dolorosa (64k+).
• Las tareas mezclan estructura de corto alcance (ideal para SSM) con recuperación de largo alcance (requiere Transformer).
• Desea realizar el despliegue con presupuestos de memoria de una sola GPU donde la caché KV del Transformer por sí sola no cabría.

Los híbridos pierden cuando:

• El contexto es corto (menos de 16k). La sobrecarga de SSM se desperdicia; un Transformer puro es suficiente.
• Las tareas requieren atención de todos-con-todos (razonamiento profundo, referencia cruzada de múltiples documentos). La escasez de capas de atención en el híbrido perjudica.
• Se está escalando a modelos de frontera de billones de parámetros. Transformer puro + MLA + MoE (estilo DeepSeek-V3) está ganando actualmente la carrera de capacidades.

El panorama competitivo

El panorama de 2026: MoE de Transformer puro domina la frontera, pero los híbridos son dueños del nicho de contexto de más de 256k. Las mejoras de seguimiento de estado de Mamba-3 pueden empujar las proporciones de los híbridos a niveles más bajos (más SSM, menos atención) en la próxima generación.

Utilícelo

code/main.py es una calculadora de memoria para arquitecturas híbridas. Dada una proporción SSM-Transformer y una configuración de tamaño oculto (hidden-size) / conteo de capas, calcula:

• Caché KV en el contexto objetivo.
• Memoria del estado SSM.
• Memoria total en el contexto N para una variedad de configuraciones de modelo.

La calculadora admite:

• Línea base de Transformer puro (la caché KV crece con N).
• Híbrido estilo Jamba 1:7.
• SSM puro (sin caché KV en absoluto).

Los números se toman directamente de los artículos de Jamba-1 y Jamba-1.5 para las configuraciones publicadas y se extrapolan para variantes hipotéticas.

Consideraciones de integración para un despliegue real:

• La mayoría de los servidores de inferencia de producción (vLLM, SGLang) son compatibles con Jamba y Mamba. Verifique la versión específica.
• En un contexto de 256k, la ventaja de memoria de Jamba se hace evidente en el rendimiento (throughput) de solicitudes concurrentes. En la misma VRAM, caben más secuencias de Jamba que de Transformer.
• Mamba-3 como modelo independiente (standalone) aún no se distribuye en producción: vista previa de investigación a 1.5B.

Póngalo en Producción

Esta lección produce outputs/skill-hybrid-picker.md. Dada una especificación de carga de trabajo (perfil de longitud de contexto, combinación de tareas, presupuesto de memoria), recomienda entre un Transformer puro, un híbrido estilo Jamba y un SSM puro, con un razonamiento explícito sobre las compensaciones (tradeoffs) de memoria y calidad.

Ejercicios

1. Ejecute code/main.py para calcular la caché KV en un contexto de 256k para un Transformer puro de 32 capas (hidden 4096, 32 cabezales) y para un híbrido Jamba-1 de la misma configuración. Verifique la reducción de memoria de ~8x que afirma el artículo de AI21.

2. Modifique la calculadora para modelar un híbrido 1:3 (4 Mamba : 1 Atención) y un híbrido 1:15 (14 Mamba : 1 Atención). Grafique la caché KV frente a la proporción. ¿En qué proporción se iguala la caché KV a la memoria del estado SSM?

3. Lea la Sección 3 del artículo de Jamba (arXiv:2403.19887). Explique por qué AI21 utiliza Mamba-1 en lugar de Mamba-2 a pesar de que Mamba-2 es más rápido. Sugerencia: la sección de ablación del híbrido documenta esto.

4. Calcule la sobrecarga de parámetros de MoE-cada-dos-capas en Jamba 1.5 Large (398B en total, 94B activos). Compare la proporción activa con DeepSeek-V3 (37B/671B) y explique por qué la arquitectura de Jamba empuja la proporción activa hacia arriba.

5. Lea la Sección 3 del artículo de Mamba-3 (arXiv:2603.15569). Explique en tres oraciones por qué una actualización de estado de valores complejos equivale a un embedding rotatorio dependente de los datos. Vincule la respuesta a la derivación de RoPE de la Fase 7 · Lección 04.

Términos Clave

Lecturas Adicionales

• Lieber et al. — Jamba: A Hybrid Transformer-Mamba Language Model (arXiv:2403.19887) — el artículo original de Jamba, ablaciones de proporción, afirmación de contexto de 256k
• AI21 — Jamba 1.5: Hybrid Transformer-Mamba at Scale (arXiv:2408.12570) — la familia ampliada, lanzamientos públicos de 398B/94B y 12B/52B
• Gu, Dao — Mamba: Linear-Time Sequence Modeling with Selective State Spaces (arXiv:2312.00752) — el artículo de SSM selectivo en el que se basa Jamba
• Dao, Gu — Mamba-2 (arXiv:2405.21060) — el sucesor simplificado con espacio de estados estructurado
• Lahoti et al. — Mamba-3 (arXiv:2603.15569, ICLR 2026) — estado de valores complejos, MIMO, la frontera del SSM puro de 2026
• Gu et al. — Efficiently Modeling Long Sequences with Structured State Spaces (arXiv:2111.00396) — el artículo de S4, el punto de partida de la genealogía SSM para LLM