SGD with momentum

Deep dive

How it works

In: parameters w, minibatch gradient g, learning rate η, momentum μ · Out: updated w · State: one velocity buffer v per parameter (4 B in fp32) · Cost: O(P) per step, memory-bound at about 20 B of traffic per parameter

PyTorch's formulation, with g the minibatch gradient, λ the weight decay and η the learning rate:

g ← g + λ·w                      (coupled L2 weight decay)
v ← g             on step 1
v ← μ·v + (1 − dampening)·g       afterwards
w ← w − η·v                       classic (heavy ball)
w ← w − η·(g + μ·v)               Nesterov: gradient "looks ahead" along v

Unrolled, w_{t+1} = w_t − η g_t + μ (w_t − w_{t−1}). This is Polyak's heavy ball. Sutskever's formulation (v ← μv − ηg; w ← w + v) is identical for constant η. It differs when η changes: in the PyTorch form a learning-rate drop takes effect immediately on the whole velocity, in Sutskever's form it acts only on new gradient contributions.

For a quadratic with eigenvalues in [m, L], the optimal heavy-ball parameters are η = 4/(√L+√m)² and μ = ((√L−√m)/(√L+√m))². The error then contracts by (√κ−1)/(√κ+1) per step, against (κ−1)/(κ+1) for gradient descent. For minibatch noise, momentum acts as an exponential moving average over about 1/(1−μ) = 10 steps.

Code

pub struct SgdMomentum {
    pub lr: f32, pub momentum: f32, pub weight_decay: f32, pub nesterov: bool,
    vel: Vec<f32>, started: bool,
}

impl SgdMomentum {
    pub fn new(n: usize, lr: f32, momentum: f32, weight_decay: f32, nesterov: bool) -> Self {
        Self { lr, momentum, weight_decay, nesterov, vel: vec![0.0; n], started: false }
    }

    /// One fused pass: read w, g, v; write w, v. The compiler auto-vectorizes this loop.
    pub fn step(&mut self, w: &mut [f32], g: &[f32]) {
        assert!(w.len() == g.len() && g.len() == self.vel.len());
        let (lr, mu, wd, nest) = (self.lr, self.momentum, self.weight_decay, self.nesterov);
        let first = !std::mem::replace(&mut self.started, true);
        for ((w, &g), v) in w.iter_mut().zip(g).zip(self.vel.iter_mut()) {
            let g = g + wd * *w;
            *v = if first { g } else { mu * *v + g }; // PyTorch seeds v with the first gradient
            *w -= lr * if nest { g + mu * *v } else { *v };
        }
    }
}

/// LARS per-layer learning rate (You et al. 2017): lr · η · ‖w‖ / ‖g + λw‖.
pub fn lars_lr(lr: f32, trust: f32, wd: f32, w: &[f32], g: &[f32]) -> f32 {
    let wn = w.iter().map(|x| x * x).sum::<f32>().sqrt();
    let gn = w.iter().zip(g).map(|(w, g)| (g + wd * w).powi(2)).sum::<f32>().sqrt();
    if wn > 0.0 && gn > 0.0 { lr * trust * wn / gn } else { lr }
}

/// Steps until f(w) = ½ Σ a_i w_i² falls below tol, eigenvalues a in [1, 100] (κ = 100).
fn steps_to_converge(lr: f32, mu: f32) -> usize {
    let a: Vec<f32> = (0..64).map(|i| 1.0 + 99.0 * i as f32 / 63.0).collect();
    let mut w = vec![1.0f32; 64];
    let mut opt = SgdMomentum::new(64, lr, mu, 0.0, false);
    for t in 0..100_000 {
        let f: f32 = a.iter().zip(&w).map(|(a, w)| 0.5 * a * w * w).sum();
        if f < 1e-8 { return t; }
        let g: Vec<f32> = a.iter().zip(&w).map(|(a, w)| a * w).collect();
        opt.step(&mut w, &g);
    }
    usize::MAX
}

