Transformers and Neural State-Space Architectures › Transformer and State-Space Hybrids

This section surveys architectures that combine attention and state-space recurrence, the design patterns they follow, and why mixing the two mechanisms is a hypothesis worth testing for long-horizon world models.

Can recurrence reduce memory cost while attention preserves useful access to selected context representations? Several 2024 hybrid language-model reports explore this design hypothesis. Static layer interleaving does not dynamically identify valuable events: it applies its assigned layers to every token. Savings come from using fewer attention layers, restricting their span, or sharing parameters. Their value for a world-model task still requires task-specific evaluation.

One pattern is layer interleaving. Jamba (AI21, 2024) uses an attention-to-Mamba ratio of 1:7 in its reported architecture. Zamba (Zyphra, 2024) repeatedly invokes a shared attention-plus-MLP module among Mamba blocks, saving parameters relative to separate copies. Griffin and Hawk (De et al., 2024) use gated linear recurrence: Griffin retains local attention, whereas Hawk is recurrent-only. Their reported quality and inference-cost comparisons concern specified language benchmarks, model scales, contexts, and hardware, not all world models. Griffin's held-out-book experiments also test length extrapolation beyond training context. A length-independent recurrence is a plausible contributing bias, but that alone neither proves the mechanism nor guarantees extrapolation.

For world models specifically, hybridization predates the recent wave of SSM-based hybrids and takes a second form: transformer dynamics inside a stochastic latent world model.

• TransDreamer (Chen et al., 2022) introduces the Transformer State-Space Model (TSSM). It replaces recurrent dynamics with transformer dynamics and modifies posterior conditioning to permit parallel training. Its 2D and 3D Hidden Order Discovery tasks involve discovering and remembering object orders through interaction. Reported gains over Dreamer are empirical, not a proof that finite recurrent state cannot solve these tasks or that memory is the only changing factor.
• TWM (Robine et al., 2023), not TSSM, uses Transformer-XL dynamics over latent states, actions, and rewards with a separate observation model. It reports results under Atari 100k's 100,000-interaction budget. Short attention paths are a possible benefit, but the reported comparison does not isolate them as the cause of sample efficiency.
• A possible chunked hybrid runs local recurrence within windows and attention over their summaries. This should not be confused with the named Block-Recurrent Transformer, whose state update uses attention and gates. For one full summary-attention pass with $M\approx T/W$ fixed-width summaries and linear within-window work, total work is $O(T+M^2)=O(T+(T/W)^2)$. Fixed $W$ leaves a quadratic term; bounded $M$, or growth no faster than order $\sqrt T$, gives linear-order total work. Recomputing summary attention at every token would change this accounting.

A third design axis changes attention's retained memory. Sliding windows restrict direct access to recent tokens; cache compression can select, merge, or quantize stored representations. Global tokens provide shared access points, but they become a recurrent cross-segment summary only if an explicit update-and-carry mechanism is implemented with the required causal restrictions. Chunked prefill is different: it schedules processing in chunks and need not reduce retained KV state or change the memory representation. These choices trade cost against available information, with consequences revisited in Part VI.

For a new task, ask what history distinctions predictions or decisions require. One useful test places a decision-critical observation at varying temporal distances while controlling its content, distractors, and task difficulty. Plot error against distance alongside aggregate prediction loss. A fixed-decay system may lose usable event information as distance grows; an attention model may fail when the relevant representation leaves its context. These are hypotheses, not guaranteed curve shapes: learned retrieval can fail inside a window, and selective recurrence can retain events through distractors. A distance-controlled experiment reveals failures that an average can hide, but our smooth-dynamics experiment below does not perform this rare-event test.