Training well: batches, epochs, and optimizers

The gradient of a batch

With $N$ examples and a per-example error $\ell_i(\mathbf{w})$, the total error is $L(\mathbf{w})=\tfrac1N\sum_i \ell_i(\mathbf{w})$. A batch $\mathcal{B}$ of $B$ randomly chosen examples gives the approximate gradient

It is an unbiased estimator. Its variance is $\frac{N-B}{N-1}\,\frac{\sigma^2}{B}$, where $\sigma^2$ is the variance of one example’s gradient: the error falls as $1/\sqrt{B}$, while the cost of the step grows as $B$. That is why cutting the noise in half takes a batch about four times larger. In the lesson, the arrow’s average error (measured over 24 batches of each size) was 20 %, 10 %, 4 %, and 1.7 % for batches of 2, 8, 32, and 128: it falls a bit faster than $1/\sqrt{B}$ because the factor $\frac{N-B}{N-1}$ matters when the batch is a large share of the 256 examples. An epoch has $\lceil N/B\rceil$ steps, and after $t$ steps $tB$ examples have been seen. The idea of updating with noise goes back to Robbins and Monro (1951).

Near a minimum with curvature $h$, gradient noise (variance $s^2=\sigma^2/B$) leaves an error floor: $\mathbb{E}[L]\approx \eta\,s^2/4$ when $\eta h\ll 1$. It is proportional to $\eta/B$: lowering the rate or raising the batch lowers it. That is why decaying the rate helps the path settle.

Momentum

With $\beta=0.9$ and a constant gradient, $\mathbf{v}\to \mathbf{g}/(1-\beta)$: the effective step is ten times larger; for the same effective step, momentum’s rate would be about ten times smaller than SGD’s (in practice the ratio is smaller: in the lesson’s canyon, 0.02 versus 0.07). For a quadratic with curvatures $\lambda_{\min}\le h\le\lambda_{\max}$ and condition number $\kappa=\lambda_{\max}/\lambda_{\min}$, the best fixed rate for SGD shrinks the distance to the bottom per step by a factor $\frac{\kappa-1}{\kappa+1}$, and well-tuned momentum by $\frac{\sqrt\kappa-1}{\sqrt\kappa+1}$ (Polyak, 1964). In the lesson’s canyon, near the bottom, $\kappa\approx 50$: about $0.96$ per step versus $0.75$ (in theory, without noise).

Adam

Everything is per weight (the square and the root act component by component). The defaults from Kingma and Ba (2015) are $\eta=0.001$, $\beta_1=0.9$, $\beta_2=0.999$, and $\varepsilon=10^{-8}$. The bias corrections compensate for $\mathbf{m}$ and $\mathbf{s}$ starting at zero. On the first step $\hat{\mathbf{m}}/\sqrt{\hat{\mathbf{s}}}=\operatorname{sign}(\mathbf{g})$, so every weight moves by $\eta$; in normal operation each weight’s step is at most on the order of $\eta$. Two consequences: the method is invariant to rescaling each gradient separately (its authors say so), and near an optimum the signal-to-noise ratio drops and the steps shrink on their own, an “automatic annealing” in Kingma and Ba’s words. With noisy batch gradients that is not enough: in the lesson’s simulation, with a constant $\eta$ the path keeps jittering around the bottom, with an amplitude that grows with $\eta$, and lowering $\eta$ at the end calms it.

Learning-rate schedules

Two real recipes: the original transformer (Vaswani et al., 2017) used Adam with $\beta_2=0.98$ and warmed up for 4000 steps; GPT-3 (Brown et al., 2020) used Adam with $\beta_2=0.95$, a linear warmup over the first 375 million tokens, and cosine decay down to 10 % of the peak value over 260 billion tokens. Its batch also grew from 32 thousand tokens to the final value (3.2 million for the 175-billion-parameter model) over the first 4 to 12 billion tokens, depending on model size, and it saw only 0.44 epochs of its web-crawl text (Table 2.2). When the batch grows, Goyal et al. (2017) propose multiplying the rate by the same factor and warming up: with that rule they trained ResNet-50 with batches of 8192 images on 256 GPUs, in one hour, with no loss of accuracy.

Hyperparameters

The weights $\mathbf{w}$ are learned with $\nabla L$; hyperparameters ($\eta$, $B$, $\beta$, $T_w$, the optimizer…) set how learning happens and have no gradient within reach. They are chosen by trying: in spaces of several dimensions, trying random combinations often pays off more per trial than sweeping a grid (Bergstra and Bengio, 2012).

In code

import math

# Eight fixed examples of y ≈ 2x + 1. Batches are taken in order so the
# result is reproducible (in practice you shuffle every epoch).
X = [0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5]
Y = [1.2, 1.8, 3.1, 4.1, 4.9, 6.2, 6.9, 8.1]

def grad(w, b, idx):                       # gradient of the batch mean squared error
    n = len(idx)
    gw = sum(2 * (w * X[i] + b - Y[i]) * X[i] for i in idx) / n
    gb = sum(2 * (w * X[i] + b - Y[i]) for i in idx) / n
    return [gw, gb]

def train(opt, eta, B, epochs, beta=0.9, b1=0.9, b2=0.999, eps=1e-8):
    p = [0.0, 0.0]                         # weights: w and b
    v = [0.0, 0.0]; m = [0.0, 0.0]; s = [0.0, 0.0]; t = 0
    for _ in range(epochs):                # one epoch = go through all the batches
        for start in range(0, len(X), B):
            g = grad(p[0], p[1], range(start, min(start + B, len(X))))
            t += 1
            if opt == "sgd":
                step = [eta * gi for gi in g]
            elif opt == "momentum":
                v = [beta * vi + gi for vi, gi in zip(v, g)]
                step = [eta * vi for vi in v]
            else:                          # Adam
                m = [b1 * mi + (1 - b1) * gi for mi, gi in zip(m, g)]
                s = [b2 * si + (1 - b2) * gi * gi for si, gi in zip(s, g)]
                step = [eta * (mi / (1 - b1 ** t)) / (math.sqrt(si / (1 - b2 ** t)) + eps)
                        for mi, si in zip(m, s)]
            p = [pi - di for pi, di in zip(p, step)]
    return [round(pi, 3) for pi in p]

# the exact least-squares solution is [1.998, 1.042]
for opt, eta in [("sgd", 0.05), ("momentum", 0.005), ("adam", 0.3)]:
    print(opt, train(opt, eta, B=4, epochs=50))
# sgd [1.999, 1.033]
# momentum [2.008, 1.043]
# adam [1.993, 1.028]

On a problem this small and well conditioned, SGD is already enough: the differences show up on landscapes like the canyon, and each optimizer needed its own rate.

What comes next

Lowering the training error is not enough: you have to measure with data the model has not seen. That is lesson 10 (generalizing and evaluating), which adds early stopping and weight decay as ways to regularize.