NB2 §3 — The edge of stability

Phenomenology (§3.1). The notebook’s key paper.

Self-stabilization (§3.2).

The central flow (§3.3). Start with the companion site, then the paper.

Rod Flow and principal flow (§3.4).

Also: A Rod Flow model for Adam at the edge of stability (follow-up for momentum methods and Adam).

Adaptive optimizers (§3.5).

The edge of stochastic stability (§3.6).

Also: Momentum further constrains sharpness at the edge of stochastic stability (ICML 2026; with momentum \beta, batch sharpness plateaus at 2(1-\beta)/\eta for small batches and 2(1+\beta)/\eta for large ones); SGD at the edge of stability: the stochastic sharpness gap (2026; predicts a \lambda_{\max} gap of roughly 1/b, quantitatively untested).

The catapult (§3.7).

Warmup and WSD (§3.8).

What happens to the sharpness S=\lambda_{\max}(H), the top eigenvalue of the loss Hessian H, during full-batch gradient descent with step \eta? It rises (progressive sharpening) to 2/\eta, the quadratic stability limit, and hovers there while the loss keeps falling, non-monotonically.

Paper: Cohen et al. 2021, Gradient descent on neural networks typically occurs at the edge of stability, the notebook’s key paper (see Readings; networks of several kinds, mostly on a 5,000-example CIFAR-10 subset).

Above 2/\eta a quadratic analysis predicts divergence along the top eigenvector u. Instead S stays at or just above 2/\eta (further above with cross-entropy); with heavy-ball momentum \beta the ceiling is (2+2\beta)/\eta. The sharpness adapts to the learning rate, which gradient flow, having no \eta in it, cannot do; Exercise 3.1 shows its small-\eta corrections cannot either. Why sharpening happens is a separate, partly open question, taken as given here (a tangent-kernel view: td2_6_curvature, §6.2).

Derivations and walkthrough. Prerequisites: Gradient descent on \tfrac12\lambda x^2 is x_{k+1}=(1-\eta\lambda)x_k: it converges iff 0<\eta\lambda<2, and for \eta\lambda>2 it flips sign every step and grows. On a positive-definite quadratic this holds along each Hessian eigenvector, so the constraint is \eta S<2. Time is t=k\eta after k steps, so gradient flow is \dot\theta=-\nabla L. A Hessian-vector product Hv=\nabla_\theta\langle\nabla L,v\rangle costs about two gradients; power iteration v\leftarrow Hv/\|Hv\| then estimates S as v^\top Hv (fine here, where the top eigenvalue is positive and dominant).

The experiment. A two-hidden-layer tanh MLP (width 64) fits 256 random \pm1 labels in \mathbb R^8 with MSE (random labels make sharpening fast), by full-batch gradient descent at three step sizes, with S measured every \Delta t=1.25. In t=k\eta all three runs follow the same gradient-flow curve until their own 2/\eta. Each then overshoots (to about 1.5\times2/\eta at \eta=0.2), falls below with a loss spike, and settles into shrinking oscillations around 2/\eta: the self-stabilization cycle of §3.2 at full size.

# A small MLP on 256 random +-1 labels; sharpness by power iteration on Hessian-vector products.
torch.manual_seed(0)
n_data, d_in, width = 256, 8, 64
X = torch.randn(n_data, d_in)
Y = torch.sign(torch.randn(n_data, 1))

def make_mlp(seed=0):
    torch.manual_seed(seed)
    return torch.nn.Sequential(torch.nn.Linear(d_in, width), torch.nn.Tanh(),
                               torch.nn.Linear(width, width), torch.nn.Tanh(),
                               torch.nn.Linear(width, 1))

def mse(net):
    return 0.5 * ((net(X) - Y) ** 2).mean()

def sharpness(net, v=None, iters=20):
    # Power iteration with Hv = d/dtheta <grad L, v>. Returns (Rayleigh quotient, last vector).
    # Pass the previous vector back in as a warm start: then a few iterations per call suffice.
    params = list(net.parameters())
    grads = torch.autograd.grad(mse(net), params, create_graph=True)
    if v is None:
        v = [torch.randn_like(p) for p in params]
    for _ in range(iters):
        norm = torch.sqrt(sum((a * a).sum() for a in v))
        v = [a / norm for a in v]
        Hv = torch.autograd.grad(grads, params, grad_outputs=v, retain_graph=True)
        lam = sum((a * b).sum() for a, b in zip(Hv, v)).item()
        v = [a.detach() for a in Hv]
    return lam, v

def gd_step(net, eta):
    loss = mse(net)
    grads = torch.autograd.grad(loss, list(net.parameters()))
    with torch.no_grad():
        for p, g in zip(net.parameters(), grads):
            p -= eta * g
    return loss.item()

def run_gd(eta, t_max, dt_measure=1.25, seed=0):
    net = make_mlp(seed)
    every = max(1, round(dt_measure / eta))
    v, ts, Ss, losses = None, [], [], []
    for k in range(round(t_max / eta) + 1):
        if k % every == 0:
            lam, v = sharpness(net, v, iters=30 if v is None else 6)
            ts.append(k * eta); Ss.append(lam)
        losses.append(gd_step(net, eta))
    return np.array(ts), np.array(Ss), np.array(losses)

t0 = time.time()
runs = {eta: run_gd(eta, t_max=250) for eta in (0.05, 0.1, 0.2)}
print(f"trained 3 runs in {time.time() - t0:.1f} s")

trained 3 runs in 17.6 s

Code (13 lines)

fig, ax = plt.subplots(1, 2, figsize=(11, 3.8))
for (eta, (ts, Ss, losses)), c in zip(runs.items(), ["C0", "C1", "C2"]):
    ax[0].plot(ts, Ss, color=c, label=f"$\\eta$ = {eta}")
    ax[0].axhline(2 / eta, color=c, ls="--", lw=1)
    ax[1].semilogy(np.arange(len(losses)) * eta, losses, color=c, lw=0.8)
ax[0].set(xlabel="$t = k\\eta$", ylabel="sharpness $S$", title="sharpness (dashed: $2/\\eta$)")
ax[1].set(xlabel="$t = k\\eta$", ylabel="training loss", title="loss: non-monotone at EoS, decreasing overall")
ax[0].legend(); plt.tight_layout(); plt.show()
for eta, (ts, Ss, _) in runs.items():
    t_cross = ts[np.argmax(Ss >= 2 / eta)]
    late = Ss[ts > t_cross + 0.5 * (ts[-1] - t_cross)]          # second half of the time spent at EoS
    print(f"eta = {eta:4}: reaches 2/eta = {2/eta:4.0f} at t = {t_cross:5.1f}; peak S = {Ss.max():5.1f}; "
          f"S over the second half of EoS: mean {late.mean():.2f}, range [{late.min():.2f}, {late.max():.2f}]")

eta = 0.05: reaches 2/eta =   40 at t = 120.0; peak S =  42.2; S over the second half of EoS: mean 40.07, range [39.71, 40.48]
eta =  0.1: reaches 2/eta =   20 at t =  50.4; peak S =  23.6; S over the second half of EoS: mean 20.03, range [19.73, 20.43]
eta =  0.2: reaches 2/eta =   10 at t =  34.8; peak S =  15.2; S over the second half of EoS: mean 10.08, range [9.85, 10.44]

Exercise 3.1 — The exact modified flow on a quadratic. Why can’t a small-step flow, however many correction terms it keeps, describe gradient descent at the edge? (If you did Exercise 1.2.c of td2_1_flows, this is the same calculation; only the last paragraph of the solution is new.) Background: for small \eta, one gradient step is, to O(\eta^3), a time-\eta step of the modified flow \dot\theta=-\nabla\big(L+\tfrac\eta4\|\nabla L\|^2\big) (implicit gradient regularization). Take gradient descent on L(x)=\tfrac12\lambda x^2 with \lambda>0, and t=k\eta.

