Momentum Methods

Standard gradient descent updates parameters $\theta_t$ using only the instantaneous gradient:

On a well-conditioned, roughly spherical loss surface this works fine. But most real loss surfaces — especially deep network losses and any quadratic with a poorly conditioned Hessian — look like long narrow valleys: curvature is very high in some directions (steep walls) and very low in others (the gentle valley floor). Gradient descent's step size is limited by the steep directions (a larger learning rate would diverge there), so progress along the shallow direction, where we actually want to travel far, is agonizingly slow. The gradient vector oscillates back and forth across the valley walls while barely creeping along the floor.

Heavy-ball momentum introduces an auxiliary velocity vector $v_t$, initialized to zero, and a momentum coefficient $\beta \in [0, 1)$ (commonly $0.9$):

The velocity is an exponentially weighted moving average of all past gradients. Along directions where consecutive gradients agree in sign (the valley floor), these contributions add constructively and the effective step size grows. Along directions where gradients keep flipping sign (the valley walls), the contributions cancel and oscillation is damped. This is exactly the physical intuition of a ball with mass and friction: $\beta$ controls how much "memory" (inertia) the ball retains, and $1-\beta$ acts like friction that keeps the velocity bounded.

Nesterov accelerated gradient (NAG) modifies where the gradient is evaluated. Instead of computing $\nabla f(\theta_t)$ at the current point, it first takes a provisional step using the existing velocity ("look ahead"), evaluates the gradient there, and only then updates:

This look-ahead gradient acts as a built-in correction: if momentum is about to carry the parameters too far, the gradient at the look-ahead point already "sees" the overshoot and pushes back before it happens, rather than after. This seemingly small reordering is precisely what elevates NAG's convergence rate from merely fast to provably optimal for smooth convex problems, as shown below.

Both methods share the same two-state update loop — maintain $(\theta_t, v_t)$ and update both each iteration — differing only in where the gradient is evaluated. The pseudocode below implements both, parameterized by a flag, and is written so that it works identically for full-batch gradient descent or for Stochastic Gradient Descent simply by substituting the stochastic gradient estimator. ```text function MOMENTUM_GD(f, grad_f, theta0, eta, beta, num_steps, nesterov=False): theta = theta0 v = zeros_like(theta0) # velocity, same shape as parameters for t in 1 .. num_steps: if nesterov: lookahead = theta - eta beta v # peek ahead using current velocity g = grad_f(lookahead) # gradient AT the look-ahead point else: g = grad_f(theta) # heavy-ball: gradient at current point v = beta v + g # accumulate velocity theta = theta - eta v # move parameters return theta ``` Two equivalent reformulations are worth knowing because they appear in real libraries.

Two-step (momentum-in-update) form, common in PyTorch's SGD-with-momentum implementation: ```text v_{t+1} = beta v_t + grad_f(theta_t) theta_{t+1} = theta_t - eta v_{t+1} ``` This is exactly heavy-ball as written above.

Nesterov via momentum-corrected gradient, an algebraically equivalent form that avoids ever materializing a separate "look-ahead" parameter copy (useful when parameters are sharded across devices): ```text g_t = grad_f(theta_t) v_{t+1} = beta v_t + g_t theta_{t+1} = theta_t - eta (beta * v_{t+1} + g_t) ``` This produces the same trajectory as the look-ahead form up to re-indexing, and is the version actually implemented in most deep learning frameworks because it needs only the current-point gradient.

Step-by-step mechanics of a single heavy-ball update: 1. Compute the gradient $\nabla f(\theta_t)$ at the current parameters (one forward+backward pass). 2. Update the velocity: decay the old velocity by $\beta$, add the new gradient. 3. Update the parameters by stepping in the direction of $-v_{t+1}$, scaled by the learning rate $\eta$. 4. Repeat until a convergence criterion (gradient norm below a threshold, or a fixed iteration budget) is met. The only structural difference for Nesterov is that step 1 is evaluated not at $\theta_t$ but at the point $\theta_t - \eta\beta v_t$ — the location the ball would reach if it coasted on pure momentum with no new gradient information.

