Adam, Momentum, Schedulers

Two model parameters are each one unit from their best value. The loss gradient asks you to change one 100 times more than the other. How big should a training step be?

The linear algebra lesson showed how a matrix can stretch different directions by different amounts. The same imbalance appears in optimization: a learning rate that barely moves one parameter can send another past its target.

A gradient gives the loss's local slope in each parameter direction. An optimizer uses those slopes to choose an update. Moving against the gradient lowers a smooth loss for a sufficiently small step, but a large step or an update based on past gradients needn't lower it.^{[1]Reference 1Deep Learninghttps://www.deeplearningbook.org/}

View larger

The same two coordinates will show what changes when an update remembers past gradients, rescales them, shrinks the weights, or changes pace. The examples use NumPy on a CPU. The final training loop uses PyTorch. No model download or GPU is needed.

A 100-to-1 valley

Let a and b be the two parameter errors. Their target is [0, 0].

Use a quadratic loss so every gradient and update is easy to check:

The coefficient 100 means the slope changes 100 times faster along a than along b. This is curvature. Feature importance requires a separate analysis of the task and model. In matrix form:

$\mathrm{diag}(100, 1)$ has condition number 100. A single step size is therefore fighting the same uneven stretch that made inverse-like solves fragile in the last chapter.

Differentiate to get

At [1, 1], the loss is 50.5 and the gradient is [100, 1]. Equal coordinate error hasn't produced equal correction requests: a asks for a step 100 times larger. What should happen if one learning rate serves both?

Run this small calculation and inspect both coordinates, not only the loss.

import numpy as np

def loss(w: np.ndarray) -> float:
    return float(0.5 * (100.0 * w[0] ** 2 + w[1] ** 2))

def gradient(w: np.ndarray) -> np.ndarray:
    return np.array([100.0 * w[0], w[1]])

w = np.array([1.0, 1.0])
g = gradient(w)
next_w = w - 0.018 * g

print("start loss", loss(w))
print("gradient", g.tolist())
print("one step", next_w.round(3).tolist())
print("next loss", round(loss(next_w), 3))

start loss 50.5
gradient [100.0, 1.0]
one step [-0.8, 0.982]
next loss 32.482

The loss falls after one step, but the coordinates tell a less comfortable story. a crosses from 1.0 to -0.8, while b only reaches 0.982.

That crossing is the zigzag: the steep error bounces across zero while the shallow error barely walks. A lower rate would calm the wall, but it would also slow the floor.

SGD gives every coordinate one learning rate

Before naming the update stochastic gradient descent (SGD), separate two choices: where the gradient comes from and how the optimizer applies it.

This walkthrough uses a deterministic two-coordinate gradient so the loss geometry stays visible. It stands in for a minibatch gradient without adding sampling noise. Real training adds that noise, so log batch size and gradient statistics when comparing optimizers.

Once the gradient source is fixed, SGD applies a single rule:

Here, $w_t$ is the parameter vector after step $t$ and $g_t$ is the minibatch gradient. $\eta$ (eta) is the learning rate, the step-size multiplier.

Try three rates for eight steps on the same valley. Predict which rate keeps a stable and which moves b farther.

import numpy as np

def loss(w: np.ndarray) -> float:
    return float(0.5 * (100.0 * w[0] ** 2 + w[1] ** 2))

def gradient(w: np.ndarray) -> np.ndarray:
    return np.array([100.0 * w[0], w[1]])

def run_sgd(lr: float, steps: int = 8) -> tuple[np.ndarray, float]:
    w = np.array([1.0, 1.0])
    for _ in range(steps):
        w -= lr * gradient(w)
    return w, loss(w)

for lr in (0.005, 0.018, 0.021):
    w, final_loss = run_sgd(lr)
    print(f"lr={lr:.3f} w={w.round(3).tolist()} loss={final_loss:.3f}")

lr=0.005 w=[0.004, 0.961] loss=0.462
lr=0.018 w=[0.168, 0.865] loss=1.781
lr=0.021 w=[2.144, 0.844] loss=230.105

The small rate is calm, but b is still 0.961 after eight steps. It also reports the lowest loss of these three runs, so a bigger rate isn't automatically faster learning.

The middle rate drives down a while b still crawls. To see the stability boundary, substitute the gradient into the update: a_next = (1 - 100 * lr) * a. Its magnitude shrinks only when abs(1 - 100 * lr) < 1, or 0 < lr < 0.02. At exactly 0.02, a flips sign without shrinking. At 0.021, the factor is -1.1, so the oscillation grows.

Momentum remembers directions that agree

SGD reacts only to the current gradient. When the steep coordinate flips sign, it forgets that the previous step was already heading back across the wall.

Momentum keeps a signed moving average of recent gradients. Use the exponential-moving-average convention here because it matches Adam's first buffer. This is a normalized teaching form of Polyak's heavy-ball idea, described in his 1964 paper on accelerating iterative methods:

The coefficient $\beta$ controls memory. With $\beta = 0.9$, the new gradient contributes 0.1 and the previous direction contributes 0.9.

PyTorch's torch.optim.SGD uses an unnormalized buffer: with default dampening, it starts at the first gradient and then computes buffer = beta * buffer + gradient. Our m is (1 - beta) times that buffer. To reproduce our updates in that convention, multiply our learning rate by (1 - beta). Check initialization and dampening too before translating another implementation. PyTorch documents these choices.

Feed momentum a prescribed four-gradient history: large alternating values for a, small positive values for b. These are inputs chosen to expose the buffers, not gradients produced by a momentum run. We'll reuse them for RMSProp and Adam. Which signs cancel, and which reinforce one another?

import numpy as np

beta = 0.9
gradients = [
    np.array([100.0, 1.00]),
    np.array([-80.0, 0.98]),
    np.array([70.0, 0.96]),
    np.array([-60.0, 0.94]),
]
m = np.zeros(2)

for step, g in enumerate(gradients, start=1):
    m = beta * m + (1.0 - beta) * g
    print(step, "gradient", g.round(2).tolist(), "memory", m.round(3).tolist())

1 gradient [100.0, 1.0] memory [10.0, 0.1]
2 gradient [-80.0, 0.98] memory [1.0, 0.188]
3 gradient [70.0, 0.96] memory [7.9, 0.265]
4 gradient [-60.0, 0.94] memory [1.11, 0.333]

The first buffer stays positive throughout this trace, even when the current gradient is negative. Opposing gradients partially cancel: on step two, 0.9 * 10 + 0.1 * (-80) = 1. The second buffer grows from 0.1 to about 0.333 because every gradient reinforces the same sign. These signed averages are what momentum subtracts from the parameters after multiplying by the learning rate.

Momentum still needs a suitable learning rate for each surface. It changes the path: useful repeated direction builds up, while alternating motion is damped. Now compare measured paths on the same loss instead of assuming that adding a buffer guarantees a win.

import numpy as np

def loss(w: np.ndarray) -> float:
    return float(0.5 * (100.0 * w[0] ** 2 + w[1] ** 2))

def gradient(w: np.ndarray) -> np.ndarray:
    return np.array([100.0 * w[0], w[1]])

def sgd(steps: int, lr: float) -> np.ndarray:
    w = np.array([1.0, 1.0])
    for _ in range(steps):
        w -= lr * gradient(w)
    return w

def momentum(steps: int, lr: float, beta: float = 0.9) -> np.ndarray:
    w = np.array([1.0, 1.0])
    m = np.zeros_like(w)
    for _ in range(steps):
        m = beta * m + (1.0 - beta) * gradient(w)
        w -= lr * m
    return w

for name, w in [
    ("sgd", sgd(40, lr=0.018)),
    ("momentum", momentum(40, lr=0.018)),
]:
    print(name, "w", w.round(4).tolist(), "loss", round(loss(w), 4))

sgd w [0.0001, 0.4836] loss 0.1169
momentum w [0.0307, 0.5324] loss 0.1889

At this unchanged numeric learning rate, SGD ends lower than our EMA-style momentum run.

Momentum changes the update scale and usually needs its own learning-rate and memory-coefficient sweep. Use those measurements to see whether its altered path improves the task.

RMSProp rescales by recent gradient magnitude

Momentum smooths signed direction, but it still sends every coordinate through one global learning rate.

RMSProp keeps a second, unsigned history: an exponential moving average of squared gradients. The method was documented by Tieleman and Hinton in their 2012 RMSProp lecture note.

The s buffer forgets old squared magnitudes. A large recent gradient makes that coordinate's denominator large, so its update is scaled down.

Unlike momentum's signed buffer, squaring keeps s from canceling when signs alternate. Each RMSProp coordinate keeps the current gradient's sign, but rescaling coordinates differently changes the direction of the whole update vector. This basic form has no Adam-style bias correction.

PyTorch calls rho alpha; its documented RMSprop defaults are lr=0.01, alpha=0.99, and eps=1e-8 (official docs).

PyTorch adds $\epsilon$ after the square root; some other libraries add it inside. Copying $\epsilon$ between them can silently change the update.

Prediction check: This fixture uses rho=0.9 so four steps make the state change easy to inspect. Before running it, predict whether the first update magnitudes will still differ by 100:1.

Use the same deliberately uneven gradient history to test that prediction:

import numpy as np

rho, lr, eps = 0.9, 0.01, 1e-8
gradients = [
    np.array([100.0, 1.00]),
    np.array([-80.0, 0.98]),
    np.array([70.0, 0.96]),
    np.array([-60.0, 0.94]),
]
s = np.zeros(2)

for step, g in enumerate(gradients, start=1):
    s = rho * s + (1.0 - rho) * g**2
    update = lr * g / (np.sqrt(s) + eps)
    print(step, "rms", np.sqrt(s).round(3).tolist(), "step", update.round(5).tolist())

1 rms [31.623, 0.316] step [0.03162, 0.03162]
2 rms [39.243, 0.431] step [-0.02039, 0.02272]
3 rms [43.313, 0.51] step [0.01616, 0.01884]
4 rms [45.259, 0.567] step [-0.01326, 0.01657]

The first row gives both coordinates nearly the same update magnitude, despite the 100:1 raw gradients. With rho=0.9, the first denominator is about sqrt(0.1) * abs(g), so the update magnitude is lr / sqrt(0.1), about 0.03162. Later, the first coordinate's denominator stays large, but its update follows the latest gradient sign. Here and below, update means the vector we subtract from the parameters.

That is the handoff to Adam: add momentum's signed m buffer to this RMSProp-like v buffer, then correct both buffers for their zero start.

Adam stores direction and squared-gradient scale

The Adam optimizer (Adaptive Moment Estimation) keeps the signed average for direction and adds a second moving average for scale: the elementwise square of gradients. Its original formulation is:^{[2]Reference 2Adam: A Method for Stochastic Optimizationhttps://arxiv.org/abs/1412.6980}

Read the four lines as two decisions. m chooses a direction from recent signs; v measures how large recent gradients have been. Bias correction turns those startup-shrunken buffers into the hatted estimates used by the update.

The symbols have precise jobs:

Kingma and Ba's usual defaults are $\beta_1 = 0.9$, $\beta_2 = 0.999$, and $\epsilon = 10^{-8}$, with $\epsilon$ added outside the square root in the denominator.^{[2]Reference 2Adam: A Method for Stochastic Optimizationhttps://arxiv.org/abs/1412.6980}

PyTorch's documented AdamW constructor starts with lr=1e-3, the same betas and eps, and weight_decay=1e-2. Tune these settings for the model and dataset, then evaluate the resulting training behavior. PyTorch AdamW docs

v is often casually called variance. That wording is misleading here: Adam averages g ** 2; it doesn't subtract the mean gradient to calculate statistical variance.

The first step is a useful sanity check. With bias correction, a nonzero gradient of 100 and one of 1 both produce an update close to one learning rate.

Prediction check: Before running the code, predict what will be equal and what will still differ in the two moment buffers.

import numpy as np

g = np.array([100.0, 1.0])
beta1, beta2, lr, eps = 0.9, 0.999, 0.01, 1e-8
m = (1.0 - beta1) * g
v = (1.0 - beta2) * g**2
m_hat = m / (1.0 - beta1)
v_hat = v / (1.0 - beta2)
update = lr * m_hat / (np.sqrt(v_hat) + eps)

print("m_hat", m_hat.tolist())
print("v_hat", v_hat.tolist())
print("update", update.round(6).tolist())

m_hat [100.0, 1.0]
v_hat [10000.0, 1.0]
update [0.01, 0.01]

View larger

Now replay the same prescribed history used for momentum and RMSProp. Look especially at steps two and four, when a's raw gradient is negative:

import numpy as np

gradients = [
    np.array([100.0, 1.00]),
    np.array([-80.0, 0.98]),
    np.array([70.0, 0.96]),
    np.array([-60.0, 0.94]),
]
beta1, beta2, lr, eps = 0.9, 0.999, 0.01, 1e-8
m = np.zeros(2)
v = np.zeros(2)

for t, g in enumerate(gradients, start=1):
    m = beta1 * m + (1.0 - beta1) * g
    v = beta2 * v + (1.0 - beta2) * g**2
    m_hat = m / (1.0 - beta1**t)
    v_hat = v / (1.0 - beta2**t)
    update = lr * m_hat / (np.sqrt(v_hat) + eps)
    print(t, "m_hat", m_hat.round(3).tolist(), "step", update.round(5).tolist())

1 m_hat [100.0, 1.0] step [0.01, 0.01]
2 m_hat [5.263, 0.989] step [0.00058, 0.00999]
3 m_hat [29.151, 0.979] step [0.00346, 0.00998]
4 m_hat [3.228, 0.967] step [0.00041, 0.00997]

Unlike RMSProp, Adam still subtracts a positive value from a on steps two and four. Its signed average hasn't turned negative yet. Momentum can keep a parameter moving against the current gradient, so smoothing doesn't guarantee descent on every step. The consistently positive b history keeps its update near 0.01.

Adam isn't finding curvature or proving which feature matters. It transforms recent gradient history into coordinate-wise steps. That distinction matters when you diagnose a run: adaptive scaling isn't a geometry oracle.

Changing a loss from a mean to a sum multiplies its gradient by the number of items being averaged. SGD's update grows by that factor at the same learning rate. Adam nearly cancels a fixed positive factor applied to the entire gradient history: m_hat scales by that factor and sqrt(v_hat) does too. This cancellation is exact only when epsilon is zero and the denominator is nonzero.^{[2]Reference 2Adam: A Method for Stochastic Optimizationhttps://arxiv.org/abs/1412.6980}

With nonzero epsilon, tiny gradients behave differently. At the first step from zero state, the update is lr * g / (abs(g) + eps). Keep lr=0.01 and eps=1e-8, then shrink the gradient:

lr, eps = 0.01, 1e-8
for g in [1.0, 1e-8, 1e-10]:
    update = lr * g / (abs(g) + eps)
    print(f"gradient {g:.0e} update {update:.8f}")

gradient 1e+00 update 0.01000000
gradient 1e-08 update 0.00500000
gradient 1e-10 update 0.00009901

The gradient equal to epsilon gets half the learning-rate-sized update. A changing scale, clipping, coupled L2, or a switch in loss reduction midway through a run also breaks the simple cancellation argument. Record the objective's reduction and denominator along with the optimizer settings.

Bias correction repairs zero-initialized state

Both moving averages begin at zero, so early values are pulled toward that initialization.

With Adam's usual coefficients, the first raw first-moment buffer is 0.1 * g, while the first raw second-moment buffer is 0.001 * g**2. Dividing those raw values would make the startup update too large. Bias correction removes that startup pull.

Compare the first update with and without correction for one scalar parameter:

import math

gradient = 4.0
beta1, beta2, lr, eps = 0.9, 0.999, 0.01, 1e-8
m = (1.0 - beta1) * gradient
v = (1.0 - beta2) * gradient**2

uncorrected_step = lr * m / (math.sqrt(v) + eps)
m_hat = m / (1.0 - beta1)
v_hat = v / (1.0 - beta2)
corrected_step = lr * m_hat / (math.sqrt(v_hat) + eps)

print("uncorrected step", round(uncorrected_step, 5))
print("corrected step", round(corrected_step, 5))

uncorrected step 0.03162
corrected step 0.01

Why divide by 1 - beta**t? If every gradient were the same value g, the first two averages would be (1 - beta) * g and (1 - beta**2) * g. After t steps, the weights on observed gradients sum to 1 - beta**t; division restores that sum to one. It removes the zero-start effect, not the uncertainty or lag in a changing gradient history.

AdamW separates data gradients from weight shrinkage

Regularization often nudges parameter magnitudes downward. Define the L2 penalty here as $\frac{\lambda}{2}\lVert w \rVert_2^2$, whose gradient is $\lambda w$.

With ordinary SGD, that penalty can produce the same shrinkage behavior as weight decay after adjusting coefficients. With adaptive optimizers, those operations aren't equivalent: putting the penalty inside Adam's gradient also sends it through coordinate-wise scaling.^{[3]Reference 3Decoupled Weight Decay Regularizationhttps://arxiv.org/abs/1711.05101}

AdamW applies a decoupled decay directly to the prior parameter value:

$\lambda$ is the weight-decay coefficient. The data-gradient update and shrinkage term are visible separately.

The difference appears even on the first step with a zero data gradient and zero moment buffers. Compare a parameter of 10.0 under the two rules:

theta = 10.0
lr = 0.1
weight_decay = 0.01
eps = 1e-8

adamw_theta = theta - lr * weight_decay * theta

# Gradient of (weight_decay / 2) * theta**2.
coupled_gradient = weight_decay * theta
coupled_adam_step = lr * coupled_gradient / (abs(coupled_gradient) + eps)
coupled_theta = theta - coupled_adam_step

print("AdamW with zero data gradient", round(adamw_theta, 4))
print("Adam plus coupled L2 first step", round(coupled_theta, 4))

AdamW with zero data gradient 9.99
Adam plus coupled L2 first step 9.9

The penalty has entered both Adam buffers in the coupled version. In AdamW, it enters neither: the parameter shrinks by the separate factor 1 - lr * weight_decay. Decoupled doesn't mean independent of the schedule, because lr still multiplies the decay. In PyTorch, a zero gradient tensor also differs from grad=None: the optimizer skips parameters with no gradient, including their decay. AdamW API

In PyTorch 2.14, torch.optim.Adam(..., decoupled_weight_decay=True) also selects AdamW behavior. The default is False, so the weight_decay argument alone still selects coupled L2. Inspect the actual constructor rather than inferring the rule from the optimizer's name. Adam API

Schedules change pace across the run

An optimizer decides how to interpret gradients at one step. A learning-rate schedule changes the global multiplier $\eta$ across many steps.

The original Transformer training recipe used Adam with $\beta_1 = 0.9$, $\beta_2 = 0.98$, and $\epsilon = 10^{-9}$, plus a schedule that increased the learning rate linearly for 4000 warmup steps and then decreased it in proportion to the inverse square root of the step number.^{[4]Reference 4Attention Is All You Need.https://arxiv.org/abs/1706.03762} Those betas and $\epsilon$ aren't Kingma and Ba's defaults.

Treat the whole recipe as evidence that one successful run existed, not as a universal setting. Cosine annealing and warm restarts were studied separately by Loshchilov and Hutter; a schedule still needs measurement on your task.

Warmup followed by cosine decay is a shape you can inspect by hand. In this zero-based convention, step counts optimizer updates, not batches or epochs. Update zero uses rate zero, update warmup_steps reaches the peak, and update total_steps - 1 reaches the final rate. Other recipes start at a small nonzero rate, so check the indexing when porting a schedule.

import math

def warmup_cosine(step: int, total_steps: int, warmup_steps: int, peak_lr: float, final_lr: float) -> float:
    if not 0 <= warmup_steps < total_steps - 1:
        raise ValueError("leave at least two schedule points at or after the peak")
    if not 0 <= step < total_steps:
        raise ValueError("step must be within the run")
    if not 0 <= final_lr <= peak_lr or not math.isfinite(peak_lr):
        raise ValueError("rates must be finite and satisfy 0 <= final_lr <= peak_lr")
    if step < warmup_steps:
        return peak_lr * step / warmup_steps
    progress = (step - warmup_steps) / (total_steps - warmup_steps - 1)
    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
    return final_lr + (peak_lr - final_lr) * cosine

values = [
    warmup_cosine(step, total_steps=10, warmup_steps=2, peak_lr=3e-4, final_lr=3e-5)
    for step in range(10)
]
print([f"{lr:.6f}" for lr in values])

['0.000000', '0.000150', '0.000300', '0.000287', '0.000249', '0.000195', '0.000135', '0.000081', '0.000043', '0.000030']

Here the peak is at index 2 and the final rate at index 9. With gradient accumulation, several batches can contribute to one optimizer update; don't advance this schedule on each of those batches. Log the actual rate used alongside the loss so an off-by-one error is visible.

A zero learning rate leaves the parameters unchanged, but an Adam step still updates its moment buffers and step count when gradients are present. Here, update zero contributes history for later updates. A zero rate is different from skipping optimizer.step().

Accumulate the objective you intended

Gradient accumulation adds gradients from several microbatches before one optimizer update. For equally sized microbatches with a mean loss, divide each loss by the number of microbatches. With unequal sizes, weight each mean by its share of the total items. Otherwise you're averaging microbatch means, giving a smaller microbatch's items more weight.

Use a scalar prediction w=1 and squared-error targets [0, 2, 4]. Split them into one item and two items. The full-batch mean gradient is (2 - 2 - 6) / 3 = -2. Predict what happens if both microbatch means receive weight 1/2.

import torch

targets = torch.tensor([0.0, 2.0, 4.0])
for mode in ["full batch", "weighted accumulation", "mean of means"]:
    w = torch.tensor(1.0, requires_grad=True)
    if mode == "full batch":
        ((w - targets) ** 2).mean().backward()
    else:
        for microbatch in [targets[:1], targets[1:]]:
            weight = len(microbatch) / len(targets) if mode == "weighted accumulation" else 0.5
            (((w - microbatch) ** 2).mean() * weight).backward()
    print(mode, "gradient", w.grad.item(), "SGD result", round((w - 0.1 * w.grad).item(), 2))

full batch gradient -2.0 SGD result 1.2
weighted accumulation gradient -2.0 SGD result 1.2
mean of means gradient -1.0 SGD result 1.1

Weighted accumulation matches the full-batch update; the mean of means changes it. For a per-token mean, count valid loss-bearing tokens rather than sequences. Normalize a final partial accumulation window by its actual item count. Accumulation matches a large batch's gradient only when the forward computation for each item is also equivalent; batch-dependent operations can change that computation.

For equal sizes across distributed training processes, effective batch size is microbatch_size * accumulation_steps * world_size, where world_size counts those processes. Each process contributes its local gradients to the distributed update. This count gives the data per optimizer update. Choose learning-rate scaling based on the optimizer and training regime. After all backward passes in a window, clip once, step once, advance a per-update schedule once, and clear gradients. With AMP, keep the scaler's scale fixed within that window and unscale only after accumulation finishes. PyTorch accumulation guidance

Gradient clipping bounds the input to the optimizer

A malformed or unusually difficult batch can create an unusually large gradient. Global-norm gradient clipping treats every parameter gradient as one long vector and measures its L2 norm.

When that norm exceeds a limit, clipping rescales all gradients together:

Here, $c > 0$ is the clipping threshold. Every coordinate gets the same scaling factor, so clipping preserves the gradient's direction before it enters the optimizer. Predict which of the two vectors below will change.

import numpy as np

def clip_global_norm(gradient: np.ndarray, max_norm: float) -> tuple[np.ndarray, float]:
    if not np.isfinite(max_norm) or max_norm <= 0:
        raise ValueError("max_norm must be finite and positive")
    norm = float(np.linalg.norm(gradient))
    if not np.isfinite(norm):
        raise ValueError("gradient norm must be finite")
    if norm == 0.0:
        return gradient.copy(), norm
    scale = min(1.0, max_norm / norm)
    return gradient * scale, norm

ordinary = np.array([3.0, 4.0])
outlier = np.array([300.0, 400.0])

for name, g in [("ordinary", ordinary), ("outlier", outlier)]:
    clipped, original_norm = clip_global_norm(g, max_norm=5.0)
    print(name, "before", original_norm, "after", round(float(np.linalg.norm(clipped)), 3), "value", clipped.round(2).tolist())

ordinary before 5.0 after 5.0 value [3.0, 4.0]
outlier before 500.0 after 5.0 value [3.0, 4.0]

Clipping bounds the optimizer's gradient input, not the full AdamW update by that same threshold. Adam state, learning rate, and weight decay also affect how far the parameters move. For example, at Adam's first step, both [300, 400] and its clipped version [3, 4] yield updates close to [lr, lr] when epsilon is negligible. Clipping still changes the history stored for later steps.

If clipping fires on nearly every batch, investigate data, loss scaling, model stability, or the threshold rather than hiding the signal.

Clip before optimizer.step(). step() uses its input gradient to update AdamW's m and v buffers as well as the parameters, so clipping first keeps one outlier from entering all three paths.

Build it: optimize the valley with AdamW

Use the equations to implement AdamW and run it beside SGD on the same loss surface. Log final loss and both coordinate errors.

The peak rates differ because adaptive normalization changes update scale, so this isn't an optimizer leaderboard. Check that each run stays finite and both coordinate errors end smaller than they started. That's a basic sanity check, not a convergence guarantee.

Even convexity alone doesn't guarantee Adam's convergence: Reddi, Kale, and Kumar gave a convex counterexample and identified a problem in the original convergence analysis. Their AMSGrad variant retains a running maximum of squared-gradient history, with guarantees under stated assumptions.^{[5]Reference 5On the Convergence of Adam and Beyondhttps://arxiv.org/abs/1904.09237} That result doesn't predict failure on this valley, or guarantee a global optimum for a neural network. Evaluate the finite training budget and held-out behavior separately.

import math
import numpy as np

def loss(w: np.ndarray) -> float:
    return float(0.5 * (100.0 * w[0] ** 2 + w[1] ** 2))

def gradient(w: np.ndarray) -> np.ndarray:
    return np.array([100.0 * w[0], w[1]])

def cosine_lr(step: int, total_steps: int, peak_lr: float) -> float:
    return peak_lr * 0.5 * (1.0 + math.cos(math.pi * step / (total_steps - 1)))

def run_sgd(steps: int = 80) -> np.ndarray:
    w = np.array([1.0, 1.0])
    for step in range(steps):
        w -= cosine_lr(step, steps, peak_lr=0.018) * gradient(w)
        assert np.isfinite(loss(w)), "SGD became non-finite"
    return w

def run_adamw(steps: int = 80) -> np.ndarray:
    w = np.array([1.0, 1.0])
    m = np.zeros_like(w)
    v = np.zeros_like(w)
    beta1, beta2, eps, decay = 0.9, 0.999, 1e-8, 0.01
    for step in range(1, steps + 1):
        g = gradient(w)
        m = beta1 * m + (1.0 - beta1) * g
        v = beta2 * v + (1.0 - beta2) * g**2
        m_hat = m / (1.0 - beta1**step)
        v_hat = v / (1.0 - beta2**step)
        lr = cosine_lr(step - 1, steps, peak_lr=0.08)
        old_w = w.copy()
        w -= lr * m_hat / (np.sqrt(v_hat) + eps)
        w -= lr * decay * old_w
        assert np.isfinite(loss(w)), "AdamW became non-finite"
    return w

for name, result in [("SGD", run_sgd()), ("AdamW", run_adamw())]:
    assert np.all(np.abs(result) < 1.0), f"{name}: a coordinate didn't shrink"
    print(name, "loss", round(loss(result), 6), "w", result.round(5).tolist())

SGD loss 0.117302 w [-0.0, 0.48436]
AdamW loss 0.023509 w [0.02158, 0.02158]

Why do AdamW's two coordinates end almost equal despite the 100:1 curvature? Starting from equal values, their gradient histories are nearly constant multiples of one another. Adam's numerator and denominator scale together, almost canceling that factor. SGD has no such normalization, so its b error remains much larger. This symmetry belongs to the constructed quadratic; real model selection still needs validation metrics.

Use PyTorch after understanding the state

In a neural-network training loop, PyTorch computes gradients by backpropagation and provides tested optimizer implementations.

Clear old gradients with zero_grad(), run backward(), clip before optimizer.step(), set this step's learning rate, then call step().

Capture the pre-clipping norm so you can see how often the guardrail fires. PyTorch's clip_grad_norm_() returns that total norm. Pass error_if_nonfinite=True so a NaN or Inf raises before it enters optimizer state.

This example uses the same two-coordinate loss, but adds clipping and runs for 60 updates rather than 80. Its final numbers therefore aren't a direct check against the unclipped NumPy run.

import math
import torch

torch.manual_seed(0)
w = torch.nn.Parameter(torch.tensor([1.0, 1.0]))
optimizer = torch.optim.AdamW([w], lr=0.08, weight_decay=0.01)
clipped_updates = 0
max_gradient_norm = 0.0

for step in range(60):
    optimizer.zero_grad()
    loss = 0.5 * (100.0 * w[0] ** 2 + w[1] ** 2)
    loss.backward()
    gradient_norm = torch.nn.utils.clip_grad_norm_([w], max_norm=20.0, error_if_nonfinite=True).item()
    max_gradient_norm = max(max_gradient_norm, gradient_norm)
    clipped_updates += int(gradient_norm > 20.0)
    lr = 0.08 * 0.5 * (1.0 + math.cos(math.pi * step / 59))
    optimizer.param_groups[0]["lr"] = lr
    optimizer.step()

final_loss = 0.5 * (100.0 * w[0] ** 2 + w[1] ** 2)
print("loss", round(final_loss.item(), 6))
print("w", [round(value, 5) for value in w.detach().tolist()])
print("max pre-clip gradient norm", round(max_gradient_norm, 3))
print("clipped updates", clipped_updates)

loss 0.047348
w [0.03062, 0.03062]
max pre-clip gradient norm 100.005
clipped updates 11

This example assigns the current update's learning rate directly before optimizer.step(). If you replace that assignment with a built-in PyTorch scheduler, call scheduler.step() after optimizer.step().

Calling it first skips the schedule's first value. Choose the scheduler's cadence deliberately: some advance per optimizer update, others per epoch, and ReduceLROnPlateau consumes a measured validation metric. PyTorch scheduler guidance

If you later use mixed precision with GradScaler, unscale the gradients before measuring or clipping them with scaler.unscale_(optimizer). Otherwise the clipping threshold applies to artificially scaled values. PyTorch AMP clipping example

If AMP skips an optimizer update because the gradients are nonfinite, a schedule defined per optimizer update should not advance for that skipped update. Count updates that actually ran, rather than every attempted batch.

Resume the history, not just the weights

AdamW's next update depends on its moments and step count as well as the current parameters. A scheduler also needs its position and learning-rate state. Test that dependency on the same valley: save after three updates, then compare six uninterrupted updates with a resumed run and a weights-only restart.

import copy
import torch

def make_run(value):
    w = torch.nn.Parameter(value.clone())
    optimizer = torch.optim.AdamW([w], lr=0.01, weight_decay=0.01)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.5)
    return w, optimizer, scheduler

def advance(run, steps):
    w, optimizer, scheduler = run
    for _ in range(steps):
        optimizer.zero_grad(set_to_none=True)
        (0.5 * (100 * w[0] ** 2 + w[1] ** 2)).backward()
        optimizer.step()
        scheduler.step()

original = make_run(torch.ones(2))
advance(original, 3)
checkpoint = copy.deepcopy({
    "weights": original[0].detach(),
    "optimizer": original[1].state_dict(),
    "scheduler": original[2].state_dict(),
})
advance(original, 3)

resumed = make_run(checkpoint["weights"])
resumed[2].load_state_dict(checkpoint["scheduler"])
resumed[1].load_state_dict(checkpoint["optimizer"])
advance(resumed, 3)
weights_only = make_run(checkpoint["weights"])
advance(weights_only, 3)
print("complete state matches:", torch.equal(original[0], resumed[0]))
print("weights-only restart matches:", torch.equal(original[0], weights_only[0]))
print("completed Adam steps:", int(resumed[1].state[resumed[0]]["step"].item()))

complete state matches: True
weights-only restart matches: False
completed Adam steps: 6

The deep copy freezes the in-memory checkpoint before the original run continues. Construct the scheduler before loading the optimizer state; the scheduler constructor can change learning rates. Here we restore the scheduler first, then the optimizer, following the PyTorch loading example. Match the original parameter groups and ordering when restoring their state.

This deterministic CPU fixture needs no random or data-loader state. A real resume also saves the update counter, data position, relevant random-generator states, and AMP scaler state when used. Later training chapters add those details, validation data, mixed precision, and distributed state. The optimizer logic you traced here remains inside that larger loop.

AdamW isn't the endpoint of optimizer research. For example, PyTorch 2.14 provides Muon, which treats two-dimensional hidden-layer weight matrices as matrices rather than independent coordinates. Its documentation recommends a method such as AdamW for other parameters, including biases and embeddings. A model can therefore use different optimizers for different parameter groups; compare the whole recipe under the same training budget and validation metric.

Optimizer state is also a memory budget

For each trained parameter, plain SGD has no moving-average buffer, momentum stores one buffer, and AdamW stores two (m and v). If both AdamW buffers use 32-bit floats, their cost is:

A model with 7 billion parameters therefore needs about 56 GB for those two buffers alone. Estimate that number before running the code:

parameters = 7_000_000_000
bytes_per_float32 = 4
moment_buffers = 2
bytes_used = parameters * bytes_per_float32 * moment_buffers

print("AdamW moment buffers in GB:", bytes_used / 1_000_000_000)

AdamW moment buffers in GB: 56.0

The 56 GB count omits parameters, gradients, activations, and any master-weight copies used by a training setup. At large scale, sharding optimizer state across devices is one reason systems such as ZeRO exist.^{[6]Reference 6ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.https://arxiv.org/abs/1910.02054}

Read optimizer failures from traces

A training curve gives you a symptom, not a cause. Collect a small evidence set before changing an optimizer:

Start with one failure hypothesis and log the quantity that could expose it. Change one mechanism, then re-run on the same data slice.

Worked checks

For these checks, name the buffer an update is using, and why one learning rate can't serve both the steep wall and the shallow floor. Use the same valley for each check. Predict first, then compare your work with the solution.

1. Starting from [a, b] = [1, 1], compute one SGD update with lr = 0.01. Does either coordinate cross zero?
2. At Adam step one, take gradient [20, 0.2], lr = 0.001, and ignore epsilon. What is the bias-corrected update?
3. A global gradient has norm 50 and threshold 5. What scaling factor does clipping apply?
4. Your run clips 95 percent of minibatches. Is "clipping works" an adequate conclusion?

Solution checks:

1. The gradient is [100, 1]; the new vector is [0, 0.99]. Neither coordinate overshoots, but the second is still slow.
2. The update is [0.001, 0.001]: first-step normalization removes magnitude differences for nonzero coordinates.
3. The factor is 5 / 50 = 0.1; the direction stays the same.
4. No. Persistent clipping is evidence that your threshold, learning rate, data, or numerical path needs inspection.