Momentum and Nesterov Accelerated Gradient

Classical (heavy-ball) momentum augments gradient descent with an auxiliary state vector $v_t \in \mathbb{R}^d$, called the velocity, that has the same shape as the parameters $\theta_t$. Instead of moving directly against the current gradient, the algorithm moves against a running, exponentially-decaying average of all past gradients:

Here $\eta > 0$ is the learning rate and $\mu \in [0,1)$ is the momentum coefficient, controlling how much of the previous velocity survives each step. Setting $\mu = 0$ recovers plain SGD exactly. Unrolling the recursion shows $v_t$ is a geometric (exponentially-weighted) sum of past gradients,

so recent gradients dominate but old ones never fully vanish — an effective "memory horizon" of about $1/(1-\mu)$ steps. For $\mu = 0.9$ this horizon is about 10 steps; for $\mu=0.99$, about 100 steps.

Nesterov Accelerated Gradient (NAG), due to Yurii Nesterov, modifies where the gradient is evaluated. Classical momentum computes the gradient at the current point $\theta_{t-1}$ and only afterward applies the momentum-augmented step. NAG instead reasons: "I already know that momentum will carry me approximately to $\theta_{t-1} + \mu v_{t-1}$ regardless of the new gradient, so I should measure the gradient there, at the lookahead point, not at $\theta_{t-1}$." Formally,

This is a correction term, not a different algorithm family: it is momentum with a smarter, anticipatory gradient evaluation point. Geometrically, NAG lets the optimizer "peek ahead," measure the slope at the point momentum is already carrying it toward, and brake early if that lookahead point has started climbing back uphill — precisely the behavior that suppresses overshoot near a minimum.

Both algorithms fit inside the same loop structure as mini-batch SGD; only the update rule inside the loop changes. Below is unified pseudocode covering plain SGD, classical momentum, and Nesterov momentum, followed by the algebraic derivation of NAG's commonly implemented equivalent form. ``` ALGORITHM Momentum-SGD(θ₀, η, μ, nesterov, num_steps) Input: θ₀ initial parameters (vector of size d) η learning rate μ momentum coefficient in [0, 1) nesterov boolean flag num_steps number of optimization steps Output: θ (trained parameters) v ← 0 # velocity, same shape as θ, O(d) memory θ ← θ₀ for t = 1 to num_steps: batch ← sample_minibatch() if nesterov: θ_lookahead ← θ + μ v # peek ahead using old velocity g ← ∇L(θ_lookahead; batch) # gradient AT the lookahead point else: g ← ∇L(θ; batch) # gradient AT the current point v ← μ v - η g # update velocity (EMA of gradients) θ ← θ + v # update parameters return θ ``` This form is mathematically exact but requires computing the gradient at a shifted point $\theta + \mu v$, which means an extra parameter copy/shift inside the forward-backward pass. Most deep learning frameworks instead use an algebraically equivalent form that evaluates the gradient at the current parameters and folds the lookahead correction into the velocity update itself — avoiding any extra forward pass: ``` ALGORITHM Momentum-SGD-Framework-Form(θ₀, η, μ, nesterov, num_steps) v ← 0 θ ← θ₀ for t = 1 to num_steps: batch ← sample_minibatch() g ← ∇L(θ; batch) # always at current θ — one gradient call v_prev ← v v ← μ v_prev - η g if nesterov: θ ← θ + μ v - η g # = θ + (v + μ (v - v_prev)) else: θ ← θ + v return θ ```

Why these two forms coincide. Define the substitution $\tilde\theta_t = \theta_t + \mu v_t$ (the point NAG would have evaluated the gradient at, one step later). Substituting into the lookahead recursion and simplifying the telescoping terms yields exactly the update $\theta_t = \theta_{t-1} + \mu v_t - \eta \nabla L(\theta_{t-1})$ used above, with $g$ evaluated at the un-shifted, current $\theta$. This re-parameterized form is what PyTorch's `SGD(..., nesterov=True)` and similar library implementations use, because it needs only one gradient evaluation per step, identical to classical momentum. The following diagram shows one training step's data flow for the framework form, making explicit which quantities are read and written and in what order.

Time complexity. Each optimizer step performs $O(d)$ arithmetic (a scalar-vector multiply and vector add for $v$, another for $\theta$), where $d$ is the number of parameters. This is asymptotically dwarfed by the cost of the forward and backward pass computing $\nabla L$, which for a network with $d$ parameters and batch size $b$ is typically $\Theta(bd)$ or worse. Momentum therefore adds no asymptotic overhead to a training step; it is a constant-factor addition, and in the framework form it uses exactly one gradient evaluation per step, identical to plain SGD.

Space complexity. Momentum requires storing one extra vector $v \in \mathbb{R}^d$, i.e., $O(d)$ additional memory — doubling the "optimizer state" memory footprint relative to plain SGD (which needs none) but still far less than storing full activation histories for backpropagation in most architectures. This is the same order of extra memory later paid by Adam and Its Variants, which stores two auxiliary vectors instead of one.

Correctness / why it accelerates — the exponentially weighted moving average view. Unrolling the momentum recursion,

