Chapter 7 of 0–21

How the model learns

Gradients, AdamW and weight decay, with toy examples and the real training curves.

Training means lowering the loss

Training repeatedly nudges all 226,816 parameters so that the loss on the training set goes down. To know which way to nudge, we need slopes.

The derivative of a function tells how fast its output changes when the input changes a little: the slope of its graph. For f(w) = (w-3)^2 the derivative is f'(w) = 2(w-3); at w = 5 it is 4 (going right increases f), at w = 1 it is −4. The gradient is the list of derivatives of the loss with respect to every parameter at once — a vector with one entry per parameter, pointing in the direction in which the loss increases fastest. Computer libraries compute it exactly and automatically.

Gradient descent moves the parameters a small step against the gradient:

where the step size \gamma is the learning rate.

Press ▶ next to “steps” to watch the path grow one step at a time.

Full-batch training

The paper computes each update from all 3,830 training examples at once — full-batch training training [1]. Then one parameter update, a step (epoch), is also one pass over the whole training set (an “epoch”), and the training process has no randomness after the start: given the starting point, every step is determined. (This will let us recompute the full-batch gradient at any saved moment of training, exactly up to floating-point rounding, in chapter 19.)

Adam, weight decay, AdamW

The rule that turns gradients into updates is the optimizer. Plain gradient descent struggles when some parameters need much bigger steps than others. Adam keeps, for every parameter, a running average of its recent gradients (“momentum”, decay 0.9 per step) and of its squared gradients (decay 0.98 here), and steps by their ratio — so each parameter moves by roughly the learning rate in its consistent direction, whatever the size of its gradient [14].

Weight decay additionally shrinks every parameter a little each step. AdamW applies it directly (“decoupled”) [13] [32]:

with \lambda the weight-decay strength and \hat m, \hat v Adam's running averages. Weight decay is a form of regularization: it prefers solutions with small weights among those that fit the data. With \gamma = 0.001 and \lambda = 1, every weight shrinks by 0.1 % per step unless the loss gradient pushes back — after 1,000 steps an unopposed weight would shrink to 0.999^{1000} \approx 37\% of its size.

The three rules side by side on a loss with two parameters, shaped like a long valley (its lowest point is marked ×). Press ▶ next to “steps” and compare where each one goes.

The same pull on the real network: each training run as a path through (size of all weights, loss). With weight decay (λ = 1) the weights first grow and then shrink while the test loss falls; without it (λ = 0) they keep growing and the test loss stays high. Chapter 19 returns to this balance.

<p class="fig-expect"><strong>What to expect:</strong> With weight decay the norm rises, then falls while the test loss collapses; without it the norm keeps growing and the test loss stays high.</p>
<p class="fig-meta">
  <a href="../chapters/19-gradients-weight-decay.html">Learn step by step →</a>
   · Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> Sec. 5.3
  · Computed by <code>checkpoint_metrics.py</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>

The settings of this experiment

Settings chosen before training are hyperparameters. The parameters start at random values (their initialization); a random seed fixes those values and the train/test split so a run can be repeated exactly.

Sources: paper Sec. 3 [1] and the paper's training code [6]. Two details matter for exact reproduction: the warm-up multiplies the learning rate by 0 at the very first step (so the first update leaves the weights unchanged, though Adam's running averages are updated), and the paper's checkpoint labelled “epoch e” holds the weights after e + 1 updates. We verified both by recomputing the paper's logged losses from its checkpoints (see Reproduce).

Overfitting

overfitting is doing much better on the training set than on unseen data. A network with 226,816 parameters and only 3,830 training examples can overfit easily: it has more than enough capacity to memorise. The next chapter shows what happens next.