NB2 §6 — Curvature: what the Hessian looks like and when the basin is chosen

The Gauss–Newton decomposition (§6.1). Also the Gauss–Newton side of §6.2, which is linear algebra rather than a paper’s result.

Bulk and outliers (§6.3). The short paper, then the long one.

The tiny subspace (§6.4). The observation, then the intervention.

When the basin is chosen (§6.5). If you read one paper of the notebook, read Fort et al.

Also: Ghorbani, Krishnan and Xiao 2019 (full Hessian densities of ImageNet-scale networks by stochastic Lanczos; without BatchNorm, outliers appear early and hold the gradient); Atanasov, Bordelon and Pehlevan 2022 (silent alignment, the matrix version of 6.5’s timing); Sagun, Bottou and LeCun 2016; Papyan 2018 and 2019.

Second-order optimizers, influence functions and Laplace approximations use the Gauss–Newton matrix in place of the Hessian: the part built from first derivatives of the network only, positive semidefinite by construction. The part it drops vanishes at interpolation for square loss, but for cross-entropy it shrinks only like one over the logit scale.

Paper: Martens, New insights and perspectives on the natural gradient method (see Readings).

Take n examples with logits z_i=f(\theta;x_i)\in\mathbb R^C, per-example loss \ell_i, L=\frac1n\sum_i\ell_i, Jacobians J_i=\partial z_i/\partial\theta\in\mathbb R^{C\times P} and logit-space loss Hessians A_i=\nabla_z^2\ell_i. Differentiating twice through the logits,

H=\nabla^2_\theta L=\underbrace{\frac1n\sum_iJ_i^\top A_i\,J_i}_{H_{\rm GN}} \;+\;\underbrace{\frac1n\sum_i\sum_{c=1}^C\frac{\partial\ell_i}{\partial z_{ic}}\,\nabla^2_\theta z_{ic}}_{H_{\rm res}} .

H_{\rm GN}\succeq0 whenever \ell is convex in the logits (square loss, softmax cross-entropy), so every negative eigenvalue of H comes from the residual H_{\rm res}. For classifiers H_{\rm GN} is also the Fisher information of the model’s predictive distribution (Martens §§8–9), which is why “Fisher”, “Gauss–Newton” and “GGN” are used interchangeably; the empirical Fisher, built at the observed labels, is a different matrix (§11). In practice the damped H_{\rm GN} is inverted approximately with K-FAC (Martens and Grosse 2015); the Data Attribution day picks up from there.

Derivations and walkthrough. Notation: as above; \theta\in\mathbb R^P; J\in\mathbb R^{nC\times P} stacks the J_i.

Square loss. For \ell=\frac12\|z-y\|^2 with one-hot y, A_i=I_C and \partial\ell/\partial z=z_i-y_i. At an interpolating minimum every residual vanishes, so H=H_{\rm GN}=\frac1nJ^\top J exactly.

Cross-entropy. With p=\mathrm{softmax}(z) and \ell=-\log p_y, \nabla_z\ell=p-e_y (e_y the y-th unit vector) and A=\mathrm{diag}(p)-pp^\top. The residual p-e_y never vanishes at finite logits. Near interpolation, with 1-p_y=\varepsilon, both A and p-e_y are of order \varepsilon, so H_{\rm GN} and H_{\rm res} shrink together; their ratio is set by the size of the logits, and improves only as the margins (correct-class logit minus the largest other logit) grow, slowly (Exercise 6.1.B).

Rank. A\mathbf 1=0, because adding a constant to every logit leaves p unchanged; when all p_c>0 the rank of A is exactly C-1. So \mathrm{rank}\,H_{\rm GN}\le n(C-1) (\le nC for square loss), and when P>n(C-1) at least P-n(C-1) eigenvalues of H_{\rm GN} are exactly zero. Part of the near-zero “bulk” of measured spectra (6.3) is this rank deficiency, which H_{\rm res} spreads into small eigenvalues of both signs.

Exercise 6.1.A — The decomposition. Notation as above.

1. Derive the decomposition with the chain rule. Show that H_{\rm GN}\succeq0 when every A_i\succeq0.
2. For square loss, show that H_{\rm res}=0 at any interpolating point, where H=\frac1nJ^\top J.
3. Take the scalar two-layer linear network f=w_1w_2 (one input x=1, target s>0, L=\frac12(w_1w_2-s)^2).
   • Compute H_{\rm GN} and H_{\rm res} at a general point, and show that H_{\rm res} has eigenvalues \pm(w_1w_2-s).
   • On the balanced line w_1=w_2=a, find where H is indefinite.
4. Show that the eigenvalues of H and H_{\rm GN} differ by at most \|H_{\rm res}\|_2, and that \lambda_{\min}(H)\ge\lambda_{\min}(H_{\rm res}).

Exercise 6.1.B — Softmax cross-entropy. Let p=\mathrm{softmax}(z)\in\mathbb R^C, \ell=-\log p_y and A=\nabla_z^2\ell; H_{\rm GN} and H_{\rm res} as in 6.1.

1. Show that \nabla_z\ell=p-e_y and A=\mathrm{diag}(p)-pp^\top.
2. Show that A=\sum_cp_c(e_c-p)(e_c-p)^\top, and deduce: A\succeq0; A\mathbf 1=0; the rank is exactly C-1 when every p_c>0; \mathrm{tr}\,A=1-\|p\|^2 and \lambda_{\max}(A)\le\max_cp_c.
3. Fisher. Show that A=\mathbb E_{\hat y\sim p}\big[\nabla_z\ell(z,\hat y)\nabla_z\ell(z,\hat y)^\top\big], so H_{\rm GN} is the Fisher information of the model’s predictive distribution. How does the empirical Fisher F_{\rm emp}=\frac1n\sum_iJ_i^\top(p_i-e_{y_i})(p_i-e_{y_i})^\top J_i scale near interpolation, compared with H_{\rm GN}?
4. Near interpolation. Let p_y=1-\varepsilon and p_c=\varepsilon q_c for c\ne y, with \sum_{c\ne y}q_c=1. Show that, exactly, p-e_y=\varepsilon\sum_{c\ne y}q_c(e_c-e_y),\qquad\text{and}\qquad A=\varepsilon\sum_{c\ne y}q_c(e_c-e_y)(e_c-e_y)^\top+O(\varepsilon^2). Write the resulting per-example contributions to H_{\rm GN} and H_{\rm res} in terms of \nabla z_c-\nabla z_y and \nabla^2z_c-\nabla^2z_y.
5. Logit scale. Suppose the network is homogeneous of degree k, f(\lambda\theta)=\lambda^kf(\theta) (e.g. a bias-free ReLU net with k layers). Along a ray \theta=\lambda\hat\theta on which every training point is classified correctly, show that \|H_{\rm res}\|/\|H_{\rm GN}\| falls like \lambda^{-k}, one over the logit scale, even though both terms vanish like \varepsilon (unless \nabla^2z=0, where the residual vanishes identically).

Code 6.1.C — The decomposition, measured. The first cell checks H=H_{\rm GN}+H_{\rm res} to machine precision and the rank of the softmax Hessians (residual is the Hessian of \theta\mapsto\sum_i\langle r_i,z_i(\theta)\rangle with the residuals r_i=\frac1n\nabla_z\ell_i held fixed). The second trains the network with cross-entropy and with square loss (128 points, 4 classes, full-batch GD). Look for:

• square loss: the residual’s share of the Hessian goes to zero; cross-entropy: it tracks one over the median margin, not the loss (from step 30 on, share × margin drifts only from about 1.6 to 1.85 while the loss falls 300-fold; the GELU net is not homogeneous, so this is beyond Exercise 6.1.B);
• at the end, the top four eigenvalues of H and H_{\rm GN} agree to 0.1–7%; further down, the residual adds to the bulk;
• the rest of the spectrum has both signs (466 positive, 193 negative) plus 17 exact zeros (adding the same vector to every last-layer weight row, and the same constant to every output bias, leaves the loss unchanged); the top four are only weakly separated (\lambda_4/\lambda_5\approx1.6).

t0 = time.time()
X, y = make_data(C=4, n_per=32)
net = MLP(d=20, h=16, C=4)
L = lambda t: ce(net(t, X), y)
theta = net.theta0.clone()
for _ in range(30):                                            # move to a generic, partly trained point
    theta = theta - 0.1 * grad(L)(theta)

for loss in (ce, mse):
    H = hessian(lambda t: loss(net(t, X), y))(theta)
    HGN, Hres = ggn(net, theta, X, loss), residual(net, theta, X, y, loss)
    err = (H - HGN - Hres).abs().max() / H.abs().max()
    print(f"{loss.__name__}: P = {net.P}, max|H - H_GN - H_res| / max|H| = {err:.1e}, "
          f"min eig H_GN = {eigs(HGN)[-1]:+.1e}, min eig H_res = {eigs(Hres)[-1]:+.2e}")

