← All tutorials

Optimizers in Practice: AdamW, SAM, Schedules and Clipping

What a real training loop adds to the update rule.

On this page

The optimizer turns each gradient into a step. Around it, AdamW changes how weight decay enters Adam's step, SAM spends a second pass to prefer flat minima, a schedule changes α per step or per epoch (one full pass over the training data), and clipping caps the gradient's length.

1 Decoupled weight decay (AdamW)

Weight decay shrinks the weights toward zero each step. Here it meets Adam's step, whose worked first step is in Adaptive Optimizers, section 4.4.

On the illustrative loss of the series, L = 5x² + 0.2y² from (−3.5, 5), at α 0.5 and weight decay λ 0.1, y after each step:

The difference is where λθ goes. This page uses PyTorch's form: the decay step is lr·λ·θ.

Loshchilov and Hutter 2019 also scale the decay by the schedule, as lr·λ·θ does here. Under L2, weights with large gradients are barely decayed; AdamW decays every weight at the same rate. Take weight A, used on every token, and weight B, a rare token's embedding row.

Illustrative values, steady state with α 1e−3, λ 0.01 and θ 1, and √v̂ 10 for A and 0.02 for B. LLaMA, Meta's large language model, trains with AdamW weight decay 0.1 (Touvron et al. 2023). Decay also meets gradient balancing, in Gradient Balancing Inside the Training Loop, section 9.2.

import torch
from torch.optim import Adam
from torch.optim import AdamW
model = torch.nn.Linear(10, 1)
p = list(model.parameters())
lr, lam = 1e-3, 1e-2  # α, λ
decay = dict(weight_decay=lam)
# Adam: λθ joins the grad
adam = Adam(p, lr, **decay)
# AdamW: decay outside it
adamw = AdamW(p, lr, **decay)

1. 1A small model, α and λ.
2. 2Adam with λ inside the gradient, the L2 line.
3. 3AdamW with the decay outside the rescaling: the step lr·weight_decay·θ.

2 Sharpness-aware minimization (SAM)

In this section ε is SAM's perturbation, a small move of the weights that SAM tests; it is unrelated to Adam's ε in section 1, the tiny constant that keeps the step from dividing by zero.

SAM wraps another optimizer and prefers flat minima, where the loss stays low when the weights move a little. Test data shifts the loss surface slightly, so a flat minimum keeps its low loss there; Foret et al. 2021 report SAM improving image classifiers on CIFAR and ImageNet. SAM tests every point within a radius ρ of the weights.

SAM estimates the worst nearby point with one step uphill, written ε; the hat marks this one-step estimate of ε. At the ravine start with ρ 0.05:

In code, perturb is the first pass; closure returns the minibatch loss. SAM's official code is in JAX, Google's array library (google-research/sam).

import torch
from torch import stack
def perturb(ps, closure, rho):
    closure().backward()
    gs = [p.grad for p in ps]
    ns = [g.norm() for g in gs]
    k = rho / stack(ns).norm()
    ds = [k * g for g in gs]
    with torch.no_grad():
        for p, d in zip(ps, ds):
            p.add_(d)  # θ + ε̂
    return ds

1. 1Reads the gradient at θ.
2. 2Scales it to length ρ, which is ε.
3. 3Moves every weight to θ + ε and returns ε.

ds = perturb(ps, closure, rho)
opt.zero_grad()
closure().backward()  # θ + ε̂
with torch.no_grad():
    for p, d in zip(ps, ds):
        p.sub_(d)  # back to θ
opt.step()  # steps with g_SAM

1. 1Moves to θ + ε with perturb.
2. 2Reads the gradient there, g_SAM, the second SAM line.
3. 3Moves the weights back to θ.
4. 4Lets the base optimizer step with g_SAM.

Each SAM step costs two forward and backward passes.

3 Learning rate schedules

A large rate moves fast early, but near a minimum a constant rate keeps bouncing, so a schedule changes α as training goes.

LLaMA used 2,000 warmup steps and a cosine to 10% of the peak. Transformers that normalize each sublayer's input (pre-LN) need less warmup than ones that normalize its output (post-LN) (The Transformer Block, section 3).

A scheduler keeps one counter and sets α from it. Step, exponential and cosine usually step once per epoch; one-cycle steps once per batch.

import torch
import torch.optim as optim
from torch.optim import lr_scheduler
net = torch.nn.Linear(4, 1)
ps = net.parameters()
opt = optim.SGD(ps, lr=0.1)
# pick one, stepped per epoch
sched = lr_scheduler.StepLR(opt, 30, 0.1)
sched = lr_scheduler.ExponentialLR(opt, 0.95)
sched = lr_scheduler.CosineAnnealingLR(
    opt, T_max=100)

1. 1Builds a model and an SGD optimizer to schedule.
2. 2Step decay, the rate times 0.1 every 30 epochs.
3. 3Exponential decay, the rate times 0.95 every epoch.
4. 4Cosine annealing over 100 epochs, the cosine schedule equation.

