Optimization & Gradients

Learning as Optimization

Core Insight

Machine learning is, fundamentally, an optimization problem. We define a loss function that measures how wrong our model's predictions are, then use mathematical optimization to find model parameters that minimize this loss.

Intuition

Imagine you're blindfolded on a hilly landscape and want to reach the lowest point (a valley). You can't see, but you can feel the slope beneath your feet. The obvious strategy: always step in the direction that goes most steeply downhill.

This is gradient descent. The "landscape" is the loss function—how wrong the model is for different parameter values. The "slope" is the gradient—a vector pointing in the direction of steepest increase. By stepping opposite to the gradient, we descend toward better solutions.

📐 Quick Math Primer

New to calculus? Here's all you need to know for this section:

Derivative (∂f/∂x): "How much does f change when I nudge x?" If ∂f/∂x = 3, then increasing x by 1 increases f by ~3.
Gradient (∇f): A collection of derivatives for ALL variables. It's like asking "which way is uphill?" for every direction at once.
Negative gradient: Points DOWNHILL. That's the direction we want to go to reduce loss!

Need more detail? See the Calculus prerequisite.

Technical

For a function f(θ) where θ represents model parameters, the gradient ∇f(θ) is a vector of partial derivatives:

∇f(θ) = [∂f/∂θ₁, ∂f/∂θ₂, ..., ∂f/∂θₙ]

Each component tells us how much f changes when we adjust that parameter slightly. The gradient descent update rule is:

θ_new = θ_old - η · ∇f(θ_old)

Where η (eta) is the learning rate—how big a step we take. Too small and progress is slow; too large and we overshoot minima or diverge.

In Practice

When training a neural network:

Forward pass: Input data flows through the network, producing predictions
Loss computation: Compare predictions to true labels using a loss function
Backward pass: Compute gradients of loss with respect to every parameter
Parameter update: Adjust each parameter in the direction that reduces loss
Repeat: Process many batches of data, gradually improving

Modern networks have billions of parameters. GPT-3 has 175 billion. Optimization must navigate a 175-billion-dimensional space to find good values—and remarkably, gradient descent works.

The Loss Function

The loss function (also called cost function or objective function) quantifies how wrong the model is. Choosing the right loss function is crucial—it defines what "good" means for your model.

Mean Squared Error (MSE)

For regression tasks—predicting continuous values:

L = (1/n) Σ (y_pred - y_true)²

Squares the errors, so large errors are penalized more heavily than small ones. Good for predictions where being "a little off" is okay but "very off" is bad.

Cross-Entropy Loss

For classification tasks—predicting categories:

L = -Σ y_true · log(y_pred)

Measures the difference between predicted probability distribution and true labels. Heavily penalizes confident wrong predictions—if you predict 99% "cat" but it's a dog, the loss is huge.

◈

The Loss Landscape

For a model with many parameters, the loss function defines a high-dimensional "landscape." Training is a journey across this landscape, seeking low points. The landscape's shape—its valleys, ridges, and plateaus—determines how easy or hard the optimization problem is.

Why Gradients Work

The Mathematical Foundation

The gradient has a remarkable property: it points in the direction of steepest ascent. This means -∇f points in the direction of steepest descent. No other direction decreases the function faster (locally).

This follows from calculus: the directional derivative (rate of change in any direction) is maximized when that direction aligns with the gradient. Taking the negative gives us the direction of maximum decrease.

Gradient as Slope in Multiple Dimensions

In 2D, the gradient is a vector with two components: how steeply the function increases in the x-direction and the y-direction.

∂f/∂x > 0 means "increasing x increases loss" → move x negative
∂f/∂y < 0 means "increasing y decreases loss" → move y positive

The gradient combines these into one vector pointing "uphill." We go the opposite direction.

Gradient Properties

Property	Meaning	Implication
∇f = 0	Gradient is zero (critical point)	Local minimum, maximum, or saddle point
Large \|∇f\|	Steep slope	Far from optimum; fast progress possible
Small \|∇f\|	Flat region	Near optimum or on a plateau; slow progress
∇f changes direction quickly	High curvature	Requires smaller learning rate

🔢 Worked Example: Gradient Descent Step-by-Step

Setup

Let's minimize a simple 1D function: f(x) = x³ - 2x² + x

Derivative: f'(x) = 3x² - 4x + 1
Starting point: x₀ = 2.0
Learning rate: η = 0.1

Iterations

Step	x	f(x)	f'(x)	Update (x - η·f'(x))
0	2.000	2.000	5.000	2.0 - 0.1×5.0 = 1.500
1	1.500	0.375	1.750	1.5 - 0.1×1.75 = 1.325
2	1.325	0.067	0.726	1.325 - 0.1×0.726 = 1.252
3	1.252	0.011	0.260	1.252 - 0.1×0.26 = 1.226
4	1.226	0.001	0.085	1.226 - 0.1×0.085 = 1.218