A = logit_hessian(net(theta, X), ce)
a = torch.linalg.eigvalsh(A)                                  # (n, C), ascending
print(f"softmax Hessians A_i: max |A_i 1| = {(A @ torch.ones(4)).abs().max():.1e}; "
      f"smallest eigenvalue / largest <= {(a[:, 0].abs() / a[:, -1]).max():.1e}; "
      f"second smallest / largest >= {(a[:, 1] / a[:, -1]).min():.1e}  (so rank C - 1 = 3)")
print(f"({time.time() - t0:.1f}s)")

ce: P = 676, max|H - H_GN - H_res| / max|H| = 3.9e-16, min eig H_GN = -2.3e-16, min eig H_res = -1.82e-01
mse: P = 676, max|H - H_GN - H_res| / max|H| = 4.4e-16, min eig H_GN = -2.5e-15, min eig H_res = -2.23e+01
softmax Hessians A_i: max |A_i 1| = 2.3e-16; smallest eigenvalue / largest <= 5.8e-15; second smallest / largest >= 2.1e-03  (so rank C - 1 = 3)
(1.8s)

Code (45 lines)

t0 = time.time()
ckpts = [0, 10, 30, 100, 300, 1000, 2000]
rec = {}
for loss in (ce, mse):
    L = lambda t: loss(net(t, X), y)
    theta, rows = net.theta0.clone(), []
    for s in range(ckpts[-1] + 1):
        if s in ckpts:
            HGN, Hres = ggn(net, theta, X, loss), residual(net, theta, X, y, loss)
            H = HGN + Hres
            z = net(theta, X)
            margin = (z.gather(1, y[:, None]).squeeze() - z.scatter(1, y[:, None], -1e9).max(1).values).median()
            rows.append((s, L(theta).item(), (torch.linalg.norm(Hres) / torch.linalg.norm(H)).item(), margin.item()))
            if s == ckpts[-1]:
                rec[loss.__name__ + "_spec"] = (eigs(H), eigs(HGN))
        theta = theta - 0.1 * grad(L)(theta)
    rec[loss.__name__] = np.array(rows)

for name in ("ce", "mse"):
    print(name)
    for s, l, r, m in rec[name]:
        extra = f"   median margin {m:5.2f}   ratio x margin {r * m:4.2f}" if name == "ce" else ""
        print(f"  step {int(s):5d}: loss {l:.2e}   ||H_res||/||H|| = {r:.3f}" + extra)
eH, eG = rec["ce_spec"]
print("CE, end of training, top 4 eigenvalues of H:", np.round(eH[:4], 5), " of H_GN:", np.round(eG[:4], 5),
      f" (relative differences {np.round(np.abs(eH[:4] - eG[:4]) / eH[:4], 3)})")
tol = 1e-10 * eH[0]                                              # numerical zeros: exact symmetries of CE
print(f"CE, end of training: {(eH > tol).sum()} positive, {(eH < -tol).sum()} negative, "
      f"{(np.abs(eH) <= tol).sum()} numerically zero eigenvalues of H (|lambda| <= {tol:.0e})")

fig, ax = plt.subplots(1, 3, figsize=(14, 3.8))
for name, mk in (("ce", "o-"), ("mse", "s-")):
    ax[0].loglog(rec[name][:, 1], rec[name][:, 2], mk, label=name)
ax[0].set(xlabel="training loss", ylabel=r"$\|H_{\rm res}\|_F/\|H\|_F$", title="residual term along training")
ax[0].invert_xaxis(); ax[0].legend()
k = np.arange(1, 41)
ax[1].semilogy(k, eH[:40], "o", label="$H$"); ax[1].semilogy(k, eG[:40], "x", label=r"$H_{\rm GN}$")
ax[1].set(xlabel="index", title="CE, end of training: top 40 eigenvalues"); ax[1].legend()
bins = np.linspace(-9, -1, 41)
ax[2].hist(np.log10(eH[eH > tol]), bins=bins, alpha=0.7, label=f"positive ({(eH > tol).sum()})")
ax[2].hist(np.log10(-eH[eH < -tol]), bins=bins, alpha=0.7, label=f"negative ({(eH < -tol).sum()})")
ax[2].set_yscale("log"); ax[2].legend()
ax[2].set(xlabel=r"$\log_{10}|\lambda|$", ylabel="count", title=f"CE, end: nonzero eigenvalues of $H$ ($P={net.P}$)")
fig.tight_layout(); plt.show()
print(f"({time.time() - t0:.1f}s)")

ce
  step     0: loss 2.10e+00   ||H_res||/||H|| = 0.903   median margin -1.42   ratio x margin -1.28
  step    10: loss 4.18e-01   ||H_res||/||H|| = 0.733   median margin  1.29   ratio x margin 0.94
  step    30: loss 7.29e-02   ||H_res||/||H|| = 0.475   median margin  3.37   ratio x margin 1.60
  step   100: loss 9.90e-03   ||H_res||/||H|| = 0.308   median margin  5.41   ratio x margin 1.66
  step   300: loss 2.24e-03   ||H_res||/||H|| = 0.250   median margin  6.97   ratio x margin 1.74
  step  1000: loss 5.10e-04   ||H_res||/||H|| = 0.212   median margin  8.54   ratio x margin 1.81
  step  2000: loss 2.26e-04   ||H_res||/||H|| = 0.196   median margin  9.36   ratio x margin 1.84
mse
  step     0: loss 2.65e+00   ||H_res||/||H|| = 0.629
  step    10: loss 1.44e-01   ||H_res||/||H|| = 0.113
  step    30: loss 6.01e-02   ||H_res||/||H|| = 0.042
  step   100: loss 2.52e-02   ||H_res||/||H|| = 0.018
  step   300: loss 1.08e-02   ||H_res||/||H|| = 0.008
  step  1000: loss 4.93e-03   ||H_res||/||H|| = 0.004
  step  2000: loss 3.21e-03   ||H_res||/||H|| = 0.003
CE, end of training, top 4 eigenvalues of H: [0.01129 0.00829 0.00713 0.00548]  of H_GN: [0.0113  0.00843 0.0074  0.00588]  (relative differences [0.001 0.017 0.038 0.073])
CE, end of training: 466 positive, 193 negative, 17 numerically zero eigenvalues of H (|lambda| <= 1e-12)

(11.1s)

Here n(C-1)=384<P=676, so H_{\rm GN} has at least 292 exactly zero eigenvalues. The bulk of H, of both signs, is mostly H_{\rm res} acting on that null space and on the small eigenvalues of H_{\rm GN}. The outliers are H_{\rm GN}’s.

Curvature lives in parameter space; the tangent kernel, which says how the outputs move under a gradient step, lives in function space. Both are built from the same Jacobian, so the Gauss–Newton matrix and the kernel weighted by the loss Hessian have the same nonzero eigenvalues. Their eigenvectors are tied only through the Jacobian.

This is linear algebra rather than a paper’s result; Martens (in the Readings, under §6.1) covers the Gauss–Newton side.

Stack the Jacobians J_i=\partial z_i/\partial\theta (n examples, C logits, P parameters) into J\in\mathbb R^{nC\times P} and the logit-space loss Hessians A_i=\nabla_z^2\ell_i into the block-diagonal A. The empirical tangent kernel is \Theta=JJ^\top: under gradient flow on L=\frac1n\sum_j\ell_j the logits move as \dot z_i=-\frac1n\sum_jJ_iJ_j^\top\nabla_z\ell_j. Since H_{\rm GN}=\frac1nJ^\top AJ=M^\top M and \frac1nA^{1/2}\Theta A^{1/2}=MM^\top with M=n^{-1/2}A^{1/2}J, with multiplicity

\mathrm{spec}_{\neq0}\,H_{\rm GN}=\mathrm{spec}_{\neq0}\Big(\tfrac1n\,A^{1/2}\,\Theta\,A^{1/2}\Big).

For square loss (A=I) at an interpolating minimum, the sharpness S=\lambda_{\max}(H) is the top kernel eigenvalue \lambda_{\max}(\Theta)/n: progressive sharpening is the top of the kernel growing, and the edge of stability holds it near 2n/\eta (\eta the step size). For cross-entropy an example fitted with confidence 1-\varepsilon_i has A_i=O(\varepsilon_i) and drops out: the top of the Hessian is the kernel restricted to the examples the model is still unsure about.

Derivations and walkthrough. Notation: as above; H is the loss Hessian, H_{\rm res}=H-H_{\rm GN} the residual of 6.1.

Directions. If v is a unit eigenvector of the weighted kernel with eigenvalue \lambda, then u=n^{-1/2}J^\top A^{1/2}v/\sqrt\lambda is one of H_{\rm GN}, and A^{1/2}Ju=\sqrt{n\lambda}\,v: u moves the loss-weighted logits along v. Neither side’s rotation forces the other’s: with square loss, J\mapsto JQ (Q orthogonal on parameter space) rotates H_{\rm GN}’s eigenvectors and leaves \Theta fixed, and J\mapsto RJ (R orthogonal on the output side) rotates \Theta’s and leaves H_{\rm GN} fixed. So a kernel distance that mixes spectrum and rotation, as in 6.5, speaks about curvature only through its spectral part, and only after the A-reweighting (Code 6.2.B shows the unweighted kernel growing while the weighted one shrinks).

