week12_notebook.ipynb. Source:
view on GitHub.
This week is where the linear-attention story collides with classical numerical-analysis. DeltaNet looks like attention with rank-1 fast-weights, but the cleanest way to read it is as online learning: the state is a parameter matrix, the per-step update is one gradient step on $\tfrac{1}{2}\|S k_t - v_t\|^2$, and the so-called learning rate $\beta_t$ is exactly the explicit-Euler step size on the gradient flow of that loss. Longhorn replaces that explicit step with the closed-form implicit (one-Newton) step on a Tikhonov-regularised version of the same objective. The two algorithms therefore differ in one place only: how $\beta_t^{\text{eff}}$ is chosen.
That single change is enough to turn a conditionally stable algorithm into an unconditionally stable one --- the same structural improvement that implicit-midpoint enjoys over forward-Euler in ODE integration. The notebook walks the algebra, runs both algorithms, generates the spectral-radius comparison figure, then sanity-checks the chunkwise reformulation that production implementations use to amortize the recurrent cost across hardware-friendly GEMMs.
The notebook is intentionally self-contained. The PDF chapter (guides/chapters/week12.tex) walks the same material in book voice; this notebook is the working copy where you should run code, perturb inputs, and read shape contracts directly. The pattern, repeated section by section, is: name the object (DeltaNet update, Longhorn solve, chunk erase product), inspect the shape contract from the docstring, run the smallest convincing example, then connect the result back to the production reference at experiments/jax/week12/.
Read each markdown block before running the next code cell. When a block says Guided modify-and-run, perturb a single value (a $\beta$ scale, a key magnitude, an $\alpha$ trust-region weight, a chunk size) and rerun just that cell. The accompanying Expected observation tells you what direction the shape, residual, or spectral-radius readout should move. If your prediction was wrong, that is useful: the boundary you just found is a thing the chapter is trying to teach.
The setup cell wires standard plotting helpers, imports the W12 modules, and configures CPU-only JAX so this notebook is safe to run on the workstation regardless of whether the research-kb daemon is holding GPU memory. Production reference paths printed below are the on-disk modules whose docstrings should be your first stop when something looks suspicious.
import sys
from pathlib import Path
# Add shared utilities to path (mirrors the W01-W11 setup pattern).
for p in ["../shared", "guides/shared"]:
resolved_path = Path(p).resolve()
if (resolved_path / "plot_utils.py").exists():
if str(resolved_path) not in sys.path:
sys.path.insert(0, str(resolved_path))
break
from plot_utils import setup_notebook, save_figure, SSM_COLORS
EXPERIMENT_DIR, FIGURE_DIR = setup_notebook(week=12)
import jax
import jax.numpy as jnp
import numpy as np
import matplotlib.pyplot as plt
from experiments.jax.week12.delta_rule import (
delta_rule_step,
delta_rule_recurrent,
delta_rule_naive,
)
from experiments.jax.week12.longhorn import (
longhorn_effective_beta,
longhorn_step,
longhorn_recurrent,
)
from experiments.jax.week12.stability_analysis import (
deltanet_spectral_radius,
longhorn_spectral_radius,
deltanet_a_stability_boundary,
)
from experiments.jax.week12.chunkwise import (
chunk_wy_representation,
apply_wy_to_state,
delta_rule_chunkwise,
)
from experiments.jax.week12.figures import make_stability_figure
print(f"JAX {jax.__version__} | backend: {jax.default_backend()} | week dir: {EXPERIMENT_DIR}")
The DeltaNet recurrence is
$$S_t = S_{t-1}(I - \beta_t k_t k_t^\top) + \beta_t v_t k_t^\top, \qquad o_t = S_t q_t.$$
The textbook delta-rule rearrangement makes the online-learning interpretation transparent:
$$S_t = S_{t-1} + \beta_t (v_t - S_{t-1} k_t)\, k_t^\top.$$
The factor $(v_t - S_{t-1} k_t)$ is the per-step retrieval error (what we want to read at $k_t$ minus what the current state would actually return), and the rank-1 outer product writes a correction along $k_t^\top$. This is literally one step of explicit gradient descent on $\tfrac{1}{2}\|S k_t - v_t\|^2$ with step size $\beta_t$. DeltaNet is an online learner whose state is its parameter vector.
Guided modify-and-run. Increase the magnitude of betas until DeltaNet starts to oscillate. The state norm should grow once the explicit-Euler step size exceeds the A-stability boundary $\beta \|k\|^2 = 2$ along any key direction. Try betas = jnp.full((B, L), 0.4) first, then 2.5 / k_norm_squared to confirm the per-step rule.
Expected observation. With moderate $\beta$ the recurrent driver and the naive Python-loop oracle agree to float32 precision (max diff $< 10^{-4}$), and the final state norm is comparable to $\|v\|/\beta$. The recurrent + naive equivalence is the test oracle that locks the code's correctness; the residual you print is a concrete proxy for that test.
B, L, d_k, d_v = 2, 16, 4, 5
key_split = jax.random.split(jax.random.key(0), 4)
queries = jax.random.normal(key_split[0], (B, L, d_k))
keys = jax.random.normal(key_split[1], (B, L, d_k))
values = jax.random.normal(key_split[2], (B, L, d_v))
betas = 0.3 * jax.nn.sigmoid(jax.random.normal(key_split[3], (B, L)))
out_rec, state_rec = delta_rule_recurrent(queries, keys, values, betas)
out_naive, state_naive = delta_rule_naive(queries, keys, values, betas)
print(f"output shape: {out_rec.shape} (expected: (B={B}, L={L}, d_v={d_v}))")
print(f"final state shape: {state_rec.shape} (expected: (B={B}, d_v={d_v}, d_k={d_k}))")
print(f"recurrent vs naive max diff: {float(jnp.max(jnp.abs(out_rec - out_naive))):.2e}")
print(f"final state norm: {float(jnp.linalg.norm(state_rec)):.3f}")
Longhorn replaces the explicit gradient step with the closed-form solve of the regularised version,
$$S_t = \arg\min_{S} \; \tfrac{1}{2}\|S k_t - v_t\|^2 + \tfrac{\alpha_t}{2}\|S - S_{t-1}\|_F^2.$$
Setting the gradient with respect to $S$ to zero gives the implicit recurrence $\alpha_t (S_t - S_{t-1}) + (S_t k_t - v_t) k_t^\top = 0$. Right-multiplying by $k_t$ to eliminate $S_t k_t$ closes the system, and the post-update state has the same rank-1 form as DeltaNet:
$$S_t = S_{t-1} + \frac{v_t - S_{t-1} k_t}{\alpha_t + \|k_t\|^2}\, k_t^\top.$$
So Longhorn = DeltaNet with effective rate $\beta_t^{\text{eff}} = 1 / (\alpha_t + \|k_t\|^2)$. The denominator self-attenuates: as the key magnitude grows the rate shrinks, keeping the per-step contraction inside its stable region no matter how large $\|k_t\|$ becomes. The cell below verifies the equivalence by feeding DeltaNet the calibrated rate explicitly.
Guided modify-and-run. Scale keys by 10x or 100x and watch the Longhorn final-state norm stay bounded. Then run the same scaled inputs through DeltaNet with a fixed $\beta = 0.5$ --- you should see the state norm explode by orders of magnitude.
Expected observation. With small $\alpha$ and large keys, Longhorn's $\beta^{\text{eff}}$ shrinks like $1/\|k\|^2$, so the effective step size never crosses the DeltaNet A-stability boundary. The equivalence cell prints a max-diff $< 10^{-4}$ between Longhorn and DeltaNet at the calibrated rate; that is the structural identity.
alphas = 0.5 + jax.nn.softplus(jax.random.normal(jax.random.key(7), (B, L)))
out_lh, state_lh = longhorn_recurrent(queries, keys, values, alphas)
# Calibrated DeltaNet: feed in beta = 1 / (alpha + ||k||^2) per step.
k_sq = jnp.einsum("blk,blk->bl", keys, keys)
betas_calibrated = 1.0 / (alphas + k_sq)
out_dn_cal, _ = delta_rule_recurrent(queries, keys, values, betas_calibrated)
beta_eff = longhorn_effective_beta(keys, alphas)
print(f"longhorn output shape: {out_lh.shape}")
print(f"effective-rate range: beta_eff in [{float(beta_eff.min()):.4f}, {float(beta_eff.max()):.4f}]")
print(f"longhorn vs calibrated DeltaNet max diff: {float(jnp.max(jnp.abs(out_lh - out_dn_cal))):.2e}")
# Stability stress test: scale keys by 10x and watch the state norms.
big_keys = 10.0 * keys
_, state_lh_big = longhorn_recurrent(queries, big_keys, values, alphas)
_, state_dn_big = delta_rule_recurrent(queries, big_keys, values, jnp.full_like(betas, 0.5))
print(f"after 10x key scaling — Longhorn |S|: {float(jnp.linalg.norm(state_lh_big)):.2f}, "
f"DeltaNet (beta=0.5) |S|: {float(jnp.linalg.norm(state_dn_big)):.2e}")
The linearised iteration along the key direction is rank-1: $E_t = E_{t-1}(I - \beta_{\text{eff}} k k^\top)$, with the only non-trivial eigenvalue $1 - \beta_{\text{eff}} \|k\|^2$. Stability requires $|1 - \beta_{\text{eff}}\|k\|^2| < 1$, i.e. $\beta_{\text{eff}}\|k\|^2 \in (0, 2)$.
For DeltaNet, $\beta_{\text{eff}} = \beta_t$ is unconstrained. Crossing $\beta\|k\|^2 = 2$ is the explicit-Euler instability mode --- the deviation from the fixed point alternates sign and grows. For Longhorn, $\beta_{\text{eff}} = 1/(\alpha + \|k\|^2)$, so $\beta_{\text{eff}}\|k\|^2 = \|k\|^2/(\alpha + \|k\|^2) < 1$ for any $\alpha > 0$. The implicit step cannot cross the A-stability boundary, regardless of the key magnitude.
The figure below renders both spectral radii. Panel A plots the DeltaNet absolute eigenvalue $|1 - \beta\|k\|^2|$ against $\beta\|k\|^2$, marking the boundary at 2. Panel B plots the Longhorn $\rho = \alpha / (\alpha + \|k\|^2)$ against $\|k\|^2/\alpha$ on a log axis, showing that it asymptotes to 1 from below but never crosses.
Guided modify-and-run. Print deltanet_spectral_radius(jnp.array(b), jnp.array(k_sq)) for a grid of $b\|k\|^2$ values straddling 2, and longhorn_spectral_radius(jnp.array(a), jnp.array(k_sq)) for $\alpha = 0.01$ and a range of large $\|k\|^2$. The DeltaNet curve crosses 1; the Longhorn curve asymptotes to it.
Expected observation. The numerical eigenvalues match the analytic formulas to float32 precision (the tests in test_stability_analysis.py cross-check this against the actual Rayleigh quotient of the iteration matrix). The figure has two panels with a clear structural difference: a V-shape that crosses 1 vs a monotone curve bounded below 1.
fig, axes = make_stability_figure()
save_figure(fig, FIGURE_DIR, "stability_comparison")
plt.show()
# Numerical sanity: probe the formulas on a grid of step-size products.
products = jnp.linspace(0.0, 3.0, 13)
print("DeltaNet rho sweep (beta * ||k||^2):")
for p in products:
rho = float(deltanet_spectral_radius(jnp.array(1.0), jnp.array(float(p))))
flag = " unstable" if rho > 1.0 else ""
print(f" beta||k||^2 = {float(p):.2f} -> rho = {rho:.3f}{flag}")
print(f"\nA-stability boundary constant: {deltanet_a_stability_boundary()} (matches docstring)")
# Longhorn never crosses 1, even for ||k||^2 / alpha = 10000.
ratios = jnp.array([0.01, 1.0, 100.0, 10000.0])
print("\nLonghorn rho sweep (||k||^2 / alpha):")
for r in ratios:
rho = float(longhorn_spectral_radius(jnp.array(1.0), jnp.array(float(r))))
print(f" ||k||^2/alpha = {float(r):>8.2f} -> rho = {rho:.6f}")
Production DeltaNet implementations replace the $O(L)$ recurrent loop with a chunkwise reformulation. Splitting the sequence into chunks of length $C$, each chunk's contribution to the state is
$$S_{c} = S_{c-1} P_c + R_c, \qquad P_c = \prod_{t \in \text{chunk}} (I - \beta_t k_t k_t^\top), \qquad R_c = \text{(chunk write contribution)}.$$
The erase product $P_c$ admits the Householder/WY representation $P_c = I - W_c^\top Y_c$, where $W_c$ and $Y_c$ are $(C, d_k)$ matrices built recursively from the chunk's keys and betas. Applying $P_c$ to the $(d_v, d_k)$ state matrix is then one $(d_v, d_k)$-shaped GEMM rather than $C$ rank-1 updates --- the matmul-friendly form that production fused kernels exploit.
delta_rule_chunkwise is a pure-JAX implementation of the chunkwise driver. The within-chunk read still uses jax.lax.scan rather than a fused parallel kernel (this notebook is in pure JAX), but the WY primitives are exposed so a future fused-kernel rewrite has a tested in-tree reference.
Guided modify-and-run. Try chunk sizes 1, 2, 4, 8, 16, then a non-divisor (3) to confirm the input-validation path. The chunkwise driver should match the recurrent driver to float32 precision for every divisor of $L$.
Expected observation. All divisors of $L$ produce identical outputs (modulo float accumulation order, $< 10^{-4}$). The WY product reconstructs the explicit erase product exactly. If you see a discrepancy, the most likely culprit is a transpose mistake in apply_wy_to_state; that path is covered by test_apply_wy_matches_dense_matmul in experiments/jax/week12/test_chunkwise.py.
# Chunkwise vs recurrent equivalence sweep.
for cs in [1, 2, 4, 8, 16]:
out_chunk, _ = delta_rule_chunkwise(queries, keys, values, betas, chunk_size=cs)
diff = float(jnp.max(jnp.abs(out_chunk - out_rec)))
print(f"chunk_size={cs:>2d}: max diff vs recurrent = {diff:.2e}")
# WY representation demo on a single chunk.
chunk_keys = keys[0, :4] # first 4 steps of batch 0
chunk_betas = betas[0, :4]
W, Y = chunk_wy_representation(chunk_keys, chunk_betas)
# Reconstruct the erase product from WY and compare against the explicit form.
P_wy = jnp.eye(d_k) - W.T @ Y
P_explicit = jnp.eye(d_k)
for t in range(4):
kt, bt = chunk_keys[t], chunk_betas[t]
P_explicit = P_explicit @ (jnp.eye(d_k) - bt * jnp.outer(kt, kt))
print(f"\nWY product matches explicit erase product: max diff = "
f"{float(jnp.max(jnp.abs(P_wy - P_explicit))):.2e}")
# apply_wy_to_state matches dense matmul.
state = jax.random.normal(jax.random.key(99), (d_v, d_k))
state_explicit = state @ P_explicit
state_via_wy = apply_wy_to_state(state, W, Y)
print(f"apply_wy_to_state matches dense matmul: max diff = "
f"{float(jnp.max(jnp.abs(state_explicit - state_via_wy))):.2e}")
The week's working skill is reading a sequence-modelling layer as an online learner with a numerical-integration heritage. Three layers stay separate. First is the mathematical object: a per-step convex problem (the squared retrieval error, optionally with a Tikhonov regulariser) whose minimiser is the next state. Second is the implementation syntax: a jax.lax.scan over time with rank-1 einsum updates inside, plus the chunkwise factorization that replaces $L$ scan steps with $L/C$ chunk applications. Third is the validation habit: a recurrent-vs-naive equivalence check (locks correctness), a Longhorn-vs-DeltaNet calibrated-rate equivalence check (locks the structural identity), and a chunkwise-vs-recurrent equivalence check (locks the chunkwise refactor).
A good reading pass is slow and mechanical. Before each code cell predict the output shape and the residual magnitude. Then check that prediction before reading the printed value. The shape contracts are documented in every function's docstring; the equivalence residuals are bounded by the float32 accumulation order ($\sqrt{L \cdot d_k} \cdot 10^{-7}$ rule of thumb). If a residual lands well above that bound, the printed value has discovered something the docstring claims is not true.
When a markdown cell talks about a spectral radius, it is making a quantitative claim about how a small deviation from the algorithm's fixed point evolves under repeated application of the linearised iteration. The fixed point in this notebook is $S^\star = v k^\top / \|k\|^2$ (the unique rank-1 state that retrieves $v$ at $k$); the linearised iteration is rank-1; and the eigenvalue along $k$ is exactly $1 - \beta_{\text{eff}} \|k\|^2$. Eigenvalues orthogonal to $k$ are 1, which is feature, not a bug --- neither algorithm forgets in directions it has not seen a key for.
Use one-edit experiments. Change one chunk size, $\beta$ scale, $\alpha$ scale, key magnitude, or sequence length. Do not change three things at once. The expected observation should be simple enough to state before rerunning. For chunk sizes that divide $L$, the chunkwise output should be identical to the recurrent output. For DeltaNet $\beta \|k\|^2 < 2$, the state norm should stay bounded; past 2 it should grow geometrically. For any $\alpha > 0$ the Longhorn state norm should stay bounded regardless of how large $\|k\|$ becomes.
When a one-edit experiment violates your prediction, the failure is informative. The notebook is intentionally written so that the most likely failure modes (transpose errors in the WY representation, calibrated-rate mismatches between DeltaNet and Longhorn, off-by-one in chunk indexing) produce visibly large residuals. Read the residual magnitude as the diagnostic; do not silently accept "close enough" if the equivalence claim says they should be float32-identical.
For each section in the notebook you should be able to fill in this checklist: input shape, output shape, state shape carried through scan, residual against the reference implementation, and the failure mode that would matter. Section 1 (DeltaNet): residual against the naive Python-loop oracle should be float32-precision; failure mode is a wrong eraser sign. Section 2 (Longhorn): residual against DeltaNet at calibrated rate should be float32-precision; failure mode is a wrong $\beta^{\text{eff}}$ formula. Section 3 (Stability): the analytic spectral-radius formulas should match the numerical Rayleigh quotient of the iteration matrix to float32-precision; failure mode is a sign error in the $1 - \beta\|k\|^2$ derivation. Section 4 (Chunkwise): residual against the recurrent driver should be float32-precision for every chunk size dividing $L$; failure mode is a transpose error in the WY $P_c = I - W_c^\top Y_c$ identity.
Pause before the summary and explain the four sections aloud in three passes. First, describe each algorithm's mathematical object without code: DeltaNet is one explicit gradient step on a squared retrieval error, Longhorn is the closed-form implicit step on the Tikhonov-regularised version, the stability analysis linearises both around the fixed-point and asks for the spectral radius of the iteration, the chunkwise factorization rewrites the $L$-step scan as $L/C$ chunk applications via the Householder/WY identity. Second, describe the array shapes and the scan carry: state is $(B, d_v, d_k)$ throughout, the scan carry is the state matrix, and the within-chunk read uses einsum with a post-update state. Third, describe the validation signals: recurrent-vs-naive, Longhorn-vs-DeltaNet-at-calibrated-rate, analytic-vs-numerical spectral radius, and chunkwise-vs-recurrent. If any of those equivalences silently fails, the test suite under experiments/jax/week12/ should catch it.
A strong answer should include at least one shape statement and one failure statement per section. For DeltaNet: "the state is $(B, d_v, d_k)$, and a sign error in the eraser would make the residual against the naive oracle blow up to $O(1)$." For Longhorn: "the effective rate is a scalar per step, and a wrong $\beta^{\text{eff}}$ formula would make the calibrated-rate equivalence fail by an $O(1)$ residual." For stability: "the iteration matrix is $d_k \times d_k$, the only non-trivial eigenvalue is $1 - \beta\|k\|^2$, and a sign error in deltanet_spectral_radius would put the analytic curve below the unit-circle reference everywhere." For chunkwise: "the WY factors are $(C, d_k)$ each, and a transpose error in apply_wy_to_state would make the chunkwise-vs-recurrent residual blow up to $O(1)$."
You should now be able to state the week in code terms: DeltaNet is online learning under explicit Euler, Longhorn is online learning under implicit-midpoint, the stability story is the textbook explicit/implicit comparison applied to the linearised retrieval-error gradient flow, and the chunkwise reformulation is the production-efficiency variant that exposes within-chunk dense matmul. The reference implementations live at experiments/jax/week12/{delta_rule,longhorn,stability_analysis,chunkwise}.py; if you want to go deeper, read each module next to its test file rather than treating the import as magic.
Run pytest experiments/jax/week12/ -v after changing any algorithm or before opening a PR. The test suite includes:
test_delta_rule.py — recurrent-vs-naive equivalence, shape contracts, beta=0 no-op, autoassociative recall.test_longhorn.py — effective-rate formula, optimality condition, equivalence to DeltaNet at calibrated rate, unconditional stability.test_stability_analysis.py — analytic spectral-radius formulas vs numerical Rayleigh quotient, A-stability boundary constant, decay rate prediction.test_figures.py — figure curves match analytic formulas; rho < 1 for Longhorn over the full plotted range.test_chunkwise.py — WY product reconstructs erase product; chunkwise driver matches recurrent for every chunk size dividing $L$.For notebook-only edits, also run make notebook-check so missing required phrases, sparse markdown, and committed outputs fail immediately.
Week 13 picks up the implicit-step thread by extending the trust-region weight $\alpha_t$ from a scalar (Longhorn) to a structured operator (Gated DeltaNet's data-dependent forgetting). The math machinery of Sections 2 and 3 transfers directly: the spectral-radius analysis is parameterised by $\beta^{\text{eff}}$, and any choice of $\alpha$ that keeps $\beta^{\text{eff}}\|k\|^2 < 1$ produces an unconditionally stable update. Week 14 lifts the chunkwise reformulation into a parallel-within-chunk variant by replacing the within-chunk scan with the WY-based block algorithm sketched in Section 4 --- the in-tree reference primitives chunk_wy_representation and apply_wy_to_state are written specifically to be reused there.