Optimizers & Schedules

Backpropagation gives, for every weight, the direction in which the loss grows. The optimizer decides how far to move each weight the other way. Every LLM (large language model) in this roadmap is trained with the same recipe:

• the AdamW optimizer (Adam with decoupled weight decay);
• a learning rate that first grows (warmup), then shrinks along a cosine curve;
• gradients clipped to a maximum size.

This page builds each piece, starting from plain gradient descent, with one small example per idea.

Gradient descent

The gradient points uphill, towards a larger loss. Gradient descent takes a small step the other way:

where:

• $\theta$ (theta) holds all the model's weights;
• $\mathcal{L}(\theta)$ is the loss, a single number that measures how wrong the model is;
• $\nabla_\theta \mathcal{L}(\theta)$ is the gradient: for each weight, how fast the loss grows when that weight grows. From now on we write it $g$;
• $\eta$ (eta) is the learning rate, a small positive number that sets the size of the step;
• the arrow $\leftarrow$ means "replace the old value with the new one".

Example. One weight, loss $\mathcal{L} = \theta^2$ (smallest at $\theta = 0$), gradient $g = 2\theta$, starting at $\theta = 3$.

• With $\eta = 0.1$: $3 \to 3 - 0.1 \cdot 6 = 2.4 \to 1.92 \to 1.536$. Each step gets closer to 0.
• With $\eta = 1.1$: $3 \to 3 - 1.1 \cdot 6 = -3.6 \to 4.32 \to -5.184$. Each step overshoots the minimum and lands further away: training diverges.

The learning rate is the most important number in training: too small and progress is slow, too large and the loss explodes.

SGD: one mini-batch at a time

Computing the gradient on the whole dataset at every step would be far too slow. Instead, each step uses a mini-batch, a small random sample of examples. Its gradient is noisy, but right on average. This is SGD (stochastic gradient descent).

Robbins and Monro (1951) showed when such noisy steps still reach the minimum: the step sizes $\eta_t$ must shrink over time, but not too fast. In symbols, $\sum_t \eta_t = \infty$ (the steps add up to an unlimited distance, so we can get anywhere) and $\sum_t \eta_t^2 < \infty$ (they shrink fast enough for the noise to die out). This is the first reason why learning rates decrease during training.

Momentum

Plain SGD follows every bump of the noisy gradient. Momentum (Polyak, 1964) smooths the path: it keeps a running direction, the velocity, and steps along it. In the form PyTorch uses:

where:

• $v$ is the velocity, one number per weight, starting at 0;
• $g$ is the gradient of the current mini-batch;
• $\mu$ (mu) is the momentum coefficient, usually $0.9$: how much of the old velocity is kept.

In words: the new velocity is 90% of the old one plus the new gradient, and the weights move along the velocity.

Example. $\mu = 0.9$, two different gradient histories.

• Steady direction, $g = 1$ at every step: the velocity goes $1,\ 1.9,\ 2.71,\ 3.44,\ 4.10, \dots$ and approaches $g / (1 - \mu) = 10$. The steps become 10 times larger than plain SGD's.
• Noise that flips sign, $g = 1, -1, 1, -1, \dots$: the velocity goes $1,\ -0.1,\ 0.91,\ -0.18,\ 0.84,\ -0.25, \dots$ and settles around $\pm 0.53$. The back-and-forth mostly cancels out.

So momentum speeds up along directions the gradient agrees on, and damps the directions where it keeps changing its mind.

Sutskever et al. (2013) write the same idea with the learning rate inside the velocity: $v \leftarrow \mu v + \eta g$, then $\theta \leftarrow \theta - v$. The two forms give the same steps while $\eta$ is constant. They differ once a schedule changes $\eta$, which is why the PyTorch documentation points out the difference.

Adam

Different weights need very different step sizes. The row of the embedding table for a rare token gets a large gradient once in a while; the gain of a LayerNorm (layer normalization) gets small, steady ones. A single learning rate cannot suit both.

Adam (Kingma and Ba, 2015; the name comes from adaptive moment estimation) gives each weight its own step size. It keeps two running averages per weight: the average gradient $m$ (the direction), and the average squared gradient $v$ (the typical size). Then it divides one by the square root of the other, so that every weight takes a step of about the same size. There are four lines.

