Gradient Descent and Optimizer Geometry: Momentum, Adam, and Loss Surfaces

The same model on the same data can behave very differently depending on the optimizer. That is hard to see from the formulas and obvious from the paths they trace on a loss surface: plain gradient descent oscillates across the steep direction and barely moves along the shallow one; momentum cancels the oscillation; Adam gives each parameter its own step size.

The example below uses a two-dimensional quadratic function, because its elliptical contours make the troublesome property — steep in one direction, flat in another — directly visible. Start point and target are fixed, so the differences between the three optimizers come entirely from how each handles curvature, gradient history and scale. The first step is worked by hand; the remaining trajectories come from NumPy.

1. The Geometry of the Loss Surface and Pathological Curvature

In deep learning, the loss surface is rarely isotropic (perfectly spherical). Instead, it is highly ill-conditioned, filled with ravines and narrow valleys. Let’s look at a canonical function demonstrating this:

L(x, y) = 1/2 * (8x^2 + y^2) + 0.8xy
grad L = [8x + 0.8y, y + 0.8x]

The Hessian matrix of this function has disparate eigenvalues. The surface is steep in the x direction (high curvature) and flatter in the y direction (low curvature). This creates a “pathological curvature” problem. A plain Gradient Descent step will bounce back and forth across the steep ravine, making excruciatingly slow progress along the flat bottom towards the minimum.

2. Hand Calculate The First Gradient Descent Step

Starting from (2.2, -2.0), the gradient is:

grad = [8*2.2 + 0.8*(-2.0), -2.0 + 0.8*2.2]
     = [16.0, -0.24]

With a learning rate of 0.08:

x_new = 2.2  - 0.08 * 16.0  = 0.92
y_new = -2.0 - 0.08 * -0.24 = -1.9808

The step taken in the x direction is massive compared to y, purely because the gradient is overwhelmingly larger in x. The lab output confirms it: step 1 for gradient descent is x=0.920000, y=-1.980800, and the loss drops from 17.840000 to 3.889516. However, if the learning rate were just slightly higher, the step in x would overshoot the valley, leading to divergence.

3. What Momentum And Adam Change

Plain gradient descent is memoryless; it only uses the current gradient. This leads to the aforementioned oscillation.

Momentum accumulates a velocity from previous gradients. Think of a heavy ball rolling down a hill. The alternating gradients in the steep x direction cancel each other out, while the consistent gradients in the flat y direction accumulate, accelerating the optimizer toward the minimum.

v_t = beta * v_{t-1} + grad_t
theta_t = theta_{t-1} - lr * v_t

Adam (Adaptive Moment Estimation) goes a step further by maintaining both first (mean) and second (uncentered variance) moments of the gradients. It dynamically scales the learning rate for each parameter individually. By dividing the update by the square root of the accumulated squared gradients, Adam normalizes the step sizes. It essentially forces the optimizer to take larger steps in flat directions and smaller steps in steep directions.

m_t = beta1 * m_{t-1} + (1-beta1) * grad_t
v_t = beta2 * v_{t-1} + (1-beta2) * grad_t^2
theta_t = theta_{t-1} - lr * m_hat / (sqrt(v_hat) + eps)

4. Optimizer Anatomy: Visualized

How do we decide which optimizer to use? The diagram below visualizes the architectural flow of these optimization algorithms.

graph TD
    A[Compute Gradient] --> B{Need history?}
    B -->|No| C[Vanilla SGD]
    B -->|Yes| D{Adaptive Scale?}
    D -->|No| E[SGD with Momentum]
    D -->|Yes| F[Compute 1st & 2nd Moments]
    F --> G[Bias Correction]
    G --> H[Adam / AdamW]
    C --> I[Apply Parameter Update]
    E --> I
    H --> I

5. Practical Python Implementation

To truly grasp these algorithms, we should build them from scratch. Here is a NumPy implementation comparing SGD, Momentum, and Adam on our quadratic surface.

import numpy as np

def grad(theta):
    x, y = theta
    return np.array([8.0 * x + 0.8 * y, y + 0.8 * x])

def run_optimizer(optimizer_name, theta_init, lr=0.08, steps=50):
    theta = np.array(theta_init)
    
    # Optimizer state
    v = np.zeros_like(theta)
    m = np.zeros_like(theta)
    beta1, beta2, eps = 0.9, 0.999, 1e-8
    
    trajectory = [theta.copy()]
    
    for t in range(1, steps + 1):
        g = grad(theta)
        
        if optimizer_name == 'SGD':
            theta -= lr * g
            
        elif optimizer_name == 'Momentum':
            v = 0.9 * v + lr * g
            theta -= v
            
        elif optimizer_name == 'Adam':
            m = beta1 * m + (1 - beta1) * g
            v = beta2 * v + (1 - beta2) * (g ** 2)
            
            # Bias correction
            m_hat = m / (1 - beta1 ** t)
            v_hat = v / (1 - beta2 ** t)
            
            theta -= lr * m_hat / (np.sqrt(v_hat) + eps)
            
        trajectory.append(theta.copy())
        
    return np.array(trajectory)

# Run test from the start point
traj_sgd = run_optimizer('SGD', [2.2, -2.0])
traj_adam = run_optimizer('Adam', [2.2, -2.0], lr=0.5)
print(f"Final Adam position: {traj_adam[-1]}")

6. What The Animation Shows

Watch whether the steep direction oscillates and whether the flatter direction progresses too slowly. Notice how Momentum swings widely like a pendulum before settling, while Adam cuts a much more direct, controlled path toward the minimum, seamlessly adjusting to the varying curvature.

7. Personal Experience / Engineer’s Perspective

In practice, the elegant math of optimizers runs into harsh hardware and systems realities. Here are a few things I’ve learned from the trenches of training large models:

• Memory Constraints: Adam is highly effective but extremely memory-hungry. Vanilla SGD only requires memory for the parameters and gradients. Adam requires storing the moving average of gradients (first moment) and the moving average of squared gradients (second moment). This essentially triples the memory footprint of your optimizer state. When training large LLMs on GPUs with limited VRAM, this is often the bottleneck, prompting engineers to use memory-efficient variants like Adafactor or 8-bit Adam.
• Weight Decay Pitfalls: There is a notorious difference between Adam and AdamW. Standard Adam applies L2 regularization to the gradient before the adaptive scaling. This inadvertently scales down the penalty for weights with high gradient variance, defeating the purpose of weight decay. Always use AdamW for Transformer architectures, which decouples weight decay from the gradient update.
• Warmup is Mandatory for Adam: Because Adam uses moving averages, the variance term (second moment) is initialized to zero and can be wildly inaccurate in the first few steps, leading to massive, destabilizing updates. A learning rate warmup (starting the LR near zero and scaling up over thousands of steps) prevents the model from blowing up early in training.

8. Practical Notes

• Plot training and validation loss before changing optimizers. Many training failures are optimizer-path problems rather than architecture problems.
• The learning rate usually matters more than the optimizer name. A well-tuned SGD with Momentum can often match or beat Adam in generalization, especially in computer vision (ResNets).
• Adam can still diverge when the learning rate is too high. Do not treat it as a silver bullet that requires no tuning.
• For noisy loss curves, try lowering the learning rate or adding warmup.

9. Optimization Trace Audit Table

An optimizer experiment should not report only the final loss. To judge whether the path is trustworthy, record coordinates, gradients, state variables, and failure modes together. The audit table below turns “it seems to converge” into reproducible numerical evidence.

The next article moves into convolution and shows how local image operations become matrix computations.