Chapters

A 21-week guide. Currently 6 of 21 chapters published.

Part I · Foundations

  1. Week 1 Published

    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.

  2. Week 2 Published

    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

  1. Week 4 Published

    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.

  2. Week 5 Published

    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

  1. Week 11 Published

    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.

  2. Week 12 Published

    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.