Exercise 6.2.A — The identity and its consequences. Notation as above. Items 2–3 show why, with square loss near interpolation, sustained progressive sharpening means the kernel is not frozen.

1. Prove the spectral identity and the eigenvector map. Show that \mathrm{rank}\,H_{\rm GN}=\mathrm{rank}(A^{1/2}J)\le\min(P,\,n(C-1)) for cross-entropy.
2. For square loss with L=\frac1{2n}\sum_i\|z_i-y_i\|^2, linearize one GD step in function space and derive r\mapsto(I-\eta\Theta/n)\,r, with r=z-y stacked. Conclude that near an interpolating minimum, GD is stable iff \eta S<2 with S=\lambda_{\max}(\Theta)/n.
3. For a model exactly linearized around a point (so J is constant), explain in two sentences why H does not change along the trajectory with square loss, while with cross-entropy it does.
4. (Optional.) A two-layer linear network. Take L=\frac12\|\Sigma_{yx}-W_2W_1\|_F^2, where W_2W_1 maps \mathbb R^d\to\mathbb R^m and the target \Sigma_{yx}=USV^\top has rank r, singular values s_1>\dots>s_r>0 and singular vectors u_\alpha (columns of U) and v_\beta (columns of V). Work at a balanced global minimum (W_2W_2^\top=(WW^\top)^{1/2}, W_1^\top W_1=(W^\top W)^{1/2}, W=\Sigma_{yx}).
   • Show that the nonzero eigenvalues of H are those of the kernel operator K[E]=W_2W_2^\top E+EW_1^\top W_1 on m\times d matrices.
   • Show that these eigenvalues are s_\alpha+s_\beta for \alpha,\beta\le r, together with s_\alpha (each with multiplicity d-r) and s_\beta (each with multiplicity m-r).
   • Conclude that S=2s_1: the sharpness of the minimum is set by the strongest mode of the target.

Exercise 6.2.B (code) — The identity, and what the weighting does. Train the network of Code 6.1.C with cross-entropy. At steps 10 and 2000, check the identity numerically, then compare the unweighted kernel \Theta/n with the weighted \frac1nA^{1/2}\Theta A^{1/2}. For the top C eigenvectors of each (vectors over (example, logit) pairs), measure the fraction explained by patterns that depend only on (class of the example, logit): the function-space side of the outliers of 6.3. Look for:

• the two spectra agree to machine precision, and \mathrm{rank}\,H_{\rm GN}=n(C-1) exactly;
• the top kernel eigenvectors are almost entirely class patterns (about 0.97, against a random baseline of 0.03), and the unweighted kernel grows during training;
• the weighted kernel, which is the curvature, shrinks about 200-fold as examples become confidently fitted, and its top eigenvectors become less class-like: the weighting concentrates the curvature on the examples still uncertain.

n, C = X.shape[0], 4


def class_score(v):
    # Fraction of ||v||^2 (v in R^{nC}, indexed by (example i, logit k)) explained by patterns that depend
    # only on (class of i, k): the "class-indicator-like" part.  A random vector scores about C^2/(nC).
    V = v.reshape(n, C)
    means = torch.stack([V[y == c].mean(0) for c in range(C)])        # (C classes, C logits)
    return (means[y] ** 2).sum().item() / (V ** 2).sum().item()


L = lambda t: ce(net(t, X), y)
# Your code: train with GD (lr 0.1) and, at steps 10 and 2000, form the stacked Jacobian J (nC x P), the
# block-diagonal A (use logit_hessian), a symmetric square root of A, and Theta = J J^T.
# 1. Check that J^T A J / n and A^1/2 Theta A^1/2 / n have the same top eigenvalues, and the rank of H_GN.
# 2. For the unweighted Theta / n and the weighted A^1/2 Theta A^1/2 / n, print the top-C eigenvalues and the
#    class_score of the top-C eigenvectors.

n, C = X.shape[0], 4


def class_score(v):
    # Fraction of ||v||^2 (v in R^{nC}, indexed by (example i, logit k)) explained by patterns that depend
    # only on (class of i, k): the "class-indicator-like" part.  A random vector scores about C^2/(nC).
    V = v.reshape(n, C)
    means = torch.stack([V[y == c].mean(0) for c in range(C)])        # (C classes, C logits)
    return (means[y] ** 2).sum().item() / (V ** 2).sum().item()


L = lambda t: ce(net(t, X), y)
theta = net.theta0.clone()
for step in range(2001):
    if step in (10, 2000):
        J = jacrev(net)(theta, X).reshape(n * C, -1)                 # stacked Jacobian (nC, P)
        A = torch.block_diag(*logit_hessian(net(theta, X), ce))      # (nC, nC)
        w, V = torch.linalg.eigh(A)
        A_half = V @ torch.diag(w.clamp(min=0).sqrt()) @ V.T
        Theta = J @ J.T                                              # empirical NTK Gram matrix
        K_w = A_half @ Theta @ A_half / n                            # loss-weighted kernel
        e_gn, e_w = eigs(J.T @ A @ J / n), eigs(K_w)
        print(f"step {step}: train loss {L(theta):.2e}; max rel. difference, top 50 eigenvalues of H_GN vs "
              f"A^1/2 Theta A^1/2 / n: {np.max(np.abs(e_gn[:50] - e_w[:50]) / e_gn[:50]):.0e}; "
              f"rank H_GN = {(e_gn > 1e-12 * e_gn[0]).sum()} (n(C-1) = {n * (C - 1)})")
        for name, K in (("unweighted Theta/n", Theta / n), ("weighted A^1/2 Theta A^1/2/n", K_w)):
            ev, U = torch.linalg.eigh(K)
            scores = [class_score(U[:, -j]) for j in range(1, C + 1)]
            print(f"   {name:30s}: top-{C} eigenvalues {np.round(ev.flip(0)[:C].numpy(), 4)}, "
                  f"class-pattern share of eigenvectors {np.round(scores, 2)}")
    theta = theta - 0.1 * grad(L)(theta)
print(f"(random-vector baseline for the class-pattern share: about {C * C / (n * C):.2f})")

step 10: train loss 4.18e-01; max rel. difference, top 50 eigenvalues of H_GN vs A^1/2 Theta A^1/2 / n: 1e-14; rank H_GN = 384 (n(C-1) = 384)
   unweighted Theta/n            : top-4 eigenvalues [26.9443 15.5958 13.4108  9.4235], class-pattern share of eigenvectors [0.97 0.97 0.97 0.97]
   weighted A^1/2 Theta A^1/2/n  : top-4 eigenvalues [2.1448 1.6786 1.2203 0.992 ], class-pattern share of eigenvectors [0.93 0.99 0.96 0.96]
step 2000: train loss 2.26e-04; max rel. difference, top 50 eigenvalues of H_GN vs A^1/2 Theta A^1/2 / n: 7e-14; rank H_GN = 384 (n(C-1) = 384)
   unweighted Theta/n            : top-4 eigenvalues [65.429  49.5037 42.4612 35.1768], class-pattern share of eigenvectors [0.98 0.97 0.97 0.98]
   weighted A^1/2 Theta A^1/2/n  : top-4 eigenvalues [0.0113 0.0084 0.0074 0.0059], class-pattern share of eigenvectors [0.66 0.81 0.77 0.8 ]
(random-vector baseline for the class-pattern share: about 0.03)

Measured Hessian spectra of classifiers show a bulk of eigenvalues near zero, which grows when the network grows, and a few large outliers tied to the C classes, which move when the data change. The outliers sit in the Gauss–Newton part and come from class means of per-example gradient vectors; the within-class scatter of those vectors forms the bulk.

Papers: Sagun et al. (2017) on bulk and outliers; Papyan (2020) on class and cross-class structure (see Readings).

How many outliers? About C dominant ones (about k for small MLPs on k Gaussian clusters), and below them a weaker cross-class layer of at most C(C-2) directions (Papyan measures about C^2-2C) that can merge into the bulk; exactly C-1 in one special case, a zero last layer. Papers report C, C-1 or about C^2 depending on what they count, the cutoff rule, the loss, normalization (BatchNorm shrinks the outliers) and the stage of training.

Derivations and walkthrough. Notation: n examples with true classes y_i; logits z_i with Jacobians J_i=\partial z_i/\partial\theta; predicted probabilities p_i=\mathrm{softmax}(z_i) with entries p_{ik}; e_k the k-th unit vector in \mathbb R^C; H_{\rm GN}=\frac1n\sum_iJ_i^\top A_iJ_i with the softmax Hessian A_i=\mathrm{diag}(p_i)-p_ip_i^\top (6.1).

