Chapter 7 · I.7

Problems

A chapter that uses mathematics has to teach that mathematics by making the reader compute. A chapter that only displays equations has failed, however correct the equations are.

M3Load-bearing

5/5 problems5/4 variants10/10 exercisesquota met, and enforced

The contract

5 worked problems, minimum.
4 distinct variants, and no variant more than half of them.
10 exercises, every one with a published solution.
At least one numeric problem — present.
At least one symbolic problem — present.
At least one limit or counterexample problem — present.
At least one complexity or shape problem — present.
At least one ▲▲▲ problem — present.

Numerical instantiationSymbolic derivationConstructed failureCost accountingDimensional algebra

Problem I.7.B01

The learning rate at which descent stops converging

numeric▲△△

Four significant figures; the threshold is exact.

STATEMENT

Find the exact learning rate above which gradient descent diverges on a quadratic, then run descent either side of it. Show that the boundary is not a gradual degradation but a change of sign in an exponent.

GIVEN

A quadratic loss with a diagonal Hessian:

\mathcal{L}(\vec{\theta}) = \tfrac{1}{2}\vec{\theta}^{\mathsf{T}}\mat{H}\vec{\theta}, \qquad \mat{H} = \mathrm{diag}(20,\ 0.2)

with the optimum at the origin and the starting point $\vec{\theta}_0 = (1, 1)$ .

FIND

The stability threshold; the condition number; the error norm after 1, 5, 10, 50 and 100 steps at $\eta \in \{0.01,\ 0.09,\ 0.0990,\ 0.11,\ 0.2\}$ ; the per-step multiplier in each eigendirection at $\eta = 0.09$ and $\eta = 0.11$ ; and the optimal rate with its contraction factor.

STRATEGY

Because the Hessian is diagonal the two coordinates never interact, so the iteration is two independent scalar recurrences. Each is multiplied by a fixed number every step, and a geometric sequence converges exactly when that number is below one in magnitude. The whole problem reduces to that comparison.

SOLUTION

Step 1 — the recurrence. Substituting $\nabla\mathcal{L} = \mat{H}\vec{\theta}$ into (I.7.1):

\vec{\theta}_{t+1} = \vec{\theta}_t - \eta\mat{H}\vec{\theta}_t = (\mat{I} - \eta\mat{H})\vec{\theta}_t

which is (I.7.3). In the eigenbasis — here the coordinate axes — each component is multiplied by $1 - \eta\lambda_i$ every step.

Step 2 — the threshold. A geometric sequence converges when $|1 - \eta\lambda| < 1$ , that is $0 < \eta < 2/\lambda$ . Every coordinate must converge, so the binding constraint is the largest eigenvalue:

\eta < \frac{2}{\lambda_{\max}} = \frac{2}{20} = 0.1

The condition number is $\kappa = 20/0.2 = 100$ .

Step 3 — either side of it. Error norms:

$\eta$	$t=1$	$t=5$	$t=10$	$t=50$	$t=100$
$0.0100$	$1.279$	$1.043$	$0.9860$	$0.9048$	$0.8186$
$0.0900$	$1.267$	$0.9702$	$0.8408$	$0.4033$	$0.1626$
$0.0990$	$1.386$	$1.280$	$1.158$	$0.5202$	$0.1914$
$0.1100$	$1.548$	$2.644$	$6.243$	$9.100\times10^{3}$	$8.282\times10^{7}$
$0.2000$	$3.150$	$2.430\times10^{2}$	$5.905\times10^{4}$	$7.179\times10^{23}$	$5.154\times10^{47}$

Step 4 — why the boundary is sharp. The multipliers:

$\eta$	steep direction	flat direction
$0.09$	$0.8000$	$0.9820$	converges
$0.11$	$1.2000$	$0.9780$	diverges

A two-percent change in $\eta$ moves the steep multiplier from $0.8$ to $1.2$ . Both are close to one; the difference is which side. Raised to the hundredth power, $0.8^{100} \approx 2 \times 10^{-10}$ and $1.2^{100} \approx 8 \times 10^{7}$ — seventeen orders of magnitude apart, from a rate difference of two percent.

Note also that the flat direction converges in both rows. Divergence is entirely the steep direction’s doing, and it is invisible in the flat one.

Step 5 — the best a scalar rate can do. Minimising $\max_i |1 - \eta\lambda_i|$ balances the two extremes, giving

\eta^* = \frac{2}{\lambda_{\max} + \lambda_{\min}} = 0.0990099, \qquad \text{contraction } \frac{\kappa - 1}{\kappa + 1} = 0.980198

which is (I.7.4). At that rate the error shrinks by a factor of a thousand in 346 steps — for a two-dimensional quadratic with an exactly known Hessian. The slowness is not the algorithm being poorly implemented; it is $\kappa = 100$ being charged for.

Answer

Threshold $\eta < 2/\lambda_{\max} = 0.1$ exactly; $\kappa = 100$ .

At $\eta = 0.09$ the error falls to $0.1626$ by step 100; at $\eta = 0.11$ it reaches $8.282 \times 10^{7}$ . The multipliers in the steep direction are $0.8000$ and $1.2000$ respectively.

Optimal $\eta^* = 0.0990099$ , contraction $0.980198$ per step, 346 steps for a thousandfold reduction.

Check — sanity

The smallest rate is not the safest choice. At $\eta = 0.01$ the error after 100 steps is $0.8186$ — barely moved, and worse than $\eta = 0.09$ by a factor of five. Stability and progress are different requirements, and satisfying the first with room to spare fails the second.

$\eta = 0.0990$ is worse than $\eta = 0.09$ at $t = 100$ . $0.1914$ against $0.1626$ , even though $0.0990$ is the theoretical optimum. The optimum minimises the asymptotic rate, and at 100 steps the iteration has not reached its asymptote: the larger multiplier in the steep direction is still costing more than the smaller one in the flat direction is saving.

The threshold uses $\lambda_{\max}$ and the rate uses both. Stability is decided by the steepest direction alone; speed is decided by their ratio. Two different eigenvalues answering two different questions, which is why quoting a single “curvature” for a model explains neither.

Where this breaks

A neural network loss is not a quadratic, and the ways it is not matter here. Its Hessian changes with $\vec{\theta}$ , so there is no single $\lambda_{\max}$ and no fixed threshold — the stable rate at initialisation may be unstable after a thousand steps, which is one reason schedules decay. Its eigenvalues span many orders of magnitude with a long tail, so $\kappa$ computed honestly is enormous and the bound above becomes vacuous rather than merely pessimistic. And it is non-convex, so there is no $\vec{\theta}^*$ to contract towards.

What survives the move to a real network is the qualitative shape and not the numbers: there is an edge, crossing it diverges rather than degrades, and the edge is set by the largest curvature rather than the typical one. A rate that “nearly worked” is not near-optimal — it is on the wrong side of a boundary.

The demonstration is also noiseless. With minibatch gradients the iteration is stochastic, the error does not contract monotonically, and a rate slightly below the deterministic threshold can still diverge because a single noisy gradient overshoots. I.7.X09 measures that noise; the interaction is what makes the usable band narrower in practice than this problem suggests.

Variation

Change $\mat{H}$ to $\mathrm{diag}(20, 20)$ , so $\kappa = 1$ . Predict the optimal rate and the contraction factor before computing them, then say how many steps a thousandfold reduction takes. The answer explains why preconditioning is worth its cost.

Problem I.7.B02

Deriving the threshold and the optimal rate

symbolic▲▲△

Symbolic throughout.

STATEMENT

