Stochastic Gradient Descent

Suppose the goal is to minimize an empirical risk defined as an average loss over a training set of $n$ examples, ``` L(θ) = (1/n) Σ_{i=1}^{n} ℓ(θ; x_i, y_i) ``` Full-batch gradient descent computes the exact gradient $\nabla L(\theta) = (1/n)\sum_i \nabla \ell(\theta; x_i, y_i)$ at every step, which requires touching all $n$ examples. Stochastic gradient descent instead draws one example (or a small random subset) $i_t$ at iteration $t$, and updates using only that example's gradient, ``` θ_{t+1} = θ_t − η_t ∇ℓ(θ_t; x_{i_t}, y_{i_t}) ``` This single-example gradient is a random variable whose expectation, over the random choice of $i_t$, equals the true gradient $\nabla L(\theta_t)$ — it is an unbiased estimator. That single algebraic fact is the entire justification for the method: on average, the noisy step points in the same direction as the expensive exact step, so accumulating many noisy steps behaves like accumulating many exact steps, provided the noise is controlled.

Mini-batch SGD generalizes this by averaging the gradient over a small batch $B_t$ of size $b$ (commonly 32–4096) rather than a single example, ``` θ_{t+1} = θ_t − η_t · (1/b) Σ_{i ∈ B_t} ∇ℓ(θ_t; x_i, y_i) ``` This is still an unbiased estimator of $\nabla L(\theta_t)$, but its variance is reduced by a factor of roughly $1/b$ compared to single-example SGD (for i.i.d. sampling with replacement), since averaging $b$ independent noisy estimates shrinks variance linearly in $b$. Mini-batching is the practical default because it exploits vectorized hardware (GPUs/TPUs) that computes a batch of gradients almost as fast as a single one, while still updating far more frequently than full-batch descent. The key conceptual shift from plain gradient descent is this: SGD does not treat $L(\theta)$ as a fixed deterministic function to be minimized exactly at every step. It treats optimization as a stochastic approximation problem — estimating the root of $\nabla L(\theta) = 0$ from a stream of noisy gradient observations — a framework formalized by the Robbins–Monro algorithm long before it became the default training method for neural networks.

The mechanics of SGD are simple, but the two design choices that matter most are (1) how batches are sampled and (2) how the learning rate is scheduled over time. A well-posed implementation shuffles the dataset once per epoch and processes it in sequential mini-batches — this approximates i.i.d. sampling without the overhead of true random sampling with replacement. ``` PSEUDOCODE: Mini-batch Stochastic Gradient Descent Input: dataset D = {(x_1,y_1), ..., (x_n,y_n)} initial parameters θ_0 learning rate schedule η_t batch size b number of epochs E loss function ℓ(θ; x, y) Output: trained parameters θ θ ← θ_0 t ← 0 for epoch = 1 to E: shuffle D uniformly at random # fresh random order each epoch partition D into batches B_1, ..., B_m of size b # m = ceil(n / b) for each batch B_k in the shuffled partition: g ← (1 / |B_k|) Σ_{(x_i,y_i) in B_k} ∇_θ ℓ(θ; x_i, y_i) # mini-batch gradient θ ← θ − η_t g # parameter update t ← t + 1 end for end for return θ ``` Two invariants are worth naming explicitly: 1. Unbiasedness invariant: at every step, $\mathbb{E}[g \mid \theta_t] = \nabla L(\theta_t)$, where the expectation is over the randomness of the batch draw. This holds because each example is drawn from (an approximation of) the uniform distribution over the dataset, and the average of per-example gradients is itself the empirical-risk gradient. 2. Diminishing-influence invariant: as training progresses, the learning rate $\eta_t$ is chosen (or decayed) so that the cumulative noise injected into the trajectory does not prevent convergence, while still allowing enough movement to make progress. This is formalized by the Robbins–Monro step-size conditions discussed below. The diagram below shows the full pipeline: how a large dataset is sliced into batches, how each batch produces a noisy gradient estimate, and how that estimate perturbs the parameter update relative to the true full-batch direction.

The lower half of the figure captures the qualitative difference from full-batch descent: the green path is smooth because it always follows the exact gradient, while the red path zig-zags because each step follows a noisy estimate — yet both paths reach the same neighborhood of the minimum, because the noise averages out over many steps.