Why the outliers track classes. Writing A_i=\sum_kp_{ik}(e_k-p_i)(e_k-p_i)^\top (Exercise 6.1.B, item 2) gives J_i^\top A_iJ_i=\sum_kp_{ik}v_{ik}v_{ik}^\top, with

v_{ik}=J_i^\top(e_k-p_i)=-\nabla_\theta\,\ell(z_i,k),

minus the gradient example i would give if it carried label k. So H_{\rm GN} is a weighted second moment of these vectors, labelled by the true class y_i and the hypothetical label k, and a second moment is means plus covariances. If the vectors cluster by (y_i,k), the group means give large eigenvalues (at most C(C-1), with C dominant), and the within-group scatter, spread over many directions, gives the bulk.

Why would they cluster? Restricted to the last-layer weights, v_{ik} is (e_k-p_i)\otimes\tilde h_i, with \tilde h_i the last hidden features plus a bias entry. Features of trained networks cluster by class, so these blocks cluster by (y_i,k); the observed clustering of logit gradients (Papyan; Fort and Ganguli) is the empirical input. Exercise 6.3.A derives the rest given that clustering. It also finds that, for cross-entropy near interpolation, the C class-level directions are spanned to leading order by the class-mean loss gradients; when these are the top directions, the gradient lives in the top subspace for a structural reason (6.4).

Exercise 6.3.A — The class-mean mechanism (derived given clustering). Notation as above, with v_{ik}=J_i^\top(e_k-p_i). The core is item 1, the class-level identity and Weyl bounds of item 3, and item 4; the counting parts are optional.

1. Show that H_{\rm GN}=\frac1n\sum_i\sum_kp_{ik}v_{ik}v_{ik}^\top, and that \sum_kp_{ik}v_{ik}=0 for every i. Deduce \sum_{k\ne y_i}p_{ik}v_{ik}=p_{iy_i}\nabla_\theta\ell_i.
2. Grouping. Group the terms by (c,k), with c=y_i the true class and k the hypothetical label. For each group define the weight w_{ck}=\frac1n\sum_{i:y_i=c}p_{ik} and the weighted mean m_{ck} (weights p_{ik}).
   • Show that H_{\rm GN}=M_{\rm pair}+W, with M_{\rm pair}=\sum_{c,k}w_{ck}m_{ck}m_{ck}^\top and W\succeq0 (the weighted within-group covariances).
   • (Optional, for the count.) Show that \mathrm{rank}\,M_{\rm pair}\le C(C-1).
3. Class level. Let \bar m_c be the w-weighted mean of m_{ck} over the wrong labels k\ne c, with total weight \bar w_c=\sum_{k\ne c}w_{ck}, and set M_{\rm class}=\sum_c\bar w_c\bar m_c\bar m_c^\top (rank \le C).
   • Show M_{\rm pair}\succeq M_{\rm class}. (Optional, for the count.) Show that the range of M_{\rm pair} is spanned by the C class directions \bar m_c together with at most C(C-2) cross-class directions m_{ck}-\bar m_c (a sum of spans, not necessarily orthogonal).
   • Use item 1 to show that \bar m_c is proportional to \sum_{i:y_i=c}p_{ic}\nabla_\theta\ell_i, a weighted class-mean gradient.
   • Conclude with Weyl’s inequality that \lambda_j(H_{\rm GN})\ge\lambda_j(M_{\rm pair}) and \lambda_j(H_{\rm GN})\le\lambda_j(M_{\rm pair})+\|W\|_2. When the within-group scatter is small in every direction, the top C(C-1) eigenvalues are those of the group means.
4. The special case at initialization. Suppose the last layer is initialized to zero, so every p_i is uniform and only the last layer has a nonzero Jacobian. Let \tilde h_i=(h_i,1) be the last hidden features with a bias entry, and M_2=\frac1n\sum_i\tilde h_i\tilde h_i^\top.
   • Show that, restricted to the last layer (zeros elsewhere), H_{\rm GN}=\frac1C\big(I_C-\tfrac1C\mathbf 1\mathbf 1^\top\big)\otimes M_2 for cross-entropy and H_{\rm GN}=I_C\otimes M_2 for square loss.
   • Assuming \lambda_{\max}(M_2) is simple and separated from the rest of M_2’s spectrum, conclude that the top eigenvalue of H_{\rm GN} is \lambda_{\max}(M_2)/C with multiplicity C-1 (cross-entropy), or \lambda_{\max}(M_2) with multiplicity C (square loss). Explain why \lambda_{\max}(M_2) is large when the features have a sizeable mean, as GELU or ReLU features typically do.

Code 6.3.B — The class-mean part reproduces the outliers. For C=2,4,8, train to 100% training accuracy and compare the spectrum of H_{\rm GN} with M_{\rm pair} and M_{\rm class} of Exercise 6.3.A; then check where the gradient sits, and the zero-last-layer case. Look for:

• M_{\rm class} (rank \le C) lower-bounds and tracks the top C eigenvalues of H_{\rm GN} (for C=8, 0.84 against 1.00 for the largest), with a visible knee after the C-th;
• M_{\rm pair} (rank \le C(C-1)) tracks the next eigenvalues too; for C=2,4 its spectrum ends where the bulk begins, while for C=8 there is no gap near index 56 and its smaller eigenvalues sit inside the bulk;
• with a zero last layer the counts are exactly C-1 for cross-entropy and C for square loss.

Code (67 lines)

t0 = time.time()


def pair_means_part(f, theta, X, y, C):
    # H_GN = (1/n) sum_{i,k} p_ik v_ik v_ik^T with v_ik = J_i^T (e_k - p_i).  Group the (i, k) terms by
    # (true class y_i, logit k), and keep only each group's weighted mean ("pair means"), or only each class's
    # mean over its wrong logits k != y ("class means").
    J = jacrev(f)(theta, X)
    p = torch.softmax(f(theta, X), -1)
    v = torch.einsum('ncp,nkc->nkp', J, torch.eye(C)[None] - p[:, None, :])   # v[i, k] = J_i^T (e_k - p_i)
    n = len(X)
    pair, cls = torch.zeros(f.P, f.P), torch.zeros(f.P, f.P)
    for c in range(C):
        idx = y == c
        wsum, msum = 0.0, 0.0
        for k in range(C):
            w = p[idx, k]
            m = (w[:, None] * v[idx, k]).sum(0) / w.sum()
            pair += (w.sum() / n) * torch.outer(m, m)
            if k != c:
                wsum, msum = wsum + w.sum(), msum + w.sum() * m
        cls += (wsum / n) * torch.outer(msum / wsum, msum / wsum)
    G = torch.einsum('nk,nkp,nkq->pq', p, v, v) / n
    return G, pair, cls


fig, ax = plt.subplots(1, 3, figsize=(14, 3.8))
for j, C in enumerate((2, 4, 8)):
    Xc, yc = make_data(C, n_per=50)
    netc = MLP(d=20, h=16, C=C)
    Lc = lambda t: ce(netc(t, Xc), yc)
    th = netc.theta0.clone()
    for _ in range(300):
        th = th - 0.1 * grad(Lc)(th)
    G, pair, cls = pair_means_part(netc, th, Xc, yc, C)
    assert torch.allclose(G, ggn(netc, th, Xc, ce))
    eG, eP, eC = eigs(G), eigs(pair), eigs(cls)
    K = C * C + 6
    k = np.arange(1, K + 1)
    ax[j].semilogy(k, eG[:K], "o-", label=r"$H_{\rm GN}$")
    ax[j].semilogy(k[:C * (C - 1)], eP[:C * (C - 1)], "x--", label=f"pair means (rank $\\leq C(C-1)={C*(C-1)}$)")
    ax[j].semilogy(k[:C], eC[:C], "+:", ms=10, label=f"class means (rank $\\leq C={C}$)")
    ax[j].axvline(C + 0.5, color="gray", ls=":")
    ax[j].set(ylim=(eG[K - 1] / 3, eG[0] * 2), xlabel="index", title=f"C = {C}, after 300 GD steps")
    ax[j].legend(fontsize=8)
    print(f"C={C}: top C+2 of H_GN / top: {np.round(eG[:C + 2] / eG[0], 2)};  class-means part: {np.round(eC[:C + 2] / eG[0], 2)}")
fig.tight_layout(); plt.show()

# Where the gradient sits: the top-C eigenspace of H_GN versus the span of the class-mean gradients (C = 8 run).
Q, _ = torch.linalg.qr(torch.stack([grad(lambda t: ce(netc(t, Xc[yc == c]), yc[yc == c]))(th) for c in range(C)], 1))
U = torch.linalg.eigh(G)[1][:, -C:]
gfull = grad(Lc)(th)
print(f"C={C}: overlap of top-{C} eigenspace with span of class-mean gradients, tr(P_U P_Q)/C = {(U.T @ Q).norm()**2 / C:.2f}; "
      f"fraction of the full gradient in the top-{C} eigenspace = {(U.T @ gfull).norm()**2 / gfull.norm()**2:.2f}")

