Gradient Descent, Step by Step

The Gradient: Your Compass Through Loss Space

Imagine you are standing on a hillside at night, unable to see the valley below. You cannot see the full landscape, but you can feel the slope beneath your feet—the gradient. At each step, you move downhill in the direction of steepest descent. Repeat this many times, and you eventually reach the valley. This is gradient descent.

Mathematically: The gradient ∇L(w) is a vector of partial derivatives:

∇L(w) = [∂L/∂w₁, ∂L/∂w₂, ..., ∂L/∂wₙ]

Each component tells you how much the loss changes when you nudge that weight. The direction of ∇L points toward steepest increase; therefore, -∇L points toward steepest decrease.

The diagram shows a contour plot of a loss function (darker = higher loss). The path traced by gradient descent zig-zags downward, with each arrow pointing in the direction of steepest descent. The red path on the right diverges because the learning rate is too high; each step overshoots the optimum.

The Mechanics: Gradient and Weight Update

The gradient descent update rule is derived from the first-order Taylor expansion of the loss function around the current weight:

L(w - α*∇L(w)) ≈ L(w) - α*||∇L(w)||² + O(α²)

For small enough α, the first-order term dominates, and we move to a point with lower loss. This justifies the update:

w_new = w_old - learning_rate * ∇L(w)

where learning_rate (α or η) controls step size. The term learning rate controls how aggressively we move in the negative gradient direction. Small steps are safe but slow; large steps are fast but risk overshooting the minimum or diverging entirely.

From Scratch: A Minimal Gradient Descent Implementation

Let's implement gradient descent for a simple linear regression problem:

import numpy as np
import matplotlib.pyplot as plt

# Generate synthetic data: y = 3*x + 2 + noise
np.random.seed(42)
X = np.random.randn(100, 1)
y_true = 3 * X + 2
y = y_true + np.random.randn(100, 1) * 0.5  # Add noise

# Initialize weights randomly
w = np.random.randn(1, 1)
b = np.random.randn(1, 1)
learning_rate = 0.01
n_iterations = 100
losses = []

for iteration in range(n_iterations):
    # Forward pass: compute predictions
    y_pred = X @ w + b

    # Compute loss (Mean Squared Error)
    loss = np.mean((y_pred - y) ** 2)
    losses.append(loss)

    # Backward pass: compute gradients
    n_samples = X.shape[0]
    dw = (2 / n_samples) * (X.T @ (y_pred - y))
    db = (2 / n_samples) * np.sum(y_pred - y)

    # Update weights: move in the opposite direction of the gradient
    w = w - learning_rate * dw
    b = b - learning_rate * db

    if (iteration + 1) % 20 == 0:
        print(f"Iteration {iteration+1}: Loss={loss:.4f}, w={w[0,0]:.4f}, b={b[0,0]:.4f}")

print(f"Final estimate: y ≈ {w[0,0]:.4f}*x + {b[0,0]:.4f}")
print(f"True parameters: y = 3*x + 2")

# Plot loss over iterations
plt.figure(figsize=(8, 4))
plt.plot(losses, 'b-', linewidth=2)
plt.xlabel('Iteration')
plt.ylabel('Loss (MSE)')
plt.title('Gradient Descent: Loss Decreases Over Time')
plt.grid(True, alpha=0.3)
plt.show()

The key steps:

1. Forward pass: Compute predictions y_pred = X*w + b.
2. Loss computation: MSE = mean((y_pred - y)^2).
3. Backward pass: Compute dw and db (partial derivatives of loss with respect to w and b).
4. Weight update: Subtract learning_rate * gradient from the current weights.

After 100 iterations, the learned weights should be close to the true values (w ≈ 3, b ≈ 2).

The Learning Rate: Too High, Too Low, Just Right

The learning rate is perhaps the single most important hyperparameter. Here is what happens at different rates:

# Simulate gradient descent with different learning rates
def gradient_descent_with_lr(learning_rate, n_iterations=100):
    """Run gradient descent with a fixed learning rate."""
    w = np.random.randn(1, 1)
    b = np.random.randn(1, 1)
    losses = []

    for iteration in range(n_iterations):
        y_pred = X @ w + b
        loss = np.mean((y_pred - y) ** 2)
        losses.append(loss)

        if np.isnan(loss):  # Diverged
            losses.extend([np.nan] * (n_iterations - iteration - 1))
            break

        n_samples = X.shape[0]
        dw = (2 / n_samples) * (X.T @ (y_pred - y))
        db = (2 / n_samples) * np.sum(y_pred - y)

        w = w - learning_rate * dw
        b = b - learning_rate * db

    return losses

# Test three learning rates
lrs = [0.001, 0.01, 0.1]  # Too slow, just right, too fast
fig, axes = plt.subplots(1, 3, figsize=(14, 4))

for idx, lr in enumerate(lrs):
    losses = gradient_descent_with_lr(lr, n_iterations=200)
    axes[idx].plot(losses, linewidth=2)
    axes[idx].set_title(f'Learning Rate = {lr}')
    axes[idx].set_xlabel('Iteration')
    axes[idx].set_ylabel('Loss')
    axes[idx].grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

Expected behavior:

• lr = 0.001 (too low): Loss decreases very slowly. Convergence requires many iterations—inefficient for large datasets or complex models.
• lr = 0.01 (just right): Loss decreases smoothly and converges in a reasonable number of iterations.
• lr = 0.1 (too high): Loss bounces around or diverges, weights oscillate wildly, eventually reaching NaN.

The optimal learning rate depends on the problem's geometry (the curvature of the loss surface). A heuristic: start with lr = 0.01 or 0.001, train for a few epochs, and adjust based on whether loss is decreasing steadily.

Stochastic Gradient Descent (SGD): Noisy but Faster

Computing the gradient over the entire dataset is expensive for large datasets. Stochastic Gradient Descent computes the gradient on a small random subset (a mini-batch) and updates weights immediately. This is noisier but much faster.

def sgd_step(X_batch, y_batch, w, b, learning_rate):
    """Perform one SGD step on a mini-batch."""
    y_pred = X_batch @ w + b
    n_samples = X_batch.shape[0]
    dw = (2 / n_samples) * (X_batch.T @ (y_pred - y_batch))
    db = (2 / n_samples) * np.sum(y_pred - y_batch)
    w = w - learning_rate * dw
    b = b - learning_rate * db
    return w, b

# Mini-batch SGD
w = np.random.randn(1, 1)
b = np.random.randn(1, 1)
learning_rate = 0.01
batch_size = 16
n_epochs = 10
losses = []

for epoch in range(n_epochs):
    # Shuffle data
    indices = np.random.permutation(X.shape[0])
    X_shuffled = X[indices]
    y_shuffled = y[indices]

    # Process mini-batches
    for i in range(0, X.shape[0], batch_size):
        X_batch = X_shuffled[i:i+batch_size]
        y_batch = y_shuffled[i:i+batch_size]
        w, b = sgd_step(X_batch, y_batch, w, b, learning_rate)

        # Record loss on full dataset
        y_pred_full = X @ w + b
        loss = np.mean((y_pred_full - y) ** 2)
        losses.append(loss)

print(f"Final estimate: y ≈ {w[0,0]:.4f}*x + {b[0,0]:.4f}")

SGD is the basis of training modern neural networks. Instead of computing gradients on all 1 million samples, you compute on a batch of 32 or 64 and update weights. This parallelizes well and introduces noise that can help escape shallow local minima (a small benefit in very deep networks).

Beyond Vanilla Gradient Descent: Momentum and Adam

Vanilla gradient descent can be slow on steep or curved loss surfaces. Momentum accumulates gradients across iterations, building velocity in directions of consistent descent and damping oscillations.

The update rule with momentum is:

v = momentum * v + gradient
w = w - learning_rate * v

The velocity v accumulates gradients; momentum (typically 0.9) controls how much of the previous velocity to retain.

Adam (Adaptive Moment Estimation) goes further. It maintains both first moments (like momentum) and second moments (adaptive per-parameter learning rates). Adam is widely used because it often works well without tuning:

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

# Create a simple neural network
model = nn.Sequential(
    nn.Linear(1, 50),
    nn.ReLU(),
    nn.Linear(50, 1)
)

# Use Adam optimizer
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()

# Convert NumPy data to PyTorch tensors
X_torch = torch.tensor(X, dtype=torch.float32)
y_torch = torch.tensor(y, dtype=torch.float32)

# Training loop
losses_adam = []
for epoch in range(100):
    optimizer.zero_grad()
    y_pred = model(X_torch)
    loss = criterion(y_pred, y_torch)
    loss.backward()
    optimizer.step()
    losses_adam.append(loss.item())

    if (epoch + 1) % 20 == 0:
        print(f"Epoch {epoch+1}: Loss={loss.item():.4f}")

print("\nAdam typically converges faster than vanilla SGD with less hyperparameter tuning.")

Adam's default learning rate (lr=0.001) works well in many cases, removing the need to manually tune the learning rate. For details on how Adam computes adaptive learning rates per parameter, see the original paper (referenced in sources).

Comparison: Optimizers in Practice

| Optimizer | Speed | Simplicity | Hyperparameters | When to use | |-----------|-------|-----------|-----------------|------------| | Vanilla SGD | Slow | Simple | lr | Baseline, well-understood problems | | SGD + Momentum | Medium | Simple | lr, momentum | General-purpose training, CNNs | | Nesterov Momentum | Medium | Simple | lr, momentum | Faster than vanilla momentum | | Adam | Fast | Simple | lr, β₁, β₂ (defaults: 0.9, 0.999) | Most neural networks, default choice | | RMSprop | Fast | Medium | lr, decay | When Adam causes instability (rare) | | Adagrad | Fast | Simple | lr, decay | Sparse feature/gradient problems |

Detailed optimizer comparison:

| Aspect | SGD | Momentum | Adam | |--------|-----|----------|------| | Per-parameter learning rate? | No (fixed α) | No (fixed α) | Yes (adaptive) | | Memory overhead | O(n) weights | O(n) weights + O(n) velocity | O(n) weights + O(n) m + O(n) v | | Convergence on pathological problems | Zigzags, slow | Builds velocity, faster | Smooth, adaptive, fastest | | Generalization (test performance) | Often best | Good | Good (sometimes slightly worse) | | Default hyperparameters | Requires tuning | Requires tuning | Works out-of-box (lr=0.001) |

For new projects, try Adam with default settings (lr=0.001, β₁=0.9, β₂=0.999). If performance is unsatisfactory or you suspect training instability, then investigate other optimizers. For production models where reproducibility and test-set generalization matter most, SGD with momentum (properly tuned) sometimes outperforms Adam because it trains toward flatter minima.

Adaptive Learning Rates: Adam's Insight

Standard gradient descent uses a fixed learning rate for all parameters. However, different parameters may require different learning rates based on the history of gradients they've received.

Adam's key innovation: Maintain two moving averages per parameter:

m_t = β₁ * m_{t-1} + (1 - β₁) * g_t        (first moment, momentum)
v_t = β₂ * v_{t-1} + (1 - β₂) * g_t²      (second moment, variance estimate)

Then compute a per-parameter, adaptive learning rate:

w_t = w_{t-1} - α * m̂_t / (√v̂_t + ε)

where m̂_t and v̂_t are bias-corrected (m_t / (1 - β₁^t), v_t / (1 - β₂^t)), and ε is a small constant (typically 1e-8) for numerical stability.

Intuition: Parameters that have received large gradients (large v_t) are damped; parameters with small gradients are amplified. This adaptive scaling allows different learning rates per parameter without manual tuning.