Derive, from the update rule alone, the exact condition under which gradient descent converges on a quadratic, the learning rate that converges fastest, and the rate it achieves. The numbers of I.7.B01 should fall out as one instance.

GIVEN

$\mathcal{L}(\vec{\theta}) = \tfrac{1}{2}(\vec{\theta} - \vec{\theta}^*)^{\mathsf{T}}\mat{H}(\vec{\theta} - \vec{\theta}^*)$ with $\mat{H}$ symmetric positive definite, eigenvalues $0 < \lambda_{\min} = \lambda_1 \le \dots \le \lambda_n = \lambda_{\max}$ .

FIND

The error recurrence; the condition on $\eta$ for convergence; the $\eta$ minimising the worst-case per-step contraction; and that contraction in terms of $\kappa = \lambda_{\max}/\lambda_{\min}$ .

STRATEGY

Change coordinates so the problem falls apart. $\mat{H}$ is symmetric, so it has an orthonormal eigenbasis; written in that basis the update decouples into $n$ independent scalar recurrences, and a scalar geometric sequence is something whose convergence can simply be read off.

SOLUTION

Step 1 — the error recurrence. The gradient is $\nabla\mathcal{L} = \mat{H}(\vec{\theta} - \vec{\theta}^*)$ . Writing $\vec{e}_t = \vec{\theta}_t - \vec{\theta}^*$ and substituting into (I.7.1):

\vec{e}_{t+1} = \vec{e}_t - \eta\mat{H}\vec{e}_t = (\mat{I} - \eta\mat{H})\vec{e}_t

so after $t$ steps $\vec{e}_t = (\mat{I} - \eta\mat{H})^t \vec{e}_0$ . The whole trajectory is one matrix raised to a power.

Step 2 — decouple. $\mat{H} = \mat{Q}\Lambda\mat{Q}^{\mathsf{T}}$ with $\mat{Q}$ orthogonal. Since $\mat{I} - \eta\mat{H} = \mat{Q}(\mat{I} - \eta\Lambda)\mat{Q}^{\mathsf{T}}$ and $\mat{Q}^{\mathsf{T}}\mat{Q} = \mat{I}$ , powers pass straight through:

(\mat{I} - \eta\mat{H})^t = \mat{Q}(\mat{I} - \eta\Lambda)^t\mat{Q}^{\mathsf{T}}

In the rotated coordinates $\tilde{\vec{e}} = \mat{Q}^{\mathsf{T}}\vec{e}$ the recurrence is $n$ separate scalar ones:

\tilde{e}_{i,t} = (1 - \eta\lambda_i)^t\,\tilde{e}_{i,0}

The directions do not interact. This is why one badly served direction cannot be rescued by another going well.

Step 3 — convergence. $(1 - \eta\lambda_i)^t \to 0$ exactly when $|1 - \eta\lambda_i| < 1$ , that is $0 < \eta < 2/\lambda_i$ . Every direction must converge, so the smallest of those bounds binds:

0 < \eta < \frac{2}{\lambda_{\max}}

The condition is necessary as well as sufficient — above it the component along $\vec{q}_{\max}$ grows geometrically, and no other direction can offset that.

Step 4 — the fastest rate. Asymptotically the slowest-decaying component dominates, so the quantity to minimise is

\rho(\eta) = \max_i |1 - \eta\lambda_i|

Each $|1 - \eta\lambda_i|$ is a V-shaped function of $\eta$ with its minimum at $1/\lambda_i$ . The maximum over $i$ is attained at one of the two extreme eigenvalues: for small $\eta$ by $1 - \eta\lambda_{\min}$ , which decreases, and for large $\eta$ by $\eta\lambda_{\max} - 1$ , which increases. The minimum of the maximum is where they cross:

1 - \eta\lambda_{\min} = \eta\lambda_{\max} - 1 \quad\Longrightarrow\quad \eta^* = \frac{2}{\lambda_{\max} + \lambda_{\min}}

Step 5 — the rate. Substituting back:

\rho(\eta^*) = 1 - \frac{2\lambda_{\min}}{\lambda_{\max} + \lambda_{\min}} = \frac{\lambda_{\max} - \lambda_{\min}}{\lambda_{\max} + \lambda_{\min}} = \frac{\kappa - 1}{\kappa + 1}

which is (I.7.4). Dividing through by $\lambda_{\min}$ is what turns two eigenvalues into one ratio: only the shape of the problem matters, not its scale.

Step 6 — check against I.7.B01. With $\lambda = (20, 0.2)$ : $\eta^* = 2/20.2 = 0.0990099$ and $\rho = 99/101 = 0.980198$ , matching the measured values exactly.

Answer

\vec{e}_{t} = (\mat{I} - \eta\mat{H})^{t}\vec{e}_0, \qquad \text{converges} \iff 0 < \eta < \frac{2}{\lambda_{\max}}

\eta^* = \frac{2}{\lambda_{\max} + \lambda_{\min}}, \qquad \rho(\eta^*) = \frac{\kappa - 1}{\kappa + 1}

Check — sanity

$\kappa = 1$ gives $\rho = 0$ . A perfectly conditioned quadratic is solved in a single step, and the formula says so: $\eta^* = 1/\lambda$ lands exactly on the optimum. Any correct derivation must reproduce this.

$\kappa \to \infty$ gives $\rho \to 1$ . Progress per step vanishes as the problem stiffens, which is the behaviour to expect and the reason the ratio rather than either eigenvalue is the quantity of interest.

$\eta^*$ lies strictly inside the stability interval. $2/(\lambda_{\max} + \lambda_{\min}) < 2/\lambda_{\max}$ since $\lambda_{\min} > 0$ . The optimum is never on the boundary, which is worth knowing: a rate tuned to the edge of stability is not the fast one.

Where this breaks

Every step used symmetry and positive definiteness. A real loss Hessian is symmetric but is not positive definite away from a minimum: at a saddle it has negative eigenvalues, and for $\lambda_i < 0$ the multiplier $1 - \eta\lambda_i$ exceeds one for every positive $\eta$ . Gradient descent moves away from a saddle along those directions, which is the behaviour wanted, but it means the analysis above describes only the neighbourhood of a minimum and says nothing about the journey there.

The asymptotic framing hides something practical too. $\rho$ governs the long-run rate, and training runs are not long-run: they are stopped after a budget. The transient matters, which is why $\eta^*$ was beaten at 100 steps by a smaller rate in I.7.B01 despite being optimal in the limit.

Momentum breaks the derivation rather than extending it. Its update depends on two previous iterates, so the recurrence is second-order and the analysis needs the roots of a quadratic characteristic polynomial rather than a single multiplier. That is what produces $\sqrt{\kappa}$ in place of $\kappa$ — I.7.T2, and the subject of I.7.B03.

Variation

Repeat Step 4 for the transient rather than the asymptote: which $\eta$ minimises $\|\vec{e}_T\|$ for a fixed budget $T$ and a given $\vec{e}_0$ ? Show that it depends on $\vec{e}_0$ , and say what that implies about tuning a learning rate on a short run.

Problem I.7.B03

A problem no learning rate solves

counterexample▲▲▲

Six decimal places on multipliers; step counts exact.

STATEMENT

Exhibit a quadratic on which every learning rate fails — not one badly chosen rate, but the whole stable range — and identify what is responsible. Then show that remembering the previous step fixes it, and by how much.

GIVEN

$\mat{H} = \mathrm{diag}(100,\ 0.01)$ , so $\kappa = 10^4$ . Start at $\vec{\theta}_0 = (1, 1)$ ; the target is an error norm below $10^{-3}$ .

