Summary: Mamba (Gu & Dao, 2024) is a Selective State Space Model that achieves Transformer-quality with linear scaling in sequence length. It uses input-dependent dynamics (selection mechanism) to filter irrelevant information — enabling million-token context and efficient long-horizon latent dynamics for planning.


Motivation: Transformer Bottleneck

Property Transformer Classical SSM (S4) Mamba (Selective SSM)
Context length O(N^2) memory O(N) O(N)
Parallel training
Reasoning ability ❌ (fixed dynamics) ✅ (input-dependent)
Inference speed O(N) O(1) O(1)
Hardware efficient FlashAttention Custom kernels PagedAttention-style kernels

Key insight: Fixed SSM dynamics (A, B, C constant) can't content-filter → poor on language. Selective SSM makes Delta, B, C input-dependent.


Selective SSM (Mamba Core)

Standard SSM (S4/Lin. Recurrence)

h_t = A h_{t-1} + B x_t
y_t = C h_t

Fixed A, B, C — no content selection.

Selective SSM (Mamba)

Delta_t = softplus(Linear_Delta(x_t)) quad (input-dependent step size)
B_t = Linear_B(x_t) quad (input-dependent)
C_t = Linear_C(x_t) quad (input-dependent)
bar{A} = exp(Delta_t A)
bar{B} = Delta_t B_t
h_t = bar{A} h_{t-1} + bar{B} x_t
y_t = C_t h_t

Discretization: Delta_t controls how much to "step" — selectively propagate or forget.


Mamba Block

Input x
  ↓
LayerNorm
  ↓
┌───────────────────────┐
│  In-Projection        │ (x → x, z for gating)
└───────────┬───────────┘
            ↓
    ┌───────────────┐
    │  Conv1D (k=3) │  (local context)
    └───────┬───────┘
            ↓
    ┌───────────────┐
    │  Selective SSM│  ↑ MAMBA CORE
    │  (x → h → y)  │
    └───────┬───────┘
            ↓
    ┌───────────────┐
    │  OUT-Projection│ (z ⊗ y)
    └───────────────┘
            ↓
Add(Residual) → Output

Gating: y = SSM(x) ⊗ SiLU(z)


Mamba-2 (2024) — Structured SSM

Improvement Effect
SSD (State Space Duality) Unifies SSM + Attention
Head dimension Multi-head like Transformer
Chunked scan Even faster, better parallelism
Hybrid Mamba-Transformer Best of both

In World Models / Planning

Use Case Why Mamba
Latent dynamics z_t to z_{t+1} for 100+ steps (linear cost)
Video world models Token-efficient video (Genie, Sora)
Value-guided diffusion Rollout 50-100 steps for planning
RL with long horizon Credit assignment over 1000+ steps

DreamerV3 + Mamba: Replace RSSM GRU with Mamba → longer effective context.


Mamba vs Transformer for Dynamics

Aspect Transformer Mamba
Rollout 100 steps 100× KV cache grow Constant memory
Training parallel
Inference cost/step O(n) O(1)
Content selection Attention Selective SSM
State size Grows with context Fixed

Related Concepts


Sources