Stochastic Gradient Descent

The gradient is computed on a minibatch rather than the full dataset, which is the “stochastic” part and is what makes it tractable. The noise this introduces is not purely a cost — it is a regulariser, and small-batch SGD generalises better than large-batch SGD on the same number of epochs, which is not what a pure optimisation view predicts.

Momentum

Plain SGD oscillates across the narrow direction of a ravine and crawls along the long one. Momentum accumulates the consistent component and cancels the oscillating one.

At $\mu = 0.9$ the effective step is roughly $1/(1-\mu) = 10$ gradients averaged, so the learning rate is doing ten times the work it appears to. This is the single most common source of confusion when porting a configuration between optimisers.

Why it fails on transformers

This is the interesting part, and it is not settled.

The observation is robust: on a language model, SGD with any learning rate reaches a substantially worse loss than AdamW, and the gap widens with scale. On a ResNet the ordering reverses. Same optimiser, same framework, opposite conclusion.

The best current explanation is Kunstner et al.’s. Token frequency is heavy- tailed — a handful of tokens account for most of the corpus, and the long tail is where most of the classes are. The gradient contribution of a rare token is small and infrequent; under SGD its parameters receive a small step, and it takes enormously many occurrences to learn. Adam’s per-parameter normalisation divides by the second moment, so a rare, consistently-signed gradient produces a full- sized step.

They show this directly by training with plain sign-SGD, which has none of Adam’s other machinery and recovers most of the gap. The adaptivity is not estimating curvature — it is discarding magnitude, and magnitude is precisely what makes rare tokens invisible to SGD.

import torch
from torch import Tensor


@torch.no_grad()
def sgd_step(
    params: list[Tensor], state: dict, lr: float, momentum: float = 0.9, wd: float = 0.0
) -> None:
    for p in params:
        if p.grad is None:
            continue
        g = p.grad

        # Decoupled: decay the weight, not the gradient. Folding it in makes
        # the effective decay depend on the momentum buffer.
        if wd:
            p.mul_(1 - lr * wd)

        buf = state.setdefault(p, torch.zeros_like(p))
        buf.mul_(momentum).add_(g)
        p.add_(buf, alpha=-lr)

If the useful part of Adam is the sign, an optimiser can take the sign and skip the second moment — halving the state.

Lion does exactly that: momentum, then $\mathrm{sign}$, with different decay rates for the update and the buffer. Found by symbolic search over optimiser programs, one state tensor instead of two, and competitive with AdamW on both vision and language.

Muon generalises the idea to matrices — where sign is the elementwise version, orthogonalisation is the matrix version — and is the strongest current challenger.

Both are readable as answers to the same question this entry ends on: what is Adam’s normalisation actually for, and how much of the machinery around it is load-bearing.