Regularization, Schedules & the Modern Training Recipe

"For the first five hundred steps they let me walk slowly so I would not trip. Then they let me run. Then, gently, over thousands of steps, they asked me to slow to a stop right at the finish. I have never trained any other way and I cannot imagine why I once did."

An Optimizer Describing Its Cosine Schedule

In Section 21.3 you decided what to fine-tune. This section decides how to run the optimization itself. The motivating fact, established in Chapter 20, is that "ResNet strikes back" showed a 2015 ResNet-50 gaining several accuracy points purely from a modern recipe, no architecture change at all. That result is this section's thesis made measurable: the recipe is a first-class part of the model. We will take each ingredient in turn, explain the problem it solves, and then assemble them, because the ingredients interact and the bundle is what works. The illustration below sums up that point: no single layer of the cake is impressive alone, but stacked together they become the recipe.

1. The Learning-Rate Schedule: Warmup and Cosine Decay Beginner

A constant learning rate is almost never optimal. Early in training the weights are far from any good solution and large steps help, but a freshly-initialized network (especially with the noisy gradients of a new head from Section 21.3) can diverge if the rate starts high. Late in training, large steps cause the loss to bounce around the minimum instead of settling into it. The modern schedule solves both ends. Warmup ramps the learning rate linearly from near zero up to its peak over the first few hundred steps, so the network stabilizes before taking big steps. Cosine decay then smoothly lowers the rate following a half-cosine curve from the peak down to near zero over the rest of training, $\eta_t = \frac{1}{2}\eta_{\max}\left(1 + \cos\frac{\pi t}{T}\right)$, so the model settles gently into the minimum.

Warmup looks like a fussy detail until you watch a run try to skip it. Launch a vision transformer straight at its peak learning rate with no warmup and the loss does not gently fail, it explodes to NaN within the first dozen steps: the freshly-initialized attention layers produce enormous gradients, one full-size step blows the weights past any usable range, and the run is dead before it has seen a single full epoch. The same network with a few hundred steps of linear warmup trains to a strong result. The fix and the catastrophe are separated by one short ramp at the very start, which is exactly why warmup, once an obscure trick, is now non-negotiable for transformers and increasingly standard for CNNs too.

Figure 21.4.1 shows the combined warmup-then-cosine curve that is now the default in essentially every vision recipe.

2. Weight Decay Done Right: AdamW Intermediate

Weight decay shrinks weights toward zero each step, a regularizer that discourages large weights and improves generalization. For plain stochastic gradient descent (SGD), the gradient-following optimizer introduced in Section 18.2, adding an L2 penalty to the loss (an extra term proportional to the sum of squared weights, which the gradient then shrinks) and applying weight decay are mathematically the same operation. For adaptive optimizers like Adam, which scale each parameter's step by a running estimate of its own gradient magnitude (also from Section 18.2), they are not, and this surprised the field for years. If the L2 penalty rides inside the gradient, Adam's per-parameter scaling shrinks it too, so parameters with large gradients are decayed less than intended. The fix, AdamW, decouples the decay: it applies the adaptive gradient step and then separately shrinks the weights by a fixed factor, $\theta \leftarrow \theta - \eta\,\lambda\,\theta$, independent of the gradient scaling. AdamW is now the default optimizer for training the vision transformers of Chapter 22 and modern CNNs alike.

3. Label Smoothing, Stochastic Depth, and EMA Intermediate

Three more regularizers round out the bundle. Label smoothing replaces the hard one-hot target with a softened one: instead of asking the model to predict probability $1.0$ for the true class and $0.0$ elsewhere, it asks for $1 - \epsilon$ on the true class and $\epsilon / (K-1)$ spread over the others, with $\epsilon$ around $0.1$. This stops the model from driving its logits (the raw pre-softmax scores from the classifier head, introduced with the cross-entropy loss in Section 18.5) to infinity chasing perfect confidence, which improves both generalization and the calibration we discussed for MixUp in Section 21.2. Stochastic depth randomly drops entire residual blocks during training (replacing them with the identity skip connection from Chapter 20), which regularizes very deep networks and, as a bonus, speeds up training since dropped blocks are not computed.