What to Notice

Loss decreases: f(x) goes from 2.0 → 0.001, approaching the minimum
Gradients shrink: As we approach the minimum, gradients get smaller (5.0 → 0.085)
Steps get smaller: Each update changes x less because gradient is smaller
Convergence: The true minimum is at x ≈ 1.215; we're approaching it

Try This

What happens with learning rate η = 0.01 (slower) or η = 0.5 (faster)? Too small: progress is slow. Too large: may overshoot and oscillate!

Challenges in Optimization

Gradient descent isn't guaranteed to find the global optimum. Understanding its failure modes helps us design better training procedures.

⚑

Local Minima

A local minimum is a point where loss is lower than all immediate neighbors, but not the lowest overall. Gradient descent stops here because the gradient is zero.

Modern insight: In high dimensions, local minima are often "good enough"—they perform nearly as well as the global minimum. The real problem is usually saddle points.

⬛

Saddle Points

Points where the gradient is zero but it's neither a minimum nor maximum—like a mountain pass. The function curves up in some directions and down in others.

Impact: Optimization slows dramatically near saddle points. Momentum-based optimizers help by carrying velocity through flat regions.

↔

Learning Rate Sensitivity

Too high: oscillation, overshooting, or divergence. Too low: extremely slow convergence, getting stuck in shallow local minima.

Solutions: Learning rate schedules (start high, decrease over time), adaptive optimizers (Adam, RMSprop) that adjust rates per-parameter.

▓

Vanishing/Exploding Gradients

In deep networks, gradients can shrink exponentially (vanish) or grow exponentially (explode) as they propagate backward through layers.

Solutions: Careful initialization, batch normalization, residual connections, gradient clipping.

Stochastic Gradient Descent

Key Technique

Stochastic Gradient Descent (SGD) computes gradients on small random subsets (mini-batches) of data rather than the full dataset. This introduces noise but enables much faster training.

Batch Gradient Descent

Use all data for each update

Exact gradient computation
Smooth convergence path
Computationally expensive per step
Can't escape sharp local minima

Stochastic (Mini-Batch) GD

Use small batches for each update

Noisy gradient estimate
Zigzag convergence path
Fast iterations, more of them
Noise helps escape bad minima

The noise in SGD is actually beneficial. It prevents getting stuck in sharp minima that don't generalize well. Sharp minima correspond to solutions that depend sensitively on specific training examples; flat minima are more robust.

Why Mini-Batches?

Consider training on 1 million images. Batch GD would require computing predictions and gradients for all 1 million before making one parameter update. With batch size 32, we make ~31,000 updates per pass through the data.

Even though each update is noisier, the total learning per compute is much higher. Modern training uses batch sizes from 32 to thousands, balanced against available GPU memory.

Modern Optimizers

Vanilla SGD has been augmented with techniques that accelerate convergence and handle the challenges of deep network optimization.

Momentum

Accumulates a "velocity" from past gradients. Helps traverse flat regions faster and dampens oscillations in narrow valleys.

v = βv - η∇f; θ = θ + v

RMSprop

Adapts learning rate per parameter based on recent gradient magnitudes. Parameters with large gradients get smaller steps.

Divides by √(running average of squared gradients)

Adam

Combines momentum and adaptive learning rates. The default choice for most deep learning applications. Robust across many problems.

Momentum + RMSprop + bias correction

AdamW

Adam with decoupled weight decay regularization. Often preferred for training transformers and large language models.

Adam + proper L2 regularization

Common Misconceptions

✗

"Neural networks find the optimal solution"

Gradient descent finds a solution where gradients are small—not necessarily the global optimum. Different random initializations lead to different final solutions.

The accurate framing: Training finds a good local solution. In deep learning, many local minima perform comparably well, so this is acceptable. The goal is finding solutions that generalize, not finding the absolute minimum.

✗

"Faster convergence always means better results"

Reaching low training loss quickly can mean overfitting. The model might find a solution that memorizes training data but doesn't generalize.

The accurate framing: What matters is test performance, not training speed. Sometimes deliberately slowing down (lower learning rate, larger batches) or adding noise improves generalization.

✗

"Zero loss is the goal"

Zero training loss often indicates overfitting—the model has memorized the training data, including its noise and errors.

The accurate framing: The goal is low loss on unseen data (generalization). Some training loss is healthy—it indicates the model hasn't over-memorized. Early stopping and regularization prevent over-optimization.

✗

"Local minima are the main problem in deep learning"

Early theory worried about local minima, but empirical research shows that in high dimensions, most critical points are saddle points, not local minima. And most local minima have similar loss values.

The accurate framing: The main challenges are saddle points (which slow convergence), choosing good learning rates, and the generalization gap between training and test performance.

⚗ Interactive Lab: Gradient Descent Visualizer