FIND

The stability limit; the step count for every $\eta$ across the stable range; the count at the optimal $\eta^*$ ; and the same for heavy-ball momentum at its own optimum, with the speed-up.

STRATEGY

The two directions decouple, as in I.7.B02, so the difficulty can be stated before anything is run: the steep direction demands $\eta < 0.02$ , and at any such $\eta$ the flat direction’s multiplier is $1 - 0.01\eta \ge 0.9998$ . One rate cannot serve both. Sweeping confirms there is no exception rather than merely no obvious one.

SOLUTION

Step 1 — the bind, stated in advance. Stability requires $\eta < 2/\lambda_{\max} = 0.02$ . At the largest permissible rate the flat direction is multiplied by

1 - \eta\lambda_{\min} > 1 - 0.02(0.01) = 0.9998

so it loses at most two parts in ten thousand per step, however $\eta$ is chosen. That is the counterexample in one line; the sweep is the demonstration.

Step 2 — the sweep.

$\eta$	steep multiplier	flat multiplier	steps to $10^{-3}$
$0.020000$	$1.000000$	$0.999800$	diverges
$0.010000$	$0.000000$	$0.999900$	$69\,075$
$0.019800$	$0.980000$	$0.999802$	$34\,885$
$0.019801$	$0.980100$	$0.999802$	$34\,883$
$0.019900$	$0.990000$	$0.999801$	$34\,709$
$0.020010$	$1.001000$	$0.999800$	diverges
$0.050000$	$4.000000$	$0.999500$	diverges

Nothing in the stable range does better than about $34\,700$ steps for a thousandfold reduction on a two-dimensional problem with an exactly known Hessian.

Step 3 — the most instructive row. At $\eta = 0.01$ the steep multiplier is exactly zero: $1 - 0.01(100) = 0$ . That direction is solved perfectly, in one step, and the run still takes $69\,075$ steps — twice as many as the rates that solve it imperfectly. Perfect handling of the direction that constrains stability buys nothing, because the cost is entirely in the other direction. This is why “tune the learning rate” is not a complete answer to slow training.

Step 4 — the optimum, and its transient. $\eta^* = 0.0199980$ gives $\rho = 0.99980002$ and $36\,272$ steps — slightly worse than $\eta = 0.0199$ at $34\,709$ . As in I.7.B01, $\eta^*$ minimises the asymptotic rate while the run is measured over a finite budget. The gap is small here because $\kappa$ is large enough that the asymptote arrives early.

Step 5 — momentum. Heavy ball keeps a running average of gradients, so the recurrence is second-order and its rate depends on $\sqrt{\kappa}$ . At the classical optimum

\beta = \left(\frac{\sqrt{\kappa} - 1}{\sqrt{\kappa} + 1}\right)^2 = 0.96078816, \qquad \eta = \frac{4}{(\sqrt{\lambda_{\max}} + \sqrt{\lambda_{\min}})^2} = 0.03921184

with $\rho = (\sqrt{\kappa} - 1)/(\sqrt{\kappa} + 1) = 0.98019802$ and 708 steps, against $36\,272$ : a speed-up of $51.2\times$ .

Step 6 — why it works. Momentum does not make the steep direction stable at a larger rate; note that $\eta = 0.0392$ is above the plain-descent threshold of $0.02$ . It works because the steep direction’s gradient alternates in sign when overshooting, so consecutive contributions to the moving average partly cancel, while the flat direction’s gradient keeps the same sign and accumulates. The average suppresses oscillation and amplifies drift, which is exactly the asymmetry the problem needs.

Answer

Stability requires $\eta < 0.02$ , and every rate in that range needs at least about $34\,700$ steps; $\eta^* = 0.0199980$ takes $36\,272$ . The obstruction is $\kappa = 10^4$ itself, not the choice of rate.

Heavy-ball momentum at $\beta = 0.96078816$ , $\eta = 0.03921184$ reaches the same target in 708 steps, a $51.2\times$ speed-up.

Check — sanity

The measured momentum rate matches the theory. $\rho = 0.98019802 = 99/101 = (\sqrt{\kappa} - 1)/(\sqrt{\kappa} + 1)$ for $\sqrt{\kappa} = 100$ , exactly. And it equals the plain-descent rate for a problem of condition number 100 — momentum turns $\kappa = 10^4$ into an effective $10^2$ , which is what “the square root” means concretely.

The speed-up is $51\times$ , not the $100\times$ the rates suggest. Both methods are measured to a fixed target from a fixed start, so both pay a transient the asymptotic rates ignore, and momentum’s is proportionally larger because it needs several steps before its average carries information. Quoting $\sqrt{\kappa}$ as the speed-up would be quoting a limit as a measurement.

Divergence appears immediately above $0.02$ . $\eta = 0.02001$ gives a steep multiplier of $1.001$ and diverges; $\eta = 0.0199$ gives $0.99$ and converges. Five parts in ten thousand separate them, which is the sharp edge of I.7.B01 in a stiffer problem.

Where this breaks

The momentum parameters here were computed from $\lambda_{\max}$ and $\lambda_{\min}$ , which nobody has for a real network. In practice $\beta$ is set to $0.9$ and left there, which is optimal for $\kappa \approx 361$ and merely serviceable elsewhere — the tuned value above is $0.96$ , and using $0.9$ on this problem gives up much of the advantage. The theory says what is available; it does not hand over the setting.

Momentum also does not help with the part of the difficulty that is noise. Averaging gradients reduces variance, but the accelerated rate of I.7.T2 assumes exact gradients, and with minibatch noise the improvement degrades towards the noise floor. On a stochastic problem momentum is worth having and the factor is smaller than $\sqrt{\kappa}$ .

Finally, this is a fixed quadratic and a neural network is not. Its curvature changes as training proceeds, so the tuned $(\eta, \beta)$ for step 1 is wrong by step 10 000 — which is one honest reading of why schedules exist, and why adaptive methods that estimate curvature per coordinate as they go, rather than being told it once, are the ones that ended up in use.

Variation

Set $\beta = 0.9$ — the default nobody tunes — and keep $\eta$ at its value above. Predict whether the run gets faster or slower before measuring, then find the $\kappa$ for which $\beta = 0.9$ is optimal. That number explains where the default came from.

Problem I.7.B04

The memory an optimiser costs before any activation is stored

complexity▲▲△

GB to two decimals, at 4 bytes per fp32 value.

STATEMENT

Count the copies of the parameter vector each optimiser holds, price them for three model sizes, and identify the size at which the choice of optimiser stops being a convergence question and becomes a hardware one.

GIVEN

A model with $N$ parameters in fp32. Every method holds the weights and the gradient; they differ in what else persists between steps. Sizes: 124M, 1.5B and 7B.

FIND

State per parameter and total copies for SGD, SGD with momentum, RMSProp and Adam; total memory at each size; the optimiser state as a share of the total; and what swapping Adam for SGD with momentum saves at 1.5B.

STRATEGY

Count in copies of $\vec{\theta}$ rather than in bytes, because the ratio is what transfers between precisions and model sizes. Weights and gradients are common to all four, so the whole difference is the third term.

SOLUTION

Step 1 — the state each method keeps.

method	state	copies	bytes/param
SGD	0	2	8
SGD + momentum	1	3	12
RMSProp	1	3	12
Adam / AdamW	2	4	16
Adam, fp32 master	4	6	24

Adam keeps both moments of (I.7.6): $m$ for direction and $v$ for scale. That is two extra copies, and with mixed precision an fp32 master copy of fp16 weights pushes it to six.

Step 2 — priced.

