Stochastic Gradient Descent › SGD with Nesterov Momentum

The chapter on Momentum will cover momentum in depth, but Nesterov momentum is closely enough tied to SGD that we introduce it here. Standard SGD with momentum maintains a velocity vector $\mathbf{v}$ that accumulates past gradients:

$$\begin{aligned} \mathbf{v}_{t+1} &= \mu \mathbf{v}_t - \eta \nabla_{\boldsymbol{\theta}} \ell(\boldsymbol{\theta}_t) \\ \boldsymbol{\theta}_{t+1} &= \boldsymbol{\theta}_t + \mathbf{v}_{t+1} \end{aligned}$$

where:

• $\mathbf{v}_t$: the velocity vector at step $t$ (initialized to zero)
• $\mu \in [0, 1)$: the momentum coefficient, controlling how much past gradients influence the current step (typically 0.9)
• $\eta$: the learning rate
• $\nabla_{\boldsymbol{\theta}} \ell(\boldsymbol{\theta}_t)$: the gradient evaluated at the current position $\boldsymbol{\theta}_t$

Momentum accumulates gradient directions over time, which helps SGD move faster in directions with consistent gradients and slows updates in directions where the gradient oscillates.

Nesterov momentum makes one key improvement: instead of computing the gradient at the current position $\boldsymbol{\theta}_t$, it computes the gradient at a "lookahead" position $\boldsymbol{\theta}_t + \mu \mathbf{v}_t$:

$$\begin{aligned} \mathbf{v}_{t+1} &= \mu \mathbf{v}_t - \eta \nabla_{\boldsymbol{\theta}} \ell(\boldsymbol{\theta}_t + \mu \mathbf{v}_t) \\ \boldsymbol{\theta}_{t+1} &= \boldsymbol{\theta}_t + \mathbf{v}_{t+1} \end{aligned}$$

where:

• $\boldsymbol{\theta}_t + \mu \mathbf{v}_t$: the "lookahead" position, approximately where we will be after applying the accumulated momentum
• $\nabla_{\boldsymbol{\theta}} \ell(\boldsymbol{\theta}_t + \mu \mathbf{v}_t)$: the gradient evaluated at the lookahead position, not at $\boldsymbol{\theta}_t$

Intuitively, Nesterov momentum first takes a step in the direction of the accumulated velocity, then computes the gradient at the resulting position, and corrects the step based on the resulting position. This "correct then move" approach is more responsive to changes in the gradient landscape than standard momentum's "move then correct" approach.

In practice, Nesterov momentum tends to converge faster and with better stability than standard momentum on convex problems. For non-convex neural networks, the empirical differences are smaller but Nesterov is still generally preferred. PyTorch's SGD implementation supports Nesterov momentum via the nesterov=True parameter.