Momentum gives gradient descent a "memory" of its recent update direction โ like a ball rolling downhill that builds up speed, it accelerates consistently in directions the gradient keeps pointing, and dampens oscillation in directions the gradient keeps flip-flopping.
Formula
\(\mathbf{v}_t\) is the "velocity" โ an exponentially weighted running average of past gradients. \(\beta\) (commonly 0.9) controls how much past gradients influence the current update; \(\beta=0\) reduces exactly to plain gradient descent.
The Physical Analogy
Picture a ball rolling down a valley-shaped loss surface. Plain gradient descent is like a ball with no mass โ it responds instantly and only to the current slope, zig-zagging across a narrow ravine. Momentum gives the ball mass and inertia โ it keeps moving in whatever direction it's been consistently pushed, smoothing out zig-zags and accelerating through consistently-sloped regions.
Why This Helps โ The Narrow Ravine Problem
In a narrow, elongated valley (common in real loss surfaces), momentum's averaging cancels out the zig-zag while reinforcing steady progress along the valley floor.
Numerical Example
With \(\beta=0.9\), starting \(\mathbf{v}_0=0\): if the gradient is consistently \(g=2\) for several steps, \(\mathbf{v}\) grows toward \(\frac{g}{1-\beta} = \frac{2}{0.1}=20\) โ the velocity "ramps up" and the effective step size becomes much larger than a single gradient's raw magnitude would suggest, accelerating progress in a consistent direction.
Code
import torch
# Manual momentum update
w = torch.tensor(10.0, requires_grad=True)
v = torch.tensor(0.0)
lr, beta = 0.1, 0.9
for step in range(20):
loss = w ** 2
loss.backward()
with torch.no_grad():
v = beta * v + w.grad
w -= lr * v
w.grad.zero_()
print(w)
import torch.optim as optim
# PyTorch handles this directly via the momentum argument
optimizer = optim.SGD([w], lr=0.1, momentum=0.9)
Common Mistakes
- Setting \(\beta\) too close to 1 (e.g. 0.999) for momentum specifically โ this can make the optimizer "overshoot" badly, since velocity barely decays and keeps pushing in a stale direction even after the true gradient has changed.
- Assuming momentum eliminates the need to tune the learning rate carefully โ it interacts with \(\eta\), and a learning rate that was well-tuned for plain SGD often needs to be reduced when momentum is added, since the effective step size grows.
Interview Relevance
Q: "What specific optimization problem does momentum solve?" It smooths out oscillation in narrow, elongated regions of the loss surface (where the gradient's direction along one axis flips repeatedly) while accelerating progress in directions the gradient consistently points โ by averaging recent gradients into a "velocity" term instead of reacting solely to the current gradient.
Practice Question
With \(\beta=0.5\) and \(\mathbf{v}_0=0\), if the gradient at step 1 is 4, what is \(\mathbf{v}_1\)? If the gradient at step 2 is also 4, what is \(\mathbf{v}_2\)?