Batches, Stochastic, and the Noise Ball

To compute the exact gradient of a neural-network loss, you need to run every training example through the model and average the per-example gradients. For GPT-2 trained on OpenWebText, that’s roughly 300 billion tokens. One full-batch step would take weeks on a single GPU.

You don’t have weeks.

So you approximate. Pick a small random batch of training examples, compute the gradient on just those, use it as a stand-in for the true gradient. Noisy, but the average of enough noisy estimates converges to the right thing. And each noisy estimate is much cheaper.

Formally: the loss is an average over training examples,

Pick a random mini-batch $B$ of $|B|$ examples. The stochastic gradient estimator is

Two facts about $\hat{\mathbf{g}}$:

• It’s unbiased: $\mathbb{E}[\hat{\mathbf{g}}] = \nabla L(\mathbf{w})$. On average, it’s the true gradient.
• It has variance proportional to $1/|B|$. Bigger batches → less noisy estimates.

Use $\hat{\mathbf{g}}$ in the GD update exactly where $\nabla L$ used to go:

That’s stochastic gradient descent.

Three sizes of batch, three tradeoffs:

• Full batch ($|B| = N$): the entire dataset. Zero variance per estimate. Infeasibly slow.
• Pure SGD ($|B| = 1$): one example at a time. Maximum noise, each step as cheap as possible.
• Mini-batch ($|B| \sim 32$–$2048$): somewhere in between. Hardware-friendly (GPUs vectorize well), variance tamed by $1/|B|$.

Mini-batch is what everyone actually uses. It’s the sweet spot. Pure SGD is too noisy; full batch is too slow. Typical: 64–512 for language models.

Here’s the consequence that surprises people. With a fixed learning rate $\eta$ and stochastic gradients, SGD does not converge to a point. It converges to a ball around the minimum, a region within which it keeps bouncing, because every step has residual noise.

Below: 120 SGD runs from the same starting point on the same quadratic, plotted as a scatter of endpoints. The blue circle is the theoretical noise-ball radius, $\eta \sigma_g / \sqrt{|B|}$. Adjust the controls and watch the cloud breathe.

Three ways to shrink the ball:

• Smaller learning rate $\eta$ → smaller residual noise per step.
• Bigger batch $|B|$ → less noisy estimates.
• Or: start with a large ball for exploration, shrink the ball (by shrinking $\eta$) when close to the minimum. That’s what a learning-rate schedule does (Lesson 10.5).

Important consequence: you cannot use “the gradient is zero” as a stopping criterion in practice; the stochastic gradient is never exactly zero even at the minimum.

Counterintuitively: the noise in SGD is helpful, not just a cost you pay for cheap gradients.

First, it helps escape saddle points. At a saddle, the exact gradient is zero and full-batch GD stalls forever. A noisy gradient estimate almost never points exactly at zero, so SGD gets nudged off the saddle and keeps moving.

Second, it acts as an implicit regularizer. SGD tends to converge to flat minima (wide, shallow basins where the loss is low over a large region). Those generalize better than sharp, narrow minima, which full-batch GD is happy to fall into. There’s a whole subfield of “SGD generalization” devoted to explaining this empirical fact.

This is why “just use full-batch if you can afford it” turns out to be wrong advice for deep learning. You want some noise. You just don’t want the noise from a mini-batch of 1.

Because the gradient never exactly zeros out, you need practical stopping criteria:

• Loss plateau: track a moving average of training loss. When it stops decreasing for $N$ evaluations, stop.
• Validation metric plateau: check a held-out set periodically; stop when it stops improving (this is early stopping, the most important criterion for generalization).
• Fixed budget: run for $N$ steps or $T$ hours, then stop. Works when you’ve tuned the run before.

In a real training loop, you use some combination of all three. nanoGPT uses “fixed step budget with a cosine schedule.” You’ll see why in Lesson 10.5.