# At a zero-initialized last layer: every p_i is uniform and only the last layer has a nonzero Jacobian.
for C in (4, 8):
    Xc, yc = make_data(C, n_per=50)
    netc = MLP(d=20, h=16, C=C, zero_last=True)
    hb = torch.cat([netc.features(netc.theta0, Xc), torch.ones(len(Xc), 1)], 1)    # features with a bias entry
    m2_top = eigs(hb.T @ hb / len(Xc))[0]
    for loss, pred in ((ce, m2_top / C), (mse, m2_top)):
        e = eigs(ggn(netc, netc.theta0, Xc, loss))
        gap = int(np.argmax(e[:2 * C] / e[1:2 * C + 1])) + 1
        print(f"zero last layer, C={C}, {loss.__name__}: biggest gap after eigenvalue {gap}; those {gap} equal to "
              f"{e[0]:.4f} (spread {np.ptp(e[:gap]):.0e}); predicted {pred:.4f}; next eigenvalue {e[gap]:.4f}")
print(f"({time.time() - t0:.1f}s)")

C=2: top C+2 of H_GN / top: [1.   0.58 0.03 0.03];  class-means part: [0.98 0.57 0.   0.  ]
C=4: top C+2 of H_GN / top: [1.   0.82 0.67 0.55 0.31 0.22];  class-means part: [0.93 0.77 0.64 0.53 0.   0.  ]
C=8: top C+2 of H_GN / top: [1.   0.82 0.67 0.57 0.38 0.35 0.33 0.25 0.14 0.13];  class-means part: [0.84 0.68 0.55 0.45 0.35 0.32 0.29 0.23 0.   0.  ]

C=8: overlap of top-8 eigenspace with span of class-mean gradients, tr(P_U P_Q)/C = 0.98; fraction of the full gradient in the top-8 eigenspace = 0.99
zero last layer, C=4, ce: biggest gap after eigenvalue 3; those 3 equal to 1.5517 (spread 1e-15); predicted 1.5517; next eigenvalue 0.5316
zero last layer, C=4, mse: biggest gap after eigenvalue 4; those 4 equal to 6.2068 (spread 4e-15); predicted 6.2068; next eigenvalue 2.1263
zero last layer, C=8, ce: biggest gap after eigenvalue 7; those 7 equal to 0.4086 (spread 4e-16); predicted 0.4086; next eigenvalue 0.2249
zero last layer, C=8, mse: biggest gap after eigenvalue 8; those 8 equal to 3.2691 (spread 4e-15); predicted 3.2691; next eigenvalue 1.7990
(6.2s)

So the class level is clean on this data: a knee after C, and a top-C eigenspace that nearly coincides with the span of the class-mean gradients. The cross-class level is clean for small C and merges into the bulk at C=8. That merging, more pronounced on real data, is one reason the papers disagree about the count.

Early in training the gradient comes to point almost entirely into the few top-curvature directions of the Hessian. Yet an optimizer restricted to that subspace stops making progress, and one that avoids it trains normally. The top component is the restoring force of an oscillation across a steep valley; progress is the slow drift along the valley floor, in the bulk.

Papers: Gur-Ari, Roberts and Dyer (2018), the observation; Song, Ahn and Yun (2025), the intervention (see Readings).

Take an eigendirection of the Hessian with curvature \lambda, and let g be the gradient’s component along it. A step -\eta g (\eta the step size) changes the loss, to second order, by

\Delta L=-\eta g^2\Big(1-\frac{\eta\lambda}{2}\Big).

At the edge of stability, \eta\lambda\approx2, the decrease cancels: the step overshoots to the opposite wall. Below the edge with SGD noise, the step lowers the loss only by removing what the noise injected, a fixed stationary floor. Either way the large gradient component along steep directions is not progress.

Derivations and walkthrough. Prerequisites: Edge of stability (td2_3): GD on \tfrac12\lambda x^2 converges iff \eta\lambda<2; in full-batch training the sharpness S=\lambda_{\max}(H) rises to about 2/\eta and hovers there (Cohen et al. 2021), with the iterate oscillating along the top eigenvector (Damian, Nichani and Lee). SGD noise floor (td2_2): near a minimum, SGD with step \eta and batch B fluctuates at a stationary excess loss of about \frac{\eta}{4B}\mathrm{tr}\,\Sigma (\Sigma the per-example gradient covariance; quadratic basin, \eta h\ll1 in the noisy directions).

Notation: H is the Hessian of the training loss L, P_k the orthogonal projector onto its top-k eigenvectors, \nabla L the full-batch gradient, g_B a minibatch gradient.

The observation (Gur-Ari, Roberts and Dyer). With k the number of classes, the fraction of the full-batch gradient in the top-k eigenspace,

f_{\rm top}=\frac{\|P_k\nabla L\|^2}{\|\nabla L\|^2},

rises close to 1 within tens of steps of SGD (in their MNIST MLP, top-20: from 0.20 to 0.94 in 50 steps), and the subspace then changes little.

The intervention (Song, Ahn and Yun). Replace the minibatch update -\eta g_B by Dom-SGD, -\eta P_kg_B, or by Bulk-SGD, -\eta(I-P_k)g_B, with P_k recomputed every step. Switched on in mid-training, once the alignment is high (at 53% and 69% training accuracy where reported), Dom-SGD stops decreasing the training loss and Bulk-SGD trains as well as SGD. Dom also fails for GD at the edge, momentum SGD, Adam and SAM. (Settings: small subsets of MNIST, CIFAR-10 and SST-2; small MLPs, CNN and Transformer; mostly square loss.) Their explanation: near the trajectory the landscape is an ill-conditioned valley, and minibatch noise keeps knocking the iterate up the steep walls, where the gradient is large and points back across. Small-step GD shows no alignment (f_{\rm top}\to0); GD at the edge aligns through its self-stabilizing oscillation. The Mysteries day’s MDL4 §11 makes the same point with a pure quadratic.

Where this holds. The valley picture is about mid-training. Late in training with cross-entropy, on fitted data, alignment can be non-oscillatory. Heuristically, if one competitor dominates an example’s loss, \ell\approx e^{-m} with margin m, then \nabla\ell\approx-\ell\nabla m and \nabla^2\ell\approx\ell(\nabla m\nabla m^\top-\nabla^2m); if the margin Hessian is negligible and \eta\ell\|\nabla m\|^2\ll2, the gradient points along a top curvature direction and a step along it realizes its full first-order decrease. Code 6.4.B panel (b) is consistent with this (there Dom-GD trains like GD and Bulk-GD barely moves) but does not establish the mechanism.

Exercise 6.4.A — A valley. Why can the gradient point almost entirely along a steep direction while all the progress is along a flat one? Take L(x,y)=\frac12ax^2+\ell(y), where x,y are coordinates in parameter space, x is steep (curvature a), and along the floor direction y the loss decreases slowly: \ell'(y)=-c<0, \ell''\approx0. Dom-GD updates only x; Bulk-GD updates only y; \eta is the step size.

1. For full GD, Dom-GD and Bulk-GD, compute the one-step change of L exactly.
2. Edge of stability. Suppose \eta a=2-\delta with 0<\delta\ll1 and the iterate starts at x_0\ne0.
   • Compute f_{\rm top} (the fraction of \|\nabla L\|^2 along x) as a function of x, and show it is close to 1 whenever |x|\gg c/a.
   • How does the loss evolve under each scheme over many steps?
3. SGD noise. Let the x-update carry noise, x_{t+1}=(1-\eta a)x_t+\eta\xi_t with \mathbb E\xi_t=0 and \mathbb E\xi_t^2=\sigma^2, independent across steps.
   • Find the stationary \mathbb E x^2, the stationary loss in x, and the mean squared gradient component a^2\mathbb Ex^2.
   • Compute the ratio a^2\mathbb Ex^2/(a^2\mathbb Ex^2+c^2), a proxy for f_{\rm top}. For Gaussian noise, explain why f_{\rm top} is typically near 1 once a\eta\sigma^2\gg c^2(2-\eta a), and why Dom-SGD stalls while Bulk-SGD keeps descending.
4. No noise, small step. Under GD with \eta a<2 and no noise, what does f_{\rm top} do over time? Compare with the small-learning-rate GD curve in Code 6.4.B.
5. What would make Bulk-GD fail in this model? Relate your answer to why Song, Ahn and Yun switch to Dom/Bulk only after the iterate is near the valley floor.

Code 6.4.B — Alignment, and what projecting does. The first cell trains a tanh network with square loss (4 overlapping classes, 256 points) in three regimes, tracking the loss, the sharpness and f_{\rm top} for the top-4 eigenspace. That eigenspace comes from a warm-started power iteration and is approximate (an exact check gave subspace overlap 0.985), so read the curves as qualitative. The second cell forks unprojected, Dom and Bulk updates from (a) the SGD run, Song et al.’s regime, and (b) a cross-entropy run on separable data after 100% training accuracy, the late regime. Look for:

• GD below the edge: f_{\rm top} drops to zero.
• GD at the edge (\eta=0.5, 2/\eta=4): the sharpness starts at 10 and drops within a few steps (a catapult; td2_3_edge_of_stability §3.7), then sits near 2/\eta from about step 80. f_{\rm top} stays near 0 until roughly step 180, then jumps to near 1 once an oscillation along the top directions has grown; the loss spikes are bursts of it. “At the edge” and “aligned” are not simultaneous here.
• SGD: f_{\rm top} is high on average (about 0.7 in this single run), driven by noise, far below the edge.
• In (a), Dom stalls (loss near 0.032) while Bulk tracks SGD (Bulk 0.032 to about 0.026, SGD 0.031 to about 0.025 over 250 steps): a modest effect, in a single run. In (b) it is the other way round.

Code (63 lines)

t0 = time.time()


def top_subspace(gL, theta, V, iters=2):
    # Top-k eigenpairs of the full-batch Hessian by warm-started block power iteration plus Rayleigh-Ritz.
    # gL is the full-batch gradient function; H V is computed by forward-over-reverse differentiation.
    HV = lambda V: vmap(lambda v: jvp(gL, (theta,), (v,))[1], in_dims=1, out_dims=1)(V)
    for _ in range(iters):
        V, _ = torch.linalg.qr(HV(V))
    ev, R = torch.linalg.eigh(V.T @ HV(V))
    return ev.flip(0), (V @ R).flip(1)


def run(net, X, y, loss, theta, lr, steps, B=None, mode="full", k=4, V=None, seed=0, every=1):
    # mode "full": (S)GD.  "dom": update projected onto the top-k Hessian eigenspace.  "bulk": projected off it.
    # The eigenspace is recomputed every `every` steps.  Returns (theta, V, rows of (loss, f_top, sharpness)).
    gL = grad(lambda t: loss(net(t, X), y))
    g = torch.Generator().manual_seed(seed)
    if V is None:
        _, V = top_subspace(gL, theta, torch.linalg.qr(torch.randn(net.P, k, generator=g))[0], iters=20)
    out = []
    for s in range(steps):
        if s % every == 0:
            ev, V = top_subspace(gL, theta, V)
            gf = gL(theta)
            frac = ((V.T @ gf).norm() ** 2 / gf.norm() ** 2).item()   # Gur-Ari et al.'s f_top (full-batch gradient)
        idx = torch.arange(len(X)) if B is None else torch.randint(len(X), (B,), generator=g)
        gb = grad(lambda t: loss(net(t, X[idx]), y[idx]))(theta)
        if mode == "dom":
            gb = V @ (V.T @ gb)
        elif mode == "bulk":
            gb = gb - V @ (V.T @ gb)
        theta = theta - lr * gb
        out.append((loss(net(theta, X), y).item(), frac, ev[0].item()))
    return theta, V, np.array(out)


Xs, ys = make_data(4, n_per=64, sep=1.0)                       # overlapping clusters: not separable at once
nets = MLP(d=20, h=16, C=4, act=torch.tanh)
regimes = {"GD, lr 0.05 (far below the edge)": (0.05, None),
           "GD, lr 0.5 (edge of stability)": (0.5, None),
           "SGD, lr 0.1, batch 8": (0.1, 8)}
states, curves = {}, {}
for name, (lr, B) in regimes.items():
    th, V, out = run(nets, Xs, ys, mse, nets.theta0.clone(), lr, 400, B, every=5)
    states[name], curves[name] = (th, V), out

fig, ax = plt.subplots(1, 3, figsize=(14, 3.6))
for i, (name, out) in enumerate(curves.items()):
    lr = regimes[name][0]
    ax[0].plot(out[:, 1], color=f"C{i}", lw=0.8, label=name)
    ax[1].semilogy(out[:, 0], color=f"C{i}", lw=0.8)
    ax[2].plot(out[:, 2], color=f"C{i}", lw=0.8)
    if 2 / lr < 12:
        ax[2].axhline(2 / lr, color=f"C{i}", ls=":", lw=1)
    print(f"{name}: mean f_top over steps 200-400 = {out[200:, 1].mean():.2f}, "
          f"final sharpness {out[-1, 2]:.2f} (2/lr = {2 / lr:.0f})")
ax[0].set(xlabel="step", ylabel=r"$f_{\rm top}$", title=r"$\|P_4\nabla L\|^2/\|\nabla L\|^2$ (MSE)")
ax[0].legend(fontsize=8)
ax[1].set(xlabel="step", title="training loss")
ax[2].set(xlabel="step", ylim=(0, 12), title=r"sharpness; dotted: $2/\eta=4$ (others: 40, 20)")
fig.tight_layout(); plt.show()
print(f"({time.time() - t0:.1f}s)")

GD, lr 0.05 (far below the edge): mean f_top over steps 200-400 = 0.00, final sharpness 6.20 (2/lr = 40)
GD, lr 0.5 (edge of stability): mean f_top over steps 200-400 = 0.80, final sharpness 4.15 (2/lr = 4)
SGD, lr 0.1, batch 8: mean f_top over steps 200-400 = 0.69, final sharpness 5.45 (2/lr = 20)

(10.0s)

Code (27 lines)

t0 = time.time()
labels = {"full": "unprojected", "dom": "Dom (top-4 only)", "bulk": "Bulk (top-4 removed)"}
fig, ax = plt.subplots(1, 2, figsize=(12, 3.6))

# (a) Song-Ahn-Yun's setting: SGD, square loss, mid-training.  Fork the SGD run of the previous cell at step 400.
th, V = states["SGD, lr 0.1, batch 8"]
for mode, ls in (("full", "-"), ("dom", "--"), ("bulk", ":")):
    _, _, out = run(nets, Xs, ys, mse, th, 0.1, 250, 8, mode=mode, V=V, seed=1, every=2)
    ax[0].semilogy(out[:, 0], ls, lw=1.2, label=labels[mode])
    print(f"(a) MSE, SGD, {mode:4s}: loss {out[0, 0]:.4f} -> mean of last 50 steps {out[-50:, 0].mean():.4f}")
ax[0].set(title="(a) MSE, SGD lr 0.1, batch 8, forked at step 400", xlabel="steps after the fork", ylabel="training loss")

# (b) Outside that setting: cross-entropy on separable data, small-lr GD, after interpolation.
Xc, yc = make_data(4, n_per=50)
netc = MLP(d=20, h=16, C=4)
thc = netc.theta0.clone()
for _ in range(100):
    thc = thc - 0.05 * grad(lambda t: ce(netc(t, Xc), yc))(thc)
for mode, ls in (("full", "-"), ("dom", "--"), ("bulk", ":")):
    _, _, out = run(netc, Xc, yc, ce, thc, 0.05, 200, None, mode=mode, every=2)
    ax[1].semilogy(out[:, 0], ls, lw=1.2, label=labels[mode])
    print(f"(b) CE, GD,  {mode:4s}: loss {out[0, 0]:.4f} -> {out[-1, 0]:.4f}, mean f_top {out[:, 1].mean():.2f}")
ax[1].set(title="(b) CE, GD lr 0.05, from step 100 (100% train accuracy)", xlabel="further steps")
for a in ax:
    a.legend()
fig.tight_layout(); plt.show()
print(f"({time.time() - t0:.1f}s)")

(a) MSE, SGD, full: loss 0.0313 -> mean of last 50 steps 0.0249
(a) MSE, SGD, dom : loss 0.0314 -> mean of last 50 steps 0.0318
(a) MSE, SGD, bulk: loss 0.0320 -> mean of last 50 steps 0.0255
(b) CE, GD,  full: loss 0.0284 -> 0.0056, mean f_top 0.99
(b) CE, GD,  dom : loss 0.0284 -> 0.0057, mean f_top 1.00
(b) CE, GD,  bulk: loss 0.0289 -> 0.0264, mean f_top 0.97

(21.0s)

Fork a run at step t, train two children with independent minibatch noise, and measure the loss on the straight line between them: once there is no barrier, the run has committed to a basin. Commitment comes early, and in a ResNet on CIFAR it coincides with the window in which the tangent kernel moves fastest. In a solvable model the basin is chosen while the relevant kernel (Gauss–Newton) eigenvalue is still in its early exponential growth.

Papers: Fort et al. (2020), the notebook’s key paper; Frankle et al. (2020) on instability to SGD noise (see Readings).

In a model linear in its parameters (a fixed kernel, the lazy limit) the loss is convex, so there is no basin to choose. The smallest model with a choice is the scalar two-layer net f=w_1w_2 with loss \frac12(w_1w_2-s)^2, s>0: its minima form two branches (w_1,w_2>0 or both <0) on either side of a saddle at the origin. From small balanced initialization, noise fixes the branch early in the escape from the saddle, while the kernel \Theta=w_1^2+w_2^2 is still one to two orders of magnitude below its final value. A conserved imbalance d=w_2^2-w_1^2 well above the noise fixes the branch at initialization, even when it is far too small to make the dynamics lazy.

