Deep Learning Optimizers Cheat Sheet
Compares SGD, Momentum, RMSprop, Adam, and AdamW update rules, and shows how to configure each optimizer in PyTorch with typical hyperparameters.
Optimizer Families
How each optimizer adapts the learning rate or update direction.
- SGD- Updates weights using the gradient of a mini-batch: w -= lr * grad; simple but sensitive to learning rate choice
- SGD with Momentum- Accumulates a velocity term from past gradients to smooth updates and speed convergence through ravines
- RMSprop- Divides the learning rate by a moving average of recent squared gradients, adapting per-parameter step size
- Adam- Combines momentum (first moment) and RMSprop-style scaling (second moment) with bias correction
- AdamW- Adam with decoupled weight decay, applied directly to weights instead of folded into the gradient -- the modern default for transformers
- Learning rate schedule- Adjusts the learning rate over training (step decay, cosine annealing, warmup) independent of the optimizer's own adaptation
Configuring Optimizers in PyTorch
Typical setup and hyperparameters for common optimizers.
import torch.optim as optim# Vanilla SGD with momentumopt = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-4)# Adam - good default for most non-transformer modelsopt = optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.999), eps=1e-8)# AdamW - standard choice for training transformersopt = optim.AdamW(model.parameters(), lr=5e-4, weight_decay=0.01)# Cosine annealing schedule on top of any optimizerscheduler = optim.lr_scheduler.CosineAnnealingLR(opt, T_max=50)
Adam Update Rule
The core math Adam performs each step (bias-corrected first/second moments).
# m, v initialized to zero; t = timestep; g = gradientm = beta1 * m + (1 - beta1) * g # 1st moment (mean)v = beta2 * v + (1 - beta2) * (g ** 2) # 2nd moment (uncentered variance)m_hat = m / (1 - beta1 ** t) # bias correctionv_hat = v / (1 - beta2 ** t)w -= lr * m_hat / (v_hat ** 0.5 + eps)
Choosing an Optimizer
Rules of thumb from practice.
- Default starting point- Adam or AdamW with lr around 1e-3 (CNNs/MLPs) or 1e-4 to 5e-4 (transformers)
- When SGD+momentum wins- Often generalizes better than Adam on large-scale image classification given enough tuning and a good LR schedule
- Gradient clipping- Clip gradient norm (e.g., to 1.0) to stabilize RNN/transformer training regardless of optimizer
- Warmup- Linearly ramp the learning rate up for the first few hundred/thousand steps before decaying -- critical for AdamW on transformers
Nesterov Accelerated Gradient
Evaluates the gradient at a look-ahead position rather than the current point, giving a correction term over classical momentum.
# Classical momentum: v = beta*v + grad(w); w -= lr*v# Nesterov: gradient is evaluated at the 'look-ahead' point w - beta*vv = beta * v + grad(w - lr * beta * v)w -= lr * v# In PyTorch, just flip a flag on SGDimport torch.optim as optimopt = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, nesterov=True)
Writing a Custom PyTorch Optimizer
Subclass torch.optim.Optimizer to implement a bespoke update rule (skeleton for a signSGD-style optimizer).
import torchfrom torch.optim import Optimizerclass SignSGD(Optimizer): def __init__(self, params, lr=1e-3): defaults = dict(lr=lr) super().__init__(params, defaults) @torch.no_grad() def step(self, closure=None): loss = None if closure is not None: with torch.enable_grad(): loss = closure() for group in self.param_groups: for p in group['params']: if p.grad is None: continue update = torch.sign(p.grad) p.add_(update, alpha=-group['lr']) return loss# usage: opt = SignSGD(model.parameters(), lr=1e-3)
Beyond Adam: Modern Optimizers
Newer optimizers used in large-scale training, and what problem each targets.
- LAMB- Layer-wise Adaptive Moments; normalizes updates by parameter norm per layer, enabling very large batch sizes (used to train BERT in minutes)
- Lion- Uses only the sign of a momentum-smoothed gradient update; smaller memory footprint than Adam and competitive on vision/language models
- Adafactor- Factorizes the second-moment matrix into row/column statistics, cutting optimizer memory roughly in half vs Adam -- common for training very large transformers
- AMSGrad- Fixes a theoretical non-convergence issue in Adam by keeping a running max of past second moments instead of an exponential average
- Shampoo- A second-order preconditioning method that approximates the full-matrix AdaGrad using Kronecker-factored per-layer preconditioners
- Sophia- Uses a diagonal Hessian estimate (via Hutchinson's estimator) with clipping, aiming for faster convergence than AdamW on LLM pretraining
Custom Warmup + Cosine Decay Schedule
Hand-rolled LambdaLR combining a linear warmup phase with cosine decay -- the standard transformer training recipe.
import mathimport torch.optim as optimdef warmup_cosine_lambda(step, warmup_steps, total_steps, min_lr_ratio=0.1): if step < warmup_steps: return step / max(1, warmup_steps) progress = (step - warmup_steps) / max(1, total_steps - warmup_steps) cosine = 0.5 * (1 + math.cos(math.pi * progress)) return min_lr_ratio + (1 - min_lr_ratio) * cosineopt = optim.AdamW(model.parameters(), lr=5e-4)scheduler = optim.lr_scheduler.LambdaLR( opt, lr_lambda=lambda step: warmup_cosine_lambda(step, warmup_steps=1000, total_steps=100_000),)# training loop: call scheduler.step() once per optimizer step, not per epoch
Gradient Accumulation
Simulate a larger effective batch size than fits in memory by accumulating gradients over several forward/backward passes before stepping.
accumulation_steps = 4opt.zero_grad()for i, (inputs, targets) in enumerate(dataloader): outputs = model(inputs) loss = criterion(outputs, targets) / accumulation_steps # scale down loss.backward() if (i + 1) % accumulation_steps == 0: torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) opt.step() opt.zero_grad()
AdamW's weight decay is decoupled from the gradient update -- if you switch from Adam to AdamW, re-tune weight_decay separately rather than reusing the same value, since it now behaves like true L2 regularization.