SGD and momentum

What it is and when you reach for it

SGD updates parameters by stepping against a gradient estimated on a random minibatch: θ ← θ − η·g, where g averages the per-example gradients of the loss. Add a momentum buffer that accumulates gradients across steps and you have the optimizer that trained AlexNet, every ResNet, and most of computer vision since. It sits at the base of the optimizer family tree: Adam is SGD with per-coordinate step sizes bolted on, and every schedule trick on the learning-rate schedules page modulates the same η that appears here. You reach for SGD with momentum when you want the fewest moving parts, when the problem is a convnet where it is still the strongest known recipe, or when you need a baseline whose behavior you can actually reason about. You reach past it, usually to Adam, when gradients across parameters differ in scale by orders of magnitude, which is the situation in transformers.

The math

Gradient descent on a quadratic, and the condition number

The clean setting is the quadratic f(x) = ½·x^TAx with A symmetric positive definite, because near any minimum every smooth loss looks like this with A the Hessian. The gradient is Ax, so gradient descent iterates x_t+1 = (I − ηA)x_t. Decompose x along the eigenvectors of A: the component along the eigenvector with eigenvalue λ_i is multiplied by (1 − ηλ_i) every step. Convergence therefore requires |1 − ηλ_i| < 1 for every eigenvalue, which pins η < 2/λ_max, and the overall rate is set by whichever direction contracts slowest. The best single η balances the two extremes, η* = 2/(λ_min + λ_max), and gives the worst-direction contraction factor

rate = (κ − 1)/(κ + 1), where κ = λ_max/λ_min is the condition number.

A small worked example makes the pain concrete. Take A = diag(1, 100), so κ = 100. The optimal step is η* = 2/101 ≈ 0.0198 and both directions contract by exactly 0.9802 per step. Halving the error takes ln 2 / ln(1/0.9802) ≈ 35 steps; reaching 10⁻⁶ of the initial error takes about 690. The steep direction is not the problem, it would converge in a handful of steps on its own; the problem is that the steep direction caps the step size while the shallow direction sets the distance, and one η must serve both. Deep-network Hessians have condition numbers far beyond 100, which is why plain gradient descent is never used bare.

Stochasticity: noise as tax and as regularizer

The minibatch gradient is an unbiased estimate of the full gradient, with covariance shrinking like Σ/B for batch size B. This costs something and buys something. The cost: near a minimum the updates no longer converge but bounce inside a noise ball whose radius grows with η and with the noise scale, roughly η·Σ/B for SGD on a quadratic. That is the fundamental reason the learning rate must decay over training, which is the subject of the schedules page, and it is why training runs like nanoGPT's end at a small fraction of their peak learning rate rather than at the peak. The benefit is subtler and partly empirical: the noise acts like a temperature that helps trajectories leave sharp, narrow basins and settle in flat ones, and flat minima correlate with better generalization in practice. Keskar et al. observed that very large batches (less noise) tend toward sharper minima and worse test accuracy, and Goyal et al.'s linear scaling rule, multiply η by k when you multiply B by k, can be read as keeping the noise scale η/B roughly constant. These are careful empirical regularities, not theorems, and both have documented exceptions at extreme scales, but the working intuition holds up well: minibatch noise is not merely tolerated, part of SGD's generalization behavior appears to depend on it.

Heavy-ball momentum, derived twice

Momentum replaces the raw gradient with a buffer that remembers past gradients:

vt = μ·vt−1 + gt          (buffer, μ ≈ 0.9)
xt+1 = xt − η·vt

