DEV Community: Harshil Rami

Blog 2: Momentum-Based Optimizers

Harshil Rami — Wed, 22 Apr 2026 18:03:58 +0000

Giving the optimizer a memory — and teaching it to look before it leaps

SGD knows where it is. Momentum knows where it's been. Nesterov knows where it's going.
That single sentence is the entire story of this post.

The ravine problem, visualized

Let's be concrete about what SGD's zig-zagging actually looks like.

Suppose your loss surface is an elongated valley — steep walls on the left and right, a gentle slope running toward the minimum far ahead. This is the classic ravine geometry, and it's not an academic toy. It shows up naturally when your features have very different scales, when layers have different learning dynamics, or when you're in the early phases of training a deep network.

SGD's update on this surface behaves as follows:

Along the steep axis (across the ravine): the gradient is large. SGD takes a big step, overshoots, corrects back, overshoots again. The updates oscillate violently.
Along the shallow axis (down the ravine, toward the minimum): the gradient is small. SGD takes tiny, tentative steps. Progress is glacial.

The result is a path that looks like a snake moving sideways more than forward. You can shrink the learning rate to tame the oscillations on the steep axis, but that makes the shallow axis even slower. There's no single learning rate that handles both directions well.

This is the fundamental limitation SGD leaves on the table, and it's what momentum is designed to fix.

Momentum: giving the optimizer velocity

The core idea behind momentum is borrowed directly from physics. Instead of updating parameters based on the current gradient alone, we maintain a velocity vector v that accumulates a running average of past gradients.

vₜ = β · vₜ₋₁ + η · ∇L(θₜ)
θₜ₊₁ = θₜ − vₜ

Where β (beta) is the momentum coefficient — typically 0.9. Some formulations absorb the learning rate differently; the semantics are equivalent.

Let's unpack what this actually does.

On the oscillating (steep) axis:
Gradients alternate sign — positive, negative, positive, negative. The velocity accumulates these with the decay factor β. Because they cancel each other out over time, the velocity along this axis stays small. Oscillations are damped.

On the consistent (shallow) axis:
Gradients consistently point in the same direction — always slightly downhill. The velocity accumulates these constructively. Each step adds to the previous. The effective step size grows, and the optimizer accelerates.

This is the momentum effect: dampening in oscillating directions, acceleration in consistent ones. The optimizer builds up speed where the surface is consistently sloped and brakes naturally where the surface is ambiguous.

Effective learning rate under momentum

With β = 0.9, a gradient that persists in the same direction for many steps produces a velocity roughly 1/(1−β) = 10× the nominal learning rate. This is why momentum often requires a slightly lower learning rate than vanilla SGD — the effective step size is larger.

More precisely, if the gradient is constant at g, the velocity converges to:

v* = η · g / (1 − β)

So momentum scales up the effective learning rate by 1/(1−β). Set β = 0.9 → 10× amplification. Set β = 0.99 → 100×. This amplification is the source of both momentum's power and its instability if misconfigured.

The ball analogy

Momentum is often described as a ball rolling down a hill. The ball doesn't instantly respond to every local slope — it carries inertia. A small bump doesn't stop it; it takes a sustained uphill slope to decelerate it meaningfully.

This analogy is accurate and useful, but it has a limit: the ball analogy suggests the optimizer might overshoot and roll up the other side of a valley. That's a real failure mode of high-momentum settings. If β is too large, the optimizer can oscillate around minima rather than settling into them, or sail through a narrow good basin entirely.

Nesterov Accelerated Gradient: look before you leap

Momentum is good. Nesterov Accelerated Gradient (NAG), proposed by Yurii Nesterov in 1983, makes one surgical improvement that turns out to matter significantly in practice.

The problem with standard momentum: the gradient is evaluated at the current position, before applying the velocity. By the time you apply the update, you're no longer at that position — you've already moved. You're using stale directional information.

NAG fixes this with a simple conceptual shift: evaluate the gradient at the position you're about to arrive at, then correct from there.

