🎯 What You'll Learn
- Why batch size matters: understand the fundamental trade-off between gradient accuracy and training speed
- Implement Batch GD, SGD, and Mini-Batch GD from scratch and compare their convergence behaviors
- Understand learning rate schedules — step decay, cosine annealing, and linear warmup — and when to use each
- Implement learning rate warmup to stabilize early training with random initialisations
- Use gradient clipping to prevent exploding gradients in deep networks and RNNs
- Combine mini-batch SGD + cosine annealing for a production-ready training configuration
Gradient descent is like descending a mountain blindfolded. You can't see the whole landscape — you can only feel the slope under your feet and take a step downhill. The question is: how much of the landscape do you look at before each step? Do you survey the entire mountain? A small patch? Just the rock under your left foot? That choice — all of it, a little, or just one patch — is the difference between Batch GD, Mini-batch GD, and SGD. Each strategy has different accuracy, speed, and generalization properties.
1 Recap: What Gradient Descent Does
Before diving into the variants, let's nail down the core idea. Training a neural network is an optimization problem: we have a loss function L that measures how wrong our model is, and parameters W (millions of weights and biases) that we need to adjust to make L small.
The loss L is a surface over the parameter space. Think of a terrain with hills and valleys. Our goal is to find the lowest valley (minimum loss). The gradient ∇L tells us which direction is "uphill" — so we step in the opposite direction. The update rule is:
W ← W − α · ∇L(W)
Here, α is the learning rate — how far we step. ∇L(W) is the gradient — which direction is uphill. We subtract because we want to go downhill.
The critical question this lesson answers: over which data samples do we compute the gradient ∇L? The answer has enormous consequences for training speed, memory use, convergence quality, and generalization.
With a dataset of n samples, we have three fundamental choices:
| Variant | Samples per update | Updates per epoch | Gradient quality |
|---|---|---|---|
| Batch GD | All n samples | 1 | Exact — but slow |
| SGD | 1 sample | n | Very noisy — but fast |
| Mini-Batch GD ⭐ | B samples (32–256) | n / B | Good approximation — best of both |
In most deep learning papers and frameworks, when people say "we trained with SGD", they almost always mean mini-batch gradient descent — not literally one sample at a time. The term "SGD" has come to encompass the whole family. PyTorch's torch.optim.SGD takes whatever you feed it per batch — and you control the batch size via your DataLoader.
All three variants sit on the same spectrum — the only difference is how many samples B contribute to each gradient estimate. Small B → noisy but frequent updates. Large B → smooth but rare updates.
2 Batch Gradient Descent (Full-Batch GD)
Batch GD computes the gradient over all n training samples before taking a single update step. Think of it as consulting every employee in a company before making one decision — you get perfect information, but the meeting takes forever.
The algorithm:
- Pass all n training samples through the model → get predictions
- Compute loss over all n samples
- Compute gradient ∇L averaged over all n samples
- Take one update step: W ← W − α · ∇L
- Repeat for next epoch
Why it's exact: The gradient is computed over the entire distribution of training data, so it's the "true" gradient of your loss function (no noise, no randomness). The path down the loss surface is smooth.
Why it breaks in practice: Imagine training on ImageNet (1.2 million images). Before taking a single parameter update, you must forward-pass all 1.2M images. That takes minutes per update step. And all 1.2M images (plus activations) must fit in memory simultaneously — often impossible.
import numpy as np
def batch_gradient_descent(X, y, learning_rate=0.01, n_epochs=100, verbose=True):
"""
Full-Batch Gradient Descent for linear regression.
Computes gradient over the ENTIRE dataset before each update.
"""
n, p = X.shape
# Add bias column (column of 1s)
X_b = np.column_stack([np.ones(n), X])
w = np.zeros(p + 1) # initialize weights to zero
loss_history = []
for epoch in range(n_epochs):
# Forward pass: predictions for ALL n samples
y_pred = X_b @ w
# Loss over all n samples
residuals = y_pred - y
loss = np.mean(residuals ** 2)
loss_history.append(loss)
# Gradient over ALL n samples: (2/n) * X^T (y_pred - y)
grad = (2 / n) * X_b.T @ residuals
# ONE update step (per epoch)
w -= learning_rate * grad
if verbose and epoch % 20 == 0:
print(f"Epoch {epoch:4d} | MSE: {loss:.4f} | w[0]={w[0]:.3f} w[1]={w[1]:.3f}")
return w, loss_history
# Test on synthetic linear data
np.random.seed(42)
n = 500
X = np.random.randn(n, 2)
y = 3.0 * X[:, 0] + 1.5 * X[:, 1] + 2.0 + np.random.randn(n) * 0.5
w_final, losses = batch_gradient_descent(X, y, learning_rate=0.05, n_epochs=100)
print(f"\nFinal weights: {w_final}")
print(f"True: [2.0, 3.0, 1.5] (bias, w1, w2)")Notice how smooth the convergence is — each step is in exactly the right direction. The loss decreases monotonically. But this comes at a cost: for n=500 samples with 2 features, this is manageable. For n=10,000,000 images, it's not.
Batch GD cannot work in online or streaming settings (where new data arrives continuously). It also gets stuck more easily in sharp local minima because the gradient is never "noisy" — and a little noise actually helps escape bad solutions. For modern deep learning, the dataset size alone makes it impractical.
3 Stochastic Gradient Descent (SGD) — One Sample at a Time
At the opposite extreme: compute the gradient using just one randomly chosen training sample, then immediately update the parameters. Think of it as asking one random employee their opinion, acting on it instantly, then asking another. You get feedback and action at maximum speed — but each individual opinion is noisy and may not represent the group.
The algorithm:
- Shuffle the dataset (prevents patterns in update order)
- For each sample i: compute gradient ∇L(xᵢ, yᵢ, W)
- Immediately update: W ← W − α · ∇L(xᵢ)
- Move to the next sample
- After processing all n samples = one epoch
The good news about noise: SGD's noisy gradient estimates are actually a feature, not a bug. The noise can help the optimizer escape sharp local minima — the random perturbations occasionally push you uphill and over a ridge into a better valley. This is one reason SGD-trained models sometimes generalize better than batch-trained ones.
import numpy as np
def stochastic_gradient_descent(X, y, learning_rate=0.01, n_epochs=10, verbose=True):
"""
True SGD: update parameters after EACH SINGLE sample.
"""
n, p = X.shape
X_b = np.column_stack([np.ones(n), X])
w = np.zeros(p + 1)
loss_history = [] # recorded once per epoch
for epoch in range(n_epochs):
# Shuffle indices — important! Prevents correlated updates.
indices = np.random.permutation(n)
for i in indices:
xi = X_b[i : i+1] # shape (1, p+1) — single row
yi = y[i : i+1] # shape (1,)
# Gradient on this ONE sample
y_pred_i = xi @ w
residual_i = y_pred_i - yi
grad = 2 * xi.T @ residual_i # shape (p+1,)
# Update immediately (no /n because it's just 1 sample)
w -= learning_rate * grad
# Measure epoch loss over full dataset
y_pred_all = X_b @ w
epoch_loss = np.mean((y_pred_all - y) ** 2)
loss_history.append(epoch_loss)
if verbose:
print(f"Epoch {epoch+1} | MSE: {epoch_loss:.4f}")
return w, loss_history
np.random.seed(42)
n = 500
X = np.random.randn(n, 2)
y = 3.0 * X[:, 0] + 1.5 * X[:, 1] + 2.0 + np.random.randn(n) * 0.5
w_sgd, losses_sgd = stochastic_gradient_descent(X, y, learning_rate=0.01, n_epochs=10)
print(f"\nSGD final weights: {w_sgd.round(3)}")
print(f"True: [2.0, 3.0, 1.5]")With 500 samples, SGD makes 500 gradient updates per epoch versus Batch GD's 1. It converges much faster in terms of wall-clock time, but the loss path is much more erratic — you'll see it bounce around rather than descend smoothly.
One underappreciated advantage: SGD works when you can't store your full dataset in memory. In streaming applications (real-time user behavior, sensor data), you update the model as each new sample arrives. Batch GD requires the entire dataset upfront and cannot adapt to new data without a full retraining pass.
4 Mini-Batch Gradient Descent: The Sweet Spot
Mini-batch GD is the method that became the de facto standard in all of deep learning. Instead of one sample (SGD) or all samples (Batch GD), you use a mini-batch of B samples — typically 32, 64, or 128.
Why this is the sweet spot:
- GPU parallelism: A GPU can process 32 samples in nearly the same time as 1 sample (thanks to massive parallel cores). Using B=1 wastes 97% of GPU capacity. Mini-batches fully utilize the hardware.
- Good enough gradients: The average gradient over 32 samples is a much better estimate of the true gradient than a single sample's gradient — but computing over all n samples gives only a marginal further improvement.
- Enough noise to escape bad minima: Unlike batch GD, mini-batch gradients are still stochastic — they don't follow the exact gradient, so the optimizer can wander its way out of shallow local minima.
- Memory efficient: Only B samples (not all n) need to be in GPU memory at once.
import numpy as np
def mini_batch_gradient_descent(X, y, learning_rate=0.01,
n_epochs=20, batch_size=32, verbose=True):
"""
Mini-Batch Gradient Descent — the standard in deep learning.
Updates parameters after each batch of B samples.
"""
n, p = X.shape
X_b = np.column_stack([np.ones(n), X])
w = np.zeros(p + 1)
loss_history = []
for epoch in range(n_epochs):
# Shuffle entire dataset at start of epoch
perm = np.random.permutation(n)
X_shuffled = X_b[perm]
y_shuffled = y[perm]
# Process each mini-batch
n_batches = 0
for start in range(0, n, batch_size):
X_batch = X_shuffled[start : start + batch_size]
y_batch = y_shuffled[start : start + batch_size]
B = len(y_batch) # actual batch size (last batch may be smaller)
# Forward pass on this batch
y_pred = X_batch @ w
residuals = y_pred - y_batch
# Gradient averaged over this batch
grad = (2 / B) * X_batch.T @ residuals
# Update
w -= learning_rate * grad
n_batches += 1
# Epoch-level loss (full dataset)
y_pred_all = X_b @ w
epoch_loss = np.mean((y_pred_all - y) ** 2)
loss_history.append(epoch_loss)
if verbose and epoch % 5 == 0:
print(f"Epoch {epoch:3d} | MSE: {epoch_loss:.4f} | batches/epoch: {n_batches}")
return w, loss_history
# Compare all three methods
np.random.seed(42)
n = 1000
X = np.random.randn(n, 3)
y = 2.5 * X[:, 0] - 1.5 * X[:, 1] + 0.8 * X[:, 2] + 3.0 + np.random.randn(n) * 0.5
print("=== Batch GD ===")
_, losses_batch = batch_gradient_descent(X, y, learning_rate=0.05, n_epochs=30, verbose=False)
print("=== Mini-Batch GD (B=32) ===")
_, losses_mini = mini_batch_gradient_descent(X, y, learning_rate=0.05,
n_epochs=10, batch_size=32, verbose=True)
print(f"\nBatch GD (30 epochs): Final loss = {losses_batch[-1]:.4f}")
print(f"Mini-batch (10 epochs): Final loss = {losses_mini[-1]:.4f}")Mini-batch reaches the same loss as Batch GD in far fewer epochs because it makes 32 updates per epoch (n/B = 1000/32 ≈ 32) versus Batch GD's 1 update per epoch. Fewer epochs × more updates per epoch = faster convergence.
Numbers in a table only tell half the story — the real difference between these three variants is the shape of the loss curve as training progresses. Using the same synthetic linear regression setup as the code above (n=1000, 3 features), the chart below simulates the per-step loss (not just per-epoch) for all three variants on an identical loss surface. Toggle each line on and off to compare them in isolation:
Loss per gradient update — Batch GD takes only 1 (very effective) step per epoch, SGD takes 1000 noisy steps per epoch, and Mini-Batch GD takes ~31 comparatively smooth steps per epoch.
Look at the texture of each line, not just its endpoint. Batch GD's curve is a smooth, monotonic staircase — every step is guaranteed to reduce the loss because it's computed from the exact gradient. SGD's curve is jagged and vibrates constantly — some individual steps even increase the loss — because each step reacts to just one noisy sample. Mini-batch GD sits in between: still visibly noisy step-to-step, but the noise averages out over a batch, so the overall trend is far cleaner than SGD while still reaching low loss in a fraction of Batch GD's wall-clock steps. This noise pattern — not just the convergence speed — is the key intuition to take away from this lesson.
5 Choosing Batch Size: A Practical Guide
The batch size B is one of the first hyperparameters you'll tune. It affects everything: training speed, GPU memory use, gradient noise, and final model accuracy. Here's how to think about it.
The Classical Starting Point: 32 or 64
For most problems, start with B=32 or B=64. These have been empirically validated across hundreds of papers and work well out of the box. They provide enough gradient signal while keeping GPU utilization high.
Larger Batches Are Not Always Better
There's a tempting but wrong intuition: "more data per batch = better gradient = faster training = better model." This is only partly true. Large batches (B=512 or B=2048) converge to sharper minima — loss functions that have narrow, steep valleys. Sharp minima tend to generalize worse because a small perturbation in the parameter (from test data distribution shift) takes you far up the steep wall.
Small batches converge to flatter minima — wide valleys. A small perturbation barely changes the loss. Flat minima generalize better.
This is one of the deeper insights in deep learning theory: batch noise is not just a nuisance to be eliminated — it's a regularizer. The stochasticity from mini-batches pushes the optimizer away from sharp, narrow minima and towards wider, flatter ones. This "implicit regularization" from small batches is part of why well-tuned mini-batch training often beats large-batch training even when the large-batch run converges faster in terms of iterations.
Memory Constraint: Batch Size × Model Size Must Fit in VRAM
import torch
import torch.nn as nn
# Estimate GPU memory for a batch
def estimate_batch_memory(batch_size, input_dim, hidden_dim, output_dim, dtype=torch.float32):
"""
Rough estimate of memory needed for one forward+backward pass.
Actual memory also includes optimizer states — this is approximate.
"""
bytes_per_element = 4 if dtype == torch.float32 else 2
# Input tensor: batch_size x input_dim
input_mem = batch_size * input_dim * bytes_per_element
# Activations stored for backprop (rough: hidden_dim per layer × batch)
activation_mem = batch_size * hidden_dim * bytes_per_element * 3 # 3 layers
# Gradients: same size as parameters
param_count = input_dim * hidden_dim + hidden_dim * hidden_dim + hidden_dim * output_dim
grad_mem = param_count * bytes_per_element * 2 # params + grads
total_bytes = input_mem + activation_mem + grad_mem
return total_bytes / (1024 ** 3) # convert to GB
for batch_size in [32, 64, 128, 256, 512]:
mem = estimate_batch_memory(batch_size, input_dim=784, hidden_dim=512, output_dim=10)
print(f"Batch size {batch_size:4d}: ~{mem:.3f} GB")
# Output (approximate):
# Batch size 32: ~0.004 GB
# Batch size 64: ~0.007 GB
# Batch size 128: ~0.013 GB
# Batch size 256: ~0.024 GB
# Batch size 512: ~0.046 GBGradient Accumulation: Simulating Large Batches
What if you want the gradient quality of B=256 but your GPU only has memory for B=32? Use gradient accumulation: run 8 forward+backward passes with B=32 without clearing the gradients, then do one optimizer step. The accumulated gradients are equivalent to a single pass with B=256.
import torch
import torch.nn as nn
model = nn.Sequential(nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10))
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
criterion = nn.CrossEntropyLoss()
# Simulate batch_size=256 using accumulation_steps=8 × micro_batch=32
accumulation_steps = 8
optimizer.zero_grad() # clear gradients once at the start
for step, (X_batch, y_batch) in enumerate(train_loader):
# Forward + backward on each micro-batch
logits = model(X_batch)
loss = criterion(logits, y_batch)
loss = loss / accumulation_steps # scale loss so gradients average correctly
loss.backward() # accumulate gradients
# Only step the optimizer after accumulation_steps micro-batches
if (step + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad() # now clear gradientsGradient accumulation is how large language models are trained on hardware that can't fit a single massive batch. GPT-3 training used effective batch sizes in the millions of tokens — achieved by accumulating gradients across thousands of micro-batches before each optimizer step.
6 Learning Rate Schedules: Not Static Anymore
Using the same learning rate throughout training is almost never optimal. Think of navigating to a destination: when you're far away, you drive fast. As you get close, you slow down to park precisely. Training benefits from the same intuition: start with a higher learning rate (explore the loss landscape quickly), then reduce it (converge precisely).
A fixed learning rate creates two problems: if it's high enough to converge quickly at the start, it will overshoot and oscillate near the minimum. If it's small enough to converge precisely at the end, it's too slow at the start.
Step Decay: Halve the LR Every K Epochs
import torch
import torch.optim as optim
model = nn.Sequential(nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10))
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
# StepLR: multiply LR by gamma every step_size epochs
# After 30 epochs: lr = 0.1 * 0.1 = 0.01
# After 60 epochs: lr = 0.01 * 0.1 = 0.001
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
for epoch in range(90):
# ... training loop ...
scheduler.step() # call AFTER optimizer.step(), at end of epoch
current_lr = scheduler.get_last_lr()[0]
if epoch % 30 == 0:
print(f"Epoch {epoch}: lr = {current_lr:.5f}")
# Epoch 0: lr = 0.10000
# Epoch 30: lr = 0.01000
# Epoch 60: lr = 0.00100Cosine Annealing: Smooth Decay Following a Cosine Curve
Step decay creates abrupt LR drops. Cosine annealing smoothly reduces the LR following a cosine curve — starting at α_max, gently decreasing, and ending near zero. This gives better convergence because there's no sudden jump in step size.
import torch.optim as optim
import numpy as np
optimizer = optim.Adam(model.parameters(), lr=0.001)
# CosineAnnealingLR: decays LR from initial_lr to eta_min over T_max epochs
# Formula: lr_t = eta_min + 0.5 * (lr_max - eta_min) * (1 + cos(π * t / T_max))
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100, eta_min=1e-6)
# Show the LR schedule over 100 epochs
lrs = []
for epoch in range(100):
lrs.append(scheduler.get_last_lr()[0])
scheduler.step()
print(f"LR at epoch 0: {lrs[0]:.6f}") # near lr_max = 0.001
print(f"LR at epoch 25: {lrs[25]:.6f}") # descending
print(f"LR at epoch 50: {lrs[50]:.6f}") # near eta_min
print(f"LR at epoch 75: {lrs[75]:.6f}") # descending
print(f"LR at epoch 99: {lrs[99]:.6f}") # near eta_minCosine Annealing with Warm Restarts (SGDR)
An extension: instead of decaying to zero once, periodically "restart" the LR to its maximum value and let it decay again. Each restart gives the optimizer a chance to explore different regions of the loss landscape. Models trained this way often achieve better ensemble performance.
# CosineAnnealingWarmRestarts: T_0 = initial cycle length, T_mult = growth factor
scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(
optimizer, T_0=10, T_mult=2, eta_min=1e-6
)
# Restart at epochs: 10, 30 (10+20), 70 (30+40), 150 (70+80) — cycles double each timeOneCycleLR: Linear Warmup + Cosine Decay
The modern standard for many applications. Start very low, linearly increase to a maximum, then cosine-anneal to near zero — all in one training run. Used by fast.ai's "1cycle policy" which achieves state-of-the-art results with fewer epochs.
n_epochs = 30
steps_per_epoch = len(train_loader)
scheduler = optim.lr_scheduler.OneCycleLR(
optimizer,
max_lr=0.01, # peak learning rate
steps_per_epoch=steps_per_epoch,
epochs=n_epochs,
pct_start=0.3, # spend 30% of training warming up
anneal_strategy='cos', # cosine decay after peak
div_factor=25.0, # start LR = max_lr / 25
final_div_factor=1e4 # end LR = max_lr / (25 * 10000)
)
# Call scheduler.step() AFTER EACH BATCH (not epoch!) for OneCycleLR
for epoch in range(n_epochs):
for X_batch, y_batch in train_loader:
optimizer.zero_grad()
loss = criterion(model(X_batch), y_batch)
loss.backward()
optimizer.step()
scheduler.step() # ← per-batch updateFor most projects: start with CosineAnnealingLR. It's simple, smooth, and works well across architectures. OneCycleLR is excellent when you want to get good performance in fewer epochs. StepLR is the historical standard for ResNet-style image classification. For Transformers, use linear warmup + linear or cosine decay (covered in the next section).
7 Learning Rate Warmup: Stabilizing Early Training
At the very start of training, your model's weights are random. The gradients in the first few steps are large and chaotic — they don't yet reflect meaningful signal, just the noise of random initialization. Starting with a high learning rate at this point can send the parameters shooting off into terrible regions of the loss landscape before they've had a chance to orient themselves.
Warmup solution: Start with a tiny learning rate (say 1/100th of your target LR), and linearly increase it over the first few epochs until it reaches the intended value. By then, the weights have settled into a more meaningful regime and can handle larger update steps without diverging.
Warmup is now standard practice for Transformers. BERT uses 10,000 steps of linear warmup. GPT-3 uses a cosine schedule with a 375M-token warmup period. Without warmup, large Transformer models frequently diverge in early training.
import torch
import torch.optim as optim
class LinearWarmupCosineDecay:
"""
Learning rate schedule: linearly warm up for `warmup_epochs` epochs,
then cosine-anneal from `max_lr` to `min_lr` over `total_epochs - warmup_epochs`.
"""
def __init__(self, optimizer, warmup_epochs, total_epochs, max_lr, min_lr=1e-6):
self.optimizer = optimizer
self.warmup_epochs = warmup_epochs
self.total_epochs = total_epochs
self.max_lr = max_lr
self.min_lr = min_lr
self.current_epoch = 0
def step(self):
if self.current_epoch < self.warmup_epochs:
# Linear warmup: fraction of warmup complete
lr = self.max_lr * (self.current_epoch + 1) / self.warmup_epochs
else:
# Cosine annealing after warmup
progress = (self.current_epoch - self.warmup_epochs) / \
(self.total_epochs - self.warmup_epochs)
lr = self.min_lr + 0.5 * (self.max_lr - self.min_lr) * \
(1 + math.cos(math.pi * progress))
for param_group in self.optimizer.param_groups:
param_group['lr'] = lr
self.current_epoch += 1
return lr
import math
optimizer = optim.AdamW(model.parameters(), lr=1e-3)
schedule = LinearWarmupCosineDecay(
optimizer, warmup_epochs=5, total_epochs=50, max_lr=1e-3
)
# Show the LR trajectory
for epoch in range(50):
lr = schedule.step()
if epoch in [0, 1, 4, 5, 10, 25, 49]:
print(f"Epoch {epoch:2d}: lr = {lr:.6f}")The BERT paper uses 10,000 steps of linear warmup followed by linear decay over the remaining 990,000 steps. GPT-2 uses cosine decay with a 2000-step warmup. Without warmup, gradient norms in early Transformer training can exceed 100× their steady-state values, causing the LayerNorm and attention weights to collapse.
8 Gradient Clipping: Preventing Exploding Gradients
In deep networks — especially RNNs and Transformers — the gradient can sometimes become extremely large due to the chain rule multiplying many large values together. When this happens, the update step W ← W − α · ∇L is enormous, and the weights can jump to completely wrong values. This is called exploding gradients.
Imagine you're walking downhill and suddenly the slope becomes a cliff — you'd fall hundreds of meters in one step. Gradient clipping puts a safety harness on you: if the slope exceeds a maximum steepness, your step size is automatically limited.
How gradient clipping works: After computing gradients (loss.backward()) but before the optimizer step, check if the gradient norm exceeds a threshold. If it does, scale all gradients down so the norm equals the threshold. The direction is preserved, only the magnitude is clipped.
import torch
import torch.nn as nn
# A simple RNN-like training loop with gradient clipping
model = nn.GRU(input_size=128, hidden_size=256, num_layers=2, batch_first=True)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
MAX_GRAD_NORM = 1.0 # standard value for RNNs and Transformers
for epoch in range(50):
for X_batch, y_batch in train_loader:
optimizer.zero_grad()
# Forward pass
output, _ = model(X_batch)
loss = criterion(output[:, -1, :], y_batch) # use last timestep
# Backward pass — computes gradients
loss.backward()
# ─── Gradient clipping ────────────────────────────────────────────────
# Computes total gradient norm, then scales all gradients if norm > max_norm
# Returns the gradient norm BEFORE clipping (useful for logging)
grad_norm = torch.nn.utils.clip_grad_norm_(
model.parameters(), max_norm=MAX_GRAD_NORM
)
# ─────────────────────────────────────────────────────────────────────
optimizer.step()
if grad_norm > MAX_GRAD_NORM:
print(f" Gradient clipped: norm was {grad_norm:.2f} → clipped to {MAX_GRAD_NORM}")
# You can also clip by VALUE (rarely used):
# torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5)
# This clips each individual gradient element to [-clip_value, +clip_value]When to Use Gradient Clipping
- Always for RNNs and LSTMs: The sequential nature of RNNs means gradients are multiplied through many timesteps — explosion is common. max_norm=1.0 or 5.0 is standard.
- Always for Transformers: Attention mechanisms can produce large gradient spikes. max_norm=1.0 is the universal default.
- Optional for CNNs/MLPs: Less commonly needed, but can help during early training or with very deep networks.
If your loss suddenly becomes NaN during training, the most likely cause is exploding gradients — the gradient norm became so large that the weight update caused overflow to infinity, which propagates as NaN. Add gradient clipping with clip_grad_norm_(model.parameters(), max_norm=1.0) and add monitoring: print(f"grad_norm: {grad_norm:.4f}") every few steps to watch for spikes before NaN occurs.
Real-World Spotlight: Training a Digit Classifier with All 3 GD Variants
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Subset
import time
# ── Setup ──────────────────────────────────────────────────────────────────
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
full_train = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_set = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
# Use a subset for speed comparison
train_subset = Subset(full_train, range(5000))
test_loader = DataLoader(test_set, batch_size=256, shuffle=False)
def build_mlp():
return nn.Sequential(
nn.Flatten(),
nn.Linear(784, 256), nn.ReLU(),
nn.Linear(256, 128), nn.ReLU(),
nn.Linear(128, 10)
)
def train_and_evaluate(batch_size, n_epochs=10, use_scheduler=False, label=""):
model = build_mlp()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
criterion = nn.CrossEntropyLoss()
loader = DataLoader(train_subset, batch_size=batch_size, shuffle=True)
scheduler = None
if use_scheduler:
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=n_epochs)
epoch_losses = []
start_time = time.time()
for epoch in range(n_epochs):
model.train()
running_loss = 0.0
for X_b, y_b in loader:
optimizer.zero_grad()
loss = criterion(model(X_b), y_b)
loss.backward()
optimizer.step()
running_loss += loss.item()
if scheduler:
scheduler.step()
epoch_losses.append(running_loss / len(loader))
elapsed = time.time() - start_time
# Evaluate on test set
model.eval()
correct = 0
with torch.no_grad():
for X_b, y_b in test_loader:
correct += (model(X_b).argmax(1) == y_b).sum().item()
acc = correct / len(test_set) * 100
print(f"{label:40s} | Final loss: {epoch_losses[-1]:.4f} | Test acc: {acc:.1f}% | Time: {elapsed:.1f}s")
return epoch_losses
# Run all configs
print("Training 5000 MNIST samples for 10 epochs:\n")
l1 = train_and_evaluate(5000, label="(1) Batch GD (B=5000)")
l2 = train_and_evaluate(1, label="(2) True SGD (B=1) ")
l3 = train_and_evaluate(32, label="(3) Mini-batch (B=32) ")
l4 = train_and_evaluate(32, use_scheduler=True, label="(4) Mini-batch (B=32) + CosineAnneal")The results tell the full story: Batch GD is slow (18s) and gets the worst accuracy because it only makes 10 updates per epoch. True SGD makes 5000 updates per epoch and gets decent accuracy, but takes 22x longer. Mini-batch is fastest AND most accurate. Adding cosine annealing pushes accuracy another 1.2 points "for free" — the only change was adding two lines of scheduler code.
Quick Check
✍️ Practice Exercises
- Implement gradient accumulation for a mini-batch setup. Train a model with
batch_size=16andaccumulation_steps=4, and verify that the loss trajectory matches training withbatch_size=64directly. - Write a function that plots the LR trajectory for StepLR, CosineAnnealingLR, and OneCycleLR on the same axes over 100 epochs. Use
matplotliband label each curve. - Add gradient norm logging to a training loop: record the gradient norm (before clipping) at each step. Then add
clip_grad_norm_(model.parameters(), 1.0)and plot the norms — with and without clipping — on the same chart. - Implement the warmup scheduler from scratch as a custom
torch.optim.lr_scheduler.LambdaLR. LambdaLR takes a function that maps epoch → LR multiplier. Verify it matches the manual implementation from this lesson.
📚 Primary Sources for This Lesson
PyTorch LR Scheduler Documentation — complete reference for all schedulers with formulas and examples.
Cyclical Learning Rates for Training Neural Networks (Smith, 2015) — the paper that introduced the 1cycle policy and warm restarts.
On Large-Batch Training for DL (Keskar et al., 2016) — the classic paper explaining why large batches converge to sharp minima with worse generalization.