1. For 0<\eta\lambda<1, find the constant a(\eta\lambda) such that the linear flow \dot x=a\,x, run for time \eta, reproduces one gradient step exactly.
2. Expand a in powers of \eta. Identify the first two terms with gradient flow and with the implicit-gradient-regularization correction -\tfrac\eta4\nabla\|\nabla L\|^2. What is the radius of convergence of the series in \eta\lambda?
3. What happens to a when 1<\eta\lambda<2, and when \eta\lambda>2? Show that no real one-dimensional autonomous ODE \dot x=F(x) can reproduce gradient descent for \eta\lambda>1, and say what this implies for any truncation of the modified-flow series at the edge of stability.

Code exercise 3.1 — Is the oscillation really along u? §3.2 assumes that at EoS the iterate oscillates with period 2 along the top Hessian eigenvector u. Check this in the MLP: train at \eta=0.2 for 1000 steps (into EoS), compute u by power iteration, run 40 more steps, and look at the projections \langle\theta_k-\bar\theta,u\rangle, where \bar\theta is the mean of those 40 iterates. Also report what fraction of each step \theta_{k+1}-\theta_k (in squared norm) lies along u.

# Code exercise 3.1: train at eta = 0.2 for 1000 steps, get u, then record 40 more iterates.
from torch.nn.utils import parameters_to_vector as p2v
# ...

from torch.nn.utils import parameters_to_vector as p2v

eta = 0.2
net = make_mlp(0)
for _ in range(1000):
    gd_step(net, eta)
lam, v = sharpness(net, None, iters=100)
u = torch.cat([t.flatten() for t in v]); u = u / u.norm()

thetas = []
for _ in range(40):
    thetas.append(p2v(net.parameters()).detach().clone())
    gd_step(net, eta)
T = torch.stack(thetas)
proj = (T - T.mean(0)) @ u
steps = T[1:] - T[:-1]
frac = ((steps @ u) ** 2 / (steps ** 2).sum(1)).mean().item()
print(f"S = {lam:.3f}, 2/eta = {2/eta}")
print("projections onto u:", np.round(proj[:10].numpy(), 4))
print(f"fraction of |step|^2 along u: {frac:.2f}")

S = 10.062, 2/eta = 10.0
projections onto u: [ 0.0422 -0.0422  0.0431 -0.0431  0.0438 -0.0438  0.0442 -0.0443  0.0445
 -0.0445]
fraction of |step|^2 along u: 0.45

Once \eta S>2, why doesn’t the iterate blow up along the top eigenvector u? Because the oscillation lowers the sharpness: through the third derivative of the loss, an oscillation of size x along u also moves the iterate by -\tfrac\eta2x^2\nabla S per step, down the gradient of the sharpness S.

Paper: Damian, Nichani & Lee, Self-stabilization (see Readings).

The push is even in x, so it survives averaging over the period-2 bounce. When \eta S>2 the oscillation grows, the push grows with it, and S is driven back to 2/\eta. It all fits in a two-parameter toy, L(x,y)=\tfrac12S(y)x^2+\ell(y) with S'>0>\ell': x oscillates along u, and y is a slow coordinate along which the loss falls as the sharpness rises.

▶ Interactive figure: eos_toy_and_catapult.html (opens full page). Panel 1 is this toy, with sliders for \eta and the initial sharpness and an optional warmup ramp (its noise uses a different random generator, so runs near a survival boundary can differ from the notebook’s).

Derivations and walkthrough. Notation: Parameters \theta, Hessian H, sharpness S(\theta)=\lambda_{\max}(H) with unit top eigenvector u, \nabla S its gradient with respect to \theta. \nabla^3L[a,b] is the third-derivative tensor contracted with two vectors (a vector). Write the iterate as a centre plus an oscillation along u: \theta=\bar\theta+xu.

The expansion. To second order, \nabla L(\bar\theta+\delta)=\nabla L+H\delta+\tfrac12\nabla^3L[\delta,\delta]+O(\delta^3); with \delta=xu this gives xSu and \tfrac12x^2\nabla^3L[u,u]. By first-order eigenvalue perturbation (S=u^\top Hu; the variation of the unit eigenvector drops out), \partial_iS=u^\top(\partial_iH)u, so \nabla S=\nabla^3L[u,u] (Damian et al., Lemma 2; needs S simple), and

\nabla L(\bar\theta+xu)=\nabla L(\bar\theta)+x\,S\,u+\tfrac12x^2\,\nabla S(\bar\theta)+O(x^3).

The last term is even in x, so each step moves \bar\theta by -\tfrac\eta2x^2\nabla S.

The four-stage cycle (Damian et al. §4). With S\approx2/\eta+s: (1) progressive sharpening raises s by about \eta\alpha per step, where \alpha=-\langle\nabla L,\nabla S\rangle>0, while x is small; (2) blowup: once s>0, |x| grows geometrically with alternating sign; (3) self-stabilization: the \tfrac12x^2\nabla S term drives s below 0; (4) return to stability: |x| shrinks, and the cycle repeats.

Their theorem (one unstable eigenvalue, progressive sharpening, bounded higher derivatives): gradient descent tracks projected gradient descent onto \{S\le2/\eta,\ \nabla L\cdot u=0\}, with loss L(\theta^\dagger)+x^2/\eta for \theta^\dagger the projected iterate. It implicitly solves \min L subject to S\le2/\eta.

The toy, concretely. Damian et al. (Appendix B) use linear S and \ell plus a third coordinate. For the code we take S(y)=y and \ell(y)=c\,y^{-5}, c=62.5: the loss keeps falling, with diminishing returns, as the sharpness grows (Exercise 3.2′ says why not linear \ell). Near x=0 the Hessian is \operatorname{diag}(y,\ell''(y)), so y is the top eigenvalue while y>30c\,y^{-7}, i.e. y\gtrsim2.6; all runs start at y_0\ge3.

Exercise 3.2 — Self-stabilization in the toy. Use L(x,y)=\tfrac12S(y)x^2+\ell(y) with S'>0 and \ell'<0.

(a) Write the gradient-descent maps for (x,y). Show that when \eta S(y)>2, |x| grows with alternating sign, and that the y-update gains a term that lowers S and depends on x only through x^2.

(b) Find the period-2 orbits (x^*,y^*)\to(-x^*,y^*)\to(x^*,y^*). Show S(y^*)=2/\eta and

x^{*2}=-\frac{2\,\ell'(y^*)}{S'(y^*)} .

(c) Check that (a) is the general mechanism: in the toy at x=0, what are the top eigenvector u, the sharpness gradient \nabla S and the term \tfrac12x^2\nabla S from the expansion above?

Code: the toy. While \eta S<2, x decays geometrically and in exact arithmetic gets absurdly small (from x_0=10^{-3}, y_0=3: about 10^{-139} by the edge at \eta=0.3), so the first blowup overshoots enormously and can kick y through zero. In a real network the component along u is constantly re-seeded (other directions, rotation of u, noise, round-off), so we add Gaussian noise of standard deviation 10^{-3} to x each step (noise=0 turns it off). The lower panel shows |x_k| at \eta=0.4: at the noise floor while \eta S<2, a blowup once \eta S crosses 2, an overshoot of x^*, then oscillation around x^* in ever smaller cycles.

Code (37 lines)

C_TOY = 62.5                      # ell(y) = C_TOY * y**-5, so S(y) = y and ell'(y) = -5 C_TOY y**-6

def toy_loss(x, y, c=C_TOY):
    return 0.5 * y * x * x + c / y**5