θ_lookahead = θₜ − β · vₜ₋₁          # project forward
vₜ = β · vₜ₋₁ + η · ∇L(θ_lookahead)  # gradient at projected position
θₜ₊₁ = θₜ − vₜ

The momentum step projects you forward to where you'll be before the gradient correction. Then you evaluate the gradient there. This means the correction accounts for the momentum-driven position, not the pre-momentum position.

Why this helps

Think of it this way. Standard momentum is like running toward a wall and only noticing the wall after you've taken your full step. Nesterov is like looking ahead as you run and starting to slow down before you hit the wall.

In regions where the momentum is carrying you toward a steep uphill, NAG detects that uphill slope earlier and applies a corrective force sooner. The update is more anticipatory than reactive.

In practice, NAG converges faster than standard momentum on convex problems — Nesterov's original theoretical analysis showed an O(1/k²) convergence rate versus SGD's O(1/k), a meaningful gap. For non-convex deep learning loss surfaces, the improvement is empirical rather than provably guaranteed, but it's consistently observed.

NAG in the equivalent update form

The two-equation NAG formulation above has an equivalent single-equation form that's more commonly implemented:

vₜ = β · vₜ₋₁ + ∇L(θₜ − β · vₜ₋₁)
θₜ₊₁ = θₜ − η · vₜ

Both are equivalent; the second form makes it clearer that the only change from standard momentum is where the gradient is evaluated.

Side-by-side comparison

	Gradient Descent	SGD	Momentum	NAG
Gradient source	Full dataset	Single sample	Mini-batch	Mini-batch (at projected pos.)
Memory	None	None	Velocity vₜ	Velocity vₜ
Oscillation handling	None	None	Dampens via averaging	Dampens + anticipates
Convergence rate (convex)	O(1/k)	O(1/√k)	O(1/k)	O(1/k²)
Typical β	—	—	0.9	0.9–0.99

The convergence rate column is worth reading carefully. SGD's O(1/√k) is actually worse than GD's O(1/k) — the variance of stochastic gradients costs you a square root. Momentum restores GD-level rates. Nesterov goes further.

Implementation and practical notes

# Standard Momentum
v = 0
for batch in dataloader:
    grad = compute_gradient(loss_fn, model, batch)
    v = beta * v + learning_rate * grad
    theta = theta - v

# Nesterov Accelerated Gradient
v = 0
for batch in dataloader:
    # Evaluate gradient at lookahead position
    theta_lookahead = theta - beta * v
    grad = compute_gradient(loss_fn, model, theta_lookahead)
    v = beta * v + learning_rate * grad
    theta = theta - v

In PyTorch, both are available via torch.optim.SGD with momentum and nesterov flags:

# Momentum
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# NAG
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9, nesterov=True)

Practical hyperparameter guidance:

β = 0.9 is the standard default. It works well across a wide range of architectures.
β = 0.99 gives stronger smoothing but risks slower response to genuine direction changes and can overshoot narrow minima.
When switching from SGD to momentum, reduce the learning rate by roughly 1/(1−β) — so if η = 0.1 for SGD, try η = 0.01 with β = 0.9.
NAG is almost always preferable to plain momentum for the same computational cost. Default to it.

Where momentum still falls short

Momentum is a major step forward. But it inherits one fundamental limitation from SGD: a single global learning rate for all parameters.

Your model's parameters live in very different regimes. The embedding layer of a language model sees only a handful of tokens in each batch — its effective gradient is sparse and noisy. The final linear layer sees a dense, consistent gradient every step. Both are updated with the same η.

This is deeply suboptimal. Sparse parameters should move faster when they do receive signal. Dense parameters can afford more conservative updates to avoid oscillation.

Momentum has no mechanism to learn this. It smooths over time but not over parameters.

That's exactly the problem Blog 3 solves. AdaGrad will introduce per-parameter learning rates, scaling each update by the history of that parameter's gradient magnitude. RMSProp will fix AdaGrad's long-term decay problem. And the combination of per-parameter scaling with momentum will eventually give us Adam.

Three things to hold onto

