State-space models (Mamba, S4, hybrids)

Why this matters

Attention's O(T^2) cost is the wall everyone hits at very long context. Since 2022, state-space models (SSMs) — culminating in Mamba and its successors — have offered a near-linear alternative. They are not "the next thing after transformers" yet, but they are the strongest non-attention architecture, and hybrid models (transformer + Mamba layers) are showing real production promise (Jamba, Zamba, Falcon-Mamba, IBM Granite-3 hybrid).

You should at least be able to discuss SSMs intelligently. This lesson gives you that fluency without weeks of math.

Learning objectives

1. Describe a state-space model in plain English.
2. Distinguish S4, S5, and Mamba.
3. Compare Mamba vs Transformer in compute, memory, and quality.
4. Recognise hybrid SSM-attention models and why they exist.
5. Use a Mamba checkpoint with HuggingFace.

1. The intuition — back to RNNs, but smart

Recall RNNs (Lesson 1.2):

h_t = f(h_{t-1}, x_t)
y_t = g(h_t)

They are fast (O(T) memory, sequential O(T) compute) but bad at long-range dependencies (vanishing gradients).

State-space models are RNNs with a specifically structured update designed to:

• Remain stable over thousands of steps.
• Be parallelisable for training (compute all h_t in O(T log T) via convolution / scan).
• Capture long dependencies as well as attention.

A continuous-time SSM:

h'(t) = A h(t) + B x(t)
y(t)  = C h(t) + D x(t)

Discretised:

h_t = Ā h_{t-1} + B̄ x_t
y_t = C  h_t   + D  x_t

Where A, B, C, D are matrices. The trick is choosing A so the model can remember things over long horizons (HiPPO theory, Voelker et al. → Albert Gu et al.).

2. The progression

S4 (Gu et al., 2021)

First SSM that worked well for sequence modelling. A parameterised by a structured matrix (HiPPO + diagonal+low-rank). Strong on Long Range Arena benchmarks. Math-heavy; not trivially "drop in" for text.

S5 (Smith et al., 2022)

Diagonal-only A. Simpler, faster, similar quality.

H3 / Hyena (2022-2023)

Add multiplicative gating like an LSTM. Bridges SSMs to language modelling but still not competitive with strong transformers at scale.

Mamba (Gu & Dao, 2023)

The breakthrough. Two big additions:

1. Selective SSM: make B, C, and the discretisation step Δ input-dependent. Now the state-space recurrence can "decide" what to remember based on the token, much like attention's content-aware mixing. (Earlier SSMs had time-invariant matrices — a deal-breaker for language.)
2. Hardware-aware scan kernel: a custom CUDA kernel that runs the recurrence efficiently in SRAM (analogous to FlashAttention).

Net result: Mamba (130M to 2.8B) matches transformer perplexity at the same parameter count, with 5× faster inference and constant memory per token (just the hidden state, not a growing KV cache).

Mamba-2 (2024)

Reformulates the SSM as a "structured masked attention," letting it use matrix multiplications and tensor cores. Faster training; same flavour.

Modern SSM-only models

• Falcon-Mamba 7B (TII 2024) — first SSM-only model competitive with transformers.
• Mamba-Codestral 7B (Mistral 2024) — code model.
• Phi-Mamba experiments — smaller scale.
• Zamba (Zyphra) — Mamba + small attention.

3. Mamba block (sketch)

class MambaBlock(nn.Module):
    def __init__(self, d, d_state=16, d_conv=4, expand=2):
        super().__init__()
        d_inner = expand * d
        self.in_proj = nn.Linear(d, 2 * d_inner)              # x and z
        self.conv = nn.Conv1d(d_inner, d_inner, kernel_size=d_conv,
                              groups=d_inner, padding=d_conv-1)
        # Selective SSM parameters (input-dependent)
        self.x_proj = nn.Linear(d_inner, d_state * 2 + 1)     # gives B, C, dt
        self.A_log = nn.Parameter(torch.randn(d_inner, d_state))   # log-A diag
        self.D = nn.Parameter(torch.ones(d_inner))            # skip term
        self.out_proj = nn.Linear(d_inner, d)

    def forward(self, x):
        # x: (B, T, d)
        xz = self.in_proj(x)
        x_, z = xz.chunk(2, dim=-1)
        x_ = self.conv(x_.transpose(1,2))[..., :x_.size(1)].transpose(1,2)
        x_ = F.silu(x_)
        # selective scan: compute B, C, dt per token; run SSM
        B, C, dt = self.x_proj(x_).split([d_state, d_state, 1], dim=-1)
        A = -torch.exp(self.A_log)
        # SSM scan (handled by mamba_ssm.selective_scan_fn in practice)
        y = selective_scan(x_, dt, A, B, C, self.D)
        return self.out_proj(y * F.silu(z))

