week04_notebook.ipynb. Source:
view on GitHub.
Week 4 is where the ODE language becomes an ML layer. The continuous matrix A gives you memory dynamics; discretization makes it step through tokens; the convolutional kernel lets the same recurrence be trained efficiently over an entire sequence.
This notebook is intentionally standalone. The PDF chapter gives a more polished narrative, but you should be able to run this notebook by itself and still understand what the computation is doing. The repeated pattern is: name the mathematical object, inspect the shape contract, run the smallest useful example, then connect the result back to the production implementation at experiments/jax/week04/s4_hippo.py.
Read each markdown block before running the next code cell. When a block says Guided modify-and-run, make the suggested small edit, rerun just that cell, and compare your observation with the inline Expected observation. These prompts are deliberately small because the goal is code literacy: you should know which axis is batch, which axis is time, where parameters live, and what failure would look like before you move to larger experiments.
This setup cell imports the S4/HiPPO production functions. The imported names are the week vocabulary: make a continuous matrix, discretize it, generate a kernel, compare recurrence with convolution, and wrap the result in an NNX model.
Guided modify-and-run. Change one small value in the next code cell, such as a sequence length, state dimension, step size, feature count, or initialization scale. Rerun the cell and check whether the reported shape or numerical comparison changes in the direction you expected.
Expected observation. A correct change should preserve the documented shape contract and should not break equivalence checks unless you intentionally violate an assumption. If a shape error appears, read it as useful feedback: the week is teaching you where the abstraction boundary really is.
import sys
from pathlib import Path
# Add shared utilities to path
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, create_tufte_figure, apply_tufte_style,
add_subtle_grid, set_tufte_title, set_tufte_labels, direct_label,
direct_label_line, SSM_COLORS, C_ACCENT, C_ALERT, C_HIGHLIGHT, C_BASELINE,
)
EXPERIMENT_DIR, FIGURE_DIR = setup_notebook(week=4)
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
from experiments.jax.week04.s4_hippo import (
make_hippo_legs, discretize_zoh, discretize_bilinear,
ssm_kernel_naive, ssm_recurrent, ssm_convolutional,
S4Layer, S4SeqModel, eigenvalue_analysis, make_delay_batch,
)
from experiments.jax.common.harness import (
ExperimentHarness, SSMParam, SimpleSeqModel, make_noisy_sine_batch,
)
from flax import nnx
print(f"JAX {jax.__version__} on {jax.default_backend()}")
print(f"Devices: {jax.devices()}")
HiPPO-LegS construction is the first concept-to-code bridge. The matrix is not arbitrary; its spectrum controls continuous-time memory. Expected observation: eigenvalues lie in the stable left half-plane before discretization.
Guided modify-and-run. Change one small value in the next code cell, such as a sequence length, state dimension, step size, feature count, or initialization scale. Rerun the cell and check whether the reported shape or numerical comparison changes in the direction you expected.
Expected observation. A correct change should preserve the documented shape contract and should not break equivalence checks unless you intentionally violate an assumption. If a shape error appears, read it as useful feedback: the week is teaching you where the abstraction boundary really is.
# HiPPO-LegS eigenvalue distribution
fig, axes = create_tufte_figure(1, 2, figsize=(10, 4))
for i, N in enumerate([16, 64]):
A, B = make_hippo_legs(N)
eigs = jnp.linalg.eigvals(A)
ax = axes[i]
ax.scatter(eigs.real, eigs.imag, c=C_ACCENT, s=20, alpha=0.7, edgecolors='white', lw=0.3)
ax.axvline(x=0, color=C_ALERT, linestyle='--', alpha=0.5)
set_tufte_labels(ax, r'Re($\lambda$)', r'Im($\lambda$)')
set_tufte_title(ax, f'HiPPO-LegS eigenvalues ($N={N}$)')
plt.tight_layout()
save_figure(fig, FIGURE_DIR, 'ex2_hippo_eigenvalues')
plt.show()
print(f"N=64: max Re(lambda) = {jnp.max(eigs.real):.4f} (should be < 0)")
This cell maps continuous eigenvalues to discrete eigenvalues. ZOH uses an exponential map, while bilinear uses a rational transform. Modify dt and observe how the discrete spectrum moves relative to the unit circle.
Guided modify-and-run. Change one small value in the next code cell, such as a sequence length, state dimension, step size, feature count, or initialization scale. Rerun the cell and check whether the reported shape or numerical comparison changes in the direction you expected.
Expected observation. A correct change should preserve the documented shape contract and should not break equivalence checks unless you intentionally violate an assumption. If a shape error appears, read it as useful feedback: the week is teaching you where the abstraction boundary really is.
# Eigenvalue migration: continuous -> discrete
A, B = make_hippo_legs(32)
dt = 0.01
Ab_zoh, _ = discretize_zoh(A, B, dt)
Ab_bil, _ = discretize_bilinear(A, B, dt)
cont_eigs = jnp.linalg.eigvals(A)
disc_eigs_zoh = jnp.linalg.eigvals(Ab_zoh)
disc_eigs_bil = jnp.linalg.eigvals(Ab_bil)
fig, axes = create_tufte_figure(1, 2, figsize=(10, 4))
# Continuous eigenvalues (left half-plane)
axes[0].scatter(cont_eigs.real, cont_eigs.imag, c=C_ACCENT, s=30, edgecolors='white', lw=0.3)
axes[0].axvline(x=0, color=C_ALERT, linestyle='--', alpha=0.5)
set_tufte_title(axes[0], 'Continuous: left half-plane')
set_tufte_labels(axes[0], r'Re($\lambda$)', r'Im($\lambda$)')
# Discrete eigenvalues (unit disk)
theta = jnp.linspace(0, 2*jnp.pi, 100)
axes[1].plot(jnp.cos(theta), jnp.sin(theta), color=C_ALERT, ls='--', alpha=0.5)
axes[1].scatter(disc_eigs_zoh.real, disc_eigs_zoh.imag, c=C_ACCENT, s=30, edgecolors='white', lw=0.3)
axes[1].scatter(disc_eigs_bil.real, disc_eigs_bil.imag, c=C_HIGHLIGHT, s=30, marker='x')
set_tufte_title(axes[1], f'Discrete ($\\Delta t={dt}$): unit disk')
set_tufte_labels(axes[1], r'Re($\bar\lambda$)', r'Im($\bar\lambda$)')
axes[1].set_aspect('equal')
# Direct labels instead of legend
direct_label(axes[1], float(disc_eigs_zoh.real[0]), float(disc_eigs_zoh.imag[0]),
'ZOH', offset=(-25, 5), fontsize=8, color=C_ACCENT)
direct_label(axes[1], float(disc_eigs_bil.real[0]), float(disc_eigs_bil.imag[0]),
'Bilinear', offset=(5, 5), fontsize=8, color=C_HIGHLIGHT)
plt.tight_layout()
save_figure(fig, FIGURE_DIR, 'ex1_discretization')
plt.show()
print(f"ZOH max |lambda| = {jnp.max(jnp.abs(disc_eigs_zoh)):.6f}")
print(f"Bilinear max |lambda| = {jnp.max(jnp.abs(disc_eigs_bil)):.6f}")
The recurrent-vs-convolutional equivalence is the key S4 trick. The same system can be unrolled as a recurrence or materialized as a kernel. Expected observation: the two outputs match to numerical tolerance.
Guided modify-and-run. Change one small value in the next code cell, such as a sequence length, state dimension, step size, feature count, or initialization scale. Rerun the cell and check whether the reported shape or numerical comparison changes in the direction you expected.
Expected observation. A correct change should preserve the documented shape contract and should not break equivalence checks unless you intentionally violate an assumption. If a shape error appears, read it as useful feedback: the week is teaching you where the abstraction boundary really is.
# Recurrent vs Convolutional equivalence
A, B = make_hippo_legs(16)
Ab, Bb = discretize_zoh(A, B, dt=0.01)
C = jax.random.normal(jax.random.key(0), (1, 16)) * 0.1
D = jnp.array(1.0)
L = 128
u = jax.random.normal(jax.random.key(1), (L,))
# Compute both modes
K = ssm_kernel_naive(Ab, Bb, C, L)
y_conv = ssm_convolutional(K, D, u)
y_rec = ssm_recurrent(Ab, Bb, C, D, u)
max_diff = jnp.max(jnp.abs(y_conv - y_rec))
fig, axes = create_tufte_figure(2, 1, figsize=(10, 5), sharex=True)
t = jnp.arange(L)
axes[0].plot(t, y_conv, alpha=0.8, lw=1.2, color=C_ACCENT)
axes[0].plot(t, y_rec, '--', alpha=0.8, lw=1.2, color=C_ALERT)
set_tufte_title(axes[0], 'Recurrent vs convolutional: identical outputs')
set_tufte_labels(axes[0], ylabel='Output')
direct_label(axes[0], float(t[-1]), float(y_conv[-1]), 'Conv (FFT)', offset=(5, 3), fontsize=8, color=C_ACCENT)
direct_label(axes[0], float(t[-1]), float(y_rec[-1]), 'Recurrent', offset=(5, -10), fontsize=8, color=C_ALERT)
axes[1].plot(t, jnp.abs(y_conv - y_rec), color=C_HIGHLIGHT, alpha=0.8, lw=1.0)
set_tufte_title(axes[1], f'Absolute difference (max = {max_diff:.2e})')
set_tufte_labels(axes[1], 'Time step', '|difference|')
axes[1].set_yscale('log')
plt.tight_layout()
save_figure(fig, FIGURE_DIR, 'ex3_equivalence')
plt.show()
print(f"Max absolute difference: {max_diff:.2e} (threshold: 1e-4)")
Training the S4 model makes the state matrix learned rather than fixed. The notebook inspects eigenvalues before and after training so you can connect optimization to stability geometry.
Guided modify-and-run. Change one small value in the next code cell, such as a sequence length, state dimension, step size, feature count, or initialization scale. Rerun the cell and check whether the reported shape or numerical comparison changes in the direction you expected.
Expected observation. A correct change should preserve the documented shape contract and should not break equivalence checks unless you intentionally violate an assumption. If a shape error appears, read it as useful feedback: the week is teaching you where the abstraction boundary really is.
# Train S4SeqModel and analyze eigenvalues before/after
model = S4SeqModel(features=16, d_state=16, rngs=nnx.Rngs(0))
pre_results = eigenvalue_analysis(model)
harness = ExperimentHarness(model, backbone_lr=1e-3, ssm_lr=1e-3, total_steps=300)
losses = harness.train(make_noisy_sine_batch, n_steps=300, batch_size=8)
post_results = eigenvalue_analysis(model)
fig, axes = create_tufte_figure(1, 3, figsize=(14, 4))
# Loss curve
axes[0].plot(losses, color=C_ACCENT, lw=1.0)
set_tufte_title(axes[0], 'S4SeqModel training')
set_tufte_labels(axes[0], 'Step', 'MSE loss')
axes[0].set_yscale('log')
# Continuous eigenvalues: before vs after
for r in pre_results:
e = r['continuous_eigenvalues']
axes[1].scatter(e.real, e.imag, c=C_ACCENT, s=10, alpha=0.3)
for r in post_results:
e = r['continuous_eigenvalues']
axes[1].scatter(e.real, e.imag, c=C_ALERT, s=10, alpha=0.3)
axes[1].axvline(x=0, color=C_ALERT, linestyle='--', alpha=0.5)
set_tufte_title(axes[1], 'Continuous eigenvalues')
set_tufte_labels(axes[1], r'Re($\lambda$)', r'Im($\lambda$)')
direct_label(axes[1], -0.5, 0.5, 'init', fontsize=8, color=C_ACCENT, offset=(3, 0))
direct_label(axes[1], -0.5, -0.5, 'trained', fontsize=8, color=C_ALERT, offset=(3, 0))
# Discrete eigenvalues: before vs after
theta = jnp.linspace(0, 2*jnp.pi, 100)
axes[2].plot(jnp.cos(theta), jnp.sin(theta), color=C_BASELINE, ls='--', alpha=0.3)
for r in pre_results:
e = r['discrete_eigenvalues']
axes[2].scatter(e.real, e.imag, c=C_ACCENT, s=10, alpha=0.3)
for r in post_results:
e = r['discrete_eigenvalues']
axes[2].scatter(e.real, e.imag, c=C_ALERT, s=10, alpha=0.3)
set_tufte_title(axes[2], 'Discrete eigenvalues (unit disk)')
set_tufte_labels(axes[2], r'Re($\bar\lambda$)', r'Im($\bar\lambda$)')
axes[2].set_aspect('equal')
plt.tight_layout()
save_figure(fig, FIGURE_DIR, 'ex5_training_eigenvalues')
plt.show()
max_disc = max(r['max_abs_discrete'] for r in post_results)
max_cont = max(r['max_real_continuous'] for r in post_results)
print(f"After training: max |lambda_disc| = {max_disc:.6f}, max Re(lambda_cont) = {max_cont:.6f}")
The delay-task comparison explains why long memory matters. Conv1D with a small kernel is local; S4 can carry information through recurrent state. Modify the delay or kernel size to find regimes where the baseline becomes competitive.
Guided modify-and-run. Change one small value in the next code cell, such as a sequence length, state dimension, step size, feature count, or initialization scale. Rerun the cell and check whether the reported shape or numerical comparison changes in the direction you expected.
Expected observation. A correct change should preserve the documented shape contract and should not break equivalence checks unless you intentionally violate an assumption. If a shape error appears, read it as useful feedback: the week is teaching you where the abstraction boundary really is.
# S4 vs Conv1D on delay task (delay=8 > kernel_size=3)
delay = 8
n_steps = 300
s4_model = S4SeqModel(features=16, d_state=16, rngs=nnx.Rngs(0))
s4_harness = ExperimentHarness(s4_model, backbone_lr=1e-3, ssm_lr=1e-3, total_steps=n_steps)
s4_losses = s4_harness.train(
lambda key, bs: make_delay_batch(key, bs, seq_len=64, delay=delay),
n_steps=n_steps, batch_size=8,
)
conv_model = SimpleSeqModel(16, kernel_size=3, rngs=nnx.Rngs(0))
conv_harness = ExperimentHarness(conv_model, backbone_lr=1e-3, total_steps=n_steps)
conv_losses = conv_harness.train(
lambda key, bs: make_delay_batch(key, bs, seq_len=64, delay=delay),
n_steps=n_steps, batch_size=8,
)
fig, ax = create_tufte_figure(figsize=(8, 4))
steps = list(range(n_steps))
ax.plot(steps, s4_losses, alpha=0.8, lw=1.2, color=C_ACCENT)
ax.plot(steps, conv_losses, alpha=0.8, lw=1.2, color=C_ALERT)
set_tufte_title(ax, f'Delay copy task (delay={delay} > kernel_size=3)')
set_tufte_labels(ax, 'Step', 'MSE loss')
ax.set_yscale('log')
# Direct labels at line endpoints with final loss values
s4_final = sum(s4_losses[-20:]) / 20
conv_final = sum(conv_losses[-20:]) / 20
direct_label_line(ax, steps, s4_losses, f'S4 ({s4_final:.3f})', position='end', fontsize=8, color=C_ACCENT)
direct_label_line(ax, steps, conv_losses, f'Conv1D ({conv_final:.3f})', position='end', fontsize=8, color=C_ALERT)
plt.tight_layout()
save_figure(fig, FIGURE_DIR, 'ex5_s4_vs_conv1d')
plt.show()
print(f"Final loss (last 20 avg) -- S4: {s4_final:.4f}, Conv1D: {conv_final:.4f}")
print(f"S4 advantage: {conv_final / s4_final:.1f}x lower loss")
This section makes the notebook intentionally self-contained. The week topic is S4 / HiPPO: Continuous to Discrete to Convolutional, and the working skill is mapping continuous dynamics into stable discrete sequence layers. The PDF may give a smoother chapter narrative, but the notebook should still answer four practical questions on its own: what object are you building, what syntax expresses it, what shape proves the code is wired correctly, and what numerical observation would make you suspicious.
For this week, keep three layers separate. First is the mathematical object: a recurrence, table, matrix, discretization, filter, or evidence record. Second is the implementation syntax: the JAX transform, Flax NNX state object, FFT call, dataclass, or array operation that expresses the object. Third is the validation habit: shape checks, equality against a reference implementation, stability checks, or a compact evidence ledger. When notebooks feel abstract, it is usually because these layers are collapsed into one imported function call. The markdown here deliberately pulls them apart before the code runs.
A good reading pass is slow and mechanical. Before a code cell, predict the output shape and the one numerical property that should hold. After the cell, check that prediction before looking at plots or final losses. If a plot is produced, ask what computation created every axis. If a scalar comparison is printed, ask which reference implementation it is comparing against. If a table row is produced, ask what would falsify it. This is the same discipline you will need when debugging the production path in experiments/jax/week04/s4_hippo.py.
Use one-edit experiments. Change one sequence length, state dimension, step size, feature count, decay rate, or row in a reading table. Do not change three things at once. The expected observation should be simple enough to say before rerunning: the leading axis should become longer, the output channel count should stay fixed, the stability magnitude should move toward zero, or the equivalence error should remain near machine precision. If the result violates that expectation, the notebook has done its job: it found the exact boundary where your mental model was incomplete.
For every major cell in this notebook, you should be able to fill in this checklist: input shape, output shape, state carried across time, parameter or evidence object being changed, reference result if one exists, and the failure mode that would matter. The answers are included inline because this is a self-teaching notebook. Use them actively: read the answer, rerun the cell, then explain why the answer follows from the code rather than memorizing the printed value.
Pause before the summary and explain the notebook aloud in three passes. First, describe the mathematical or methodological object without code. Second, describe the exact array or Python object that represents it. Third, describe the validation signal printed by the notebook. This deliberate repetition is useful because later weeks reuse the same ideas with more dimensions and less scaffolding.
A strong answer should include at least one shape statement and one failure statement. For example: "the sequence axis is length L, the state axis is size N, and the equivalence check would fail if the scan composition were ordered incorrectly." For reading weeks, replace shape with schema: "the evidence row has a family, claim, limitation, and follow-up, and the failure would be mixing evidence with opinion." This habit turns notebook reading into implementation readiness.
When you adapt the code, prefer small edits that preserve the contract. If a small edit breaks the notebook, do not immediately patch around the error. Read the message, identify the violated shape or state assumption, and then decide whether the production implementation should prevent that violation with validation or whether the notebook simply taught you an expected boundary.
You should now be able to state the week in code terms: what object is being represented, what shape each major tensor carries, which production file contains the durable implementation, and what observation confirms the implementation is behaving correctly. The notebook uses experiments/jax/week04/s4_hippo.py as the reference path; if you want to go deeper, read that file next to the tests rather than treating the import as magic.
Run pytest experiments/jax/week04/ -v after changing the implementation or adapting an example. For notebook-only edits, also run make notebook-check so sparse markdown, committed outputs, and missing syntax supplements fail immediately.
Week 5 diagonalizes the state matrix. You give up dense state interactions to gain simpler kernels, easier stability analysis, and faster implementations.