Concept	What it means	Why it matters
Velocity accumulation	Past gradients persist via β decay	Accelerates in consistent directions, dampens oscillations
Lookahead gradient (NAG)	Gradient evaluated at projected position	Earlier correction, better convergence rate
Effective LR scaling	Velocity → η/(1−β) effective step	Must tune LR down when adding momentum

What's next

In Blog 3, we shift from when the optimizer has seen a gradient to which parameters have seen large gradients. AdaGrad introduces a per-parameter accumulator — parameters that receive frequent, large gradients get smaller effective learning rates; sparse parameters get larger ones. RMSProp then fixes AdaGrad's fatal flaw: the accumulator grows without bound, eventually shrinking all learning rates to zero.

If momentum gave the optimizer a memory across time, adaptive methods give it a memory across parameters. Both are necessary. Neither is sufficient alone.

This is Blog 2 of an 8-part series on optimization algorithms for deep learning.

Blog 1: Foundations of Gradient Descent

Harshil Rami — Wed, 22 Apr 2026 17:54:43 +0000

How neural networks learn — and why the obvious approach breaks immediately

Every optimizer you'll ever use — Adam, AdamW, Lion, LAMB — is an answer to a problem that gradient descent creates. To understand why those answers exist, you need to feel the problem first.

The loss surface is a landscape you can't see

Imagine you're blindfolded, standing somewhere on a hilly terrain. Your only tool is a stick: you can poke the ground around you and measure the slope. Your goal is to reach the lowest valley.

That's optimization.

The "terrain" is your loss surface — a high-dimensional function L(θ) mapping your model's parameters θ to a scalar loss. You can't see the whole surface. You can only evaluate the gradient at your current position and take a step.

The question every optimizer tries to answer: which direction, and how far?

Gradient Descent: the right idea, the wrong scale

Gradient Descent (GD) is the foundational answer. Given a loss function L(θ), we compute the gradient over the entire dataset and update:

θ ← θ − η · ∇L(θ)

Where η (eta) is the learning rate — a scalar controlling step size.

The update rule is clean. The gradient points in the direction of steepest ascent, so we move opposite to it. Mathematically, this is the direction of maximum local decrease in loss.

The intuition is correct. The implementation is catastrophically expensive.

To compute ∇L(θ) exactly, you need to pass your entire dataset through the model. For ImageNet-scale data or modern LLM corpora, this means billions of examples per update. You'd compute one parameter update per epoch. On a 100M parameter model. That's not slow — it's dead on arrival.

GD also has a subtle failure mode people underappreciate: when your dataset has redundant structure (and it almost always does), successive gradients are nearly identical. You're paying full-dataset cost for almost zero additional information after the first few passes.

Stochastic Gradient Descent: embrace the noise

The fix seems almost too simple: instead of computing the gradient over all N samples, pick one sample at random and update on that alone.

θ ← θ − η · ∇Lᵢ(θ)    for a randomly sampled i

This is Stochastic Gradient Descent (SGD). The gradient estimate is now noisy — it's a single-sample approximation of the true gradient. But that noise turns out to be a feature, not a bug.

Why noisy updates help:

Escaping shallow local minima. A noisy gradient doesn't always point exactly downhill. This stochasticity gives the optimizer a kind of thermal energy — it can jitter out of shallow basins that would trap a deterministic update.
Better generalization (empirically). The noise in SGD acts as implicit regularization. Models trained with SGD often generalize better than those trained with exact gradient methods, particularly in overparameterized regimes. There's a growing body of theory around this — the "flat minima" hypothesis suggests noisy SGD preferentially finds wider, flatter basins that transfer better to test data.
Speed per effective update. One SGD step is O(1) in data cost. You can make N updates in the time GD makes one, seeing every sample along the way.

The cost of noise:

SGD's gradient estimate has high variance. Updates zig-zag erratically, especially in directions where the loss surface has high curvature along one axis and low curvature along another (the classic "ravine" geometry). The path to the minimum looks like a drunk person's walk rather than a confident descent.

You can reduce the learning rate to smooth this out, but then you lose the speed advantage. You're always trading variance against convergence rate.

