Initialization and Optimizers: Where to Start, and How to Step

Backpropagation, by Hand trained a two-layer network and found two things plain gradient descent doesn’t handle on its own. Where training starts matters: two hidden units could represent XOR, but most random starts never found it. And the step itself is crude: one learning rate for every weight, and no memory of the steps before. Learning Is Optimization had already shown the cost of that. On a badly scaled problem, one learning rate has to be small enough for the steepest direction, so it crawls along all the others.

This post is about those two choices, made at every training run, each derived and then tested with numbers on the router or on Learning Is Optimization’s problem. Dropout and Normalization covers the other half of making deep networks train, dropout and normalization layers. The two choices are:

• Initialization. How large the starting weights should be, derived from what happens to a signal passing through many layers. Get it wrong and a deep network can’t learn at all.
• The optimizer. Momentum, which remembers past gradients; Adam, which gives every weight its own step size; the warmup that starts training gently; and AdamW, the version of weight decay that nearly every large language model uses.

Table of contents

Why depth breaks the signal

Take one layer of a network, $z = x W$, with $n$ inputs and weights drawn at random with mean 0 and variance $\sigma^2$. (In this section $n$ is a layer’s number of inputs and $z$ a pre-activation, as in Backpropagation, by Hand, both local; they aren’t the dataset size and logits of the notation in AI Engineering, from the Inside.) Each output is a sum of $n$ terms $x_i W_{ij}$. With independent, zero-mean weights, the variances add:

Now pass $z$ through a ReLU. $z$ is symmetric around zero, so the ReLU zeroes half of it and keeps the other half unchanged, which halves its mean square: $\mathbb{E}[\operatorname{ReLU}(z)^2] = \tfrac12 \operatorname{Var}(z)$. Put together, one ReLU layer multiplies the signal’s mean square by

and $L$ layers multiply it by that factor to the power $L$. There’s no middle ground. If the factor is 0.5, twenty layers shrink the signal’s mean square by $0.5^{20} \approx 10^{-6}$, its size by a thousandth. If it’s 2, they grow it by the same amounts. Only a factor of exactly 1 keeps it level, which means

That’s He initialization, from He et al. (2015). It’s what Backpropagation, by Hand’s mlp.py used. The older $\sigma^2 = 1/n$ is the same argument for activations that don’t halve the signal, like tanh near zero. Strictly, $1/n_\text{in}$ is LeCun’s initialization; Xavier (or Glorot) initialization averages the input and output counts, $\sigma^2 = 2/(n_\text{in} + n_\text{out})$, and for the square layers here the two are the same.

