01/ 03Case study
Efficient Sequence Modeling: From S4 to Mamba-3
Five sequence architectures implemented from first principles, trained at matched budgets, and measured against their own papers' claims - OOM points included.
- Python
- PyTorch
- NumPy
- SciPy
- Hugging Face
- Matplotlib
- pytest
Lab/See it run
Tokens inA sequence of tokens enters the model. Every architecture in the study reads this same input through one shared skeleton.
01/Overview
What this is.
ssmbench is a from-scratch experimental study of how efficient sequence models evolved: a causal Transformer baseline alongside S4D, Mamba, Mamba-2 and Mamba-3, all implemented in pure PyTorch behind one shared training skeleton, then trained and benchmarked on three designed tasks plus a 1K-32K token efficiency sweep on a single 6 GB RTX 4050.
Every architecture shares the same block plumbing, so a benchmark difference isolates the mixer, not the scaffolding. Every number is committed as JSON with its environment block, and every dual compute form (FFT convolution vs recurrent step, chunked vs sequential scan) is paired with an equivalence test at 1e-3 tolerance.
02/Problem
What motivated it.
Attention scales quadratically with sequence length, and long context is where it becomes unusable. State-space models promise O(T) computation with a fixed-size recurrent state - but each generation of SSM architecture carries its own trade-off, and the claims are easy to repeat and hard to verify.
The gap this project attacks: what survives when you implement each generation from first principles, train the models at matched budgets, and measure the claims yourself - including the wall-clock cost that operation-count notation hides.
03/System
How it is built.
- transformer_minimal.py
- Pre-norm causal Transformer with SDPA attention; the quality and scaling baseline. No KV cache: decode recomputes the full prefix, which the efficiency sweep measures.
- s4_minimal.py
- S4D: complex diagonal A with S4D-Lin initialization, exact ZOH via a numerically safe (e^z-1)/z, FFT-convolution training form and an O(1)-state recurrent decode step.
- mamba_minimal.py
- Mamba-1: input-dependent Δ, B, C (selectivity) as an explicit time-loop selective scan, with a verified incremental step() for decoding.
- mamba2_minimal.py
- Scalar-per-head A reduces the recurrence to SSD: chunked decayed-attention matmuls with a carried (N, P) state, plus a sequential_scan reference used by the equivalence tests.
- mamba3_minimal.py
- Complex diagonal transitions A = -exp(a) + iθ with a data-dependent exponential-trapezoidal rule; reduces exactly to Mamba-1/2 when λ = 1.
- shared skeleton
- blocks.py (classifier heads), synthetic + IMDb data pipelines, a config-driven trainer that emits JSON, and a CUDA-event efficiency harness.
04/Approach
The mechanisms that matter.
Math-verified discretization
Exact ZOH with a Taylor fallback keeps (e^z-1)/z stable for tiny z. The recurrence is asserted against its parallel form, never eyeballed.
Chunked SSD with a sequential reference
Mamba-2's decayed-attention form is computed block-wise in chunk groups to bound transient memory (about 1 GB saved at 32K), and tested equivalent to the naive sequential scan for aligned (32) and misaligned (64) chunk sizes at atol/rtol 1e-3.
Complex transitions for state tracking
Mamba-3 pairs A = -exp(a) + iθ with a data-dependent trapezoidal λ. The rotation phase is what lets the state track parity; a test pins the reduction to exponential-Euler when λ = 1, θ = 0.
OOM as data
The efficiency sweep records out-of-memory points as first-class results (oom_train: true) instead of hiding them, including the WDDM shared-memory spill past the 6 GB card.
Slow on purpose
The selective scans are plain Python loops. That turns the gap between O(T) operation counts and real GPU execution into a measured, teachable finding.
05/Evaluation
Setup and workloads.
- Hardware: RTX 4050 Laptop (6 GB), CUDA 13; seed 42 for every run; all numbers committed under results/ with environment blocks.
- Tasks: cumulative-XOR parity (L=128), selective copying (train L=64, test L=128), IMDb sentiment (12.5k samples, max_len 512), each at matched training budgets per config.
- Efficiency sweep: CUDA-event timing, 5 repeats, batch 8, d_model 256 / 2 layers, sequence lengths 1K to 32K; per-model training-step latency, peak memory, throughput and decode cost, with the Python-loop models capped at 4K.
06/Results
Measured outcomes.
- 0.983
- parity accuracy, Mamba-3 (baselines 0.51-0.56)
- 67%
- selective copying, Mamba-2 (S4D 6.4%)
- 0.836
- IMDb accuracy, Mamba-2 (AUC 0.917)
- 186 s
- Transformer train step at 32K (S4D 13.5 s)
| Model | Parity acc | Sel. copy | IMDb acc | 1K step | Long-context limit |
|---|---|---|---|---|---|
| Transformer | 0.533 | 0.205 | 0.813 | 30 ms | 32K · 186 s/step |
| S4D | 0.506 | 0.064 | 0.827 | 19.5 ms | 32K · 13.5 s/step |
| Mamba | 0.559 | 0.649 | skipped | 2 632 ms | 4K · 248 s/step |
| Mamba-2 | 0.563 | 0.670 | 0.836 | 78.5 ms | 16K · OOM at 32K |
| Mamba-3 | 0.983 | 0.215 | skipped | 10 880 ms | OOM at 2K |
All values from results/*.json (parity_*, selective_copy_*, imdb_*, efficiency_gpu.json). IMDb rows for Mamba-1/3 were skipped on wall-clock budget, documented in the README. Points past 6 GB VRAM ran via WDDM shared-memory spill, which also distorts near-OOM timings - recorded as data.
07/Findings
What the results say.
- 01
Complex input-dependent transitions are what restore state tracking: Mamba-3 is the only model above chance on cumulative parity (0.983 vs 0.51-0.56).
- 02
LTI fails selectively even when complex: S4D's fixed rotation cannot encode per-bit rotations, collapsing to 6.4% on selective copying.
- 03
Ranking flips with the task: selective SSMs win selective copying, but the Transformer and S4D win exact copying in distribution (1.00 at L=32) - the task formulation, not the architecture, decides.
- 04
O(T) operations are not fast wall-clock: the Python-loop selective scan is 87x slower than S4D's FFT at 1K (2 632 ms vs 19.5 ms) - kernel launches and autograd dominate.
- 05
The absolute-PE Transformer reproduces the published copying failure exactly: 1.00 in distribution at L=32, 0.036 at L=64.
08/Decisions
Why it was built this way.
One skeleton, five mixers
Every backbone shares blocks.py, so a benchmark difference isolates the mixer rather than the plumbing.
Slow scans on purpose
Keeping the selective scan as a Python loop turns the theory-vs-wall-clock gap into a measured, teachable result.
Verify math, not shapes
Each dual compute form has an equivalence test at atol/rtol 1e-3: FFT vs recurrent, chunked vs sequential, and the λ=1 reduction.
Chunk-grouped SSD
Processing decay matrices in bounded groups cut roughly 1 GB of transient memory at 32K.
Honest failures
Full S4 (NPLR), fused kernels and Mamba-3's MIMO form are declared not-implemented; OOM is recorded as data, not a crash.
09/Validation
How the claims are checked.
- A 14-test suite (tests/test_models.py) asserts S4D's FFT-convolution equals its recurrent step; Mamba-2's chunked scan equals the sequential reference for aligned (32) and misaligned (64) chunk sizes at atol/rtol 1e-3; Mamba-3 reduces to exponential-Euler exactly when λ=1, θ=0; and Mamba's step() equals a full forward.
- Dataset contract tests pin task semantics: parity is cumulative XOR, selective-copy masking, exact-copy BOS/SEP/EOS structure.
- Every benchmark JSON embeds its environment block (GPU, torch, CUDA version) and the seed, so numbers are attributable.
10/Limitations
Stated plainly.
- Single seed (42) on one 6 GB laptop GPU: rankings are measured, not statistical.
- No fused CUDA/Triton kernels: wall-clock results reflect the Python-loop implementations by design.
- Full S4 (NPLR/Cauchy) is not implemented - S4D only; Mamba-3's MIMO form and chunked dual are documented but not implemented.
- IMDb skips Mamba-1/3 on wall-clock budget; no language-modeling/perplexity evaluation.
- Near-OOM timings on Windows are distorted by WDDM shared-memory spill (documented in the README).
11/Interview TL;DR
- Built
- Five sequence models (Transformer, S4D, Mamba, Mamba-2, Mamba-3) from first principles in pure PyTorch, trained on three designed tasks and swept 1K-32K on a 6 GB RTX 4050.
- Why
- To verify each SSM generation's claims the only way that sticks: implement them, train at matched budgets, and measure - including the wall-clock cost the papers omit.
- How it works
- One shared block skeleton with swappable mixers; every dual compute form paired with an equivalence test; every number committed as JSON with its environment block.
- Key decision
- Deliberately slow Python-loop scans, which turned the O(T)-ops vs wall-clock gap into a measured finding (87x) instead of a footnote.
- Strongest result
- Mamba-3's complex transitions are the only thing above chance on cumulative parity (98.3% vs 51-56%), while LTI S4D collapses to 6% on selective copying.
- Main limitation
- Single seed on one laptop GPU with no fused kernels - rankings are measured, not statistical.
- Next
- Fused selective-scan kernels to close the wall-clock gap, and multi-seed runs for error bars.
Likely questions
Q01Why does ZOH produce (e^z-1)/z, and how do you keep it stable?
Exact ZOH for diagonal A gives B-bar = ΔB(e^z-1)/z with z = Δa. For |z| < 1e-4 the code switches to a Taylor series (1 + z/2 + z²/6) to avoid 0/0.
Q02Why can't LTI models do parity?
Cumulative parity needs input-dependent state rotations. S4D is complex but its rotation is fixed; real-transition eigenvalues cannot rotate at all. Only input-dependent complex transitions (Mamba-3) encode per-bit rotations.
Q03How do you prove chunked SSD equals the sequential scan?
A test runs both forms on identical inputs and asserts allclose at atol/rtol 1e-3, for aligned (32) and misaligned (64) chunk sizes.
Q04If Mamba is O(T), why is it 87x slower than S4D?
Operation count is not wall-clock. The Python-loop scan pays per-step kernel launches and autograd graph overhead, while S4D trains as one FFT convolution.
Q05What happens to the Transformer at 32K?
30 ms becomes 186 s per training step with ~9.9 GB peak memory (via WDDM spill): the quadratic term dominates, while O(T)-state models hold longer or OOM on the 6 GB card.