Optimizers Explained: SGD / Momentum / Adam / AdamW

You computed a gradient (L1-06 / L3-02). Now what?

The optimizer decides HOW to use the gradient to update parameters.

The choice of optimizer can mean 2× faster training, or training that doesn’t converge at all.

Vanilla Gradient Descent

The simplest update rule:

$$\theta_{t+1} = \theta_t - \eta \cdot \nabla L(\theta_t)$$

Where $\eta$ is the learning rate (typically 0.001 to 0.1).

for batch in data:
    loss = compute_loss(model, batch)
    grad = compute_gradient(loss, model.parameters)
    for p in model.parameters:
        p -= learning_rate * p.grad

Works but slow and gets stuck in plateaus.

SGD (Stochastic Gradient Descent)

Instead of full dataset, use mini-batches:

$$\theta_{t+1} = \theta_t - \eta \cdot \nabla L(\theta_t; \text{mini-batch})$$

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

for batch in data:
    optimizer.zero_grad()
    loss = compute_loss(model, batch)
    loss.backward()
    optimizer.step()

Faster than full-batch + noise helps escape local minima.

Problem: noisy gradients oscillate around the minimum.

Momentum

Add momentum — accumulate past gradients:

$$v_{t+1} = \beta v_t + (1-\beta) \nabla L$$ $$\theta_{t+1} = \theta_t - \eta \cdot v_{t+1}$$

Where $\beta = 0.9$ typically. Acts like a heavy ball rolling downhill — momentum smooths through small bumps.

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

Pros: smooths oscillation, accelerates in consistent direction. Cons: can overshoot near minimum.

Nesterov Accelerated Gradient (NAG)

A subtle improvement on momentum: compute gradient at the “look-ahead” position:

$$v_{t+1} = \beta v_t + \nabla L(\theta_t - \beta v_t)$$

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

Slightly better than vanilla momentum. Rarely makes a big difference in practice.

AdaGrad (2011)

Different parameters need different learning rates. AdaGrad adapts per-parameter:

$$\theta_{t+1} = \theta_t - \frac{\eta}{\sqrt{G_t + \epsilon}} \nabla L$$

Where $G_t = \sum_{i=0}^t (\nabla L_i)^2$ — accumulated squared gradients.

Effect:

• Parameters with large gradient history → smaller updates
• Parameters with small gradient history → larger updates

Pros: handles sparse features well. Cons: learning rate monotonically decreases → eventually too small to learn.

RMSProp (2012)

AdaGrad’s problem: $G_t$ grows forever. RMSProp uses exponential moving average of squared gradients:

$$G_t = \gamma G_{t-1} + (1-\gamma)(\nabla L)^2$$

Now $G_t$ stays bounded. Learning rate doesn’t decay to zero.

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

Used in: early RL (DQN), still occasionally seen.

Adam (2014)

The most popular optimizer. Combines momentum + RMSProp:

$$m_t = \beta_1 m_{t-1} + (1-\beta_1) \nabla L \quad\text{(momentum)}$$ $$v_t = \beta_2 v_{t-1} + (1-\beta_2) (\nabla L)^2 \quad\text{(RMSProp)}$$ $$\theta_{t+1} = \theta_t - \frac{\eta \cdot \hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$$

(With bias correction $\hat{m}_t$ and $\hat{v}_t$ for early steps.)

optimizer = torch.optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999))

Defaults ($\beta_1=0.9, \beta_2=0.999, \epsilon=10^{-8}$) work for most cases.

Pros: fast convergence, adaptive, robust to hyperparameters. Cons: sometimes worse generalization than SGD on convex problems.

AdamW (2019)

Adam with decoupled weight decay:

$$\theta_{t+1} = \theta_t - \eta\left(\frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda \theta_t\right)$$

The L2 regularization is applied separately from the gradient update — this fixes a subtle bug in original Adam where L2 mixed badly with adaptive rates.

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

Strongly preferred for Transformer training. Used in GPT, BERT, LLaMA.

Comparison

Optimizer	Best for	Notes
SGD	Convex problems, small models	Slow but stable
SGD + Momentum	CNN training (classical)	Often best for ResNet etc.
Adam	Most things, prototypes	Fast, works on most setups
AdamW	Transformers, modern LLMs	Default for new projects
RMSProp	Specific RL setups	Mostly legacy now

Learning Rate Schedules

The learning rate $\eta$ shouldn’t be fixed. Common schedules:

Step Decay

scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
# Drop LR by 10× every 30 epochs

Cosine Annealing

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
# Smooth cosine decay from initial LR to 0 over T_max epochs

Often used for fine-tuning. Smoother than step decay.

Warmup + Cosine (used in LLM training)

from transformers import get_cosine_schedule_with_warmup

scheduler = get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=1000,
    num_training_steps=total_steps
)

• First N steps: linear ramp from 0 to peak LR (warmup)
• Then: cosine decay to 0

Standard for modern Transformer training.

Why Warmup

In early training, Adam’s $\hat{v}_t$ is small (few samples) — leads to huge effective learning rates.

Warmup prevents this initial instability.

How to Choose LR

The most important hyperparameter.

Rule of thumb (Adam family)

1e-3  → typical default
1e-4  → fine-tuning (e.g., pretrained models)
1e-5  → very fine adjustment
3e-4  → "Karpathy's constant" — works for many setups

LR Range Test

Train for a few iterations across many LRs, plot loss:

from torch_lr_finder import LRFinder

lr_finder = LRFinder(model, optimizer, criterion, device="cuda")
lr_finder.range_test(train_loader, end_lr=1, num_iter=100)
lr_finder.plot()  # → pick the LR just before loss explodes

A diagnostic, not a final answer.

Common Pitfalls

1. Too-high learning rate

Symptoms:

• Loss explodes (NaN)
• Loss oscillates wildly
• Model never converges

Fix: lower LR by 10×.

2. Too-low learning rate

Symptoms:

• Loss barely decreases over many epochs
• Stuck in plateau

Fix: increase LR by 10×.

3. Forgetting weight decay

For modern Transformer training, AdamW + weight decay 0.01 to 0.1 is the norm. Forgetting this hurts generalization.

4. Wrong optimizer for the task

• Don’t use Adam for simple linear models (SGD is enough + better generalization)
• Don’t use SGD for Transformers (much slower convergence than AdamW)

Code Recipe: Transformer Training

import torch
from torch.optim import AdamW
from transformers import get_cosine_schedule_with_warmup

# Different weight decay for different param groups
decay_params = []
no_decay_params = []
for name, param in model.named_parameters():
    if "bias" in name or "LayerNorm" in name:
        no_decay_params.append(param)
    else:
        decay_params.append(param)

optimizer = AdamW(
    [
        {"params": decay_params, "weight_decay": 0.01},
        {"params": no_decay_params, "weight_decay": 0.0}
    ],
    lr=1e-4,
    betas=(0.9, 0.95),     # different for LLM training
)

scheduler = get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=2000,
    num_training_steps=100000
)

for batch in loader:
    loss = compute_loss(model, batch)
    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # gradient clipping
    optimizer.step()
    scheduler.step()

This is how every modern LLM is trained (in broad outline).

Beyond Adam

Recent optimizers:

• Lion (Google 2023): sign-based update, less memory, slightly better in some tasks
• Sophia (Stanford 2023): uses Hessian information
• Shampoo: full-matrix preconditioning

In practice: AdamW remains the default for 95% of cases.

Next recommended: L3-01 Perceptron to MLP or L3-03 Training Tricks.