Gradient Descent Cheat Sheet
The mechanics of gradient descent optimization, covering batch, stochastic, and mini-batch variants plus momentum and adaptive learning rate methods.
Gradient Descent From Scratch
Minimize MSE loss for linear regression.
import numpy as npdef gradient_descent(X, y, lr=0.01, epochs=1000): n, m = X.shape weights = np.zeros(m) bias = 0.0 for epoch in range(epochs): y_pred = X @ weights + bias error = y_pred - y # Gradients of MSE loss w.r.t. weights and bias grad_w = (2 / n) * X.T @ error grad_b = (2 / n) * np.sum(error) # Update parameters in the direction that reduces loss weights -= lr * grad_w bias -= lr * grad_b if epoch % 100 == 0: loss = np.mean(error ** 2) print(f"Epoch {epoch}: loss={loss:.4f}") return weights, bias
PyTorch Optimizers
SGD with momentum and Adam in a training loop.
import torch.nn as nnimport torch.optim as optimmodel = nn.Linear(10, 1)criterion = nn.MSELoss()# SGD with momentum: smooths updates using a moving average of past gradientsoptimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)# Adam: adapts the learning rate per parameter using 1st/2nd moment estimatesoptimizer = optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999))for epoch in range(100): optimizer.zero_grad() # clear old gradients output = model(X_batch) loss = criterion(output, y_batch) loss.backward() # compute gradients via backprop optimizer.step() # update weights
Gradient Descent Concepts
Variants and core vocabulary of the algorithm.
- Batch gradient descent- computes the gradient over the entire dataset per update; stable but slow for large data
- Stochastic gradient descent (SGD)- updates using one sample at a time; noisy but fast and can escape local minima
- Mini-batch gradient descent- updates using small batches (e.g. 32-256); standard in deep learning, balances speed and stability
- Learning rate- step size for each update; too high diverges, too low converges slowly
- Momentum- accumulates a moving average of past gradients to smooth updates and speed convergence
- Adam- combines momentum with per-parameter adaptive learning rates; a common default optimizer
- Learning rate schedule/decay- reduces the learning rate over training to fine-tune convergence
- Vanishing/exploding gradients- gradients shrink or grow uncontrollably in deep networks, hindering training
Optimizer Comparison
Trade-offs between common optimizers.
- SGD- simple, generalizes well, but sensitive to learning rate and can be slow to converge
- SGD + Momentum- accelerates convergence and dampens oscillations in ravines
- RMSprop- adapts learning rate per parameter using a moving average of squared gradients; good for RNNs
- Adam- combines momentum and RMSprop-style adaptive rates; fast convergence, a common default choice
- AdamW- Adam with decoupled weight decay; often preferred for transformer training
Learning Rate Schedulers & Warmup
Anneal the learning rate over training instead of holding it fixed, with a warmup phase for large-batch/transformer training.
import torch.optim as optimfrom torch.optim.lr_scheduler import CosineAnnealingLR, OneCycleLR, LambdaLRoptimizer = optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)# Cosine annealing: smoothly decays lr from initial value to ~0 over T_max stepscosine = CosineAnnealingLR(optimizer, T_max=1000, eta_min=1e-6)# One-cycle: ramps lr up then back down within a single training run,# often converges faster than a fixed schedule (Smith, 2018)one_cycle = OneCycleLR(optimizer, max_lr=1e-3, total_steps=1000)# Linear warmup then decay, the standard transformer recipedef warmup_then_decay(step, warmup_steps=500, total_steps=10000): if step < warmup_steps: return step / max(1, warmup_steps) return max(0.0, (total_steps - step) / max(1, total_steps - warmup_steps))warmup_sched = LambdaLR(optimizer, lr_lambda=warmup_then_decay)for step in range(total_steps): train_step() optimizer.step() warmup_sched.step() # call once per optimizer step, not per epoch
Gradient Clipping & Numerical Gradient Checking
Stabilize training against exploding gradients, and verify backprop correctness with finite differences.
import torchimport torch.nn.utils as utils# Clip by global norm: rescales all gradients together if their combined# L2 norm exceeds max_norm -> preserves gradient direction, standard for RNNs/transformersloss.backward()utils.clip_grad_norm_(model.parameters(), max_norm=1.0)optimizer.step()# Clip by value: hard-clamps each gradient element independently (cruder,# can distort direction, but simple and cheap)utils.clip_grad_value_(model.parameters(), clip_value=0.5)# Numerical gradient check (central difference) to validate a hand-written# gradient implementation before trusting it in productiondef numerical_gradient(f, w, eps=1e-5): grad = torch.zeros_like(w) for i in range(w.numel()): orig = w.view(-1)[i].item() w.view(-1)[i] = orig + eps f_plus = f() w.view(-1)[i] = orig - eps f_minus = f() w.view(-1)[i] = orig grad.view(-1)[i] = (f_plus - f_minus) / (2 * eps) return grad
Second-Order Methods (Newton / L-BFGS)
Use curvature (Hessian) information to converge in far fewer iterations than plain gradient descent, for small/medium models.
import torchfrom scipy.optimize import minimize# PyTorch's LBFGS: approximates the inverse Hessian from a history of# gradients/updates (quasi-Newton) -> needs a closure since it may# re-evaluate the loss multiple times per stepoptimizer = torch.optim.LBFGS(model.parameters(), lr=1.0, max_iter=20, history_size=10)def closure(): optimizer.zero_grad() output = model(X_train) loss = criterion(output, y_train) loss.backward() return lossoptimizer.step(closure)# scipy L-BFGS-B for classic (non-deep) models, e.g. custom loss with# an analytic gradient -- converges in far fewer iterations than SGD# when the full-batch Hessian is well-conditioneddef loss_and_grad(w): y_pred = X @ w err = y_pred - y loss = (err ** 2).mean() grad = (2 / len(y)) * X.T @ err return loss, gradresult = minimize(loss_and_grad, x0=w0, jac=True, method="L-BFGS-B")
Adam's Bias-Corrected Moment Estimates
Implement Adam manually to see why the m_hat/v_hat correction terms matter, especially early in training.
import numpy as npdef adam_step(w, grad, m, v, t, lr=1e-3, beta1=0.9, beta2=0.999, eps=1e-8): # 1st moment (mean) and 2nd moment (uncentered variance) of the gradient m = beta1 * m + (1 - beta1) * grad v = beta2 * v + (1 - beta2) * (grad ** 2) # Bias correction: m and v are initialized at 0, so early estimates are # biased toward 0, especially with beta1/beta2 close to 1. Without this # correction the first few updates would be far too small. m_hat = m / (1 - beta1 ** t) v_hat = v / (1 - beta2 ** t) w = w - lr * m_hat / (np.sqrt(v_hat) + eps) return w, m, v
Advanced Optimizer Variants
Beyond vanilla SGD/Adam — what each variant actually changes and when to reach for it.
- Nesterov momentum- evaluates the gradient at the 'look-ahead' position (current position + momentum step) rather than the current position, giving a corrective effect that plain momentum lacks
- AdaGrad- accumulates the sum of squared past gradients per parameter; learning rate monotonically shrinks, good for sparse features but stalls on long training runs
- RAdam- rectifies Adam's adaptive learning rate variance in early training, removing the need for a manual warmup phase
- Lookahead- wraps a base optimizer (e.g. Adam), periodically interpolating 'fast' weights back toward a slower-moving average of past weights to reduce variance
- LAMB- layer-wise adaptive learning rates scaled by the ratio of weight norm to update norm; enables stable training with very large batch sizes
- Condition number- ratio of the largest to smallest eigenvalue of the loss Hessian; ill-conditioned (elongated) loss surfaces cause plain gradient descent to zig-zag, which momentum/Adam mitigate
- Learning rate range test- sweep lr exponentially over a few hundred steps and plot loss vs lr to pick a good max_lr for OneCycle/warmup schedules (fast.ai / Leslie Smith technique)
If training loss oscillates wildly or diverges (NaN), your learning rate is almost always too high — halve it before assuming there's a bug in your model architecture, and consider gradient clipping for RNNs/transformers where exploding gradients are common.