Hybrid Attention-SSM Architectures
Why the strongest linear-time models are not pure state space models, what a small number of full attention layers restores, and how to reason about the mixing ratio.
By 2024 the interesting result about state space models was not that they replaced attention. It was that a few layers of attention, sprinkled into an otherwise recurrent stack, recovered nearly all of the transformer's quality while keeping most of the linear-time model's inference economics. Nearly every production long-context model that is not a pure transformer now has this shape.
What each half contributes
A selective state space layer compresses history into a fixed-size state. Cost per token is constant in context length, memory is constant, and throughput at long context is excellent. What it cannot do is recall a specific token from far back with precision, because the information had to fit through a bottleneck of dimension \(N\).
An attention layer keeps every token available and can retrieve any of them exactly. Cost per token grows with context, and the KV cache grows linearly with sequence length, which at long context becomes the dominant memory consumer during serving.
The empirical finding is that the model needs the second property only occasionally. Copying an exact string, resolving a coreference to a distant antecedent, and performing in-context retrieval are localised abilities that a small number of attention layers can supply for the whole stack, because once one layer has moved the needed token's information into the residual stream, subsequent recurrent layers can carry it forward.
Ratios in practice
Published hybrids cluster around one attention layer for every six to eight recurrent layers. Jamba (Lieber et al., 2024, arXiv:2403.19887) uses a 1:7 ratio combined with mixture of experts, and reports a KV cache far smaller than a comparable transformer at 256k context, which is the point of the exercise: the cache only exists for the attention layers. Zamba, Samba and several later systems adopt similar ratios with different placement rules.
Where the attention layers sit matters less than that they exist, but the common finding is that neither the very first nor the very last layer is the best place, and spreading them evenly through the middle outperforms clustering them.
Reasoning about the tradeoff
The quantity to think about is the serving KV cache. In a pure transformer it is proportional to layers times context length. In a 1:7 hybrid it is proportional to one eighth of that, since only attention layers hold a cache. At 256k context that difference decides how many concurrent requests fit on a device, and concurrency is what determines cost per token far more than raw FLOPs.
Against this, hybrids give up the clean theoretical story. A pure state space model has constant memory and a well-understood recurrent form. A hybrid is a transformer with most of its attention removed, which means it inherits the transformer's context-length memory scaling with a smaller constant rather than escaping it.
When it breaks
The recall ceiling moves but does not vanish. A hybrid still funnels most information through fixed-size states between its attention layers. Tasks requiring many simultaneous precise retrievals degrade earlier than in a full transformer, and the degradation appears at a context length that depends on the ratio, which makes it easy to miss during evaluation on shorter inputs.
Two kernels, two sets of problems. A hybrid needs both an efficient attention implementation and an efficient scan implementation, plus a serving stack that handles two different cache types in the same request. Continuous batching, prefix caching and speculative decoding were all built assuming uniform layers, and each needs adaptation. This engineering cost is the main reason hybrids diffused more slowly than their benchmark results suggested they would.
Distillation across the boundary is possible and lossy. Converting a pretrained transformer into a hybrid by replacing most attention layers and fine-tuning is much cheaper than pretraining, and it consistently loses some long-range recall relative to the parent. It is a good way to get a cheaper serving profile from an existing model, not a way to get the same model cheaper.
Benchmarks under-report the difference. Most standard benchmarks use short contexts where the architectures are equivalent, and long-context benchmarks that test retrieval of a single fact are easy for hybrids. The discriminating tests involve many facts, aggregation across a long document, or precise reproduction, and they are less commonly run.
12 flashcards for this concept
Click a card to reveal the answer.