Stochastic and Mini-Batch Gradient Descent
Imagine training an image classifier on 14 million labeled photographs. Computing the exact gradient of the loss requires a forward and backward pass over every image. This must happen before you can update the weights even once. If each pass takes a few milliseconds per image, a single gradient step costs hours. You may also need tens of thousands of steps to converge. This is not a hypothetical inconvenience. It is the central computational bottleneck that made deep learning impractical for decades. A simple statistical trick eventually made it tractable. The trick is to stop insisting on the exact gradient. Instead of averaging the loss gradient over the entire dataset, we average it over a small random subset. This subset is called a mini-batch. We treat this noisy estimate as good enough to take a step. This substitution, stochastic gradient descent (SGD) and its practical workhorse variant mini-batch gradient descent, is arguably the single algorithmic decision that makes large-scale deep learning possible. It converts an intractable per-step cost into a cheap, parallelizable, GPU-friendly operation. This comes at the price of introducing noise into the optimization trajectory. This article builds directly on Gradient Descent for Deep Learning. There, the update rule and the loss landscape were introduced assuming access to the true gradient. Here we relax that assumption. We ask: what happens when every gradient is a random variable rather than a fixed vector? We show that the noisy gradient is still an unbiased estimator of the true gradient. We explain why this unbiasedness is the mathematical bedrock that lets the algorithm still converge. We also show how the choice of batch size trades off gradient variance against computational efficiency and hardware utilization. By the end, you will understand precisely why practitioners almost never compute the full-batch gradient. You will learn how to reason about the variance-cost tradeoff when choosing a batch size. You will also understand why mini-batch SGD is the substrate on top of which every modern optimizer is built — Momentum and Nesterov Accelerated Gradient, AdaGrad and RMSProp, and Adam.
- Full-batch gradient descent computes the exact gradient over the whole dataset; this is accurate but computationally prohibitive at scale.
- Stochastic gradient descent (SGD) estimates the gradient from a single random example; mini-batch SGD estimates it from a small random subset (typically 32–8192 examples).
- The mini-batch gradient is an unbiased estimator of the true gradient — its expected value equals the full-batch gradient, which is why the algorithm still converges.
- Batch size controls a variance/throughput tradeoff: larger batches give lower-variance gradients but cost more compute per step and use hardware parallelism better.
- Per-epoch cost is the same order across batch sizes, but wall-clock time, generalization, and memory usage differ sharply with batch size.
- Mini-batch SGD is the base algorithm that momentum, adaptive learning rates, and modern optimizers all extend.
Replacing the exact gradient with a cheap, unbiased random estimate over a small batch lets each optimization step be nearly free, and averaging over many noisy steps still drives the parameters toward a minimum.
The idea
In Gradient Descent for Deep Learning, the parameter update is ``` θ ← θ − η ∇θ L(θ) ``` where the loss `L(θ) = (1/N) Σᵢ ℓ(f(xᵢ; θ), yᵢ)` is an average over all `N` training examples, and `∇θ L(θ)` is the full-batch gradient: an exact sum over every example's per-example loss gradient. Computing it requires one full forward/backward pass through the entire dataset for every single parameter update.
Stochastic gradient descent (SGD) replaces this exact average with the gradient of the loss on a single, uniformly random training example `(x_i, y_i)`: ``` θ ← θ − η ∇θ ℓ(f(xᵢ; θ), yᵢ), i ~ Uniform{1, …, N} ``` This is a stochastic approximation to the true gradient — a random vector whose value changes every time you draw a new example, but whose statistical properties are well controlled.
Mini-batch gradient descent interpolates between these two extremes. At each step we sample a mini-batch `B` of `m` examples (with `1 ≪ m ≪ N`, typically drawn without replacement from a shuffled dataset) and compute ``` g = (1/m) Σ_{i∈B} ∇θ ℓ(f(xᵢ; θ), yᵢ) θ ← θ − η g ``` Three regimes fall out of this single family, indexed by batch size `m`:
- `m = 1`: pure stochastic gradient descent — maximal noise, minimal per-step compute.
- `1 < m < N`: mini-batch SGD — the practical default in essentially all modern training.
- `m = N`: full-batch gradient descent — zero gradient noise, maximal per-step compute.
The crucial mathematical fact that justifies all of this is unbiasedness. If the mini-batch `B` is drawn uniformly at random (with or without replacement, for large `N` the distinction barely matters), then ``` E_B[ g ] = E_B[ (1/m) Σ_{i∈B} ∇θ ℓᵢ ] = (1/N) Σᵢ ∇θ ℓᵢ = ∇θ L(θ) ``` The mini-batch gradient is a random variable whose expected value is exactly the true full-batch gradient. Every noisy step points, on average, in the correct direction. The noise averages out over many steps even though any single step may be a poor estimate. This is the entire theoretical foundation on which stochastic optimization rests, and it is why "using less data per step" does not mean "using a wrong direction" — it means "using a randomized direction with the right mean and controllable variance."
How it works
The mechanism has four moving parts: shuffling, batching, gradient estimation, and the update rule. An epoch is one full pass over the shuffled dataset, subdivided into `N/m` mini-batches. ``` Algorithm: Mini-Batch Stochastic Gradient Descent Input: dataset D = {(x_1,y_1), ..., (x_N,y_N)}, batch size m, learning rate η, number of epochs E, initial parameters θ_0 Output: trained parameters θ θ ← θ_0 for epoch = 1 to E: D_shuffled ← random_permutation(D) # break correlation between adjacent samples B_1, ..., B_{N/m} ← split D_shuffled into consecutive chunks of size m for k = 1 to N/m:
# --- forward pass --- for each (x_i, y_i) in B_k: ŷ_i ← f(x_i; θ) ℓ_i ← loss(ŷ_i, y_i)
# --- gradient estimate (unbiased) --- g ← (1/m) * Σ_{(x_i,y_i) in B_k} ∇θ ℓ_i # backward pass, one batched op
# --- parameter update --- θ ← θ − η * g end for end for return θ ``` Notes on correctness-critical details:
- Shuffling each epoch is not cosmetic. Without it, mini-batches are drawn in a fixed, non-random order, correlated batches produce biased local gradient estimates, and the unbiasedness argument above (which relies on uniform sampling) breaks down. Re-shuffling every epoch approximates independent uniform sampling with replacement well enough in practice.
- Batch composition matters, not batch size alone. A batch of `m` examples from the same class gives an extremely biased local estimate even though the expectation over all possible batches is still unbiased — shuffling is what guarantees batches are representative in expectation.
- The forward and backward pass over a mini-batch is a single batched tensor operation on modern hardware (GPU/TPU), computed via matrix multiplication across the batch dimension. This is the computational reason mini-batches of size 32–512 are dramatically faster per example processed than single-example SGD: hardware parallelism is wasted on `m = 1`.
- Last batch handling: if `N` is not divisible by `m`, the final mini-batch of an epoch has fewer than `m` examples; implementations either keep it (slightly higher-variance step) or drop it — a minor but real edge case.
Complexity & correctness
Per-step time complexity. Let `c` be the cost of one forward+backward pass for a single example (proportional to the network's FLOP count). A mini-batch of size `m` costs `O(m·c)` per step in raw arithmetic, but due to hardware parallelism (SIMD/GPU cores processing the batch dimension concurrently) the wall-clock cost per step often scales sublinearly in `m` up to the hardware's parallel-capacity limit, then linearly beyond it.
Per-epoch time complexity. Regardless of `m`, one epoch touches every example exactly once, so total arithmetic work per epoch is always `O(N·c)` — the same order for SGD, mini-batch, and full-batch. What changes is:
- Number of updates per epoch: `N/m`. Small `m` gives many noisy updates; large `m` gives few precise updates.
- Wall-clock time per epoch: smaller `m` means more sequential steps, each with fixed overhead (kernel launch, memory transfer) that doesn't amortize — so very small batches can be slower in practice despite identical FLOP counts.
- Space complexity: `O(m·d)` for storing a batch's activations for backpropagation, where `d` is the size of the largest intermediate layer. This is the direct, often binding, memory constraint on how large `m` can be on a given GPU.
Variance analysis — why noise shrinks with batch size. Let `∇ℓᵢ` denote the per-example gradient, with `Var[∇ℓᵢ] = σ²` (coordinate-wise, informally). For a mini-batch of `m` i.i.d. samples, ``` Var[ g ] = Var[ (1/m) Σᵢ ∇ℓᵢ ] = σ² / m ``` This is the standard reduction-in-variance-of-the-mean argument: averaging `m` independent noisy estimates shrinks variance by a factor of `m`, and the standard deviation (noise magnitude) shrinks by `√m`. This has two direct consequences: 1. To halve the gradient noise's standard deviation, you must quadruple the batch size — variance reduction has diminishing returns per unit of compute. 2. Since `g` is unbiased (`E[g] = ∇L(θ)`, shown above) with variance `σ²/m`, mini-batch SGD is exactly gradient descent plus zero-mean noise that shrinks as `m` grows. As `m → N`, `Var[g] → 0` and the method recovers exact full-batch gradient descent.
Correctness argument (why SGD still converges). Treat the parameter sequence as a stochastic process: ``` θ_{t+1} = θ_t − η g_t, E[g_t | θ_t] = ∇L(θ_t) ``` Under standard assumptions (L-smooth loss, gradients with bounded variance, appropriately decaying or small enough learning rate), a classical descent-lemma argument gives, for L-smooth `L`, ``` L(θ_{t+1}) ≤ L(θ_t) − η(1 − Lη/2)‖∇L(θ_t)‖² + (Lη²/2)·Var[g_t] ``` Taking expectations and summing over `t = 0..T-1`, the first term telescopes into a decrease in loss proportional to the sum of squared true gradient norms, while the variance term contributes an additive "noise floor" of order `η·σ²/m`. Rearranging shows the average squared gradient norm over `T` steps is bounded by ``` O(1/(ηT)) + O(η σ² / m) ``` This says two things precisely: (1) the deterministic part of the bound shrinks like ordinary gradient descent (`1/T` with fixed `η`), confirming the noisy method still makes genuine progress on average; (2) there is an irreducible noise floor set by `ησ²/m` that does not vanish as `T → ∞` unless `η` is decayed or `m` is increased. This is exactly why SGD with a constant learning rate does not converge to an exact minimum but instead oscillates in a small neighborhood of one — while full-batch GD (`m = N`, `σ_eff = 0`) can converge exactly. Practical schedules decay `η` over training precisely to shrink this noise floor.
Honest tradeoffs.
| Aspect | Small batch (e.g. 32) | Large batch (e.g. 4096) |
|---|---|---|
| Gradient variance | High | Low |
| Updates per epoch | Many | Few |
| Hardware utilization | Often poor (parallel units idle) | Excellent |
| Memory per step | Low | High |
| Generalization (empirical) | Often better — noise acts as implicit regularization | Can generalize worse without tuning (needs learning-rate warmup, more epochs) |
| Wall-clock per epoch | Can be slower (overhead-dominated) | Often faster, until memory-bound |
Worked example
Consider a toy 1-D linear regression: fit `ŷ = θx` to minimize squared error, with a dataset of `N = 6` points:
| index i | xᵢ | yᵢ | per-example loss ℓᵢ(θ) = (θxᵢ − yᵢ)² | ∇θℓᵢ = 2xᵢ(θxᵢ − yᵢ) |
|---|---|---|---|---|
| 1 | 1 | 2 | (θ−2)² | 2(θ−2) |
| 2 | 2 | 3 | (2θ−3)² | 4(2θ−3) |
| 3 | 1 | 1 | (θ−1)² | 2(θ−1) |
| 4 | 3 | 7 | (3θ−7)² | 6(3θ−7) |
| 5 | 2 | 5 | (2θ−5)² | 4(2θ−5) |
| 6 | 1 | 0 | θ² | 2θ |
Let `θ = 1` (current parameter), `η = 0.05`.
Full-batch gradient (average over all 6 points): ``` ∇ℓ₁=2(1−2)=−2, ∇ℓ₂=4(2−3)=−4, ∇ℓ₃=2(1−1)=0, ∇ℓ₄=6(3−7)=−24, ∇ℓ₅=4(2−5)=−12, ∇ℓ₆=2(1)=2 ∇L(θ) = (−2−4+0−24−12+2)/6 = (−40)/6 ≈ −6.667 θ_new = 1 − 0.05·(−6.667) = 1 + 0.333 = 1.333 ```
Mini-batch SGD, `m = 2`, batch `B = {i=4, i=5}` (drawn at random this step): ``` g = (∇ℓ₄ + ∇ℓ₅)/2 = (−24 + −12)/2 = −18 θ_new = 1 − 0.05·(−18) = 1 + 0.9 = 1.9 ``` This single step overshoots dramatically relative to the true direction — batch `{4,5}` happens to contain the two steepest, most "surprising" points, so its estimate has high variance and swings far from the full-batch answer.
Mini-batch SGD, `m = 2`, a different random batch `B = {i=1, i=6}`: ``` g = (∇ℓ₁ + ∇ℓ₆)/2 = (−2 + 2)/2 = 0 θ_new = 1 − 0.05·0 = 1.0 (no movement this step) ``` This batch happens to be "balanced" and produces zero net signal.
Checking unbiasedness directly. Enumerate all `C(6,2) = 15` possible batches of size 2, average their gradient estimates `g_B`, and you will find ``` (1/15) Σ_B g_B = ∇L(θ) ≈ −6.667 ``` exactly matching the full-batch gradient — confirming, on this concrete instance, that the mini-batch estimator is unbiased even though any individual batch (like `{4,5}` above) can be wildly off from the true gradient. Over many epochs with re-shuffling, these individual swings cancel on average, and `θ` drifts toward the same region full-batch GD would reach, just via a noisier, cheaper path.
Common pitfalls & traps
| Pitfall | Why it breaks things |
|---|---|
| Not shuffling data (or shuffling only once, not per epoch) | Batches become correlated across epochs; gradient estimates lose the i.i.d. structure that guarantees unbiasedness, causing systematic drift or cyclic behavior |
| Using the same fixed learning rate `η` regardless of batch size | Gradient variance scales as `σ²/m`; halving `m` without adjusting `η` can make training unstable, while doubling `m` without raising `η` can needlessly slow convergence |
| Treating batch size purely as a memory-fit parameter | Ignoring its effect on generalization — very large batches often need learning-rate warmup and more epochs to match small-batch test accuracy |
| Dropping the last partial batch inconsistently between train/eval code | Silent bias if the dropped examples are non-random (e.g., always the same tail of a sorted dataset before shuffling was added) |
| Confusing "epoch" with "step/iteration" when reading loss curves | Loss curves plotted against steps look very different from those plotted against epochs when comparing different batch sizes; comparing across `m` values requires the same x-axis |
| Assuming smaller batch always generalizes better | True often, but not universally — depends on architecture, learning-rate schedule, and normalization layers (e.g., BatchNorm statistics themselves depend on `m`) |
Pioneers & history
The stochastic approximation idea predates deep learning by decades, originating in statistics and control theory before being adopted wholesale by the neural-network community once datasets and models grew too large for full-batch methods.
Origins · 1 minHerbert Robbins and how this was discovered
Robbins and Monro's 1951 paper on finding the root of a noisy function established the mathematical template — decaying step sizes applied to noisy, unbiased measurements — that underlies every convergence proof for SGD used in deep learning today.
Where this matters
Every major deep learning framework — PyTorch, TensorFlow, JAX — implements mini-batch SGD as the default `DataLoader`/optimizer loop, and every subsequent optimizer builds on it rather than replacing it. Large-scale distributed training systems (data-parallel training across hundreds of GPUs/TPUs) partition each mini-batch across devices, compute local gradients, and average them — a direct hardware exploitation of the unbiasedness and additivity of the mini-batch gradient estimator. Batch normalization layers, learning-rate warmup schedules, and large-batch training recipes at companies training foundation models are all engineering responses to the variance/compute tradeoff analyzed above. The concept directly foreshadows Momentum and Nesterov Accelerated Gradient, which smooths the noisy mini-batch gradient sequence using an exponential moving average, and AdaGrad and RMSProp, which adapt the learning rate per-parameter based on the accumulated magnitude of these same noisy gradients — both are refinements layered on top of, not replacements for, mini-batch SGD.
Ask the tutor: How should I adjust the learning rate if I double my mini-batch size to better use my GPU's memory?
Why not just always use the largest batch size that fits in memory?
Larger batches reduce gradient variance and improve hardware utilization, but per-epoch arithmetic cost stays the same, and you get fewer parameter updates per epoch. Empirically, very large batches often need careful learning-rate warmup and more epochs to reach the same test accuracy as smaller batches, a phenomenon linked to the "generalization gap" in large-batch training. Small-batch noise can also act as an implicit regularizer, helping escape sharp minima that generalize poorly. In practice, batch size is chosen jointly with the learning rate and schedule, not maximized in isolation — the largest batch that fits in memory is a starting point for tuning, not automatically the best choice.
Is mini-batch SGD guaranteed to converge to a global minimum?
No, and neither is full-batch gradient descent on non-convex deep network losses. What is guaranteed, under smoothness and bounded-variance assumptions, is that the expected squared norm of the true gradient shrinks toward a noise floor as training proceeds, meaning the algorithm reliably finds points where the (expected) gradient is small — stationary points, which may be local minima or saddle points. In practice, the noise introduced by mini-batching is believed to help escape sharp local minima and saddle points, which is one reason mini-batch SGD often generalizes better than exact full-batch descent despite having weaker convergence guarantees on paper.
How is mini-batch SGD different from pure stochastic gradient descent with batch size 1?
They are the same family of algorithm with different batch sizes, not different algorithms. Batch size 1 (pure SGD) has the highest gradient variance (`σ²/1 = σ²`) and updates the parameters after every single example, giving very frequent but very noisy steps. Mini-batch SGD averages over `m` examples, reducing variance to `σ²/m` at the cost of computing `m` times more forward/backward work per update. In practice, batch size 1 rarely uses hardware parallelism efficiently, so mini-batch sizes of 32–1024 dominate in real training pipelines, reserving the term "SGD" informally to mean "any batch size less than the full dataset."
Does batch size affect how the learning rate should be chosen?
Yes, and this is one of the most common tuning mistakes. Because gradient variance scales as `σ²/m`, larger batches yield smoother, more reliable gradient estimates, which typically tolerate — and often require — a proportionally larger learning rate to make comparable progress per epoch. A common heuristic ("linear scaling rule") increases the learning rate roughly linearly with batch size within a moderate range, often paired with a warmup phase to avoid instability early in training. Blindly reusing a learning rate tuned for batch size 32 at batch size 4096 typically causes very slow convergence, since the effective step size relative to the (now much smaller) gradient noise is too small.
Why does mini-batch SGD work well with GPUs specifically?
GPUs are throughput-oriented processors with thousands of parallel cores optimized for the same operation applied to many data points simultaneously (SIMD-style parallelism). A mini-batch's forward and backward pass is naturally expressed as batched matrix multiplications, where the batch dimension maps directly onto this parallelism. Processing one example at a time (batch size 1) leaves most of the GPU idle, whereas batch sizes of 32–1024 keep cores fully occupied, often making a mini-batch of 256 examples nearly as fast per step as a batch of 32, up to memory limits. This hardware fit is a major practical reason mini-batch SGD, rather than pure SGD or full-batch GD, became the default.
How does mini-batch size interact with Batch Normalization layers?
Batch Normalization computes normalization statistics (mean and variance) across the examples in the current mini-batch, so its behavior is directly coupled to batch size. Very small batches produce noisy, unreliable batch statistics, which can destabilize training or hurt generalization, while very large batches give smoother statistics closer to the true dataset statistics. This is a concrete example of batch size affecting not just optimization noise but also architectural components downstream, and it motivates alternatives like Layer Normalization or Group Normalization when small batch sizes are unavoidable (e.g., due to memory constraints with large models).
Does increasing batch size always reduce the number of epochs needed to converge?
Not necessarily. Larger batches give lower-variance gradient estimates per step, which can mean each step is "worth more," but there are also fewer steps per epoch, and the diminishing-returns relationship (`Var ∝ 1/m`) means doubling batch size does not double the useful signal — it only reduces the standard deviation of the noise by a factor of `√2`. Beyond a certain point, adding more examples to a batch yields progressively smaller reductions in noise for a proportionally larger compute cost, so past a critical batch size, total training time (in wall-clock terms, accounting for compute) often stops improving or even worsens.