def toy_run(etas, y0=3.0, x0=1e-3, noise=1e-3, seed=0, c=C_TOY):
    # Gradient descent on L = y x^2/2 + c y^-5 with step etas[k] at step k, plus noise on x.
    # Returns arrays (x_k, y_k) of the state *before* step k; NaN after a divergence.
    rng = np.random.default_rng(seed)
    xi = noise * rng.standard_normal(len(etas))
    xs, ys = np.full(len(etas), np.nan), np.full(len(etas), np.nan)
    x, y = x0, y0
    with np.errstate(over="ignore", invalid="ignore"):
        for k, eta in enumerate(etas):
            xs[k], ys[k] = x, y
            x, y = (1 - eta * y) * x + xi[k], y - eta * (0.5 * x * x - 5 * c / y**6)
            if not (np.isfinite(x) and np.isfinite(y)) or y <= 0 or abs(x) > 1e3:
                break
    return xs, ys

# ---- demo
K = 4000
fig, ax = plt.subplots(2, 1, figsize=(10, 5.4), sharex=True)
for eta, col in zip((0.3, 0.4, 0.5), ("C0", "C1", "C2")):
    xs, ys = toy_run(np.full(K, eta))
    t = np.arange(K) * eta
    ax[0].plot(t, ys, color=col, lw=1, label=f"$\\eta$ = {eta}")
    ax[0].axhline(2 / eta, color=col, ls="--", lw=0.8)
    if eta == 0.4:
        ax[1].semilogy(t, np.abs(xs), color=col, lw=0.6, label="$|x_k|$, $\\eta$ = 0.4")
        ax[1].axhline(math.sqrt(10 * C_TOY * (eta / 2) ** 6), color="k", ls=":", lw=1, label="$x^*$")
    sigma2 = 10 * C_TOY * (eta / 2) ** 6      # x*^2 = -2 ell'(2/eta) / S'
    print(f"eta = {eta}: mean S (last 1000 steps) = {ys[-1000:].mean():.4f}, 2/eta = {2/eta:.4f};  "
          f"mean x^2 = {np.mean(xs[-1000:]**2):.5f}, x*^2 = {sigma2:.5f}")
ax[0].set(ylabel="$S = y$", title="toy: sharpness rises, overshoots, then settles or hovers at $2/\\eta$ (dashed)")
ax[1].set(xlabel="$t = k\\eta$", ylim=(1e-5, 3), xlim=(0, 600)); ax[0].legend(); ax[1].legend()
plt.tight_layout(); plt.show()

eta = 0.3: mean S (last 1000 steps) = 6.6683, 2/eta = 6.6667;  mean x^2 = 0.00701, x*^2 = 0.00712
eta = 0.4: mean S (last 1000 steps) = 4.9996, 2/eta = 5.0000;  mean x^2 = 0.04005, x*^2 = 0.04000
eta = 0.5: mean S (last 1000 steps) = 3.9998, 2/eta = 4.0000;  mean x^2 = 0.15266, x*^2 = 0.15259

Exercise 3.2′ (optional) — When does the orbit attract? These parts explain what “hovers at 2/\eta” means. Take the toy with S(y)=y. Its period-2 orbit, from Exercise 3.2(b), is (\pm x^*,y^*) with y^*=2/\eta and x^{*2}=2|\ell'(y^*)|. Write P=|\ell'(y^*)|, C=\ell''(y^*), D=\eta C, Q=\eta^2x^{*2}=2\eta^2P.

(d) Linearize the two-step map at the period-2 orbit. Show that the orbit is linearly (exponentially) attracting iff |1-D+Q|<1, \;2-D+\tfrac Q2>0 and (D-\tfrac Q2)^2+\tfrac{Q^2}4-2D+2>0, which for small D,Q reduces to

\ell''(y^*)>2\eta\,|\ell'(y^*)| .

(e) In Damian et al.’s linear model (\ell''=0, so D=0), what are the eigenvalues of the two-step Jacobian? Show that the long-run averages are still pinned: if y_k stays bounded then \frac1K\sum_kx_k^2\to x^{*2} exactly, and if |x_k| stays bounded away from 0 and \infty then \frac1K\sum_k\log|1-\eta S(y_k)|\to0. What does the second statement say about the average of \eta S?

(f) Check the condition of (d) exactly for \ell=c\,y^{-5}.

At the edge the iterate bounces every step, but its time average moves smoothly. Can we write an ODE for that average, and predict the loss we measure? Yes: average the gradient over the oscillation, and choose the oscillation’s size so that the sharpness stays at 2/\eta.

Paper: Cohen, Damian, Talwalkar, Kolter & Lee, Understanding optimization in deep learning with central flows (see Readings). Start with their companion site, Part I in particular.

Let w(t) be the time-averaged iterate (with t=k\eta), S(w) its sharpness, \nabla S the sharpness gradient, and \sigma^2 the variance of the oscillation along the top eigenvector. While S<2/\eta the central flow is gradient flow. Once S reaches 2/\eta,

\frac{dw}{dt}=-\Big[\nabla L(w)+\tfrac12\sigma^2(t)\,\nabla S(w)\Big],\qquad \sigma^2=\max\!\Big(0,\ \frac{-2\langle\nabla L,\nabla S\rangle}{\|\nabla S\|^2}\Big).

On the edge, w follows gradient flow projected orthogonally to \nabla S, sliding along the surface S=2/\eta, and the predicted time-averaged loss is L(w)+\sigma^2/\eta. It is a heuristic, tested on real networks over long runs (full batch, no momentum; better for MSE than cross-entropy; degrading as \eta grows; possibly poor for ReLU; blind to large loss spikes).

Derivations and walkthrough. From §3.2: With u the unit top eigenvector at w, \nabla L(w+xu)=\nabla L(w)+xSu+\tfrac12x^2\nabla S(w)+O(x^3). Cohen et al. count time in steps, so they write dw/dt=-\eta[\nabla L+\tfrac12\sigma^2\nabla S].

Several unstable directions. Then \sigma^2 becomes the oscillation’s covariance \Sigma\succeq0, the drift is -\nabla L-\tfrac12\nabla\operatorname{tr}(H\Sigma) (gradient with \Sigma fixed), and the loss is L+\operatorname{tr}\Sigma/\eta. \Sigma lives only in directions whose eigenvalue sits at 2/\eta and keeps them from crossing it: a complementarity problem, the matrix version of “either \sigma^2=0 or the constraint is tight”.

Exercise 3.3 — Derive the one-direction central flow. Let w be the centre of the oscillation, and suppose the oscillation is \delta=xu with \mathbb E[x]=0 and \mathbb E[x^2]=\sigma^2.

(a) Show that \mathbb E\,\nabla L(w+\delta)=\nabla L(w)+\tfrac12\sigma^2\nabla S(w)+O(\mathbb E|x|^3). Gradient descent moves w by -\eta times this per step: write the flow in t=k\eta.

(b) Suppose S(w)=2/\eta. Find the \sigma^2 for which \frac{d}{dt}S(w(t))=0. Show \sigma^2\ge0 iff \langle\nabla L,\nabla S\rangle\le0, i.e. iff plain gradient flow would increase S. What should the flow do when this fails?

(c) Show that while the constraint is active the flow is \dot w=-\Pi^\perp\nabla L, with \Pi^\perp=I-\nabla S\nabla S^\top/\|\nabla S\|^2, and that L(w(t)) is non-increasing.

(d) Show that the time-averaged loss is \mathbb E\,L(w+\delta)\approx L(w)+\sigma^2/\eta.

(e) Evaluate \sigma^2 for the §3.2 toy L(x,y)=\tfrac12S(y)x^2+\ell(y) at (0,y^*), where S(y^*)=2/\eta, and compare with the orbit amplitude x^{*2}=-2\ell'(y^*)/S'(y^*) from Exercise 3.2(b).

Code: central flow against gradient descent. Does the central flow track gradient descent when the edge is curved? Add a coordinate z to the toy: L(x,y,z)=\tfrac12S(y,z)\,x^2+\ell(y,z) with

S(y,z)=y\,(1+\kappa z),\qquad \ell(y,z)=c\,y^{-5}+\tfrac m2z^2,\qquad \kappa=0.3,\ m=0.002 ,

