Week 6: Stochastic Gradient Descent

Finite-Sum Objectives#

Nearly all supervised learning objectives have the form:

$f(w) = \frac{1}{n} \sum_{i=1}^n f_i(w)$

where $f_i(w) = \ell(f_w(x_i), y_i)$ is the loss on the $i$-th training example. This is the ERMEmpirical Risk Minimization objective from Week 1.

The full gradient is $\nabla f(w) = \frac{1}{n} \sum_{i=1}^n \nabla f_i(w)$, which requires scanning all $n$ examples — one full pass through the dataset — for a single parameter update. When $n$ is large (millions or billions), this is impossibly slow.

SGD replaces the full sum with an estimate computed from a small, randomly chosen mini-batch.

Mini-Batch SGD#

At iteration $k$, sample a mini-batch $\mathcal{B}_k \subset \{1, \ldots, n\}$ of size $B$ uniformly at random (without replacement). The mini-batch gradient estimate is:

$g_k = \frac{1}{B} \sum_{i \in \mathcal{B}_k} \nabla f_i(w_k)$

The SGD update is then $w_{k+1} = w_k - \eta_k \, g_k$.

Two key properties of $g_k$:

1. Unbiased: $\mathbb{E}[g_k] = \nabla f(w_k)$. On average, the mini-batch gradient points in the right direction. The expectation is over the random choice of $\mathcal{B}_k$.

2. Noisy: $\mathrm{Var}(g_k)$ is proportional to $1/B$. A smaller batch has higher variance (noisier gradient estimate) but is cheaper to compute.

Computational cost per step:

• Full-batch GD: $\mathcal{O}(n d)$ to compute $\nabla f$, then $\mathcal{O}(d)$ to update
• Mini-batch SGD (size $B$): $\mathcal{O}(B d)$ to compute $g_k$, then $\mathcal{O}(d)$ to update

SGD makes roughly $n/B$ times as many updates in the time it takes GD to do one update.

Noise and Variance#

The variance of the mini-batch gradient estimate is:

$\mathrm{Var}(g_k) \approx \frac{\sigma^2}{B} \quad \text{where } \sigma^2 = \frac{1}{n} \sum_{i=1}^n \|\nabla f_i(w_k) - \nabla f(w_k)\|^2$

The term $\sigma^2$ measures how different individual example gradients are from the average. When examples are "similar" (homogeneous data), $\sigma^2$ is small and SGD works well even with $B=1$. When examples are very different (heterogeneous data), you need larger batches to reduce variance.

The batch-size tradeoff:

• Small $B$: cheap per iteration, many noisy updates per epoch → better exploration
• Large $B$: expensive per iteration, fewer but more accurate updates → reliable descent
• $B = n$: recovers full-batch GD (zero variance, maximum cost per step)

In practice, $B$ is typically chosen as a power of 2 (32, 64, 128, 256) to fit well with GPU memory and vectorised computation.

Step Sizes Under Noise#

With a fixed step size $\eta$, SGD does not converge to the exact optimum — it oscillates in a ball around $w^*$ whose radius is proportional to $\eta \sigma^2$. To converge exactly, we must decay the step size:

$\eta_k = \frac{\eta_0}{1 + \gamma k} \quad \text{or} \quad \eta_k = \frac{\eta_0}{\sqrt{k}}$

The $\mathcal{O}(1/\sqrt{k})$ schedule is common in theory; in practice, step-decay (halve $\eta$ every $T$ epochs) and cosine annealing (Week 9) are more popular.

A useful heuristic: for SGD with batch size $B$, a rule of thumb is $\eta \propto B$. Doubling the batch size approximately allows doubling the learning rate (because the gradient estimate is twice as accurate).

Epoch vs Step#

• Step (or iteration): one mini-batch processed, one parameter update
• Epoch: one full pass through the entire dataset ($\lceil n/B \rceil$ steps)

At the end of each epoch, the data is typically shuffled (random permutation) to avoid periodicity effects. Shuffling ensures each mini-batch is an independent random sample, which is important for the unbiasedness property.

Tracking loss per epoch is more meaningful than loss per step, because epoch-level loss estimates the true objective while step-level loss is computed on a single mini-batch and is noisier.

Browser lab#

Compare GD, SGD(B=1), and SGD(B=16) on a finite-sum quadratic objective. Plot loss vs epoch and observe the noise–convergence tradeoff.

import numpy as np
import matplotlib.pyplot as plt

rng = np.random.default_rng(42)
d, n = 3, 200

w_true = np.array([1.5, -0.8, 0.6])
X = rng.normal(size=(n, d))
y = X @ w_true + rng.normal(scale=0.3, size=n)

def loss(w, X, y):
    return 0.5 * np.mean((X @ w - y)**2)

def run_method(X, y, w0, n_epochs, method, B=None, eta=None):
    w = w0.copy()
    n = len(X)
    epoch_losses = []
    for ep in range(n_epochs):
        perm = rng.permutation(n)
        for start in range(0, n, B or n):
            idx = perm[start:start + (B or n)]
            g = X[idx].T @ (X[idx] @ w - y[idx]) / len(idx)
            w -= eta * g
        epoch_losses.append(loss(w, X, y))
    return epoch_losses

w0 = np.zeros(d)
n_epochs = 40
L = np.linalg.eigvalsh(X.T @ X / n).max()
eta = 1.0 / L

losses_gd = run_method(X, y, w0, n_epochs, "GD", B=n, eta=eta)
losses_sgd1 = run_method(X, y, w0, n_epochs, "SGD1", B=1, eta=eta*0.05)
losses_sgd16 = run_method(X, y, w0, n_epochs, "SGD16", B=16, eta=eta*0.5)

plt.figure(figsize=(8, 5))
plt.plot(losses_gd, label="GD (full batch)", linewidth=2)
plt.plot(losses_sgd1, alpha=0.8, label="SGD B=1 (noisy)", linewidth=1)
plt.plot(losses_sgd16, alpha=0.8, label="SGD B=16", linewidth=1.5)
plt.yscale("log")
plt.xlabel("Epoch"); plt.ylabel("Loss")
plt.title("GD vs SGD: Loss per Epoch")
plt.legend(); plt.grid(True, alpha=0.3)
plt.show()