Why Adam Beats SGD: The Physics of Momentum

Topic Overview

What is Momentum?

Momentum is an optimization technique that accelerates gradient descent by accumulating a velocity vector based on the direction of past gradients. Instead of updating parameters using only the current gradient, it carries forward a fraction of the previous update direction.

Why does it exist?

Standard Gradient Descent is shortsighted. It only looks at the local slope and takes a step. In high-dimensional spaces, this causes severe “zigzagging” across the walls of narrow valleys (ravines) in the loss landscape. Momentum exists to smooth out these oscillations by introducing physical inertia.

What problem does it solve?

It solves the problem of slow convergence in ill-conditioned loss landscapes. It dampens oscillations in high-curvature directions while amplifying consistent movement in low-curvature directions. Practically, it turns hundreds of wasted zigzag steps into a direct, fast path toward the minimum.

Why should engineers care?

Without momentum, training a ResNet or a Vision Transformer (ViT) would take 5x to 10x longer, and might fail to converge entirely. Furthermore, you cannot understand the Adam optimizer—the default for LLMs—without understanding momentum, because Adam’s first moment (\(m\)) is literally just momentum. If you don’t understand momentum, you don’t understand modern deep learning optimization.

Prerequisites

1. Gradient Descent

Why it is needed: Momentum is an extension of vanilla gradient descent. If you don’t understand how \(\theta_{t+1} = \theta_t - \eta \nabla L\) works, you cannot understand what momentum is modifying.
Where it appears: It serves as the base update equation that momentum wraps around.

Historical Motivation

In the 1960s, researchers optimizing convex quadratic functions faced an engineering wall.

The Problem: Second-order methods (like Newton’s method, which uses the Hessian matrix) converged incredibly fast but required calculating and inverting massive matrices. For a neural network with a million parameters, the Hessian is a million-by-million matrix. Impossible to store, let alone invert. First-order methods (Vanilla GD) were cheap (O(N) memory) but agonizingly slow, taking millions of steps to navigate narrow valleys.

The Emergence: In 1964, Boris Polyak published the “Heavy Ball Method.” The engineering insight was brilliant: Instead of paying the computational cost of calculating the curvature (Hessian), let’s just simulate a physical system that naturally responds to curvature. By adding a momentum term, the optimizer behaves like a heavy ball rolling down a hill. Physics dictates that a heavy ball naturally builds speed in straight valleys and loses energy bouncing off walls. This gave second-order convergence speeds at first-order computational costs.

Intuitive Explanation

Imagine you are trying to walk down to the bottom of a steep, narrow canyon.

Vanilla Gradient Descent: You are a blindfolded person. You feel the slope under your feet. If the left wall is steep, you take a huge step right. Now the right wall is steep, so you take a huge step left. You end up violently bouncing wall-to-wall, slowly inching forward. You are exhausted (wasted compute) and moving slowly.

Gradient Descent with Momentum: You are a heavy bowling ball. When you hit the left wall, the slope pushes you right. But because you are heavy, you don’t just stop when the slope flattens out at the center of the canyon. Your inertia carries you across the center, up the right wall a bit, but not as high as before. Meanwhile, along the length of the canyon (the gentle slope downward), every bounce adds to your forward speed. You smoothly accelerate down the canyon floor, ignoring the walls.

Momentum translates the geometric shape of the loss landscape into physical forces.

Visual Understanding

Draw a 2D contour plot of a loss landscape (concentric ellipses, much longer vertically than horizontally—this is a ravine).

Without Momentum: 1. Start at the top left. 2. Draw an arrow pointing straight down (towards the center, steep gradient). Move there. 3. Now you are on the right side of the ravine. Draw an arrow pointing left. Move there. 4. The path looks like extreme, jagged zigzags. The progress down the ravine is microscopic compared to the distance traveled across it.

With Momentum (\(\beta = 0.9\)): 1. Start at the top left. Draw a small arrow down. The velocity vector (\(v\)) points down. Move. 2. You are now on the right. The gradient arrow points sharply left. But the velocity vector is a weighted average: 90% of the old “down” vector + 10% of the new “left” vector. 3. The resulting velocity arrow points down-left. You move down and slightly left. 4. Next step: 90% down-left + 10% right = mostly down. 5. The Visual Result: The path smoothly unwinds into a gentle, widening curve that plunges straight down the center of the ravine.