shows $v_t$ is (up to the $-\eta$ scale) an exponentially weighted moving average of the gradient history with decay $\mu$. Consider a coordinate direction where consecutive gradients alternate in sign (oscillation, typical of high curvature): the alternating terms partially cancel in the sum, shrinking the effective step in that direction. Consider a coordinate where gradients repeatedly point the same way (low curvature, a long shallow ravine): the terms reinforce, and the effective step size approaches $\eta/(1-\mu)$ instead of $\eta$ — a genuine speed-up of up to $10\times$ at $\mu=0.9$. This single argument explains both the oscillation-damping and the acceleration properties simultaneously, without any convexity assumption.

Correctness / why Nesterov's lookahead helps — a first-order Taylor argument. Near a minimum, the momentum term $\mu v_{t-1}$ is often already an accurate rough estimate of the coming step. Classical momentum computes the gradient at $\theta_{t-1}$ and is blind to what that momentum is about to do; if the momentum is about to carry the iterate past the minimum, the gradient at $\theta_{t-1}$ cannot warn about it. NAG evaluates the gradient at $\theta_{t-1}+\mu v_{t-1}$, the point momentum is about to reach, so if that point has already started climbing back uphill, the new gradient opposes and brakes the velocity before the overshoot compounds. Formally, for an $L$-smooth convex function, Nesterov's method achieves the convergence rate

versus $O(1/k)$ for vanilla gradient descent — a provable quadratic-vs-linear improvement in iterations needed to reach a given accuracy, established by Nesterov's 1983 analysis for convex smooth objectives, and this is known to be optimal among first-order methods in that setting (Nesterov's lower bound matches this rate up to constants). Deep network losses are non-convex, so this guarantee does not transfer verbatim, but the same lookahead mechanism empirically reduces oscillatory overshoot and is observed to improve or match momentum's practical convergence speed at negligible extra cost.

Honest tradeoffs. Momentum introduces a new hyperparameter $\mu$ that interacts with $\eta$: an effective step size of roughly $\eta/(1-\mu)$ means that increasing $\mu$ without shrinking $\eta$ can destabilize training (overshoot, divergence). Momentum can also overshoot narrow minima, temporarily increasing loss before correcting — a real cost when using early stopping or noisy validation signals. Very high $\mu$ (e.g. $0.999$) creates a long memory horizon that reacts sluggishly to true changes in gradient direction late in training, motivating the coefficient schedules used in practice and the switch many practitioners make to adaptive methods (see Adam and Its Variants) or explicit decay schedules (see Learning-Rate Schedules).

Consider the simple ill-conditioned quadratic loss commonly used to illustrate momentum's benefit:

This mimics a "steep ravine": curvature along $\theta_2$ is 20× that along $\theta_1$. Start at $\theta_0 = (5, 1)$, use $\eta = 0.05$, and compare plain SGD ($\mu=0$) against momentum ($\mu = 0.8$), using the exact (non-Nesterov) update for clarity.

Reading the table: plain SGD's $\theta_2$ coordinate flips sign every single step ($1 \to -1 \to 1 \to -1$), because the curvature in that direction (20) is so high that $\eta \cdot 20 = 1$ overshoots the minimum every time — classic oscillation, with almost no net progress in $\theta_2$ after 4 steps ($1.000 \to 1.000$ in magnitude). Momentum, by contrast, accumulates a velocity in $\theta_2$ that partially cancels the alternating gradient signs (the geometric series argument from the previous section in action): its $\theta_2$ coordinate collapses toward zero much faster ($1.000 \to 0.024$ in magnitude) while its $\theta_1$ coordinate — the low-curvature direction where gradients keep the same sign — moves further per step (compare $4.073$ for SGD vs. $2.812$ for momentum at step 4), exactly the acceleration effect predicted. Now trace Nesterov's correction on the first step from the same start, $\theta_0=(5,1)$, $v_0=(0,0)$, $\mu=0.8$, $\eta=0.05$:

The divergence from classical momentum appears from step 2 onward, once $v_1\neq 0$:

Notice NAG's $\theta_2$-coordinate lands exactly at $0.000$ after step 2 — because it measured the gradient at the lookahead point $-0.80$ where the restoring force was already strong, braking the velocity precisely enough to cancel the overshoot rather than merely damping it, illustrating concretely the anticipatory correction argued for above.

Momentum-SGD (often written `SGD(lr=..., momentum=0.9)`) remains a default, frequently competitive optimizer in production deep learning frameworks such as PyTorch, TensorFlow, and JAX-based libraries, and is the standard choice for training large-scale convolutional networks in computer vision, where it often generalizes better than adaptive methods at comparable or better final accuracy. Nesterov momentum specifically is used in many published state-of-the-art vision training recipes (including variants of ResNet and other ImageNet-scale architectures) because its anticipatory correction stabilizes training with the large learning rates needed for large-batch, distributed training. The velocity-accumulation mechanism introduced here is also the direct conceptual ancestor of the first-moment estimate inside Adam and Its Variants, and momentum coefficients are routinely tuned jointly with Learning-Rate Schedules (e.g., warmup and cosine decay) in modern large-model training pipelines, since the two interact through the effective step size $\eta/(1-\mu)$.

Ask the tutor: How exactly does the algebraic substitution θ̃ₜ = θₜ + μvₜ prove that the two Nesterov update forms are equivalent, step by step?