Skip to content

Optimizers and Regularization

What This Is

This page is about one practical training question:

  • which knob is the smallest fix that actually improves validation — optimizer, learning rate, weight decay, dropout, or none of the above

The trap this page is trying to prevent is adding regularization without knowing which problem you are solving. Overfitting, underfitting, and optimization instability all look different in the curves, and they call for different fixes. Reaching for dropout when the real problem is a too-high learning rate makes both worse.

The working skill: read the curve, name the problem, apply the smallest fix, and see whether validation actually moved.

When You Use It

  • training curves look unstable, oscillating, or diverging
  • the validation loss has stalled and training loss keeps dropping (overfitting)
  • both losses are flat and bad (underfitting)
  • gradient norms explode or vanish
  • a fine-tune is drifting from the pretrained behavior faster than it should
  • you need a defensible choice between AdamW and SGD for a new task

The Three Forces

Every training curve is driven by three forces, and most fixes belong to one of them:

Force Controls Visible when Typical fix
Optimization how the weights move per step instability, divergence, exploding gradients LR, clipping, warmup, optimizer choice
Capacity how much the model can memorize overfitting vs underfitting gap width, depth, weight decay, dropout
Data what the model sees everything upstream better sampling, augmentation, cleaning

Almost every debugging mistake comes from reaching into the wrong column.

Optimizer Choice

A short, opinionated ladder:

  1. AdamW — widely used starting point for transformers and fine-tuning. Adaptive per-parameter learning rates and decoupled weight decay.
  2. SGD with momentum — established baseline for CNN training; it can outperform adaptive methods on some tasks after schedule tuning.
  3. Lion — lower optimizer-state memory than AdamW, but requires its own tuning and is not a universal upgrade.
  4. RMSprop / Adagrad — legacy; rarely a good first pick today.
  5. Sophia, Muon, distributed Shampoo — cutting edge; expect some volatility and a larger tuning budget.

Example starting points from common recipe families—not cross-task defaults:

  • LLMs / transformers: AdamW, LR 1e-4 to 5e-4, weight decay 0.1, β₂ = 0.95
  • vision from scratch: SGD with momentum 0.9, LR 0.1 with cosine schedule, weight decay 5e-4
  • fine-tuning: AdamW, LR 1e-5 to 5e-5, weight decay 0.01, warmup 500–1000 steps

AdamW vs Adam — Why The W Matters

Plain Adam does not include weight decay in its original algorithm. Implementing L2 regularization by adding it to the gradient couples the penalty to Adam's adaptive scaling; AdamW decouples the shrinkage from that gradient update:

# AdamW-style update (simplified)
g_t = gradient
m_t = β1 * m_{t-1} + (1 - β1) * g_t
v_t = β2 * v_{t-1} + (1 - β2) * g_t^2
θ_t = θ_{t-1} - lr * (m_t / (sqrt(v_t) + ε)  +  wd * θ_{t-1})
                                                ^^^^^^^^^^
                                       decoupled weight decay

AdamW is a strong first comparison when weight decay is intended. In PyTorch, use torch.optim.AdamW(...). The distinction is developed in Decoupled Weight Decay Regularization.

Weight Decay vs L2 — Subtle But Important

People often say "L2 regularization" and "weight decay" interchangeably. They are only equivalent for plain SGD.

  • L2 penalty: adds λ ||θ||² to the loss. Gradients include 2λ θ.
  • Weight decay: multiplies θ by (1 - lr * λ) every step.

For adaptive optimizers, these two generally diverge. AdamW applies decoupled weight decay; whether it improves held-out performance and what coefficient to use are empirical questions.

Practical tuning:

  • weight decay 0.1 for LLM pretraining
  • 0.05 for LLM fine-tune
  • 0.01 for most vision fine-tunes
  • 5e-4 for vision from scratch with SGD + momentum

Learning Rate Choice

The single most important hyperparameter. A wrong LR makes every other fix look ineffective.

LR finder

Run a short exponential LR scan over a range appropriate to the optimizer and plot a smoothed loss against LR. Save model and optimizer state first and restore it afterward; the scan itself changes both.

# minimal LR finder sketch
lrs = np.logspace(-7, 0, num=200)
losses = []
for i, batch in enumerate(train_loader):
    if i >= len(lrs):
        break
    for g in optimizer.param_groups:
        g["lr"] = lrs[i]
    loss = step(batch)
    losses.append(loss.item())
