Optimizers and learning-rate schedules

What an optimizer holds and does

An optimizer from torch.optim is given the parameters to update, for example torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9). optimizer.step() reads each parameter's .grad and updates the parameter; optimizer.zero_grad() clears the gradients.

SGD with momentum keeps a running velocity so updates keep moving in consistent directions. It is simple and, when well tuned, generalises well; it is common for image models. Adam keeps running averages of each parameter's gradient and squared gradient and gives every parameter its own effective step size, so it works reasonably with little tuning; lr=1e-3 is a common starting point. AdamW is Adam with decoupled weight decay: the decay is applied directly to the weights rather than mixed into the gradient, which regularises more predictably. It is the usual choice for transformers.

The learning rate matters more than the choice of optimizer. Schedulers in torch.optim.lr_scheduler (such as StepLR, CosineAnnealingLR, OneCycleLR) change it during training. Call scheduler.step() after optimizer.step(), at the interval the scheduler expects (per epoch for most, per batch for OneCycleLR).

One training step

Each batch goes through the same five stages: a forward pass produces predictions, the loss scores them, backward computes gradients, the optimizer step updates weights, and gradients are cleared before the next batch.

AdamW with a step schedule

The learning rate halves every 10 epochs, so after 30 epochs it is 1e-3 x 0.5^3 = 1.25e-4. A real loop would iterate over batches; one synthetic batch keeps the example short.

import torch
from torch import nn

model = nn.Linear(10, 1)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)
x, y = torch.randn(64, 10), torch.randn(64, 1)

for epoch in range(30):
    optimizer.zero_grad()
    loss = nn.functional.mse_loss(model(x), y)
    loss.backward()
    optimizer.step()
    scheduler.step()        # after optimizer.step(), once per epoch here

print(scheduler.get_last_lr())   # about [0.000125]

Two hikers in fog

SGD is a hiker who always takes the same stride length down the slope. Adam is a hiker who shortens the stride on rocky, steep ground and lengthens it on gentle, consistent ground, separately for each direction. The adaptive hiker needs less advice to get started; the steady one, well guided, often finds a better valley.

Quick check: How does AdamW differ from Adam?

• AdamW does not use momentum
• AdamW uses a fixed learning rate that cannot be scheduled
• AdamW only works on GPUs
• AdamW applies weight decay directly to the weights instead of adding it to the gradient