Expectation, variance and sampling

Problem

The loss you actually care about is the average over every input the model will ever see, which is an expectation over a distribution you can't enumerate. What you can do is draw a sample, a dataset, and average over that. Then, because even the dataset is too big to process at every step, you average over a minibatch of a few hundred examples.

Each of those steps replaces an exact quantity with an estimate. You need to know whether the estimate is right on average (is it biased?), how much it wobbles (its variance), and how both change with the number of samples. Those three questions decide batch sizes, learning rates and how many evaluation examples you need to trust a number.

Intuition

The expectation is the balance point of a distribution: if the probability mass were weights on a ruler, it is where the ruler balances. A fair die balances at 3.5, even though it never shows 3.5.

The variance is how spread out the mass is around that balance point: the average squared distance from it. Its square root, the standard deviation, is in the original units and is the typical size of a deviation.

Averages of independent samples are less spread out than single samples, because errors in different directions cancel. They cancel slowly, though: the typical error of an average of n samples is σ/√n, so 100 times the samples buys only 10 times the accuracy.

The law of large numbers in one picture: the sample mean converges to the expectation, and the one-standard-error funnel says how fast.

Mechanics

Expectation and linearity

E[X] = Σ_x x · p(x)    or    E[X] = ∫ x · p(x) dx

More generally, E[f(X)] = Σ f(x) p(x): average the function's values, weighted by probability. A loss is exactly this, E_{(x, y)}[ℓ(model(x), y)]. Expectation is linear: E[aX + bY + c] = aE[X] + bE[Y] + c, always, whether or not X and Y are independent. That one fact is behind almost every "unbiased" claim in machine learning, including the minibatch one below.

Variance and standard deviation

Var[X] = E[(X − E[X])²] = E[X²] − E[X]²     σ = √Var[X]

Variance is not linear. Shifting doesn't change it and scaling squares: Var[aX + c] = a² Var[X]. For independent X and Y, variances add: Var[X + Y] = Var[X] + Var[Y]. Combine the two and the average of n independent copies, X̄ = (X₁ + ⋯ + Xₙ)/n, has

Var[X̄] = n · Var[X] / n² = σ² / n     standard error = σ / √n

Covariance and correlation

Cov[X, Y] = E[(X − E[X])(Y − E[Y])]     ρ = Cov[X, Y] / (σ_X σ_Y)

Covariance is positive when X and Y tend to be above their means together, and negative when one tends to be high as the other is low. The correlation ρ rescales it to [−1, 1]. For the five students below (hours studied and score), Cov = 15 and ρ = 0.973. In general Var[X + Y] = Var[X] + Var[Y] + 2 Cov[X, Y], so correlated errors don't cancel, which is why an ensemble of near-identical models helps less than the count suggests.

Independent variables have zero covariance, but zero covariance does not mean independent: for X uniform on {−2, −1, 0, 1, 2} and Y = X², the covariance is 0 even though Y is a function of X. Covariance only measures linear association.

Monte Carlo estimates and the central limit theorem

Any expectation can be estimated by sampling: E[f(X)] ≈ (1/n) Σ f(xᵢ) with xᵢ drawn from p. The estimate is unbiased (its expectation is the true value, by linearity), and the law of large numbers says it converges as n grows. Its standard error is σ_f/√n, independent of the dimension of X, which is why Monte Carlo works for integrals over millions of parameters where grids are hopeless.

The central limit theorem adds the shape: for large n, the sample mean is approximately Normal(μ, σ²/n), whatever the original distribution looked like. That is what justifies "mean ± 2 standard errors" as a 95% interval on an evaluation score.

Exact distributions, computed by convolving the die's PMF with itself. Bars are scaled to each panel's tallest; the narrowing is the σ/√n law.

Why minibatch gradients are unbiased but noisy

With N training examples and per-example gradients g₁, …, g_N, the full gradient is their mean ḡ. A minibatch of b examples drawn uniformly at random gives ĝ = (1/b) Σ_{i ∈ batch} gᵢ. Every example is equally likely to be in the batch, so by linearity E[ĝ] = ḡ: unbiased. Its variance, sampling without replacement, is

Var[ĝ] = (σ_g² / b) · (N − b) / (N − 1)

where σ_g² is the variance of the per-example gradients. Doubling the batch halves the variance, so it halves the noise in each step's direction, at twice the compute. The (N − b)/(N − 1) factor only matters when the batch is a large fraction of the data, and makes the variance exactly 0 at b = N.

Worked example

Each dot is one minibatch's gradient. They balance on the full gradient at every batch size, and bigger batches pull them in.

Implementation

Estimating an expectation by sampling in PyTorch, reproducibly:

import torch

torch.manual_seed(0)                         # same draws every run
rolls = torch.randint(1, 7, (10_000,)).float()
rolls.mean()                                 # close to 3.5
rolls.var(correction=0)                      # close to 35/12 ≈ 2.917
rolls.std() / rolls.numel() ** 0.5           # standard error ≈ 0.017

with torch.random.fork_rng():                # sample without moving the main stream
    probe = torch.randn(3)

z = torch.distributions.Normal(2.0, 0.5)
z.mean, z.variance                           # (tensor(2.), tensor(0.2500))
z.sample((100_000,)).mean()                  # close to 2.0

var divides by n − 1 by default (Bessel's correction), which makes the sample variance an unbiased estimate of the population variance; pass correction=0 for the plain average of squared deviations used on this page.

Checking minibatch unbiasedness by brute force:

from itertools import combinations

x = [1, 2, 3, 4, 5, 6]
y = [2, 3, 7, 8, 9, 13]
w = 1.0
g = [2 * xi * (w * xi - yi) for xi, yi in zip(x, y)]   # per-example gradients
full = sum(g) / len(g)                                  # -31.0

for b in (1, 2, 3, 6):
    batches = [sum(c) / b for c in combinations(g, b)]
    mean = sum(batches) / len(batches)
    var = sum((v - mean) ** 2 for v in batches) / len(batches)
    print(b, mean, round(var, 2))   # mean is -31.0 every time; var 751.67, 300.67, 150.33, 0.0

Tradeoffs

• Bigger batches: less noise, worse returns. Noise falls as 1/√b, so going from 256 to 1,024 halves it for four times the compute per step. Past a "critical batch size", extra examples per step barely speed up training.
• Noise is not all bad. Small-batch gradient noise helps escape sharp minima and acts as a regularizer. Large-batch training usually needs a larger learning rate and warmup to match.
• Unbiased needs truly random batches. Batches drawn from sorted or clustered data (all of one class, one user, one day) give biased gradient estimates. Shuffle every epoch.
• Seeds give reproducibility, not certainty. A result from one seed is one sample. Report the mean and spread over several seeds; with only a few runs, the standard error is wide.
• Heavy tails break the √n intuition. If the variance is infinite or enormous (rare huge losses or gradients), averages converge slowly and erratically. Gradient clipping is partly a fix for that.

Check yourself

Connections

References

• Blitzstein and Hwang. Introduction to Probability, 2nd ed., chapters 4, 7 and 10. CRC Press, 2019.
• Goodfellow, Bengio and Courville. Deep Learning, sections 3.8 and 5.9 (stochastic gradient descent). MIT Press, 2016.
• Bottou, Curtis and Nocedal. Optimization Methods for Large-Scale Machine Learning. SIAM Review, 2018. arXiv:1606.04838
• McCandlish et al. An Empirical Model of Large-Batch Training. 2018. arXiv:1812.06162 (the critical batch size).
• PyTorch v2.15.0: torch/random.py and Normal's moments.