and c=62.5. Now \nabla S=(1+\kappa z,\ \kappa y) rotates, and the flow slides along the curved surface y(1+\kappa z)=2/\eta: lowering z lets y grow at fixed sharpness, lowering c\,y^{-5}. In gradient descent the x^2 kick does the sliding (its z-component is -\tfrac\eta2\kappa y\,x^2). We integrate the central flow (Euler, dt=\eta/20, projected back onto the surface each step) and compare y, z, oscillation variance and loss, averaged over matching windows because the iterates circle the central point slowly.

Code (48 lines)

eta, kap, m, K = 0.4, 0.3, 0.002, 4000
S3  = lambda y, z: y * (1 + kap * z)
gS3 = lambda y, z: np.array([1 + kap * z, kap * y])            # grad S in (y, z)
l3  = lambda y, z: C_TOY / y**5 + 0.5 * m * z * z
gl3 = lambda y, z: np.array([-5 * C_TOY / y**6, m * z])         # grad ell in (y, z)

# gradient descent on (x, y, z), same noise floor on x as before
rng = np.random.default_rng(0)
x, y, z = 1e-3, 3.0, 2.0
G = np.zeros((K, 5))                                          # columns: x, y, z, loss, S
for k in range(K):
    G[k] = x, y, z, 0.5 * S3(y, z) * x * x + l3(y, z), S3(y, z)
    gy, gz = gl3(y, z) + 0.5 * x * x * gS3(y, z)
    x, y, z = (1 - eta * S3(y, z)) * x + 1e-3 * rng.standard_normal(), y - eta * gy, z - eta * gz

# central flow in t = k*eta
dt, w, CF = eta / 20, np.array([3.0, 2.0]), []
for i in range(K * 20):
    gl, gS = gl3(*w), gS3(*w)
    on_edge = S3(*w) >= (2 / eta) * (1 - 1e-9)                # tolerance: the projection lands exactly on the edge
    s2 = max(0.0, -2 * gl @ gS / (gS @ gS)) if on_edge else 0.0
    if i % 20 == 0:
        CF.append((w[0], w[1], s2, l3(*w), l3(*w) + s2 / eta))
    w = w - dt * (gl + 0.5 * s2 * gS)
    if s2 > 0:
        w[0] = (2 / eta) / (1 + kap * w[1])                   # back onto S = 2/eta
CF = np.array(CF)

idx, x2m = running_mean(G[:, 0] ** 2, 200)
_, lossm = running_mean(G[:, 3], 200)
t = np.arange(K) * eta
fig, ax = plt.subplots(1, 3, figsize=(13, 3.4))
ax[0].plot(t, G[:, 1], color="C0", lw=1.2, label="GD $y_k$"); ax[0].plot(t, CF[:, 0], "k--", lw=0.8, label="central flow $y$")
ax[0].plot(t, G[:, 2], color="C2", lw=0.8, label="GD $z_k$"); ax[0].plot(t, CF[:, 1], "k:", label="central flow $z$")
ax[0].set(xlabel="$t$", title="trajectory"); ax[0].legend(fontsize=7)
ax[1].plot(t[idx], x2m, color="C1", label="GD $x^2$, 200-step mean"); ax[1].plot(t, CF[:, 2], "k--", label="$\\sigma^2(t)$")
ax[1].set(xlabel="$t$", ylim=(0, 0.02), title="oscillation variance"); ax[1].legend(fontsize=7)
ax[2].plot(t[idx], lossm, color="C3", label="GD loss, 200-step mean"); ax[2].plot(t, CF[:, 4], "k--", label="$L(w)+\\sigma^2/\\eta$")
ax[2].plot(t, CF[:, 3], "k:", label="$L(w)$"); ax[2].set(xlabel="$t$", ylim=(0, 0.05), title="loss"); ax[2].legend(fontsize=7)
plt.tight_layout(); plt.show()
# Early: 200-step windows. Late, the iterates circle the central point in slow, wide loops,
# so compare means over 1000-step windows (CF values averaged over the same window).
for k, half in ((400, 100), (800, 100), (2000, 500), (3000, 500)):
    w_ = slice(k - half, k + half)
    gd = [G[w_, i].mean() for i in (1, 2)] + [(G[w_, 0] ** 2).mean(), G[w_, 3].mean()]
    cf = [CF[w_, i].mean() for i in (0, 1, 2, 4)]
    print(f"t = {k*eta:5.0f} ({2*half}-step window): y {gd[0]:.3f} vs CF {cf[0]:.3f}; z {gd[1]:.3f} vs {cf[1]:.3f}; "
          f"x^2 {gd[2]:.5f} vs sigma^2 {cf[2]:.5f}; loss {gd[3]:.5f} vs L+sigma^2/eta {cf[3]:.5f}")

t =   160 (200-step window): y 5.900 vs CF 5.849; z -0.512 vs -0.482; x^2 0.00405 vs sigma^2 0.00446; loss 0.01921 vs L+sigma^2/eta 0.02061
t =   320 (200-step window): y 6.544 vs CF 6.516; z -0.788 vs -0.775; x^2 0.00272 vs sigma^2 0.00280; loss 0.01265 vs L+sigma^2/eta 0.01293
t =   800 (1000-step window): y 7.389 vs CF 7.378; z -1.079 vs -1.074; x^2 0.00233 vs sigma^2 0.00227; loss 0.00983 vs L+sigma^2/eta 0.00970
t =  1200 (1000-step window): y 7.756 vs CF 7.749; z -1.186 vs -1.182; x^2 0.00221 vs sigma^2 0.00221; loss 0.00916 vs L+sigma^2/eta 0.00916

Late in the run, 1000-step averages of all four quantities agree with the central flow to a few percent; just after the first excursion, variance and loss are off by about 10%. The measured loss is more than twice L(w): much of what you see during EoS is the oscillation. The toy satisfies the flow’s assumptions by construction, so this checks the derivation, not the claim about networks.

The central flow needs the top eigenvector, the sharpness gradient and, with several unstable directions, a constrained matrix problem. Is there a continuous-time model of the edge needing only gradients and Hessian-vector products? Rod Flow is one. And principal flow is what the small-step modified-flow series of §3.1 becomes if every order is kept.

Papers: Regis & Chewi, Rod Flow (2026); Rosca, Wu, Qin & Dherin, the principal flow (2023) (see Readings).

Rod Flow tracks the rod joining two consecutive iterates (midpoint and half-length) instead of the iterate: at the edge the iterates jump back and forth, but the rod moves smoothly, and its length carries the instability. Principal flow applies the exact quadratic flow of Exercise 3.1, rate \frac1\eta\log(1-\eta\lambda), eigendirection by eigendirection; it captures the onset of instability but, having no third derivatives, not the stabilization.

Derivations and walkthrough. Notation: Gradient-descent iterates w_{k+1}=w_k-\eta\nabla L(w_k), continuous time t=k\eta, rod midpoint \bar w_k=\tfrac12(w_{k+1}+w_k), half-length \delta_k=\tfrac12(w_{k+1}-w_k)=-\tfrac\eta2\nabla L(w_k), and \Sigma_k=\delta_k\delta_k^\top (the rod’s squared length, as a matrix). The midpoint’s update is exact,

\bar w_{k+1}-\bar w_k=-\tfrac\eta2\big[\nabla L(\bar w_k+\delta_k)+\nabla L(\bar w_k-\delta_k)\big],

with gradients at the rod’s ends, where the iterates are, not Taylor-expanded around the centre. \Sigma_k also obeys an exact identity (Exercise 3.4(a)). As ODEs, with one backward-error correction for the midpoint, in t=k\eta and with L_\pm=L(\bar w\pm\delta):