Per-step cost. Both heavy-ball and Nesterov require exactly one gradient evaluation per step (the same as plain gradient descent), plus $O(d)$ extra arithmetic to update and apply the velocity vector, where $d$ is the number of parameters. Memory overhead is one extra length-$d$ vector to store $v_t$. Since gradient computation (a forward and backward pass through a model) already costs $O(d)$ or more, momentum's overhead is asymptotically free — it does not change the per-step complexity class of the base algorithm at all. The entire benefit comes from needing fewer steps, not cheaper steps.

Convergence rate on strongly convex quadratics. Consider the model problem

with eigenvalues of $A$ ranging over $[\mu, L]$ ($\mu$ = strong convexity constant, $L$ = smoothness constant), and condition number $\kappa = L/\mu$. Analyzing gradient descent's update per eigen-direction with eigenvalue $\lambda$ gives a scalar recursion $\theta_{t+1} = (1 - \eta\lambda)\theta_t$, so convergence requires $|1 - \eta\lambda| < 1$ for every $\lambda \in [\mu, L]$ simultaneously. The best achievable worst-case contraction factor, optimized over $\eta$, is

meaning the number of iterations to reach accuracy $\varepsilon$ scales as $O(\kappa \log(1/\varepsilon))$. Heavy-ball momentum turns the same per-direction recursion into a second-order linear recurrence:

Its characteristic polynomial is $r^2 - (1+\beta-\eta\lambda) r + \beta = 0$. Choosing the classical Polyak parameters

makes both roots of the characteristic polynomial complex with modulus

Compare the two contraction rates: gradient descent needs $O(\kappa)$ iterations while heavy-ball needs $O(\sqrt{\kappa})$ iterations to reach the same accuracy — the condition number enters as a square root instead of linearly. For a valley with $\kappa = 10{,}000$ (entirely plausible in deep learning), this is the difference between roughly $10{,}000$ and $100$ iterations — a two-order-of-magnitude speedup, purely from remembering one extra vector. This local, per-quadratic-direction analysis is honest about heavy-ball's caveat: the clean $O(\sqrt{\kappa})$ rate is proven rigorously only for exact quadratics (or more generally, functions well approximated by them near a strongly convex minimum). For general smooth convex (not necessarily strongly convex) functions, heavy-ball's global worst-case guarantee is not better than plain gradient descent, and there exist convex functions on which heavy-ball with fixed parameters fails to converge at all. This is precisely the gap Nesterov's method closes.

Nesterov's global guarantee. For smooth convex functions (not necessarily strongly convex, Lipschitz gradient constant $L$), Nesterov's accelerated gradient method with the time-varying momentum schedule $\beta_t = \frac{t-1}{t+2}$ and $\eta = 1/L$ achieves

compared to plain gradient descent's $O(1/t)$ rate on the same class of functions. This $O(1/t^2)$ rate is not merely fast — Nemirovski and Yudin's classical lower bound shows no first-order method (using only gradient evaluations) can do better than $\Omega(1/t^2)$ on this function class in the worst case. Nesterov's method is therefore optimal: it is impossible to improve its asymptotic rate for general smooth convex functions using only gradient information. In the strongly convex quadratic case, the momentum-coefficient analog of the Polyak choice again yields the $O(\sqrt{\kappa})$ per-iteration contraction rate, matching heavy-ball's local rate but with a rigorously provable global guarantee that heavy-ball lacks.

Why the look-ahead matters, intuitively. Heavy-ball's velocity can drive the iterate past the minimum along the low-curvature direction before the gradient (measured at the pre-overshoot point) has a chance to react — creating a lag between accumulation of momentum and correction. Nesterov's method evaluates the gradient after the momentum step is provisionally taken, so the correction term already accounts for where the point is about to end up, producing tighter control over the trajectory and enabling a stronger convergence proof (typically via a carefully constructed "estimate sequence" or Lyapunov function that decreases monotonically — a proof technique this article will not reproduce in full, but whose conclusion, the $O(1/t^2)$ rate above, is the one to remember).

Tradeoffs. Momentum is not universally superior. In the stochastic setting (mini-batch gradients), momentum also accumulates gradient noise, not just signal; without an accompanying decayed or well-tuned learning rate, large $\beta$ can amplify variance rather than damp it, and pathological cases (e.g. very large $\beta$ combined with too-large $\eta$) can cause divergence that plain SGD would have avoided. Momentum also introduces one more hyperparameter to tune, and the elegant $O(\sqrt{\kappa})$ theory is proved for deterministic, convex, quadratic-like problems — real neural loss landscapes are nonconvex, so the guarantee is best read as strong intuition and an excellent practical heuristic rather than a literal proof for deep learning.