Mathematical Foundations

Let’s formalize the physics. We introduce a new state variable: Velocity (\(\mathbf{v}\)).

Step 1: Update the Velocity \[\mathbf{v}_t = \beta \mathbf{v}_{t-1} + (1-\beta) \nabla_\theta \mathcal{L}_t\]

\(\mathbf{v}_t\): The velocity vector at step \(t\). This is the direction and magnitude we will actually move.
\(\beta\): The friction/momentum coefficient (typically \(0.9\)).
\(\beta \mathbf{v}_{t-1}\): The inertia. We keep 90% of our previous speed.
\((1-\beta)\): A scaling factor. We add 10% of the new gradient. (Note: Some frameworks like PyTorch SGD omit this \((1-\beta)\) scaling and just use \(+ \nabla L\). We will cover this in Implementation).
\(\nabla_\theta \mathcal{L}_t\): The current gradient.

Step 2: Update the Parameters \[\theta_{t+1} = \theta_t - \eta \mathbf{v}_t\] * We step using the velocity, not the raw gradient.

The Exponentially Weighted Moving Average (EWMA) If we unroll the velocity equation over 3 steps: \(\mathbf{v}_3 = (1-\beta)\mathbf{g}_3 + \beta(1-\beta)\mathbf{g}_2 + \beta^2(1-\beta)\mathbf{g}_1\) * The gradient from 1 step ago has weight \(0.1\). * The gradient from 2 steps ago has weight \(0.09\) (\(0.9 \times 0.1\)). * The gradient from 3 steps ago has weight \(0.081\). This is a decaying exponential. With \(\beta=0.9\), the velocity is effectively an average of the last \(\frac{1}{1-\beta} = 10\) gradients.

Nesterov Accelerated Gradient (NAG) Standard momentum has a flaw: it builds up speed, but it’s “looking backward.” By the time it takes a step, it might have already built too much speed and overshoot the minimum.

Nesterov fixes this with a simple trick: Look ahead before you compute the gradient. \[\mathbf{v}_t = \beta \mathbf{v}_{t-1} + \eta \nabla_\theta \mathcal{L}(\theta_t - \beta \mathbf{v}_{t-1})\] * Instead of computing the gradient at our current position \(\theta_t\), we compute it at \(\theta_t - \beta \mathbf{v}_{t-1}\) (where we would end up if we just followed inertia). * Why? If the minimum is nearby, the look-ahead gradient will point backward, acting as a brake. It provides anticipatory correction.

Worked Numerical Example

Let’s manually compute 3 steps of Momentum in a 2D parameter space. * Learning rate \(\eta = 1.0\) * Momentum \(\beta = 0.5\) (Using the scaled version with \((1-\beta)\) for clean math) * Initial position \(\theta_0 = [0, 0]\) * Initial velocity \(\mathbf{v}_0 = [0, 0]\)

Iteration 1: * Current gradient \(\nabla L_1 = [4, 0]\) (Steep slope in x-direction) * Update velocity: \(\mathbf{v}_1 = 0.5[0, 0] + 0.5[4, 0] = \mathbf{[2, 0]}\) * Update params: \(\theta_1 = [0, 0] - 1.0[2, 0] = \mathbf{[-2, 0]}\)

Iteration 2: * Current gradient \(\nabla L_2 = [-1, 6]\) (Slope reverses in x, strong slope in y) * Update velocity: \(\mathbf{v}_2 = 0.5[2, 0] + 0.5[-1, 6] = [1, 0] + [-0.5, 3] = \mathbf{[0.5, 3]}\) * Update params: \(\theta_2 = [-2, 0] - 1.0[0.5, 3] = \mathbf{[-2.5, -3]}\)

Iteration 3: * Current gradient \(\nabla L_3 = [-2, 6]\) (Slope continues) * Update velocity: \(\mathbf{v}_3 = 0.5[0.5, 3] + 0.5[-2, 6] = [0.25, 1.5] + [-1, 3] = \mathbf{[-0.75, 4.5]}\) * Update params: \(\theta_3 = [-2.5, -3] - 1.0[-0.75, 4.5] = \mathbf{[-1.75, -7.5]}\)

