SGD and Momentum

Learning Objectives

By the end of this section, you will be able to:

1. Explain the difference between batch, mini-batch, and stochastic gradient descent and why mini-batch SGD is the standard in practice
2. Identify the ravine problem where vanilla SGD oscillates in high-curvature directions and crawls in low-curvature directions
3. Derive the momentum update rule and explain how exponential averaging of gradients accelerates convergence
4. Implement SGD with momentum from scratch in NumPy and compare it against vanilla SGD
5. Use PyTorch's SGD optimizer with the $\texttt{momentum}$ parameter to train a real network
6. Describe Nesterov accelerated gradient and when it helps

Where We Left Off

In Chapter 8, we derived backpropagation and used it to compute gradients for all 31 parameters of our 2×2 diagonal flip network. Then we applied the simplest possible update rule:

$w_{\text{new}} = w_{\text{old}} - \eta \cdot \frac{\partial L}{\partial w}$

After one step with $\eta = 0.1$, the loss dropped from 0.896 to about 0.63 — a 30% improvement. That felt good. But here's the question that Chapter 8 left unanswered: is this the best we can do?

The answer is no. The simple update rule we used — called vanilla gradient descent — has fundamental problems that get worse as networks get larger. This chapter introduces optimizers: smarter update rules that converge faster, more reliably, and with less tuning.

What changes in this chapter: In Chapter 8, we focused on computing gradients (the hard part). Now we take those gradients as given and focus on using them more intelligently. The gradient tells us which direction is downhill — the optimizer decides how far and how fast to step.

Stochastic Gradient Descent

Before improving the update rule, we need to address a practical issue: how many training examples do we use per gradient step?

Three Flavors of Gradient Descent

Batch gradient descent computes the gradient using all training examples, then takes one step. The gradient is exact but you only get one update per pass through the data. For our 16-image dataset, that's one step per epoch.

Stochastic gradient descent (SGD) computes the gradient from a single random example and updates immediately. Each gradient is noisy — it may point in a slightly wrong direction — but you get 16 updates per epoch instead of 1. In practice, the noise can actually help escape local minima.

Mini-batch SGD splits the difference: use $B$ examples per step (typically 32–512). The gradient is smoother than single-sample SGD but you still get multiple updates per epoch. This is what nearly all modern training uses.

The Noise Trade-Off

Mini-batch gradients are noisy estimates of the true gradient. Consider our 16-image dataset. The true gradient (averaged over all 16 images) points directly toward the minimum. But the gradient from image #3 alone might point slightly off to the side, because image #3 has different features than the average.

This noise has consequences. Instead of walking smoothly downhill, SGD takes a random walk with a downhill drift — zig-zagging toward the minimum rather than marching straight there. The drift gets you there eventually, but the zig-zags waste steps.

The Ravine Problem

The deepest problem with vanilla SGD isn't noise — it's geometry. Real loss surfaces have different curvatures in different directions. Imagine a long, narrow valley (a ravine):

• Along the valley floor (the long axis), the surface slopes gently. The gradient is small. Progress is slow.
• Across the valley walls (the short axis), the surface is steep. The gradient is large. Updates overshoot and bounce back and forth.

Mathematically, consider the loss function $L(w_1, w_2) = w_1^2 + 10 \cdot w_2^2$. The curvature in the $w_2$ direction (second derivative = 20) is ten times larger than in the $w_1$ direction (second derivative = 2). The contour lines form elongated ellipses — a ravine.

The Learning Rate Dilemma

This creates an impossible dilemma for vanilla SGD:

There is no single learning rate that works well for both directions simultaneously. You are forced to choose a conservative (small) learning rate to prevent the steep direction from diverging, which makes the gentle direction painfully slow.

The ravine problem in one sentence: vanilla SGD can only take one-size-fits-all steps, but different directions need different step sizes. This mismatch is the #1 reason vanilla SGD is slow on real loss surfaces, where curvature ratios of 100:1 or 1000:1 are common.

This is where momentum comes in. It solves both problems — noise and ravines — with a single, elegant idea.

Momentum: The Physics Intuition