1. Average of the gradient (the direction):

2. Average of the squared gradient (the typical size):

3. Bias correction, explained below:

4. The step:

where:

• $g$ is the gradient and $g^2$ its square. All operations here are elementwise: each weight is handled on its own, with its own $m$ and $v$;
• $m$ and $v$ start at 0. Here $v$ is Adam's second average, not the momentum velocity of the previous section: the two papers use the same letter;
• $\beta_1$ (beta one), usually $0.9$, and $\beta_2$, usually $0.999$, set how long the averages remember: $m$ averages over roughly the last 10 steps, $v$ over the last 1,000;
• $t$ is the step number, starting from 1, and $\beta_1^t$ is $\beta_1$ to the power $t$;
• $\hat m$ and $\hat v$ ("m hat", "v hat") are the corrected averages;
• $\epsilon$ (epsilon), typically $10^{-8}$, only avoids a division by zero.

Why bias correction is needed. $m$ and $v$ start at 0, so after the first step $m = (1 - \beta_1)\, g = 0.1\, g$: ten times too small. For $v$ it is a thousand times too small ($1 - \beta_2 = 0.001$). Dividing by $1 - \beta^t$ fixes exactly that: at $t = 1$ it divides $m$ by $0.1$ and $v$ by $0.001$. As $t$ grows, $\beta^t$ goes to 0 and the correction disappears.

Example. Two weights, two steps each. Weight A gets gradients $0.5$ then $0.4$; weight B gets $50$ then $40$, a hundred times larger.

	weight A	weight B
step 1: $m$, $v$	$0.05$, $0.00025$	$5$, $2.5$
step 1: $\hat m$, $\hat v$	$0.5$, $0.25$	$50$, $2500$
step 1: $\hat m / \sqrt{\hat v}$	$1$	$1$
step 2: $\hat m / \sqrt{\hat v}$	$0.988$	$0.988$

Both weights move by the same amount, about $\eta$, although B's gradients are 100 times larger: multiplying $g$ by 100 multiplies $\hat m$ by 100 and $\sqrt{\hat v}$ by 100 too, and the ratio does not change.

If instead the gradient flips, $0.5$ then $-0.5$, the second ratio is about $-0.05$: the average direction $\hat m$ is almost 0, so Adam barely moves. Like momentum, it hesitates when the gradient disagrees with itself.

Weight decay, done right: AdamW

Weight decay pulls every weight a little towards zero at each step, to keep the weights small. With plain SGD, this is the same as adding a penalty to the loss, $\tfrac{\lambda}{2}\lVert\theta\rVert^2$ (L2 regularization), where $\lambda$ (lambda) is the decay strength and $\lVert\theta\rVert^2$ is the sum of the squares of all the weights. The gradient of that penalty is $\lambda \theta$.

With Adam the two are not the same. If $\lambda\theta$ is added to the gradient, it also gets divided by $\sqrt{\hat v}$, so weights with large gradients are barely decayed. Loshchilov and Hutter (2019) decouple the two, and call the result AdamW: the decay is applied directly to the weights, outside Adam's adaptive step. PyTorch's AdamW first shrinks the weights,

then takes the Adam step above. In words, each step multiplies every weight by $1 - \eta\lambda$.

Example. With $\eta = 0.001$ and $\lambda = 0.01$, each step multiplies every weight by $0.99999$: a tiny pull. Over 10,000 steps it adds up to $0.99999^{10000} \approx 0.905$, so a weight that the gradient does not defend loses about 10% of its size.

AdamW is the usual optimizer for transformers. Weight decay is often applied only to matrices. nanoGPT, for example, decays every 2-D parameter (linear weights and embeddings) and leaves biases and LayerNorm gains alone.

Learning-rate schedules: warmup, then cosine decay

Warmup. At the start of training the weights are random, and Adam's $\hat v$ is estimated from just a handful of gradients. A full-size step at this point can push the model somewhere it never recovers from. So the learning rate starts near 0 and grows in a straight line up to its maximum over the first steps. Goyal et al. (2017) used such a gradual warmup to make early training with very large batches stable.