Analysis: Look at the X-dimension (the oscillating direction): Gradients were: \(4 \to -1 \to -2\). Velocities were: \(2 \to 0.5 \to -0.75\). The velocity is heavily dampened and lagging behind the raw gradient. The violent oscillation is smoothed out.

Look at the Y-dimension (the consistent downward direction): Gradients were: \(0 \to 6 \to 6\). Velocities were: \(0 \to 3 \to 4.5\). The velocity is accumulating. We are moving faster and faster in the correct direction.

Computational Interpretation

What happens on the hardware when you add momentum?

Memory Representation: Vanilla GD requires 1 state tensor: Parameters (\(\theta\)). Momentum requires 2 state tensors: Parameters (\(\theta\)) and Velocity (\(\mathbf{v}\)). * Memory Cost: It strictly doubles the optimizer state memory. For a 1 Billion parameter model in FP32, this is an extra 4 GB of VRAM just for the velocity buffers.

Compute Operations (per step): 1. Scaled addition: \(\mathbf{v} = \beta \cdot \mathbf{v} + (1-\beta) \cdot \text{grad}\) (Element-wise multiply + add) 2. Scaled subtraction: \(\theta = \theta - \eta \cdot \mathbf{v}\) (Element-wise multiply + subtract)

Complexity: * Time: \(O(N)\) where \(N\) is the number of parameters. This is identical asymptotically to Vanilla GD. The constant factor is roughly 3x (a few extra element-wise operations). * Bottleneck: Memory bandwidth, not compute. The GPU must read \(\theta\), read \(\mathbf{v}\), read \(\text{grad}\), and write back \(\theta\) and \(\mathbf{v}\). Optimizer steps are notoriously memory-bound, which is why frameworks fuse them into a single CUDA kernel to avoid writing intermediate results to global memory.

Implementation Perspective

1. NumPy (From Scratch)

We will use the mathematically pure \((1-\beta)\) formulation.

import numpy as np

def sgd_momentum_step(theta, grad, velocity, lr=0.01, beta=0.9):
    """
    theta: Current parameters (e.g., weights)
    grad: Gradient of loss w.r.t theta
    velocity: Previous velocity state
    """
    # 1. Update velocity (EWMA of gradients)
    # Shape: matches theta exactly (element-wise)
    velocity = beta * velocity + (1.0 - beta) * grad
    
    # 2. Update parameters using velocity
    theta = theta - lr * velocity
    
    return theta, velocity

# Dummy data: 3x3 weight matrix
theta = np.random.randn(3, 3)
v = np.zeros_like(theta) # Initialize velocity to 0!

grad1 = np.random.randn(3, 3) # Simulated gradient
theta, v = sgd_momentum_step(theta, grad1, v)

2. PyTorch (Production Standard)

CRITICAL ENGINEERING NOTE: Look closely at PyTorch’s torch.optim.SGD. By default, it does not use the \((1-\beta)\) scaling we used above. It uses \(v_t = \text{mom} \cdot v_{t-1} + \text{grad}\). If you set \(\text{momentum}=0.9\), the effective window is much larger than 10 steps, and you usually need a slightly lower learning rate than the theoretical formula to compensate for the un-scaled gradient injections.

import torch
import torch.nn as nn
import torch.optim as optim

# A simple layer
layer = nn.Linear(10, 2)

# Initialize optimizer
# dampening=0 (default) means NO (1-beta) scaling. 
# nesterov=True enables the look-ahead trick.
optimizer = optim.SGD(
    layer.parameters(), 
    lr=0.01, 
    momentum=0.9, 
    nesterov=True
)

# Simulate a training loop step
optimizer.zero_grad()
dummy_input = torch.randn(4, 10)
out = layer(dummy_input)
loss = out.sum() # Dummy loss

loss.backward() # Computes gradients

# Inside this call, PyTorch does:
# if nesterov: grad = grad + momentum * v_old
# v_new = momentum * v_old + grad
# param = param - lr * v_new
optimizer.step()

Where It Appears In AI