The first view is the averaging view. Unrolling the recursion, v_t = g_t + μ·g_t−1 + μ²·g_t−2 + …, a geometrically weighted sum of the gradient history with an effective memory of about 1/(1 − μ) steps. Along a direction where consecutive gradients agree, the geometric series compounds: with a constant gradient g the buffer converges to g/(1 − μ), so at μ = 0.9 the steady-state step is ten times the plain-SGD step. Along a direction where the iterate overshoots and the gradient flips sign every step, the terms cancel and the buffer stays small. That single mechanism is the whole trick: momentum amplifies persistent gradient directions by up to 1/(1 − μ) and damps oscillating ones, which is exactly the medicine the ill-conditioned quadratic needs. On the quadratic, tuned heavy ball (Polyak, 1964) achieves rate (√κ − 1)/(√κ + 1) with μ* set to the square of that quantity. For the κ = 100 example the rate improves from 0.980 to 9/11 ≈ 0.818: error halves in about 3.5 steps instead of 35, and 10⁻⁶ arrives in roughly 70 steps instead of 690. The square root of the condition number is the entire headline.

The second view is physical. Rewrite the update by eliminating v: x_t+1 = x_t + μ·(x_t − x_t−1) − η·∇f(x_t). This is a discretization of a ball with mass rolling on the loss surface with friction, m·ẍ + c·ẋ = −∇f(x), where μ plays the role of (1 − friction). A frictionless ball oscillates forever; an overdamped one is plain gradient descent; the useful regime is in between, where inertia carries the iterate through narrow ravines and shallow noise without stalling. The physical picture also predicts the failure mode: too much mass (μ too close to 1) and the ball orbits the minimum instead of settling, which is exactly the ringing you see on loss curves when momentum is set too high for the learning rate.

Nesterov's lookahead

Nesterov's variant evaluates the gradient not at the current point but at the point the momentum is about to carry you to: v_t = μ·v_t−1 + ∇f(x_t − η·μ·v_t−1). If the buffer is about to overshoot, the lookahead gradient already points back, so the correction arrives one step earlier than heavy ball's. On smooth convex problems this is the accelerated method with the optimal O(1/t²) rate, versus O(1/t) for gradient descent. Deep learning frameworks implement the Sutskever et al. rearrangement, which keeps the parameters at the lookahead point implicitly and applies the update g_t + μ·v_t instead of v_t; it is equivalent up to a reparameterization of what "the parameters" means and is what the nesterov=True flag toggles below. In stochastic, non-convex practice the gap between Nesterov and heavy ball is small; it is worth a try, not a reorganization of your life.

The momentum and learning-rate interaction

Because the buffer steady-states at g/(1 − μ), the quantity a run actually experiences is the effective learning rate η_eff = η/(1 − μ). Move μ from 0.9 to 0.99 at fixed η and you have silently multiplied the effective step by ten; most "momentum 0.99 diverged" reports are really "effective learning rate 10× diverged". When sweeping momentum, hold η/(1 − μ) constant so you are comparing the averaging horizon, not the step size. There is a second, sneakier interaction with schedules: PyTorch multiplies η into the update at the end (the buffer stores gradients), while the Sutskever formulation folds η into the buffer itself (v = μv − η·g). At constant η they coincide; under a decaying schedule they do not, because the PyTorch form rescales the entire accumulated history by today's η while the folded form lets old steps keep the η they were taken with. The difference is usually small but it is real, and it is one reason optimizer trajectories from different frameworks drift apart under aggressive schedules.

Implementation, twice

The PyTorch version mirrors torch.optim.SGD's documented algorithm exactly, including the two easy-to-miss details: the first momentum step initializes the buffer to the raw gradient (not (1 − dampening)·g), and weight decay is added to the gradient before the buffer update, so decay rides the momentum too. The JAX version is a pure init/update pair over pytrees in the optax style; note that optax's trace transform has no dampening knob, its buffer is the plain accumulator, so the two libraries agree exactly when dampening is zero.

import torch