Imagine a heavy ball rolling down our ravine. Unlike vanilla SGD (which teleports to a new position each step), the ball has inertia. It keeps rolling in the direction it was already moving. This gives it two superpowers:

1. Acceleration along the valley floor: the gradient consistently points along the valley. Each step adds more velocity in the same direction. The ball speeds up, taking ever-larger steps — exactly what we want.
2. Dampening across the walls: the gradient alternates direction (left wall, right wall, left wall...). The velocity in this direction oscillates too, but each oscillation partially cancels the previous one. The ball averages out the bouncing and rolls smoothly along the floor.

The same logic applies to noisy gradients. If noise pushes the ball left on step 5 and right on step 6, the velocity absorbs both pushes and barely changes direction. The signal (consistent downhill) accumulates while the noise (random zig-zags) cancels.

The Mathematics of Momentum

The momentum update introduces a velocity vector $\mathbf{v}$ that accumulates past gradients. At each step $t$:

Step 1: Update the velocity

$\mathbf{v}_t = \beta \cdot \mathbf{v}_{t-1} + \nabla L(\mathbf{w}_t)$

Step 2: Update the weights

$\mathbf{w}_{t+1} = \mathbf{w}_t - \eta \cdot \mathbf{v}_t$

Here $\beta$ is the momentum coefficient (typically 0.9) and $\eta$ is the learning rate. When $\beta = 0$, this reduces to vanilla SGD.

Unrolling the Velocity

Let's expand $\mathbf{v}_t$ to see what it really computes:

$\mathbf{v}_t = \nabla L_t + \beta \cdot \nabla L_{t-1} + \beta^2 \cdot \nabla L_{t-2} + \beta^3 \cdot \nabla L_{t-3} + \cdots$

This is an exponential moving average of gradients. The current gradient has weight 1, the previous gradient has weight $\beta = 0.9$, two steps ago has weight $\beta^2 = 0.81$, and so on. Gradients from the distant past fade away exponentially.

The Effective Window

The sum of the weights is $1 + \beta + \beta^2 + \cdots = \frac{1}{1 - \beta}$. For $\beta = 0.9$, this is $\frac{1}{0.1} = 10$. So the velocity effectively averages the last ~10 gradients, weighted toward the most recent ones.

Why This Solves the Ravine Problem

Consider what happens in each direction of our ravine:

Along the valley floor (consistent direction): every gradient points roughly the same way. The velocity accumulates: $v \approx g + 0.9g + 0.81g + \cdots \approx 10g$. The effective step is 10× larger than vanilla SGD. We race along the valley floor.

Across the valley walls (oscillating direction): the gradient alternates sign. The velocity alternates too: $v \approx g - 0.9g + 0.81g - \cdots$. The terms partially cancel, giving $v \approx g \cdot \frac{1}{1+\beta} \approx 0.53g$. The effective step is smaller than vanilla SGD. The oscillation is dampened.

Momentum is adaptive without being per-parameter. A single $\beta$ value automatically adjusts the effective step size in each direction based on gradient consistency. Consistent directions get amplified; oscillating directions get dampened. No extra hyperparameters needed.

Interactive: SGD vs Momentum

The visualization below shows vanilla SGD (orange) and SGD with momentum (blue) optimizing the same loss surface. Both start at the green dot (5, 2) and aim for the white circle at the origin (0, 0).

Try these experiments:

1. Default settings (lr=0.015, β=0.9, noise=0): Press play. Watch momentum race ahead while SGD barely moves in $w_1$.
2. Increase learning rate to ~0.04: SGD starts oscillating in $w_2$. Momentum stays smooth.
3. Add noise (noise=5): SGD path becomes jagged. Momentum path stays smoother because it averages out the noise.
4. Set β=0: The blue path matches orange exactly — confirming that $\beta = 0$ is vanilla SGD.
5. Set β=0.99: Very heavy momentum. The ball overshoots initially but then corrects. Too much momentum is slow to change direction.

NumPy: SGD vs Momentum from Scratch

Let's implement momentum from scratch and see the acceleration in numbers. We minimize $f(w) = w^2$ starting at $w = 5$. The gradient is $\frac{df}{dw} = 2w$. We use the same learning rate (0.01) for both vanilla SGD and momentum, so the only difference is the velocity buffer.

