AdamW decay groups and update-indexed schedules

Last updated: 6 Oct 20265 min read

tutorial

An optimizer schedule needs an exact update counter, a declared parameter-group policy and a resume rule; epoch labels alone are too vague.

Count parameter updates rather than batches

Gradient accumulation can process several microbatches before one optimizer step. A learning-rate schedule intended to warm over five parameter updates should advance only when those five updates occur, not after each microbatch. Record completed update count in the checkpoint and define whether the logged rate belongs to the update just made or the next one. With 47 planned updates and a five-update warmup, a resume after update 23 must use the rate for update 24, not restart at the first warmup value. Effective-batch accounting supplies the same update boundary.

Separate decay from the data gradient

AdamW applies decoupled weight shrinkage rather than adding a squared-weight penalty into Adam’s adaptive gradient statistics. That does not mean every parameter should decay. A common policy excludes biases and normalization offsets, while matrix weights receive a declared coefficient. Enumerate groups before constructing the optimizer and assert every trainable parameter belongs to exactly one group. Otherwise a parameter may silently receive two updates or no update. The code demonstrates the policy with a single linear layer; a production model needs a complete audited parameter inventory.

Specify the schedule equation

The example uses a linear warmup followed by a cosine descent to a nonzero floor. It computes the learning rate explicitly from zero-based update index, so it does not depend on an implicit scheduler step order. The first warmup update receives one-fifth of peak rate; the fifth receives peak rate; the final planned update receives the floor. If training ends early, the rate may never reach the floor. If training extends beyond the budget, define whether to clamp at the floor or recompute the plan. Store peak, floor, warmup count and total planned updates.

Resume the complete optimizer state

Loading model weights without Adam moments changes the next parameter update even if learning rate is restored. A coherent checkpoint includes optimizer state, completed update count, scaler state for mixed precision and any sampler position needed to reconstruct the next batch. Restore the optimizer before setting the rate for the next update, then compare one fixed-batch step with an uninterrupted reference. The resume lesson gives the broader state inventory. A plotted learning-rate line is no proof of resume parity.

Compare schedules fairly

Hold data order, effective batch, total accepted examples, architecture and evaluation windows fixed when comparing a constant rate with warmup and decay. Report class-specific validation behavior and multiple seeds rather than naming the schedule that happened to achieve one best run. An overly small floor may stop useful late learning; an overly high peak may destabilize early updates. Checkpoint selection must follow a separately declared rule, or schedule tuning and model selection become entangled.

Implementation

import math
import torch
from torch import nn

receipt_classifier = nn.Linear(9, 3)
optimizer = torch.optim.AdamW([
    {"params": [receipt_classifier.weight], "weight_decay": 0.04},
    {"params": [receipt_classifier.bias], "weight_decay": 0.0},
], lr=0.001)

planned_updates = 47
warmup_updates = 5
peak_rate = 0.001
floor_rate = 0.00003

def rate_for_update(update_index: int) -> float:
    if not 0 <= update_index < planned_updates:
        raise ValueError("update outside planned budget")
    if update_index < warmup_updates:
        return peak_rate * (update_index + 1) / warmup_updates
    progress = (update_index - warmup_updates + 1) / (planned_updates - warmup_updates)
    return floor_rate + 0.5 * (peak_rate - floor_rate) * (1 + math.cos(math.pi * progress))

assert math.isclose(rate_for_update(0), 0.0002)
assert math.isclose(rate_for_update(4), peak_rate)
assert math.isclose(rate_for_update(46), floor_rate)
for group in optimizer.param_groups:
    group["lr"] = rate_for_update(23)

Performance and operating cost

The schedule calculation is O(1) per optimizer update. AdamW stores first and second moments for each trainable parameter, so optimizer state can exceed model-weight storage; excluding bias from decay does not remove its moment storage. The warmup and cosine arithmetic is negligible beside a forward and backward pass. The material cost is experimental: comparing schedules fairly requires multiple training runs, and a wrong update index can waste those runs while still producing a smooth-looking learning-rate chart.

Common Mistakes

Do not advance an update-indexed schedule for every accumulation microbatch.
Do not load weights alone and call the next AdamW update an exact resume.
Do not apply weight decay twice by combining decoupled decay with an unexamined L2 penalty.