EngineersofAI LogoCoursesBlogResearch LabAI LettersThe LabCode BankInteractive 3DKodrEarnestJobs
Skip to main content
Interactive 3D/Mamba State Space Model
Transformer Self-Attention
Transformer Self-Attention - n×n matrix (n=32)
Memory: O(n²) = 0.0 MB at n=32
Every position attends to every other. Memory grows quadratically.
Mamba SSM (Selective State Space)
Mamba SSM - fixed-size hidden state h_t (d_state=64)
h_0
→
Δ,B,C
SSM step
x_0
→
h_0
Current state h_0 (32 of 64 dims shown)
Memory: O(d_state) = 0.3 KB - constant regardless of sequence length
Fixed-size state updated recurrently. Memory is constant in n.
Seq length
256
Attn memory
256 KB
Mamba mem
0.3 KB
Speedup
1024×
Transformer: O(n²) memory
At n=1M: ~3.7 TB just for attention matrix.
Mamba: linear scaling
At n=1M with d_state=64: only 0.3 KB. Constant.
Controls
Seq length256
648192
d_state64
16256
Mamba (Gu & Dao, 2023) is a selective SSM. The A, B, C, Δ matrices are input-dependent - the model selects what to remember.

Unlike RNNs: parallel training via selective scan. Unlike Transformers: O(n) inference memory.

Mamba State Space Model - Interactive Visualization

Mamba is a state space model (SSM) that processes sequences recurrently using a learned hidden state, scaling linearly with sequence length rather than quadratically like attention. Its key innovation is selectivity - the state update matrices A, B, C depend on the input, letting it selectively remember or forget information like an LSTM but faster.

  • Recurrent state visualization - watch how the hidden state h evolves as each token in a sequence is processed
  • Selectivity mechanism - see how input-dependent A, B, C matrices allow selective memory retention vs forgetting
  • Linear vs quadratic scaling - compare Mamba memory and compute usage to full attention as sequence length grows
  • Parallel scan - understand how Mamba trains in parallel despite being recurrent at inference time
  • Compare Mamba, LSTM, and transformer attention on the same sequence to see when each architecture excels

Part of the EngineersOfAI Interactive 3D - free interactive visualizations covering every major concept in machine learning and AI engineering. Hover any element for a plain-English explanation. No code required.