PyTorch Optimizers

An optimizer adjusts a model's weights after each training step to reduce the loss. Think of the optimizer as a hiker trying to reach the lowest point in a valley — the loss landscape. The gradients tell the hiker which direction is downhill. The optimizer decides how big a step to take and in which direction.

What Optimizers Do

Training Step:
  1. Forward pass  → compute predictions
  2. Loss          → measure error
  3. backward()    → compute gradients
  4. Optimizer     → update weights using gradients
  5. zero_grad()   → clear gradients for next step

SGD – Stochastic Gradient Descent

SGD is the simplest optimizer. It moves each weight in the direction that reduces the loss by a fixed step size called the learning rate.

import torch.optim as optim

optimizer = optim.SGD(model.parameters(), lr=0.01)

The update rule: weight = weight - lr × gradient

Weight update diagram:
  Current weight: 0.8
  Gradient:       0.5  (loss increases as weight increases)
  Learning rate:  0.01
  New weight:     0.8 - 0.01 × 0.5 = 0.795

SGD with Momentum

Plain SGD can oscillate back and forth before converging. Momentum helps SGD continue moving in a consistent direction by accumulating a "velocity" that resists sharp direction changes — like a ball rolling downhill that builds up speed.

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

Adam – Adaptive Moment Estimation

Adam is the most popular optimizer for deep learning. It adapts the learning rate individually for each parameter based on recent gradient history. Parameters that rarely receive updates get a larger effective learning rate; those updated frequently get a smaller one.

optimizer = optim.Adam(model.parameters(), lr=0.001)

Adam works well out of the box for most tasks. It is the default starting point for new projects.

AdamW – Adam with Weight Decay

AdamW adds weight decay (L2 regularization) to Adam in a mathematically correct way. Standard Adam's weight decay implementation has a known bug where regularization interacts incorrectly with the adaptive learning rate. AdamW fixes this. It is the preferred optimizer for transformer models.

optimizer = optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.01)

RMSprop

RMSprop adapts learning rates based on a moving average of recent squared gradients. It works well for RNNs and problems where gradients vary dramatically in size.

optimizer = optim.RMSprop(model.parameters(), lr=0.001)

Optimizer Comparison

Optimizer   Speed     Stability   Best For
──────────────────────────────────────────────────────
SGD         Fast      Low         CNNs with tuning
SGD+Momentum Fast     Medium      General training
Adam        Medium    High        Most tasks, quick start
AdamW       Medium    High        Transformers, NLP
RMSprop     Medium    High        RNNs, noisy gradients

Using the Optimizer in Training

import torch
import torch.nn as nn
import torch.optim as optim

model = nn.Linear(10, 1)
optimizer = optim.Adam(model.parameters(), lr=0.001)
loss_fn = nn.MSELoss()

X = torch.rand(32, 10)
y = torch.rand(32, 1)

for epoch in range(100):
    optimizer.zero_grad()       # 1. clear old gradients
    pred = model(X)             # 2. forward pass
    loss = loss_fn(pred, y)     # 3. compute loss
    loss.backward()             # 4. compute gradients
    optimizer.step()            # 5. update weights

    if epoch % 10 == 0:
        print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

Learning Rate — The Most Important Hyperparameter

The learning rate controls how big each weight update step is. Too large: training is unstable, loss oscillates or explodes. Too small: training is very slow, may get stuck.

Learning rate too high:
  Loss: 5 → 8 → 3 → 9 → 2  ← jumps around, no convergence

Learning rate too low:
  Loss: 5.000 → 4.999 → 4.998  ← barely moving

Good learning rate:
  Loss: 5 → 3.5 → 2.1 → 1.2 → 0.5  ← steady decrease

Setting Per-Parameter Learning Rates

You can give different layers different learning rates by passing a list of parameter groups to the optimizer.

import torch.nn as nn
import torch.optim as optim

model = nn.Sequential(nn.Linear(10, 50), nn.ReLU(), nn.Linear(50, 1))

optimizer = optim.Adam([
    {'params': model[0].parameters(), 'lr': 0.0001},   # small lr for layer 0
    {'params': model[2].parameters(), 'lr': 0.001}     # larger lr for layer 2
])

Summary

Optimizers use gradients to update model weights after each training step. SGD is the simplest optimizer. Adam is the most widely used — it adapts learning rates per parameter and works well without tuning. AdamW is preferred for transformers. Always call zero_grad() before backward() to prevent gradient accumulation. The learning rate is the single most important hyperparameter to tune. When starting a new project, begin with Adam at lr=0.001 and adjust from there.