State Space Models & Hybrid Architectures: Mamba-2, Jamba, Griffin
State space models keep a fixed-size state instead of a key–value cache that grows with every token, so their memory and per-token cost do not grow with context length.
How to study this pagePredict · Check · Generalize · Continue
PredictSay what you expect.
Before you move a control in the demo, say what you expect to change.
CheckCheck against the math and the papers.
Compare what the demo shows with the key equation and with the papers cited below.
GeneralizeName what stays the same.
Identify the part of the idea that survives a change of notation or model, so you can recognize it in another paper or system.
ContinueFollow the connections.
Use the prerequisites and related concepts at the end of the page to fill the next gap.
Why It Matters for Modern Models
State space models keep a fixed-size state instead of a key–value cache that grows with every token, so their memory and per-token cost do not grow with context length.
Mamba-2's state space duality shows that SSMs and attention are two ways of computing structured weighted sums over the past, which lets SSM layers run largely as matrix multiplications on GPUs (Dao & Gu, 2024).
Jamba interleaves Mamba and attention layers, adds mixture-of-experts layers, and reports strong results for contexts up to 256K tokens (Lieber et al., 2024).
Griffin and RecurrentGemma mix gated linear recurrences with local attention, keeping a fixed-size state while still attending exactly to recent tokens (De et al., 2024; Botev et al., 2024).
Selectivity, the input-dependent update, is what made SSMs competitive on language; without it the layer is a fixed linear filter whose weights cannot depend on token content (Gu & Dao, 2024).
Questions to Think About
Why does a time-invariant SSM behave like a fixed convolution filter, and how do input-dependent parameters let it retrieve a specific earlier token?
Why combine local attention for recent tokens with a recurrence for longer-range memory, instead of using either alone?
A fixed-size state never grows, but it must compress everything seen so far; why does this make exact copying and recall from long inputs harder than for transformers (Jelassi et al., 2024)?
Why do SSMs still need specialized GPU kernels, such as Mamba's parallel scan, and how does Mamba-2 recast most of the computation as matrix multiplication?
Interactive Demos
Core Math
If you prefer intuition first, try the demos above, then come back for the math.
Key Equation
ht=Aht−1+Bxt,yt=Cht
A state space model (SSM) layer summarizes the past in a fixed-size state that is updated once per token, instead of attending over all previous tokens. Hybrid architectures interleave such layers with attention layers.
Discrete linear SSM:
ht=Aht−1+Bxt,yt=Cht
The state ht has a fixed size, so generation needs constant memory per step and O(T) time for a sequence of length T.
The same layer as a convolution:
Unrolling the recurrence from h−1=0 gives
yt=k=0∑tKkxt−k,Kk=CAkB
so the output is a weighted sum over past inputs, as in attention, but with weights fixed by A, B and C rather than computed from the content. Dao & Gu (2024) make the link precise: an SSM layer is a form of linear attention with a structured mask.
Selectivity:
In Mamba, B, C and the step size Δ used to discretize A are computed from the current input, while A itself does not depend on the input (Gu & Dao, 2024). Input-dependent parameters let the layer decide what to store and what to forget, which a fixed kernel cannot do. The convolution view then no longer applies, and Mamba evaluates the recurrence with a parallel scan.
Hybrids interleave layer types:
H(ℓ+1)=H(ℓ)+Mixℓ(H(ℓ)),Mixℓ∈{SSM or recurrence,attention}
with an MLP after each mixing layer. Jamba uses one attention layer for every seven Mamba layers (Lieber et al., 2024); Griffin alternates two gated-recurrence blocks with one local-attention block (De et al., 2024).
Key Papers
Mamba: Linear-Time Sequence Modeling with Selective State Spaces
Which limitations does it address, or leave to a later idea?
Linear complexity: a limitation of State Space Models that Attention addresses. Open Attention
State-space models such as Mamba were developed as an alternative to attention's O(n²) cost: they process a sequence in O(n) time and keep a fixed-size state during generation (Gu & Dao, 2024). The price is that the whole past must be compressed into that state, which makes exact copying from long inputs harder than for transformers (Jelassi et al., 2024).
Fixed-state memory: a limitation of Long Context that State Space Models addresses. Open Long Context
Even with RoPE scaling and KV compression, transformer decoding still pays KV costs that grow with T. SSM/hybrid architectures were invented to get fixed-size state and linear-time sequence modeling for very long contexts.
Both SSMs and linear attention can be written as recurrences with particular transition matrices. The "selective" in Mamba means that the recurrence depends on the input (Gu & Dao, 2024), the same idea as the data-dependent gates of gated linear attention.
Position via state: an analogy between State Space Models and RoPE. Open RoPE
SSMs encode position implicitly through state evolution, RoPE encodes position explicitly through rotation. Both solve "where am I in the sequence?" differently.
Iterative → Recurrent: an analogy between State Space Models and Diffusion. Open Diffusion
Both apply one learned update many times: a diffusion model applies its denoiser once per noise level, and a state-space model applies its state update once per token. The analogy is loose: diffusion steps refine a whole sample, while SSM steps move along the sequence.
Which ideas are its mathematical duals?
Kernel ↔ Recurrence: Attention is its mathematical dual. Open Attention
Attention applies a (data-dependent) kernel over history; SSMs apply a structured kernel computed by recurrence/scan. Two mathematically connected ways to implement "weighted sums over the past," with different computational tradeoffs.