AI Component	Exact Usage
Computer Vision (ResNet/ViT)	Standard training uses `SGD(momentum=0.9)`. It is the gold standard because it generalizes slightly better than adaptive methods on image data.
Adam Optimizer (LLMs, NLP)	The \(m\) variable in Adam is exactly momentum (an EWMA of gradients). Adam is essentially “Momentum + RMSprop”.
BatchNorm / LayerNorm	While not optimization momentum, the running mean/variance tracked during training uses an identical EWMA formula (`running_mean = 0.1 * batch_mean + 0.9 * running_mean`).
Reinforcement Learning (PPO)	Policy gradient methods have massive variance. Momentum (via Adam) is strictly required to smooth out the noisy gradients and prevent the policy from collapsing.
Contrastive Learning (SimCLR)	Because of the massive batch sizes (e.g., 4096), gradients are noisy. High momentum (0.9 to 0.99) is critical to push the encoder through the noise.

Deep Dive Into Real Models

In a Transformer (AdamW)

When you train GPT-3 or LLaMA, you use AdamW. Let’s look at the exact update equations for a single parameter \(\theta\):

Momentum step: \(m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t\) (Usually \(\beta_1 = 0.9\))
Variance step: \(v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2\) (Usually \(\beta_2 = 0.999\))
Bias correction: \(\hat{m}_t = m_t / (1 - \beta_1^t)\)
Update: \(\theta_t = \theta_{t-1} - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}\)

The Role of Momentum here: In LLMs, the loss landscape is incredibly non-stationary (the optimal weights for token 1 are different from token 10,000). The momentum term \(m_t\) acts as a low-pass filter. It ignores the high-frequency noise of individual token predictions and locks onto the low-frequency “true” gradient direction. Without the \(m_t\) term (using just the variance term), LLM training would violently diverge.

In a CNN (SGD + Nesterov)

For a ResNet-50 on ImageNet: The loss landscape features long, flat plateaus. Vanilla SGD gets “stuck” on these plateaus because gradients approach zero. Momentum solves this: even if the gradient is \(0.0001\), the velocity might be \(0.5\) from previous steps. The model coasts over the flat plateau purely on inertia, whereas vanilla GD would halt completely.

Failure Modes

1. The Fine-Tuning Trap (Forgetting to Zero Velocity)

Cause: You pre-train a model, save it, then load it for fine-tuning on a new dataset, but you forget to re-initialize the optimizer state (or you load the old optimizer checkpoint).
Symptom: The model immediately diverges, loss spikes to NaN, or performance is terrible.
Why? The velocity buffer \(\mathbf{v}\) still contains the momentum from the old dataset. The first step of fine-tuning will violently throw the parameters in the direction of the old task.
Solution: When fine-tuning, ALWAYS instantiate a fresh optimizer (optim.SGD(model.parameters(), ...)) so the velocity buffers start at zero.

2. Overshooting Minima (Too much Momentum)

Cause: Setting \(\beta = 0.99\) or higher on a small, simple dataset (like MNIST).
Symptom: Loss decreases, then oscillates wildly and never converges, bouncing around the minimum.
Solution: Reduce \(\beta\) to 0.9, or use Nesterov momentum which provides an automatic “braking” mechanism near the bottom of minima.

3. PyTorch Scaling Confusion

Cause: Copying math from a paper (which uses \(1-\beta\) scaling) directly into PyTorch SGD.
Symptom: Training is much slower than expected.
Solution: Remember PyTorch SGD lacks the \((1-\beta)\) scaling. To match papers, you often need to scale your learning rate down by roughly \((1-\beta)\), or use a custom optimizer.

Engineering Insights

1. Memory is the Bottleneck In distributed training (DDP/FSDP), optimizer states account for the majority of VRAM. For a 7B parameter model: * Weights (FP16): 14 GB * Gradients (FP16): 14 GB * Adam States (\(m\) and \(v\) in FP32): 56 GB Momentum is elegant, but it is incredibly expensive at scale. This is why techniques like 8-bit Optimizers (bitsandbytes) were invented—to quantize the momentum buffers to save 75% of that 56 GB.

2. Kernel Fusion Never write your own momentum update in a PyTorch training loop using param.data -= lr * v. This forces Python to issue multiple read/write commands to the GPU. PyTorch’s native optimizer.step() runs a single fused C++/CUDA kernel that reads the gradient, updates the velocity, and updates the weight in one fast pass without round-tripping to global memory.

Interview Questions

