Hybrid Architectures — Attention + SSM

Two cost curves in one stack

Attention and SSMs sit at opposite ends of the same trade. A softmax-attention layer compares every token to every other, so prefill costs O(N^2 · d) and decoding needs a KV cache that grows linearly with N. An SSM layer instead carries a fixed-size hidden state h: [d_inner, d_state] and updates it once per token, so prefill is O(N · d · d_state) and decoding needs O(1) memory in N. Neither curve is strictly better: attention is expensive but lossless, while the recurrent state is cheap but lossy — a fixed-capacity summary whose overwritten contents are gone for good. A hybrid pays attention’s price on a small fraction of layers and inherits most of its benefit.

What attention is uniquely good at

The capability that separates the two is associative recall: given a key seen earlier, reproduce the value that followed it. Attention does this exactly — the softmax over QK^T / sqrt(d_k) is a soft dictionary lookup over an ever-growing table. It is the mechanism behind induction heads, in-context learning, verbatim copying, and retrieving a fact from the middle of a long document.

A recurrent layer must instead store those associations in d_inner × d_state numbers, and recall accuracy degrades once the number of key–value pairs to be held approaches that capacity — an information-theoretic wall, not a training artifact. Pure SSMs are strong on perplexity and weak precisely on the tasks that need exact lookup.

What SSMs are uniquely good at

The recurrent side wins on everything that scales with context length rather than context content. During decoding the state is a small constant, so per-token latency is flat whether you are at token 1,000 or token 100,000 — no growing cache to re-read, no growing bandwidth bill. During prefill the cost is linear in N, and selective SSMs express the scan as a chunked matmul, so they still saturate hardware rather than crawling along a sequential loop. They are also good at the bulk of what a language model does: local syntax, topic tracking, gradual accumulation of context. Most layers in a transformer do that kind of work — which is exactly why most layers can be replaced without a quality collapse.

The interleaving recipe

The standard construction repeats a block of one attention layer plus r recurrent layers, giving an attention fraction f = 1 / (1 + r). Reported ratios cluster tightly: r = 7 (f = 1/8) is most common, with 1:5 and 1:3 used when recall matters more than throughput. Below one attention layer per eight, long-context retrieval degrades; far above it you pay transformer prices for diminishing returns. Everything else stays familiar — each layer keeps its pre-norm and MLP (or MoE) sublayer, and the two layer types are drop-in replacements because both map X: [N, d] → [N, d].

Cost model for a mixed stack

Write L for total layers, f for the attention fraction, c_a for KV-cache bytes per token per attention layer, and s for state bytes per recurrent layer. Then the two quantities you actually care about are:

prefill  ≈  f·L·(4·N^2·d)  +  (1−f)·L·O(N·d·d_state)
                  quadratic term          linear term

memory   =  f·L·c_a·N       +  (1−f)·L·s
                  grows with N            constant in N

Two things fall out immediately. First, the quadratic term does not disappear — it is merely scaled by f, so a 1:7 hybrid with full attention is still O(N^2), just 8× cheaper. Second, the memory term splits into a context-dependent part and a fixed part, and which one dominates depends entirely on N.

Worked example: a 24-layer SLM at 32k context

Take L = 24, d = 2048, GQA with 4 KV heads of 128 (d_kv = 512), fp16, and Mamba-2-style recurrent layers with d_inner = 4096, d_state = 128. Per layer: c_a = 2 × 512 × 2 B = 2 KB per token; s = 4096 × 128 × 2 B = 1 MB.

N = 32,768   all-attention (f = 1)
  KV cache = 24 × 2 KB × 32,768        ≈ 1.57 GB
  attn FLOPs = 24 × 4·N^2·d          ≈ 2.11e14

N = 32,768   1:7 hybrid (f = 1/8 → 3 attn, 21 SSM)
  KV cache = 3 × 2 KB × 32,768         ≈ 196 MB
  SSM state = 21 × 1 MB                   =  21 MB
  total ≈ 217 MB          (7.2× smaller)
  attn FLOPs = 3 × 4·N^2·d           ≈ 2.64e13   (8× smaller)

The headline is memory: 1.57 GB of KV cache is a deployment blocker on a laptop, 217 MB is not — and at batch-1 decode, where the cache is re-read every token, the same factor shows up directly as latency.