In practice you import from mamba_ssm import Mamba (or Mamba2) and replace the attention block with it.

4. Mamba vs Transformer — the trade-offs

Mamba's killer feature: constant memory and per-token cost for inference. A 1M-context Mamba literally costs the same per token as a 1k-context Mamba.

Mamba's weakness: empirical "in-context recall" — when you put a fact at position 1000 and ask about it, attention models retrieve it more reliably. The hidden state of an SSM has finite capacity.

5. Why hybrid models are the practical winner

Combine the strengths:

• A few attention layers → strong in-context recall, retrieval.
• Many Mamba layers → cheap long-context, low memory.

Examples

• Jamba (AI21, 2024) — Transformer/Mamba/MoE hybrid with 256k context.
• Zamba 2 (Zyphra, 2024) — small hybrid models.
• Granite-3 hybrid (IBM, 2024).
• Falcon-3 hybrid (TII, 2025).
• Llama-4 has investigated SSM layers in some configurations.
• Hymba (NVIDIA 2024) — hybrid attention + Mamba per layer.

A hybrid layer pattern might look like [A, M, M, M, A, M, M, M, ...] — one attention every 4-7 layers, the rest Mamba. Maintains transformer-class quality, gains huge inference savings.

6. Using Mamba with HuggingFace

from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("state-spaces/mamba-2.8b-hf")
m   = AutoModelForCausalLM.from_pretrained("state-spaces/mamba-2.8b-hf",
                                           torch_dtype="bfloat16",
                                           device_map="auto")
ids = tok("The cat sat on the", return_tensors="pt").input_ids.to(m.device)
out = m.generate(ids, max_new_tokens=100, do_sample=True, top_p=0.9)
print(tok.decode(out[0]))

Or for hybrids: tiiuae/falcon-mamba-7b-instruct, ai21labs/AI21-Jamba-1.5-Mini, Zyphra/Zamba2-2.7B-instruct.

7. Open research questions

• Can SSMs match transformers on strong in-context learning? Hybrids partially solve this.
• Can SSMs do needle-in-haystack at 1M context as well as long-context transformers? Empirically, hybrids do; pure SSMs struggle.
• Is the right ratio of attention to Mamba layers small (1:8) or large (1:4)? Active research.
• Can MoE + SSM combine cleanly? Recent papers say yes (Jamba, Zamba 2).

If you go into research, this is a fertile area. If you go into product, you will use hybrids more than build them.

Hands-on lab (4 hours, GPU)

mamba_lab.ipynb:

1. Install: pip install mamba-ssm causal-conv1d.
2. Build a tiny model: 8 Mamba layers, d=512. Train on tiny_shakespeare 5000 steps. Compare PPL with your nano-GPT (Lesson 2.4) at similar param count.
3. Generate 4096 tokens. Measure tokens/sec — should be roughly constant regardless of position (vs. transformer slowing as KV grows).
4. Run NIAH on state-spaces/mamba-2.8b-hf at 4k, 8k, 16k. Compare to Qwen2.5-3B-Instruct (transformer). Discuss recall difference.
5. Try the Jamba 1.5 Mini hybrid (ai21labs/AI21-Jamba-1.5-Mini). Run NIAH at 16k, 64k, 128k.
6. Bonus: build a hybrid mini-model — alternate one attention block and four Mamba blocks. Train and compare to all-attention and all-Mamba baselines.

Common pitfalls

1. Forgetting selective_scan kernel install — running on CPU triggers slow Python fallback.
2. Comparing pure Mamba on tasks where attention dominates (e.g., complex few-shot retrieval) — it loses; check your benchmark fits.
3. Treating SSM "state" like a KV cache — different semantics; you cannot trim or page it the same way.
4. Hyperparam transfer from transformer recipes — some learning rates and batch sizes don't translate.
5. Misunderstanding Mamba-2 = Mamba — Mamba-2 is structurally different (matrix-form state); kernels and shapes differ.

Self-check

1. What does the "selective" in "selective SSM" make input-dependent?
2. Why does Mamba inference have constant per-token cost?
3. What is the Mamba weakness on which transformers usually win?
4. What are hybrid models and why do they exist?
5. Name two production-ready hybrid models.

References

• Gu et al. (2021), "Efficiently Modeling Long Sequences with Structured State Spaces" (S4).
• Smith et al. (2022), "Simplified State Space Layers for Sequence Modeling" (S5).
• Gu & Dao (2023), "Mamba: Linear-Time Sequence Modeling with Selective State Spaces."
• Dao & Gu (2024), "Mamba-2: Transformers are SSMs."
• Lieber et al. (2024), "Jamba: A Hybrid Transformer-Mamba Language Model."
• TII (2024), "Falcon Mamba 7B."
• Zyphra (2024), "Zamba 2."
• NVIDIA (2024), "Hymba: A Hybrid-head Architecture for Small Language Models."