fn main() {
    let (l, m) = (100.0f32, 1.0f32);
    let gd = steps_to_converge(1.0 / l, 0.0);
    let hb = steps_to_converge(4.0 / (l.sqrt() + m.sqrt()).powi(2),
                               ((l.sqrt() - m.sqrt()) / (l.sqrt() + m.sqrt())).powi(2));
    println!("gradient descent: {gd} steps, heavy ball: {hb} steps"); // about 900 vs about 80
    assert!(hb * 5 < gd);
}

Complexity and constants

• O(P) per step. In fp32 the update reads w, g, v and writes w, v: 20 B/param. It is purely memory-bound: 1B params ≈ 20 GB of traffic ≈ 7 ms at 3 TB/s HBM. Fuse the update into one kernel (PyTorch foreach/fused), or the step costs 3-5x that.
• Noise averaging window ≈ 1/(1−μ) steps: 10 at 0.9, 100 at 0.99.
• Linear scaling rule (Goyal et al.): batch 256 → lr 0.1; batch 8192 → lr 3.2, with 5 epochs of linear warmup. ResNet-50 reached 76.3% top-1 in 1 hour on 256 GPUs. Beyond about 8k, plain SGD loses accuracy; LARS reached 32k batch.

Efficiency

• Cost: in fp32 the update reads w, g, v and writes w, v: 20 B per parameter. 1B parameters is 20 GB of traffic, about 7 ms at 3 TB/s HBM (arithmetic from the page). Optimizer state is 4 B per parameter, half of Adam's. Bound: memory bandwidth, not FLOPs.
• Levers:
  • fused or foreach update kernel — one pass instead of several, step time order of 3-5x lower — any GPU training
  • large batch with linear lr scaling and 5 epochs of warmup — ResNet-50 to 76.3% top-1 in 1 hour on 256 GPUs (Goyal et al. 2017) — data-parallel vision training
  • LARS per-layer trust ratio — batch up to 32k without accuracy loss (You et al. 2017) — convnets beyond batch 8k
  • weights in bf16 with v kept in fp32 — half the weight and gradient traffic — mixed-precision runs
  • no weight decay on biases and norm scales — better accuracy, no extra cost (two parameter groups) — all recipes
• Naive → standard → state of the art:
  • per-tensor, per-op update with separate kernels · 3-5x the fused step (order of)
  • fused SGD with momentum 0.9, Nesterov optional · about 7 ms per 1B parameters
  • large-batch SGD with LARS, or Lion and Muon for other model families · one state buffer, same bandwidth bound

Implementation notes

• Exclude biases and BatchNorm/LayerNorm scales from weight decay. It is cheap to do with two parameter groups and consistently helps.
• Zero-initialize the last BatchNorm γ in each residual block (Goyal et al.). Every block then starts as identity, which stabilizes large-lr starts.
• Keep v in fp32 even when weights are bf16 (see mixed precision).
• With data parallelism, average gradients (ring all-reduce) before the update. All replicas then hold identical v without communicating it.
• Gradient accumulation over k micro-batches: divide the loss by k, or the effective lr is k times too large.

Pitfalls

• Loading weights but not the momentum buffers on resume gives a loss spike and effectively a warm restart.
• Increasing μ without lowering lr: the effective step lr/(1−μ) jumps 10x going from 0.9 to 0.99.
• Weight decay in SGD is L2 (coupled). Porting λ to AdamW, or the reverse, does not carry over.
• Large batch without warmup diverges in the first few hundred steps.
• dampening ≠ 0 combined with Nesterov is rejected by PyTorch, and it silently changes the effective lr in other code.

Variants and alternatives

In production

torch.optim.SGD (foreach, fused), the torchvision ResNet recipe (lr 0.1, μ 0.9, wd 1e-4, 90 epochs, step decay), Detectron2 and MMDetection defaults, Optax sgd, NVIDIA's LARS in large-batch MLPerf ResNet submissions.

Implementations

Sources

• Polyak, Some methods of speeding up the convergence of iteration methods, USSR Computational Mathematics and Mathematical Physics, 1964. paper
• Nesterov, A method for solving the convex programming problem with convergence rate O(1/k²), 1983.
• Sutskever, Martens, Dahl, Hinton, On the importance of initialization and momentum in deep learning, ICML 2013. paper
• Goyal et al., Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour, 2017. paper
• You, Gitman, Ginsburg, Large Batch Training of Convolutional Networks (LARS), 2017. paper