# Adam implementation from scratch
def adam_step(w, g, m, v, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8, t=1):
    """
    Perform one Adam optimization step.

    Args:
        w: Current weights
        g: Gradient
        m: First moment (momentum) accumulator
        v: Second moment (variance) accumulator
        learning_rate: Step size
        beta1, beta2: Exponential decay rates
        epsilon: Numerical stability constant
        t: Current timestep (for bias correction)

    Returns:
        w_new: Updated weights
        m_new, v_new: Updated accumulators
    """
    # Update biased first and second moments
    m_new = beta1 * m + (1 - beta1) * g
    v_new = beta2 * v + (1 - beta2) * (g ** 2)

    # Bias correction
    m_corrected = m_new / (1 - beta1 ** t)
    v_corrected = v_new / (1 - beta2 ** t)

    # Update weights
    w_new = w - learning_rate * m_corrected / (np.sqrt(v_corrected) + epsilon)

    return w_new, m_new, v_new

# Example: Optimizing a non-convex loss
np.random.seed(42)
w = np.array([0.5, -0.3, 0.8])  # Initial weights
m = np.zeros_like(w)
v = np.zeros_like(w)

# Simulated gradients (time-varying, non-uniform)
gradients = [
    np.array([0.1, -0.5, 0.2]),
    np.array([0.2, -0.4, 0.3]),
    np.array([0.05, -0.6, 0.1]),
]

for t, g in enumerate(gradients, start=1):
    w, m, v = adam_step(w, g, m, v, t=t)
    print(f"Step {t}: w={w}, v_estimate={v}")

Compare this with standard SGD, which would use the same fixed learning rate for all parameters regardless of gradient magnitude history.

Learning Rate Scheduling: Adjusting on the Fly

A fixed learning rate may not be optimal throughout training. Learning rate scheduling decreases the learning rate as training progresses, allowing finer adjustments near the optimum.

from torch.optim.lr_scheduler import StepLR, ReduceLROnPlateau

# Example: Reduce learning rate by a factor of 0.1 every 30 epochs
scheduler = StepLR(optimizer, step_size=30, gamma=0.1)

for epoch in range(100):
    # Training step
    optimizer.zero_grad()
    y_pred = model(X_torch)
    loss = criterion(y_pred, y_torch)
    loss.backward()
    optimizer.step()

    # Step the scheduler
    scheduler.step()

    if (epoch + 1) % 20 == 0:
        print(f"Epoch {epoch+1}: Loss={loss.item():.4f}, LR={scheduler.get_last_lr()[0]:.6f}")

Common schedules:

• StepLR: Multiply learning rate by gamma every N epochs.
• ExponentialLR: Multiply by gamma^epochs.
• ReduceLROnPlateau: Reduce when validation loss plateaus (requires monitoring validation loss).

Learning rate scheduling can accelerate convergence by 10–20% on some problems, but be cautious: poorly chosen schedules can hurt. Start with a constant learning rate, then add scheduling only if convergence is slow.

Batch Size: Another Critical Hyperparameter

Learning rate is not the only hyperparameter affecting convergence. Batch size—the number of samples processed before a weight update—also matters.

Small batches (e.g., 16):

• Noisier gradient estimates (some randomness in descent direction).
• More frequent weight updates (more iterations per epoch).
• Can escape shallow local minima due to noise.
• Higher variance in training loss curve.

Large batches (e.g., 512):

• Cleaner gradient estimates (closer to true gradient).
• Fewer weight updates per epoch (slower training).
• More stable, smoother training curves.
• Risk of getting stuck in sharp minima that don't generalize.

