Training Dynamics 2.8: further readings, and a map of the day

1. Dropout as a geometric-mean ensemble (§1).

Also: Srivastava et al., Dropout (JMLR 2014), the standard reference. Dropout as an adaptive L^2 penalty is derived in the Mysteries day’s MDL1 §4, exercise 4.4.

2. Layer cushion, noise robustness and basin size (§2).

3. Geometry-adaptive harmonic bases (§3).

Also: Mohan et al., bias-free denoising CNNs (ICLR 2020) (why the Jacobian reading works); Kamb & Ganguli, an analytic theory of creativity (2024) (locality and equivariance alone predict trained diffusion models’ outputs); Belrose et al., statistics of increasing complexity (2024) (td2_7’s “low-order statistics first”, measured on real networks).

4. Critical periods and loss of plasticity (readings only). Does it matter when a network sees something? Yes. A deficit early in training, such as blurred images for the first epochs, lowers final accuracy in a way later training does not repair (a critical period), and a network trained on a long sequence of tasks gradually loses the ability to learn new ones (loss of plasticity). Early training builds something later training cannot easily undo.

Two hooks to this day: critical periods already appear in deep linear networks, where the mode equations of td1_linear_networks §1 make them tractable (Kleinman et al.), and plasticity loss is tied to changes in loss-landscape curvature (Lyle et al.; td2_6). The Mysteries day lists critical periods among its strange facts in MDL4 §6.

Also: Kleinman, Achille & Soatto, critical periods in deep linear networks (ICLR 2024); Lyle et al., understanding plasticity (ICML 2023) (plasticity loss tracks curvature, often without saturated units); Ash & Adams, on warm-starting (NeurIPS 2020) (the everyday version, and “shrink and perturb”).

5. Fine-tuning and RL update dynamics (§5).

Also: Cai et al., predictability of RL dynamics (2025) (the top singular component of each update recovers about 99% of the accuracy); Schulman and Thinking Machines Lab, LoRA without regret (2025) (rank-1 LoRA suffices for policy-gradient RL); Shenfeld, Pari & Agrawal, RL’s Razor (2025) (forgetting tracks the KL divergence to the base model); Malladi et al., a kernel view of fine-tuning (ICML 2023) (prompt-based fine-tuning as kernel regression with the pretrained tangent kernel); CRW, some interesting papers on RLVR (blog, 2026) (a short, informal survey of these papers and more).

Dropout trains a randomly thinned subnetwork at each step; at test time it runs the full network once, with each weight scaled by its keep probability. Why does that one pass stand in for averaging all the subnetworks? For a logistic unit it is exactly their normalized geometric mean, which is close to their arithmetic mean.

Paper: Baldi & Sadowski, Understanding dropout, §3–4 (see Readings).

Derivations and walkthrough. Setup: One logistic unit, \sigma(x)=1/(1+e^{-x}), with inputs I_i and weights w_i. Dropout gates each input with \delta_i\sim\mathrm{Bernoulli}(p_i), independently, so the input sum is S=\sum_i\delta_iw_iI_i and the output is O=\sigma(S). Index the 2^n subnetworks by k, with probabilities P_k, sums S_k and outputs O_k. The arithmetic mean is E=\sum_kP_kO_k; the weighted geometric means are G=\prod_kO_k^{P_k} and G'=\prod_k(1-O_k)^{P_k}; the normalized weighted geometric mean is \mathrm{NWGM}=G/(G+G'), a probability that treats the output and its complement symmetrically. The test-time rule computes \sigma\big(\sum_ip_iw_iI_i\big)=\sigma(\mathbb E S).

Exercise 1.1 — The averaging identity. Show that \mathrm{NWGM}=\sigma(\mathbb E S) exactly, for any probabilities P_k (they need not come from independent gates).

Exercise 1.2 — How far is the geometric mean from the arithmetic one? Suppose the spread of S across subnetworks is small: S=\mu+\varepsilon with \mu=\mathbb ES, \mathbb E\varepsilon=0 and \mathrm{Var}\,S=s^2 small. Write V=\mathrm{Var}\,O. To leading order in s^2, show that

E-\mathrm{NWGM}\;\approx\;\tfrac12\sigma''(\mu)\,s^2\;\approx\;\frac{V\,(1-2E)}{2E(1-E)}.