Per-step time complexity. If computing the gradient of a single example costs $O(d)$ (where $d$ is the number of parameters, as in a linear model or the backward pass of a network with $d$ weights), then:

• Full-batch gradient descent costs $O(nd)$ per update, since it touches all $n$ examples.
• Mini-batch SGD with batch size $b$ costs $O(bd)$ per update.
• Single-example SGD ($b=1$) costs $O(d)$ per update.

This is the entire computational argument for SGD's scalability: to complete one epoch (one full pass over the data), full-batch GD performs a single update costing $O(nd)$, while mini-batch SGD performs $n/b$ updates each costing $O(bd)$, for the same total $O(nd)$ work per epoch — but it makes $n/b$ parameter updates instead of just one. Since progress toward the minimum is roughly proportional to the number of updates (not the amount of data examined per update, up to noise), SGD converges in far less wall-clock time even though the total floating-point operations per epoch are comparable.

Space complexity. Full-batch GD needs $O(n d)$ or at least $O(n)$ auxiliary memory if gradients across the whole dataset must be aggregated per step (or streamed with $O(d)$ working memory if implemented carefully), whereas mini-batch SGD only ever needs $O(bd)$ memory for the current batch's activations and gradients. Since $b \ll n$ typically, SGD has a dramatically smaller memory footprint, which is often the deciding factor for training on GPUs with limited memory.

Correctness — why noisy steps still converge. The classical convergence guarantee for SGD on a convex, $L$-smooth objective with bounded gradient variance $\sigma^2$ rests on the Robbins–Monro conditions on the learning rate schedule: ``` Σ_t η_t = ∞ (steps must be large enough, cumulatively, to reach any point) Σ_t η_t^2 < ∞ (steps must shrink fast enough to damp the noise) ``` A schedule like $\eta_t = c/t$ or $\eta_t = c/\sqrt{t}$ satisfies these. The proof sketch is a stochastic Lyapunov / supermartingale argument: define $D_t = \mathbb{E}\|\theta_t - \theta^*\|^2$, the expected squared distance to a minimizer $\theta^*$. Using the update rule and unbiasedness, ``` D_{t+1} = D_t − 2η_t E[(θ_t − θ)ᵀ ∇L(θ_t)] + η_t² E‖g_t‖² ``` Convexity gives $(\theta_t - \theta^*)^\top \nabla L(\theta_t) \ge L(\theta_t) - L(\theta^*) \ge 0$, and bounded variance/gradient norm gives $\mathbb{E}\|g_t\|^2 \le \sigma^2 + G^2$ for constants $\sigma, G$. So: ``` D_{t+1} ≤ D_t − 2η_t (L(θ_t) − L(θ)) + η_t² (σ² + G²) ``` Summing over $t$ and using $\sum \eta_t^2 < \infty$ bounds the accumulated noise term, while $\sum \eta_t = \infty$ ensures the negative (progress) term can drive $\mathbb{E}[L(\theta_t) - L(\theta^*)]$ to zero along a subsequence. Formally this yields convergence in expectation of the running average of iterates, $\mathbb{E}[L(\bar\theta_T) - L(\theta^*)] = O(1/\sqrt{T})$ for convex, non-smooth objectives with a suitably tuned schedule — slower than the $O(1/T)$ rate of full-batch gradient descent on the same class of problems, but each SGD step is $n/b$ times cheaper, so it reaches a given accuracy in less total work whenever $n$ is large.

The honest tradeoff. SGD's per-iteration convergence rate is worse than full-batch GD's — noisy gradients simply cannot drive the error down as fast, iteration for iteration, as exact gradients can. What makes SGD win in practice is that each of its iterations is vastly cheaper, so more iterations fit in the same wall-clock budget; the crossover point favors SGD whenever $n$ is large, which is essentially always true in modern machine learning. A second, more subtle correctness point: mini-batch SGD does not converge to the exact minimizer with a constant learning rate — it converges to a stationary distribution of iterates fluctuating in a neighborhood of size $O(\eta \sigma^2)$ around $\theta^*$, unless the rate is decayed to zero. Practitioners often accept this residual noise floor deliberately, since it also has a regularizing effect that improves generalization.