Beginner

Q: What is the primary physical analogy for momentum in gradient descent? A: A heavy ball rolling down a hill. It accumulates inertia in consistent downward directions, allowing it to pass over small bumps, while its mass dampens oscillations when it hits the sides of a valley.

Intermediate

Q: In the equation \(v_t = \beta v_{t-1} + (1-\beta)g_t\), what does setting \(\beta = 0.9\) actually mean in terms of the history of the gradients? A: It means the velocity is an Exponentially Weighted Moving Average. The current gradient gets a weight of \(0.1\), the previous gets \(0.09\), the one before that \(0.081\), etc. Effectively, the velocity represents the average direction of the last \(1 / (1 - 0.9) = 10\) gradients, with the most recent gradients having the highest influence.

Advanced

Q: Explain the difference between Standard Momentum and Nesterov Accelerated Gradient (NAG). Why does NAG theoretically converge faster? A: Standard momentum computes the gradient at the current position \(\theta_t\), then adds it to the historical velocity. NAG computes the gradient at a “look-ahead” position \(\theta_t - \beta v_{t-1}\) (where momentum would take us). NAG converges faster because if the look-ahead position has crossed a minimum, the gradient computed there will point backward, acting as an anticipatory brake and reducing overshooting.

Research

Q: Modern vision models often generalize better with SGD+Momentum than with Adam, even though Adam converges faster. Why might this be? A: This is an active area of research (often called “Adam vs SGD generalization gap”). One prominent theory is that the adaptive learning rate of Adam causes the optimizer to find sharp, narrow minima in the loss landscape. SGD with momentum, due to its uniform learning rate and inertia, tends to converge to flatter, wider minima, which have better out-of-distribution generalization properties.

Research Perspective

Current Limitations: Momentum assumes the loss landscape is relatively smooth and that past gradients are a good predictor of future gradients. In highly non-convex landscapes (like those in RL or diffusion models), this assumption breaks down. Momentum can actually trap the optimizer in saddle points because the accumulated velocity points directly into the flat region, and the gradients perpendicular to the velocity are too small to escape.

Modern Directions: 1. Lion (EvoLved Sign Momentum): Recent research (by Google) discovered that you can entirely throw away the magnitude of the gradients and the velocity. Lion only uses the sign of the momentum update (\(\text{sign}(\beta v + g)\)). This drastically reduces compute and memory while matching AdamW performance. 2. Scheduled Momentum: Instead of a constant \(\beta=0.9\), recent papers explore cyclical momentum (increasing and decreasing \(\beta\) in a sinusoidal pattern) to actively push the optimizer out of sharp minima.

Connections

Prerequisites: * Gradient Descent (Base update mechanism) * Exponentially Weighted Averages (Signal processing concept)

Enables: * Fast convergence in ill-conditioned spaces * Plateau crossing * Noise dampening

Used by: * SGD Optimizer (Vision standard) * Adam / AdamW Optimizer (The \(m\) term) * BatchNorm (Running stats tracking)

Extended by: * Nesterov Accelerated Gradient (Look-ahead) * AdamW (Decoupled weight decay combined with momentum) * Lion (Sign-momentum)

Related concepts: * Learning Rate Scheduling (Often interacts with momentum; e.g., reducing LR while maintaining momentum) * Gradient Clipping (Often applied before momentum accumulation in RL to prevent velocity explosions)

Final Mental Model

To permanently remember Momentum, visualize the “Blindfolded Bowling Ball.”

Vanilla Gradient Descent is a blindfolded person taking careful, equal-sized steps based solely on the slope under their feet. They zigzag wildly in canyons and stop dead on flat ground.

Momentum replaces the person with a heavy bowling ball. You give it a push (the gradient). It starts rolling. As it rolls down a straight canyon, it accelerates (accumulates velocity). When it hits a wall, it doesn’t instantly reverse like the person; its heavy mass dampens the bounce, and it smoothly redirects its energy down the canyon floor.

The \(\beta\) parameter is simply the friction of the floor. High \(\beta\) (0.99) is an icy floor—the ball rolls forever. Low \(\beta\) (0.5) is a carpet—the ball stops quickly. Deep learning engineers usually set it to 0.9: a polished wood floor.

Neural Nexus

Search This Blog