Stochastic gradient descent, the noise ball and batch size

Before this lesson: the linear rate of gradient descent (lesson 03), the descent lemma and strong convexity (lesson 02), expectation and variance (probability and statistics track, lesson 02), and the fact that the mean of $B$ independent samples has variance $\sigma^2/B$ (probability and statistics track, lesson 06). Time: about 75 minutes, plus the notebook.

---

Why this matters #

Every large model is trained with stochastic gradients. Two numbers then shape training: the learning rate $\eta$ and the batch size $B$. This lesson explains four things you see in practice. The loss stops falling at a constant learning rate, and it drops again when you decay the learning rate. A larger global batch in data-parallel training (DDP, FSDP) needs a larger learning rate. Adding more GPUs stops reducing the number of steps at some point. And SGD, with a slower rate per step than gradient descent, still wins on large datasets. Adam (lesson 06) and the schedules of lesson 07 build on these facts.

1. Predict first #

2. The picture #

All examples agree far away, and disagree near the minimum #

Take the simplest problem: estimate a mean. The data are numbers $a_1, \dots, a_n$ with mean $\bar a$ and variance $\sigma^2 = \frac1n\sum_i (a_i - \bar a)^2$. The loss

has derivative $f'(x) = x - \bar a$, so its minimizer is $x^\star = \bar a$. One random example $i$ gives the stochastic gradient $g = x - a_i$. Its average over $i$ is $x - \bar a = f'(x)$: the stochastic gradient is right on average.