method	124M	1.5B	7B
SGD	$0.99$ GB	$12.00$ GB	$56.00$ GB
SGD + momentum	$1.49$ GB	$18.00$ GB	$84.00$ GB
RMSProp	$1.49$ GB	$18.00$ GB	$84.00$ GB
Adam / AdamW	$1.98$ GB	$24.00$ GB	$112.00$ GB

Step 3 — where it stops being a preference. At 7B, Adam’s four copies come to $112$ GB. The largest single accelerator commonly available holds $80$ GB, so the model cannot be trained on one device with Adam before a single activation is stored — and activations are the larger term at any useful batch, by I.6.B05. SGD with momentum brings it to $84$ GB, which is still over. The decision has moved from “which converges better” to “which fits”, and no amount of gradient checkpointing addresses it, because checkpointing trades activation memory and this is not activation memory.

Step 4 — the share.

\text{SGD } 0\%, \qquad \text{momentum } 33.3\%, \qquad \text{Adam } 50.0\%

Half of Adam’s four-way footprint is state that exists only to choose a step size per coordinate.

Step 5 — the swap. At 1.5B:

24.00\ \text{GB} \longrightarrow 18.00\ \text{GB}, \qquad 6.00\ \text{GB saved},\ 1.33\times

A third of the optimiser footprint, for one moment estimate. Whether that is a good trade depends on whether the per-coordinate scaling is buying convergence on the model in question — which is an empirical matter, and is exactly the question the framing of Adam as a default discourages anyone from asking.

Answer

Copies: SGD 2, momentum 3, RMSProp 3, Adam 4, Adam with an fp32 master 6.

At 1.5B: SGD $12.00$ GB, momentum $18.00$ GB, Adam $24.00$ GB. At 7B, Adam needs $112.00$ GB — over a single 80 GB device before any activation.

Adam’s state is $50\%$ of its own footprint. Swapping to SGD with momentum at 1.5B saves $6.00$ GB, a factor of $1.33$ .

Check — sanity

The ratios are precision-independent. Counting in copies means fp16 halves every number in the table and changes none of the ratios. That is why the copy count is the quantity to remember and the GB figures are an illustration.

Adam is exactly twice SGD. Four copies against two, at every size. A useful sanity check on any memory estimate: if switching from SGD to Adam does not roughly double the parameter-side footprint, something else is being counted.

The 124M row is unremarkable and that is the point. $1.98$ GB against $0.99$ GB — a gigabyte, on hardware with tens. The choice genuinely does not matter at that scale, which is why the habit of reaching for Adam forms on small models and then goes unexamined on large ones.

Where this breaks

This counts the optimiser and nothing else. Activations are usually the larger term — I.6.B05 measured $3.22$ GB of them for a 48-layer model at batch $32 \times 512$ , against $0.20$ GB of parameters — and they scale with batch while this table does not. A memory plan built from this table alone will be wrong in the direction that matters.

The fp32-master row also understates modern practice. A production mixed-precision setup may hold fp16 weights, fp16 gradients, an fp32 master copy, and fp32 moments, and sharded training distributes those copies across devices so that the per-device figure is smaller than any row here while the total is larger. ZeRO-style sharding exists precisely because this table is the binding constraint at scale.

Finally, the memory saved by dropping to SGD is not free. Adam’s per-coordinate scaling is what makes one learning rate serviceable across parameters whose gradients differ by orders of magnitude, and removing it usually means tuning a schedule that Adam absorbed. The trade is real in both directions.

Variation

Recompute the 7B row in bf16 for weights and gradients with fp32 moments — the common mixed-precision arrangement. Does Adam fit on an 80 GB device then? State what you had to assume about the master copy to get your answer.

Problem I.7.B05

Adam's state, shape by shape, and what its division buys

shape▲▲△

Exact shapes; algebra exact.

STATEMENT

Give the shape of every tensor Adam holds for a concrete model, then establish the property those shapes exist to provide: that the step taken in each coordinate is unchanged when the loss is multiplied by a constant.

GIVEN

The network of I.5.B03 — widths $(784, 512, 256, 128, 10)$ , four affine layers, $567\,434$ parameters. Adam as in (I.7.6), with $\beta_1 = 0.9$ , $\beta_2 = 0.999$ , $\varepsilon = 10^{-8}$ .

FIND

The shape of $m$ and $v$ for each parameter tensor and the totals; then, for a loss rescaled to $c\mathcal{L}$ , the effect on $g$ , $m$ , $v$ and the update.

STRATEGY

Every moment estimate is an average of things with the shape of the gradient, and a gradient has the shape of its parameter — so the shapes follow without any new argument. The scale question is answered by tracking a factor of $c$ through (I.7.6) and seeing where it cancels.

SOLUTION

Step 1 — the shapes. $m$ and $v$ are elementwise running averages of $g$ and $g^2$ , so each has exactly the shape of the parameter it belongs to:

tensor	shape	entries	$m$	$v$
$\mat{W}^{(1)}$	$784 \times 512$	$401\,408$	same	same
$\vec{b}^{(1)}$	$1 \times 512$	$512$	same	same
$\mat{W}^{(2)}$	$512 \times 256$	$131\,072$	same	same
$\vec{b}^{(2)}$	$1 \times 256$	$256$	same	same
$\mat{W}^{(3)}$	$256 \times 128$	$32\,768$	same	same
$\vec{b}^{(3)}$	$1 \times 128$	$128$	same	same
$\mat{W}^{(4)}$	$128 \times 10$	$1\,280$	same	same
$\vec{b}^{(4)}$	$1 \times 10$	$10$	same	same
total		$567\,434$	$567\,434$	$567\,434$

Two extra copies, $1\,134\,868$ values, $4.54$ MB in fp32 against the model’s own $2.27$ MB — the $2\times$ of I.7.B04, on a concrete model.

Nothing is per-layer or per-tensor: the state is per scalar parameter. That is what “adaptive” means here, and it is the whole reason the state is this large.

Step 2 — rescale the loss. Replace $\mathcal{L}$ by $c\mathcal{L}$ for a constant $c > 0$ . Gradients are linear, so every $g_t$ becomes $cg_t$ . Both moments are linear in their inputs:

m_t = \beta_1 m_{t-1} + (1-\beta_1)cg_t \Longrightarrow cm_t, \qquad v_t = \beta_2 v_{t-1} + (1-\beta_2)c^2g_t^2 \Longrightarrow c^2v_t

The first moment picks up $c$ and the second picks up $c^2$ , because one averages the gradient and the other its square.

Step 3 — the cancellation. Bias correction divides both by constants, so the factors survive it: $\hat{m}_t \to c\hat{m}_t$ and $\hat{v}_t \to c^2\hat{v}_t$ . The step is

\frac{c\,\hat{m}_t}{\sqrt{c^2\hat{v}_t} + \varepsilon} = \frac{c\,\hat{m}_t}{c\sqrt{\hat{v}_t} + \varepsilon}

For $\varepsilon$ negligible against $c\sqrt{\hat{v}_t}$ the $c$ cancels and the update is unchanged. Adam’s step is invariant to rescaling the loss, and by the same argument to rescaling any individual coordinate’s gradient.

Step 4 — what plain descent does instead. Its update is $-\eta cg_t$ , which scales with $c$ directly. Doubling the loss doubles every step, so a rate tuned for one loss scale is wrong for another — and since a sum over a batch and a mean over a batch differ by a factor of $B$ , that is not a hypothetical.