class SGD:
    """From-scratch mirror of torch.optim.SGD.

    Per parameter, per step:
        g = grad + weight_decay * p
        if momentum != 0:
            buf = g                                (first step)
            buf = momentum * buf + (1-dampening)*g (afterwards)
            g = g + momentum * buf   if nesterov else buf
        p -= lr * g
    """

    def __init__(self, params, lr, momentum=0.0, dampening=0.0,
                 weight_decay=0.0, nesterov=False):
        if nesterov and (momentum <= 0 or dampening != 0):
            raise ValueError("nesterov requires momentum > 0 and dampening = 0")
        self.param_groups = [dict(params=list(params), lr=lr,
                                  momentum=momentum, dampening=dampening,
                                  weight_decay=weight_decay,
                                  nesterov=nesterov)]
        self.state = {}  # param -> momentum buffer

    @torch.no_grad()
    def step(self):
        for group in self.param_groups:
            mu, tau = group["momentum"], group["dampening"]
            for p in group["params"]:
                if p.grad is None:
                    continue
                g = p.grad
                if group["weight_decay"] != 0:
                    # decay enters BEFORE the buffer: it gets momentum too
                    g = g.add(p, alpha=group["weight_decay"])
                if mu != 0:
                    buf = self.state.get(p)
                    if buf is None:
                        # torch initializes buf = g, not (1-tau)*g
                        buf = self.state[p] = g.clone().detach()
                    else:
                        buf.mul_(mu).add_(g, alpha=1.0 - tau)
                    g = g.add(buf, alpha=mu) if group["nesterov"] else buf
                p.add_(g, alpha=-group["lr"])

    def zero_grad(self):
        for group in self.param_groups:
            for p in group["params"]:
                p.grad = None

from typing import Any, NamedTuple

import jax
import jax.numpy as jnp


class SGDState(NamedTuple):
    trace: Any  # pytree of momentum buffers, mirrors params


def sgd(lr, momentum=0.0, nesterov=False, weight_decay=0.0):
    """Pure init/update pair in the optax GradientTransformation style.

    Matches optax.sgd (= torch.optim.SGD with dampening=0):
    the trace is an accumulator, t <- g + momentum * t.
    Zero-init makes the first step's trace equal g, same as torch.
    """

    def init(params):
        return SGDState(trace=jax.tree_util.tree_map(jnp.zeros_like, params))

    def update(grads, state, params=None):
        if weight_decay:
            grads = jax.tree_util.tree_map(
                lambda g, p: g + weight_decay * p, grads, params)
        new_trace = jax.tree_util.tree_map(
            lambda g, t: g + momentum * t, grads, state.trace)
        if nesterov:
            direction = jax.tree_util.tree_map(
                lambda g, t: g + momentum * t, grads, new_trace)
        else:
            direction = new_trace
        updates = jax.tree_util.tree_map(lambda d: -lr * d, direction)
        return updates, SGDState(trace=new_trace)

    return init, update


def apply_updates(params, updates):
    return jax.tree_util.tree_map(lambda p, u: p + u, params, updates)

The verification harness runs both against the reference optimizers on the Rosenbrock function, a classic curved valley that punishes any disagreement in the update rule within a few dozen steps. The whole parameter trajectory is compared, not just the endpoint, so a first-step initialization bug or a misplaced dampening factor cannot hide.

import torch

def rosenbrock(p):
    x, y = p
    return (1 - x) ** 2 + 100 * (y - x ** 2) ** 2

def trajectory(opt_cls, steps=300, **kw):
    p = torch.tensor([-1.5, 2.0], requires_grad=True)
    opt = opt_cls([p], **kw)
    out = []
    for _ in range(steps):
        opt.zero_grad()
        rosenbrock(p).backward()
        opt.step()
        out.append(p.detach().clone())
    return torch.stack(out)

configs = [
    dict(lr=1e-3),
    dict(lr=1e-3, momentum=0.9),
    dict(lr=1e-3, momentum=0.9, nesterov=True),
    dict(lr=1e-3, momentum=0.9, dampening=0.5),
    dict(lr=1e-3, momentum=0.9, weight_decay=1e-2),
]
for kw in configs:
    ours = trajectory(SGD, **kw)
    ref = trajectory(torch.optim.SGD, **kw)
    assert torch.allclose(ours, ref, atol=1e-7), kw
print("all 5 trajectories match torch.optim.SGD exactly")

import jax
import jax.numpy as jnp
import optax

def rosenbrock(p):
    x, y = p
    return (1 - x) ** 2 + 100 * (y - x ** 2) ** 2

