Learning-rate schedules and batch size

The convex last-iterate bound has the shape of a WSD loss curve. For SGD $x_{t+1}=x_t-\gamma\eta_tg_t$ on a convex $f$, with base rate $\gamma$, schedule shape $\eta_t$, $D=\|x_1-x_\star\|$ and $\mathbb E\|g_t\|^2\le G_t^2$, the bound of Defazio et al. reads $$\mathbb E[f(x_T)-f(x_\star)]\le\frac{D^2+\gamma^2\sum_{t\le T}\eta_t^2G_t^2}{2\gamma\sum_{t\le T}\eta_t}+\frac\gamma2\sum_{k=1}^{T-1}\frac{\eta_k}{\sum_{t=k+1}^T\eta_t}\cdot\frac{\sum_{t=k}^T\eta_t^2G_t^2}{\sum_{t=k}^T\eta_t}.$$ The first term is the classical bound for the averaged iterate; the sum is the extra error from reporting the last one. The right side is $\mathcal T_1/\gamma+\gamma\mathcal T_2$, so $\gamma^\star=\sqrt{\mathcal T_1/\mathcal T_2}$; truncated at step $t$, it is the curve compared with loss curves. With constant $G$, a constant schedule gives $\frac{DG}{\sqrt T}\sqrt{1+H_{T-1}}$ ($H_n$ harmonic), while WSD with cooldown fraction $c$ gives at most $\frac{2DG}{\sqrt{(2-c)T}}\sqrt{\kappa+\ln\frac{2-c}{c}}$ with $\kappa$ bounded, so there is no $\ln T$ term. At $T=10^5$ (my computation, $G=D=1$), $\sqrt T$ times the bound is 3.62 for constant, 2.16 for WSD with 20% cooldown, 2.06 for cosine and 2.00 for linear decay.

An online-gradient inequality holds against any reference point fixed by $x_k$. Expand $\|x_{t+1}-u\|^2$, use convexity in the cross term, take conditional expectations and sum from $k$ to $T$ (with $\gamma$ absorbed into $\eta$): $$\sum_{t=k}^T\eta_t\,\mathbb E[f(x_t)-f(u)]\le\tfrac12\mathbb E\|x_k-u\|^2+\tfrac12\sum_{t=k}^T\eta_t^2\,\mathbb E\|g_t\|^2.$$ This is Zinkevich's online-gradient inequality. The reference point $u$ may be random, as long as it is determined by $x_k$: the step-$t$ gradient is unbiased given the past, so conditioning still turns the cross term into $f(x_t)-f(u)$. With $u=x_\star$ and $k=1$ it is the averaged-iterate bound, the first term of Step 1.

The last iterate is the average plus one correction per suffix (the steps from $k$ to $T$), each taking the suffix's first iterate as reference point. For any sequence $q_t$ and positive weights $w_t$, a telescoping identity (Defazio et al.) gives exactly $$q_T=\frac{\sum_{t\le T}w_tq_t}{\sum_{t\le T}w_t}+\sum_{k=1}^{T-1}\frac{w_k}{\sum_{t>k}w_t}\cdot\frac{\sum_{t\ge k}w_t(q_t-q_k)}{\sum_{t\ge k}w_t}.$$ Take $q_t=\mathbb E[f(x_t)-f(x_\star)]$ and $w=\eta$. The first piece is Step 2 at $x_\star$. Each correction $\sum_{t\ge k}\eta_t\mathbb E[f(x_t)-f(x_k)]$ is Step 2 with $u=x_k$, where the distance term vanishes and only the suffix's noise $\frac12\sum_{t\ge k}\eta_t^2G_t^2$ is left, divided by its total step. This completes the proof.

The sharp drop at the cooldown comes from the last-iterate term. With a constant step the $k$-th correction is $G^2/(T-k)$: short suffixes take full steps and do no averaging, and the corrections sum to $G^2H_{T-1}\approx G^2\ln T$. Under WSD a suffix starting inside the cooldown has steps shrinking to zero and a correction of about $\frac4{3cT}$, a constant over all $cT$ of them, and the stable-phase suffixes add at most a constant plus about $\ln\frac{2-c}{c}$, a logarithm of $c$ alone. At fixed $\gamma$, during the stable phase the distance term falls like $1/t$ while the last-iterate term grows like $\gamma G^2\ln t$, so the bound flattens; once the step shrinks, the last-iterate term falls quickly. In my decomposition of the exact bound ($T=400$, 20% cooldown), over the cooldown the last-iterate term falls from 0.081 to 0.042 while the other terms barely move. The drop needs gradients that do not vanish (with $G_t\propto t^{-1/2}$ or $t^{-1}$ it disappears), and the worst-case lower bound has the same shape, so the drop is not an artifact of the bounding.

