Edge of stability: two solvable toys

Companion to Training Dynamics NB2 §3 (td2_3_edge_of_stability). Time is counted in steps k here; the notebook's continuous time is t = kη.

1. Self-stabilization (§2.2, §2.8)

Gradient descent on L(x, y) = ½ y x² + c y⁻⁵, c = 62.5. The coordinate y is the curvature along x, which is the sharpness S near x = 0 (for y ≳ 2.6); ℓ(y) = c y⁻⁵ pushes it up (progressive sharpening). Each step also adds noise of size 10⁻³ to x, a floor for the unstable direction (in a real network: other directions, minibatch noise, round-off).

sharpness S = y_k2/η_k

x_k (flips sign every step at EoS)±σ, central-flow amplitude

loss L(x_k, y_k) (log scale)

learning rate η

initial sharpness y₀

linear warmup η_k = η·min(1, (k+1)/T)

warmup length T (steps)

noise floor on x

presets:

mean of S, last 1000 steps
2/η (final)
mean x², last 1000 steps
σ² = −2ℓ′/S′ = 10c(η/2)⁶
peak loss

2. The catapult (§2.7)

The exact maps of the width-n model f = vᵀu/√n, one input, target 0, L = ½f²: f′ = f(1 − ηλ + η²f²/n), λ′ = λ + (ηf²/n)(ηλ − 4). Start at f₀ = 1.

loss ½f_k² (log scale)½n (catapult scale f² ~ n)

ηλ_k2 (stability edge)4 (divergence edge in this model)

ηλ₀

width n

presets:

final ηλ
peak loss (step)
early-growth estimate of steps to the f² ~ n scale, ln n / (2 ln(ηλ₀ − 1)) (ignores curvature feedback)

Both panels iterate the maps exactly; nothing is fitted. A check script (checks/eos_toy_and_catapult_check.py) runs the simulation code below under Node and compares it with a Python implementation.