Hybrid architectures (Jamba & recurrent attention)

Flip a layer from SSM to attention and its bar jumps — even at this toy sequence length of 8, an attention layer already costs noticeably more than an SSM layer at the same width. Flip all four and the total climbs every time; flip them all back and it falls. The stack's total cost is just the sum of whatever each layer happens to be.

A hybrid stack's total cost is the sum of each layer's own cost, and only the count of attention layers matters — not where they sit in the stack:

• $k$ — how many of the $L$ layers are full attention.
• $L$ — the total number of layers in the stack.
• $n$, $d$ — sequence length and model width, shared by every layer.
• $s$ — the SSM's fixed hidden-state size.

1. Each attention layer contributes its own full n²d term

   There's no sharing between layers — every attention layer independently pays the full quadratic cost at whatever $n$ the model is run at.

2. Every SSM layer instead contributes a linear nds term

   Swapping an attention layer for an SSM layer doesn't shrink the existing attention layers' cost — it just removes one $n^2d$ term from the sum and replaces it with a much smaller $nds$ one.

3. Even one leftover attention layer keeps the whole stack quadratic in n

   $\text{cost}(k,n)$ is still $\Theta(n^2)$ for any $k \geq 1$ — a hybrid with even a single attention layer scales the same way the all-attention stack does, just with a far smaller constant in front. Only $k=0$ (pure SSM) removes the quadratic term entirely.

The bold curve (all 4 layers attention) curves upward fast — pure O(n²). The dashed line (all 4 layers SSM) stays nearly flat by comparison — O(n). The middle curve is the hybrid with just one attention layer: still curving upward, since it's still quadratic in shape, but starting from a much smaller base — three of its four layers are paying the cheap linear cost instead.

$d=8$, SSM state size $s=4$, $n=64$, $L=4$ layers:

1. All 4 layers attention

   Each layer costs $64^2\times8=32{,}768$. Total: $4\times32{,}768=131{,}072$.

2. All 4 layers SSM

   Each layer costs $64\times8\times4=2{,}048$. Total: $4\times2{,}048=8{,}192$. Ratio to all-attention: $131{,}072/8{,}192=16\times$ reduction.

3. Hybrid: 1 attention layer, 3 SSM layers

   $32{,}768 + 3\times2{,}048 = 32{,}768+6{,}144=38{,}912$. Ratio to all-attention: $131{,}072/38{,}912\approx3.4\times$ reduction, while keeping one layer able to retrieve any token exactly, however far back it sits.