1import numpy as np
2
3# ── Minimize f(w) = w² starting at w=5 ──
4# Gradient: df/dw = 2w
5# Minimum: w* = 0
6
7# === Vanilla SGD ===
8w_sgd = 5.0
9lr = 0.01
10
11print("=== Vanilla SGD (lr=0.01) ===")
12for step in range(8):
13    grad = 2 * w_sgd
14    w_sgd = w_sgd - lr * grad
15    print(f"Step {step}: w={w_sgd:.4f}")
16
17# === SGD + Momentum ===
18w_mom = 5.0
19v = 0.0
20beta = 0.9
21
22print("\n=== SGD + Momentum (lr=0.01, β=0.9) ===")
23for step in range(8):
24    grad = 2 * w_mom
25    v = beta * v + grad
26    w_mom = w_mom - lr * v
27    print(f"Step {step}: v={v:.3f}  w={w_mom:.4f}")
28
29print(f"\nAfter 8 steps:")
30print(f"  Vanilla SGD: w = {w_sgd:.4f}")
31print(f"  Momentum:    w = {w_mom:.4f}")

The results tell the story:

Momentum moved 3.4× further in the same number of steps. The velocity built up from 10 to 45.3, making each step progressively larger. This is the power of accumulated momentum in a consistent direction.

PyTorch: Training with Momentum

Now let's apply momentum to real neural network training. We train our diagonal flip network from Chapters 7–8 using both vanilla SGD and SGD with momentum, on all 16 binary images. The models start with identical random weights. The only difference is one line: $\texttt{momentum=0.9}$.

1import torch
2import torch.nn as nn
3
4torch.manual_seed(42)
5
6# ── Network (same architecture as Chapters 7-8) ──
7class FlipNet(nn.Module):
8    def __init__(self):
9        super().__init__()
10        self.layer1 = nn.Linear(4, 3)
11        self.layer2 = nn.Linear(3, 4)
12        self.relu = nn.ReLU()
13
14    def forward(self, x):
15        return self.layer2(self.relu(self.layer1(x)))
16
17# ── Dataset: all 16 binary 2×2 images ──
18X = torch.tensor([[int(b) for b in f"{i:04b}"]
19                   for i in range(16)], dtype=torch.float32)
20Y = X[:, [0, 2, 1, 3]]  # diagonal flip
21
22# ── Train with vanilla SGD ──
23torch.manual_seed(42)
24model_sgd = FlipNet()
25opt_sgd = torch.optim.SGD(model_sgd.parameters(), lr=0.05)
26
27# ── Train with SGD + Momentum ──
28torch.manual_seed(42)
29model_mom = FlipNet()
30opt_mom = torch.optim.SGD(model_mom.parameters(),
31                           lr=0.05, momentum=0.9)
32
33for epoch in range(30):
34    loss_sgd = loss_mom = 0.0
35    for i in range(16):
36        # --- Vanilla SGD ---
37        pred = model_sgd(X[i])
38        lv = torch.mean((pred - Y[i]) ** 2)
39        opt_sgd.zero_grad()
40        lv.backward()
41        opt_sgd.step()
42        loss_sgd += lv.item()
43        # --- Momentum ---
44        pred = model_mom(X[i])
45        lm = torch.mean((pred - Y[i]) ** 2)
46        opt_mom.zero_grad()
47        lm.backward()
48        opt_mom.step()
49        loss_mom += lm.item()
50    if epoch % 5 == 0:
51        print(f"Epoch {epoch:2d}:  SGD={loss_sgd/16:.4f}"
52              f"  Momentum={loss_mom/16:.4f}")

The difference is dramatic. With momentum, the network reaches low loss in roughly half the epochs. The momentum buffer $\mathbf{v}$ for each of the 31 parameters accumulates gradient information across the 16 training images, effectively remembering "which way is downhill on average" rather than being yanked around by individual images.

Nesterov Accelerated Gradient

