Chapters
A 21-week guide. Currently 6 of 21 chapters published.
Part I · Foundations
-
JAX Primitives Through Dynamical Systems
Internalize five JAX primitives — jit, grad, vmap, lax.scan, lax.associative_scan — by implementing dynamical systems you already know: damped oscillator, coupled oscillator ring, EMA. The scan signature (carry, x) → (carry, y) is exactly the SSM state update.
-
Flax NNX Modules and Training Loop
Build the minimal training pipeline reused for every SSM experiment in the curriculum: NNX modules, dual-optimizer pattern, nnx.scan for layer stacking, and the SSMParam variable type.
Part II · SSM Core
-
S4 / HiPPO: Continuous to Discrete
Implement S4 from continuous-time SSM through HiPPO initialization, ZOH/bilinear discretization, and the recurrent/convolutional duality that makes S4 practical. The pedagogical bridge that underpins every selective SSM in production.
-
Diagonal SSMs (S4D) and Stability Analysis
S4D — the diagonal simplification of S4 that drops kernel cost from O(N²L) to O(NL) and makes stability guaranteed by construction via log-parameterized eigenvalues.
Part III · Beyond SSMs
-
Hyena, Long Convolution, and Linear Attention
Two parallel attempts to escape attention's O(L²): Hyena (eclipsed by mechanistic evaluation on associative recall) and the linear-attention lineage (which became the production frontier). The contrast sets up the W12 Unification Quadruple.
-
Delta-Rule and Online Learning
DeltaNet as one explicit gradient step on per-token retrieval error; Longhorn as the closed-form implicit step. The Unification Quadruple (UIA, FWP, SSD, Longhorn) and chunkwise WY factorization that ships in Qwen 3.5/3.6 and Kimi Linear.