Derivations and walkthrough. Notation: \theta_A,\theta_B are the final parameters of two children trained from a common checkpoint; \mathcal E is the error (test or training); \Theta_t=J_tJ_t^\top is the empirical tangent-kernel Gram matrix at step t, with J_t the Jacobian of the network’s outputs on a data set.

Operational definition (Frankle et al.). Copy the parameters at step t, train two children to the end with independent minibatch noise, and evaluate the error along the straight line between them. The error barrier is

\max_{\alpha\in[0,1]}\mathcal E\big((1-\alpha)\theta_A+\alpha\theta_B\big)-\tfrac12\big(\mathcal E(\theta_A)+\mathcal E(\theta_B)\big),

and the run is committed (stable, in their terms) at step t if the barrier is near zero, under 2% in their experiments. With the loss in place of \mathcal E it is the loss barrier, used below.

What is measured.

• Frankle et al.: stable at initialization for LeNet, after about 1.5–3% of training for CIFAR ResNet-20/VGG-16, and after about 16–20% for ImageNet ResNet-50/Inception-v3 (test error; on training error ResNet-50 commits later and Inception-v3 not within their runs).
• Fort et al.: in ResNet-20 on CIFAR, commitment within 2–3 epochs coincides with the window in which the kernel distance 1-\langle\Theta_t,\Theta_{t'}\rangle_F/(\|\Theta_t\|_F\|\Theta_{t'}\|_F) changes fastest. Training linearized around the epoch-3–4 kernel gives at least 3× lower error than around the initial one.
• A related early-phase signature: in the critical-periods experiments of Achille, Rovere and Soatto, the Fisher information (the Gauss–Newton matrix, for cross-entropy) rises rapidly early in training and then decreases (td2_8_further_readings §4), the same shape as the weighted kernel in Code 6.2.B.

How curvature enters. By 6.2, H_{\rm GN} and the loss-weighted kernel have the same nonzero eigenvalues, so growth of kernel eigenvalues is growth of curvature outliers; Fort et al.’s unweighted cosine relates to curvature only through that spectral part.

The solvable model. For f=w_1w_2 (Exercise 6.1.A, item 3) both branches compute the same function, so here “the basin” is a choice between symmetric copies; in real networks the copies are related by permutations and the choice can also select different functions. The kernel \Theta=w_1^2+w_2^2 is also the nonzero eigenvalue of the Gauss–Newton matrix. Gradient flow conserves the imbalance d; balanced means d=0, and for |d|\gg s the dynamics are lazy (the kernel barely moves; cf. the linear-network regimes of td1_linear_networks). Exercise 6.5.A and Code 6.5.B add two findings to the summary above: the barrier between branches is large while within a branch it is fourth order in the difference; and with two modes of different strength, commitment happens in two stages, of which a whole-kernel cosine distance (dominated by the largest eigenvalue) registers only the first.

So the model supports a spectral statement: a basin is chosen while the corresponding kernel (Gauss–Newton) eigenvalue is in its early exponential growth, before it is an outlier in absolute terms. That is consistent with Fort et al.’s timing for the dominant directions without explaining it, and with the silent alignment of Atanasov, Bordelon and Pehlevan: in linear networks from small initialization, the kernel aligns to the task while still small, before the loss moves.

Exercise 6.5.A — Basins of the scalar linear network. Take L=\frac12(w_1w_2-s)^2 with s>0, branches w_1,w_2>0 and w_1,w_2<0, kernel \Theta=w_1^2+w_2^2 and imbalance d=w_2^2-w_1^2. The loss barrier between two endpoints is the maximum loss on the straight line between them minus the mean of the endpoint losses; Fort et al.’s kernel distance between times t,t' is 1-\cos(\Theta_t,\Theta_{t'}).

1. Let the children end on opposite branches, at (a,s/a) and (-b,-s/b) with a,b>0. Show that the loss along the linear path has barrier \frac{s^2}2\big(\frac{(a+b)^2}{4ab}\big)^2\ge\frac{s^2}2.
2. On the same branch, at (a,s/a) and (b,s/b), show w_1w_2=s\big[1+\alpha(1-\alpha)\frac{(a-b)^2}{ab}\big] along the path, and a barrier of \frac{s^2(a-b)^4}{32a^2b^2}. Evaluate it for s=1, a=1.1, b=0.9.
3. Choosing the branch. Gradient flow conserves w_2^2-w_1^2. Near the origin, use u=(w_1+w_2)/\sqrt2 and v=(w_1-w_2)/\sqrt2, and linearize:
   • show that \dot u=su and \dot v=-sv;
   • argue that from a small random initialization with gradient noise, the branch is fixed by the sign of u once its growth outruns the noise;
   • show that this happens while the loss is still close to its initial value s^2/2, early in a window in which \Theta grows by orders of magnitude;
   • (optional) with du=su\,dt+\sigma\,dW and u(0)\sim\mathcal N(0,a_0^2), show that two children forked at time t end on opposite branches with probability \arccos\rho(t)/\pi, where \rho(t) is the correlation of the two children’s limits of e^{-st}u(t) as t\to\infty (the “bracket” of the hint), and show that \rho(t)=1-\frac{q}{a_0^2+q}e^{-2st} with q=\sigma^2/2s (Code 6.5.B plots this prediction).
4. Convexity. For a model linear in \theta with a convex loss, show that the barrier is at most \frac12|L(\theta_A)-L(\theta_B)|.
5. Imbalance removes the choice. Under gradient flow, suppose d=w_2(0)^2-w_1(0)^2>0 with w_1(0)^2\ll d.
   • Show that w_2 never changes sign, so the branch is \mathrm{sign}\,w_2(0), fixed at initialization.
   • Show that at the minimum reached \Theta=\sqrt{d^2+4s^2}, while initially \Theta\approx d. Conclude that two thresholds differ: imbalance well above the noise fixes the branch, while laziness (a nearly frozen kernel) needs d\gg s. In between, the kernel can change by a large factor (about 2s/d when d\ll s) although nothing is being chosen.
6. Two modes. Take L=\sum_{m=1,2}\frac12(w_1^{(m)}w_2^{(m)}-s_m)^2 with s_1>s_2: two decoupled modes of a linear network on two orthonormal input directions, with kernel \Theta=\mathrm{diag}(\Theta_1,\Theta_2), \Theta_m=(w_1^{(m)})^2+(w_2^{(m)})^2.
   • From small balanced initialization, when is each mode’s branch chosen?
   • Compare the time at which each mode’s branch is fixed with the time at which \Theta_m grows from its small initial value to O(s_m).
   • The direction \Theta/\|\Theta\|_F ends at (s_1,s_2)/\sqrt{s_1^2+s_2^2}. When does Fort et al.’s kernel distance register mode 2, and why does it barely register mode 2’s commitment?

Code 6.5.B — Forks of a two-mode linear network, and imbalance. The model of Exercise 6.5.A, item 6, with s_1=1, s_2=0.25, trained by noisy GD w\leftarrow w-\eta\nabla L+\sqrt\eta\,\sigma\xi (\xi standard normal, \eta=0.05, \sigma=0.02) from three starts: small balanced (rich), w_2(0)=\pm0.3 (d=0.09: committed but not lazy), and w_2(0)=\pm3 (lazy). For 60 parents per setting it forks 4 pairs of children at several times; children train 1000 steps, with the noise switched off for the last 200. It reports the loss barrier, the fraction of pairs on opposite branches per mode (against \arccos\rho(t)/\pi from item 3), the per-mode \Theta_m, and Fort et al.’s kernel velocity (cosine distance ten steps apart, per step). This shows a mechanism that makes commitment early in a model where every step can be followed; it is not evidence that the same mechanism operates in a ResNet. Look for:

• Balanced: disagreement roughly follows the prediction (tails off by a few percentage points either way, e.g. mode 1 at step 60: 0.05 measured vs 0.02 predicted; mode 2 at step 120: 0.03 vs 0.10). Mode 1 commits by about step 60–100, mode 2 by about step 120 (a few percent of pairs persist to step 200).
• Over steps 0–60, \Theta_1 grows about a hundredfold but stays tens of times below its final 2s_1; mode 2 commits while \Theta_2 is a few percent of 2s_2.
• Kernel velocity: noise for about 20 steps, then mode 1’s escape (steps 20–60, overlapping its commitment), at its minimum (about 10^{-7} per step) while mode 2 commits, rising again over steps 150–300 as \Theta_2 catches up.
• Imbalanced: no opposite-branch pairs; within-branch barriers of 10^{-5} to 10^{-6}. With w_2(0)=\pm0.3 the kernel still moves a lot early, so a moving kernel does not imply a choice; with \pm3 it barely moves.

