DL0171 Jamba Hybrid SSM-Transformer Architecture

Jamba (AI21 Labs, 2024) is a decoder-only language model whose 32 layers are built from a repeating 8-layer block that contains seven Mamba mixers and one attention mixer, with a mixture-of-experts feed-forward on every other layer (16 experts, top-2 routing, 52B total and 12B active parameters). Every layer still writes into the same residual stream, so the only thing that changes from layer to layer is which token-mixing operator runs: a selective state space recurrence that is linear in sequence length and keeps a fixed-size state, or full self-attention that is quadratic in prefill and keeps a growing KV cache. The blend exists because the two operators fail in opposite directions. A pure Mamba stack of this size does not reliably form induction heads, so it degrades at few-shot in-context learning, format copying, and verbatim retrieval, while a pure Transformer of the same depth carries a KV cache eight times larger and loses long-context throughput. Four attention layers out of thirty-two are enough to restore the copy and in-context behaviour, and the remaining twenty-eight Mamba layers deliver the memory and speed: at a 256K-token context the KV cache is 4 GB instead of 32 GB, and long-context throughput is roughly 3x that of a comparable all-attention MoE model.

(1) Two Mixers, One Residual Stream: Mamba and attention layers are interchangeable drop-ins at the token-mixing position, so no fusion or adapter machinery is needed to combine them.
(2) A 1:7 Attention Ratio: one attention layer per eight-layer block was chosen after ablations showed 1:3 and 1:7 score alike, so the cheaper ratio wins.
(3) MoE Buys Capacity Without FLOPs: replacing the MLP on every second layer with 16 experts and top-2 routing raises total parameters to 52B while keeping 12B active per token.
(4) Attention Supplies Exact Recall: the few attention layers are the only components that can address an arbitrary earlier token exactly, which is what induction heads and few-shot copying need.
(5) Mamba Supplies Position: the recurrence is inherently ordered, so Jamba ships with no explicit positional encoding and RoPE gave no measurable gain.
(6) The Payoff Is Memory, Not Perplexity: a 256K context fits in 4 GB of KV cache, letting a 52B model serve 140K tokens on a single 80 GB GPU in int8.

Figure 1: The whole architecture is one repeated block. Attention appears at layer 4 of every 8, giving 4 attention layers out of 32, and the MoE feed-forward alternates with a dense MLP so that capacity grows without raising the per-token FLOP count.

The reason neither operator survives alone is a difference in what they can store. A selective SSM compresses the entire prefix into a state of fixed size, so its recall is lossy and content-addressed by whatever the gating learned to keep. Theoretical and empirical work on copying shows that a fixed-state recurrent model needs state size proportional to the string it must reproduce, whereas attention copies with a constant number of heads. That is exactly the gap seen in Jamba’s ablations: a pure Mamba model trained on the same 250B tokens tracks the Transformer on log-probability benchmarks yet collapses on few-shot tasks where the model must imitate the label format shown in the prompt, because it never develops induction heads. Inserting attention into one layer in eight repairs this, and the repaired model then inherits Mamba’s cost profile for the other seven eighths of its depth. Ordering matters too, since spreading the attention layers evenly through the stack lets every group of Mamba layers be followed by an exact-lookup step rather than concentrating all lookup capability at one depth.

Mathematical Formulation:

Where:

• is the SSM state at step , the layer input and the layer output; is the skip term.
• are the discretized transition, input, and output matrices. In a selective SSM they depend on the current token, which is what lets the layer decide what to keep and what to forget.
• is the input-dependent step size; a large overwrites the state with the new token, a small one carries the old state forward.
• is sequence length, the model width, and the state dimension (16 in Mamba), so SSM cost is linear in while attention prefill is quadratic.
• is bytes per element (2 in fp16), the number of attention layers only, the KV heads under GQA, and the head dimension; the leading 2 counts keys and values.
• The hybrid changes exactly one factor in , namely , which drops from 32 to 4.

KV Cache At A 256K Context (fp16, 8 KV heads, ):

The 28 Mamba layers contribute a per-sequence state of only a few megabytes that does not grow with at all, so the entire cache curve of the hybrid is set by its four attention layers. That is what turns long context from a memory problem into a compute problem: batch size at 256K stops being limited by cache residency, and decode throughput at 128K measures roughly 3x a comparable all-attention MoE because far fewer bytes move per generated token.

Figure 2: Only the attention layers create a cache that grows with context. Cutting them from 32 to 4 shifts the whole line down by a constant factor of 8x, while a pure SSM stack is flat because its state size is independent of . The hybrid buys most of the SSM memory profile at the price of four layers.