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 and to decorrelate consecutive mini-batches. Mini-batches drawn from a shuffled epoch are not independent — within an epoch they partition the data, so they are negatively correlated, but each is an unbiased estimate of the full gradient, which is what the convergence theory needs.

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.

Knowledge check#

Browser lab: GD vs SGD noise tradeoff#

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()

print("Notice: full-batch GD is smooth but makes one update per epoch; SGD B=1 is noisy yet far cheaper per step.")

Key takeaways#

• SGD replaces the full gradient with a mini-batch average that is unbiased but noisy, with variance scaling as $1/B$.
• Cost per update falls from $\mathcal{O}(nd)$ to $\mathcal{O}(Bd)$, so SGD takes roughly $n/B$ updates in the time full-batch GD takes one.
• A fixed step size leaves SGD oscillating in a ball of radius proportional to $\eta$; decaying $\eta$ is what buys convergence to the optimum.
• Batch size and learning rate trade off multiplicatively: the linear scaling rule lets you roughly double $\eta$ when you double $B$.
• Epoch-level loss is the meaningful curve; step-level loss is a noisy estimate from a single batch.
• The noise structure introduced here is the input to momentum (Week 7) and Adam (Week 8).

Further reading#

• Léon Bottou, Frank E. Curtis & Jorge Nocedal, Optimization Methods for Large-Scale Machine Learning (SIAM Review, 2018). Section 5 is the standard analysis of SGD, mini-batching, and step-size schedules.
• Ian Goodfellow, Yoshua Bengio & Aaron Courville, Deep Learning (MIT Press, 2016). §8.3 covers mini-batch and stochastic gradient methods.
• Herbert Robbins & Sutton Monro, "A Stochastic Approximation Method" (Annals of Mathematical Statistics, 1951). The original stochastic-approximation result behind SGD.