In 1983, Yurii Nesterov proposed a subtle but powerful improvement to momentum. The idea: instead of computing the gradient at the current position, compute it at the position you're about to jump to.

Standard Momentum vs Nesterov

Standard momentum: "I'm at position $\mathbf{w}$. What's the gradient here? OK, combine it with my velocity and move."

Nesterov momentum: "My velocity will carry me to roughly $\mathbf{w} - \eta \beta \mathbf{v}$. What's the gradient there? That's a better gradient because it accounts for where I'm actually going."

The Nesterov update:

$\mathbf{v}_t = \beta \cdot \mathbf{v}_{t-1} + \nabla L(\mathbf{w}_t - \eta \beta \mathbf{v}_{t-1})$

$\mathbf{w}_{t+1} = \mathbf{w}_t - \eta \cdot \mathbf{v}_t$

This "look-ahead" gradient gives Nesterov momentum a corrective quality. If the velocity is about to overshoot the minimum, the look-ahead gradient will already be pointing back — so the velocity starts slowing down before the overshoot happens instead of after.

In PyTorch: One Flag

Enabling Nesterov momentum in PyTorch requires exactly one change:

$\texttt{torch.optim.SGD(params, lr=0.05, momentum=0.9, nesterov=True)}$

In practice, Nesterov gives a small but consistent improvement over standard momentum. It is especially helpful when the loss surface has sharp turns — Nesterov "sees around the corner" while standard momentum overshoots and has to backtrack.

Connection to Modern Training

SGD with momentum is not just a historical artifact — it remains deeply relevant in modern deep learning:

Momentum as the Foundation

Adam (Section 2) extends momentum with per-parameter adaptive learning rates. Adam maintains two exponential moving averages: the first moment (gradient mean, like momentum) and the second moment (gradient variance). Understanding momentum is essential to understanding Adam.

Large-scale training of models like GPT, LLaMA, and BERT actually uses Adam or AdamW (Adam with weight decay), which build directly on the momentum concept. The velocity buffer idea is unchanged — Adam just adds a second buffer for gradient magnitude.

SGD vs Adam in Practice

Interestingly, well-tuned SGD with momentum often achieves better generalization than Adam on computer vision tasks. This is why many image classification papers (ResNets, EfficientNets) still use SGD+momentum. The hypothesis: Adam's per-parameter adaptation finds sharper minima that generalize worse.

Connection to Transformer Training

When training transformers, momentum appears in several places:

• AdamW optimizer: the standard optimizer for transformer training maintains a momentum buffer $m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t$ with $\beta_1 = 0.9$ — this is exactly the exponential moving average we studied in this section
• Learning rate warmup: transformers typically warm up the learning rate from 0 to the target value over the first 1,000–4,000 steps, then decay it. This interacts with momentum because the velocity buffer needs time to build up
• Gradient accumulation: when training with batch sizes of 1M+ tokens (e.g. Chinchilla, LLaMA), gradients are accumulated across micro-batches. The accumulated gradient acts like a batch gradient, and the optimizer's momentum provides additional temporal smoothing on top

Summary

1. SGD processes mini-batches instead of the full dataset, giving noisy but frequent gradient updates. Mini-batch SGD is the default in practice.
2. The ravine problem: when curvatures differ across directions, vanilla SGD oscillates in steep directions and crawls in gentle ones. No single learning rate works well for all directions.
3. Momentum adds a velocity buffer $v_t = \beta v_{t-1} + \nabla L$ that is an exponential moving average of past gradients. Consistent directions accelerate; oscillating directions dampen.
4. The standard setting is $\beta = 0.9$, which averages the last ~10 gradients. This gives roughly a 3–5× convergence speedup over vanilla SGD.
5. In PyTorch: just add $\texttt{momentum=0.9}$ to your SGD optimizer. For Nesterov's improvement, add $\texttt{nesterov=True}$.
6. Momentum is the foundation of Adam, which adds per-parameter adaptive learning rates. We cover Adam in Section 2.

Looking ahead: Momentum gives every parameter the same effective learning rate (just smoothed). But what if some parameters need a larger step than others? That's the idea behind adaptive learning rate methods like Adam, which we explore in Section 2.