SGD with Momentum: Accelerating Convergence

While Stochastic Gradient Descent (SGD) offers computational advantages over batch gradient descent, its updates can be noisy, and its convergence can be slow, especially when navigating elongated valleys or ravines in the loss landscape. Imagine the loss surface as a hilly terrain. Vanilla SGD takes steps based only on the gradient at the current point, which can lead to erratic zig-zagging if the slope changes rapidly in different directions. This is particularly problematic in narrow ravines where the gradient points steeply across the valley but only gently along its bottom. SGD might bounce back and forth across the narrow axis, making very slow progress towards the minimum along the valley floor.

To address this, we can incorporate the concept of momentum, drawing inspiration from physics. Think of a ball rolling down a hill. It doesn't just move based on the slope at its current position; it also has momentum built up from its previous movement. This momentum helps it smooth out its path, roll over small bumps, and accelerate faster down consistent slopes.

SGD with Momentum applies a similar idea to parameter updates. Instead of using only the current gradient to update the parameters, we maintain a "velocity" vector, which is essentially an exponentially decaying moving average of past gradients. This velocity term represents the accumulated "momentum" of the parameter updates.

How Momentum Works

At each step $t$, we update the velocity $v_t$ and then use this velocity to update the parameters $\theta_t$. The process looks like this:

1. Compute the current mini-batch gradient: $\nabla J(\theta_{t-1})$
2. Update the velocity: The new velocity $v_t$ is a combination of the previous velocity $v_{t-1}$ (decayed by a factor $\beta$) and the current gradient (scaled by the learning rate $\eta$). $$v_t = \beta v_{t-1} + \eta \nabla J(\theta_{t-1})$$
3. Update the parameters: The parameters are updated by moving in the direction of the current velocity. $$\theta_t = \theta_{t-1} - v_t$$

Here:

• $\theta$ represents the model parameters.
• $v$ is the velocity vector, typically initialized to zeros.
• $\beta$ is the momentum coefficient, a hyperparameter usually between 0 and 1 (commonly set to 0.9). It controls how much influence past gradients have on the current update. A higher $\beta$ means past gradients contribute more.
• $\eta$ is the learning rate.
• $\nabla J(\theta_{t-1})$ is the gradient of the loss function $J$ with respect to the parameters $\theta$, computed using the parameters from the previous step $t-1$.

Benefits of Momentum

Incorporating momentum brings several advantages:

• Accelerated Convergence: In areas where the gradient consistently points in the same direction (like along the bottom of a ravine), the velocity term accumulates, leading to larger steps and faster progress towards the minimum.
• Dampened Oscillations: When gradients oscillate back and forth (like across a narrow ravine), the momentum term averages these opposing gradients. Positive and negative components tend to cancel out, reducing the zig-zagging behavior and keeping the updates focused on the overall downhill direction.
• Improved Navigation: The accumulated velocity can help the optimizer push past small local minima or flat regions like saddle points, where vanilla SGD might get stuck due to near-zero gradients.

Consider the difference in paths taken by SGD and SGD with Momentum in a typical loss landscape ravine:

Comparison of optimization paths for SGD (red) and SGD with Momentum (blue) on a hypothetical loss surface. Momentum takes a more direct path towards the minimum, avoiding the oscillations seen with standard SGD.

The Momentum Hyperparameter ($\beta$)

The momentum coefficient $\beta$ determines the "memory" of the optimizer.

• If $\beta = 0$, the update rule reduces back to standard SGD (velocity depends only on the current gradient).
• As $\beta$ approaches 1, past gradients are given more weight, leading to smoother updates and potentially faster acceleration, but also increasing the risk of overshooting the minimum if the learning rate is too high.

Common starting values for $\beta$ are 0.9 or sometimes 0.99. Like the learning rate $\eta$, the optimal value for $\beta$ is problem-dependent and often needs to be tuned through experimentation. It's important to remember that $\beta$ and $\eta$ interact, so adjusting one may require adjusting the other.

Implementation

Most deep learning frameworks provide straightforward implementations of SGD with momentum. For example, in PyTorch, you can enable momentum by simply specifying the momentum argument when creating an SGD optimizer instance:

import torch.optim as optim

# Assume 'model_parameters' holds your network's parameters
# Set learning rate and momentum coefficient
learning_rate = 0.01
momentum_beta = 0.9

optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum_beta)

# --- Inside your training loop ---
# optimizer.zero_grad()
# # Forward pass, calculate loss
# loss = criterion(outputs, labels)
# # Backward pass
# loss.backward()
# # Update weights
# optimizer.step()

By accumulating a velocity based on past gradients, SGD with Momentum offers a significant improvement over vanilla SGD, often leading to faster convergence and more stable training, particularly for the complex, non-convex loss landscapes encountered in deep learning. It forms the basis for many subsequent, more advanced optimization algorithms.