# Optimizers and learning-rate schedules — PyTorch

Source: https://www.skillbyai.com/en/pytorch/pt-optimizers

> Create SGD, Adam and AdamW optimizers, explain how they differ, and attach a learning-rate scheduler correctly.

## 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.

![Diagram of 5 connected stages from left to right: forward pass, loss, backward, optimizer step, zero gradients.](assets/figures/pytorch/section-5-map.svg) — Figure 5.1: forward pass -> loss -> backward -> optimizer step -> zero gradients

## 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.

```python
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.

**Quiz:** 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
- [x] AdamW applies weight decay directly to the weights instead of adding it to the gradient

*Answer:* AdamW applies weight decay directly to the weights instead of adding it to the gradient. AdamW decouples weight decay from the adaptive gradient update. Momentum-like running averages and scheduling work the same as in Adam.