Transferring a learning rate across schedules and horizons. With constant $G$, $\gamma^\star$ is $D/G$ times a function of the schedule shape alone, so the unknown $D$ and $G$ cancel in ratios, and only the ratios transfer. The bound reproduces the $1/\sqrt T$ scaling of the best rate and the factor of about 2 between cosine and WSD ($0.92\,T^{-0.49}$ against $0.47\,T^{-0.49}$) found in LLM sweeps. Retuned per cooldown fraction, longer cooldowns are better all the way to linear decay, since $\gamma^\star$ grows with the fraction and compensates for the smaller steps; at fixed $\gamma$ the best fraction is interior. From 20% cooldown to linear decay it predicts $\gamma^\star(1)\approx e^{0.7}\gamma^\star(0.2)$ (my check: $e^{0.67}$), which a 124M grid confirms. To extend a WSD run, keep the stable phase and lower the rate at the old cooldown start by $\rho=0.525$ for twice the horizon or 0.375 for four times, about 0.01 better than continuing at the same rate.

A second-order expansion of one noisy step gives the noise scale. Expand the loss after $\theta\mapsto\theta-\epsilon G_{est}$ to second order and average over the batch, using $\mathbb E\,G_{est}G_{est}^\top=GG^\top+\Sigma/B$ ($G_{est}$ has mean $G$ and covariance $\Sigma/B$): $$\mathbb E\,L(\theta-\epsilon G_{est})=L(\theta)-\epsilon|G|^2+\frac{\epsilon^2}2\Big(G^\top HG+\frac{\operatorname{tr}(H\Sigma)}B\Big).$$ The estimate is unbiased, so the noise enters only through the curvature-weighted variance in the quadratic term, which lowers the safe step and with it the gain. Minimizing over $\epsilon$ gives $\epsilon_{opt}(B)=\epsilon_{max}/(1+\mathcal B_{noise}/B)$ and $\Delta L_{opt}(B)=\Delta L_{max}/(1+\mathcal B_{noise}/B)$, with $\mathcal B_{noise}=\operatorname{tr}(H\Sigma)/G^\top HG$ and the noiseless values $\epsilon_{max}=|G|^2/G^\top HG$ and $\Delta L_{max}=\frac12|G|^4/G^\top HG$. Below $\mathcal B_{noise}$ the gain is linear in $B$, so $k$ small steps match one $k$ times larger; above it the gain saturates, and linear learning-rate scaling breaks down. With $H\propto I$ this becomes the simple noise scale $\mathcal B_{simple}=\operatorname{tr}\Sigma/|G|^2$, the batch at which the gradient's noise equals its signal.

Steps trade against examples on a hyperbola. A batch $B$ needs $1+\mathcal B/B$ steps and $B+\mathcal B$ examples per unit of progress, so reaching a fixed loss in $S$ steps and $E$ examples obeys $$\Big(\frac S{S_{min}}-1\Big)\Big(\frac E{E_{min}}-1\Big)=1,\qquad\mathcal B_{crit}=\frac{E_{min}}{S_{min}}\approx\mathcal B_{noise},$$ and training at $\mathcal B_{crit}$ costs twice the minimum of each. If the noise scale varies along the run, the best schedule keeps $B^2/\mathcal B$ at a constant $r$, the number of examples one step is worth, so $B(s)=\sqrt{r\,\mathcal B(s)}$, and the square root is why adaptive batches gain little unless the noise scale varies a lot (about 4% predicted on SVHN). Across eight tasks, $\mathcal B_{simple}$ tracks the measured $\mathcal B_{crit}$ within a factor of about 20 while both span six orders of magnitude and grow during a run. It scales like $B/\epsilon$ (a run with a larger step per batch size fluctuates higher up the walls of the loss valley, where $|G|$ is larger), so it must be read off a well-tuned run, and for Adam the tuned rate rises like $B^\alpha$, $\alpha\in[0.5,1]$, then flattens. For Adam-trained OLMo, Merrill et al. (2025) find it underestimates the directly measured critical batch by orders of magnitude, though the shape (small, rising, then flat) survives.

The speedrun's batch and schedule records. On fixed hardware, time is roughly a cost per token plus a cost per step; with $r$ the per-step cost in tokens, Step 7 makes time proportional to $r(1+\mathcal B/B)+(B+\mathcal B)$, minimized at $B^\star=\sqrt{r\,\mathcal B_{noise}}$, the paper's square-root rule with $r$ set by the hardware (my reading, not from any record). Three kinds of record fit it: smaller batches (record 21, 64K to 48K tokens per GPU, about 3.5 s; again at record 39), which help whenever the old batch sat above $B^\star$; a batch ramp as the noise scale grows (record 46, 65 steps and about 1.8 s, though that record also moved to a newer PyTorch nightly, so part may be the upgrade); and Adam stepping every other step for the embeddings and head (record 39), a larger batch for rows with their own noise scale. On the schedule side, record 2 adopted the trapezoid along with other changes (45 to 31.4 minutes, not separable), and later records lengthened the cooldown (0.45 of the run at record 26, 0.80 at record 92), as the bound predicts once the rate is retuned. The speedrun decays to a floor of 0.1 to 0.2 of peak (with a weight EMA over the last ~300 steps in ANVIL2), where the bound favors zero; averaging the tail is another way to cancel the last-iterate penalty, a reading the paper does not test.

Consequences. Neither theory covers the speedrun's optimizers (both are for SGD), the noise-scale derivation is local, greedy and quadratic, and only ratios and shapes transfer.