Mini-Batch Gradient Descent: the practical compromise

The resolution in practice is obvious in retrospect: use a small batch of B samples.

θ ← θ − η · (1/B) · Σᵢ∈Bₜ ∇Lᵢ(θ)

Where Bₜ is a randomly sampled mini-batch of size B at step t.

This is Mini-Batch Gradient Descent (MBGD) — and when practitioners say "SGD" today, this is almost always what they mean. Typical batch sizes range from 32 to 512, though the right choice depends on your model, hardware, and regularization goals.

What mini-batching buys you:

Variance reduction. Averaging over B samples reduces gradient variance by a factor of B compared to single-sample SGD, without B× the compute cost (thanks to parallelism on GPU/TPU).
Hardware efficiency. GPUs are throughput machines — they saturate at batch sizes that fully utilize memory bandwidth. A single-sample forward pass wastes most of your compute budget.
Enough noise to generalize. Mini-batch gradients are still noisy enough to provide the regularization benefits of SGD, unlike full-batch gradients.

The batch size isn't free:

Larger batches reduce noise, which sounds good, but:

They can converge to sharper minima with worse generalization (the "large-batch training problem" — Keskar et al., 2017).
They require proportionally larger learning rates to maintain the same effective update magnitude, but scaling LR linearly with batch size breaks down at large B.
Beyond a critical batch size, you're paying compute cost without improving convergence speed.

This last point becomes the central tension in Blog 6, when we look at LARS and LAMB — optimizers specifically designed to handle very large batches in distributed LLM training.

The update rule in full

Here's where we stand after mini-batch SGD. The complete training loop:

for epoch in range(num_epochs):
    shuffle(dataset)
    for batch in get_batches(dataset, batch_size=B):
        grad = compute_gradient(loss_fn, model, batch)
        θ = θ - learning_rate * grad

This works. It's what trains ResNets, early language models, and still forms the backbone of large-scale training in certain regimes (SGD with momentum remains competitive with Adam on image classification tasks).

But watch what happens on a ravine-shaped loss surface:

The gradient along the short axis (high curvature) is large → big oscillating steps
The gradient along the long axis (low curvature, toward the minimum) is small → slow progress

The optimizer zig-zags across the ravine instead of marching down it. You need a very small learning rate to prevent divergence on the steep axis, which makes the shallow axis painfully slow.

This is exactly the problem Blog 2 solves. Momentum will give the optimizer memory — a velocity vector that accumulates in persistent directions and dampens oscillations. Nesterov will take that one step further, looking ahead before committing to the update.

Three things to hold onto

Before moving to momentum, here are the three tensions that the rest of this series is spent resolving:

Problem	What causes it	Solved by
Slow, expensive updates	Full-dataset gradient	SGD / Mini-batch
High-variance, zig-zagging path	Single/small-batch noise	Momentum (Blog 2)
Uniform learning rate for all params	LR is a global scalar	AdaGrad, RMSProp (Blog 3)

Every optimizer from here on is a targeted intervention on one of these failure modes — or a combination of several at once.

Key equations, plain-English summary

Method	Update rule	One sentence
GD	`θ ← θ − η · ∇L(θ)`	Exact gradient, entire dataset, one step per epoch
SGD	`θ ← θ − η · ∇Lᵢ(θ)`	Noisy gradient, one sample, fast but erratic
MBGD	`θ ← θ − η · (1/B)·Σ∇Lᵢ(θ)`	Averaged gradient, batch of B, the practical default

What's next

In Blog 2, we add velocity. Momentum accumulates past gradients into a running average, smoothing the zig-zagging path and accelerating convergence in consistent directions. Nesterov takes the lookahead step — evaluating the gradient at a projected future position rather than the current one.

If SGD is someone walking blindfolded downhill, momentum is that same person carrying a ball that's already rolling. It takes more to change direction. That turns out to be exactly what you want.

This is Blog 1 of an 8-part series on optimization algorithms for deep learning. Each post covers one family of optimizers, following a problem → limitation → next solution arc.