\frac{d\bar w}{dt}=-\tfrac12\big[\nabla L_++\nabla L_-\big]-\tfrac\eta8\big[\nabla^2L_++\nabla^2L_-\big]\big[\nabla L_++\nabla L_-\big],\qquad \frac{d\Sigma}{dt}=\tfrac\eta4\big[\nabla L_+\nabla L_+^\top+\nabla L_-\nabla L_-^\top\big]-\tfrac2\eta\Sigma .

To close the system, \delta is \sqrt{\lambda_{\max}(\Sigma)} times \Sigma’s top eigenvector (the authors observe \Sigma stays nearly rank one). On a quadratic this reproduces the 2/\eta threshold (Exercise 3.4), on a quartic it gives a finite stable rod length, and expanding \nabla L_\pm to second order heuristically recovers the central flow.

Exercise 3.4 — Rod Flow on a quadratic.

(a) Verify the exact midpoint update above, and \delta_k=-\tfrac\eta2\nabla L(w_k). Then show that \Sigma_k=\delta_k\delta_k^\top satisfies, exactly,

\Sigma_{k+1}-\Sigma_k=\tfrac{\eta^2}4\big[\nabla L_+\nabla L_+^\top+\nabla L_-\nabla L_-^\top\big]-2\Sigma_k,\qquad L_\pm=L(\bar w_k\pm\delta_k),

and that dividing by the time step \eta gives the \Sigma equation above.

(b) Take L=\tfrac12Sw^2 in one dimension and \Sigma=\delta^2. Write the \Sigma equation at \bar w=0 and show \Sigma grows iff \eta S>2.

(c) Substitute \delta=0 into the \bar w equation and show it becomes the modified flow \dot{\bar w}=-\nabla\big(L+\tfrac\eta4\|\nabla L\|^2\big). Is \delta=0 a consistent state of the full system away from a stationary point?

Does Adam have an edge of stability? Yes, but in the preconditioned sharpness \lambda_{\max}(P^{-1}H), where H is the loss Hessian and P is Adam’s diagonal preconditioner. In full-batch training with \beta_1=0.9 it sits at 38/\eta while the raw sharpness keeps rising.

Paper: Cohen, Ghorbani, Krishnan, Agarwal et al., Adaptive gradient methods at the edge of stability (v2; see Readings).

With P frozen, Adam on a quadratic is heavy-ball momentum on a rescaled problem, stable iff \lambda_{\max}(P^{-1}H)<2(1+\beta_1)/\big((1-\beta_1)\eta\big)=38/\eta. The edge holds because the instability’s large gradients inflate P, pulling the preconditioned sharpness back down. With minibatches it sits below the threshold, more so for small batches (and at very small \eta, even in full batch).

Derivations and walkthrough. Prerequisites: Heavy-ball momentum is x_{k+1}=x_k-\eta\nabla L(x_k)+\beta(x_k-x_{k-1}). On \tfrac12\lambda x^2 its characteristic polynomial is z^2-(1+\beta-\eta\lambda)z+\beta, and for 0\le\beta<1 both roots lie inside the unit circle iff 0<\eta\lambda<2(1+\beta) (derivation and pictures: Goh, Why momentum really works). Adam with gradient g_k: first moment m_k=\beta_1m_{k-1}+(1-\beta_1)g_k (an exponential moving average, EMA), second moment v_k=\beta_2v_{k-1}+(1-\beta_2)g_k^2 (coordinatewise), step \theta_{k+1}=\theta_k-\eta P_k^{-1}m_k with P_k=\operatorname{diag}(\sqrt{v_k})+\epsilon I.

Exercise 3.5 — The 38/\eta threshold.

(a) Show that EMA momentum, m_k=\beta_1m_{k-1}+(1-\beta_1)\nabla L(\theta_k), \theta_{k+1}=\theta_k-\eta m_k, is heavy-ball momentum with step \eta(1-\beta_1) and momentum \beta_1. Deduce that on \tfrac12\lambda\theta^2 with \lambda>0 it is stable iff \lambda<2(1+\beta_1)/\big((1-\beta_1)\eta\big), and evaluate this at \beta_1=0.9.

(b) Now use a fixed positive diagonal preconditioner, \theta_{k+1}=\theta_k-\eta P^{-1}m_k, on \tfrac12\theta^\top H\theta with H positive definite. Show that stability is governed by the eigenvalues of P^{-1}H, which are real and positive even though P^{-1}H is not symmetric, and that the condition is \lambda_{\max}(P^{-1}H)<38/\eta.

(c) Explain why, if Adam’s raw sharpness keeps rising, its preconditioned sharpness can stay at 38/\eta instead of overshooting and diverging. Be careful about which direction is unstable.

# Check the 38/eta threshold for EMA momentum (beta1 = 0.9), plain and with a frozen diagonal preconditioner.
def ema_final(eta_H, P, beta=0.9, steps=4000, seed=0):
    th = np.random.default_rng(seed).standard_normal(len(P)); m = np.zeros_like(th)
    for _ in range(steps):
        m = beta * m + (1 - beta) * (eta_H @ th)        # eta folded into H: eta_H = eta * H
        th = th - m / P
        if np.linalg.norm(th) > 1e100:
            return np.inf                               # diverged
    return np.linalg.norm(th)

for r in (37.5, 38.5):
    print(f"1-D, eta*lambda = {r}: |theta| after 4000 steps = {ema_final(np.array([[r]]), np.ones(1)):.2e}")
A = np.random.default_rng(1).standard_normal((3, 3)); H = A @ A.T + 0.1 * np.eye(3)   # positive definite
P = np.array([0.5, 2.0, 5.0])
lam_pre = np.max(np.linalg.eigvals(np.diag(1 / P) @ H).real)
for r in (0.98, 1.02):
    eta = r * 38 / lam_pre
    print(f"3-D, eta * lambda_max(P^-1 H) = {r*38:.2f} (raw eta*S = {eta*np.linalg.eigvalsh(H).max():.1f}): "
          f"|theta| = {ema_final(eta * H, P):.2e}")

1-D, eta*lambda = 37.5: |theta| after 4000 steps = 1.37e-92
1-D, eta*lambda = 38.5: |theta| after 4000 steps = inf
3-D, eta * lambda_max(P^-1 H) = 37.24 (raw eta*S = 58.5): |theta| = 5.81e-92
3-D, eta * lambda_max(P^-1 H) = 38.76 (raw eta*S = 60.9): |theta| = inf

Code: Adam’s edge in the MLP. Full-batch Adam on the §3.1 network (\beta_1=0.9, \beta_2=0.999, \epsilon=10^{-7}, no bias correction, as in Cohen et al.) at two step sizes. Every 100 steps we freeze P and compute \lambda_{\max}(P^{-1/2}HP^{-1/2}) and the raw S by Lanczos on Hessian-vector products.

Code (50 lines)

import scipy.sparse.linalg as sla

def top_eig_lanczos(net, Pis=None):
    # Largest eigenvalue of P^{-1/2} H P^{-1/2} (or of H if Pis is None) by Lanczos on Hessian-vector products.
    # Power iteration converges too slowly here (small spectral gap), so we use scipy's eigsh.
    params = list(net.parameters()); shapes = [q.shape for q in params]; sizes = [q.numel() for q in params]
    grads = torch.autograd.grad(mse(net), params, create_graph=True)
    pis = None if Pis is None else torch.cat([q.flatten() for q in Pis]).double().numpy()
    def matvec(x):
        x = np.asarray(x).ravel() * (1 if pis is None else pis)
        vs = [torch.tensor(c, dtype=torch.float32).view(sh) for c, sh in zip(np.split(x, np.cumsum(sizes)[:-1]), shapes)]
        Hv = torch.autograd.grad(grads, params, grad_outputs=vs, retain_graph=True)
        y = torch.cat([h.flatten() for h in Hv]).double().numpy()
        return y * (1 if pis is None else pis)
    N = sum(sizes)
    op = sla.LinearOperator((N, N), matvec=matvec, dtype=np.float64)
    return sla.eigsh(op, k=1, which="LA", tol=1e-5, return_eigenvectors=False)[0]

