Lab 03 — State Space Models (S4 / Mamba)

Phase: 02 — Model Architecture Mastery | Difficulty: ⭐⭐⭐⭐⭐ | Time: 4–5 hours

Mamba (Gu & Dao, 2023) achieves linear time/memory in sequence length vs quadratic for attention.

What you build

discretize_zoh — Zero-Order Hold discretization: Ā = exp(ΔA), B̄ = (Ā-I)/A·B
s4_recurrence — Sequential SSM scan over a sequence
SelectiveSSM — Mamba's input-dependent B, C, Δ (selection mechanism)
MambaBlock — full block: in_proj + depthwise conv + SelectiveSSM + gating
compare_memory_usage — O(T²) attention vs O(T·d_state) SSM

Concept	What to understand
SSM state equation	h_t = Ā·h_{t-1} + B̄·x_t, y_t = C·h_t + D·x_t
ZOH discretization	Converts continuous-time A, B to discrete Ā, B̄
Selective scan	B, C, Δ are functions of x — allows content-based memory
Linear recurrence	O(N) sequential scan vs O(N²) attention matrix
Hardware-efficient	Mamba uses parallel scan + recomputation for GPU efficiency

File	Purpose
`lab.py`	Stubs with `TODO` markers
`solution.py`	Complete working implementation
`test_lab.py`	pytest suite
`requirements.txt`	Dependencies

pip install -r requirements.txt
python lab.py
pytest test_lab.py -v