Down the
Loss Landscape

You have parameters \(\theta \in \mathbb{R}^d\) and a dataset of \(N\) examples. You want the \(\theta\) that minimizes the average loss:

The one tool every first-order optimizer uses is the gradient \(g=\nabla_\theta L(\theta)\): the direction of steepest ascent. So we move the other way. Everything that follows — momentum, adaptivity, bias correction — is a different answer to one question: given the gradients I've seen so far, how should I actually move?

Two facts make this hard in deep learning, and they motivate almost every algorithm below:

1. The landscape is ill-conditioned. Curvature differs wildly across directions. A single global step size is too big for steep directions (you oscillate) and too small for flat ones (you crawl). The condition number \(\kappa=\lambda_{\max}/\lambda_{\min}\) of the Hessian can be \(10^6\)+.

2. You can't afford the true gradient. Computing \(\nabla L\) over all \(N\) examples per step is wasteful. So we estimate it from a mini-batch — fast, but noisy.

The simplest possible rule. Compute the full gradient, step against it, repeat.

\(\eta\) is the learning rate — the single most important hyperparameter in all of deep learning. For a convex quadratic with Hessian \(H\), GD converges only if \(\eta < 2/\lambda_{\max}(H)\), and its rate is governed by the condition number: error shrinks by a factor \(\big(\frac{\kappa-1}{\kappa+1}\big)\) each step. Large \(\kappa\) ⇒ painfully slow. This single inequality is why every other optimizer exists.

Instead of the full gradient, estimate it from a random mini-batch \(\mathcal{B}\) of \(b\) examples:

This estimate is unbiased (\(\mathbb{E}[\hat g_t]=\nabla L\)) but noisy, with variance \(\propto 1/b\). Robbins & Monro showed back in 1951^[1] that as long as the step sizes satisfy \(\sum_t \eta_t=\infty\) and \(\sum_t \eta_t^2<\infty\), this noisy process still converges. Modern practice mostly uses a constant-then-decayed \(\eta\) rather than that schedule, but the principle stands.

# The canonical training loop. Everything fancier swaps out the optimizer.
import torch

model = MyNet()
opt = torch.optim.SGD(model.parameters(), lr=0.1)   # plain SGD

for x, y in dataloader:                 # each batch = one noisy gradient estimate
    opt.zero_grad()                     # clear stale grads (PyTorch accumulates!)
    loss = loss_fn(model(x), y)
    loss.backward()                     # autograd fills p.grad for every param
    opt.step()                          # theta <- theta - lr * p.grad

Give the marble mass. Accumulate an exponentially-weighted velocity of past gradients and step with that, not the raw gradient^[2]:

The momentum coefficient \(\mu\) (typically 0.9) controls how much history persists. In flat directions, consecutive gradients agree and add up, so velocity grows and you accelerate. In oscillating (steep) directions, gradients flip sign each step and cancel, damping the zig-zag. That is exactly the disease SGD had.

opt = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
# Note: lr dropped 10x vs the plain-SGD example because of the eta/(1-mu) factor.

# From scratch, the heavy-ball update is just two lines per parameter:
v = mu * v - lr * g          # accumulate velocity (EMA of -gradient)
theta = theta + v            # step with velocity

Nesterov's twist^[3]: don't measure the gradient where you are — measure it where momentum is about to take you. Evaluate the gradient at the look-ahead point \(\theta_t+\mu v_{t-1}\):

The look-ahead acts as a correction term: if the velocity is about to overshoot, the gradient at the future point already points back, so NAG brakes earlier than plain momentum. On smooth convex problems this earns the famous \(O(1/t^2)\) rate versus GD's \(O(1/t)\). In deep learning the practical gain is modest but real, and Sutskever et al.^[4] showed careful momentum + good init was enough to train deep nets that were thought to need second-order methods.

opt = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9,
                      nesterov=True)   # the only change from plain momentum

A completely different idea: instead of one shared \(\eta\), give every parameter its own learning rate, shrinking it for parameters that have already seen large gradients^[5]. Accumulate the sum of squared gradients elementwise:

(\(\odot\) and the square are elementwise.) Frequent, large-gradient directions get damped; rare, small-gradient directions keep a large effective rate. This was a breakthrough for sparse features (NLP, recommendation) where rare features deserve big updates.