Code (77 lines)

t0 = time.time()
rng = np.random.default_rng(0)
s_modes = np.array([1.0, 0.25])            # two decoupled modes of a linear net, s1 > s2
eta, sigma, a0 = 0.05, 0.02, 0.01          # step size, gradient-noise level, init scale


def noisy_gd(w, steps, noise=sigma):
    # w[..., m, 0], w[..., m, 1]: first- and second-layer weight of mode m.  Loss sum_m (w1 w2 - s_m)^2 / 2.
    for _ in range(steps):
        r = w[..., 0] * w[..., 1] - s_modes
        g = np.stack([r * w[..., 1], r * w[..., 0]], -1)
        w = w - eta * g + np.sqrt(eta) * noise * rng.standard_normal(w.shape)
    return w


def loss(w):
    return 0.5 * ((w[..., 0] * w[..., 1] - s_modes) ** 2).sum(-1)


def kernel(w):
    # The NTK Gram matrix on the two (orthonormal) input directions is diagonal: Theta = diag(w1^2 + w2^2).
    return (w ** 2).sum(-1)


def init(n, w2_scale):
    w = a0 * rng.standard_normal((n, 2, 2))
    if w2_scale > a0:                        # imbalanced: second layer starts at +-w2_scale
        w[..., 1] = w2_scale * np.sign(rng.standard_normal((n, 2)))
    return w


setups = {"balanced, init 0.01": a0, "imbalanced, w2(0) = ±0.3": 0.3, "strongly imbalanced (lazy), w2(0) = ±3": 3.0}
forks = [0, 10, 20, 40, 60, 80, 120, 160, 200, 300, 400]
n_par, n_pair, T = 60, 4, 1000
alpha = np.linspace(0, 1, 21).reshape(-1, 1, 1, 1, 1)
fig, ax = plt.subplots(1, 3, figsize=(15, 3.6))
for j, (name, w2s) in enumerate(setups.items()):
    parents, last, rows = init(n_par, w2s), 0, []
    for tf in forks:
        parents, last = noisy_gd(parents, tf - last), tf
        kids = [noisy_gd(noisy_gd(np.repeat(parents[:, None], n_pair, 1), T - 200), 200, noise=0.0) for _ in range(2)]
        Lp = loss((1 - alpha) * kids[0] + alpha * kids[1])
        barrier = Lp.max(0) - 0.5 * (Lp[0] + Lp[-1])
        opposite = (np.sign(kids[0][..., 0]) != np.sign(kids[1][..., 0])).mean((0, 1))   # per mode
        rows.append((tf, barrier.mean(), *opposite, np.median(loss(parents))))
    rows = np.array(rows)
    print(name)
    for tf, b, o1, o2, l in rows:
        print(f"  fork at step {int(tf):3d}: mean loss barrier {b:.1e}; pairs on opposite branches: "
              f"mode 1 {o1:.2f}, mode 2 {o2:.2f}; parent loss {l:.3f}")
    ls = ["-", "--", ":"][j]
    if j == 0:
        tt = np.linspace(0, forks[-1], 200)
        for m, mk in ((0, "o"), (1, "s")):
            ax[0].plot(rows[:, 0], rows[:, 2 + m], mk, color=f"C{m}", label=f"mode {m + 1} ($s={s_modes[m]}$), measured")
            q = sigma**2 / (2 * s_modes[m])                    # linearized-escape prediction (Exercise 6.5.A, item 3)
            rho = 1 - q / (a0**2 + q) * np.exp(-2 * s_modes[m] * eta * tt)
            ax[0].plot(tt, np.arccos(rho) / np.pi, "-", color=f"C{m}", lw=1, label=r"$\arccos\rho(t)/\pi$")

    # Kernel velocity along fresh parent trajectories: Fort et al.'s cosine distance per step, 10 steps apart.
    w, vel, eig = init(200, w2s), [], []
    for k in range(40):
        K1 = kernel(w); w = noisy_gd(w, 10); K2 = kernel(w)
        cos = (K1 * K2).sum(-1) / np.linalg.norm(K1, axis=-1) / np.linalg.norm(K2, axis=-1)
        vel.append(np.median(1 - cos) / 10); eig.append(np.median(K2, 0))
    eig = np.array(eig)
    if j < 2:
        for m in range(2):
            ax[1].semilogy(10 * np.arange(1, 41), eig[:, m], ls, color=f"C{m}", label=f"$\\Theta_{m + 1}$, {name}")
    ax[2].semilogy(10 * np.arange(1, 41), np.array(vel) + 1e-9, ls, color="C2", label=name)
    print(f"  kernel velocity, median over 200 runs, steps 10/100/200/300/400: {np.round(np.array(vel)[[0, 9, 19, 29, 39]], 7)}")
ax[0].set(xlabel="fork step", ylabel="fraction of child pairs", title="balanced init: children on opposite branches")
ax[0].legend(fontsize=7)
ax[1].set(xlabel="step", title=r"per-mode kernel eigenvalues $\Theta_m=w_1^2+w_2^2$ (median)"); ax[1].legend(fontsize=7)
ax[2].set(xlabel="step", title="whole-kernel velocity: (1 - cos) per step"); ax[2].legend(fontsize=7)
fig.tight_layout(); plt.show()
print(f"({time.time() - t0:.1f}s)")

balanced, init 0.01
  fork at step   0: mean loss barrier 1.9e-01; pairs on opposite branches: mode 1 0.35, mode 2 0.46; parent loss 0.531
  fork at step  10: mean loss barrier 1.3e-01; pairs on opposite branches: mode 1 0.24, mode 2 0.37; parent loss 0.531
  fork at step  20: mean loss barrier 7.9e-02; pairs on opposite branches: mode 1 0.14, mode 2 0.30; parent loss 0.531
  fork at step  40: mean loss barrier 4.8e-02; pairs on opposite branches: mode 1 0.08, mode 2 0.25; parent loss 0.528
  fork at step  60: mean loss barrier 3.1e-02; pairs on opposite branches: mode 1 0.05, mode 2 0.20; parent loss 0.514
  fork at step  80: mean loss barrier 1.9e-02; pairs on opposite branches: mode 1 0.03, mode 2 0.13; parent loss 0.432
  fork at step 120: mean loss barrier 1.1e-03; pairs on opposite branches: mode 1 0.00, mode 2 0.03; parent loss 0.037
  fork at step 160: mean loss barrier 1.4e-03; pairs on opposite branches: mode 1 0.00, mode 2 0.04; parent loss 0.029
  fork at step 200: mean loss barrier 9.6e-04; pairs on opposite branches: mode 1 0.00, mode 2 0.03; parent loss 0.023
  fork at step 300: mean loss barrier 4.5e-05; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.003
  fork at step 400: mean loss barrier 6.3e-05; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.000
  kernel velocity, median over 200 runs, steps 10/100/200/300/400: [6.8321e-03 2.0000e-07 9.0000e-07 3.9000e-06 3.4000e-06]
imbalanced, w2(0) = ±0.3
  fork at step   0: mean loss barrier 4.5e-05; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.531
  fork at step  10: mean loss barrier 4.6e-05; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.480
  fork at step  20: mean loss barrier 3.4e-05; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.398
  fork at step  40: mean loss barrier 4.8e-05; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.124
  fork at step  60: mean loss barrier 3.8e-05; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.021
  fork at step  80: mean loss barrier 4.5e-05; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.011
  fork at step 120: mean loss barrier 3.8e-05; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.003
  fork at step 160: mean loss barrier 4.0e-05; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.001
  fork at step 200: mean loss barrier 4.0e-05; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.000
  fork at step 300: mean loss barrier 4.0e-05; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.000
  fork at step 400: mean loss barrier 4.1e-05; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.000
  kernel velocity, median over 200 runs, steps 10/100/200/300/400: [1.6472e-03 1.1200e-05 3.4000e-06 2.5000e-06 2.1000e-06]
strongly imbalanced (lazy), w2(0) = ±3
  fork at step   0: mean loss barrier 1.4e-06; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.531
  fork at step  10: mean loss barrier 8.9e-07; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.000
  fork at step  20: mean loss barrier 1.4e-06; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.000
  fork at step  40: mean loss barrier 1.3e-06; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.000
  fork at step  60: mean loss barrier 1.1e-06; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.000
  fork at step  80: mean loss barrier 1.1e-06; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.000
  fork at step 120: mean loss barrier 1.2e-06; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.000
  fork at step 160: mean loss barrier 1.3e-06; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.000
  fork at step 200: mean loss barrier 1.1e-06; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.000
  fork at step 300: mean loss barrier 1.1e-06; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.000
  fork at step 400: mean loss barrier 1.4e-06; pairs on opposite branches: mode 1 0.00, mode 2 0.00; parent loss 0.000
  kernel velocity, median over 200 runs, steps 10/100/200/300/400: [5.1e-06 1.0e-06 7.0e-07 9.0e-07 1.0e-06]

(5.1s)