def trajectory(init_fn, update_fn, steps=300):
    params = jnp.array([-1.5, 2.0])
    state = init_fn(params)

    def step(carry, _):
        params, state = carry
        grads = jax.grad(rosenbrock)(params)
        updates, state = update_fn(grads, state, params)
        params = optax.apply_updates(params, updates)
        return (params, state), params

    _, traj = jax.lax.scan(step, (params, state), None, length=steps)
    return traj

for kw in [dict(momentum=0.0), dict(momentum=0.9),
           dict(momentum=0.9, nesterov=True)]:
    init, update = sgd(lr=1e-3, **kw)
    ref = optax.sgd(1e-3, momentum=kw.get("momentum") or None,
                    nesterov=kw.get("nesterov", False))
    ours = trajectory(init, update)
    theirs = trajectory(ref.init, ref.update)
    assert jnp.allclose(ours, theirs, atol=1e-7), kw
print("all trajectories match optax.sgd exactly")

Using it on a real shape of problem

The two-parameter valley proves correctness; a small classifier shows the optimizer in its natural habitat. The loop below trains a 784→256→10 MLP, the standard MNIST shape covered in more depth on the neural network page, on synthetic data with the same dimensions. With lr = 0.1 and momentum = 0.9, expect cross-entropy to fall from about ln 10 ≈ 2.30 to well under 0.1 within a few hundred steps on separable synthetic clusters; exact curves are machine and seed dependent. Rerun with momentum = 0 at the same lr and the early descent is visibly slower; rerun with momentum = 0.99 at the same lr and it will oscillate or diverge, which is the effective-learning-rate interaction from the math section showing up on schedule.

import torch

torch.manual_seed(0)
X = torch.randn(4096, 784)
y = (X[:, :10].argmax(dim=1))          # synthetic 10-class labels

model = torch.nn.Sequential(
    torch.nn.Linear(784, 256), torch.nn.ReLU(),
    torch.nn.Linear(256, 10),
)
opt = SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=1e-4)

for step in range(500):
    idx = torch.randint(0, 4096, (128,))
    loss = torch.nn.functional.cross_entropy(model(X[idx]), y[idx])
    opt.zero_grad()
    loss.backward()
    opt.step()
    if step % 100 == 0:
        print(step, round(loss.item(), 3))
# expect ~2.3 at step 0, falling below ~0.3 by step 500 (machine-dependent)

import jax
import jax.numpy as jnp
import optax

key = jax.random.PRNGKey(0)
kx, k1, k2 = jax.random.split(key, 3)
X = jax.random.normal(kx, (4096, 784))
y = X[:, :10].argmax(axis=1)           # synthetic 10-class labels

params = {
    "w1": jax.random.normal(k1, (784, 256)) * (2.0 / 784) ** 0.5,
    "b1": jnp.zeros(256),
    "w2": jax.random.normal(k2, (256, 10)) * (2.0 / 256) ** 0.5,
    "b2": jnp.zeros(10),
}

def loss_fn(params, xb, yb):
    h = jax.nn.relu(xb @ params["w1"] + params["b1"])
    logits = h @ params["w2"] + params["b2"]
    return optax.softmax_cross_entropy_with_integer_labels(logits, yb).mean()

init, update = sgd(lr=0.1, momentum=0.9, weight_decay=1e-4)
state = init(params)

@jax.jit
def train_step(params, state, xb, yb):
    loss, grads = jax.value_and_grad(loss_fn)(params, xb, yb)
    updates, state = update(grads, state, params)
    return optax.apply_updates(params, updates), state, loss

for step in range(500):
    key, sub = jax.random.split(key)
    idx = jax.random.randint(sub, (128,), 0, 4096)
    params, state, loss = train_step(params, state, X[idx], y[idx])
    if step % 100 == 0:
        print(step, round(float(loss), 3))
# expect loss ~2-3 at step 0, well below 0.3 by step 500 (machine-dependent)

Applications

