Phase 10 - Lesson 21

Jamba — SSM-Transformer Híbrido

Modelos de espaço de estados (SSMs) e transformers querem coisas diferentes. Os transformers compram qualidade via atenção a um custo quadrático. Os SSMs compram inferência em tempo linear e memória constante via recorrência, mas ficam atrás em qualidade. O Jamba da AI21 (março de 2024) e o Jamba 1.5 (agosto de 2024) colocam ambos no mesmo modelo: 1 camada Transformer para cada 7 camadas Mamba, MoE em cada dois blocos e uma janela de contexto de 256k que cabe em uma única GPU de 80GB. O Mamba-3 (ICLR 2026) aprimora a parte de SSM com espaços de estados de valores complexos e projeções MIMO. Esta lição lê ambas as arquiteturas de ponta a ponta e explica por que a receita híbrida sobreviveu a três anos de escala quando as tentativas de contexto longo de SSM puro e Transformer puro não sobreviveram.

Type: Learn Languages: Python (stdlib, calculadora de mistura de camadas) Prerequisites: Phase 10 · 14 (arquiteturas de modelos abertos), Phase 10 · 17 (atenção esparsa nativa) Time: ~60 minutos

Objetivos de Aprendizado

• Explicar as três primitivas em um bloco Jamba — camadas Transformer, camadas Mamba, MoE — e a receita de intercalação 1:7:par.
• Apresentar como é a recorrência de um SSM em alto nível e por que ela permite inferência com memória constante.
• Calcular a pegada do cache KV de um modelo Jamba em um contexto de 256k e comparar com o que um modelo Transformer puro precisaria.
• Nomear as três inovações do Mamba-3 (discretização exponencial-trapezoidal, atualização de estado com valores complexos, MIMO) e o problema que cada uma visa resolver.

O Problema

A atenção é quadrática no comprimento da sequência. Os modelos de espaço de estados são lineares. Essa diferença se acumula: com 256k tokens, um mapa de atenção de um Transformer tem 65B de entradas por cabeça; o estado recorrente de um SSM tem tamanho fixo, independentemente do comprimento da sequência.

Modelos SSM puros (Mamba, Mamba-2) igualam a perplexidade de Transformers em escalas pequenas, mas ficam para trás em tarefas de rastreamento de estado (state-tracking) e falham em algumas categorias de recuperação em contexto (in-context retrieval). A intuição: os SSMs comprimem o histórico em um estado fixo e, quando o histórico é longo, há vazamento de informações. A atenção lembra de tudo exatamente, mas paga um custo quadrático.

A solução óbvia: use ambos. Coloque camadas Transformer onde a recuperação exata importa. Use camadas SSM em outros lugares. Ajuste a proporção. O Jamba é o primeiro modelo de nível de produção a entregar essa receita híbrida em escala (52B no total, 12B ativos, contexto de 256k, uma única GPU de 80GB). O Jamba 1.5 estende a família para 398B no total / 94B ativos. O Mamba-3 (ICLR 2026) é a melhor linha de base de SSM puro atual em torno da qual os híbridos podem ser reconstruídos.

Esta lição lê os três artigos e produz o modelo mental para "escolher a proporção correta".

O Conceito

Um SSM em uma página

Um modelo de espaço de estados processa uma sequência x_1, ..., x_N por meio de um estado de tamanho fixo h:

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

A cada passo, o estado evolui através de uma dinâmica linear A, recebe a entrada B x_t e emite a saída C h_t. A, B, C podem ser aprendidos. Observe a propriedade crítica: calcular y_t precisa apenas de h_{t-1} e x_t, e não de qualquer x anterior. A memória é constante. A inferência é O(1) por token.

O truque para a qualidade da modelagem é a estrutura de A. O S4 (Gu 2021) utilizou uma matriz altamente estruturada que podia ser avaliada eficientemente como uma convolução longa durante o treinamento. O Mamba (Gu, Dao 2023) substituiu os parâmetros fixos A, B, C por parâmetros dependentes de dados (a parte "seletiva"). O Mamba-2 (2024) simplificou ainda mais a estrutura. O Mamba-3 (2026) adiciona complexidade novamente em locais específicos.

A propriedade fundamental: para um LLM decodificador, uma camada SSM é um substituto direto para uma camada de atenção, com um estado por camada de tamanho fixo em vez de um cache KV crescente.

O bloco Jamba

Um bloco Jamba intercala camadas de acordo com dois números:

• l: a proporção de atenção para Mamba. O Jamba usa l = 8, o que significa 1 camada Transformer para cada 7 camadas Mamba (7 Mamba + 1 Atenção = 8 camadas por grupo).
• e: a frequência de MoE. O Jamba usa e = 2, o que significa que uma a cada duas camadas aplica MoE.