sched = lr_scheduler.CosineAnnealingWarmRestarts(
    opt, T_0=10)
sched = lr_scheduler.OneCycleLR(
    opt, max_lr=0.1, total_steps=1000)
for step in range(1000):
    opt.step()
    sched.step()  # per batch

1. 1Warm restarts every 10 epochs.
2. 2One-cycle up to 0.1 and back over 1000 steps.
3. 3Steps the optimizer, then the one-cycle schedule, once per batch.

4 Gradient clipping

Clipping guards against exploding gradients, gradients that suddenly become very large. It caps either the length ‖g‖ of the gradient vector (Vectors and Matrices, section 2) or each coordinate. For g = (−35, 2):

Clipping sits between backward() and opt.step().

import torch
import torch.nn.functional as F
from torch.nn.utils import clip_grad_norm_
from torchvision import datasets, models, transforms
data = datasets.CIFAR10('.', download=True,
    transform=transforms.ToTensor())
loader = torch.utils.data.DataLoader(data, batch_size=128)
net = models.resnet18(num_classes=10)
opt = torch.optim.SGD(net.parameters(), lr=0.1)
x, y = next(iter(loader))
loss = F.cross_entropy(net(x), y)
loss.backward()
clip_grad_norm_(net.parameters(), 1.0)
opt.step()  # after the clip

1. 1A ResNet-18 for the 10 CIFAR-10 image classes, its optimizer, and batches of 128 images.
2. 2One batch, its cross-entropy loss (the standard classification loss) and its gradients.
3. 3Gradient length capped at τ 1.0, the norm clipping line.
4. 4Steps with the clipped gradient.

The through-time case is in RNN Basics: Recurrence and Backpropagation Through Time, section 2.2; its interaction with balancing is in Gradient Balancing Inside the Training Loop, section 10.

5 AdamW, a schedule and clipping in one PyTorch loop

The optimizer is built once, before the loop, and keeps its state for the parameters from then on; each batch runs forward, backward, the optional clip and opt.step(), and the scheduler changes α per batch or per epoch.

The setup trains a ResNet-18 on the 50,000 CIFAR-10 training images with AdamW and a cosine schedule over every batch.

import torch
import torch.nn.functional as F
from torch.nn.utils import clip_grad_norm_
from torch.optim import AdamW, lr_scheduler
from torchvision import datasets, models, transforms
data = datasets.CIFAR10('.', download=True,
    transform=transforms.ToTensor())
loader = torch.utils.data.DataLoader(
    data, batch_size=128, shuffle=True)
net = models.resnet18(num_classes=10)
opt = AdamW(net.parameters())  # once
epochs = 30
sched = lr_scheduler.CosineAnnealingLR(
    opt, T_max=epochs * len(loader))

1. 1CIFAR-10 in shuffled batches of 128 images.
2. 2The ResNet-18 and AdamW, built once, before the loop, with PyTorch's defaults lr 1e-3, betas (0.9, 0.999) and eps 1e-8 (AdamW documentation).
3. 3A cosine schedule over every batch of 30 epochs.

One update per batch, 391 batches per epoch.

for epoch in range(epochs):
    for x, y in loader:
        loss = F.cross_entropy(net(x), y)
        opt.zero_grad()
        loss.backward()
        clip_grad_norm_(net.parameters(), 1.0)
        opt.step()  # every batch
        sched.step()  # per batch

1. 1The cross-entropy loss on one batch of images.
2. 2Clears old gradients and fills new ones.
3. 3Caps the gradient's total length at 1.0 before the step (section 4).
4. 4One optimizer update per batch.
5. 5One schedule tick per batch, since T_max counts batches; with T_max in epochs it would step once per epoch, after the inner loop.

SAM replaces the single opt.step() with its two-pass step from section 2.

6 Failures in the training loop

7 References

1. Kingma, D. P., and Ba, J. (2015). Adam: a method for stochastic optimization. ICLR.
2. Goyal, P., et al. (2017). Accurate, large minibatch SGD: training ImageNet in 1 hour. arXiv:1706.02677.
3. Loshchilov, I., and Hutter, F. (2017). SGDR: stochastic gradient descent with warm restarts. ICLR.
4. Smith, L. N. (2018). A disciplined approach to neural network hyper-parameters: part 1, learning rate, batch size, momentum, and weight decay. arXiv:1803.09820.
5. Loshchilov, I., and Hutter, F. (2019). Decoupled weight decay regularization. ICLR.
6. Foret, P., et al. (2021). Sharpness-aware minimization for efficiently improving generalization. ICLR.
7. Touvron, H., et al. (2023). LLaMA: open and efficient foundation language models. arXiv:2302.13971.
8. PyTorch documentation, torch.optim.lr_scheduler (when schedulers step).
9. PyTorch documentation, torch.optim.AdamW (default settings), and torch.nn.utils.clip_grad_norm_.