Stochastic gradient descent

Full gradient descent needs the gradient of the loss over the entire dataset at every step — an $O(n)$ sweep per update. Modern datasets made that absurd, and the fix changed the character of the algorithm: estimate the gradient from a random minibatch, and accept that every step is now a draw from a distribution. This lesson treats SGD as what it is — a stochastic process — because that is the only frame in which its behaviour (fast early, noisy late, schedule-dependent forever) makes sense.

The estimator

The training loss is a finite sum, $f(x) = \frac{1}{n} \sum_{i=1}^n f_i(x)$, one term per example. Sampling a minibatch $B$ of size $b$ uniformly and averaging gives the minibatch gradient

which is unbiased — $\mathbb{E}[g_B(x)] = \nabla f(x)$ — with covariance shrinking linearly in the batch size:

where $\Sigma(x)$ is the covariance of a single-example gradient. Two consequences before any dynamics: the direction is right on average, and averaging four times more examples only halves the noise — standard-error scaling, $\sigma/\sqrt{b}$. Batch size buys noise reduction at a steep, sublinear exchange rate.

The update is otherwise unchanged:

That decomposition — gradient step plus zero-mean kick scaled by the learning rate — is the whole algorithm. Everything below is about the tug-of-war between the two terms.

idx = rng.permutation(n)                    # each epoch: shuffle, then slice
for s in range(0, n, b):
    batch = idx[s : s + b]
    x -= lr(t) * grad(x, batch)             # unbiased, noisy, cheap
    t += 1

The noise ball: what a constant learning rate converges to

Far from the optimum, the true gradient dwarfs the noise and SGD behaves like slightly drunk gradient descent — early progress is nearly as fast, at $1/b$-th the cost per step. Near the optimum the gradient vanishes but the noise does not, and the two terms balance. For a $\mu$-strongly convex, $L$-smooth loss with (bounded) gradient-noise variance $\sigma^2/b$, constant-$\eta$ SGD satisfies, in the limit,

— it does not converge. It reaches a stationary distribution: a fuzzy cloud around the minimum, the noise ball, whose radius scales with $\eta$ and $\sigma^2/b$. On the loss curve this is unmistakable: a fast drop, then a plateau that no amount of further training lowers. The plateau is not a bug, a bad init, or a local minimum — it is the equilibrium radius of a stochastic process, and the only levers that lower it are the ones in the formula: smaller $\eta$, bigger $b$, or less noise.

Schedules: how to shrink the ball

If a constant $\eta$ ends in a ball of radius $\propto \eta$, the classical move is to decay it. The Robbins–Monro conditions (1951 — SGD is older than the transistor) say the schedule must satisfy

— enough total step to travel anywhere, but square-summable so the accumulated noise is finite. $\eta_t = \eta_0 / t$ qualifies; $\eta_t = \eta_0/\sqrt{t}$ fails the second condition but achieves the optimal $O(1/\sqrt{t})$ convex rate with averaging, and is what practice actually approximates. The realities of deep learning added an engineering layer on top:

The unifying logic is exactly the noise ball: large $\eta$ early to cross the landscape while gradients dominate; small $\eta$ late to shrink the equilibrium cloud once noise dominates. Every schedule is a policy for spending that budget.

Batch size is not a free parameter

The variance formula tempts a simple story — bigger batch, less noise, better — but the exchange rates matter:

• Compute per update grows linearly in $b$, while noise shrinks only as $1/\sqrt{b}$. Past a problem-dependent critical size, extra examples per batch stop buying wall-clock progress (Shallue et al. call this the end of “perfect scaling”).
• The linear-scaling rule — when you multiply $b$ by $k$, multiply $\eta$ by $k$ too (Goyal et al., training ImageNet at $b = 8192$) — falls straight out of the update: $k$ small noisy steps ≈ one $k$-times-larger, $\sqrt{k}$-less-noisy step, provided the gradient does not change much across them. It breaks near the stability ceiling $2/L$ from the previous lesson, which is why large-batch recipes need warmup.
• The noise is not purely adversarial. Small-batch SGD reliably finds solutions that generalize slightly better in many settings (Keskar et al.); one mechanism is that the noise ball cannot fit inside sharp, narrow minima, so the process preferentially settles in wide, flat basins. The full flat-vs-sharp story — and its caveats — is picked up in the non-convex landscapes lesson, but the operational summary stands: the noise is a regularizer you were given for free, and “reduce it to zero” is not automatically the goal.

Exercises