plt.semilogx(lrs, losses)

The steep downward region provides candidates for a short validation sweep; it is not a substitute for held-out evaluation.

Scale with batch size

Some large-batch SGD recipes scale LR approximately linearly with batch size and add warmup. This is a recipe-specific heuristic, not a law for adaptive optimizers or every task; see the ImageNet experiments in Goyal et al..

Gradient Clipping

A cheap safeguard that prevents rare huge gradients from destroying the model:

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

When to try it:

  • transformers or RNNs when gradient norms spike or the reference recipe uses clipping; 1.0 is a common starting threshold
  • on any task where loss spikes show up in the training curve
  • during LR warmup, when the model is most vulnerable

When not to use it:

  • if your gradients are already well-behaved and clipping is not firing; it adds overhead for no benefit
  • with very aggressive clipping (max_norm < 0.5) where you are hiding a real LR problem

Learning Rate Warmup

Start at a small LR and linearly increase to the target over the first N steps:

from torch.optim.lr_scheduler import LambdaLR

def warmup_linear(step, warmup_steps):
    return min(1.0, step / warmup_steps)

scheduler = LambdaLR(optimizer, lambda s: warmup_linear(s, warmup_steps=1000))

Warmup can stabilize early optimization when the target learning rate is too aggressive at initialization. Choose its length relative to the total schedule and validate it; 500–2000 steps can be unreasonable for either a very short fine-tune or a very long pretraining run.

Gradient Accumulation

If your GPU cannot hold the batch size you need, accumulate gradients over several micro-batches:

