Gradient Descent & Its Variants
Comparing Batch, Stochastic, and Mini Batch Gradient Descent optimization strategies.
What is Gradient Descent?
Gradient Descent finds parameter weights $\theta$ that minimize a loss function $J(\theta)$ by taking steps proportional to the negative gradient:
$$\theta_{\text{new}} = \theta_{\text{old}} - \eta \nabla_{\theta} J(\theta)$$
How many data samples should we process before computing $\nabla_{\theta} J(\theta)$ and updating weights?
┌──────────────────────────┬──────────────────────────┬──────────────────────────┐
│ 1. BATCH GRADIENT DESCENT│ 2. STOCHASTIC (SGD) │ 3. MINI BATCH GD │
├──────────────────────────┼──────────────────────────┼──────────────────────────┤
│ Uses ALL N samples. │ Uses 1 single sample. │ Uses B samples (32-256). │
│ Smooth deterministic path│ High gradient noise. │ Industry standard! │
│ Slow on large datasets. │ Fast updates, poor GPU. │ Vectorized GPU speed. │
└──────────────────────────┴──────────────────────────┴──────────────────────────┘
1. Batch Gradient Descent
Computes loss and gradients across 100 percent of the training dataset:
$$\nabla_{\theta} J(\theta) = \frac{1}{N} \sum_{i=1}^N \nabla_{\theta} L(f(x^{(i)}), y^{(i)})$$
- Pros: Smooth, stable convergence path directly toward a local minimum.
- Cons: If you have 10 million samples, computing 10 million forward passes just to take one single step is painfully slow and exhausts RAM.
2. Stochastic Gradient Descent (SGD - Batch Size = 1)
Computes loss and updates weights after every single sample:
$$\theta = \theta - \eta \nabla_{\theta} L(f(x^{(i)}), y^{(i)})$$
- Pros: Updates weights immediately. High gradient noise helps jump out of shallow local minima.
- Cons: Extremely noisy steps that bounce around wildly. Fails to utilize modern GPU parallel tensor hardware.
3. Mini-Batch Gradient Descent (The Production Standard)
Splits data into small mini batches of size $B$ (typically $B = 32, 64, 128, 256$):
$$\theta = \theta - \eta \frac{1}{B} \sum_{i=1}^B \nabla_{\theta} L(f(x^{(i)}), y^{(i)})$$
Mini Batch Optimization Pipeline:
Dataset (1,000,000 Samples) ──► Split into Mini Batches (Batch Size 64)
- Batch 1: Compute 64 Forward Passes ──► Compute Gradient ──► Update Weights!
- Batch 2: Compute 64 Forward Passes ──► Compute Gradient ──► Update Weights!
- Pros:
- GPU Speed: Matrix operations on batch tensor blocks (e.g. $[64 \times 768]$) saturate GPU cores efficiently.
- Healthy Gradient Noise: Small mini batch noise acts as implicit regularization to find flat, robust minima.
Say this out loud
Batch Gradient Descent computes gradients across the full dataset before taking one step, making it slow on large data. Stochastic Gradient Descent updates weights after every single sample, creating noisy updates that miss GPU parallel speed. Mini Batch Gradient Descent updates weights every 32 to 256 samples, combining fast GPU matrix parallelization with helpful gradient noise.
Followups to expect
- How does learning rate interact with mini batch size? When increasing mini batch size by factor K, scale the learning rate up by K (Linear Scaling Rule) alongside warmup to maintain stable weight updates per epoch.
- What is an Epoch vs an Iteration? One Epoch is one complete pass through the entire dataset. One Iteration (Step) is a single weight update on one mini batch.
Check yourself
Why is Batch Gradient Descent impractically slow for training machine learning models on 10 million data points?