Cosine decay. After warmup, the learning rate follows half a cosine wave, from its maximum down to a minimum. The schedule comes from SGDR (stochastic gradient descent with warm restarts; Loshchilov and Hutter, 2017):

where:

• $s$ is the current step and $\eta(s)$ the learning rate used at that step;
• $\eta_{\max}$ and $\eta_{\min}$ are the largest and smallest learning rates;
• $r$ is the progress through the decay, from 0 at the end of warmup to 1 at the last step.

How to read it: $\cos \pi r$ goes from $1$ (at $r = 0$) to $-1$ (at $r = 1$), so the bracket $1 + \cos \pi r$ goes from 2 to 0, and $\eta$ slides from $\eta_{\max}$ down to $\eta_{\min}$.

Example. With $\eta_{\max} = 0.001$ and $\eta_{\min} = 0.0001$:

progress $r$	learning rate $\eta$
0	0.001
0.25	0.00087
0.5	0.00055
0.75	0.00023
1	0.0001

From $r = 0$ to $0.25$ the rate drops by $0.00013$, from $0.25$ to $0.5$ by $0.00032$: it decays slowly at first, fastest in the middle, and gently again at the end.

nanoGPT's training script uses exactly this shape, with $\eta_{\min} = \eta_{\max} / 10$.

Gradient clipping

A single bad batch can produce a gradient hundreds of times larger than usual, and one step with it can undo hours of training. Pascanu et al. (2013) proposed a simple guard: whenever the gradient is too long, shrink it. Its length, the norm, is measured over all parameters together, as if every gradient were laid end to end in one long vector:

where:

• $\lVert g \rVert$ is the norm: square every gradient entry $g_i$ of every tensor, add them all up, take the square root;
• $c$ is the threshold, the largest norm allowed;
• $\min(1, \cdot)$ means: if the norm is already below $c$, multiply by 1, so nothing changes.

Example. Two gradient entries, $3$ and $4$, have norm $\sqrt{9 + 16} = 5$. With $c = 1$ both are multiplied by $1/5$ and become $0.6$ and $0.8$: the new norm is 1, and the direction is the same (the second entry is still $4/3$ of the first). With $c = 10$ the factor is $\min(1, 10/5) = 1$, and nothing changes.

Scaling everything by one factor keeps the gradient's direction and only shortens it. Clipping each tensor on its own would change the direction. PyTorch's clip_grad_norm_ adds $10^{-6}$ to the norm in the denominator, and returns the norm before clipping. Training loops log that number: a norm that keeps hitting the threshold is an early sign of trouble. nanoGPT clips at $c = 1.0$ by default.

References

• H. Robbins, S. Monro, A Stochastic Approximation Method, Annals of Mathematical Statistics 22(3), 400–407, 1951.
• B. T. Polyak, Some methods of speeding up the convergence of iteration methods, USSR Computational Mathematics and Mathematical Physics 4(5), 1–17, 1964. The heavy-ball (momentum) method.
• I. Sutskever, J. Martens, G. Dahl, G. Hinton, On the importance of initialization and momentum in deep learning, ICML, 2013.
• D. P. Kingma, J. Ba, Adam: A Method for Stochastic Optimization, ICLR, 2015.
• I. Loshchilov, F. Hutter, Decoupled Weight Decay Regularization, ICLR, 2019. AdamW.
• I. Loshchilov, F. Hutter, SGDR: Stochastic Gradient Descent with Warm Restarts, ICLR, 2017. Cosine annealing.
• P. Goyal et al., Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour, technical report, 2017. Linear scaling and gradual warmup.
• R. Pascanu, T. Mikolov, Y. Bengio, On the difficulty of training recurrent neural networks, ICML, 2013. Gradient-norm clipping.
• I. Goodfellow, Y. Bengio, A. Courville, Deep Learning, MIT Press, 2016. Chapter 8: Optimization for Training Deep Models.
• PyTorch documentation: torch.optim.SGD, torch.optim.AdamW and torch.nn.utils.clip_grad_norm_, the exact update rules the exercises reproduce.