accum = 4
num_batches = len(train_loader)
optimizer.zero_grad()
for i, batch in enumerate(train_loader):
    group_start = (i // accum) * accum
    group_size = min(accum, num_batches - group_start)
    loss = forward(batch) / group_size
    loss.backward()
    if (i + 1) % accum == 0 or (i + 1) == len(train_loader):
        optimizer.step()
        optimizer.zero_grad()

Most updates now average four micro-batches; the final group may be smaller. Exact equivalence to one large batch can break with batch-dependent layers or stochastic operations. Step-based schedulers should advance when optimizer.step() runs, and mixed precision requires the usual scaler sequence.

Exponential Moving Average (EMA)

Keep a second copy of the weights that is the exponential moving average of the training weights. Use it at inference for smoother, often better predictions:

from copy import deepcopy

class EMA:
    def __init__(self, model, decay=0.999):
        self.shadow = deepcopy(model).eval()
        for p in self.shadow.parameters():
            p.requires_grad = False
        self.decay = decay

    @torch.no_grad()
    def update(self, model):
        for s, p in zip(self.shadow.parameters(), model.parameters()):
            s.mul_(self.decay).add_(p, alpha=1 - self.decay)

EMA is common in diffusion and semi-supervised training. It adds a second weight copy and can help or hurt depending on decay and training length, so compare EMA and current weights on the same validation set.

Regularization Tools

Each changes the training problem differently. Add one at a time in an order justified by the observed failure:

Weight decay

A common first regularizer when parameter shrinkage is appropriate. Covered above.

Dropout

Drop a fraction of activations during training. Strong for feed-forward networks and transformer FFNs; weaker for CNNs.

self.drop = nn.Dropout(0.1)

Rules:

  • 0.1 is a common transformer starting point, but attention and FFN dropout should follow the architecture recipe and held-out evidence
  • dropout in LSTM/GRU only works between stacked layers, not within one layer
  • combining dropout and batch-normalized blocks changes activation statistics; validate the ordering and the held-out effect

Label smoothing

Replace a one-hot target with a mixture of the target distribution and a uniform distribution. In PyTorch this assigns 1 - ε + ε/K to the target and ε/K to each other class. It may reduce overconfidence, but calibration and accuracy must be measured:

loss_fn = nn.CrossEntropyLoss(label_smoothing=0.1)

Data augmentation

Often worth more than any other regularizer. Covered in Data Augmentation and Vision Augmentation and Shift Robustness.

Mixup and CutMix

Blend two training examples and their labels. Strong regularizer for vision; does not always help NLP.

Early stopping

The simplest regularizer. Watch the validation metric and stop when it plateaus or reverses. Pair with checkpointing so you keep the best epoch, not the last.

Stochastic weight averaging

Average weights from a late region of training. It can improve generalization when checkpoints occupy a compatible region; implemented via torch.optim.swa_utils.

Reading Curves Honestly

  • Training down, validation flat or up → capacity problem (overfit). Start with weight decay and augmentation; add dropout if still overfitting.
  • Both flat and bad → underfit or optimization problem. Try a higher LR, larger model, or longer training.
  • Training oscillates wildly → LR too high, missing gradient clipping, or missing warmup.
  • Validation worse than training even at epoch 1 → inspect expected generalization gap, augmentation differences, leakage, and train/validation shift.
  • Loss spikes then recovers → inspect the batch, gradient norm, scaler state, and LR; clipping is one possible intervention.
  • Loss drops then explodes far into training → inspect data batches, numerical range, mixed-precision scaler, optimizer state, and LR schedule.

Do not reach for dropout before you have named the problem.

Interaction With Mixed Precision

torch.amp changes the numerical regime and can interact with everything above:

  • weight decay still works but the scaling matters; use AdamW
  • gradient clipping must happen after scaler.unscale_
  • reduced-precision gradients or intermediates can underflow before the float32 optimizer update; use the AMP scaler workflow for float16 and consider bfloat16 where supported

See Mixed Precision Training for the full recipe.

What To Inspect

  • the LR trajectory (print optimizer.param_groups[0]['lr'] per step)
  • gradient norm and clipping frequency over time; sustained clipping means the threshold, LR, and data should be reviewed together
  • the training vs validation gap, not the absolute numbers
  • effective batch size after accumulation and DDP
  • whether weight decay is being applied to parameters it should not be (biases, LayerNorm)
  • the shape of the validation curve around the best epoch — is it a sharp minimum or a plateau?

A common refinement: exclude biases and LayerNorm parameters from weight decay.

def group_params_for_wd(model, wd):
    decay, no_decay = [], []
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if p.dim() <= 1 or name.endswith(".bias") or "norm" in name.lower():
            no_decay.append(p)
        else:
            decay.append(p)
    return [
        {"params": decay, "weight_decay": wd},
        {"params": no_decay, "weight_decay": 0.0},
    ]

This exclusion is common in transformer recipes, but reproduce the target model's recipe or validate the grouping.

Failure Pattern

Piling regularizers on top of an optimization problem. A training run that is unstable at LR 1e-3 does not need more dropout — it needs a lower LR, gradient clipping, or warmup. Regularization applied to an unstable optimizer produces a model that is both undertrained and underperforming.

Another failure: ignoring interactions. LR and weight decay, batch size, and warmup length can interact. Start with a small controlled sweep, then test the interactions suggested by the curves.

Common Mistakes

  • using plain Adam when AdamW was meant
  • applying weight decay to LayerNorm and bias parameters
  • turning up dropout to "fix" underfitting
  • using a from-scratch LR for a fine-tune
  • skipping warmup on transformer training
  • changing batch size without reconsidering the LR and schedule
  • stacking regularizers (dropout + weight decay + label smoothing + augmentation) without ablating
  • reading training loss as the primary signal
  • leaving gradient clipping off "because it wasn't needed last time"
  • using the same LR across all parameter groups on a fine-tune

Practice

  1. Run an LR finder on your current task and pick an LR from the curve.
  2. Train the same model with and without weight decay excluded from biases and norms. Report the delta.
  3. Introduce one training instability (LR too high) and show how warmup, clipping, and a lower LR each fix it.
  4. Compare AdamW and SGD + momentum on the same task. Report validation metric and training time.
  5. Add gradient accumulation to double your effective batch size. Test whether the LR should change rather than assuming a scaling rule.
  6. Train with and without label smoothing. Report validation accuracy and calibration (ECE).
  7. Add EMA and compare EMA vs last-step weights on validation.
  8. Show one curve pattern that calls for regularization and one that calls for a lower LR. Defend the diagnosis.

Runnable Example

Run the controlled optimizer comparison from the repository root:

.venv/bin/python examples/deep-learning-recipes/optimizer_regularization_demo.py

Inspect train and validation curves for each recipe and keep the split, seed, and model fixed. Attribute a gain only to the one optimizer or regularization choice that changed.

Longer Connection

Continue with Learning Rate Schedulers for the schedule side of LR choice, PyTorch Optimization Recipes for end-to-end examples, and Optimization, Regularization, and PEFT for the full track that combines these choices into a defended workflow.