So the gap is second order in the spread, and to this order vanishes at E=\tfrac12.

Numerical check. Exact enumeration of all 2^{12} subnetworks of a 12-input unit. The last column is the paper’s own second-order estimate of the gap, its eq. (20), V(1-2E)/(1-2V); in this small-spread regime it is smaller than the exact gap by about the factor 1/(2E(1-E))\ge2. The paper’s conclusion, that the NWGM is close to E, is unaffected.

sigma = lambda s: 1 / (1 + np.exp(-s))
n = 12
masks = np.array(list(itertools.product([0, 1], repeat=n)))   # all 2**12 subnetworks


def means(w, I, p, bias=0.0):
    # Arithmetic mean E, variance V, NWGM and sigma(E[S]) over every dropout mask, exactly.
    P = np.prod(np.where(masks == 1, p, 1 - p), axis=1)        # probability of each mask
    S = masks @ (w * I) + bias                                 # input sum of each subnetwork
    O = sigma(S)
    E = P @ O
    V = P @ (O - E) ** 2
    G, Gc = np.exp(P @ np.log(O)), np.exp(P @ np.log(1 - O))
    return E, V, G / (G + Gc), sigma(P @ S)


# Exercise 1.1: the identity is exact, even with large weights and unequal keep probabilities.
E, V, nwgm, sig_ES = means(rng.normal(0, 1, n), rng.normal(0, 1, n), rng.uniform(0.3, 0.9, n))
print(f"NWGM = {nwgm:.12f}   sigma(E[S]) = {sig_ES:.12f}   arithmetic mean E = {E:.6f}\n")

# Exercise 1.2: small weights (small spread), with a bias to move E away from 1/2.
w, I, p = rng.normal(0, 0.1, n), rng.normal(0, 1, n), np.full(n, 0.5)
print(" bias      E      E - NWGM    V(1-2E)/(2E(1-E))   B-S eq.(20) V(1-2E)/(1-2V)")
for b in [0.0, 1.0, 2.0, 3.0]:
    E, V, nwgm, _ = means(w, I, p, bias=b)
    print(f"{b:5.1f}  {E:.4f}  {E - nwgm:+.3e}   {V*(1-2*E)/(2*E*(1-E)):+.3e}          {V*(1-2*E)/(1-2*V):+.3e}")

NWGM = 0.289411894326   sigma(E[S]) = 0.289411894326   arithmetic mean E = 0.319655

 bias      E      E - NWGM    V(1-2E)/(2E(1-E))   B-S eq.(20) V(1-2E)/(1-2V)
  0.0  0.4210  +6.489e-04   +6.439e-04          +3.152e-04
  1.0  0.6622  -1.222e-03   -1.216e-03          -5.457e-04
  2.0  0.8412  -1.534e-03   -1.540e-03          -4.120e-04
  3.0  0.9349  -8.902e-04   -9.012e-04          -1.098e-04

In what precise sense is a trained network robust to noise, and does training make it so? Arora et al. measure how strongly noise injected at one layer is attenuated by the layers above it, on the training data. Their key quantity, the layer cushion, is small at random initialization and grows during training, most in the higher layers.

This is a quantitative replacement for “flat minima generalize”, whose raw flatness depends on how the network is parametrized (td2_4 §4.2; the Mysteries day’s MDL1 §4, under “Why flatness could matter, and why sharpness needs a metric”). A network that attenuates injected noise can be compressed a lot, since compression error acts like injected noise, and the compressed network gets a generalization bound the authors report as orders of magnitude better than parameter counting. td2_6 §6.5 covers when the basin is chosen.

Papers: Arora, Ge, Neyshabur & Zhang, Stronger generalization bounds for deep nets via a compression approach; Dziugaite & Roy on nonvacuous PAC-Bayes bounds (see Readings).

Derivations and walkthrough. Notation: Layer i computes x^i=A_i\phi(x^{i-1}), with weight matrix A_i, elementwise activation \phi (ReLU in the paper) and input x^0. \lVert A\rVert_F is the Frobenius norm, \lVert A\rVert_2 the spectral norm, and \eta\sim\mathcal N(0,I) a standard Gaussian noise vector (not a learning rate here), so \mathbb E[\eta\eta^\top]=I.

