Hybrid Architectures

Dive into Deep Learning · §12.7

Hybrid Architectures
fixed-state capacity · attention-cache growth · recall with one attention layer · deployed configurations

Finite-State Recall Limit

Five upgrades this chapter — gate, linearize, select, edit, learn — and none changed the state’s size.

  • Reproducing k tokens from vocabulary V needs k \log_2 V bits, held from write to query.
  • Copy bound (Jelassi et al. 2024): a finite-state model errs with probability \ge 1 - |\mathcal{S}|/D^L on uniformly random length-L strings — a counting argument, no training trick appears in it.

Measured: a 410M transformer beats a 2.8B Mamba on phone-book lookup past ~70 entries.

The diagnosis: associative recall

Where the wall bites in real text (Zoology, Arora et al. 2024):

  • 82% of the perplexity gap between efficient models and attention sits on the recall slice — tokens requiring retrieval of a binding seen once.
  • Attention solves MQAR at width independent of length; a fixed state only while bindings fit.

Production symptom: fuzzy recall — MMLU 29 vs 46 at matched 8B; ablate a hybrid’s few attention layers and needle-in-a-haystack drops to ~0.

Attention-Layer Cache Cost

Per layer, at production width (16-bit):

  • Full attention (GQA 8 \times 128): 4 KB per token — 512 MB per layer at 128K context.
  • Mamba-2 state: ~2 MB, constant in context length.
image/svg+xml Matplotlib v3.10.8, https://matplotlib.org/

The hybrid cache equals the attention fraction times the transformer cache.

Jamba Cache at 256K Context

model KV cache
Llama-2-70B-class 128 GB
Mixtral 32 GB
Jamba (4 attn of 32 layers) 4 GB

Recurrence removes the growing term; the constant state can still be quantized, but nothing grows. The surviving attention cache takes GQA / MLA / quantization on top.

Three matched models, one letter apart

Same depth, width, heads, MLPs, embeddings — only the mixer string differs:

LAYOUTS = {'linear': 'GGGG', 'attention': 'AAAA', 'hybrid': 'GGAG'}
  • 'G' = ScalarGatedMixer: scalar-per-head decay, Mamba-2’s rung of the decay ladder (not GLA’s per-coordinate gate), trained via the quadratic dual.
  • 'A' = the causal attention of ch. 11, through d2l.TransformerBlock’s attn_factory hook.
  • Printed parameter counts: 58,544 / 57,984 / 58,404 — spread < 1%.

Three lessons before training

A fresh gate must retain. Default init → a \approx 0.5: state half-life of one token, chance recall at every load. Bias -4.5 → a \approx 0.99. (S4D, Mamba, Gated DeltaNet all ship this trick.)

Mask before exp. The dual’s \exp(\textrm{cum}_i - \textrm{cum}_j) is large and positive above the diagonal; trained decays → 0 make it overflow: inf * 0 = nan. Mask with -\infty in log space first — only the LM panel triggers it, never the recall task.

Normalize the read-out. \mathbf{S}^\top \mathbf{q} grows with what the state has accumulated (\sqrt{T} at init); per-head RMSNorm before W_o; otherwise, recurrent-block outputs dominate the residual stream before it reaches the downstream attention layer.

Recall with One Attention Layer

  • Linear: degradation, then collapse at 64 pairs — the crowding heuristic’s cliff for 8 \times 8 head states (a diagnostic analogy, not a theorem about this trained stack).
  • Hybrid = 3/4 the same recurrence + one mid-stack attention layer: 1.000 at every load (PyTorch); JAX dips to ~0.92–0.94 mid-sweep, then 0.99–1.00 at the loads that break the linear stack. Attention: 1.000 throughout.

Perplexity misses the deficit

Same three stacks as character LMs on The Time Machine:

     model  val loss  bits/char
    linear      1.78       2.57
 attention      1.83       2.65
    hybrid      1.76       2.54
  • All within ~0.1 nat; the pure-linear model matches or beats pure attention — while collapsing on recall. (One pass, one small corpus, one seed: a demonstration.)
  • At scale the dissociation is the rule: the perplexity gap sits on the rare recall tokens (Zoology); ratio sweeps move recall, barely perplexity (2507.06457). MQAR and needle probes exist because the training objective did not see the deficit.

Persistent decode state, measured

  • KV cache + recurrent state only — parameters and activations are the same across stacks.
  • Attention: grows linearly. Linear: 4 KiB at every length.
  • Hybrid: slope = attention fraction (here 1/4).

The hybrid matches attention recall while retaining one quarter of the attention stack’s context-dependent state.

Design rules, measured

  • Ratio: shipped 5–25% attention; ablation optimum ~8% (NVIDIA); DeltaNet-family knee at 3:1 (Kimi ablation). Recall climbs to ~1:3; perplexity flat throughout.
  • Placement: evenly spread, not first in any report so far — one front attention layer broke Samba’s length extrapolation. A heuristic from a young design space, not a law.
  • Parallel (Hymba, Falcon-H1): fuse per layer; beats sequential in controlled ablation.
  • Sharing (Zamba2): re-enter two weight-tied attention blocks along the backbone.
image/svg+xml Matplotlib v3.10.8, https://matplotlib.org/

Interactions Between Components

AI21 tested Mamba-2 in Jamba — the better standalone model — and rejected it.

  • With attention layers pooling the full context, the recurrence no longer needs the bigger state.
  • No quality difference 1:3 vs 1:7 at their scale → shipped the cheaper.

Evaluate the recurrent component together with the attention component. Once quality saturates, choose the less expensive configuration.

The recipe table

model layers attention recurrence context
Jamba 32 4 full, GQA Mamba-1 256K
Nemotron-H 52 4 full, GQA Mamba-2 8K
Granite 4.0-H 40 4 full, NoPE Mamba-2 128K
Qwen3-Next 48 12, GQA 16{:}2 gated DeltaNet 262K
Kimi Linear 27 7 MLA, NoPE KDA 1M

Every row: evenly spread, none first; the surviving attention wears every compression of ch. 11.

Distillation: duality warm-starts, data does the rest

Duality is exact for linear attention only; softmax attention is not a recurrence in disguise. Conversion is a learned approximation:

  • MOHAWK: match mixing matrices → hidden states → fine-tune; Phi-1.5 → Mamba with ~3B tokens (<1% of scratch).
  • Mamba-in-the-Llama: init recurrence from attention projections, keep ~1/4 of attention layers — recall survives at 20× training length; remove them all and it does not.

Architecture Comparison

Per layer: decode work / persistent state / training depth / exact?

architecture decode state depth exact?
LSTM/GRU O(d^2) O(d) T no parallel form
S4D, Mamba-1 O(dN) dN O(\log T) yes
SSD / GLA / DeltaNet O(h d_k d_v) h d_k d_v T/C yes
softmax attention O(td) ↑ 2td ↑ O(1) by definition
1-in-4 hybrid mix const + cache/4 mix per component

The fixed state lost exact recall (a counting bound, plus our measured collapse) and won the economics (by a factor that grows with context). Production stopped choosing. Kernels → Computational Performance; pretrained stacks → the Language Models part.