The crossover point

The fixed SSM state and the growing KV cache trade places at a specific context length. Setting the two memory terms equal gives

N* = (1−f)·s / (f·c_a)

With the numbers above: N* = (0.875 × 1 MB) / (0.125 × 2 KB) = 3,584 tokens. Below ~3.6k tokens the recurrent state is the larger cost and the hybrid saves you little; above it, the three attention layers dominate and keep growing while the SSM half stays flat.

That is a useful design compass: hybrids are not a general win, they are a long-context win. At 2k a plain transformer of the same size is simpler and roughly as cheap; the gap opens at 8k and becomes decisive past 32k.

Sliding-window hybrids: making the whole stack linear

If you want to kill the N^2 term outright rather than divide it by eight, make the attention layers local. A sliding window of size w turns each attention layer into O(N · w · d) compute with a bounded w-token cache, so the whole model becomes linear in time and constant in decode memory.

The division of labour is clean: the window handles precise short-range work and the recurrent layers carry long-range information forward, so the window never has to reach far. The cost is that unbounded exact recall is gone again — anything the state failed to keep is unrecoverable. Designs that need needle-in-a-haystack retrieval usually keep at least one global attention layer.

Where to place the attention layers

Placement is not free. The empirical pattern across published hybrids is that attention layers earn most in the middle of the stack and least at the very bottom: early layers do local, positional feature mixing that a recurrent layer already handles more cheaply, while mid-stack layers are where retrieval and induction-style circuits form.

A second consequence: several hybrids report they can drop explicit position encodings entirely. A recurrent layer is inherently ordered, so positional information reaches the attention layers through the residual stream and RoPE becomes optional — removing one of the usual obstacles to length extrapolation.

Families in the wild

The design space has converged on a small number of shapes:

Zamba’s trick is the telling one: with so few attention layers you can share one set of weights across all of them at almost no quality cost — the model needs the attention operation repeatedly, but not many distinct copies of it.

Why so few attention layers suffice

The intuition is that exact recall is a routing problem, not a representation problem. An induction circuit — find the previous occurrence of the current token, copy what came next — needs a way to move information from an arbitrary earlier position to here. One or two attention layers provide that channel; the surrounding recurrent layers then compute freely on what was retrieved. The recurrent stack is a good compressor and a bad index; a few attention layers give it an index, and once the index exists, adding more mostly duplicates a capability the model already has. That is the diminishing-returns curve the 1:7 ratio sits on.

Training and systems consequences

Hybrids are harder to build than either pure stack. You need two kernel families — a FlashAttention-style fused attention and a chunked selective scan — and both must be fast, or the model inherits the slower one’s ceiling. Layer costs are also heterogeneous, which breaks the pipeline-parallel assumption that every stage has equal work: an even split leaves the attention stage as the bubble-producing straggler. And the recurrent state accumulates over thousands of steps, so it is more sensitive to low precision than an attention score — keeping the scan state in fp32 while the rest runs bf16 is the usual compromise.

What this means for CPU and small models

Batch-1 decoding on a CPU is memory-bandwidth-bound, and the KV cache is usually the biggest thing re-read per token. At 32k the example above reads 1.57 GB per token; at a realistic 40–50 GB/s of effective bandwidth that is ~30 ms of cache traffic before any weight is touched. The hybrid’s 217 MB costs 4–5 ms.

Two secondary wins follow. Latency stays flat: the recurrent state does not grow, so tokens-per-second barely degrades as the conversation lengthens — a far better experience than a transformer that visibly slows down. And a few megabytes of state can stay resident in last-level cache, whereas a multi-gigabyte KV cache guarantees DRAM round trips. For a CPU-served SLM that combination is the strongest argument for hybrids.

Pitfalls

Assuming linearity you did not buy. A 1:7 hybrid with global attention is still O(N^2); only sliding-window or fully-linear attention makes the stack truly linear.

Undersizing the state. Shrinking d_state to save memory is the fastest way to destroy recall, and the damage never shows up on perplexity — only on retrieval evals.

Short-context evaluation. Hybrids and transformers look identical at 2k tokens; test at 32k and up, on tasks that require exact lookup.

Treating the layers as interchangeable. Uniform pipeline splits and uniform quantization both assume homogeneous layers; here they differ in cost and in numerical sensitivity.