def run_adam(eta, steps=1500, b1=0.9, b2=0.999, eps=1e-7, every=100):
    net = make_mlp(0); params = list(net.parameters())
    m = [torch.zeros_like(q) for q in params]; vv = [torch.zeros_like(q) for q in params]
    rec = []
    for k in range(steps + 1):
        g = torch.autograd.grad(mse(net), params)
        with torch.no_grad():
            for i in range(len(params)):
                m[i].mul_(b1).add_((1 - b1) * g[i]); vv[i].mul_(b2).add_((1 - b2) * g[i] ** 2)
        if k > 0 and k % every == 0:
            Pis = [1 / (vi.sqrt() + eps).sqrt() for vi in vv]       # frozen P = diag(sqrt v) + eps
            rec.append((k, top_eig_lanczos(net), top_eig_lanczos(net, Pis)))
        with torch.no_grad():
            for i, q in enumerate(params):
                q -= eta * m[i] / (vv[i].sqrt() + eps)
    return np.array(rec)

t0 = time.time()
fig, ax = plt.subplots(1, 2, figsize=(11, 3.4))
for eta, col in ((1e-3, "C0"), (3e-3, "C1")):
    r = run_adam(eta)
    ax[0].plot(r[:, 0], eta * r[:, 2], color=col, label=f"$\\eta$ = {eta:g}")
    ax[1].plot(r[:, 0], r[:, 1], color=col, label=f"$\\eta$ = {eta:g}")
    late = r[r[:, 0] > 500]
    print(f"eta = {eta:g}, steps 600-1500: eta * lambda_max(P^-1 H) mean {eta * late[:, 2].mean():.1f}, "
          f"range [{eta * late[:, 2].min():.1f}, {eta * late[:, 2].max():.1f}] (threshold 38);  "
          f"raw S = {late[:, 1].mean():.1f}, so eta * S = {eta * late[:, 1].mean():.3f}")
ax[0].axhline(38, color="k", ls="--", lw=1)
ax[0].set(xlabel="step", ylabel="$\\eta\\,\\lambda_{\\max}(P^{-1}H)$", title="preconditioned sharpness (dashed: 38)")
ax[1].set(xlabel="step", ylabel="raw sharpness $S$", title="raw sharpness: far below $2/\\eta$")
ax[0].legend(); plt.tight_layout(); plt.show()
print(f"({time.time() - t0:.1f} s)")

eta = 0.001, steps 600-1500: eta * lambda_max(P^-1 H) mean 37.7, range [36.4, 38.5] (threshold 38);  raw S = 25.5, so eta * S = 0.025
eta = 0.003, steps 600-1500: eta * lambda_max(P^-1 H) mean 37.4, range [36.0, 38.6] (threshold 38);  raw S = 17.4, so eta * S = 0.052

(7.6 s)

At both step sizes \eta\lambda_{\max}(P^{-1}H) hovers near 38 (about 36 to 38.6), while \eta S is a few hundredths: the edge is in the preconditioned quantity. A caveat on how it gets there: here the loss is near zero after a few hundred steps, so v (hence P) decays, raising \lambda_{\max}(P^{-1}H) toward the edge, and each renewed oscillation’s gradient burst pushes it back. That is not progressive sharpening; in Cohen et al.’s longer runs Adam is still training at the edge while the raw sharpness climbs.

Under minibatch SGD, does the sharpness still rise to 2/\eta? The full-batch \lambda_{\max} settles lower, by a gap that grows as the batch size b shrinks. What sits at 2/\eta, for every batch size, is the batch sharpness

\mathrm{BS}(\theta)=\mathbb E_B\!\left[\frac{g_B^\top H_B\,g_B}{\|g_B\|^2}\right],

where a minibatch B has gradient g_B and Hessian H_B: each minibatch’s curvature along its own gradient, with the expectation outside the ratio.

Paper: Andreyev & Beneventano, Edge of Stochastic Stability (v5; see Readings).

The subtlety is which average: requiring the expected batch-loss change to be negative gives a ratio of expectations instead, which also stabilizes, but above 2/\eta and at a batch-size-dependent level. Why batch sharpness settles at 2/\eta (an SGD analogue of §3.2) is open.

Derivations and walkthrough. Notation: A minibatch B of size b has loss L_B, gradient g_B=\nabla L_B and Hessian H_B; the SGD step is \theta'=\theta-\eta g_B; S_B=g_B^\top H_Bg_B/\|g_B\|^2, so \mathrm{BS}=\mathbb E[S_B]. The paper’s theory is one-sided; Exercise 3.6(c–d) examines what it shows.

Exercise 3.6 — Which average? Fix \theta and draw a minibatch B.

