Nesterov momentum: looking ahead

The Problem with Standard Momentum

Standard momentum works well, but it has a subtle flaw: it computes the gradient at the wrong place.

Nesterov momentum is the version used in many high-performance training setups, including the SGD configurations that train ImageNet models. The improvement in convergence speed is consistent enough that it became the standard choice whenever momentum SGD is used.

The update sequence in standard momentum is:

1. Evaluate gradient at current position $\theta_t$
2. Update velocity: $v_t = \beta v_{t-1} + \alpha \nabla L(\theta_t)$
3. Step: $\theta_{t+1} = \theta_t - v_t$

The problem: the velocity $\beta v_{t-1}$ was already "committed" before you evaluated the gradient. You know you're going to move by approximately $\beta v_{t-1}$ regardless — that's the momentum. So why evaluate the gradient at $\theta_t$, which is where you currently are, rather than at the position you're about to be?

The Nesterov Idea: Look Before You Step

The Nesterov Accelerated Gradient ⓘ ( NAG ⓘ ) swaps the order:

1. First, take the momentum step to the look-ahead position:

   $$\theta_{\text{look}} = \theta_t - \beta \cdot v_{t-1}$$

   $θ_look$

   look-ahead position — where pure momentum would take you

   $β$

   momentum coefficient

   $vₜ₋₁$

   previous velocity

2. Evaluate the gradient at the look-ahead position: $\nabla L(\theta_{\text{look}})$

3. Update velocity and parameters:

   $$v_t = \beta \cdot v_{t-1} + \alpha \cdot \nabla L(\theta_{\text{look}})$$

   $vₜ$

   updated velocity

   $α$

   learning rate

   $∇L(θ_look)$

   gradient at the look-ahead position

   $$\theta_{t+1} = \theta_t - v_t$$

   $θₜ₊₁$

   updated parameters

The gradient is computed after the "free" momentum step, giving a more accurate signal about where the objective is going.

The Physical Analogy

Standard momentum: running downhill. You look at the slope where you are standing, then leap in that direction with accumulated velocity. You may overshoot the valley floor before realizing you should have braked.

Nesterov: running downhill. You project yourself forward to where your momentum will carry you, look at the slope there, then apply the corrected step. You see the upslope coming before you reach it, and you brake earlier.

The difference is subtle but consistently beneficial.

Numerical Comparison

Consider minimizing $L(\theta) = 5\theta^2$ starting at $\theta_0 = 1.0$. True minimum at $\theta = 0$. Use $\alpha = 0.1$, $\beta = 0.9$, $v_0 = 0$.

Standard momentum:

• Value: $\nabla L(\theta_0) = 10\theta_0 = 10$
• Value: $v_1 = 0.9 \cdot 0 + 0.1 \cdot 10 = 1.0$
• Value: $\theta_1 = 1.0 - 1.0 = 0.0$ (happened to land on minimum in one step)
• Value: $\nabla L(\theta_1) = 0$
• Value: $v_2 = 0.9 \cdot 1.0 + 0 = 0.9$
• Value: $\theta_2 = 0 - 0.9 = -0.9$ (overshot past the minimum!)

Nesterov momentum:

• Value: $\theta_{\text{look}} = 1.0 - 0.9 \cdot 0 = 1.0$ (same as current at start)
• Value: $v_1 = 0.9 \cdot 0 + 0.1 \cdot 10 = 1.0$; $\theta_1 = 0.0$
• Value: $\theta_{\text{look}} = 0.0 - 0.9 \cdot 1.0 = -0.9$
• Value: $\nabla L(-0.9) = 10 \cdot (-0.9) = -9$ (gradient points back toward 0!)
• Value: $v_2 = 0.9 \cdot 1.0 + 0.1 \cdot (-9) = 0.9 - 0.9 = 0.0$
• Value: $\theta_2 = 0.0 - 0.0 = 0.0$ (stays at minimum!)

Nesterov recognized it was about to overshoot (the look-ahead gradient was negative) and applied a braking correction. Standard momentum overshot to -0.9.

The Equivalent Reformulation

In practice, computing $\nabla L(\theta_{\text{look}})$ requires storing the look-ahead parameters. An equivalent formulation that avoids this uses a change of variables $\phi = \theta - \beta v$:

In PyTorch, this is handled automatically:

optimizer = torch.optim.SGD(
    model.parameters(),
    lr=0.01,
    momentum=0.9,
    nesterov=True    # Enable Nesterov look-ahead
)

One boolean change. In practice, Nesterov almost always performs equal to or better than standard momentum — it is the preferred default when using SGD with momentum.