Step 5 — where the invariance fails. $\varepsilon$ is an additive constant and does not scale, so the cancellation is approximate. When $c$ is small enough that $c\sqrt{\hat{v}} \lesssim \varepsilon$ , the denominator is dominated by $\varepsilon$ and the step becomes $\approx c\hat{m}/\varepsilon$ — linear in $c$ again, like plain descent. So the invariance holds over a wide range and breaks at the bottom of it, which is one reason $\varepsilon$ is a number worth knowing rather than a formality: it sets where the adaptive behaviour stops.

Answer

$m$ and $v$ each have the shape of the parameters, tensor for tensor: $567\,434$ entries apiece, $1\,134\,868$ values of state, $4.54$ MB in fp32 against $2.27$ MB of weights.

Under $\mathcal{L} \to c\mathcal{L}$ : $g \to cg$ , $m \to cm$ , $v \to c^2v$ , and the update $\hat{m}/(\sqrt{\hat{v}} + \varepsilon)$ is unchanged while $c\sqrt{\hat{v}} \gg \varepsilon$ . Plain descent’s update scales with $c$ .

Check — sanity

The state has no shape of its own. Every entry of $m$ and $v$ corresponds to exactly one scalar parameter, so a memory estimate that treats optimiser state as a per-layer overhead is wrong by orders of magnitude.

The exponents on $c$ are 1 and 2. They must be: $m$ averages $g$ and $v$ averages $g^2$ . Getting $c$ and $c$ would mean the square had been dropped, and the cancellation would then leave a stray factor rather than none.

The step’s magnitude is about $\eta$ . With $\hat{m} \approx \sqrt{\hat{v}}$ for a steady gradient, the ratio is near 1 and each coordinate moves by roughly $\eta$ per step regardless of its gradient’s size. That is the property that makes a single $\eta$ serviceable across parameters of very different scales, and it is also why Adam’s $\eta$ is typically $10^{-3}$ where SGD’s might be $10^{-1}$ — they are measured in different things.

Where this breaks

Scale invariance is a property of the step, not of training. Rescaling the loss leaves Adam’s trajectory unchanged but not anything coupled to the loss value: gradient clipping by norm triggers at different points, a fixed convergence threshold on the loss means something different, and in mixed precision the gradient scaler exists precisely because the raw magnitude decides whether fp16 overflows or flushes to zero. “Adam does not care about the loss scale” is true of one equation and false of the system around it.

The invariance is also per-coordinate, which is a double-edged property. A parameter whose gradient is genuinely tiny because it barely matters is stepped just as far as one whose gradient is large, since only the ratio survives. That is the intended behaviour and it is also why Adam can keep moving parameters that plain descent would leave alone — sometimes usefully, sometimes as an argument for decoupled weight decay, which I.7.X04 takes up.

Finally, $\varepsilon$ sits outside the square root in (I.7.6) and some implementations place it inside. The two differ, and the difference shows up exactly in the small-gradient regime where the invariance is already breaking down. It is not a detail worth arguing about often, but it is a real difference between two things both called Adam.

Variation

Work out what happens under a rescaling of the parameters rather than the loss — replace $\theta$ by $a\theta$ with the model reparameterised to compensate. Is Adam invariant to that too? The answer is different from Step 3’s, and it is the reason normalisation layers still matter when Adam is in use.

Exercises

Every one has a published solution. A hidden solution is a solution; a missing one is an abandonment.

I.7.X01More momentum is not monotonically betternumeric▲▲△

On the quadratic of I.7.B01 — $\mat{H} = \mathrm{diag}(20, 0.2)$ , $\kappa = 100$ — compare plain descent with momentum at $\beta \in \{0.6694, 0.9, 0.99\}$ . Predict the ordering before looking, then explain the one that surprises you.

Hint

By I.7.X03 the buffer settles at $g/(1-\beta)$ for a steady gradient. What does that do to the effective step size, and what does an effective step size that is too large do?

Solution

The measurements. Steps to reach an error norm of $10^{-3}$ :

method	$\eta$	$\beta$	steps
plain descent	$0.099010$	$0$	$363$
momentum, tuned	$0.165289$	$0.6694$	$\mathbf{58}$
momentum, default	$0.165289$	$0.9000$	$132$
momentum, default $\eta$	$0.099010$	$0.9000$	$127$
momentum, $\beta = 0.99$	$0.165289$	$0.9900$	$\mathbf{1\,214}$

The tuned case behaves as advertised. $58$ steps against $363$ , a factor of $6.3$ , against the asymptotic prediction $\sqrt{\kappa} = 10$ . The shortfall is the transient: momentum needs several steps before its average carries information, and 58 steps is not many multiples of that.

The surprise is the last row. At $\beta = 0.99$ momentum takes $1\,214$ steps — 3.3 times worse than using no momentum at all. More memory made it worse, decisively.

Why. By I.7.X03 a steady gradient drives the buffer to $g/(1-\beta)$ , so the effective step is $\eta/(1-\beta)$ . At $\beta = 0.99$ that is a hundredfold amplification: an $\eta$ of $0.165$ behaves like $16.5$ , far past the stability threshold of $0.1$ . The iteration does not diverge only because the oscillation it induces is itself damped by the averaging — so instead of blowing up, it lurches around the valley, and the lurching is slower than walking.

The optimum is not near the default. For $\kappa = 100$ the tuned $\beta$ is $0.6694$ , well below the $0.9$ everyone writes. The default is optimal for $\kappa \approx 361$ , and the further a problem is from that, the worse it serves — in both directions.

What to take from this. $\beta$ and $\eta$ are not independent knobs, and raising $\beta$ is a learning-rate increase wearing a different name. A configuration search that varies them separately is searching a space where the interesting direction is diagonal. It also explains a common experience: adding momentum to a run tuned without it makes things worse, and the fix is to lower $\eta$ at the same time rather than to conclude that momentum does not help.

I.7.X02What bias correction is correcting, and for how longgradient▲▲△

Feed Adam a constant gradient $g = 0.1$ and record $m$ , $v$ and the step, with and without the bias correction of (I.7.6), for the first ten steps. Say by what factor the uncorrected first step is wrong and in which direction, then find how long the correction still matters.

Hint

Both moments start at zero, so both are biased towards zero — but one of them is under a square root in a denominator. Work out which way that pushes the step.

Solution

The correction works exactly. With a constant gradient the corrected step is $1.000000$ at every step from the first — because $\hat{m} = g$ and $\hat{v} = g^2$ exactly, so the ratio is $g/|g| = 1$ . That is the property the correction is for: the step size should not depend on how long training has been running.

The uncorrected step is wrong, and grows.

$t$	$m$ raw	$v$ raw	step raw	step corrected
1	$0.01000000$	$0.0000100000$	$3.162268$	$1.000000$
2	$0.01900000$	$0.0000199900$	$4.249582$	$1.000000$
5	$0.04095100$	$0.0000499001$	$5.797132$	$1.000000$
10	$0.06513216$	$0.0000995512$	$6.527874$	$1.000000$

The direction is the part worth getting right. The uncorrected step is $3.1623\times$ too large, not too small — which is the opposite of what the phrase “biased towards zero” suggests at first reading.

Both moments are biased towards zero, but by different amounts and in different places. After one step $m$ is $(1-\beta_1)g = 0.1g$ , a tenth of its target, while $v$ is $(1-\beta_2)g^2 = 0.001g^2$ , a thousandth of its target. The step divides by $\sqrt{v}$ , and $\sqrt{0.001} = 0.0316$ , so the denominator is shrunk 31.6-fold while the numerator is shrunk 10-fold. Net: $31.6/10 = 3.16$ times too large.

How long it matters. The correction factors are $1 - \beta^t$ :