Consider a toy 1-D least-squares problem to make the noise concrete. Suppose the true model is $y = 2x$, and the training set has $n = 4$ points:

The model is $\hat{y} = \theta x$, and the per-example squared loss is $\ell_i(\theta) = (\theta x_i - y_i)^2$, with per-example gradient $\nabla \ell_i(\theta) = 2 x_i (\theta x_i - y_i)$. The full-batch gradient at $\theta$ is the average over all four points: ``` ∇L(θ) = (1/4) Σ_i 2 x_i (θ x_i − y_i) ``` Start at $\theta_0 = 1.0$, learning rate $\eta = 0.02$ (constant, for clarity of the trace).

Full-batch step at $\theta_0 = 1.0$: compute each per-example gradient.

Full-batch gradient: $\nabla L(\theta_0) = (-2.2 - 7.2 - 19.8 - 29.6)/4 = -58.8/4 = -14.7$. Full-batch update: $\theta_1 = 1.0 - 0.02 \times (-14.7) = 1.0 + 0.294 = 1.294$.

Single-example SGD, same starting point $\theta_0 = 1.0$, drawing examples in order $i = 2, 4, 1, 3$ (one epoch, mimicking a shuffle):

After one epoch, single-example SGD lands at $\theta \approx 1.82$, compared to the full-batch method's single step to $\theta_1 = 1.294$ — SGD made four updates in the time full-batch made one, and moved further toward the true value $\theta^* = 2$ despite each individual step being noisier and based on incomplete information. This illustrates the core tradeoff numerically: each SGD step is a worse estimate of the true gradient direction than the full-batch gradient, but because four cheap steps fit in the budget of one expensive step, SGD's net progress per unit of computation is larger. Continuing for more epochs with a decaying learning rate would let $\theta$ settle near $2.0$, while the noise from single-example estimates would keep it oscillating slightly if the rate stayed constant — exactly the noise-floor phenomenon predicted by the convergence analysis above.

The theoretical seed of SGD lies in statistics and stochastic approximation, decades before its widespread use in deep learning. Herbert Robbins and Sutton Monro formalized the general problem of finding a root of a function known only through noisy observations, giving the first rigorous convergence proof for a Robbins–Monro-type recursive stochastic algorithm — the mathematical ancestor of every modern SGD convergence proof.

The Widrow–Hoff least-mean-squares (LMS) rule of 1960 is arguably the first widely used single-example stochastic gradient method in engineering practice, applied to adaptive signal filters long before "machine learning" was a common phrase. When backpropagation was popularized in the 1980s, mini-batch and online (single-example) updates were the natural computational fit for the hardware of the era, since full-batch computation over large datasets was infeasible. The 2012 AlexNet result cemented mini-batch SGD (with added momentum, foreshadowing Momentum Methods) as the default training procedure for deep networks, a status it has retained even as more sophisticated adaptive-rate variants have emerged.

Mini-batch SGD (and its refinements) is the training engine behind essentially every large-scale machine learning system in production. Deep learning frameworks such as PyTorch and TensorFlow implement SGD as a first-class optimizer primitive, and virtually every large neural network — convolutional networks for vision, transformers for language, recommendation systems at web scale — is trained by some descendant of the mini-batch SGD loop. Search-ranking and ad-click prediction models, trained on data streams too large to ever fit in memory, use online variants of SGD that treat the entire training process as one continuous stochastic optimization. Federated learning systems, which train models across many devices without centralizing data, are built on local mini-batch SGD updates aggregated across a network, since the fundamentally cheap, streaming nature of SGD makes it well suited to decentralized and privacy-sensitive settings. Understanding SGD's cost-noise tradeoff is also the prerequisite for understanding why practitioners add momentum, adaptive learning rates (as in Adam or RMSProp), and learning rate warmup/decay schedules — all of these techniques exist specifically to manage the noise that plain SGD introduces, which is exactly the direction that Momentum Methods develops further.

Ask the tutor: How does adding momentum change the noisy zig-zag trajectory of SGD, and why does it accelerate convergence in directions of consistent gradient sign?