Gradient Descent with Momentum

As we saw, standard Gradient Descent updates weights based solely on the gradient calculated at the current position. While this works, it can lead to inefficient training, especially in certain types of loss landscapes. Imagine a narrow ravine or valley in the loss surface. Gradient Descent tends to oscillate back and forth across the steep walls of the ravine, making slow progress along the bottom towards the minimum. Similarly, it can get stuck in local minima or slow down considerably on saddle points.

The Intuition: Adding Inertia

To address these issues, we can introduce the concept of momentum, borrowing intuition from physics. Imagine a ball rolling down a hill. Instead of just calculating the steepest slope at its current location and moving slightly in that direction (like standard GD), the ball has momentum. Its current velocity influences its movement in the next step. If it's already moving quickly in a certain direction, it tends to keep going that way, even if the immediate slope changes slightly. This momentum helps it:

1. Accelerate down consistent slopes.
2. Dampen oscillations by averaging out opposing gradient directions over time.
3. Overshoot small local minima or plateaus due to its built-up speed.

Gradient Descent with Momentum applies this physical analogy to the optimization process. It introduces a "velocity" vector, $v$, which accumulates an exponentially decaying moving average of past gradients. The weight update then considers both the current gradient and this velocity.

How Momentum Works

The update rule for Gradient Descent with Momentum involves two steps at each iteration $t$:

1. Update the velocity vector $v_t$:

   $$v_t = \beta v_{t-1} + \eta \nabla L(w_{t-1})$$

   Here:

   • $v_t$ is the velocity vector at time step $t$.
   • $v_{t-1}$ is the velocity vector from the previous time step (initialized to zero).
   • $\beta$ is the momentum coefficient, a hyperparameter usually between 0.5 and 0.99 (commonly 0.9). It determines how much past velocity is retained. A higher $\beta$ means past gradients have more influence.
   • $\eta$ is the learning rate.
   • $\nabla L(w_{t-1})$ is the gradient of the loss function with respect to the weights $w$ at the previous time step $t-1$.

   This step essentially calculates the new velocity by taking a fraction ($\beta$) of the old velocity and adding the scaled current gradient ($\eta \nabla L(w_{t-1})$). If the current gradient points consistently in the same direction as the previous velocity, the velocity magnitude increases. If the gradient direction oscillates, the velocity tends to be dampened because opposing gradient terms partially cancel out over time.

2. Update the weights $w_t$:

   $$w_t = w_{t-1} - v_t$$

   The weights are updated by moving in the direction of the newly computed velocity vector $v_t$.

By incorporating the velocity $v_t$, which holds information about recent gradient history, the updates become smoother and faster, particularly along directions where the gradient is consistent.

Visualizing the Effect of Momentum

Consider optimizing a function with a narrow valley. Standard Gradient Descent might bounce between the walls, while Momentum takes a more direct route.

Optimization paths on a sample loss surface. Standard Gradient Descent (pink) oscillates across the narrow valley, making slow progress towards the minimum (near origin). Momentum (blue) dampens these oscillations and accelerates along the valley floor, taking a more direct path.

The Momentum Coefficient $\beta$

The momentum coefficient $\beta$ controls the influence of past gradients.

• If $\beta = 0$, the update rule reduces to standard Gradient Descent ($v_t = \eta \nabla L(w_{t-1})$, $w_t = w_{t-1} - \eta \nabla L(w_{t-1})$).
• As $\beta$ approaches 1, past gradients are given more weight, leading to smoother updates but potentially causing the optimizer to overshoot the minimum if the learning rate is also high.

Common values for $\beta$ are 0.9 or even higher, like 0.99. It often works well to start with 0.9 and potentially increase it during training. This parameter, like the learning rate, may require tuning for optimal performance on a specific problem.

Implementation in PyTorch

Using momentum in frameworks like PyTorch is straightforward. When defining the optimizer, you simply specify the momentum parameter.

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

# Assume 'model' is your defined neural network (nn.Module)
# Assume 'learning_rate' is your chosen learning rate (e.g., 0.01)

# Define the optimizer using SGD with momentum
momentum_coefficient = 0.9
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum_coefficient)

# --- Inside your training loop ---
# loss.backward() # Compute gradients

# optimizer.step() # Update weights using SGD with momentum
# optimizer.zero_grad() # Reset gradients for next iteration

By simply adding the momentum argument to the optim.SGD constructor, the optimizer automatically implements the velocity calculation and weight update steps described earlier.

Gradient Descent with Momentum is a significant improvement over standard gradient descent, often leading to faster convergence and better navigation of complex loss surfaces. While it helps overcome many challenges, further refinements led to the development of adaptive methods like RMSprop and Adam, which we will explore next.