Take the 2-D quadratic

an intentionally ill-conditioned "ravine": steep along $\theta_1$ ($L=100$), shallow along $\theta_2$ ($\mu=1$), so $\kappa = 100$. Start at $\theta_0 = (1, 1)$; the minimum is $\theta^* = (0,0)$.

Plain gradient descent needs $\eta < 2/L = 0.02$ for stability along $\theta_1$; take $\eta = 0.018$.

Notice $\theta_1$ flips sign every step (it overshoots the minimum along the steep direction repeatedly — the oscillation across the ravine wall) while $\theta_2$, the direction we actually need to traverse, barely moves ($0.982 \to 0.929$ after four steps). This is exactly the zig-zag failure mode motivating momentum: reaching $f < 10^{-3}$ here takes on the order of hundreds of iterations because progress is bottlenecked by the slow $\theta_2$ direction.

Heavy-ball momentum, using Polyak's optimal parameters for this quadratic:

(rows marked show the characteristic transient overshoot of underdamped heavy-ball — complex characteristic roots mean the steep-direction error briefly grows in magnitude while oscillating in sign, before decaying geometrically at rate $\rho_{HB} = (\sqrt{\kappa}-1)/(\sqrt{\kappa}+1) \approx 0.818$ per step; contrast with GD's $\rho_{GD} = (\kappa-1)/(\kappa+1) \approx 0.980$ per step.)* Even though the early iterations look worse, the asymptotic contraction factor per step is $0.818$ (momentum) versus $0.980$ (plain GD). Solving $\rho^t = 10^{-3}$ for each: plain GD needs $t \approx \log(10^{-3})/\log(0.980) \approx 342$ iterations; heavy-ball needs $t \approx \log(10^{-3})/\log(0.818) \approx 35$ iterations — roughly a 10× reduction, consistent with the $\sqrt{\kappa} = 10$ factor predicted by the theory ($\kappa = 100 \Rightarrow \sqrt\kappa = 10$).

Nesterov, with the same $\eta,\beta$, differs only in evaluating the gradient at the look-ahead point $\theta_t - \eta\beta v_t$ each step; on this exact quadratic (where the Hessian is constant) the trajectories are extremely close to heavy-ball's, and the same $O(\sqrt\kappa)$ asymptotic rate applies — the practical distinction between the two methods shows up mainly on non-quadratic or stochastic objectives, where Nesterov's look-ahead correction more reliably prevents runaway oscillation.

Momentum's mathematical roots trace to control theory and mechanics — literally modeling optimization as a physical system with mass and friction — before being formalized for numerical optimization in the 1960s and then proven optimal for a broad function class in the 1980s.

Momentum is not a niche technique — it is a load-bearing component of essentially every large-scale learning system in production today. Training image classifiers, language models, and recommendation systems almost universally uses SGD with heavy-ball or Nesterov momentum, or an adaptive optimizer (Adam, AdamW) whose update rule is built directly on top of an exponentially weighted momentum term for the first moment of the gradient. Deep learning frameworks such as PyTorch and TensorFlow expose momentum as a first-class, default-on hyperparameter in their SGD optimizer classes precisely because of the reliability gains documented above. Beyond deep learning, momentum-style acceleration appears throughout numerical computing: accelerated proximal gradient methods (FISTA) used in sparse recovery and compressed sensing are a direct application of Nesterov's technique to nonsmooth composite objectives; large-scale convex solvers in operations research and control use accelerated first-order methods when second-order (Newton-type) methods are too expensive per iteration due to high dimensionality. Any setting with a smooth, expensive-to-differentiate objective and a poorly conditioned landscape — physics simulations, large-scale logistic regression, matrix factorization for recommender systems — is a candidate for momentum acceleration. Understanding momentum is also the essential stepping stone to understanding adaptive methods and the broader landscape of modern optimizer design, since nearly every widely used optimizer is, at its core, "momentum plus some coordinate-wise rescaling."

Ask the tutor: How does momentum interact with adaptive per-coordinate learning rates in optimizers like Adam?