Research · Deep Learning
Mamba vs Transformer
A practical comparison of selective state-space models and attention-based Transformers — what changes, what doesn't, and what we learned shipping ResMamba (IEEE).
1. Why this comparison matters
Transformers dominated sequence modeling for nearly a decade because self-attention captures arbitrary token relationships in parallel. The cost is well known: attention is O(n²) in sequence length, KV-cache grows linearly with context, and inference latency blows up on long inputs. Mamba — built on a selective state-space model (SSM) — claims O(n) training, constant-memory inference, and Transformer-quality results on language and vision benchmarks. That tradeoff is what makes Mamba vs Transformer one of the most active design questions in 2025–2026 ML systems.
2. The architectures, side by side
| Dimension | Transformer | Mamba (Selective SSM) |
|---|---|---|
| Core op | Self-attention (softmax QKᵀ) | Selective state-space scan |
| Train complexity | O(n²) in sequence length | O(n) with parallel scan |
| Inference memory | Grows with KV-cache | Constant (recurrent state) |
| Long context (>32K) | Expensive, needs tricks | Native |
| In-context recall | Strong (induction heads) | Weaker without attention |
| Hardware fit | Mature kernels everywhere | Custom scan kernels (improving) |
Transformers compute, for every layer, a softmax over the full token window. Mamba instead maintains a recurrent hidden state and uses input-dependent selection parameters (Δ, B, C) to decide what to keep and what to forget — a data-dependent SSM, evaluated with a hardware-aware parallel scan.
3. Where Mamba wins
- Long context. Inference is constant-memory, so 100K+ token windows are practical without KV-cache tricks.
- Throughput. Selective scan saturates GPU memory bandwidth better than attention at long sequence lengths.
- Edge / streaming. Recurrent inference makes Mamba a strong fit for audio, time-series, and on-device workloads where attention's KV-cache is the bottleneck.
4. Where Transformers still win
- In-context recall. Pure SSMs underperform attention on tasks that require copying or pointing to a specific earlier token — induction heads, retrieval, exact lookups.
- Tooling and ecosystem. Attention has a decade of kernels, quantization, fine-tuning recipes, and serving infrastructure.
- Multi-modal fusion. Cross-attention between modalities is still the dominant interface in vision-language models.
5. Hybrid designs: Jamba and friends
The most interesting recent work doesn't pick a side. AI21's Jamba: a hybrid Transformer-Mamba language model interleaves Mamba blocks with sparse attention layers and a Mixture-of-Experts router. The Mamba blocks carry the long context cheaply, while a thin layer of attention preserves the in-context recall behavior Transformers are known for. The pattern generalizes: most production-grade long-context models shipping in 2025–2026 are hybrids, not pure SSMs.
6. What we learned building ResMamba
ResMamba — our IEEE-published architecture — applies the same hybrid intuition to a residual stack. A few takeaways that generalize beyond the paper:
- Residual connections around Mamba blocks materially stabilize training at depth; without them, deep SSM stacks are gradient-fragile.
- The selection parameters (Δ, B, C) benefit from a warmup schedule — initializing them too sharp causes the scan to collapse to a near-identity early in training.
- Mixing in a small fraction of attention heads (≈10–15% of layers) recovers most of the recall gap at a fraction of the FLOPs of a pure Transformer.
- On our benchmark suite, the hybrid matched a parameter-matched Transformer on perplexity while running ~2.4× faster at 32K context.
7. How to choose for your project
- Short context, retrieval-heavy: stay on Transformer. The ecosystem advantage is real.
- Long context, streaming, or edge: a Mamba or hybrid backbone is worth the engineering investment.
- Production LLM with long context: follow the Jamba-style hybrid pattern — Mamba for bulk capacity, a few attention layers for recall.
Further reading
- Gu & Dao, “Mamba: Linear-Time Sequence Modeling with Selective State Spaces” (2023)
- Lieber et al., “Jamba: a hybrid Transformer-Mamba language model” (AI21, 2024)
- Vaswani et al., “Attention Is All You Need” (2017)
- Mandal et al., “ResMamba” — IEEE