← All tutorials

SGD, Momentum and Nesterov

How memory speeds up and calms the step.

On this page

Each optimizer adds one idea to the one before; each arrow names it.

An optimizer turns each step's gradient into a weight update. It runs once per training step on one minibatch, a small random batch of training examples, and keeps a little state, such as momentum's velocity, between steps.

1 Gradient descent and SGD

Every figure in this series uses one loss, L = 5x² + 0.2y², started at (−3.5, 5). I call its shape, a long narrow valley, the ravine; it stands for the ill-conditioned loss of a real network, where one direction is far steeper than another. Its numbers are illustrative, computed by script. Its first step at learning rate α 0.15:

Step downhill by an amount proportional to the slope (Gradient Descent, section 1.1).

Steep places get big steps, flat places tiny ones. SGD computes ∇L on a minibatch instead of the whole data set, so each step is cheaper and a little noisy. Below, SGD fits the mean squared error of a linear model (Linear Regression, section 2).

import numpy as np
rng = np.random.default_rng(0)
X = rng.normal(size=(256, 3))
y = X @ rng.normal(size=3)
w = np.zeros(3)
lr, bs = 0.1, 32  # α, batch
for i in range(0, 256, bs):
    xb = X[i:i + bs]
    yb = y[i:i + bs]
    err = xb @ w - yb
    grad = 2 * xb.T @ err / bs
    w = w - lr * grad

1. 1Data, zero weights, α and the batch size.
2. 2Takes the next minibatch of rows.
3. 3Slope of the batch's mean squared error.
4. 4Steps every weight by the same α.

1.1 The steep direction sets the learning rate

Along one axis of curvature h (how fast the slope grows), each step multiplies the distance to the minimum by one number.

Even at α 0.19, just under the limit, y needs 59 steps to reach 1% of its start. One α cannot suit both axes; momentum below and per-axis step sizes (Adaptive Optimizers) are the two fixes.

2 Momentum

Momentum keeps a velocity v, a running sum of past gradients in which each older gradient is shrunk by the momentum factor μ every step.

Sutskever et al. 2013 train deep networks with momentum 0.9 and higher, so the velocity can reach 10 times one gradient. On the ravine at α 0.05, μ 0.9:

Directions that agree add up; directions that flip cancel. The same mechanism speeds y and makes x overshoot.

import numpy as np

theta = np.array([-3.5, 5.0])
h = np.array([10.0, 0.4])    # curvature of x and y
v = np.zeros(2)
lr, mu = 0.05, 0.9           # α, μ
for t in range(60):
    grad = h * theta
    v = mu * v + grad
    theta = theta - lr * v

1. 1The ravine, its start, zero velocity, α and μ.
2. 2Reads today's slope.
3. 3Shrinks v by μ and adds the slope.
4. 4Steps by α times the velocity.

2.1 Momentum keeps moving after the slope vanishes

The last pushed step is 0.878. Once the slope is gone, each step is μ times the one before.

In multitask training a gradient balancer edits each task's gradient and can remove a direction; the same glide carries it back into the step (Gradient Balancing Inside the Training Loop, section 8.1).

3 Nesterov momentum

Nesterov reads the slope where the velocity is about to carry the weights, not where they are now.

The brake starts one step earlier, because Nesterov reads a slope ten times gentler at the lookahead than momentum reads at θ. At α 0.15 and μ 0.9 Nesterov diverges on this ravine while momentum does not, because its earlier reading of the steep slope makes it more sensitive to a large α.

4 Saddle points and minibatch noise

Exact gradient descent started on the x axis stops at the centre. A minibatch gradient is never exactly zero, so its noise nudges the weights along y and they escape. The DQN agent learns from batches of 64 earlier moves drawn at random from its memory, so every gradient it follows is this kind of noisy minibatch gradient (Deep Q-Networks, section 6).

5 SGD and momentum in PyTorch

Both run once per step; momentum keeps one buffer the size of the weights.

opt = torch.optim.SGD(model.parameters(), lr=0.05, momentum=0.9)
opt_nag = torch.optim.SGD(model.parameters(), lr=0.05, momentum=0.9, nesterov=True)

1. 1The momentum buffer is the velocity v.
2. 2The same with Nesterov momentum.

6 Failure modes

7 References

1. Robbins, H., and Monro, S. (1951). A stochastic approximation method. Annals of Mathematical Statistics.
2. Polyak, B. T. (1964). Some methods of speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics.
3. Nesterov, Y. (1983). A method for solving the convex programming problem with convergence rate O(1/k²). Doklady Akademii Nauk SSSR 269(3), 543-547; English translation in Soviet Mathematics Doklady 27.
4. Sutskever, I., et al. (2013). On the importance of initialization and momentum in deep learning. ICML.
5. PyTorch documentation, torch.optim (SGD, momentum and nesterov settings).