$t$	$1 - \beta_1^t$	$1 - \beta_2^t$
1	$0.10000000$	$0.00100000$
10	$0.65132156$	$0.00995512$
100	$0.99997344$	$0.09520785$
1000	$1.00000000$	$0.63230458$
5000	$1.00000000$	$0.99327889$

The first moment’s correction is finished by step 100. The second’s is still correcting by a factor of $1.6$ at step 1000 and only settles around 5000 — because $\beta_2 = 0.999$ has a memory of a thousand steps by I.7.X03.

Why this connects to warm-up. Without the correction the opening steps would be several times too large, which is precisely the damage a warm-up exists to prevent. The correction removes the systematic part of that, and a warm-up is still used — because the variance of the early estimates is high even once their bias is gone, and an estimate built from three samples is unreliable whether or not it is unbiased. Bias correction and warm-up address different halves of the same problem, which is why doing one does not remove the need for the other.

I.7.X03How many gradients a momentum buffer rememberssymbolic▲▲△

Unroll the momentum recurrence $v_t = \beta v_{t-1} + g_t$ into a sum over past gradients, identify the weights, and derive the effective window length. Then say what $\beta = 0.9$ and $\beta = 0.99$ each mean in steps, and what the buffer’s magnitude is for a constant gradient.

Hint

Substitute the recurrence into itself a few times. The coefficients form a sequence you can recognise and sum.

Solution

Unrolling. The recurrence is the first half of (I.7.5). With $v_0 = 0$ :

v_t = g_t + \beta g_{t-1} + \beta^2 g_{t-2} + \dots = \sum_{k=0}^{t-1}\beta^k g_{t-k}

The buffer is a weighted sum of every gradient so far, with weights decaying geometrically. Nothing is forgotten exactly; contributions become negligible.

The window. The natural measure is the mean age of a contribution, weighted by its coefficient:

\bar{k} = \frac{\sum_k k\beta^k}{\sum_k \beta^k} = \frac{\beta/(1-\beta)^2}{1/(1-\beta)} = \frac{\beta}{1-\beta}

so the buffer averages over roughly $1/(1-\beta)$ recent steps — the usual statement, differing from $\bar{k}$ by one.

$\beta$	$1/(1-\beta)$	$\bar{k}$
$0.9$	$10$	$9$
$0.99$	$100$	$99$
$0.999$	$1000$	$999$

Moving from $0.9$ to $0.99$ is not a small adjustment: it multiplies the memory by ten.

Magnitude for a constant gradient. If $g_t = g$ for every step, the sum is geometric:

v_t = g\sum_{k=0}^{t-1}\beta^k = g\,\frac{1 - \beta^t}{1 - \beta} \xrightarrow{t \to \infty} \frac{g}{1-\beta}

The buffer settles at $g/(1-\beta)$ , which is $10g$ at $\beta = 0.9$ . The step is $\eta v$ , so momentum multiplies the effective step size by $1/(1-\beta)$ in any direction where the gradient is consistent. Raising $\beta$ without lowering $\eta$ is a learning-rate increase in disguise, and it is a common way to make a stable configuration diverge while believing only the smoothing changed.

Why the asymmetry helps. In a direction where the gradient oscillates in sign, consecutive terms cancel and $v$ stays small. In a direction where it is consistent, they accumulate to $g/(1-\beta)$ . Momentum damps the first and amplifies the second, which is exactly the discrimination the ill-conditioned problem of I.7.B03 requires — and it is why the improvement there was a factor of fifty rather than a modest smoothing.

The variant worth knowing. Some formulations use $v_t = \beta v_{t-1} + (1-\beta)g_t$ , which normalises the weights to sum to one and removes the $1/(1-\beta)$ gain. PyTorch’s SGD uses the first form and its Adam uses the second, so the same $\beta$ means a different effective step in the two — a fact worth checking before porting a configuration between them.

I.7.X05A batch that is not a samplecounterexample▲▲△

A dataset of 1000 examples is sorted by label — all 500 of class 0, then all 500 of class 1 — and read in batches of 100 without shuffling. Show that the resulting gradient estimate is biased, that no learning rate corrects it, and what the loss curve looks like.

Hint

What is the expectation of the gradient over the batches actually presented, against the expectation over a random batch? Compare each individual batch, not the average over an epoch.

Solution

The unbiasedness that is assumed. With $B$ examples drawn uniformly at random, the estimator of (I.7.2) satisfies $\mathbb{E}[\hat{g}_B] = \nabla\mathcal{L}$ — the batch gradient is the full gradient plus zero-mean noise. Every convergence argument for stochastic descent begins here.

What sorting does. Batches 1 to 5 contain only class 0; batches 6 to 10 only class 1. Batch $k$ ‘s gradient is

\hat{g}_k = \nabla\mathcal{L}_{\text{class }0} \quad (k \le 5), \qquad \hat{g}_k = \nabla\mathcal{L}_{\text{class }1} \quad (k > 5)

Each is the exact gradient of a different objective — the loss on one class — and neither equals $\nabla\mathcal{L}$ . The error is not zero-mean noise around the right answer; it is the right answer to the wrong question, five times in a row.

Why no learning rate corrects it. A scalar multiplies the direction and cannot change it. Every one of the first five steps moves towards the minimiser of the class-0 loss, which for a classifier is roughly “predict class 0 always”. Shrinking $\eta$ makes the journey there slower, not different.

What the curve looks like. Loss falls through batches 1–5 as the model learns to predict class 0, then jumps sharply at batch 6 and falls again as it unlearns that and predicts class 1, then jumps at the start of the next epoch. A sawtooth with period one epoch, its teeth aligned to the label boundary. The average over an epoch is unbiased — the two halves do cancel — but the parameters never sit at that average, and with momentum the buffer arrives at each boundary carrying ten steps of exactly the wrong direction.

Why it is easy to hit. Nothing errors. The data loads, the loss decreases within each run of batches, and the sawtooth is easy to read as “needs a lower learning rate” — which makes it worse, by lengthening the time spent on each side. The fix is one flag, and the diagnosis is the hard part.

The general statement. Any correlation between batch membership and the gradient breaks the assumption: sorting by label, by length, by acquisition date, by hospital site. The last is the one that matters in the medical-imaging work of Book VIII — batches that are accidentally single-site are not samples of the population the model is meant to serve, and the bias survives every learning rate, every optimiser and every schedule.

I.7.X06What happens exactly at the stability thresholdlimit▲▲△

Take the recurrence $e_{t+1} = (1 - \eta\lambda)e_t$ and examine the three boundary cases: $\eta\lambda = 1$ , $\eta\lambda = 2$ , and $\eta\lambda \to 2^-$ . Say what the iteration does in each, and what the third implies about tuning a learning rate by trying larger values until something breaks.

Hint

The multiplier is $1 - \eta\lambda$ . Evaluate it, and remember it is a signed quantity — the sign carries information the magnitude does not.

Solution

$\eta\lambda = 1$ : exact, in one step. The multiplier is $0$ , so $e_1 = 0$ regardless of $e_0$ . This is Newton’s method for a quadratic: the step $\eta = 1/\lambda$ lands on the optimum exactly. It is also the reason a perfectly conditioned problem is trivial — with all eigenvalues equal, one $\eta$ solves every direction at once.

$\eta\lambda = 2$ : perpetual oscillation. The multiplier is $-1$ , so $e_t = (-1)^t e_0$ . The error alternates in sign and never changes magnitude. The iterate bounces between two points forever: the loss is constant, the gradient is non-zero, and nothing is converging or diverging. Neither the loss curve nor the gradient norm reveals progress, because there is none.