(a) Expand L_B(\theta') to second order. Show that the step decreases its own batch loss iff S_B<2/\eta.

(b) Show that \mathbb E_B\big[L_B(\theta')-L_B(\theta)\big]<0 (to the same order) iff \mathbb E[g_B^\top H_Bg_B]/\mathbb E\|g_B\|^2<2/\eta. Write this ratio of expectations as a weighted average of the S_B. When does it differ from \mathrm{BS}=\mathbb E[S_B]?

(c) Treat L_B as exactly quadratic near \theta. Show that the same batch’s gradient after its own step satisfies \|\nabla L_B(\theta')\|^2/\|g_B\|^2\ge(1-\eta S_B)^2, and deduce \mathbb E_B\big[\|\nabla L_B(\theta')\|^2/\|g_B\|^2\big]\ge(1-\eta\,\mathrm{BS})^2.

(d) Does (c) show that SGD is unstable whenever \mathrm{BS}>2/\eta? Test it on two equally likely quadratic batches H_1=\operatorname{diag}(2.3,1), H_2=\operatorname{diag}(1,2.3) (both with minimum at 0), \eta=1, \theta=(1,1): compute \mathrm{BS}(\theta) and \mathbb E_B\|\theta'\|^2. What, then, does (c) say, and does it explain why \mathrm{BS} rather than the ratio in (b) settles at 2/\eta?

Code: batch sharpness in the MLP. SGD at \eta=0.3 (2/\eta\approx6.67) on the §3.1 network, batch sizes b=16 and 64. Every 100 steps: full-batch \lambda_{\max}, \mathrm{BS}, and the ratio of expectations from Exercise 3.6(b), each expectation over 32 random batches. A small, noisy run.

Code (44 lines)

ALL = torch.arange(n_data)

def loss_on(net, idx):
    return 0.5 * ((net(X[idx]) - Y[idx]) ** 2).mean()

def batch_stats(net, b, nb, gen):
    # returns (mean of S_B, ratio of expectations) over nb random batches
    params = list(net.parameters()); rq, num, den = [], 0.0, 0.0
    for _ in range(nb):
        idx = torch.randperm(n_data, generator=gen)[:b]
        g = torch.autograd.grad(loss_on(net, idx), params, create_graph=True)
        gd = [t.detach() for t in g]
        Hg = torch.autograd.grad(g, params, grad_outputs=gd)
        q = sum((a * c).sum() for a, c in zip(Hg, gd)).item(); gg = sum((a * a).sum() for a in gd).item()
        rq.append(q / gg); num += q; den += gg
    return float(np.mean(rq)), num / den

def run_sgd(eta, b, steps=6000, every=100, nb=32):
    net = make_mlp(0); gen = torch.Generator().manual_seed(1); v, rec = None, []
    for k in range(steps + 1):
        if k % every == 0:
            lam, v = sharpness(net, v, iters=20 if v is None else 6)
            rec.append((k, lam, *batch_stats(net, b, nb, gen)))
        idx = torch.randperm(n_data, generator=gen)[:b]
        grads = torch.autograd.grad(loss_on(net, idx), list(net.parameters()))
        with torch.no_grad():
            for q, g in zip(net.parameters(), grads):
                q -= eta * g
    return np.array(rec)

t0 = time.time(); eta = 0.3
fig, ax = plt.subplots(1, 2, figsize=(11, 3.4), sharey=True)
for i, b in enumerate((16, 64)):
    r = run_sgd(eta, b)
    ax[i].plot(r[:, 0], r[:, 1], label="full-batch $\\lambda_{\\max}$")
    ax[i].plot(r[:, 0], r[:, 2], label="batch sharpness $\\mathbb{E}[S_B]$")
    ax[i].plot(r[:, 0], r[:, 3], lw=0.8, label="ratio of expectations")
    ax[i].axhline(2 / eta, color="k", ls="--", lw=1)
    ax[i].set(xlabel="step", title=f"SGD, b = {b} (dashed: $2/\\eta$)")
    late = r[r[:, 0] >= 3600]
    print(f"b = {b:2d}, steps 3600-6000: lambda_max {late[:, 1].mean():.2f}, batch sharpness {late[:, 2].mean():.2f} "
          f"(+-{late[:, 2].std():.2f}), ratio of expectations {late[:, 3].mean():.2f};  2/eta = {2/eta:.2f}")
ax[0].legend(fontsize=7); plt.tight_layout(); plt.show()
print(f"({time.time() - t0:.1f} s)")

b = 16, steps 3600-6000: lambda_max 4.03, batch sharpness 6.22 (+-0.51), ratio of expectations 7.14;  2/eta = 6.67
b = 64, steps 3600-6000: lambda_max 6.57, batch sharpness 6.88 (+-0.49), ratio of expectations 7.05;  2/eta = 6.67

(18.6 s)

The run shows the gap clearly: \lambda_{\max} is about 4.0 at b=16 and 6.6 at b=64. It cannot tell batch sharpness from the ratio of expectations (6.2 and 7.1 at b=16; 6.9 and 7.05 at b=64; fluctuations about \pm0.5); that needs longer, larger runs (the paper’s App. J).

What if the step is too large for the curvature \lambda_0 already at initialization, \eta\lambda_0>2? In a wide two-layer linear network the dynamics reduce to two exact scalar maps with three phases: lazy (\eta\lambda_0<2), catapult (2<\eta\lambda_0<4: the loss spikes, the curvature drops below 2/\eta, and training converges somewhere flatter) and divergent (\eta\lambda_0>4).

Paper: Lewkowycz, Bahri, Dyer, Sohl-Dickstein & Gur-Ari, The catapult mechanism (see Readings).

Like self-stabilization (§3.2), this is “oscillation squared lowers curvature”, but with no sharpening pressure pushing back, so the drop is one-shot. The coupling is weak, O(1/n) at width n, so the loss must grow to order n (about \log n steps) before the drop: the spike is part of the mechanism. Cohen et al. 2021 suggest EoS may be a series of micro-catapults; the first drop in the §3.1 plot resembles one. The edge 4/\lambda_0 is specific to this model: ReLU networks tolerate up to about 12/\lambda_0 (tanh, closer to 4/\lambda_0).

Derivations and walkthrough. The model: f=v^\top u/\sqrt n with weights u,v\in\mathbb R^n (here u is a weight vector, not the eigenvector of earlier sections; their f(x)=n^{-1/2}v^\top ux with a single input x=1), target 0, loss L=\tfrac12f^2, and the neural tangent kernel \lambda=\|\nabla f\|^2=(\|u\|^2+\|v\|^2)/n, which at a minimum (f=0) is the sharpness. With standard initialization u,v\sim\mathcal N(0,I), f_0=O(1) and \lambda_0\approx2. One gradient step with constant \eta gives, exactly,

f_{k+1}=f_k\Big(1-\eta\lambda_k+\frac{\eta^2f_k^2}n\Big),\qquad \lambda_{k+1}=\lambda_k+\frac{\eta f_k^2}n\big(\eta\lambda_k-4\big).

The \lambda-update is driven by f^2 as the toy’s y is driven by x^2 in §3.2. The phase boundaries are sharp as n\to\infty with \eta\lambda_0 held away from 2 and 4.

▶ Interactive figure: eos_toy_and_catapult.html (opens full page). Panel 2 iterates these maps, with sliders for \eta\lambda_0 and the width n.

Exercise 3.7 — The catapult maps.

(a) Derive the two maps from one gradient step on (u,v).

(b) Show that at f=0 the Hessian of L is \nabla f\,\nabla f^\top and its top eigenvalue is \lambda, so \eta\lambda<2 is the quadratic stability condition.

(c) Read the three phases off the maps, for large n and f_0=O(1). In the catapult phase, estimate how many steps the growth lasts, and say why a run that converges must end with \eta\lambda_\infty\le2. In the lazy phase, how far does \lambda move?

Code (48 lines)

def catapult(f0, a0, n, eta=1.0, K=200):
    # Iterate the exact maps; returns arrays f_k and a_k = eta*lambda_k (NaN after divergence).
    f, a = np.full(K + 1, np.nan), np.full(K + 1, np.nan)
    f[0], a[0] = f0, a0
    for k in range(K):
        b = eta**2 * f[k] ** 2 / n
        f[k + 1], a[k + 1] = f[k] * (1 - a[k] + b), a[k] + b * (a[k] - 4)
        if not np.isfinite(f[k + 1]) or abs(f[k + 1]) > 1e100:
            f[k + 1] = a[k + 1] = np.nan; break
    return f, a

# 1. the maps agree with explicit gradient descent on (u, v)
n = 1000; rng = np.random.default_rng(0)
u, v = rng.standard_normal(n), rng.standard_normal(n)
f0, lam0 = v @ u / np.sqrt(n), (u @ u + v @ v) / n
eta = 3.0 / lam0
fm, am = catapult(f0, eta * lam0, n, eta, K=40)
uu, vv, err = u.copy(), v.copy(), 0.0
for k in range(40):
    fk, ak = vv @ uu / np.sqrt(n), eta * (uu @ uu + vv @ vv) / n
    err = max(err, abs(fk - fm[k]) / (1 + abs(fm[k])), abs(ak - am[k]))
    uu, vv = uu - eta * fk * vv / np.sqrt(n), vv - eta * fk * uu / np.sqrt(n)
print(f"f0 = {f0:.3f}, lambda0 = {lam0:.3f}; max |maps - explicit GD| over 40 steps: {err:.1e}")

# 2. the three phases
fig, ax = plt.subplots(1, 2, figsize=(11, 3.6))
for a0, col in ((1.5, "C0"), (3.0, "C1"), (4.1, "C3")):
    f, a = catapult(1.0, a0, n)
    ax[0].semilogy(0.5 * f**2, color=col, label=f"$\\eta\\lambda_0$ = {a0}")
ax[0].axhline(0.5 * n, color="gray", ls=":", lw=1)
ax[0].set(xlabel="step $k$", ylabel="loss $f^2/2$", xlim=(0, 60), ylim=(1e-12, 1e12), title=f"loss (dotted: $n/2$, n = {n})")
ax[0].legend()
grid = np.linspace(0.2, 4.4, 400)
fin, conv = [], []
for a0 in grid:
    f, a = catapult(1.0, a0, n, K=400)
    ok = np.isfinite(f[-1]) and abs(f[-1]) < 1e-6            # converged within 400 steps
    fin.append(a[-1] if np.isfinite(f[-1]) else np.nan); conv.append(ok)
fin, conv = np.array(fin), np.array(conv)
ax[1].plot(grid[conv], fin[conv], ".", ms=2.5, color="C0", label="converged")
ax[1].plot(grid[~conv], fin[~conv], "x", ms=3, color="C3", label="not converged in 400 steps")
ax[1].plot(grid, grid, color="gray", lw=0.8)
for x_, c_ in ((2, "C2"), (4, "C3")):
    ax[1].axvline(x_, color=c_, ls="--", lw=1)
ax[1].axhline(2, color="C2", ls=":", lw=1)
ax[1].set(xlabel="$\\eta\\lambda_0$", ylabel="final $\\eta\\lambda$", ylim=(0, 4.5),
          title="final curvature (gray: no change; missing: diverged)")
ax[1].legend(fontsize=7); plt.tight_layout(); plt.show()

f0 = 1.870, lambda0 = 2.002; max |maps - explicit GD| over 40 steps: 1.3e-15

The right-hand plot (with a=\eta\lambda, b=\eta^2f^2/n): the catapult over-flattens. For \eta\lambda_0 between about 3 and 4 the final \eta\lambda is often far below 2, sometimes below 0.1, in lobes set by how the O(1)-per-step drop of a happens to end. Runs landing at very small a converge slowly (red crosses): \|u\|^2+\|v\|^2\ge2|u^\top v| gives a\ge2\sqrt b, so at a\approx0 the multiplier 1-a+b is close to 1. They are simply very flat, not near a phase boundary.

Why does ramping the learning rate up at the start (warmup) help, and why does the loss drop sharply in the final decay of a warmup–stable–decay (WSD) schedule? Both have readings in terms of oscillation across sharp directions.

Papers: Kalra & Barkeshli, Why warmup the learning rate?; Wen et al., Understanding warmup-stable-decay learning rates (see Readings).

Warmup: a large step from a sharp start throws the network into a catapult (§3.7) or divergence; ramping \eta up lets the sharpness follow 2/\eta down to where the target step is tolerable. WSD: the loss is a river valley, a flat river between steep hillsides. At a large step the iterate bounces between the hillsides while progressing along the river; the bouncing hides the progress, and the decay removes it.

Derivations and walkthrough. Notation: \eta_k is the step at step k, \eta_{\rm trgt} the target step after warmup, S the sharpness (top Hessian eigenvalue), and time t=\sum_{j<k}\eta_j.

Warmup. Kalra and Barkeshli use the linear version of the §3.2 model: x_{k+1}=(1-\eta_k\lambda_k)x_k, \lambda_{k+1}=\lambda_k+\eta_k(\alpha-\beta x_k^2), with x the displacement along the top eigenvector, \lambda the sharpness, \alpha>0 a sharpening rate and \beta>0 the strength of the x^2 kick. Exercise 3.8(a) runs the experiment in our toy. (For Adam an earlier, compatible explanation is that its second-moment average starts from few gradients, so the adaptive step has high variance until warmup lets it settle: Liu et al.’s RAdam paper (ICLR 2020).)

WSD. In Wen et al. the bouncing comes from SGD noise (Exercise 3.8(b)); with full-batch small-step gradient descent their iterates track the river. A deterministic analogue, from §3.3 (a connection neither paper makes): at the edge the measured loss is L(w)+\sigma^2/\eta (w the averaged iterate, \sigma^2 the oscillation variance), and decaying \eta raises the edge 2/\eta, so the oscillation and its excess loss die out until sharpening reaches the new edge. For the SGD side, see NB2 §2.

Exercise 3.8 — Warmup in the toy, and a river valley.

(a) Code. Start the §3.2 toy (L=\tfrac12yx^2+c\,y^{-5}, c=62.5, sharpness S=y) sharp, at y_0=8, and train at \eta_{\rm trgt}=0.4, so 2/\eta_{\rm trgt}=5. Compare jumping straight to \eta_{\rm trgt} with linear warmups \eta_k=\eta_{\rm trgt}\min\big(1,(k+1)/T\big) for T\in\{100,400,1600\}, using toy_run with 4000 steps. For each, plot S_k=y_k and 2/\eta_k against t, and report whether it diverged and its peak loss. Then repeat from y_0=5.5: when does warmup help, and can it hurt? Panel 1 of the interactive figure in §3.2 has the same experiment with sliders.

(b) Derivation. Take L(h,r)=\tfrac12\gamma h^2-r: a hill coordinate h with curvature \gamma and a river coordinate r along which the loss falls at unit slope. Run SGD with noise only in the hill direction: h_{k+1}=h_k-\eta\gamma h_k+\eta\xi_k, \xi_k\sim\mathcal N(0,s^2) i.i.d. (so s^2 is the noise variance), and r_{k+1}=r_k+\eta. Find the stationary mean of the hill loss \tfrac12\gamma h^2 at constant \eta, for 0<\eta\gamma<2 (exactly, and for \eta\gamma\ll1). Then explain what a WSD schedule (long constant phase at \eta_{\max}, short decay to \eta_{\min}) achieves compared with (i) \eta_{\max} throughout, (ii) \eta_{\min} throughout.

# Exercise 3.8(a): warmup in the toy. Use toy_run and toy_loss from §3.2.
eta_trgt, K, y0 = 0.4, 4000, 8.0
schedules = {"no warmup": np.full(K, eta_trgt)}
# add linear warmups for T in (100, 400, 1600), run each, and plot y_k and 2/eta_k against t = cumsum(eta)

Solution (code)

eta_trgt, K = 0.4, 4000
for y0 in (8.0, 5.5):
    schedules = {"no warmup": np.full(K, eta_trgt)}
    for T in (100, 400, 1600):
        schedules[f"warmup T={T}"] = eta_trgt * np.minimum(1, (np.arange(K) + 1) / T)
    fig, ax = plt.subplots(1, len(schedules), figsize=(13, 2.8), sharey=True)
    for i, (name, etas) in enumerate(schedules.items()):
        xs, ys = toy_run(etas, y0=y0)
        t = np.concatenate([[0], np.cumsum(etas)[:-1]])
        diverged = np.isnan(ys).any()
        peak = np.nanmax(toy_loss(xs, ys))
        ax[i].plot(t, ys, lw=0.8, label="$S_k$")
        ax[i].plot(t, 2 / etas, "--", lw=1, label="$2/\\eta_k$")
        ax[i].set(title=f"{name}: " + ("DIVERGED" if diverged else f"peak loss {peak:.2g}"),
                  xlabel="$t$", ylim=(0, 10))
        kc = np.argmax(etas * ys > 2)                                  # first crossing of the edge
        over = np.nanmax((etas * ys)[kc:kc + 200])
        print(f"y0 = {y0}, {name:16s}: " + ("diverged" if diverged else
              f"peak loss = {peak:.3g}, S at crossing = {ys[kc]:.2f}, max eta*S after crossing = {over:.2f}, "
              f"final S = {ys[-1]:.3f} (2/eta = {2/eta_trgt})"))
    ax[0].set_ylabel(f"$y_0$ = {y0}"); ax[0].legend(fontsize=7)
    plt.tight_layout(); plt.show()

y0 = 8.0, no warmup       : diverged
y0 = 8.0, warmup T=100    : diverged
y0 = 8.0, warmup T=400    : peak loss = 15.3, S at crossing = 8.04, max eta*S after crossing = 2.33, final S = 4.983 (2/eta = 5.0)
y0 = 8.0, warmup T=1600   : peak loss = 2.61, S at crossing = 8.14, max eta*S after crossing = 2.12, final S = 5.016 (2/eta = 5.0)

y0 = 5.5, no warmup       : peak loss = 2.66, S at crossing = 5.50, max eta*S after crossing = 2.24, final S = 4.981 (2/eta = 5.0)
y0 = 5.5, warmup T=100    : peak loss = 3.78, S at crossing = 5.66, max eta*S after crossing = 2.29, final S = 4.981 (2/eta = 5.0)
y0 = 5.5, warmup T=400    : peak loss = 3.54, S at crossing = 5.98, max eta*S after crossing = 2.22, final S = 4.981 (2/eta = 5.0)
y0 = 5.5, warmup T=1600   : peak loss = 1.57, S at crossing = 6.61, max eta*S after crossing = 2.12, final S = 5.019 (2/eta = 5.0)