opt = torch.optim.Adagrad(model.parameters(), lr=0.01)
# G accumulates forever -> the LR you set is an *upper bound* that only decays.

Fix AdaGrad's vanishing rate by replacing the ever-growing sum with an exponential moving average (EMA) of squared gradients. Old gradients decay away, so the denominator stops growing without bound^[6]:

With \(\rho\approx 0.9\), the average tracks a moving window of recent curvature. RMSProp keeps AdaGrad's per-parameter adaptivity but learns indefinitely. Famously, it was never formally published — it comes from Geoff Hinton's Coursera lecture slides, and the whole field cites a lecture^[6].

opt = torch.optim.RMSprop(model.parameters(), lr=0.001, alpha=0.9)
# 'alpha' is the EMA decay (the rho above).

Zeiler's AdaDelta^[7] was developed alongside RMSProp to fix the same AdaGrad decay, plus a subtler complaint: in \(\eta\,g/\sqrt{G}\), the units don't match those of \(\theta\). The fix keeps a second EMA — of the squared updates — and uses its root to supply the step scale, eliminating \(\eta\) entirely:

Then \(\mathbb{E}[\Delta\theta^2]_t=\rho\,\mathbb{E}[\Delta\theta^2]_{t-1}+(1-\rho)\Delta\theta_t^2\). The ratio of two RMS quantities is dimensionless × gradient — the units now match \(\theta\). It's elegant but in practice tends to be conservative and slow to warm up, and Adam superseded it for most uses.

opt = torch.optim.Adadelta(model.parameters(), rho=0.95)  # no lr needed in theory
# PyTorch still exposes an lr multiplier (default 1.0) for practical control.

Adam^[8] is the single most consequential optimizer in modern deep learning. The idea is almost embarrassingly simple: RMSProp + momentum. Keep an EMA of the gradient (the first moment, like momentum) and an EMA of its square (the second moment, like RMSProp), then add one crucial correction.

Both \(m\) and \(v\) start at zero, so early in training they're biased toward zero — especially \(v\), with \(\beta_2=0.999\) decaying very slowly. Adam's signature move is bias correction:

Defaults that work shockingly often: \(\eta=10^{-3}\), \(\beta_1=0.9\), \(\beta_2=0.999\), \(\epsilon=10^{-8}\). Adam is scale-invariant in a useful sense: multiply all gradients by a constant and the update barely changes, because the \(\sqrt{\hat v}\) in the denominator rescales it away. That robustness is why it "just works" out of the box on architectures where SGD needs careful tuning.

Writing Adam from scratch is a beloved interview exercise. Here is the whole thing — no framework:

import numpy as np

def adam(grad_fn, theta, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8, steps=1000):
    m = np.zeros_like(theta)        # 1st moment  (EMA of g)
    v = np.zeros_like(theta)        # 2nd moment  (EMA of g**2)
    for t in range(1, steps + 1):
        g  = grad_fn(theta)
        m  = b1 * m + (1 - b1) * g
        v  = b2 * v + (1 - b2) * g * g
        m_hat = m / (1 - b1**t)     # bias correction -- the part people forget
        v_hat = v / (1 - b2**t)
        theta = theta - lr * m_hat / (np.sqrt(v_hat) + eps)
    return theta

# In practice, just:  opt = torch.optim.Adam(model.parameters(), lr=1e-3)

Reddi, Kale & Kumar^[9] showed Adam can fail to converge even on simple convex problems. The culprit: because \(v_t\) is an EMA, it can shrink when a run of small gradients follows a big one, so the effective step can suddenly grow at the wrong moment and undo progress. Their fix is one line — never let the denominator decrease:

By taking a running maximum of the second moment, AMSGrad guarantees a non-increasing effective learning rate and recovers Adam's convergence proof. In practice the gain is often small — Adam with good hyperparameters usually behaves — but this is the canonical "Adam is not theoretically sound" story, and interviewers love it.

opt = torch.optim.Adam(model.parameters(), lr=1e-3, amsgrad=True)

If there is one optimizer to understand cold for a frontier-lab interview, it's AdamW^[10] — the default for training essentially every modern transformer. Its entire contribution is fixing how Adam handles weight decay, and the distinction is subtle enough that it trips up most candidates.