A sequência de camadas dentro de um bloco:

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

Cada bloco Jamba possui 8 camadas. Com 4 blocos de profundidade (32 camadas no total), você obtém 28 camadas Mamba e 4 camadas de Atenção. 16 delas utilizam MoE.

Por que a proporção 1:7

A AI21 executou ablações: qual proporção de atenção para Mamba oferece a melhor perplexidade por parâmetro E recuperação em contexto em suas avaliações de contexto longo?

• Muita atenção (1:1): a qualidade aumenta, mas a memória e a velocidade degradam.
• Pouca atenção (1:15): a memória é excelente, mas a recuperação em contexto falha.
• Ponto ideal: 1:7 ou 1:8.

A intuição: as camadas Transformer lidam com a recuperação exata e o rastreamento de estado. As camadas Mamba lidam com a maior parte do processamento de forma barata.

Codificação posicional

As camadas Mamba são, por si só, cientes de posição (via recorrência). As camadas de atenção nos híbridos originais baseados em Mamba não utilizavam RoPE — as camadas SSM forneciam informações de posição. O Jamba 1.5 adiciona RoPE às camadas de atenção para generalização de contexto mais longo, um refinamento a posteriori baseado em avaliação empírica de contexto longo.

O orçamento de memória

Para um formato Jamba-1 (32 camadas: 28 Mamba + 4 Atenção, hidden 4096, 32 cabeças de atenção):

• Cache KV (apenas camadas de atenção): 2 * 4 * 32 * 128 * 256k * 2 = 8.4 GB em BF16 de 256k. Apenas as 4 camadas de atenção contribuem.
• Estado SSM: 28 * hidden * state_size por prefixo de token, mas este é um tamanho fixo por camada, não escalonando com o comprimento da sequência. Um estado típico do Mamba é 16 por recurso (feature), hidden 4096: 28 * 4096 * 16 * 2 = 3.7 MB no total.

Compare com um Transformer puro com 32 camadas, mesmo hidden, MHA completo com 32 cabeças: 2 * 32 * 32 * 128 * 256k * 2 = 128 GB em BF16 de 256k. Uma redução de 8x no cache KV. Mesmo contra a linha de base GQA(8) usada pela maioria dos modelos de 2024 (2 * 32 * 8 * 128 * 256k * 2 = 32 GB), o híbrido 1:7 do Jamba com 16 GB ainda é 2x menor.

É isso o que a AI21 quer dizer com "contexto de 256k em uma única GPU de 80GB". O cache KV de um Transformer puro com MHA completo não caberia; mesmo uma linha de base GQA não deixa espaço para pesos e ativações; o do Jamba deixa.

Mamba-3: a linha de base de SSM puro em 2026

O Mamba-3 (ICLR 2026, arXiv:2603.15569) introduz três inovações no lado do SSM puro:

1. Discretização exponencial-trapezoidal. Substitui a discretização do método de Euler no Mamba-2 por uma recorrência mais expressiva. Operação semelhante a uma convolução aplicada na entrada de estado dentro da recorrência central, em vez de uma convolução externa em x_t.

2. Atualização de estado com valores complexos. Os Mambas anteriores reduziram a matriz de estado de complexa (S4) para diagonal real (Mamba) e depois para identidade escalonada (Mamba-2). O Mamba-3 adiciona novamente valores complexos — o equivalente a um embedding rotativo (rotary embedding) dependente de dados no estado. Isso restaura capacidades de rastreamento de estado que as simplificações de valores reais anteriores haviam custado.

3. Projeções multi-entrada multi-saída (MIMO). Em vez de projeções escalares por recurso, usa projeções com valores de matriz. Melhora o poder de modelagem e a utilização do hardware em tempo de inferência sem aumentar a latência de decodificação.

Com 1,5B de parâmetros, o Mamba-3 melhora a precisão média downstream em 0,6 pontos em relação ao Gated DeltaNet; a variante MIMO adiciona mais 1,2 para um ganho total de 1,8 pontos. Com o mesmo tamanho de estado, o Mamba-3 se iguala ao Mamba-2 com metade do estado.

O Mamba-3 ainda não está sendo entregue em um híbrido de produção em escala — mas é o candidato óbvio para o lado SSM do próximo modelo da classe Jamba.

Quando recorrer a um híbrido

Os híbridos vencem quando:

• O contexto é longo o suficiente para que o cache KV do Transformer puro se torne doloroso (64k+).
• As tarefas misturam estrutura de curto alcance (boa para SSM) com recuperação de longo alcance (precisa de Transformer).
• Você deseja implantar em orçamentos de memória de GPU única onde o cache KV do Transformer sozinho não caberia.

Os híbridos perdem quando:

• O contexto é curto (menos de 16k). A sobrecarga do SSM é desperdiçada; o Transformer puro é suficiente.
• As tarefas precisam de atenção de todos-para-todos (raciocínio profundo, referência cruzada de múltiplos documentos). A esparsidade das camadas de atenção no híbrido prejudica.
• Você está escalando para modelos de fronteira de trilhões de parâmetros. Transformer puro + MLA + MoE (estilo DeepSeek-V3) está atualmente vencendo a corrida de capacidade.

O cenário competitivo

O cenário de 2026: MoE de Transformer puro domina a fronteira, mas os híbridos são donos do nicho de contexto de mais de 256k. As melhorias de rastreamento de estado do Mamba-3 podem empurrar as proporções híbridas para patamares mais baixos (mais SSM, menos atenção) na próxima geração.

Use-o

code/main.py é uma calculadora de memória para arquiteturas híbridas. Dada uma proporção SSM-Transformer e uma configuração de tamanho oculto (hidden-size) / contagem de camadas, ela calcula:

• Cache KV no contexto alvo.
• Memória do estado SSM.
• Memória total no contexto N para uma variedade de formatos de modelo.

A calculadora suporta:

• Linha de base de Transformer puro (o cache KV cresce com N).
• Híbrido estilo Jamba 1:7.
• SSM puro (sem nenhum cache KV).

Os números são diretos dos artigos do Jamba-1 e Jamba-1.5 para os formatos publicados e extrapolados para variantes hipotéticas.

Considerações de integração para uma implantação real:

• A maioria dos servidores de inferência de produção (vLLM, SGLang) suporta Jamba e Mamba. Verifique a versão específica.
• No contexto de 256k, a vantagem de memória do Jamba aparece no throughput de requisições concorrentes. Na mesma VRAM, você consegue encaixar mais sequências do Jamba do que sequências do Transformer.
• O Mamba-3 como um modelo independente (standalone) ainda não está disponível em produção — pré-visualização de pesquisa com 1,5B.

Coloque em Produção

Esta lição produz outputs/skill-hybrid-picker.md. Dada uma especificação de carga de trabalho (perfil de comprimento de contexto, mix de tarefas, orçamento de memória), ele recomenda entre um Transformer puro, um híbrido estilo Jamba e um SSM puro, com justificativa explícita sobre as compensações (tradeoffs) de memória e qualidade.

Exercícios

1. Execute code/main.py para calcular o cache KV em um contexto de 256k para um Transformer puro de 32 camadas (hidden 4096, 32 cabeças) e para um híbrido Jamba-1 com o mesmo formato. Verifique a redução de memória de ~8x reivindicada no artigo da AI21.

2. Modifique a calculadora para modelar um híbrido 1:3 (4 Mamba : 1 Atenção) e um híbrido 1:15 (14 Mamba : 1 Atenção). Plote o cache KV versus a proporção. Em qual proporção o cache KV se iguala à memória do estado SSM?

3. Leia a Seção 3 do artigo do Jamba (arXiv:2403.19887). Explique por que a AI21 usa o Mamba-1 em vez do Mamba-2, apesar de o Mamba-2 ser mais rápido. Dica: a seção de ablação do híbrido documenta isso.

4. Calcule a sobrecarga de parâmetros do MoE em cada duas camadas no Jamba 1.5 Large (398B no total, 94B ativos). Compare a proporção ativa com o DeepSeek-V3 (37B/671B) e explique por que a arquitetura do Jamba empurra a proporção ativa para cima.

5. Leia a Seção 3 do artigo do Mamba-3 (arXiv:2603.15569). Explique em três frases por que uma atualização de estado com valores complexos é equivalente a um embedding rotativo dependente de dados. Relacione a resposta com a derivação do RoPE da Fase 7 · Lição 04.

Termos-Chave

Leitura Adicional

• Lieber et al. — Jamba: A Hybrid Transformer-Mamba Language Model (arXiv:2403.19887) — o artigo original do Jamba, ablações de proporção, alegação de contexto de 256k
• AI21 — Jamba 1.5: Hybrid Transformer-Mamba at Scale (arXiv:2408.12570) — a família ampliada, lançamentos públicos de 398B/94B e 12B/52B
• Gu, Dao — Mamba: Linear-Time Sequence Modeling with Selective State Spaces (arXiv:2312.00752) — o artigo do SSM seletivo no qual o Jamba se baseia
• Dao, Gu — Mamba-2 (arXiv:2405.21060) — o sucessor simplificado com espaço de estados estruturado
• Lahoti et al. — Mamba-3 (arXiv:2603.15569, ICLR 2026) — estado com valores complexos, MIMO, a fronteira do SSM puro de 2026
• Gu et al. — Efficiently Modeling Long Sequences with Structured State Spaces (arXiv:2111.00396) — o artigo do S4, o ponto de partida da genealogia SSM para LLMs