Now look at the step direction $-g = a_i - x$ of each example. If $x$ is to the right of all data points, every $a_i - x$ is negative: every example says "move left", like the full gradient. If $x$ is inside the data range, some examples say "left" and others "right". This interval is the region of confusion (Bertsekas' example, in Bottou, Curtis and Nocedal §3.3).

Interactive figure: loading…

Two facts from this picture are the core of the lesson:

1. The noise does not shrink near $x^\star$. The variance of $g = x - a_i$ is $\sigma^2$ at every $x$. What shrinks is the signal $|f'(x)| = |x - \bar a|$.
2. So with a fixed step, SGD cannot settle at $x^\star$. It reaches a noise ball around $x^\star$ and stays there.

The noise ball in two dimensions #

Now take $f(x) = \tfrac12(x_1^2 + 10x_2^2)$, so $\mu = 1$, $L = 10$ and $x^\star = 0$. We add Gaussian noise to the exact gradient: $g = \nabla f(x) + \xi$ with $\xi \sim \mathcal N(0, I/B)$. This is the noisy quadratic: a model of a minibatch gradient with batch size $B$.

Interactive figure: loading…

3. Definitions and main results #

Stochastic gradients and minibatches #

Training minimizes the empirical risk, an average over the $n$ training examples:

where $f_i$ is the loss on example $i$ and $x$ are the parameters (in a neural network, $\theta$). SGD is

where $g_k$ is a stochastic gradient: a random vector computed from a few examples. We assume two properties. Write $\mathbb E_k$ for the expectation given $x_k$ (only the new random sample is random).

• Unbiased: $\mathbb E_k[g_k] = \nabla f(x_k)$.
• Bounded variance: $\mathbb E_k\|g_k - \nabla f(x_k)\|^2 \le \sigma^2$.

A minibatch gradient averages $B$ example gradients, $g = \frac1B\sum_{j=1}^B \nabla f_{i_j}(x)$, with indices $i_1, \dots, i_B$ drawn independently and uniformly. One full pass over the data is an epoch.

Proposition (minibatch mean and variance). If one example gradient has mean $\nabla f(x)$ and variance $\sigma^2$, the minibatch gradient has mean $\nabla f(x)$ and variance $\sigma^2/B$.

Proof.

1. Write $\delta_j = \nabla f_{i_j}(x) - \nabla f(x)$. Each $\delta_j$ has mean 0 and $\mathbb E\|\delta_j\|^2 = \sigma^2$, and they are independent.
2. Mean: $\mathbb E[g] = \frac1B\sum_j \mathbb E[\nabla f_{i_j}(x)] = \nabla f(x)$.
3. Variance: $\mathbb E\|g - \nabla f(x)\|^2 = \frac{1}{B^2}\,\mathbb E\big\|\sum_j \delta_j\big\|^2 = \frac{1}{B^2}\sum_{j,l}\mathbb E[\delta_j^\top\delta_l]$.
4. For $j \ne l$, independence gives $\mathbb E[\delta_j^\top\delta_l] = \mathbb E[\delta_j]^\top\mathbb E[\delta_l] = 0$. Only the $B$ terms with $j = l$ remain: $\frac{1}{B^2}\cdot B\sigma^2 = \sigma^2/B$. $\blacksquare$

In practice, batches are often drawn without replacement. Then the variance is $\frac{n - B}{B(n - 1)}\sigma^2$: almost $\sigma^2/B$ when $B \ll n$, and $0$ at $B = n$, where the minibatch is the whole dataset (Garrigos and Gower, Lemma 6.5).

The exact noise ball in one dimension #

Go back to the mean example with a constant step $\eta \in (0, 2)$. SGD is $x_{k+1} = x_k - \eta(x_k - a_{i_k}) = (1 - \eta)x_k + \eta a_{i_k}$. Subtract $x^\star = \bar a$:

1. The new noise $a_{i_k} - \bar a$ has mean 0 and variance $\sigma^2$, and it is independent of $e_k$.
2. Square and take expectations. The cross term has mean 0, so $\mathbb E e_{k+1}^2 = (1 - \eta)^2\,\mathbb E e_k^2 + \eta^2\sigma^2$.
3. This recursion has the fixed point $m = (1-\eta)^2 m + \eta^2\sigma^2$, that is $m = \dfrac{\eta^2\sigma^2}{1 - (1 - \eta)^2} = \dfrac{\eta^2\sigma^2}{\eta(2 - \eta)}$.
4. Since $(1 - \eta)^2 < 1$, $\mathbb E e_k^2$ converges to this fixed point:

The error does not go to zero. It settles at a level proportional to $\eta$ (for small $\eta$).

A decaying step reaches $x^\star$. Take $\eta_k = 1/(k + 1)$ for $k = 0, 1, \dots$ Then $x_1 = a_{i_0}$ (whatever $x_0$ is), $x_2 = \tfrac12(a_{i_0} + a_{i_1})$, and in general $x_k$ is the average of the first $k$ samples. Its variance is $\sigma^2/k$. So $\mathbb E e_k^2 = \sigma^2/k$: the error goes to zero, at the rate $O(1/k)$. This rate cannot be improved in general: with $k$ samples of Gaussian data, no estimator of the mean has a smaller error than $\sigma^2/k$, and for strongly convex problems no method that sees $k$ noisy gradients can guarantee better than order $1/k$ (Agarwal et al. 2012).

Theorem 1 (SGD with a constant step converges to a noise ball) #

Setting. $f$ is $\mu$-strongly convex and $L$-smooth with minimizer $x^\star$ (lesson 02). The stochastic gradients are unbiased with variance at most $\sigma^2$. The step is constant, $0 < \eta \le 1/L$. Then

With minibatches of size $B$, replace $\sigma^2$ by $\sigma^2/B$.

Why this is true. Each step contracts the squared error like gradient descent, by a factor $1 - \eta\mu$, and adds fresh noise of size $\eta^2\sigma^2$. With a squared error $m$, the two balance when $\eta\mu\, m \approx \eta^2\sigma^2$, that is at $m \approx \eta\sigma^2/\mu$.

Proof.

1. Expand the square: $\|e_{k+1}\|^2 = \|e_k - \eta g_k\|^2 = \|e_k\|^2 - 2\eta\, g_k^\top e_k + \eta^2\|g_k\|^2$.
2. Take $\mathbb E_k$. Unbiasedness gives $\mathbb E_k[g_k^\top e_k] = \nabla f(x_k)^\top e_k$. Split $g_k$ into mean plus noise; the cross term has mean 0, so $\mathbb E_k\|g_k\|^2 = \|\nabla f(x_k)\|^2 + \mathbb E_k\|g_k - \nabla f(x_k)\|^2 \le \|\nabla f(x_k)\|^2 + \sigma^2$.
3. Strong convexity between $x_k$ and $x^\star$: $f^\star \ge f(x_k) + \nabla f(x_k)^\top(x^\star - x_k) + \tfrac\mu2\|e_k\|^2$. Rearranged: $\nabla f(x_k)^\top e_k \ge f(x_k) - f^\star + \tfrac\mu2\|e_k\|^2$.
4. Smoothness: the descent lemma at the point $x_k - \nabla f(x_k)/L$ gives $f^\star \le f(x_k) - \|\nabla f(x_k)\|^2/(2L)$, so $\|\nabla f(x_k)\|^2 \le 2L\,(f(x_k) - f^\star)$.
5. Put 2, 3 and 4 into 1:

For $\eta \le 1/L$ the middle term is $\le 0$. Drop it.

6. Take full expectations and unroll the recursion $m_{k+1} \le (1 - \eta\mu)m_k + \eta^2\sigma^2$, with $m_k = \mathbb E\|e_k\|^2$. The geometric sum $\sum_{j \ge 0}(1 - \eta\mu)^j$ is at most $1/(\eta\mu)$:

With $\sigma = 0$ this is the linear rate of gradient descent in distance (lesson 03).

Reading the bound. The first term is the transient: it shrinks like gradient descent with step $\eta$. The second term is the floor.

• Halving $\eta$ halves the floor, but the transient becomes about twice as slow.
• Doubling $B$ halves the floor in the same number of steps, but each step costs twice as many example gradients.
• Per example gradient, the two options cost about the same: minibatch SGD with step $\eta$ behaves like single-example SGD with step $\eta/B$ (Bottou, Curtis and Nocedal §4.2).

Robbins–Monro conditions #

To reach $x^\star$ itself, the step must go to zero, but not too fast (Robbins and Monro 1951):

• The first condition lets the iterates travel any distance. If $\sum\eta_k$ were finite, the total distance would be bounded, and SGD could stop before reaching $x^\star$.
• The second condition keeps the total injected noise $\sum_k \eta_k^2\sigma^2$ finite.

$\eta_k = c/(k + 1)$ satisfies both: $\sum 1/k$ diverges and $\sum 1/k^2 = \pi^2/6$. In general $c/k^a$ satisfies both exactly when $1/2 < a \le 1$.

Theorem 2 (the $O(1/k)$ rate with a decaying step) #

Take the setting of Theorem 1 and $\eta_k = c/(k_0 + k)$ with $c\mu > 1$ and $\eta_0 = c/k_0 \le 1/L$. Then

This is a variant of Theorem 4.7 of Bottou, Curtis and Nocedal, written for the distance.

Proof by induction

Step 5 of Theorem 1 holds for every step $\eta_k \le 1/L$. With $m_k = \mathbb E\|e_k\|^2$ it gives $m_{k+1} \le (1 - \eta_k\mu)m_k + \eta_k^2\sigma^2$. Write $\hat k = k_0 + k$, so $\eta_k = c/\hat k$.

1. Base case: $m_0 = \|e_0\|^2 \le \nu/k_0$, because $\nu \ge k_0\|e_0\|^2$.
2. Induction step: assume $m_k \le \nu/\hat k$. Since $1 - \eta_k\mu \ge 0$,

$$m_{k+1} \le \Big(1 - \frac{c\mu}{\hat k}\Big)\frac{\nu}{\hat k} + \frac{c^2\sigma^2}{\hat k^2} = \frac{(\hat k - 1)\,\nu}{\hat k^2} - \frac{(c\mu - 1)\,\nu - c^2\sigma^2}{\hat k^2} .$$

3. The second fraction is $\ge 0$ because $\nu \ge c^2\sigma^2/(c\mu - 1)$. Drop it.
4. $(\hat k - 1)(\hat k + 1) = \hat k^2 - 1 \le \hat k^2$, so $\dfrac{\hat k - 1}{\hat k^2} \le \dfrac{1}{\hat k + 1}$. Hence $m_{k+1} \le \nu/(\hat k + 1)$. $\blacksquare$

The condition $c\mu > 1$ matters. Even without noise, $\eta_k = c/k$ on $f(x) = \tfrac\mu2 x^2$ gives $e_k = e_0\prod_{j=1}^{k}(1 - c\mu/j)$, and $e_k^2$ decays like $k^{-2c\mu}$. If $c\mu < 1/2$ this is slower than $1/k$. A step that decays too early is as bad as no decay.

Batch size: the linear scaling rule #

Linear scaling rule (Goyal et al. 2017, §2.1): when you multiply the batch size by $c$, multiply the learning rate by $c$.

Why. Compare $c$ small steps with one large step over the same $cB$ examples:

• $c$ steps of size $\eta$ move by $\eta\,(g_k + g_{k+1} + \dots + g_{k+c-1})$, each $g$ a mean over $B$ examples;
• one step of size $c\eta$ with a batch of $cB$ examples moves by $c\eta\cdot\frac1c(g_k + \dots + g_{k+c-1})$ if the gradients are computed at the same point.

The two agree when the gradient changes little during the $c$ small steps. In Theorem 1, the floor depends only on $\eta/B$, which the rule keeps fixed. And $c$ contractions $(1 - \eta\mu)^c$ match one contraction $1 - c\eta\mu$ to first order.

When it fails.

1. When $c\eta$ reaches the curvature limit: gradient descent is unstable for steps above $2/L$ (lesson 01), whatever the noise.
2. Early in training, where the gradient changes fast. Goyal et al. needed a gradual warmup (5 epochs) to train ResNet-50 with batch 8192 (lesson 07).
3. When $B$ is far above the critical batch size (next).

The gradient noise scale and the critical batch size #

McCandlish et al. (2018) ask: for which batch size does a larger batch stop helping? Let $G$ be the true gradient, $H$ the Hessian and $\Sigma$ the covariance matrix of the per-example gradients. A step $x - \varepsilon\, g$ with a batch of size $B$ changes the loss, to second order and in expectation, by

Minimize over $\varepsilon$ (a parabola in $\varepsilon$):

Here $\varepsilon_{\max} = \|G\|^2/G^\top HG$ and $\Delta f_{\max} = \tfrac12\|G\|^4/G^\top HG$ are the best step and the best decrease with the exact gradient. $B_{\text{noise}}$ is the gradient noise scale. If $H$ is a multiple of the identity, it becomes the simple noise scale

the total gradient variance divided by the squared gradient norm. It is the batch size at which noise and signal in one step have about the same size.

• $B \ll B_{\text{noise}}$: $\varepsilon_{\text{opt}}$ and the progress per step grow linearly with $B$. This is the linear scaling regime: doubling $B$ halves the number of steps, for the same number of examples.
• $B \gg B_{\text{noise}}$: the progress per step saturates at $\Delta f_{\max}$. Extra examples are wasted.

Over a whole training run, the number of steps $S$ and the number of examples $E$ needed to reach a target loss follow

The critical batch size $B_{\text{crit}} = E_{\min}/S_{\min}$ is the batch size where training takes twice the minimum number of steps and twice the minimum number of examples. McCandlish et al. find $B_{\text{crit}} \approx B_{\text{noise}}$. As the loss falls, $\|G\|$ shrinks and $B_{\text{noise}}$ grows: a large batch is more useful late in training than early.

Interactive figure: loading…

GD or SGD per unit of computation #

Count example gradients (one $\nabla f_i$ is one unit of work) for a finite sum of $n$ terms, to reach $\mathbb E\|e_k\|^2 \le \varepsilon$:

• GD: $n$ per step. With $\eta = 1/L$ and $\sigma = 0$, Theorem 1 gives $\|e_k\|^2 \le (1 - 1/\kappa)^k\|e_0\|^2 \le e^{-k/\kappa}\|e_0\|^2$, so about $\kappa\ln(\|e_0\|^2/\varepsilon)$ steps: about $n\kappa\ln(1/\varepsilon)$ in total.
• SGD with $\eta_k = c/(k_0 + k)$: one per step and about $\nu/\varepsilon$ steps (Theorem 2). This does not depend on $n$.

GD wins for high accuracy on a small dataset. SGD wins when $n$ is large. In learning, the accuracy only needs to reach the statistical error, about $\varepsilon \sim 1/n$. Then GD costs about $n\kappa\ln n$ and SGD about $\nu n$: a fixed number of passes over the data (Bottou, Curtis and Nocedal §4.4). Redundant data widen the gap: copy the dataset 10 times, and a GD step costs 10 times more, while SGD does not change.

Interactive figure: loading…

Answers to section 1. (1) The level halves: the floor $\eta\sigma^2/\mu$ is proportional to $\eta$. But the contraction factor $1 - \eta\mu$ moves closer to 1, so reaching the floor takes about twice as many steps. This is why decaying the learning rate makes the loss drop. (2) Multiply the learning rate by 8, with a warmup. This stops working when $8\eta$ reaches the curvature limit (about $2/L$), or when the global batch is far above $B_{\text{crit}}$: then more GPUs cost more but hardly reduce the number of steps.

4. Worked example #

Problem. $f(x) = \tfrac12(x_1^2 + 10x_2^2)$, so $\mu = 1$, $L = 10$, $x^\star = 0$. Start at $x_0 = (1, 1)$, so $\|e_0\|^2 = 2$. The stochastic gradient is $g = \nabla f(x) + \xi$ with $\xi \sim \mathcal N(0, I_2/B)$. One example has variance $\sigma^2 = \operatorname{tr} I_2 = 2$.

(a) The exact floor. The coordinates do not interact. For a coordinate with curvature $\lambda$, the error follows $e_{k+1} = (1 - \eta\lambda)e_k - \eta\xi_k$ with $\operatorname{Var}\xi_k = 1/B$. As in the 1D derivation of section 3:

For $\eta = 0.05$, $B = 1$: the flat coordinate ($\lambda = 1$) gives $0.05/(1 \cdot 1.95) = 0.0256$; the steep one ($\lambda = 10$) gives $0.05/(10 \cdot 1.5) = 0.0033$. The sum is $0.0290$.

$\eta$	$B$	exact floor	bound $\eta\sigma^2/(\mu B)$
0.05	1	$0.0256 + 0.0033 = 0.0290$	0.1
0.025	1	$0.0127 + 0.0014 = 0.0141$	0.05
0.05	4	$0.0290/4 = 0.0072$	0.025

The bound is about 3.5 times too large here, because it uses only $\mu$ (the flattest direction) for all the noise. But it scales correctly: proportional to $\eta$ and to $1/B$. The transient factor $(1 - \eta\mu)^k = 0.95^k$ shrinks by $10^3$ after $\ln 10^3/(-\ln 0.95) \approx 135$ steps.

(b) Robbins–Monro steps. Take $\eta_k = 2/(20 + k)$: $c = 2$, $k_0 = 20$, $c\mu = 2 > 1$ and $\eta_0 = 0.1 = 1/L$. Then

and Theorem 2 gives $\mathbb E\|e_k\|^2 \le 40/(20 + k)$: $0.039$ at $k = 1000$ and $0.0040$ at $k = 10^4$. The exact value, from the second-moment recursion of each coordinate, is about $1.44/(20 + k)$: $0.0014$ and $0.00014$. The bound is loose in the constant but has the right $1/k$ shape. For comparison, the constant step $\eta = 0.05$ stays at $0.029$ forever.

(c) Linear scaling. Keep $\eta/B = 0.0125$ and spend $E = 4000$ examples, so $S = E/B$ steps (exact second moments):

$(\eta, B)$	(0.0125, 1)	(0.025, 2)	(0.05, 4)	(0.1, 8)	(0.2, 16)
steps $S$	4000	2000	1000	500	250
$\mathbb E\|e_S\|^2$	0.0070	0.0070	0.0072	0.0078	1.63

Up to $B = 8$, the same number of examples gives almost the same error in 8 times fewer steps. At $B = 16$ the step $0.2 = 2/L$ sits at the stability limit of the steep coordinate: $|1 - 0.2 \cdot 10| = 1$. That coordinate no longer contracts, the noise accumulates, and the error grows to $1.63$.

(d) Noise scale. Here $H = \operatorname{diag}(1, 10)$ and the per-example covariance is $\Sigma = I$. At a point where $G = (1, 1)$: $B_{\text{noise}} = \operatorname{tr}(H\Sigma)/G^\top HG = 11/11 = 1$ and $B_{\text{simple}} = \operatorname{tr}\Sigma/\|G\|^2 = 2/2 = 1$. Closer to the minimum, at $G = (0.1, 0.1)$: $11/0.11 = 100$ and $2/0.02 = 100$. Far from the minimum a batch of 1 is already enough; near it, batches up to about 100 still pay off.

(e) Cost. Suppose this were a finite sum of $n = 10^6$ examples with the same noise, and we want $\mathbb E\|e\|^2 \le 10^{-4}$. For GD with $\eta = 1/L = 0.1$, Theorem 1 with $\sigma = 0$ gives $\|e_k\|^2 \le 0.9^k\|e_0\|^2$, so it needs $\ln(2 \cdot 10^4)/(-\ln 0.9) \approx 94$ steps of $10^6$ example gradients: about $9.4 \cdot 10^7$. SGD with the steps of (b) needs, by the bound, $40/10^{-4} = 4 \cdot 10^5$ example gradients (with the exact constant $1.44$, about $1.4 \cdot 10^4$).

5. In code #

NumPy checks the floors of the worked example with 20,000 independent SGD runs, then measures minibatch noise on a least-squares problem and estimates the simple noise scale. The estimate uses two batch sizes: since $\mathbb E\|g_B\|^2 = \|G\|^2 + \operatorname{tr}(\Sigma)/B$, the squared norms at a small and a large batch give unbiased estimates of $\|G\|^2$ and $\operatorname{tr}\Sigma$. In data-parallel training the small batch is the per-GPU batch and the large one is the global batch after the all-reduce, so the estimate is almost free (McCandlish et al., Appendix A.1).

import numpy as np

# 1) Noise ball on f(x) = (x1^2 + 10 x2^2)/2 with gradient noise N(0, I/B).
#    R independent SGD runs at once (one row per run).
rng = np.random.default_rng(0)
lam = np.array([1.0, 10.0])
R, K = 20_000, 2_000

def floor(eta, B):
    x = np.ones((R, 2))                                   # x0 = (1, 1) in every run
    for _ in range(K):
        g = lam * x + rng.standard_normal((R, 2)) / np.sqrt(B)   # unbiased, variance 2/B
        x -= eta * g
    return (x ** 2).sum(1).mean()

for eta, B in [(0.05, 1), (0.025, 1), (0.05, 4)]:
    exact = (eta / (B * lam * (2 - eta * lam))).sum()
    print(eta, B, round(floor(eta, B), 4), round(exact, 4), eta * 2 / B)
# 0.05 1 0.0288 0.029 0.1        <- simulated, exact, bound eta sigma^2 / (mu B)
# 0.025 1 0.0141 0.0141 0.05
# 0.05 4 0.0072 0.0072 0.025

# 2) Minibatch noise is tr(Sigma)/B, and the simple noise scale grows near the minimum.
rng = np.random.default_rng(0)
n, d = 10_000, 5
X = rng.standard_normal((n, d))
y = X @ rng.standard_normal(d) + 0.5 * rng.standard_normal(n)
w_hat = np.linalg.lstsq(X, y, rcond=None)[0]          # minimizer of ||Xw - y||^2 / (2n)

def stats(w):
    per_example = X * (X @ w - y)[:, None]            # gradient of each term, shape (n, d)
    G = per_example.mean(0)                           # full gradient
    tr_sigma = ((per_example - G) ** 2).sum(1).mean() # trace of the per-example covariance
    return per_example, G, tr_sigma

def batch_grads(per_example, B, draws=20_000):        # minibatch gradients, with replacement
    return per_example[rng.integers(0, n, size=(draws, B))].mean(1)

per_example, G, tr_sigma = stats(np.zeros(d))
for B in [1, 8, 64]:
    var = ((batch_grads(per_example, B) - G) ** 2).sum(1).mean()   # E|G_B - G|^2
    print(B, round(var * B / tr_sigma, 2))            # about 1: the variance is tr(Sigma)/B
# 1 0.98
# 8 1.0
# 64 1.0

for w in [np.zeros(d), w_hat + 0.05]:                 # far from and near the minimizer
    per_example, G, tr_sigma = stats(w)
    b_s, b_b = 8, 64
    g_s = (batch_grads(per_example, b_s) ** 2).sum(1).mean()   # E|G_Bsmall|^2
    g_b = (batch_grads(per_example, b_b) ** 2).sum(1).mean()   # E|G_Bbig|^2
    G2 = (b_b * g_b - b_s * g_s) / (b_b - b_s)                 # estimates |G|^2
    S = (g_s - g_b) / (1 / b_s - 1 / b_b)                      # estimates tr(Sigma)
    print(round(S / G2, 1), round(tr_sigma / (G @ G), 1))      # estimated vs exact B_simple
# 6.5 6.3
# 100.1 103.2

The noise scale grows from about 6 far from the minimizer to about 100 near it, as in part (d).

6. Where it shows up #

• Training large models. In DDP or FSDP the global batch is (per-GPU batch) × (number of GPUs) × (gradient-accumulation steps). Scaling out raises $B$. The linear scaling rule with warmup keeps the loss curve per example roughly unchanged while $B$ stays below $B_{\text{crit}}$. Kaplan et al. (2020) fit $B_{\text{crit}}$ as a power law in the loss for language models and used it to choose batch sizes. The noise-scale estimate of section 5 costs almost nothing in a DDP run.
• Learning-rate decay. Step decay, cosine decay and the cooldown of warmup–stable–decay (WSD) schedules all shrink the noise ball at the end of training (lesson 07).
• ML for science. A physics-informed loss is a Monte Carlo estimate of an integral over collocation points; resampling them every step makes the gradient stochastic, with a variance that falls like one over the number of points. In diffusion and flow-matching training each loss term also samples a time $t$ and a Gaussian noise, so $\sigma^2$ includes these sources too. For small full-batch problems (an inverse problem with a few thousand residuals), deterministic methods with linear rates win (quasi-Newton and Gauss–Newton methods, lesson 10).

7. Common mistakes #

• Reading the plateau as convergence. At a constant $\eta$ the error stops at a floor proportional to $\eta\sigma^2/B$. Decay $\eta$ or grow $B$ to go lower.
• "Unbiased means it converges." An unbiased gradient is right on average, but the iterates still jump around $x^\star$ unless the step goes to zero.
• Decaying the step badly. $c/k^2$ stops before $x^\star$ (the total travel is finite). $c/k$ with $c\mu < 1/2$ decays slower than $1/k$, even without noise.
• Comparing per iteration instead of per example. An SGD step costs $B$ example gradients; a GD step costs $n$.
• Scaling the learning rate with $B$ without limit. Linear scaling holds only for $B \lesssim B_{\text{noise}}$ and $\eta$ below the curvature limit, and it needs warmup.
• Forgetting that the noise scale changes. $B_{\text{noise}}$ grows as the loss falls, so a batch that is too large early can be the right size late.

8. Check yourself #

1. In the 1D mean example with $\sigma^2 = 1$ and $\eta = 0.1$, what is the exact noise floor, and what does Theorem 1 give?
   Answer

   Exact: $\eta\sigma^2/(2 - \eta) = 0.1/1.9 \approx 0.053$. Theorem 1: $\eta\sigma^2/\mu = 0.1$ (here $\mu = L = 1$).
2. Do $\eta_k = (k+1)^{-0.7}$ and $\eta_k = (k+1)^{-0.4}$ satisfy the Robbins–Monro conditions?
   Answer

   $(k+1)^{-0.7}$: yes, $\sum k^{-0.7} = \infty$ and $\sum k^{-1.4} < \infty$. $(k+1)^{-0.4}$: no, $\sum k^{-0.8} = \infty$, so the second condition fails. In general $c/k^a$ works exactly for $1/2 < a \le 1$.
3. Where does the proof of Theorem 1 use $\eta \le 1/L$?
   Answer

   In step 5. Smoothness turns the $\eta^2\|\nabla f\|^2$ term into at most $2\eta^2L(f - f^\star)$, and $\eta \le 1/L$ makes $-2\eta(1 - \eta L)(f - f^\star) \le 0$, so the term can be dropped.
4. Batch 256 with learning rate 0.1 works well. You move to batch 2048. What learning rate does the linear scaling rule give, and what can break it?
   Answer

   $0.1 \cdot 2048/256 = 0.8$. It breaks if $0.8$ exceeds the curvature (stability) limit, early in training without warmup, or if 2048 is far above the critical batch size, where the extra examples give little progress.
5. At some point in training $B_{\text{simple}} \approx 50$, and your batch is 4096. What does the noise-scale model predict?
   Answer

   Take $B_{\text{noise}} \approx 50$. Then $\Delta f_{\text{opt}} = \Delta f_{\max}/(1 + 50/4096) \approx 0.988\,\Delta f_{\max}$; with half the batch it is about $0.976\,\Delta f_{\max}$. Almost the same progress per step for half the compute: most examples are wasted (until $B_{\text{noise}}$ grows later in training).
6. $n = 10^4$ examples, $\kappa = 10$, and you need $\varepsilon = 10^{-10}$ (a high-accuracy fit). GD or SGD?
   Answer

   GD (or conjugate gradients, L-BFGS): about $n\kappa\ln(1/\varepsilon) = 10^4 \cdot 10 \cdot 23 \approx 2.3 \cdot 10^6$ example gradients. SGD needs about $\nu/\varepsilon$, of order $10^{10}\nu$.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. What does "unbiased stochastic gradient" mean?

   Show answer

2. What is the variance of the mean of $B$ independent example gradients with variance $\sigma^2$?

   Show answer

3. State the noise-ball bound for SGD with a constant step.

   Show answer

4. Why does SGD with a constant step not converge to $x^\star$?

   Show answer

5. In 1D mean estimation, what is the exact noise floor of SGD with step $\eta$?

   Show answer

6. State the Robbins–Monro conditions and say why each one is needed.

   Show answer

7. What is the linear scaling rule, and when does it fail?

   Show answer

8. What is the simple noise scale, and what does it predict?

   Show answer

9. Per example gradient, why does SGD beat GD on a large dataset?

   Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Why does SGD not converge with a constant learning rate?" — The gradient noise does not vanish at $x^\star$, so the error settles at about $\eta\sigma^2/(\mu B)$. Decay the learning rate or grow the batch to go lower.
• "You double the batch size. What do you do with the learning rate?" — Double it while $B \ll B_{\text{noise}}$ and the step stays below the curvature limit, with warmup. Beyond $B_{\text{crit}}$ larger batches mostly waste compute.
• "SGD is $O(1/k)$ and GD is linear. Why use SGD?" — A step costs $B$ instead of $n$ example gradients. To statistical accuracy $\varepsilon \sim 1/n$, SGD needs $O(n)$ example gradients, GD $O(n\kappa\log n)$.
• "What is the gradient noise scale, and how do you measure it?" — $B_{\text{simple}} = \operatorname{tr}\Sigma/\|G\|^2$, the batch size at which noise and signal in one step are about equal. Estimate it from the squared gradient norms at two batch sizes, for example per GPU and after the all-reduce.

10. Practice #

The notebook for this lesson runs in Colab. It has four exercises, each with a check cell and a folded solution:

1. Simulate SGD for 1D mean estimation and compare the measured noise floor with $\eta\sigma^2/(2 - \eta)$.
2. Write minibatch SGD for least squares; compare a constant step with a decaying step $c/(k_0 + k)$.
3. Measure the variance of minibatch gradients against $B$ for logistic regression on the scikit-learn digits dataset, and estimate the simple noise scale.
4. Find the number of steps to a target error against the batch size, with the linear scaling rule, and find the batch size where linear scaling breaks.

11. Watch and read #

25. Stochastic Gradient DescentMIT OpenCourseWare · 18.065, Suvrit Sra · 53 min

Why SGD makes fast progress far from the minimum and slows down near it (the region of confusion).

Stochastic Gradient Descent, Clearly Explained!!!StatQuest with Josh Starmer · 11 min

A gentle picture of SGD and minibatches on a small regression problem.

Mini Batch Gradient Descent (C2W2L01)DeepLearningAI · Andrew Ng · 11 min

How minibatches and epochs are organized in practice.

• Bottou, Curtis and Nocedal, Optimization Methods for Large-Scale Machine Learning (SIAM Review, 2018): §3.3 (why SGD, the region of confusion), §4.2 (Theorems 4.6 and 4.7) and §4.4 (work complexity).
• Garrigos and Gower, Handbook of Convergence Theorems for (Stochastic) Gradient Methods: §5 (SGD theorems) and Lemma 6.5 (minibatch variance).
• McCandlish, Kaplan, Amodei et al., An Empirical Model of Large-Batch Training: §2 (the noise scale and the critical batch size) and Appendix A.1 (the estimator of section 5).
• Goyal et al., Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour: §2 (the linear scaling rule and gradual warmup).
• Robbins and Monro, A Stochastic Approximation Method (Annals of Mathematical Statistics, 1951): the original step-size conditions.
• Agarwal, Bartlett, Ravikumar and Wainwright, Information-theoretic lower bounds on the oracle complexity of stochastic convex optimization (IEEE Transactions on Information Theory, 2012): why no method beats the order $1/k$ with noisy gradients.