Table of Contents
What an Optimizer Does in PyTorch
Training a neural network means choosing parameters, usually called weights and biases, that make the loss small. After you compute gradients with backpropagation, an optimizer updates each parameter using those gradients. Conceptually, each update takes the form
$$\theta \leftarrow \theta - \alpha \cdot \text{update}(\nabla_\theta L),$$
where $\theta$ are the parameters, $L$ is the loss, and $\alpha$ is the learning rate. The key difference between optimizers is how they turn raw gradients into an update step, for example by using momentum, adaptive scaling, or both.
In PyTorch, an optimizer is an object from torch.optim that is constructed with the model parameters, then called each iteration after gradients exist.
The Learning Rate Is the Most Important Knob
The learning rate controls step size. If it is too large, training can diverge, oscillate, or produce NaNs. If it is too small, training can be painfully slow and may appear stuck.
A useful mental model is that the loss surface has steep and flat directions. A single global learning rate has to work across all of them, which is why many modern optimizers include adaptation or momentum.
Rule of thumb: if the loss explodes or becomes NaN early, try lowering the learning rate first. If the loss decreases extremely slowly and gradients are not tiny, try increasing it.
Core PyTorch Pattern: zero_grad, backward, step
An optimizer does not compute gradients. It only uses whatever gradients are currently stored in param.grad. A typical update sequence is:
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
for x, y in loader:
optimizer.zero_grad() # clear old gradients
pred = model(x)
loss = loss_fn(pred, y)
loss.backward() # compute new gradients
optimizer.step() # update parameters
The optimizer.step() call uses the optimizer’s internal state, such as momentum buffers, which is why you should keep the same optimizer instance throughout training unless you intentionally reset it.
SGD and Momentum
Stochastic Gradient Descent updates parameters by moving directly opposite the gradient. In PyTorch, plain SGD is torch.optim.SGD(..., lr=...).
Momentum adds a running “velocity” that smooths updates across steps, which often helps convergence and reduces zigzagging. In PyTorch, you enable it with momentum=0.9 (a common default).
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)If you are learning from scratch, SGD with momentum is a great baseline for many vision and classical deep learning setups, but it can require more learning rate tuning than adaptive methods.
Adam and AdamW
Adam is an adaptive optimizer that keeps running estimates of the first and second moments of gradients. Practically, it adjusts step sizes per parameter, which often makes training easier at the beginning and reduces sensitivity to the exact learning rate choice.
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)AdamW is a variant that handles weight decay in a way that is usually preferred for modern neural networks, especially transformers. In PyTorch, AdamW is a separate optimizer:
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
Important: if you want L2 style regularization via weight decay with Adam-like optimizers, prefer AdamW over setting weight_decay on Adam for most modern use cases.
Weight Decay and Learning Rate Interactions
Weight decay gently pushes weights toward zero during training. It is often controlled by a hyperparameter called weight_decay. Larger learning rates generally make weight decay act “stronger” in effect because updates are larger overall, so these two settings interact. If you increase the learning rate substantially, you may need to revisit weight decay.
Do not apply weight decay blindly to every parameter. It is common to exclude biases and normalization parameters in more advanced setups, but that parameter grouping pattern is a separate topic you can adopt later when you need it.
Parameter Groups for Different Learning Rates
PyTorch optimizers can assign different hyperparameters to different sets of parameters. This is useful when you want, for example, a smaller learning rate for a pretrained backbone and a larger one for a new head.
optimizer = torch.optim.AdamW([
{"params": model.backbone.parameters(), "lr": 1e-4},
{"params": model.head.parameters(), "lr": 1e-3},
], weight_decay=0.01)The optimizer will manage state separately for each parameter tensor, regardless of grouping, but the hyperparameters can differ per group.
Choosing a Starting Learning Rate
There is no universal best value, but beginners can start with reliable defaults.
For Adam, a common starting point is lr=1e-3.
For AdamW, common starting points are lr=3e-4 or lr=1e-3 depending on model size and task.
For SGD with momentum, a common starting point is lr=0.1 for moderately sized models and normalized inputs, but it is highly problem dependent.
Batch size affects the effective gradient noise. Larger batches often allow larger learning rates, but you should treat this as a tuning relationship rather than a rule you can apply mechanically.
Do not compare optimizers with different learning rates and conclude one is better after a single run. Each optimizer typically needs its own learning rate tuning to be judged fairly.
Practical Signs Your Learning Rate Is Wrong
If the learning rate is too high, the loss may bounce wildly, fail to decrease, or become NaN. Model parameters can also grow without bound, which you might notice as exploding outputs.
If the learning rate is too low, loss will decrease very slowly, training accuracy may improve at a crawl, and you may see almost no change across many steps even though nothing is “broken.”
A useful practice is to plot loss versus iteration, not just per epoch, because learning rate issues show up early at the iteration level.
What “Learning Rate Schedules” Are, at a High Level
A learning rate schedule changes the learning rate during training, for example warming up early, then decaying later. This chapter focuses on what the learning rate is and how it interacts with the optimizer. The mechanics of schedulers and common schedule types are covered separately.
Minimal Checklist for Beginners
When choosing an optimizer and learning rate for a first attempt, pick Adam or AdamW, set a reasonable initial learning rate, and verify that the loss decreases on a small subset of data. If it does not, lower the learning rate and try again before changing many other things at once.