🎯 What You'll Learn
- Why training a model is fundamentally an optimization problem, and what the loss landscape looks like
- Understand derivatives and partial derivatives — the mathematical tools that tell us which direction to move
- Derive and implement the gradient descent update rule from scratch in NumPy
- Understand the learning rate and diagnose too-large vs too-small values using loss curves
- Compare Batch GD, Mini-batch GD, and SGD — and understand why mini-batch is the standard in practice
Every neural network trainer — PyTorch's optimizer.step(), Keras's model.fit() — is calling a variant of gradient descent under the hood. The concepts in this lesson are the foundation for understanding Adam, RMSprop, momentum, and backpropagation. If you can implement gradient descent from scratch, the rest of DL optimization becomes much clearer.
1 The Optimization Problem
Training any ML model reduces to the same fundamental problem: find the parameter values that minimize a loss function.
For linear regression (last lesson), the parameters are the weights β (β₀, β₁, ..., βₙ) and the loss function is MSE. Every model you'll meet later just swaps in its own parameters and its own loss — logistic regression (Lesson 15) uses a loss suited to yes/no answers, and a neural network (Phase 4) has millions of parameters — but the optimization problem is always the same shape: find parameter values θ that make L(θ) small.
Visualize the loss function as a hilly landscape over the parameter space. The vertical axis is the loss value. Your goal is to find the lowest point (the valley) in this landscape:
- Convex loss landscapes (linear/logistic regression): one global minimum, shaped like a bowl. Any path downhill leads to the optimal solution.
- Non-convex loss landscapes (neural networks): many local minima, saddle points, and flat plateaus. Gradient descent may find a local minimum, which is usually good enough in practice.
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
# Visualize a simple convex loss landscape: L(w) = (w - 3)^2
w_values = np.linspace(-1, 7, 100)
loss_values = (w_values - 3) ** 2
plt.figure(figsize=(8, 4))
plt.plot(w_values, loss_values, 'steelblue', linewidth=2)
plt.xlabel('Parameter w')
plt.ylabel('Loss L(w)')
plt.title('Convex Loss Landscape — One Global Minimum at w=3')
plt.axvline(x=3, color='red', linestyle='--', alpha=0.5, label='Optimal w=3')
plt.scatter([3], [0], color='red', zorder=5, s=80)
plt.legend()
plt.savefig('loss_landscape.png', dpi=100)
print("Saved: loss_landscape.png")The normal equation β = (XᵀX)⁻¹Xᵀy gives the exact optimal parameters for linear regression in one step. But it requires inverting a p×p matrix, which is O(p³) — impractical when you have millions of features, and impossible for neural networks where the loss isn't a quadratic in the parameters. Gradient descent, while approximate and iterative, scales to any model and any differentiable loss function.
2 Derivatives: The Direction of Steepest Change
The derivative of a function f(x) at a point x tells us the instantaneous rate of change of f as x changes. Geometrically, it's the slope of the tangent line to the function at that point.
If the derivative is positive at x, f is increasing — moving x right increases f. If the derivative is negative, f is decreasing — moving x right decreases f. The derivative's sign tells us which direction to move x to decrease f, and its magnitude tells us how steep the slope is.
import numpy as np
# Example: f(x) = x^2, analytical derivative f'(x) = 2x
def f(x):
return x ** 2
def df(x):
"""Analytical derivative of f(x) = x^2."""
return 2 * x
# Numerical derivative (finite difference approximation)
def numerical_df(x, h=1e-5):
"""Approximate derivative using central differences."""
return (f(x + h) - f(x - h)) / (2 * h)
x_points = np.array([-3.0, -1.0, 0.0, 1.0, 3.0])
print(f"{'x':>6} {'f(x)':>8} {'f\'(x) analytic':>16} {'f\'(x) numeric':>14}")
for x in x_points:
print(f"{x:>6.1f} {f(x):>8.2f} {df(x):>16.4f} {numerical_df(x):>14.4f}")
# Output:
# x f(x) f'(x) analytic f'(x) numeric
# -3.0 9.00 -6.0000 -6.0000
# -1.0 1.00 -2.0000 -2.0000
# 0.0 0.00 0.0000 0.0000 ← minimum: derivative = 0
# 1.0 1.00 2.0000 2.0000
# 3.0 9.00 6.0000 6.0000At x=0, the derivative is 0 — this is the minimum of f(x)=x². At a minimum (or maximum), the gradient is zero. This is why we call them "critical points" — gradient descent is searching for a point where the gradient equals zero.
3 Partial Derivatives: Gradients in Multiple Dimensions
ML models have multiple parameters. The gradient is the vector of partial derivatives — it generalises the derivative to multi-variable functions and points in the direction of steepest ascent of the loss.
For linear regression with parameters β₀ (intercept) and β₁ (slope), the MSE loss is:
L(β₀, β₁) = (1/n) Σᵢ (yᵢ − β₀ − β₁xᵢ)²
Taking partial derivatives with respect to each parameter:
∂L/∂β₀ = −(2/n) Σᵢ (yᵢ − β₀ − β₁xᵢ) ∂L/∂β₁ = −(2/n) Σᵢ xᵢ(yᵢ − β₀ − β₁xᵢ)
import numpy as np
def mse_gradients(X, y, beta_0, beta_1):
"""
Compute partial derivatives of MSE w.r.t. beta_0 and beta_1.
X: feature vector (1D), y: target vector, beta_0, beta_1: scalars
"""
n = len(y)
predictions = beta_0 + beta_1 * X
residuals = y - predictions # yᵢ - ŷᵢ
grad_b0 = (-2/n) * np.sum(residuals) # ∂L/∂β₀
grad_b1 = (-2/n) * np.sum(X * residuals) # ∂L/∂β₁
return grad_b0, grad_b1
# Test with some data
np.random.seed(42)
X_data = np.random.randn(50)
y_data = 3.0 * X_data + 2.0 + np.random.randn(50) * 0.5 # true: b0=2, b1=3
# Current (wrong) parameters: β₀=0, β₁=0
g0, g1 = mse_gradients(X_data, y_data, beta_0=0, beta_1=0)
print(f"∂L/∂β₀ at (0,0): {g0:.4f}") # large negative → should increase β₀
print(f"∂L/∂β₁ at (0,0): {g1:.4f}") # large negative → should increase β₁
# At the optimal parameters
g0_opt, g1_opt = mse_gradients(X_data, y_data, beta_0=2.0, beta_1=3.0)
print(f"∂L/∂β₀ at (2,3): {g0_opt:.4f}") # ≈ 0
print(f"∂L/∂β₁ at (2,3): {g1_opt:.4f}") # ≈ 0The gradient vector ∇L points in the direction that increases the loss the fastest. To minimize the loss, we move in the opposite direction: negative gradient. This is why gradient descent subtracts the gradient from the current parameters.
4 The Chain Rule: Gradients Through Composed Functions
The chain rule is how we compute derivatives of composed functions — functions applied to functions. In a neural network, the loss L depends on the final layer output, which depends on every preceding layer, which depends on the input weights. The chain rule lets us compute ∂L/∂(any weight) by multiplying derivatives along the chain.
If z = f(g(x)), then dz/dx = f'(g(x)) · g'(x)
import numpy as np
# Example: three functions applied in sequence
# z = w * x (a weighted input)
# a = sigmoid(z) (an S-shaped squashing function, defined below —
# it stars in Lesson 15, but here it's just
# "some function we need the derivative of")
# L = (a - y)^2 (squared error loss)
def sigmoid(z):
return 1 / (1 + np.exp(-z))
def sigmoid_derivative(z):
s = sigmoid(z)
return s * (1 - s) # dsigmoid/dz
# Chain rule: dL/dw = dL/da · da/dz · dz/dw
def chain_rule_gradient(x, w, y):
z = w * x # linear step
a = sigmoid(z) # activation
L = (a - y) ** 2 # loss
dL_da = 2 * (a - y) # derivative of L w.r.t. a
da_dz = sigmoid_derivative(z) # derivative of sigmoid w.r.t. z
dz_dw = x # derivative of z w.r.t. w
dL_dw = dL_da * da_dz * dz_dw # chain rule product
return dL_dw, L
x, w, y_true = 2.0, 0.5, 1.0
grad_w, loss = chain_rule_gradient(x, w, y_true)
print(f"Loss: {loss:.6f}")
print(f"dL/dw: {grad_w:.6f}")
print(f" (negative gradient → increase w to reduce loss)")
# Verify numerically
h = 1e-7
_, L_plus = chain_rule_gradient(x, w + h, y_true)
_, L_minus = chain_rule_gradient(x, w - h, y_true)
numerical_grad = (L_plus - L_minus) / (2 * h)
print(f"Numerical dL/dw: {numerical_grad:.6f}")Backpropagation — the algorithm that trains neural networks — is nothing more than the chain rule applied systematically to every layer. Starting from the loss and working backwards (hence "back" propagation), it computes ∂L/∂(every weight) using the chain rule. This is why understanding the chain rule is prerequisite to understanding deep learning.
5 The Gradient Descent Algorithm
With the gradient in hand, the update rule is simple: subtract a small fraction of the gradient from the current parameters. Repeat until convergence.
θ ← θ − α · ∇L(θ)
where θ represents all parameters, α is the learning rate (a small positive scalar), and ∇L(θ) is the gradient of the loss with respect to all parameters.
Step-by-step algorithm:
- Initialize parameters θ randomly (or to zeros)
- Compute predictions ŷ = f(X; θ)
- Compute loss L(y, ŷ)
- Compute gradient ∇L(θ) via calculus (or backpropagation in neural nets)
- Update: θ ← θ − α · ∇L(θ)
- Go to step 2. Repeat until loss converges or maximum iterations reached.
One iteration of θ ← θ − α·∇L(θ): the gradient at θₜ points uphill (steepest ascent), so we step in the opposite direction — downhill, toward θₜ₊₁. Repeat enough times and θ converges to θ*, where the slope is flat.
import numpy as np
def gradient_descent_linear_regression(X, y, learning_rate=0.01, n_iterations=1000):
"""
Gradient descent for linear regression: y = beta_0 + beta_1 * x
"""
n = len(y)
beta_0, beta_1 = 0.0, 0.0 # initialize to zeros
loss_history = []
for i in range(n_iterations):
# Forward pass: compute predictions
y_pred = beta_0 + beta_1 * X
# Compute MSE loss
loss = np.mean((y - y_pred) ** 2)
loss_history.append(loss)
# Compute gradients (partial derivatives of MSE)
residuals = y - y_pred
grad_b0 = (-2/n) * np.sum(residuals)
grad_b1 = (-2/n) * np.sum(X * residuals)
# Update rule: θ ← θ − α · ∇L
beta_0 -= learning_rate * grad_b0
beta_1 -= learning_rate * grad_b1
if i % 100 == 0:
print(f"Iter {i:4d} | Loss: {loss:.4f} | β₀: {beta_0:.4f} | β₁: {beta_1:.4f}")
return beta_0, beta_1, loss_history
# Test on salary data
np.random.seed(42)
X = np.random.uniform(0, 15, 100)
y = 30000 + 4500 * X + np.random.normal(0, 5000, 100)
b0, b1, losses = gradient_descent_linear_regression(X, y, learning_rate=0.001, n_iterations=1000)
print(f"\nFinal: β₀={b0:.0f}, β₁={b1:.1f}")
print(f"True: β₀=30000, β₁=4500")6 Learning Rate: The Most Critical Hyperparameter
The learning rate α controls how large a step we take in the direction of the negative gradient at each iteration. Choosing it poorly breaks training:
- Too large (α = 0.1): the update overshoots the minimum. Parameters oscillate wildly or diverge — loss increases instead of decreasing.
- Too small (α = 0.0001): convergence is extremely slow. You might need 100× more iterations to reach the same performance.
- Just right (α = 0.01): loss decreases smoothly and quickly to the minimum.
Drag the slider below to watch gradient descent take steps on the convex loss L(w) = (w − 3)², starting from w = −5. Watch what happens to the path as α crosses 0.5 (oscillation begins) and 1.0 (it diverges off the chart):
α = 0.10 — smooth, fast convergence toward w = 3.
import numpy as np
np.random.seed(42)
X = np.random.uniform(0, 10, 50)
y = 2.0 * X + 5.0 + np.random.randn(50)
def run_gd(lr, n_iter=200):
"""Run gradient descent with given learning rate, return loss curve."""
b0, b1 = 0.0, 0.0
losses = []
n = len(y)
for _ in range(n_iter):
y_hat = b0 + b1 * X
loss = np.mean((y - y_hat) ** 2)
losses.append(loss)
grad_b0 = (-2/n) * np.sum(y - y_hat)
grad_b1 = (-2/n) * np.sum(X * (y - y_hat))
b0 -= lr * grad_b0
b1 -= lr * grad_b1
return losses
losses_small = run_gd(lr=0.0001)
losses_good = run_gd(lr=0.01)
losses_large = run_gd(lr=0.15) # might diverge
print("Final loss with lr=0.0001:", round(losses_small[-1], 4)) # high — slow
print("Final loss with lr=0.01: ", round(losses_good[-1], 4)) # low — good
print("Final loss with lr=0.15: ", round(losses_large[-1], 4)) # may be inf/nanLearning Rate Schedules
In practice, using a fixed learning rate throughout training is rarely optimal. Learning rate schedules adjust α over time:
# Step decay: reduce lr by factor 0.1 every 50 epochs
def step_decay_lr(initial_lr, epoch, drop=0.1, epochs_drop=50):
return initial_lr * (drop ** (epoch // epochs_drop))
# Exponential decay
def exp_decay_lr(initial_lr, epoch, decay_rate=0.95):
return initial_lr * (decay_rate ** epoch)
# Example: use in training loop
initial_lr = 0.1
for epoch in range(200):
lr = step_decay_lr(initial_lr, epoch, drop=0.5, epochs_drop=50)
if epoch % 50 == 0:
print(f"Epoch {epoch}: lr = {lr:.5f}")
# Epoch 0: lr = 0.10000
# Epoch 50: lr = 0.05000
# Epoch 100: lr = 0.02500
# Epoch 150: lr = 0.01250If you see nan or inf in your loss during training, the first thing to do is reduce the learning rate by 10×. Numerical overflow from large gradients multiplied by a large learning rate is by far the most common cause of NaN loss. In neural networks, gradient clipping (torch.nn.utils.clip_grad_norm_) also prevents this.
A loss surface that's steep in one direction and shallow in another — exactly why a single learning rate has to compromise, and why it zigzags on the way down.
7 Batch vs Mini-Batch vs Stochastic Gradient Descent
The three variants of gradient descent differ in how much data they use to compute the gradient at each update step:
| Variant | Data used per update | Gradient quality | Speed per epoch | Memory use |
|---|---|---|---|---|
| Batch GD | Full dataset (n samples) | Exact | Slow (1 update per epoch) | High |
| SGD | 1 sample | Noisy | Fast (n updates per epoch) | Very low |
| Mini-Batch GD ⭐ | Batch of b samples (32–512) | Good approximation | Fast + parallelisable | Moderate |
import numpy as np
np.random.seed(42)
n = 1000
X = np.random.randn(n, 1)
y = 3.5 * X.ravel() + 2.0 + np.random.randn(n)
def compute_gradients(X_batch, y_batch, b0, b1):
nb = len(y_batch)
y_hat = b0 + b1 * X_batch
residuals = y_batch - y_hat
g0 = (-2/nb) * np.sum(residuals)
g1 = (-2/nb) * np.sum(X_batch * residuals)
return g0, g1
# ── 1. Batch Gradient Descent ──────────────────────────────────────────
def batch_gd(X, y, lr=0.01, n_epochs=50):
b0, b1 = 0.0, 0.0
losses = []
for _ in range(n_epochs):
g0, g1 = compute_gradients(X.ravel(), y, b0, b1)
b0 -= lr * g0
b1 -= lr * g1
losses.append(np.mean((y - b0 - b1 * X.ravel()) ** 2))
return b0, b1, losses
# ── 2. Stochastic Gradient Descent ─────────────────────────────────────
def sgd(X, y, lr=0.01, n_epochs=5):
b0, b1 = 0.0, 0.0
losses = []
for _ in range(n_epochs):
indices = np.random.permutation(len(y)) # shuffle each epoch
for i in indices:
g0, g1 = compute_gradients(np.array([X[i, 0]]), np.array([y[i]]), b0, b1)
b0 -= lr * g0
b1 -= lr * g1
losses.append(np.mean((y - b0 - b1 * X.ravel()) ** 2))
return b0, b1, losses
# ── 3. Mini-Batch Gradient Descent ─────────────────────────────────────
def mini_batch_gd(X, y, lr=0.01, n_epochs=5, batch_size=32):
b0, b1 = 0.0, 0.0
n = len(y)
losses = []
for _ in range(n_epochs):
indices = np.random.permutation(n)
for start in range(0, n, batch_size):
batch_idx = indices[start:start + batch_size]
X_batch = X[batch_idx, 0]
y_batch = y[batch_idx]
g0, g1 = compute_gradients(X_batch, y_batch, b0, b1)
b0 -= lr * g0
b1 -= lr * g1
losses.append(np.mean((y - b0 - b1 * X.ravel()) ** 2))
return b0, b1, losses
# Compare final parameters
b0_batch, b1_batch, _ = batch_gd(X, y, lr=0.01, n_epochs=50)
b0_sgd, b1_sgd, _ = sgd(X, y, lr=0.01, n_epochs=10)
b0_mini, b1_mini, _ = mini_batch_gd(X, y, lr=0.01, n_epochs=10, batch_size=32)
print(f"True: b0=2.000, b1=3.500")
print(f"Batch GD: b0={b0_batch:.3f}, b1={b1_batch:.3f}")
print(f"SGD: b0={b0_sgd:.3f}, b1={b1_sgd:.3f}")
print(f"Mini-batch: b0={b0_mini:.3f}, b1={b1_mini:.3f}")Mini-batch GD with batch sizes of 32–512 is the universal standard in neural network training. It balances three properties: (1) enough data per batch that gradients are a good approximation of the true gradient, (2) enough noise to escape sharp local minima (unlike batch GD, which follows the exact gradient), and (3) batch sizes that fit in GPU memory and exploit parallelism efficiently. PyTorch's DataLoader handles mini-batching automatically.
8 Full Implementation: Gradient Descent from Scratch
Let's build a complete, reusable gradient descent implementation for linear regression, with loss tracking and convergence detection:
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
class LinearRegressionGD:
"""
Linear Regression trained via Mini-Batch Gradient Descent.
Implements the same interface as sklearn's LinearRegression.
"""
def __init__(self, learning_rate=0.01, n_epochs=100,
batch_size=32, tolerance=1e-6, verbose=False):
self.lr = learning_rate
self.n_epochs = n_epochs
self.batch_size = batch_size
self.tol = tolerance
self.verbose = verbose
self.loss_history_ = []
def fit(self, X, y):
n, p = X.shape
# Add bias column
X_b = np.column_stack([np.ones(n), X])
p_b = p + 1
# Initialize weights randomly (small values)
self.weights_ = np.random.randn(p_b) * 0.01
prev_loss = float('inf')
for epoch in range(self.n_epochs):
# Shuffle data each epoch (prevents patterns in updates)
perm = np.random.permutation(n)
X_shuffled, y_shuffled = X_b[perm], y[perm]
# Process mini-batches
for start in range(0, n, self.batch_size):
X_batch = X_shuffled[start:start + self.batch_size]
y_batch = y_shuffled[start:start + self.batch_size]
nb = len(y_batch)
# Forward pass
y_hat = X_batch @ self.weights_
# Gradient of MSE: (2/nb) * Xᵀ(ŷ - y)
grad = (2/nb) * X_batch.T @ (y_hat - y_batch)
# Update step
self.weights_ -= self.lr * grad
# Track epoch loss (on full dataset)
y_hat_full = X_b @ self.weights_
epoch_loss = np.mean((y - y_hat_full) ** 2)
self.loss_history_.append(epoch_loss)
if self.verbose and epoch % 10 == 0:
print(f"Epoch {epoch:4d} | MSE: {epoch_loss:.4f}")
# Early stopping
if abs(prev_loss - epoch_loss) < self.tol:
print(f"Converged at epoch {epoch}")
break
prev_loss = epoch_loss
self.intercept_ = self.weights_[0]
self.coef_ = self.weights_[1:]
return self
def predict(self, X):
return X @ self.coef_ + self.intercept_
def score(self, X, y):
y_pred = self.predict(X)
ss_res = np.sum((y - y_pred) ** 2)
ss_tot = np.sum((y - y.mean()) ** 2)
return 1 - ss_res / ss_tot
# Test
from sklearn.datasets import make_regression
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
X, y = make_regression(n_samples=500, n_features=5, noise=20, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# Scale features (important for GD convergence speed)
scaler = StandardScaler()
X_train_s = scaler.fit_transform(X_train)
X_test_s = scaler.transform(X_test)
model = LinearRegressionGD(learning_rate=0.05, n_epochs=100, batch_size=32, verbose=True)
model.fit(X_train_s, y_train)
print(f"\nTrain R²: {model.score(X_train_s, y_train):.4f}")
print(f"Test R²: {model.score(X_test_s, y_test):.4f}")
# Compare to sklearn
from sklearn.linear_model import LinearRegression as SkLearnLR
sk_model = SkLearnLR().fit(X_train_s, y_train)
print(f"sklearn R²: {sk_model.score(X_test_s, y_test):.4f}")Real-World Spotlight: Neural Network Training is Gradient Descent
optimizer.step() implements exactly the update rule θ ← θ − α · ∇L(θ) — but with millions of parameters and automatic gradient computation via autograd.
# PREVIEW ONLY — this is PyTorch, the deep learning library you'll
# learn from scratch in Phase 4 (Lesson 39). Don't type it now; just
# spot the update rule you built in this lesson inside a real framework.
import torch
import torch.nn as nn
# A PyTorch training loop — gradient descent in action
model = nn.Sequential(
nn.Linear(10, 64),
nn.ReLU(),
nn.Linear(64, 1)
)
# Adam optimizer = adaptive gradient descent variant
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()
# Training loop
for epoch in range(100):
for X_batch, y_batch in dataloader: # mini-batch GD
# Forward pass
y_pred = model(X_batch)
loss = criterion(y_pred.squeeze(), y_batch)
# Backward pass = compute gradients via backpropagation (chain rule!)
optimizer.zero_grad() # clear previous gradients
loss.backward() # dL/d(every weight) via autograd
# Update step = θ ← θ − α · ∇L
optimizer.step()
if epoch % 10 == 0:
print(f"Epoch {epoch}: Loss = {loss.item():.4f}")
# The 'loss.backward()' call computes ∂L/∂w for EVERY weight in the network
# using the chain rule — exactly what you implemented manually in this lesson.
# torch.optim.Adam, SGD, RMSprop all call the same optimizer.step() interface.Understanding gradient descent deeply means you can diagnose training problems: NaN loss (learning rate too large), plateauing loss (learning rate too small, or wrong optimizer), oscillating loss (noisy mini-batches need a larger batch size). All of DL optimization is built on this foundation.
Quick Check
✍️ Practice Exercises
- Implement gradient descent for logistic regression from scratch. The gradient of the binary cross-entropy loss is
(1/n) * X.T @ (sigmoid(X @ w) - y). Train it on a binary classification dataset and compare to sklearn's LogisticRegression. - Write an experiment that plots loss curves for 5 different learning rates (0.0001, 0.001, 0.01, 0.1, 1.0) on the same axes. Identify the "Goldilocks" learning rate for your dataset.
- Extend the
LinearRegressionGDclass with L2 regularization (Ridge penalty). The gradient becomes:grad = (2/nb) * X_batch.T @ (y_hat - y_batch) + 2 * lambda_ * weights(don't regularize the bias term). - Compare convergence speed of Batch GD vs Mini-Batch GD (batch_size=32) on 10,000 samples. Plot loss vs wall-clock time (not epoch number) to show that mini-batch is faster in practice.
📚 Primary Source for This Lesson
AI Notes: Optimization — deeplearning.ai
An excellent visual and mathematical walkthrough of gradient descent, learning rates, and gradient descent variants. Also highly recommended: CS231n Notes on Optimization (Stanford) — covers gradient computation, numerical gradients, and the chain rule with clear visualisations.