Watch gradient descent in action on different loss landscapes. Adjust the learning rate, choose different surfaces, and see how optimization behaves in various scenarios.

Loss Surface

Learning Rate (η) 0.030

Step 0

Loss —

|∇f| —

Click on the surface to place the starting point, then watch gradient descent navigate to a minimum.

See how different learning rates affect convergence on the same problem. Watch all three simultaneously to build intuition.

η = 0.001 (Too Low)

Steps: 0

η = 0.03 (Good)

Steps: 0

η = 0.15 (Too High)

Steps: 0

Observations

Too low: Makes progress but very slowly—may need thousands of steps
Good: Steady progress toward minimum with reasonable speed
Too high: Overshoots, oscillates, may diverge entirely

Compare vanilla SGD with momentum-enhanced optimization. Momentum helps navigate narrow valleys and escape shallow local minima.

Vanilla SGD

Steps: 0, Loss: —

SGD + Momentum (β=0.9)

Steps: 0, Loss: —

Surface Type

This is what the visualizations above would look like in actual PyTorch code. The sliders you adjust correspond to the hyperparameters below.

                                    Basic Gradient Descent
                                    import torch

# Define a simple loss function (quadratic bowl)
def loss_fn(x):
    return x[0]**2 + x[1]**2

# Initialize parameters
x = torch.tensor([3.0, 4.0], requires_grad=True)

# Hyperparameters
learning_rate = 0.03  # ← Your slider controls this!
num_steps = 100

# Gradient descent loop
for step in range(num_steps):
    # Forward pass: compute loss
    loss = loss_fn(x)
    
    # Backward pass: compute gradients
    loss.backward()
    
    # Update parameters (gradient descent step)
    with torch.no_grad():
        x -= learning_rate * x.grad  # x_new = x_old - η∇f
        x.grad.zero_()  # Reset gradients for next iteration
    
    print(f"Step {step}: x = {x.data}, loss = {loss.item():.4f}")
                                

                                    SGD with Momentum
                                    import torch
import torch.optim as optim

# Define model parameters
params = torch.tensor([3.0, 4.0], requires_grad=True)

# Using PyTorch's built-in optimizer
optimizer = optim.SGD([params], lr=0.03, momentum=0.9)

for step in range(100):
    optimizer.zero_grad()  # Reset gradients
    
    loss = params[0]**2 + params[1]**2  # Loss function
    loss.backward()  # Compute gradients
    
    optimizer.step()  # Update with momentum
    
    # Momentum internally maintains velocity:
    # v_t = β * v_{t-1} + ∇f(x)
    # x_t = x_{t-1} - lr * v_t
                                

                                    Adam Optimizer (Industry Standard)
                                    import torch
import torch.optim as optim

model = MyNeuralNetwork()  # Your model
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Adam combines:
# - Momentum (first moment estimate)
# - RMSprop (second moment estimate for adaptive learning rates)
# - Bias correction for early steps

for batch in dataloader:
    optimizer.zero_grad()
    outputs = model(batch.inputs)
    loss = criterion(outputs, batch.targets)
    loss.backward()
    optimizer.step()  # Adam handles everything
                                

Try it yourself: Copy this code into a Jupyter notebook or Google Colab to experiment with different learning rates and see the actual numbers.

Key Insights

Learning rate is critical: The single most important hyperparameter. Too high causes instability; too low wastes computation.
Surface shape matters: Narrow valleys cause oscillation in vanilla SGD. Saddle points slow convergence dramatically.
Momentum provides acceleration: By accumulating velocity, momentum moves faster through flat regions and dampens oscillations.
Multiple runs, different results: Different starting points can lead to different final solutions—ensemble methods exploit this.

Learning Objectives

Learning as Optimization

The Loss Function

Mean Squared Error (MSE)

Cross-Entropy Loss

The Loss Landscape

Why Gradients Work

The Mathematical Foundation

Gradient as Slope in Multiple Dimensions

Gradient Properties

🔢 Worked Example: Gradient Descent Step-by-Step

Setup

Iterations

What to Notice

Try This

Challenges in Optimization

Local Minima

Saddle Points

Learning Rate Sensitivity

Vanishing/Exploding Gradients

Stochastic Gradient Descent

Batch Gradient Descent

Stochastic (Mini-Batch) GD

Why Mini-Batches?

Modern Optimizers

Common Misconceptions

"Neural networks find the optimal solution"

"Faster convergence always means better results"

"Zero loss is the goal"

"Local minima are the main problem in deep learning"

⚗ Interactive Lab: Gradient Descent Visualizer

η = 0.001 (Too Low)

η = 0.03 (Good)

η = 0.15 (Too High)

Observations

Vanilla SGD

SGD + Momentum (β=0.9)

Basic Gradient Descent

SGD with Momentum

Adam Optimizer (Industry Standard)

Key Insights

✓ Check Your Understanding