The classic confusion: L2 regularization ≠ weight decay. For plain SGD they're identical. For adaptive optimizers they are not. With L2, you add \(\lambda\theta\) to the gradient — which then gets fed through Adam's adaptive denominator:

See the problem? The decay term \(\lambda\theta\) gets divided by \(\sqrt{\hat v}\) too. Parameters with large historical gradients (large \(\hat v\)) get their regularization weakened — the opposite of what you want. Weight decay becomes entangled with gradient magnitude. AdamW decouples them: do the Adam step, then shrink the weights separately, untouched by the adaptive scaling:

# The standard recipe for training transformers today:
opt = torch.optim.AdamW(model.parameters(),
                        lr=3e-4,            # the "Karpathy constant" for many LLMs
                        betas=(0.9, 0.95),  # b2=0.95 common at scale (vs 0.999)
                        weight_decay=0.1)   # decoupled -- this is real weight decay

# Pro move: never decay biases or LayerNorm/embedding gains.
decay, no_decay = [], []
for n, p in model.named_parameters():
    (no_decay if p.ndim < 2 else decay).append(p)   # 1D params -> no decay
opt = torch.optim.AdamW(
    [{"params": decay, "weight_decay": 0.1},
     {"params": no_decay, "weight_decay": 0.0}], lr=3e-4, betas=(0.9, 0.95))

Nadam^[11] does to Adam what Nesterov did to momentum: replace the plain first-moment with a look-ahead version. The update folds the next step's momentum into the current one:

The benefit mirrors NAG's — slightly faster, slightly better-damped convergence on many problems — at no extra memory cost. It's a solid default where available, though AdamW dominates the LLM world.

opt = torch.optim.NAdam(model.parameters(), lr=2e-3)

AdamW is the workhorse, but the bleeding edge has moved on in specific directions. Know these names and the one idea behind each — they show up the moment an interview turns to "what are you excited about."

Lion — sign of the momentum

Discovered by program search rather than human design^[12], Lion (EvoLved Sign Momentum) is startlingly simple and uses half the memory of Adam (it keeps only one momentum buffer, no second moment):

The \(\operatorname{sign}\) makes every update the same magnitude per coordinate — pure direction, no scale. That demands a much smaller LR (≈3–10× lower than AdamW) and larger weight decay, but it can match or beat AdamW on vision and language at lower cost.

LARS / LAMB — layer-wise rates for huge batches

When you scale batch size to tens of thousands to train faster on many GPUs, per-layer gradient norms vary enormously and training destabilizes. LARS^[13] (for SGD) and LAMB^[14] (its Adam-based successor) rescale each layer's update by the ratio of its weight norm to its update norm — a "trust ratio." LAMB famously cut BERT pre-training from days to 76 minutes by enabling batch sizes of 32k+.

Shampoo, Sophia, Muon — toward (cheap) second order

Adam's \(\sqrt{v}\) is a diagonal curvature approximation — it ignores interactions between parameters. The frontier of optimizer research is capturing more of the curvature without paying for the full \(d\times d\) Hessian:

Shampoo^[17] keeps small preconditioner matrices per tensor dimension (a Kronecker-factored approximation of the curvature) and has powered record-setting training runs. Sophia^[18] uses a light, periodically-estimated diagonal Hessian to halve the steps needed for LLM pre-training. Muon^[19] orthogonalizes the momentum update for 2-D weight matrices (via a few Newton–Schulz iterations) and has become a popular, genuinely faster alternative to AdamW for the matrix parameters of transformers, often paired with AdamW on the 1-D params.

The optimizer chooses direction and per-parameter scaling; the schedule chooses the global step size over time. At scale, the schedule matters as much as the optimizer choice.

Warmup — why Adam needs it

Start \(\eta\) near zero and ramp up linearly over the first few thousand steps. The reason is precise: early on, \(\hat v_t\) is estimated from very few gradients, so its variance is huge; dividing by an unreliable \(\sqrt{\hat v}\) produces wild steps. Warmup buys time for the second-moment estimate to stabilize. (This is also the motivation behind RAdam, which derives a warmup automatically from the variance of the adaptive rate.)

Cosine decay

After warmup, anneal \(\eta\) along a half-cosine to (near) zero^[20]. Large early steps explore; small late steps settle into a minimum. It's the default LLM schedule:

Linear scaling rule & gradient clipping

Double the batch size → roughly double the LR^[21], because a larger batch gives a lower-variance gradient you can trust with a bigger step (holds up to a "critical batch size," then breaks). And gradient clipping — capping \(\lVert g\rVert\) to a threshold — is non-negotiable for transformers/RNNs: a single bad batch can produce an enormous gradient that blows up the run; clipping caps the damage.

opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)
warmup, total = 2000, 100_000
sched = torch.optim.lr_scheduler.SequentialLR(opt, schedulers=[
    torch.optim.lr_scheduler.LinearLR(opt, 1e-3, 1.0, total_iters=warmup),
    torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=total - warmup)],
    milestones=[warmup])

for step, (x, y) in enumerate(dataloader):
    opt.zero_grad()
    loss_fn(model(x), y).backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # clip!
    opt.step()
    sched.step()

Here is the whole family released from the same starting point on the same landscape, each running its true update rule live in your browser. Switch landscapes to see how the ranking flips — there is no universally best optimizer, only best for a given surface and budget. This is the figure to keep in your head.

Reading the race: on the ravine, adaptive methods (AdaGrad, RMSProp, Adam family) drive nearly straight down while SGD zig-zags. On the saddle, methods that amplify tiny gradients escape the plateau first — and AdaDelta may barely move. On Beale, the curved valley punishes anything without adaptivity; plain SGD/momentum stall on the flat plateau while Adam & friends thread the canyon to \((3,0.5)\).

Tap each to reveal a crisp, interview-ready answer. If you can deliver these from memory, you understand optimizers better than most candidates.

1. [1] Robbins & Monro (1951). A Stochastic Approximation Method. Ann. Math. Stat.
2. [2] Polyak (1964). Some methods of speeding up the convergence of iteration methods (heavy-ball momentum).
3. [3] Nesterov (1983). A method for solving the convex programming problem with rate O(1/k²).
4. [4] Sutskever, Martens, Dahl & Hinton (2013). On the importance of initialization and momentum in deep learning. ICML.
5. [5] Duchi, Hazan & Singer (2011). Adaptive Subgradient Methods (AdaGrad). JMLR.
6. [6] Tieleman & Hinton (2012). RMSProp. Coursera "Neural Networks for ML," Lecture 6e.
7. [7] Zeiler (2012). ADADELTA: An Adaptive Learning Rate Method. arXiv:1212.5701.
8. [8] Kingma & Ba (2015). Adam: A Method for Stochastic Optimization. ICLR.
9. [9] Reddi, Kale & Kumar (2018). On the Convergence of Adam and Beyond (AMSGrad). ICLR.
10. [10] Loshchilov & Hutter (2019). Decoupled Weight Decay Regularization (AdamW). ICLR.
11. [11] Dozat (2016). Incorporating Nesterov Momentum into Adam (Nadam). ICLR Workshop.
12. [12] Chen et al. (2023). Symbolic Discovery of Optimization Algorithms (Lion). NeurIPS.
13. [13] You, Gitman & Ginsburg (2017). Large Batch Training of CNNs (LARS). arXiv:1708.03888.
14. [14] You et al. (2020). Large Batch Optimization for Deep Learning (LAMB). ICLR.
15. [15] Keskar et al. (2017). On Large-Batch Training... Sharp vs Flat Minima. ICLR.
16. [16] Wilson et al. (2017). The Marginal Value of Adaptive Gradient Methods. NeurIPS.
17. [17] Gupta, Koren & Singer (2018). Shampoo: Preconditioned Stochastic Tensor Optimization. ICML.
18. [18] Liu et al. (2023). Sophia: A Scalable Stochastic Second-order Optimizer. arXiv:2305.14342.
19. [19] Jordan et al. (2024). Muon: an optimizer for the hidden layers of neural networks.
20. [20] Loshchilov & Hutter (2017). SGDR: Stochastic Gradient Descent with Warm Restarts (cosine). ICLR.
21. [21] Goyal et al. (2017). Accurate, Large Minibatch SGD (linear scaling + warmup). arXiv:1706.02677.

Dates and venues are given from memory of the literature; verify exact citations against the original papers when you submit work. The animations implement each method's published update rule with hyperparameters tuned per landscape for legibility, not to declare a "winner."