The backward pass has the same problem. Backpropagation, by Hand’s rules send the gradient down through each layer as $\delta_X = (\delta_A \odot \phi'(Z))\, W^\top$: the same weights and the same ReLU mask, so the same kind of factor per layer, except that $n$ now counts the layer’s outputs, because $W^\top$ sums over them. With square layers the two factors are equal; when widths change, a fan-in initialization leaves the backward pass off by a fixed ratio of widths rather than an exponential, which is what Glorot’s average splits.

Measure it. depth_signal.py pushes a batch of random inputs through 20 layers of width 256, then sends a random gradient back down, for each of five initializations (median of five seeds):

def signal(scale, act, depth=20, width=256, seed=0):
    """Push a batch of 512 random inputs through `depth` layers. Return the size
    of the activations after each layer, and of the gradient arriving at each
    layer on the way back, starting from a random unit-size gradient at the top."""
    rng = np.random.default_rng(seed)
    x = rng.normal(size=(512, width))
    Ws = [rng.normal(0, scale(width), (width, width)) for _ in range(depth)]
    phi, dphi = {"relu": (lambda z: np.maximum(z, 0), lambda z: (z > 0) * 1.0),
                 "tanh": (np.tanh, lambda z: 1 - np.tanh(z) ** 2)}[act]
    rms = lambda a: np.sqrt((a ** 2).mean())
    h, zs, forward = x, [], []
    for W in Ws:
        zs.append(h @ W)
        h = phi(zs[-1])
        forward.append(rms(h))
    g, backward = rng.normal(size=h.shape), []
    for W, z in zip(Ws[::-1], zs[::-1]):
        g = (g * dphi(z)) @ W.T                        # "Backpropagation, by Hand": mask, then W transposed
        backward.append(rms(g))
    return forward, backward[::-1]                      # backward[0] reaches layer 1

Full script: depth_signal.py

                          activation size after layer         gradient reaching 
                            1         5        10        20   layer 20   layer 1
N(0, 0.01^2), ReLU   1.13e-01  1.98e-05  3.68e-10  1.37e-19   1.14e-01  1.21e-19
N(0, 1), ReLU        1.13e+01  1.98e+05  3.68e+10  1.37e+21   1.14e+01  1.21e+21
Xavier 1/n, ReLU     7.07e-01  1.89e-01  3.35e-02  1.14e-03   7.15e-01  1.00e-03
He 2/n, ReLU         1.00e+00  1.07e+00  1.07e+00  1.16e+00   1.01e+00  1.03e+00
Xavier 1/n, tanh     6.28e-01  3.21e-01  2.27e-01  1.58e-01   9.73e-01  1.88e-01

With $n = 256$, the formula predicts the factor on the signal’s size (the square root of the mean-square factor) exactly.

• $\sigma = 0.01$: $\sqrt{256 \times 0.0001 / 2} = 0.113$ per layer. After 20 layers the activations are $10^{-19}$, and so is the gradient reaching layer 1. The network’s output doesn’t depend on its input, and its first layers get no signal to learn from.
• $\sigma = 1$: $\sqrt{128} = 11.3$ per layer, to $10^{21}$.
• Xavier with ReLU: $\sqrt{1/2} = 0.707$ per layer, so 20 layers lose a factor of 1,000. Survivable at 10 layers, crippling at 100.
• He with ReLU: 1.00 per layer. The signal and the gradient both stay near 1 all the way through.
• Xavier with tanh decays, because tanh squashes large values, but only slowly, because tanh has slope 1 near zero. This is the pairing Xavier was derived for, and it does far better than Xavier with ReLU.

Does it matter for training? Give the four-way router of Cross-Entropy and the Softmax Gradient ten hidden layers of 64 ReLU units, and train from each initialization:

def init(sizes, rng, std):
    return [{"W": rng.normal(0, std(m), (m, n)), "b": np.zeros(n)}
            for m, n in zip(sizes[:-1], sizes[1:])]

if __name__ == "__main__":
    inits = {"N(0, 0.01^2)": lambda m: 0.01, "Xavier 1/n": lambda m: np.sqrt(1 / m),
             "He 2/n": lambda m: np.sqrt(2 / m)}
    sizes = [None] + [64] * 10 + [4]                         # ten hidden ReLU layers
    out = {}
    for seed in range(5):
        ((Xtr, ytr), _, _), _ = four_way(*split(seed))
        sizes[0] = Xtr.shape[1]
        for name, std in inits.items():
            for opt_name, make in [("SGD 0.1", lambda: SGD(0.1)), ("Adam 0.001", lambda: Adam(1e-3))]:
                rng = np.random.default_rng(seed)
                layers, opt = init(sizes, rng, std), make()
                start, g0 = loss_and_grads(layers, Xtr, ytr)
                first = g0[0]["W"]
                early = None
                for t in range(1000):
                    if t == 100:
                        early = loss_and_grads(layers, Xtr, ytr)[0]
                    idx = rng.choice(len(ytr), 16, replace=False)
                    _, grads = loss_and_grads(layers, Xtr[idx], ytr[idx])
                    for i, (L, G) in enumerate(zip(layers, grads)):
                        for k in L:
                            opt.step(L[k], G[k], (i, k))
                out.setdefault((name, opt_name), []).append(
                    (np.abs(first).max(), start, early, loss_and_grads(layers, Xtr, ytr)[0]))

Full script: deep_init.py

               optimizer    largest first-layer gradient   train loss at 0 / 100 / 1,000 steps
N(0, 0.01^2)   SGD 0.1                           7.3e-15          1.386   1.387   1.387
N(0, 0.01^2)   Adam 0.001                        7.3e-15          1.386   1.386   1.386
Xavier 1/n     SGD 0.1                           6.9e-04          1.386   1.376   0.001
Xavier 1/n     Adam 0.001                        6.9e-04          1.386   0.091   0.000
He 2/n         SGD 0.1                           2.5e-02          1.399   0.700   0.000
He 2/n         Adam 0.001                        2.5e-02          1.399   0.011   0.000

(adam_eps.py, below, imports init from deep_init.py. deep.py, also below, has the network’s forward and backward passes.)

• Too small is fatal, even for Adam. From $\sigma = 0.01$, the first layer’s largest gradient is $7 \times 10^{-15}$, and after 1,000 steps the loss is still 1.386. That’s $\log 4$ from Cross-Entropy and the Softmax Gradient: the logits are all near zero, so the network predicts a uniform distribution over the four routes and never leaves it. Adam, which rescales small gradients, can’t rescue it either; the Adam section below shows why.
• He trains fastest. After 100 steps of SGD it’s at 0.700, while Xavier, whose gradients are 36 times smaller at the first layer, has barely moved from 1.386. Both memorize the training set eventually.

Two practical notes. Deep-learning libraries initialize layers automatically, but their defaults differ, so it’s worth checking which one a model uses. And modern language models have another reason to care: they’re so deep that even small per-layer drift adds up, which is one reason they rely on residual connections and normalization (Dropout and Normalization and The Residual Stream) as well as careful initialization.

The same argument sets the output layer’s scale, and with it the first loss Cross-Entropy and the Softmax Gradient described. The last layer is linear, so the variance argument says its weights should have $\sigma^2 = 1/n$ at most, and many models start it smaller still or at zero. Then the logits start near zero and the first loss is near $\log V$, or $\log 4 = 1.386$ for the router, plus about $\sigma_z^2/2$ for logits of standard deviation $\sigma_z$. If they start large, the first loss is far above that, and training spends its first steps shrinking them. (deep.py below uses He for every layer, output included, but the router’s sparse 0/1 inputs keep its logits small anyway: the He network’s first loss is 1.399, just above $\log 4$.)

Momentum: remember the last step

Learning Is Optimization left gradient descent crawling. On the token-cost fit with raw word counts, the Hessian’s eigenvalues differ by a factor of $\kappa = 128$. Any step size small enough for the steep direction is 128 times too small for the flat one, and descent zigzags across the valley instead of rolling along it.

Momentum (Polyak’s “heavy ball”, 1964) adds a memory. Keep a running velocity $v$, a decaying sum of past gradients, and step along it:

Gradients that point the same way step after step, along the valley floor, add up, up to a factor of $1/(1 - \beta)$, which is 10 for the common $\beta = 0.9$. Gradients that flip sign each step, across the valley, cancel.

How much faster? On a quadratic, each eigen-direction with curvature $\lambda$ evolves on its own. (Here $\lambda$ is a Hessian eigenvalue, as in Learning Is Optimization; in the AdamW section it will be the weight-decay strength.) Write $e_t$ for the error along it; momentum turns Learning Is Optimization’s one-step recurrence into a two-step one:

Its solutions are powers $r^t$ of the roots of $r^2 - (1 + \beta - \eta\lambda) r + \beta = 0$. When those roots are complex, both have size exactly $\sqrt\beta$, whatever $\lambda$ is. That happens for every curvature in the range $(1 - \sqrt\beta)^2 < \eta\lambda < (1 + \sqrt\beta)^2$. So choose $\eta$ and $\beta$ to make that range just cover the problem, from $\lambda_\text{min}$ to $\lambda_\text{max}$. Dividing the two edges gives $\sqrt{\kappa} = (1 + \sqrt\beta)/(1 - \sqrt\beta)$, and every direction then shrinks by

The number of steps to a given accuracy goes from order $\kappa$ to order $\sqrt\kappa$. For $\kappa = 128$ that’s a factor of about 11. On Learning Is Optimization’s problem: momentum.py imports that post’s data from a_loss.py, a copy of its script, builds the fit’s Hessian and least-squares solution, and counts the steps each method takes to reach it:

def steps_to(update, tol=1e-6, limit=100_000):
    """Steps until the parameters are within tol (relative) of the least-squares fit."""
    th, state = np.zeros(2), {}
    for t in range(1, limit + 1):
        th = update(th, state, t)
        if np.linalg.norm(th - theta_star) < tol * np.linalg.norm(theta_star):
            return t
    return None

def gd(eta):
    return lambda th, s, t: th - eta * grad(th)

def heavy_ball(eta, beta):
    def step(th, s, t):
        s["v"] = beta * s.get("v", 0) + grad(th)     # velocity: a running sum of gradients
        return th - eta * s["v"]
    return step

k = hi / lo
eta_gd = 2 / (hi + lo)                                # the best fixed step for plain descent
eta_hb = 4 / (np.sqrt(hi) + np.sqrt(lo)) ** 2         # Polyak's best choice for momentum
beta_hb = ((np.sqrt(k) - 1) / (np.sqrt(k) + 1)) ** 2

Full script: momentum.py

eigenvalues 0.2407 and 30.72, kappa 128
gradient descent, eta 0.0646:              882 steps   (theory: error shrinks by 0.9844 per step)
momentum, eta 0.1099, beta 0.701:     108 steps   (theory: 0.8374 per step)
momentum, beta 0.5, same effective step:     736 steps
momentum, beta 0.9, same effective step:     630 steps
momentum, beta 0.99, same effective step:    2302 steps
first t with sqrt(beta)^t < 1e-6: 78; with t * sqrt(beta)^t < 1e-6: 104

• 882 steps against 108. Plain descent at its best step shrinks the error by 0.9844 per step, exactly as predicted, and needs 882 steps to get within a millionth of the answer. Tuned momentum needs 108. The rate 0.8374 alone would give 78. The gap is not noise, and the theory accounts for most of it: tuned this way, the two extreme directions sit right on the edge of the band, where the two roots coincide and the error decays like $t\,(\sqrt\beta)^t$ instead of $(\sqrt\beta)^t$. That predicts 104, the last line.
• Momentum takes steps plain descent can’t. Its step size, 0.110, is past plain descent’s stability limit of $2/\lambda_\text{max} = 0.065$ (see Learning Is Optimization). The velocity’s memory is what keeps a step that large stable.
• Untuned, the gain is small or negative. Holding the effective step $\eta / (1 - \beta)$ at plain descent’s value, $\beta = 0.5$ or $0.9$ helps a little, and $\beta = 0.99$ makes things worse: the velocity overshoots and oscillates for a long time before it settles.

Watch it happen on Learning Is Optimization’s widget, now with a momentum control. In raw units, plain descent at $\eta = 0.05$ zigzags onto the valley floor and then crawls along it. Set $\beta = 0.7$ and $\eta = 0.11$, a step that makes plain descent explode, and the path swings along the floor and settles:

Real networks aren’t quadratics, and nobody computes $\kappa$ for them, but the mechanism carries over: momentum averages away the back-and-forth and builds speed along consistent directions. In practice $\beta = 0.9$ is the default, and the learning rate is tuned around it. A common variant, Nesterov momentum (nesterov=True in PyTorch), evaluates the gradient at the point the velocity is about to carry the parameters to; it’s the version with guarantees beyond quadratics.

Adam: a step size for every weight

Momentum still uses one learning rate for every weight. Learning Is Optimization found the other fix for a badly scaled problem: rescale each direction, which is what standardizing the feature did. Adam (Kingma and Ba, 2014) does that automatically, per weight, while training. For each parameter it keeps two running averages, of the gradient and of the squared gradient:

and steps by their ratio (here $\epsilon$ is a small constant, not a regularizer):

with defaults $\beta_1 = 0.9$, $\beta_2 = 0.999$ and $\epsilon = 10^{-8}$. $m$ is momentum, written as an average instead of a sum. $\sqrt v$ is the typical size of this parameter’s gradient. Dividing by it makes the step independent of the gradient’s scale. Multiply a feature by 100 and its weight’s gradients grow 100 times, but so does $\sqrt v$, and the step stays about $\eta$. That’s not quite Learning Is Optimization’s standardization. The weight now needs to move 100 times less, and Adam still steps $\eta$. But plain descent’s step would have grown 100 times while its target shrank 100 times, a mismatch of $10^4$; Adam’s is $10^2$. In Learning Is Optimization’s terms Adam is a diagonal preconditioner, a cheap and partial one, updated every step.

Bias correction. Both averages start at zero, so early on they’re too small. With a gradient that’s a constant $g$, the average after $t$ steps is

It reaches $g$ only as $\beta_1^t \to 0$. Dividing by $(1 - \beta_1^t)$ undoes that exactly: $\hat m_t = m_t / (1 - \beta_1^t)$, and likewise $\hat v_t = v_t / (1 - \beta_2^t)$. It matters more for $v$, because $\beta_2 = 0.999$ takes about a thousand steps to forget its zero start:

import numpy as np

# Adam's moving averages on a gradient that is always 0.001: a small, perfectly
# consistent gradient. Both averages start at 0, so both start too small, and
# v (averaged with 0.999) far more than m (0.9). Without correction, the first
# steps are too large.
b1, b2, eta, g = 0.9, 0.999, 0.01, 0.001
m = v = 0.0
for t in range(1, 6):
    m = b1 * m + (1 - b1) * g
    v = b2 * v + (1 - b2) * g * g
    raw, fixed = m / np.sqrt(v), (m / (1 - b1 ** t)) / np.sqrt(v / (1 - b2 ** t))
    print(f"t={t}: m {m:.2e}  v {v:.2e}   step without correction {eta * raw:.5f}"
          f"   with correction {eta * fixed:.5f}")

# The first corrected step is eta * g / |g| for every parameter, whatever the gradient's size.
for g in (1e-6, 1e-2, 10.0):
    m, v = (1 - b1) * g, (1 - b2) * g * g
    print(f"gradient {g:>7g}: first Adam step {eta * (m / (1 - b1)) / np.sqrt(v / (1 - b2)):.4f}"
          f"   first gradient-descent step {eta * g:.1e}")

# How long does the uncorrected step stay too large? Its ratio to the
# corrected one is (1 - b1^t) / sqrt(1 - b2^t).
t = np.arange(1, 1001)
ratio = (1 - b1 ** t) / np.sqrt(1 - b2 ** t)
print(f"uncorrected / corrected: peak {ratio.max():.2f} at t={t[ratio.argmax()]}, "
      f"{ratio[99]:.2f} at t=100, {ratio[999]:.2f} at t=1000")

Full script: bias_correction.py

t=1: m 1.00e-04  v 1.00e-09   step without correction 0.03162   with correction 0.01000
t=2: m 1.90e-04  v 2.00e-09   step without correction 0.04250   with correction 0.01000
t=3: m 2.71e-04  v 3.00e-09   step without correction 0.04950   with correction 0.01000
t=4: m 3.44e-04  v 3.99e-09   step without correction 0.05442   with correction 0.01000
t=5: m 4.10e-04  v 4.99e-09   step without correction 0.05797   with correction 0.01000
gradient   1e-06: first Adam step 0.0100   first gradient-descent step 1.0e-08
gradient    0.01: first Adam step 0.0100   first gradient-descent step 1.0e-04
gradient      10: first Adam step 0.0100   first gradient-descent step 1.0e-01
uncorrected / corrected: peak 6.57 at t=12, 3.24 at t=100, 1.26 at t=1000

Without correction, $v$ starts much further below its true value than $m$ does, so the first steps are too large, not too small: 3.2 times at the first step, peaking at 6.6 times around step 12, and still 3.2 times at step 100 (the last line). Corrected, a consistent gradient gives a step of exactly $\eta$ from the first step.

The three gradient lines show what Adam’s learning rate means. At the first step, $\hat m_1 / \sqrt{\hat v_1} = g / |g| = \pm 1$, so every parameter moves by exactly $\eta$, whether its gradient was $10^{-6}$ or 10. Plain descent would have moved them by amounts differing by $10^7$. In Adam, $\eta$ is roughly the largest distance a parameter moves in one step, in the parameter’s own units. When gradients are consistent the step is near $\eta$; when they’re noise, it’s smaller, as the warmup section shows. (The exception is a sudden spike after a long run of near-zero gradients, which can move a parameter up to about $3\eta$ with the default $\beta$s.)

Why couldn’t Adam rescue the tiny initialization, then? Its gradients were about $10^{-14}$, below $\epsilon = 10^{-8}$, so the denominator was $\epsilon$ and each step was about $\eta g / \epsilon \approx 10^{-9}$: a million times more than plain descent’s, and still nothing. Removing $\epsilon$ doesn’t reliably help either:

import numpy as np
from router_data import split
from router4 import four_way
from deep import loss_and_grads
from deep_init import init
from optim import Adam

# Is Adam's epsilon what keeps it stuck at the tiny start? Remove it and see.
sizes = [None] + [64] * 10 + [4]
for eps in (1e-8, 1e-30):
    finals = []
    for seed in range(5):
        ((X, y), _, _), _ = four_way(*split(seed))
        sizes[0] = X.shape[1]
        rng = np.random.default_rng(seed)
        layers, opt = init(sizes, rng, lambda m: 0.01), Adam(1e-3, eps=eps)
        for t in range(1000):
            idx = rng.choice(len(y), 16, replace=False)
            _, grads = loss_and_grads(layers, X[idx], y[idx])
            for i, (L, G) in enumerate(zip(layers, grads)):
                for k in L:
                    opt.step(L[k], G[k], (i, k))
        finals.append(loss_and_grads(layers, X, y)[0])
    print(f"Adam eps {eps:g}: train loss after 1,000 steps, 5 seeds: " + " ".join(f"{v:.3f}" for v in finals))

Full script: adam_eps.py

Adam eps 1e-08: train loss after 1,000 steps, 5 seeds: 1.386 1.386 1.386 1.386 1.386
Adam eps 1e-30: train loss after 1,000 steps, 5 seeds: 0.169 1.386 1.386 0.067 1.386

Two of five seeds escape, three stay at $\log 4$. The input’s signal was multiplied by about $10^{-13}$ on its way to the output, so those gradients carry almost no information about the input, and rescaling them to full size mostly amplifies noise. The fix is the initialization, not the optimizer.

When the optimizer matters

Here’s the full optimizer code, and a network that uses it:

# optim.py
import numpy as np

class SGD:
    def __init__(self, eta):
        self.eta = eta
    def step(self, p, g, key):
        p -= self.eta * g

class Momentum:
    def __init__(self, eta, beta=0.9):
        self.eta, self.beta, self.v = eta, beta, {}
    def step(self, p, g, key):
        v = self.v[key] = self.beta * self.v.get(key, 0) + g
        p -= self.eta * v

def decays(key):
    """Weight decay applies to weight matrices only, not to biases or norm gains."""
    name = key[-1] if isinstance(key, tuple) else key
    return name == "W"

class Adam:
    def __init__(self, eta, beta1=0.9, beta2=0.999, eps=1e-8, decay=0.0, decoupled=True):
        self.eta, self.b1, self.b2, self.eps = eta, beta1, beta2, eps
        self.decay, self.decoupled, self.m, self.v, self.t = decay, decoupled, {}, {}, {}
    def step(self, p, g, key):
        if self.decay and not self.decoupled and decays(key):
            g = g + self.decay * p                              # L2 folded into the gradient
        t = self.t[key] = self.t.get(key, 0) + 1
        m = self.m[key] = self.b1 * self.m.get(key, 0) + (1 - self.b1) * g
        v = self.v[key] = self.b2 * self.v.get(key, 0) + (1 - self.b2) * g * g
        m_hat, v_hat = m / (1 - self.b1 ** t), v / (1 - self.b2 ** t)   # bias correction
        if self.decay and self.decoupled and decays(key):
            p -= self.eta * self.decay * p                      # AdamW: decay applied directly
        p -= self.eta * m_hat / (np.sqrt(v_hat) + self.eps)

Full script: optim.py

# deep.py
import numpy as np
from softmax_basics import softmax, cross_entropy

def init(sizes, rng):
    """He initialization for every ReLU layer: std sqrt(2 / fan_in)."""
    return [{"W": rng.normal(0, np.sqrt(2 / m), (m, n)), "b": np.zeros(n)}
            for m, n in zip(sizes[:-1], sizes[1:])]

def loss_and_grads(layers, X, y):
    """Softmax cross-entropy of a ReLU network, and its gradient ("Backpropagation, by Hand"'s rules)."""
    hs, h = [X], X
    for i, L in enumerate(layers):
        z = h @ L["W"] + L["b"]
        h = z if i == len(layers) - 1 else np.maximum(z, 0)
        hs.append(h)
    g = (softmax(hs[-1]) - np.eye(hs[-1].shape[1])[y]) / len(y)      # p - y, "Cross-Entropy and the Softmax Gradient"
    grads = [None] * len(layers)
    for i in reversed(range(len(layers))):
        grads[i] = {"W": hs[i].T @ g, "b": g.sum(0)}
        g = (g @ layers[i]["W"].T) * (hs[i] > 0)                         # back through ReLU
    return cross_entropy(hs[-1], y), grads

Full script: deep.py

(softmax_basics.py and router4.py are from Cross-Entropy and the Softmax Gradient. Adam’s decay options are for the weight-decay section below. As in practice, decay applies only to weight matrices, not to biases or, in Dropout and Normalization, norm gains.)

Race them on a problem with one answer: Cross-Entropy and the Softmax Gradient’s four-way linear router with an L2 penalty, trained full-batch. It’s convex, so a better optimizer can only mean a lower objective sooner. Each optimizer gets its best learning rate at each checkpoint, from a grid of half-decades: an oracle choice that no single run could make, but a fair one for asking how fast each optimizer can go. Then add one ordinary feature an engineer might add, the question’s length in characters, and don’t standardize it. race.py holds the objective and the training loop, and its main block runs the learning-rate grid; race_both.py runs it without and then with the length feature:

# The four-way router of "Cross-Entropy and the Softmax Gradient" with an L2 penalty: a convex objective, so
# "better optimizer" can only mean "lower objective sooner".
LAM = 0.003
def objective(W, b, X, y):
    return cross_entropy(X @ W + b, y) + 0.5 * LAM * (W ** 2).sum()

def run(opt, X, y, checkpoints=(30, 300, 3000)):
    W, b, out = np.zeros((X.shape[1], 4)), np.zeros(4), []
    for t in range(1, max(checkpoints) + 1):
        G = (softmax(X @ W + b) - np.eye(4)[y]) / len(y)
        opt.step(W, X.T @ G + LAM * W, "W")
        opt.step(b, G.sum(0), "b")
        if t in checkpoints:
            out.append(objective(W, b, X, y))
    return out

Full script: race.py

Full script: race_both.py

bag of words:
gradient descent  after   30 steps 0.3052 (eta 10)   after  300 steps 0.3011 (eta 10)   after 3000 steps 0.3011 (eta 3.2)
momentum 0.9      after   30 steps 0.3105 (eta 3.2)   after  300 steps 0.3011 (eta 10)   after 3000 steps 0.3011 (eta 0.32)
Adam              after   30 steps 0.3158 (eta 0.1)   after  300 steps 0.3011 (eta 0.1)   after 3000 steps 0.3011 (eta 0.01)
plus raw question length:
gradient descent  after   30 steps 1.3734 (eta 0.0032)   after  300 steps 1.3021 (eta 0.0032)   after 3000 steps 0.8829 (eta 0.0032)
momentum 0.9      after   30 steps 1.3237 (eta 0.0032)   after  300 steps 0.8877 (eta 0.0032)   after 3000 steps 0.2948 (eta 0.01)
Adam              after   30 steps 0.3495 (eta 0.1)   after  300 steps 0.2917 (eta 0.32)   after 3000 steps 0.2914 (eta 0.01)

• On the bag of words, it barely matters. All three reach the same minimum, 0.3011, within 300 steps, and gradient descent is if anything the fastest. Every feature is 0 or 1, so the problem is already well scaled, and the learning rate matters more than the optimizer.
• One unscaled feature breaks gradient descent. Question lengths run to tens of characters, so that one weight’s curvature is about a thousand times the words’ (below). Gradient descent’s step is capped by it, near 0.0037, and after 3,000 steps the objective is still 0.883. Momentum at the same step does better. Adam reaches 0.2917 in 300 steps, because it rescales the length weight’s steps separately from the words’.

import numpy as np
from router_data import split
from router4 import four_way

# The largest curvature of race.py's objective at the start (W = 0), where the
# softmax Hessian block is (diag(p) - p p^T) with p uniform, largest eigenvalue 1/4.
tr, va, te = split(0)
((X, y), _, _), _ = four_way(tr, va, te)
for name, Z in [("bag of words", X), ("plus raw length", np.c_[X, [len(r["question"]) for r in tr]])]:
    Zb = np.c_[Z, np.ones(len(Z))]
    top = 0.25 * np.linalg.eigvalsh(Zb.T @ Zb / len(Zb))[-1] + 0.003
    print(f"{name:16s} largest curvature {top:8.2f}, so gradient descent is stable only below 2 / that = {2 / top:.4f}")

Full script: curvature.py

bag of words     largest curvature     0.51, so gradient descent is stable only below 2 / that = 3.9118
plus raw length  largest curvature   543.91, so gradient descent is stable only below 2 / that = 0.0037

The curves, median of five splits:

race_curves.py prints the curves behind the chart: each optimizer at its best learning rate for 3,000 steps from race.py’s grid, plus momentum at gradient descent’s learning rate, median of five splits.

Full script: race_curves.py

gradient descent: 1.3833 1.3821 1.3815 1.3804 1.3792 1.3769 1.3729 1.3661 1.3541 1.3333 1.2981 1.2400 1.1496 1.0234 0.8678
momentum 0.9, same eta: 1.3833 1.3863 1.3809 1.3796 1.3707 1.3541 1.3187 1.2596 1.1643 1.0320 0.8725 0.7008 0.5382 0.4079 0.3314
momentum 0.9, eta 0.01: 1.4433 1.7060 1.8667 6.4745 13.3828 12.4024 21.6123 14.2705 13.1169 6.2789 6.9338 0.4115 0.3242 0.3016 0.2945
Adam: 1.3531 1.2584 1.1951 1.0186 0.8411 0.5947 0.4377 0.3681 0.3287 0.3085 0.2982 0.2935 0.2917 0.2914 0.2914
steps: 1 2 3 6 10 18 32 56 100 178 316 562 1000 1778 3162

Momentum at its best step for 3,000 steps, $\eta = 0.01$, nearly catches Adam, reaching 0.2945 (0.2948 in the race), but only by surviving a wild start: the objective climbs to 21.6 in the first 32 steps while the velocity builds up, before it settles. That curve is printed but not drawn; it would flatten the others.

Standardizing the length feature would fix gradient descent too, as Learning Is Optimization showed. Adam’s advantage is that it doesn’t need to be told. Inside a deep network, the “features” each layer sees are the previous layer’s activations, their scales differ from layer to layer and change as training proceeds, and nobody can standardize them by hand. Adam adapts to each weight’s scale continuously, and that’s much of why it became the default optimizer for training language models.

Why training starts with a warmup

Almost every large training run starts with a warmup: the learning rate climbs from near zero to its peak over the first few thousand steps. Llama 2’s runs used 2,000 warmup steps out of roughly 500,000 (2 trillion tokens at 4 million tokens per step), then a cosine decay. Learning Is Optimization promised a reason. Adam supplies one directly. Give it gradients that are pure noise, so that no direction is right, and measure the size of its step:

import numpy as np

# 100,000 parameters whose gradients are pure noise of size 1: no direction is
# right, so ideally they shouldn't move much. Two things about Adam's early steps:
#   the typical step, m_hat / sqrt(v_hat), in units of eta;
#   the per-parameter rate 1 / sqrt(v_hat), whose true value here is 1. With few
#   samples, v_hat can be tiny by chance, and then that parameter's rate is huge.
rng = np.random.default_rng(0)
b1, b2, n = 0.9, 0.999, 100_000
m, v = np.zeros(n), np.zeros(n)
for t in range(1, 5001):
    g = rng.normal(size=n)
    m = b1 * m + (1 - b1) * g
    v = b2 * v + (1 - b2) * g * g
    m_hat, v_hat = m / (1 - b1 ** t), v / (1 - b2 ** t)
    if t in (1, 2, 5, 10, 100, 1000, 5000):
        step = m_hat / np.sqrt(v_hat)
        print(f"step {t:5d}: typical update {np.sqrt((step ** 2).mean()):.3f} x eta;  "
              f"parameters with a rate over 3x its true value: {(1 / np.sqrt(v_hat) > 3).mean():6.2%}")
print(f"long-run update size, sqrt((1 - b1) / (1 + b1)) = {np.sqrt((1 - b1) / (1 + b1)):.3f}")

Full script: warmup_noise.py

step     1: typical update 1.000 x eta;  parameters with a rate over 3x its true value: 25.97%
step     2: typical update 0.708 x eta;  parameters with a rate over 3x its true value: 10.42%
step     5: typical update 0.452 x eta;  parameters with a rate over 3x its true value:  0.97%
step    10: typical update 0.329 x eta;  parameters with a rate over 3x its true value:  0.02%
step   100: typical update 0.231 x eta;  parameters with a rate over 3x its true value:  0.00%
step  1000: typical update 0.229 x eta;  parameters with a rate over 3x its true value:  0.00%
step  5000: typical update 0.229 x eta;  parameters with a rate over 3x its true value:  0.00%
long-run update size, sqrt((1 - b1) / (1 + b1)) = 0.229

Two things go wrong at the start, one per average.

• The step is too big. At the first step every parameter moves by the full $\eta$, in a random direction. By step 100 the typical move on noise is $0.23\,\eta$. That long-run value follows from $\hat m$: it averages about $(1 + \beta_1)/(1 - \beta_1) = 19$ noisy gradients, which shrinks their size by $\sqrt{19}$. At step 1 it has averaged one, so Adam’s steps on noise are 4.4 times their later size, for every parameter at once, at the moment the network is least trained.
• Some rates are wildly off. Each parameter’s effective rate is $\eta / \sqrt{\hat v}$. With one or two samples, $\hat v$ is just the square of a gradient or two, and when that happens to be small, the rate is huge. At step 1, a quarter of the parameters have a rate more than 3 times its true value; by step 10, almost none do. This is the effect Liu et al. (2019) identified. Their fix, RAdam, turns the rescaling off until $v$ has seen enough gradients.

Warmup handles both: a small learning rate until the averages are reliable.

That’s not the only reason. Early gradients in deep Transformers can be very large, particularly with some placements of the normalization layers, and warmup protects against those too (Xiong et al., 2020; The Residual Stream). A warmup is cheap insurance against both.

A related guard is gradient clipping: if the whole gradient’s norm exceeds a threshold, scale it down to that threshold before the step. It stops one bad batch from producing one enormous update. Llama 2 and GPT-3 (Brown et al., 2020) both clipped at a norm of 1.0. The same worry explains the $\beta_2 = 0.95$ both used instead of 0.999: a shorter memory lets $\hat v$ catch up quickly when gradients suddenly grow, so a stale, small $\hat v$ can’t turn a gradient spike into a huge step.

Weight decay, decoupled: AdamW

Features and Overfitting’s L2 penalty adds $\tfrac{\lambda}{2}\|W\|^2$ to the loss, which adds $\lambda W$ to the gradient. Folding that into Adam’s gradient means the penalty’s pull passes through Adam’s rescaling: the decay term is divided by $\sqrt{\hat v}$ like everything else. AdamW (Loshchilov and Hutter, 2017) instead applies the decay directly to the weights, outside the rescaling:

(Here $\lambda$ is the decay strength, as in Features and Overfitting.) Every step shrinks every weight by the same fraction, $\eta\lambda$, before Adam’s step, however large or small its gradients are. “Decoupled” means decoupled from Adam’s rescaling, not from the learning rate: in PyTorch’s AdamW, as here, the shrink per step is $\eta\lambda$, so changing $\eta$ changes the decay too. In optim.py, it’s the decoupled switch. Here are the ways to regularize the router, on mini-batches of 16 as real training uses, each penalty chosen on validation and the chosen one reported on test:

class L2SGD(SGD):
    """Features and Overfitting's L2 penalty with plain mini-batch gradient descent."""
    def __init__(self, eta, lam):
        super().__init__(eta)
        self.lam = lam
    def step(self, p, g, key):
        super().step(p, g + (self.lam * p if key == "W" else 0), key)

def train(opt, X, y, seed, steps=STEPS, batch=16):
    rng = np.random.default_rng(seed)
    W, b = np.zeros((X.shape[1], 4)), np.zeros(4)
    for _ in range(steps):
        i = rng.choice(len(y), batch, replace=False)
        G = (softmax(X[i] @ W + b) - np.eye(4)[y[i]]) / batch
        opt.step(W, X[i].T @ G, "W")                    # the optimizer adds any decay itself
        opt.step(b, G.sum(0), "b")
    return W, b

METHODS = [("no decay", lambda lam: Adam(0.01), [0]),
           ("SGD + L2", lambda lam: L2SGD(0.5, lam), [0.0003, 0.001, 0.003]),
           ("Adam + L2", lambda lam: Adam(0.01, decay=lam, decoupled=False), [0.0001, 0.0003, 0.001, 0.003]),
           ("AdamW", lambda lam: Adam(0.01, decay=lam), [0.1, 0.3, 1])]

Full script: decay.py

Common words are those in 5 or more of the 120 training questions, and every other word is rare. Each $\pm$ is the standard deviation of the 20 per-split differences scaled by $\sqrt{1/20 + 40/120}$, which corrects for the overlap between splits.

3,000 steps, mini-batches of 16
no decay  λ 0      val CE 0.881   weight size: common words 0.852, rare words 0.927, ratio 0.92
SGD + L2  λ 0.0003 val CE 0.914   weight size: common words 0.769, rare words 0.410, ratio 1.88
SGD + L2  λ 0.001  val CE 0.888   weight size: common words 0.626, rare words 0.328, ratio 1.91
SGD + L2  λ 0.003  val CE 0.893   weight size: common words 0.496, rare words 0.245, ratio 2.02
Adam + L2 λ 0.0001 val CE 0.912   weight size: common words 0.818, rare words 0.523, ratio 1.57
Adam + L2 λ 0.0003 val CE 0.907   weight size: common words 0.747, rare words 0.428, ratio 1.75
Adam + L2 λ 0.001  val CE 0.889   weight size: common words 0.624, rare words 0.333, ratio 1.88
Adam + L2 λ 0.003  val CE 0.894   weight size: common words 0.497, rare words 0.249, ratio 2.00
AdamW     λ 0.1    val CE 0.822   weight size: common words 0.613, rare words 0.664, ratio 0.92
AdamW     λ 0.3    val CE 0.819   weight size: common words 0.502, rare words 0.525, ratio 0.96
AdamW     λ 1      val CE 0.894   weight size: common words 0.307, rare words 0.312, ratio 0.98
chosen on validation: SGD + L2 λ 0.001, Adam + L2 λ 0.001, AdamW λ 0.3
on test, at the chosen λ:  no decay CE 0.870 acc 0.691, SGD + L2 CE 0.896 acc 0.640, Adam + L2 CE 0.897 acc 0.631, AdamW CE 0.822 acc 0.688
test, AdamW minus SGD + L2: CE -0.074 ± 0.042, accuracy +0.047 ± 0.040
test, AdamW minus no decay: CE -0.048 ± 0.047, accuracy -0.004 ± 0.017
test, SGD + L2 minus no decay: CE +0.026 ± 0.041, accuracy -0.051 ± 0.034

And the same experiment trained four times as long, since Features and Overfitting found that a fixed step count can act as hidden early stopping (python decay.py 12000):

12,000 steps, mini-batches of 16
no decay  λ 0      val CE 1.128   weight size: common words 1.342, rare words 1.511, ratio 0.89
SGD + L2  λ 0.0003 val CE 0.912   weight size: common words 0.768, rare words 0.423, ratio 1.82
SGD + L2  λ 0.001  val CE 0.887   weight size: common words 0.624, rare words 0.329, ratio 1.90
SGD + L2  λ 0.003  val CE 0.892   weight size: common words 0.496, rare words 0.246, ratio 2.02
Adam + L2 λ 0.0001 val CE 0.958   weight size: common words 0.904, rare words 0.516, ratio 1.76
Adam + L2 λ 0.0003 val CE 0.916   weight size: common words 0.771, rare words 0.428, ratio 1.81
Adam + L2 λ 0.001  val CE 0.890   weight size: common words 0.628, rare words 0.333, ratio 1.89
Adam + L2 λ 0.003  val CE 0.894   weight size: common words 0.498, rare words 0.249, ratio 2.01
AdamW     λ 0.1    val CE 0.966   weight size: common words 1.025, rare words 1.093, ratio 0.94
AdamW     λ 0.3    val CE 0.855   weight size: common words 0.764, rare words 0.818, ratio 0.93
AdamW     λ 1      val CE 0.867   weight size: common words 0.340, rare words 0.361, ratio 0.94
chosen on validation: SGD + L2 λ 0.001, Adam + L2 λ 0.001, AdamW λ 0.3
on test, at the chosen λ:  no decay CE 1.098 acc 0.699, SGD + L2 CE 0.894 acc 0.634, Adam + L2 CE 0.892 acc 0.630, AdamW CE 0.856 acc 0.676
test, AdamW minus SGD + L2: CE -0.037 ± 0.060, accuracy +0.042 ± 0.040
test, AdamW minus no decay: CE -0.242 ± 0.098, accuracy -0.022 ± 0.029
test, SGD + L2 minus no decay: CE -0.205 ± 0.095, accuracy -0.065 ± 0.035

This isn’t the story usually told, so read it carefully.

• Adam + L2 isn’t broken here. It is L2. Its weights and its validation loss match plain SGD + L2 almost exactly. On a problem it can converge on, any optimizer that follows the penalized gradient ends near the same minimum; Adam only changes the path. Under L2, rare words end up with weights about half the size of common ones: a well-evidenced weight pushes back against the penalty harder.
• AdamW keeps Adam’s own pattern. Even with no decay at all, Adam’s rare and common words grow to the same size (ratio 0.92), because Adam moves every weight about the same distance per step, whatever its gradients. Decoupled decay shrinks all of them by the same fraction, so the pattern survives. That’s a lighter penalty on rare words than L2’s.
• No decay is the honest baseline, and at 3,000 steps it’s hard to beat. Its validation loss, 0.881, is already better than L2’s 0.888, because stopping at 3,000 steps is itself a regularizer. On test, AdamW ties it (cross-entropy $-0.048 \pm 0.047$). AdamW does beat L2 on test, by $0.074 \pm 0.042$ in cross-entropy and $0.047 \pm 0.040$ in accuracy, suggestive rather than decisive.
• Train four times as long and decay earns its place. Without decay, the weights keep growing and the test cross-entropy climbs to 1.098; AdamW and L2 now beat it by 0.24 and 0.21. AdamW and L2 tie ($-0.037 \pm 0.060$), and AdamW’s weights are still growing (0.50 to 0.76 at the chosen $\lambda$).
• Decay buys calibrated probabilities, not more right answers. No version of it improves accuracy over no decay, and L2 lowers it by about 0.05 to 0.07. The gain is all in how well the model’s confidence matches how often it is right, which is what cross-entropy rewards and accuracy ignores.

The general lesson for AdamW: $\lambda$, $\eta$ and the length of training set the regularization together, because Adam’s steps don’t shrink as the gradients do. Loshchilov and Hutter’s case for it is about the realistic setting, where training never converges and the learning rate follows a schedule. There, L2 through Adam gives weights with large gradient histories almost no decay, and AdamW makes weight decay behave the same for every weight and easier to tune. Llama 2 used AdamW with $\beta_2 = 0.95$, $\epsilon = 10^{-5}$ and weight decay 0.1, with a cosine schedule that decays the learning rate to 10% of its peak. GPT-3 used the same $\beta$s with $\epsilon = 10^{-8}$. Values like these are common in language-model pretraining, though not universal.

What Adam costs

Adam keeps two numbers, $m$ and $v$, for every parameter, in FP32 for accuracy. With mixed-precision training, which keeps BF16 copies for the fast math and an FP32 master copy for the updates (PyTorch from the Inside), the bill per parameter is (the ZeRO paper counts FP16 copies; BF16 is the same size):

# Bytes of training state per parameter, mixed-precision training (ZeRO's accounting).
params = 8e9                                   # an 8B model
state = {"BF16 weights": 2, "BF16 gradients": 2,
         "FP32 master weights": 4, "Adam m (FP32)": 4, "Adam v (FP32)": 4}
for name, b in state.items():
    print(f"{name:20s} {b:2d} bytes/param  {b * params / 1e9:5.0f} GB")
total = sum(state.values())
print(f"{'total':20s} {total:2d} bytes/param  {total * params / 1e9:5.0f} GB, "
      f"{total * params / 80e9:.1f} H100s of memory before a single activation")

Full script: adam_memory.py

BF16 weights          2 bytes/param     16 GB
BF16 gradients        2 bytes/param     16 GB
FP32 master weights   4 bytes/param     32 GB
Adam m (FP32)         4 bytes/param     32 GB
Adam v (FP32)         4 bytes/param     32 GB
total                16 bytes/param    128 GB, 1.6 H100s of memory before a single activation

Twelve of the sixteen bytes are optimizer state by ZeRO’s accounting: the FP32 master copy plus Adam’s $m$ and $v$. Adam’s own share is the 8 bytes of $m$ and $v$; plain SGD with momentum would still need the master copy and the velocity, 8. An 8B model needs 128 GB before storing a single activation, more than one H100’s 80 GB. The weights alone, the 16 GB that inference needs (as in AI Engineering, from the Inside), are an eighth of it. That’s why training even a mid-sized model needs several GPUs, and why ZeRO-style sharding (Rajbhandari et al., 2020) splits the optimizer state across them rather than copying it to each. It also motivates optimizers that keep less state, one of which, Muon, comes up in Anatomy of a Frontier Model.

Where this leads

• Dropout and Normalization covers the other tools that keep deep networks trainable: dropout, and the normalization layers (BatchNorm, LayerNorm, RMSNorm) that re-standardize activations inside the network, the in-network version of Learning Is Optimization’s feature scaling.
• PyTorch from the Inside builds these pieces in PyTorch: autograd instead of hand-written backward passes, torch.optim.AdamW instead of optim.py, and the mixed-precision formats behind the memory bill above.
• Part VII: Training, fine-tuning and alignment trains real models with AdamW, warmup and a cosine schedule, where the learning rate turns out to be the hyperparameter that matters most.

Takeaways

• Each ReLU layer multiplies a signal’s mean square by $n\sigma^2/2$, in both the forward and backward passes. Over many layers anything but 1 vanishes or explodes. He initialization, $\sigma^2 = 2/n$, makes it 1; Xavier’s $1/n$ is the version for activations that don’t halve the signal.
• Too-small initialization is fatal: a 10-layer router starting at $\sigma = 0.01$ stayed at a loss of $\log 4$, predicting uniformly, even with Adam. Gradients that small carry no usable signal, so no optimizer can rescale its way out.
• Momentum cuts the steps needed on a quadratic from order $\kappa$ to order $\sqrt\kappa$: 882 to 108 on Learning Is Optimization’s problem. It also stays stable at steps plain descent can’t take, but needs tuning.
• Adam divides each parameter’s step by its own gradient scale, so its step is about $\eta$ in the parameter’s units: a partial, per-parameter version of standardization. Bias correction removes the zero start of both averages; without it the early steps are up to 6.6 times too large.
• Adam wins when parameters have very different gradient scales. On 0/1 features it doesn’t; add one unscaled feature and it reaches the minimum in 300 steps while gradient descent hasn’t in 3,000.
• Early in training, Adam’s steps on noise are about 4× their long-run size, and some parameters’ rates are far too high because $\hat v$ rests on one or two samples. Warmup, gradient clipping and $\beta_2 = 0.95$ protect the start.
• AdamW decays every weight by the same fraction per step. At convergence, Adam + L2 is plain L2; AdamW is a different penalty, and its strength depends on $\eta$, $\lambda$ and the length of training together. On the router, decay improved calibration, not accuracy, and only mattered once training ran long enough to overfit.
• Mixed-precision Adam needs 16 bytes per parameter, 12 of them optimizer state (including the FP32 master copy): 128 GB for an 8B model.

References

Start with Goh’s Why Momentum Really Works, an interactive essay that makes this post’s momentum section visual. Then read Kingma and Ba for Adam and Loshchilov and Hutter for AdamW, both short. He et al. §2.2 is the initialization derivation.

• Glorot and Bengio, Understanding the Difficulty of Training Deep Feedforward Neural Networks (AISTATS, 2010): signal variance through layers, and Xavier initialization.
• He, Zhang, Ren and Sun, Delving Deep into Rectifiers (ICCV, 2015): the $2/n$ initialization for ReLU networks.
• Polyak, Some Methods of Speeding Up the Convergence of Iteration Methods (USSR Computational Mathematics and Mathematical Physics, 1964): the heavy-ball method and its $\sqrt\kappa$ rate.
• Sutskever, Martens, Dahl and Hinton, On the Importance of Initialization and Momentum in Deep Learning (ICML, 2013): both of this post’s topics, for deep networks.
• Goh, Why Momentum Really Works (Distill, 2017): an interactive treatment of momentum on quadratics.
• Kingma and Ba, Adam: A Method for Stochastic Optimization (ICLR, 2015): Adam and its bias correction.
• Loshchilov and Hutter, Decoupled Weight Decay Regularization (ICLR, 2019): AdamW.
• Liu et al., On the Variance of the Adaptive Learning Rate and Beyond (ICLR, 2020): the variance of Adam’s early steps, as a reason for warmup.
• Xiong et al., On Layer Normalization in the Transformer Architecture (ICML, 2020): normalization placement and the need for warmup.
• Brown et al. (OpenAI), Language Models are Few-Shot Learners (NeurIPS, 2020): GPT-3’s training settings, Adam with $\beta_1 = 0.9$, $\beta_2 = 0.95$, $\epsilon = 10^{-8}$, and gradient clipping at a norm of 1.0.
• Touvron et al. (Meta AI), Llama 2: Open Foundation and Fine-Tuned Chat Models (2023): AdamW settings, warmup, cosine schedule and gradient clipping for pretraining.
• Rajbhandari, Rasley, Ruwase and He, ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (SC, 2020): the 16 bytes per parameter of mixed-precision Adam.