Exponential moving average (EMA) of the weights maintains a slowly-updated shadow copy, $\theta_{\text{EMA}} \leftarrow \alpha\,\theta_{\text{EMA}} + (1-\alpha)\,\theta$ with $\alpha$ near $0.9999$, and uses that averaged copy for evaluation. Because SGD bounces around the minimum, the average sits closer to the center of the basin than any single iterate, typically buying a fraction of a point of accuracy for almost no cost. Label smoothing in PyTorch is a one-argument change.

import torch
import torch.nn as nn

# Label smoothing is a single argument to the standard loss.
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

# A minimal model-EMA: keep a shadow copy, update it after each optimizer step.
class ModelEMA:
    def __init__(self, model, decay=0.9999):
        self.decay = decay
        self.shadow = {k: v.detach().clone()
                       for k, v in model.state_dict().items()}

    @torch.no_grad()
    def update(self, model):
        for k, v in model.state_dict().items():
            if v.dtype.is_floating_point:
                self.shadow[k].mul_(self.decay).add_(v, alpha=1 - self.decay)
# Usage: opt.step(); ema.update(model)   then evaluate with ema.shadow weights.

4. Mixed Precision and the Assembled Recipe Advanced

The last ingredient is about speed, not accuracy. Mixed-precision training runs most operations in 16-bit floating point (half the memory and roughly twice the throughput on modern GPUs with tensor cores, the specialized hardware units that multiply low-precision matrices very fast) while keeping a 32-bit master copy of the weights and a loss scaler to prevent tiny gradients from underflowing to zero. It typically halves memory use and noticeably speeds up training with no accuracy loss, which is why it is standard, and it is the same mechanism first met in Section 18.6. PyTorch 2.x provides it through the device-generic torch.amp.autocast and a gradient scaler. With every ingredient introduced, the assembled recipe below is close to what timm and other production frameworks run by default.

import torch
from torch.amp import autocast, GradScaler         # PyTorch 2.x device-generic AMP
from torch.optim.lr_scheduler import OneCycleLR   # warmup + cosine-like decay

model = model.cuda()
# Split params: decay weights, no-decay for norms/biases (subsection 2 insight).
decay, no_decay = [], []
for name, p in model.named_parameters():
    (no_decay if p.ndim <= 1 or name.endswith("bias") else decay).append(p)
opt = torch.optim.AdamW(
    [{"params": decay, "weight_decay": 0.05},
     {"params": no_decay, "weight_decay": 0.0}], lr=1e-3)

steps_per_epoch, epochs = len(loader), 100
sched = OneCycleLR(opt, max_lr=1e-3, total_steps=steps_per_epoch * epochs,
                   pct_start=0.05)               # 5% of steps as warmup
scaler = GradScaler("cuda")
criterion = torch.nn.CrossEntropyLoss(label_smoothing=0.1)

for epoch in range(epochs):
    for images, labels in loader:               # loader already does RandAugment
        images, labels = images.cuda(), labels.cuda()
        opt.zero_grad()
        with autocast("cuda"):                   # run forward in fp16
            loss = criterion(model(images), labels)
        scaler.scale(loss).backward()            # scaled backward to avoid underflow
        scaler.step(opt); scaler.update()
        sched.step()                             # per-step cosine schedule

# timm's train.py reproduces the full modern recipe from flags:
#   python train.py /imagenet \
#     --model convnext_tiny --opt adamw --weight-decay 0.05 \
#     --sched cosine --warmup-epochs 5 --epochs 300 \
#     --smoothing 0.1 --mixup 0.8 --cutmix 1.0 \
#     --aa rand-m9-mstd0.5 --drop-path 0.1 --model-ema --amp

# Programmatic equivalents of the key pieces:
from timm.scheduler import CosineLRScheduler
from timm.loss import SoftTargetCrossEntropy   # for MixUp/CutMix soft labels
from timm.utils import ModelEmaV2
from timm.data import Mixup