Optimizers and learning-rate schedules

Mini-batch SGD, the baseline

With a mini-batch $\mathcal{B}$ of $B$ examples, the gradient estimate $\hat{\mathbf{g}} = \frac{1}{B}\sum_{i\in\mathcal{B}}\nabla_\theta\ell_i$ is unbiased for the full gradient and has variance proportional to $1/B$ (1.10 Optimization: from gradient descent to Adam). The update is $\theta_{t} = \theta_{t-1} - \eta\,\hat{\mathbf{g}}_t$. Two problems drive everything else in this lesson: the gradient is noisy, and its scale differs wildly between parameters and over training — a single global step size $\eta$ is a poor fit for both.

Momentum and Nesterov

Momentum accumulates an exponentially-weighted average of past gradients and moves along it:

with $\beta\approx0.9$ (this is PyTorch's convention; some texts write $\mathbf{v}_t=\beta\mathbf{v}_{t-1}+(1-\beta)\hat{\mathbf{g}}_t$, which just rescales $\eta$). If the gradient points the same way for many steps, $\mathbf{v}$ grows toward $\hat{\mathbf{g}}/(1-\beta)$ — a $10\times$ larger effective step for $\beta=0.9$ — while components that flip sign cancel. In a narrow ravine, momentum damps the side-to-side oscillation and accelerates along the valley floor, which is why it converges in far fewer steps on ill-conditioned problems.

Nesterov momentum evaluates the gradient at the look-ahead point $\theta_{t-1}-\eta\beta\mathbf{v}_{t-1}$ instead of at $\theta_{t-1}$: it peeks where momentum is about to carry it and corrects early. In practice it is a small, consistent improvement for SGD on vision models.

Adaptive per-parameter step sizes

AdaGrad divides each coordinate's step by the root of its accumulated squared gradients, $\mathbf{s}_t = \mathbf{s}_{t-1}+\hat{\mathbf{g}}_t^{2}$, $\theta_t = \theta_{t-1} - \eta\,\hat{\mathbf{g}}_t/(\sqrt{\mathbf{s}_t}+\epsilon)$ (all operations element-wise). Rarely-updated parameters (rare words in an embedding table) get large steps, frequent ones get small steps — but $\mathbf{s}_t$ only grows, so the learning rate decays to zero whether or not you are done.

RMSProp fixes that with an exponential average, $\mathbf{s}_t = \rho\,\mathbf{s}_{t-1} + (1-\rho)\,\hat{\mathbf{g}}_t^{2}$, so the denominator tracks the recent gradient magnitude.

Adam: momentum + RMSProp + bias correction

Adam (Kingma & Ba, 2015) keeps both a first-moment and a second-moment estimate and corrects their start-up bias:

with defaults $\beta_1=0.9$, $\beta_2=0.999$, $\epsilon=10^{-8}$ (LLM recipes use $\beta_2=0.95$ for stability).

Deriving the bias correction. Unroll the first-moment recursion from $\mathbf{m}_0=\mathbf{0}$: $\mathbf{m}_t = (1-\beta_1)\sum_{i=1}^{t}\beta_1^{\,t-i}\,\hat{\mathbf{g}}_i$. If the gradients have a common mean $\mathbb{E}[\hat{\mathbf{g}}_i]=\mathbf{g}$, then $\mathbb{E}[\mathbf{m}_t] = \mathbf{g}\,(1-\beta_1)\sum_{i=1}^{t}\beta_1^{\,t-i} = \mathbf{g}\,(1-\beta_1^{\,t})$ by the geometric series. So $\mathbf{m}_t$ underestimates the true mean by the factor $(1-\beta_1^{t})$, which starts at $0.1$ and tends to $1$; dividing by it removes the bias. The same argument gives $(1-\beta_2^{t})$ for the second moment.

Math check — why it matters. At $t=1$: $\mathbf{m}_1 = 0.1\,\hat{\mathbf{g}}_1$ and $\mathbf{v}_1 = 0.001\,\hat{\mathbf{g}}_1^{2}$. Without correction the first step would be $\eta\cdot 0.1\,\hat g/\sqrt{0.001\,\hat g^{2}} = \eta\cdot 3.16\,\mathrm{sign}(\hat g)$ — three times the intended step, in every coordinate, on the very first iteration. With correction, $\hat{\mathbf{m}}_1=\hat{\mathbf{g}}_1$, $\hat{\mathbf{v}}_1=\hat{\mathbf{g}}_1^{2}$ and the step is $\eta\,\mathrm{sign}(\hat g)$ as designed.

Two consequences of the form of the update. First, each coordinate's step is roughly $\eta\times(\text{mean gradient})/(\text{RMS gradient})$, a number of order $\eta$ regardless of the gradient's scale — Adam is invariant to rescaling the loss, and $\eta$ directly controls how far parameters move per step (typical: $10^{-4}$ to $10^{-3}$; $3\times10^{-4}$ is the folk default). Second, the update is close to a sign-based method in early training, which is why warmup is needed (below).

Weight decay: AdamW versus L2

Adding an L2 penalty $\frac{\lambda}{2}\|\theta\|^2$ to the loss adds $\lambda\theta$ to the gradient. Under plain SGD that is identical to weight decay, $\theta_t = (1-\eta\lambda)\,\theta_{t-1} - \eta\hat{\mathbf{g}}_t$: shrink, then step. Under Adam it is not identical: the penalty term enters $\mathbf{m}_t$ and $\mathbf{v}_t$ and gets divided by $\sqrt{\hat{\mathbf{v}}_t}$, so parameters with large gradients are decayed less — the opposite of what you want. AdamW (Loshchilov & Hutter, 2019) decouples the two:

It is the default optimizer for transformers; typical $\lambda$ is $0.01$ for fine-tuning and $0.1$ for LLM pretraining, usually not applied to biases and normalization parameters. Note that PyTorch's torch.optim.Adam(weight_decay=…) implements the coupled L2 version — use AdamW.

Learning-rate schedules

The schedule usually matters more than the choice between Adam variants.

• Warmup. Increase $\eta$ linearly from ~0 to $\eta_{\max}$ over the first few hundred to few thousand steps. Early on, Adam's $\hat{\mathbf{v}}$ is a poor estimate from few samples and the loss landscape near a random initialization is sharp; a full-size step there can throw the model into a bad region it never recovers from. Warmup is near-universal for transformers.
• Cosine decay to a floor: $\eta_t = \eta_{\min} + \tfrac12(\eta_{\max}-\eta_{\min})\big(1+\cos(\pi t/T)\big)$. Smooth, one hyperparameter, the standard for pretraining; the related warm restarts (SGDR) repeat the cosine cycle.
• Step decay (divide by 10 at fixed epochs) is the classic vision recipe; one-cycle (Smith) warms up to a high peak then anneals, often converging in far fewer epochs; warmup–stable–decay (WSD) holds $\eta$ constant then decays sharply at the end, letting you stop training at any time by running the decay phase — popular for LLMs since 2024.
• Learning-rate range test. Increase $\eta$ geometrically each step for a few hundred iterations and plot loss against $\eta$: the loss falls, flattens, then explodes. Pick $\eta_{\max}$ roughly an order of magnitude below the explosion, where the curve is steepest. Five minutes of compute replaces a grid search.

Gradient clipping

Clip the global norm of the gradient: $\hat{\mathbf{g}}\leftarrow\hat{\mathbf{g}}\cdot\min\!\big(1,\ c/\|\hat{\mathbf{g}}\|_2\big)$ with $c=1.0$ typical. It leaves most steps untouched and caps the rare huge gradient from a bad batch that would otherwise produce a loss spike or NaN (4.8 Training dynamics and debugging). Logging the pre-clip norm is one of the most useful training diagnostics.

Batch size and the linear scaling rule

Doubling the batch halves the gradient variance, so you can afford a bigger step. For SGD the linear scaling rule (Goyal et al., 2017) — multiply $\eta$ by $k$ when the batch grows by $k$, with warmup — reproduces small-batch results up to a critical batch size, beyond which extra samples per step stop buying faster convergence (McCandlish et al.'s gradient noise scale predicts where). For Adam the empirical scaling is closer to $\sqrt{k}$. Rule of thumb: tune $\eta$ at your actual batch size, and when you must change the batch, re-tune or scale, never keep $\eta$ fixed blindly.

Landscape intuition

In millions of dimensions, bad local minima are rare; saddle points and long flat plateaus are the obstacles, and momentum and adaptive steps are what get you through them. SGD noise tends to settle in flat minima, which generalize better than sharp ones — one reason small-batch SGD sometimes beats large-batch training at equal steps and why the LR decay phase (noise reduction) produces the final drop in loss (4.7 Regularization in deep nets).

Practical defaults

Where this shows up later

• The schedule and clipping values you will see in nanoGPT and GPT-2 reproduction — 7.8 Build GPT from scratch (nanoGPT walkthrough); AdamW's memory cost (two extra copies of every parameter) in the VRAM budget — 4.12 Compute, hardware and scale basics.
• Reading loss curves to tell "LR too high" from "LR too low" — 4.8 Training dynamics and debugging; weight decay as a regularizer — 4.7 Regularization in deep nets.
• Discriminative learning rates for fine-tuning pretrained backbones — 5.4 Transfer learning and training CNNs in practice; much smaller learning rates for LoRA and full fine-tuning of LLMs — 8.2 Fine-tuning and PEFT; PPO's clipped objective as a trust-region idea — 4.11 Reinforcement learning essentials.