SGD with momentum is still the reference recipe for convolutional networks. The original ResNet training used SGD with momentum 0.9, weight decay 10⁻⁴, and lr 0.1 divided by 10 on a step schedule, and the modern torchvision and timm ImageNet recipes are refinements of the same core (cosine decay, label smoothing, longer training) rather than replacements of the optimizer. Detection and segmentation stacks built on those backbones, Detectron2 and YOLO-family trainers among them, default to SGD with momentum for the same reason. The generalization folklore deserves a careful statement: Wilson et al. (2017) reported that adaptive methods can generalize worse than tuned SGD on vision benchmarks, and that observation shaped practice for years, but later work showed much of the gap closes when Adam's hyperparameters, especially ε and weight decay, are tuned as aggressively as SGD's. The honest summary is that on convnets, SGD with momentum remains at least as good as anything else while being simpler, whereas on transformers Adam-family optimizers win clearly, plausibly because attention and layer-norm gradients are badly scaled across parameters in a way per-coordinate adaptation fixes and a single global η cannot. Beyond vision, SGD variants train large-scale recommendation models, and SGD itself is the backbone of federated learning (local SGD steps averaged across devices) and of most convergence theory: when a paper proves something about neural network training, the object of proof is almost always SGD.

Against the real libraries

The production implementations to compare against are torch.optim.SGD and optax.sgd. Mathematically they compute exactly what the from-scratch versions above compute, which is the point of the trajectory test. What they add is engineering. PyTorch ships three execution paths for the same update: a per-parameter loop, a foreach path that batches the arithmetic across all parameters into a few large multi-tensor kernels, and a fused path that does the whole update in one kernel per dtype and device; on a model with hundreds of parameter tensors the difference is a meaningful slice of step time. It also handles the bookkeeping this page ignores: parameter groups with per-group hyperparameters, state_dict save and restore for resumable training, sparse gradients, a maximize flag, a capturable mode for CUDA graphs, and a differentiable mode that lets you backprop through the optimizer itself for meta-learning. Optax's added value is compositionality: optax.sgd is literally chain(trace(momentum), scale(-lr)), and the same chaining gives you gradient clipping, schedules, and weight decay as independent, testable transforms, a design the schedules page leans on heavily. The from-scratch version is genuinely enough for research code with a single parameter group and no checkpointing, and writing it is the fastest way to stop treating optimizers as black boxes. The verification recipe is the one shown above and it generalizes: fixed seed, identical initial parameters, run both optimizers for a few hundred steps, and assert the full trajectories agree to float32 tolerance across every flag combination you claim to support.

Traps and misconceptions

"Momentum is just a bigger learning rate." No. The steady-state amplification 1/(1 − μ) applies only along directions where gradients persist; along oscillating directions momentum shrinks the step. A bigger η scales both. That asymmetry is why momentum helps on ill-conditioned losses where a bigger η alone diverges.

Sweeping momentum without rescaling η. Because the realized step scales like η/(1 − μ), comparing μ = 0.9 against μ = 0.99 at fixed η is really comparing two effective learning rates an order of magnitude apart. Hold η/(1 − μ) constant when you change μ, or the sweep measures nothing.

Assuming all frameworks compute the same momentum. PyTorch keeps gradients in the buffer and multiplies η at the end; the Sutskever formulation folds η into the buffer. Identical at constant η, different under a schedule, since PyTorch rescales the whole accumulated history by the current η. Also, PyTorch's first momentum step sets buf = g while some texts write buf = (1 − μ)·g; and its dampening is a separate knob that optax does not have at all. Trajectory-level tests catch every one of these.

Weight decay here is not the decoupled kind. In torch.optim.SGD (and the mirror above), decay is added to the gradient before the momentum update, so the decay term is itself momentum-smoothed and scaled by any schedule. That is L2 regularization, not the decoupled weight decay of AdamW; the distinction matters more for Adam and is treated on the Adam page, but do not assume weight_decay= means the same thing across optimizers.

Expecting Nesterov to look like the textbook. Frameworks do not evaluate the gradient at a lookahead point; they use the Sutskever rearrangement g + μ·v, which is the same method expressed at shifted coordinates. If you diff a framework's update against the textbook two-point form and see disagreement, that is the reparameterization, not a bug, and PyTorch additionally requires dampening = 0 when nesterov is on.