$\eta\lambda \to 2^-$ : convergence at an arbitrarily slow rate. The multiplier approaches $-1$ from inside, so the error still contracts — but by a factor approaching $1$ . Steps to reach a target grow without bound as the threshold is approached, while the iteration remains formally convergent throughout.

What that implies about tuning. The common procedure is to raise $\eta$ until training becomes unstable and then back off. This limit says the procedure has a trap: just below the threshold the run is stable and useless, contracting so slowly that it is indistinguishable from a plateau over any finite budget. It looks like a learning-rate problem of the opposite kind — too small — and the natural response of raising $\eta$ further crosses into divergence.

The sign is the tell. Below $\eta\lambda = 1$ the multiplier is positive and the error decays monotonically. Above it the multiplier is negative and the error alternates, so the parameter oscillates around the optimum while approaching it. A trajectory that zigzags is already past $1/\lambda$ in that direction and is in the upper half of the stable band — which is visible in a parameter trace and invisible in a loss curve.

Why the loss hides it. For a quadratic the loss is $\tfrac{1}{2}\lambda e^2$ , which is even in $e$ . Both signs give the same loss, so the alternation that distinguishes the two regimes is exactly what the loss discards. This is a small instance of a general point: a scalar summary of a high-dimensional trajectory loses the information needed to diagnose it, and monitoring only the loss is monitoring the one quantity engineered to be flat.

I.7.X07One bad gradient, and every step after itcounterexample▲▲△

A single batch produces an infinite gradient in one coordinate. Trace what Adam does on that step and on every subsequent one, and explain why the loss goes to NaN several steps after the batch responsible. Then say why skipping the bad step is not sufficient.

Hint

Follow the value through (I.7.6), and remember that $m$ and $v$ persist between steps while the gradient does not.

Solution

The step itself. With $g = \infty$ in one coordinate:

m \leftarrow \beta_1 m + (1-\beta_1)\infty = \infty, \qquad v \leftarrow \beta_2 v + (1-\beta_2)\infty^2 = \infty

The update is $\hat{m}/(\sqrt{\hat{v}} + \varepsilon) = \infty/\infty =$ NaN, and that coordinate’s parameter becomes NaN.

Why it spreads. A NaN parameter makes every activation downstream of it NaN, so the next forward pass produces a NaN loss and the next backward pass NaN gradients everywhere. On the following step every $m$ and $v$ in the model becomes NaN, and every parameter follows. One coordinate at step $t$ is the whole model by step $t + 2$ .

Why the timing misleads. The observable is the loss, and the loss is a scalar computed from the forward pass. At step $t$ the gradient is infinite but the parameters are still finite, so the loss printed for step $t$ looks ordinary. The NaN appears in the loss at step $t + 1$ or later, and the batch on screen when it appears is not the batch responsible. People go looking for corrupt data in the wrong file.

Why skipping is not sufficient. The natural fix — detect a non-finite gradient and skip the update — leaves $m$ and $v$ untouched only if the skip happens before they are written. If the moments have already been updated, they are non-finite and persist: the next step divides by a NaN $\sqrt{\hat{v}}$ even though its gradient is perfectly good. The optimiser state must be repaired or reset, not merely the step abandoned. This is the sense in which the failure outlives its cause, and it is what distinguishes this from an ordinary bad batch.

What actually protects against it. Gradient clipping bounds the magnitude before it reaches the moments, which is the cheap general defence. In mixed precision a loss scaler both provokes and manages the problem: it multiplies the loss to keep fp16 gradients out of the underflow range, detects the resulting overflows, skips the step, and lowers the scale — the last part being the piece that makes it a control loop rather than a filter.

The assumption this tests. The chapter assumes the gradient is finite along the path taken. That is not a mathematical nicety: it is the condition under which the optimiser’s state remains a summary of past gradients rather than an absorbing state that no future gradient can leave.

I.7.X08Square-root returns, and where they stop being worth buyingcomplexity▲▲△

Price what a larger batch buys. Give the noise, the arithmetic and the noise per unit of arithmetic across $B$ from 1 to 4096, then identify the regime in which raising $B$ stops being a good trade and say what the budget does and does not decide.

Hint

By (I.7.2) noise falls as $1/\sqrt{B}$ while cost rises as $B$ . Form the product and read off which direction it moves.

Solution

The trade.

$B$	standard error	vs $B{=}1$	arithmetic	error $\times$ cost
1	$1.000000$	$1.0000$	1	$1.00$
4	$0.500000$	$0.5000$	4	$2.00$
16	$0.250000$	$0.2500$	16	$4.00$
64	$0.125000$	$0.1250$	64	$8.00$
256	$0.062500$	$0.0625$	256	$16.00$
1024	$0.031250$	$0.0312$	1024	$32.00$
4096	$0.015625$	$0.0156$	4096	$64.00$

Quadrupling $B$ halves the noise and costs four times the arithmetic, at every scale. The last column — noise per unit of compute — doubles each time, which is the trade stated in the least flattering way that is still true: measured per FLOP, a large batch is strictly worse at reducing noise.

Why anyone uses large batches anyway. The table prices arithmetic, and arithmetic is not the binding constraint on modern hardware. A batch of 256 does not take 256 times as long as a batch of 1 — it takes barely longer, because the device was idle. The square-root law is about statistical efficiency; large batches win on hardware efficiency, and the two pull in opposite directions.

Where it stops. Comparing noise against the signal it corrupts:

$B$	noise/signal	regime
1	$1.000000$	noise-dominated: more $B$ helps a lot
16	$0.250000$	mixed: linear scaling (I.7.7) roughly holds
256	$0.062500$	mixed: linear scaling roughly holds
4096	$0.015625$	signal-dominated: more $B$ buys little
65536	$0.003906$	signal-dominated: more $B$ buys little

Once the noise is small against the gradient, further averaging refines a number that was already accurate enough to step with. That boundary is McCandlish and colleagues’ critical batch size, and the useful thing about it is that it is a property of the problem at a point in training, not of the hardware — it generally grows as training proceeds and the gradient shrinks.

What the budget does not decide. With a fixed budget of $10^6$ example-visits:

$B$	steps	noise per step
32	$31\,250$	$0.17678$
256	$3\,906$	$0.06250$
2048	$488$	$0.02210$

Same compute, three different training runs: many noisy steps or few clean ones. The budget constrains their product and says nothing about which is better — that is decided by where the problem sits relative to its critical batch, and by whether the noise is doing anything useful, which by I.7.P02 it sometimes is.

I.7.X09Measuring the 1/B law instead of assuming itprobability▲▲△

Build a population of per-example gradients with known mean and variance, draw four thousand minibatches at each of six sizes, and compare the measured variance of the batch mean against $\sigma^2/B$ . Check unbiasedness at the same time, and say why the population was given a deliberately non-normal shape.

Hint

Two claims are being tested, and they need different evidence: that the mean of the batch means sits at $\mu$ , and that their spread falls as $1/B$ .

Solution

The population. $N = 10\,000$ values with $\mu = -0.000960$ and $\sigma^2 = 7.859247$ . The shape is a sawtooth with occasional large spikes — deliberately not a bell curve, so that agreement cannot be an artefact of the per-example gradients happening to be normal. Real per-example gradients are not normal either; they are heavy-tailed, with a few examples contributing far more than the rest.

The measurements. 4000 batches at each size:

$B$	measured $\operatorname{Var}$	$\sigma^2/B$	ratio	mean of means
1	$7.845540$	$7.859247$	$0.998$	$0.087900$
2	$4.076844$	$3.929624$	$1.037$	$0.017487$
8	$0.998829$	$0.982406$	$1.017$	$-0.012431$
32	$0.248467$	$0.245601$	$1.012$	$0.020018$
128	$0.061120$	$0.061400$	$0.995$	$0.007109$
512	$0.014889$	$0.015350$	$0.970$	$0.002993$

Unbiasedness. Every “mean of means” sits near $\mu = -0.000960$ , and the agreement tightens as $B$ grows — not because the estimator becomes less biased, which it never was, but because the estimate of its mean becomes more precise. Unbiasedness is a property at every $B$ ; only our ability to see it improves.

The $1/B$ law. Ratios between $0.970$ and $1.037$ across three orders of magnitude in $B$ . The variance falls by a factor of $512$ from $B{=}1$ to $B{=}512$ , matching the prediction to within four percent at the worst point.

Where the residual error comes from. The ratios wander by a few percent because 4000 trials estimate a variance with a standard error of roughly $\sqrt{2/(4000-1)} \approx 2.2\%$ — so the observed spread is what sampling noise alone predicts. Measuring a variance is itself a noisy business, and quoting three decimal places on the ratio would be quoting the noise.

Why the non-normal shape matters. The derivation in I.7.X10 uses only linearity of expectation and independence — nothing about the distribution’s shape. A demonstration on Gaussian data would be consistent with a much weaker claim that happened to need normality. Using a spiky population tests the theorem that was actually proved.

What this does not establish. The law assumes independent draws. Sorted batches (I.7.X05), duplicated examples, or any correlation within a batch break the variance half while leaving the unbiasedness half intact — and no measurement of this kind detects that, because the sampling here was constructed to be independent. The experiment confirms the theorem under its hypotheses; it says nothing about a pipeline that violates them.

I.7.X10Why a minibatch gradient is the right one on averageproof▲▲△

Prove that the minibatch gradient is an unbiased estimator of the full-batch gradient, and derive its variance. State precisely which property of the sampling each half of the result requires, so it is clear what breaks when the sampling is wrong.

Hint

Linearity of expectation gives the first half and needs almost nothing. Independence is what the second half needs, and it is a strictly stronger assumption.

Solution

Setup. The full objective is the mean over $N$ examples, $\mathcal{L}(\theta) = \frac{1}{N}\sum_{i=1}^{N}\ell_i(\theta)$ , so $\nabla\mathcal{L} = \frac{1}{N}\sum_i \nabla\ell_i$ . A batch is $B$ indices $i_1, \dots, i_B$ drawn uniformly at random with replacement, and

\hat{g}_B = \frac{1}{B}\sum_{j=1}^{B}\nabla\ell_{i_j}

Unbiasedness. For a single uniformly drawn index,

\mathbb{E}[\nabla\ell_{i}] = \sum_{k=1}^{N}\frac{1}{N}\nabla\ell_k = \nabla\mathcal{L}

By linearity of expectation — which requires no independence at all —

\mathbb{E}[\hat{g}_B] = \frac{1}{B}\sum_{j=1}^{B}\mathbb{E}[\nabla\ell_{i_j}] = \nabla\mathcal{L}

The only property used is that each index is marginally uniform. This is why unbiasedness survives sampling without replacement, and why it fails under the sorted batches of I.7.X05 — there the marginal distribution of an index within batch 1 is uniform over class 0 only.

Variance. Write $\Sigma = \operatorname{Var}[\nabla\ell_i]$ for the per-example covariance. Now independence is needed: for independent draws the variance of a sum is the sum of the variances, so

\operatorname{Var}[\hat{g}_B] = \frac{1}{B^2}\sum_{j=1}^{B}\Sigma = \frac{\Sigma}{B}

which is (I.7.2). The standard error is $\sigma/\sqrt{B}$ .

What each half needs. Unbiasedness needs marginal uniformity. Variance needs independence between the draws. The gap between them is the whole reason correlated batches are dangerous: a batch of near-duplicate examples can still be unbiased while having a variance far above $\Sigma/B$ , so the $1/B$ accounting overstates how much noise the batch removed. Deduplication is a variance intervention, not only a data-hygiene one.

Sampling without replacement. Real training shuffles and takes disjoint batches, which is sampling without replacement. Unbiasedness is unaffected. The variance gains a finite-population correction $\frac{N-B}{N-1}$ , which is negligible while $B \ll N$ and becomes decisive as $B \to N$ — at $B = N$ the variance is zero, because the batch is the dataset. The $1/B$ law is the large- $N$ limit, and full-batch descent is its other end.

The consequence for practice. Halving the noise costs four times the arithmetic, so the returns are square-root. That single fact is why batch size is a trade rather than a setting to maximise, and why the critical batch size of McCandlish and colleagues exists: past the point where the gradient is already accurate, additional examples buy accuracy nobody needed.

I.7.X04Two things both called weight decay, which differ under Adamcounterexample▲▲▲

Adding $\tfrac{\lambda}{2}\|\theta\|^2$ to the loss and subtracting $\eta\lambda\theta$ from each update are the same operation under plain gradient descent. Show they are the same, then show they are not the same under Adam, and say which coordinates the difference falls on.

Hint

Under Adam the penalty’s gradient does not go straight into the update — it passes through the same division by $\sqrt{\hat{v}}$ as everything else.

Solution

Under plain descent they coincide. The penalty contributes $\nabla(\tfrac{\lambda}{2}\|\theta\|^2) = \lambda\theta$ , so

\theta \leftarrow \theta - \eta(g + \lambda\theta) = \theta - \eta g - \eta\lambda\theta

which is the update with a separate $-\eta\lambda\theta$ term. The two descriptions are one algorithm, which is why the names are used interchangeably.

Under Adam they diverge. L2 regularisation puts $\lambda\theta$ into $g$ before the moments are formed, so it is divided by $\sqrt{\hat{v}}$ along with everything else:

\theta \leftarrow \theta - \eta\,\frac{\widehat{m[g + \lambda\theta]}}{\sqrt{\widehat{v[g + \lambda\theta]}} + \varepsilon}

Decoupled decay (AdamW) applies it outside:

\theta \leftarrow \theta - \eta\,\frac{\hat{m}[g]}{\sqrt{\hat{v}[g]} + \varepsilon} - \eta\lambda\theta

Which coordinates differ. Consider a parameter whose gradient is persistently large, so $\sqrt{\hat{v}}$ is large. Under L2 its decay term is divided by that large number and is therefore weak. A parameter with a small gradient has small $\sqrt{\hat{v}}$ and its decay is strong.

So L2-under-Adam regularises inversely to gradient magnitude: the parameters doing the most work are penalised least, and the ones already doing little are pushed towards zero hardest. That is the opposite of what a practitioner adding weight decay intends, and it is not a small effect — $\sqrt{\hat{v}}$ can differ by orders of magnitude across a network.

Decoupled decay applies $\eta\lambda\theta$ uniformly, which is what the name describes.

Why this counts as a counterexample. The claim being refuted is the reasonable one that weight decay and L2 are two names for one thing. They are, under any optimiser whose update is proportional to the gradient — SGD, momentum, Nesterov. They are not, under any optimiser that rescales the gradient per coordinate, and Adam is the one everybody uses. The identity fails exactly where it is most relied on.

The practical residue. Loshchilov and Hutter’s AdamW is the decoupled form, and it is now the default in most training code. A configuration ported from an SGD recipe by copying $\lambda$ across will not reproduce its regularisation under Adam, and the discrepancy is largest for the parameters that matter most.