# Experiment with batch sizes
for batch_size in [16, 64, 256]:
    dataset = TensorDataset(X_train_tensor, y_train_tensor)
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    model = model_with_dropout
    optimizer = optim.Adam(model.parameters(), lr=0.01)
    losses = []

    for epoch in range(50):
        for X_batch, y_batch in loader:
            optimizer.zero_grad()
            logits = model(X_batch)
            loss = criterion(logits, y_batch)
            loss.backward()
            optimizer.step()
            losses.append(loss.item())

    print(f"Batch size {batch_size}: Final loss = {losses[-1]:.4f}")

Practical guidance:

• Start with batch_size = 32 or 64.
• Larger batch sizes (128–512) for stable, repeatable results.
• Smaller batch sizes (16–32) for faster convergence on small datasets or when regularization via noise is desired.

Visualizing the Loss Landscape: Understanding Optimization

To build intuition, let's visualize how gradient descent navigates the loss surface:

# Create a 2D loss landscape for visualization
def loss_landscape(w1, w2):
    """Define a simple non-convex loss surface."""
    return (w1 - 2)**2 + (w2 - 1)**2 + 0.5 * np.sin(3 * w1) * np.sin(3 * w2)

# Compute loss over a grid
w1_grid = np.linspace(-2, 5, 100)
w2_grid = np.linspace(-2, 4, 100)
W1, W2 = np.meshgrid(w1_grid, w2_grid)
L = loss_landscape(W1, W2)

# Gradient descent trajectory
def gradient_landscape(w1, w2):
    """Compute gradients numerically."""
    dw1 = 2*(w1 - 2) + 1.5 * np.cos(3*w1) * np.sin(3*w2)
    dw2 = 2*(w2 - 1) + 1.5 * np.sin(3*w1) * np.cos(3*w2)
    return dw1, dw2

# Simulate GD trajectory
w = np.array([0.0, 0.0])
trajectory = [w.copy()]
lr = 0.05
for _ in range(100):
    dw1, dw2 = gradient_landscape(w[0], w[1])
    w = w - lr * np.array([dw1, dw2])
    trajectory.append(w.copy())

trajectory = np.array(trajectory)

# Plot
plt.figure(figsize=(10, 8))
contour = plt.contour(W1, W2, L, levels=20, cmap='viridis')
plt.clabel(contour, inline=True, fontsize=8)
plt.plot(trajectory[:, 0], trajectory[:, 1], 'ro-', markersize=5, linewidth=2, label='GD path')
plt.xlabel('w1')
plt.ylabel('w2')
plt.title('Gradient Descent on a Non-Convex Loss Surface')
plt.legend()
plt.show()

This visualization shows how gradient descent navigates the surface, taking steps proportional to the local slope. The path may not be straight because the loss landscape has curves and ridges.

Common Mistake: Setting the Learning Rate Once and Forgetting It

Many beginners pick a learning rate, run training, and accept whatever results. This is wasteful.

Best practice:

1. Start with a small learning rate (0.001 or 0.0001).
2. Train for a few iterations and check if loss is decreasing. If loss is nearly flat, increase learning rate.
3. If loss oscillates or diverges, decrease learning rate.
4. Once you have a stable, decreasing loss curve, you can fine-tune via validation metrics or learning rate scheduling.
5. Log learning rate and loss at every iteration during development. Plotting these together will reveal if your learning rate is working.

The learning rate is not a one-time choice—it is part of the scientific process of tuning your model. Treat it with the care it deserves.

Advanced: Second-Order Methods (Brief Mention)

While SGD, momentum, and Adam dominate, a class of second-order methods uses the Hessian (matrix of second derivatives) to adapt step sizes more intelligently. Methods like L-BFGS compute an approximate Hessian and take larger, more informed steps. These are powerful for convex problems and small networks but computationally expensive for deep learning. For curiosity, PyTorch's LBFGS optimizer:

optimizer = optim.LBFGS(model.parameters(), lr=0.1)
# Note: LBFGS requires a closure function; see PyTorch docs for details

For your projects: Adam is the default. If performance is unsatisfactory, try SGD with momentum or RMSprop before exploring second-order methods.