Definitions (from the paper). The noise sensitivity of a map M at input x is the expected squared noise-to-signal ratio of the output under isotropic input noise scaled by \lVert x\rVert:

\psi(M,x)=\mathbb E_{\eta\sim\mathcal N(0,I)}\frac{\lVert M(x+\eta\lVert x\rVert)-M(x)\rVert^2}{\lVert M(x)\rVert^2}.

The layer cushion \mu_i is the largest \mu such that \mu\lVert A_i\rVert_F\lVert\phi(x^{i-1})\rVert\le\lVert A_i\phi(x^{i-1})\rVert for every training input. The stable rank of A is \lVert A\rVert_F^2/\lVert A\rVert_2^2.

Exercise 2.1 — Cushion is inverse noise sensitivity.

1. For a linear map M, show \psi(M,x)=\lVert M\rVert_F^2\lVert x\rVert^2/\lVert Mx\rVert^2. Conclude that 1/\mu_i^2 is the largest noise sensitivity of A_i over the training inputs, that \psi is at least the stable rank, and so \mu_i\le1/\sqrt{\text{stable rank}}. When is the bound attained?
2. Take x\sim\mathcal N(0,I_h) independent of A\in\mathbb R^{h'\times h}. Show \mathbb E\big[\lVert Ax\rVert^2/\lVert x\rVert^2\big]=\lVert A\rVert_F^2/h, so a typical input has cushion about 1/\sqrt h. What must happen during training for the cushion to grow well above that?

Numerical check of both parts, on a 256\times256 layer with singular values 1/k and on a random Gaussian layer.

h = 256
cushion = lambda M, X: np.linalg.norm(M @ X, axis=0) / (np.linalg.norm(M) * np.linalg.norm(X, axis=0))

# 2.1(1): a layer with a decaying spectrum. Monte Carlo noise sensitivity against the formula,
# then the bound 1/sqrt(stable rank) and inputs that attain it (along the top right singular vector).
U, _ = np.linalg.qr(rng.normal(size=(h, h)))
Vt, _ = np.linalg.qr(rng.normal(size=(h, h)))
B = U @ np.diag(1 / np.arange(1, h + 1)) @ Vt.T
x = rng.normal(size=h)
eta = rng.normal(size=(h, 4000))
psi_mc = np.mean(np.linalg.norm(B @ (eta * np.linalg.norm(x)), axis=0) ** 2) / np.linalg.norm(B @ x) ** 2
print(f"noise sensitivity at a random x: Monte Carlo {psi_mc:.0f}, formula {np.linalg.norm(B)**2 * (x @ x) / np.linalg.norm(B @ x)**2:.0f}")
X = rng.normal(size=(h, 1000))
print(f"bound 1/sqrt(stable rank) = {np.linalg.norm(B, 2) / np.linalg.norm(B):.3f}; cushion with inputs along the top "
      f"singular vector {cushion(B, Vt[:, :1]).min():.3f}, in the top-5 subspace {cushion(B, Vt[:, :5] @ rng.normal(size=(5, 1000))).min():.3f}, "
      f"random inputs {cushion(B, X).min():.3f}")

# 2.1(2): a random layer (large stable rank) with random inputs: cushion about 1/sqrt(h).
A = rng.normal(0, 1 / np.sqrt(h), (h, h))
print(f"random layer, random inputs: min cushion {cushion(A, X).min():.3f}, mean {cushion(A, X).mean():.3f}, 1/sqrt(h) = {1/np.sqrt(h):.3f}")

noise sensitivity at a random x: Monte Carlo 142, formula 142
bound 1/sqrt(stable rank) = 0.781; cushion with inputs along the top singular vector 0.781, in the top-5 subspace 0.172, random inputs 0.022
random layer, random inputs: min cushion 0.053, mean 0.062, 1/sqrt(h) = 0.062

What function does a trained image denoiser, the network inside a diffusion model, compute once it has enough data to stop memorizing? Locally, it shrinks the noisy image along an image-dependent basis of harmonics that oscillate along contours and in smooth regions, and these bases appear even on data for which they are suboptimal. It is an inductive-bias result like the spectral bias of td2_7 §7.3, read off from the network’s Jacobian.

Paper: Kadkhodaie, Guth, Simoncelli & Mallat, Generalization in diffusion models arises from geometry-adaptive harmonic representations (see Readings).

Derivations and walkthrough. The optimal denoiser: A clean image x\in\mathbb R^d is observed as y=x+\sigma z, with z\sim\mathcal N(0,I) and noise level \sigma. Write p_\sigma for the density of y, and \nabla f(y) for the Jacobian of a map f at y. The MSE-optimal denoiser is f^\star(y)=\mathbb E[x\mid y]=y+\sigma^2\nabla\log p_\sigma(y) (the Miyasawa or Tweedie identity, which is why a denoiser gives the score). Its Jacobian is \nabla f^\star(y)=\sigma^{-2}\,\mathrm{Cov}[x\mid y], symmetric and positive semidefinite, and the optimal error is \mathrm{MSE}=\mathbb E\,\mathrm{tr}\,\mathrm{Cov}[x\mid y]=\sigma^2\,\mathbb E\,\mathrm{tr}\,\nabla f^\star(y). Exercise 3.1 checks the last two for Gaussian data.

What a bias-free network computes. The paper’s denoisers are ReLU networks with no additive biases, so they are positively homogeneous and f(y)=\nabla f(y)\,y wherever f is differentiable. Idealizing the Jacobian as symmetric positive semidefinite (approximately true for trained networks), with \nabla f(y)=\sum_k\lambda_k(y)e_k(y)e_k(y)^\top,

f(y)=\sum_k\lambda_k(y)\,\langle y,e_k(y)\rangle\,e_k(y):

shrinkage by factors \lambda_k(y) along a basis e_k(y), both of which can depend on y. The exercise works out the case where they cannot.

Exercise 3.1 — The Gaussian case: shrinkage in a fixed basis. Let x\sim\mathcal N(0,\Sigma) in \mathbb R^d, with eigendecomposition \Sigma=\sum_k\lambda_kv_kv_k^\top, observed as y=x+\sigma z. Here \lambda_k are the data variances, not the shrinkage factors \lambda_k(y).

1. Show f^\star(y)=\Sigma(\Sigma+\sigma^2I)^{-1}y. What are the basis and the shrinkage factors, and do they depend on y?
2. Check \nabla f^\star=\sigma^{-2}\mathrm{Cov}[x\mid y] directly.
3. Show \mathrm{MSE}=\sigma^2\sum_k\lambda_k/(\lambda_k+\sigma^2). Read \mathrm{tr}\,\nabla f^\star as an effective number of dimensions kept, and compare the MSE with d\sigma^2, the error of returning y unchanged.

Numerical check by Monte Carlo, on 64-dimensional Gaussian data with variances 1/k^2 and noise level \sigma=0.1.

d, s = 64, 0.1
lam = 1 / np.arange(1, d + 1) ** 2                       # a decaying prior spectrum
Q, _ = np.linalg.qr(rng.normal(size=(d, d)))
Sigma = Q @ np.diag(lam) @ Q.T
J = Sigma @ np.linalg.inv(Sigma + s**2 * np.eye(d))     # the optimal (linear) denoiser and its Jacobian

Xs = (Q * np.sqrt(lam)) @ rng.normal(size=(d, 100_000))   # clean samples
Ys = Xs + s * rng.normal(size=Xs.shape)                  # noisy observations
mse = np.mean(np.sum((J @ Ys - Xs) ** 2, axis=0))
print(f"Monte Carlo MSE {mse:.4f}   sigma^2 tr J {s**2 * np.trace(J):.4f}   "
      f"returning y: d sigma^2 = {d * s**2:.2f}   effective dimension tr J = {np.trace(J):.1f} of {d}")
print("shrinkage factors along the top 6 principal directions:",
      np.round(np.sort(np.linalg.eigvalsh(J))[::-1][:6], 3))

Monte Carlo MSE 0.1367   sigma^2 tr J 0.1367   returning y: d sigma^2 = 0.64   effective dimension tr J = 13.7 of 64
shrinkage factors along the top 6 principal directions: [0.99  0.962 0.917 0.862 0.8   0.735]

Post-training moves a pretrained model a small distance; in which directions? Measured in weights, RL fine-tuning of language models changes few parameters (as counted in bfloat16), yet each update matrix has full numerical rank, is dominated by one singular direction, and avoids the pretrained matrices’ “principal weights”, where supervised fine-tuning (SFT) does not.

Papers: Mukherjee, Yuan, Hakkani-Tür & Peng on RL fine-tuning small subnetworks; Zhu et al., The path not taken (see Readings).

The question this day adds, from td2_6 §6.4: how does the update sit relative to the curvature (the Hessian or Gauss–Newton eigenbasis)? None of these papers measures that directly at language-model scale, and Exercise 5.2 shows why a weight-matrix proxy for it needs care. The Mysteries day’s MDL4 §12 reads Mukherjee et al. and Shenfeld et al. as a pair, including what “unchanged” means in bfloat16.

Derivations and walkthrough. Notation: \Delta W is the change in one weight matrix from fine-tuning. Its numerical rank counts singular values above a small tolerance; its stable rank is \lVert\Delta W\rVert_F^2/\lVert\Delta W\rVert_2^2 (Frobenius over spectral norm, squared). Cai et al. keep only the top singular component of each update, rescaled to the update’s norm, and recover about 99% of the trained model’s accuracy.

Exercise 5.1 — Full rank and dominated by one direction. Let \Delta W=uv^\top+\epsilon N with unit vectors u,v\in\mathbb R^n, a dense Gaussian N with entries of variance 1/n, and small \epsilon. What are its rank, numerical rank and stable rank? Which of these do Mukherjee et al. and Cai et al. measure?

Numerical check of Exercise 5.1, plus the measurement issue behind the sparsity numbers: a dense update much smaller than the bfloat16 spacing of the weights is mostly invisible when both checkpoints are rounded to bfloat16. Shenfeld et al. found that in their setting the sparsity came from bfloat16 and vanished in float32, at the same performance.

u, v = rng.normal(size=512), rng.normal(size=512)
dW = np.outer(u / np.linalg.norm(u), v / np.linalg.norm(v)) + 0.02 * rng.normal(size=(512, 512)) / np.sqrt(512)
sv = np.linalg.svd(dW, compute_uv=False)
print(f"rank-1 plus small dense noise: numerical rank {np.linalg.matrix_rank(dW)} of 512, "
      f"stable rank {np.sum(sv**2) / sv[0]**2:.2f}, top direction's share of ||dW||_F^2 {sv[0]**2 / np.sum(sv**2):.0%}\n")


def to_bf16(a):
    # Round finite float32 values to the nearest bfloat16 (ties to even); return as float32.
    b = np.asarray(a, dtype=np.float32).view(np.uint32).astype(np.uint64)
    b = (b + 0x7FFF + ((b >> 16) & 1)) >> 16
    return (b.astype(np.uint32) << 16).view(np.float32)


W0 = rng.normal(0, 0.02, (512, 512)).astype(np.float32)
for scale in [1e-6, 1e-5, 1e-4]:
    W1 = W0 + scale * rng.normal(size=(512, 512)).astype(np.float32)   # a dense update: every entry moves
    print(f"dense update of size {scale:.0e}: entries unchanged after rounding both to bfloat16 "
          f"{np.mean(to_bf16(W0) == to_bf16(W1)):.0%}, in float32 {np.mean(W0 == W1):.0%}")

rank-1 plus small dense noise: numerical rank 512 of 512, stable rank 1.20, top direction's share of ||dW||_F^2 83%

dense update of size 1e-06: entries unchanged after rounding both to bfloat16 97%, in float32 0%
dense update of size 1e-05: entries unchanged after rounding both to bfloat16 80%, in float32 0%
dense update of size 1e-04: entries unchanged after rounding both to bfloat16 30%, in float32 0%

Exercise 5.2 — Principal directions are not curvature directions. Take a single linear layer f(x)=Wx with loss L(W)=\tfrac12\mathbb E\lVert Wx-y\rVert^2 and input second moment \Sigma_x=\mathbb E[xx^\top].

1. Show that L(W+\Delta W)=L(W)+\langle\nabla L,\Delta W\rangle+\tfrac12\mathrm{tr}(\Delta W\,\Sigma_x\Delta W^\top) exactly. What are the eigen-directions of this curvature, and their eigenvalues?
2. Let W=\sum_i s_iu_iv_i^\top be the SVD. Give a case where every highest-curvature update is orthogonal (in the Frobenius inner product) to W’s top singular direction u_1v_1^\top. When are “along W’s principal directions” and “along high curvature” close to the same thing?