Part II — Linear Models · Lecture 3

Gradient Descent That Actually Works

Every model from here on is fitted by gradient descent, so it is worth doing properly. Convergence rates from the condition number, why stochastic gradients have a noise floor, and where momentum and Adam come from — each derived, not quoted.

By the end you can

Derive the convergence rate \(\big(\tfrac{\kappa-1}{\kappa+1}\big)^t\) and say what limits it
Explain the SGD noise floor and why every schedule decays the learning rate
Derive momentum and explain the \(\kappa \to \sqrt{\kappa}\) improvement
Build AdaGrad, RMSProp and Adam one modification at a time
Diagnose a failing training run from its loss curve

Assumed

Phase A, Lecture 6: Taylor, Newton, the \(\eta<2/L\) bound
Phase A, Lecture 11: eigenvalues, condition number, preconditioning
Phase A, Lecture 13: mini-batch gradients are unbiased with variance σ²/B

1What Phase A established

For \(f(\vw)=\tfrac12\vw^\top\mA\vw\) with \(\mA\succ0\), gradient descent is \(\vw_{t+1}=(\mI-\eta\mA)\vw_t\), and in the eigenbasis the coordinates decouple:

\[z_i^{(t)} = (1-\eta\lambda_i)^{t}\,z_i^{(0)}.\]

Three consequences, all from Phase A Lecture 11: stability needs \(\eta<2/\lambda_{\max}\); speed is governed by the slowest coordinate; and with the best fixed \(\eta\) the error contracts by \(\frac{\kappa-1}{\kappa+1}\) per step, where \(\kappa=\lambda_{\max}/\lambda_{\min}\). Let me finish that argument, because the optimal \(\eta\) was asserted rather than derived.

The optimal step size and the κ rate0 / 5

\[\text{contraction}(\eta) = \max_i \big|1-\eta\lambda_i\big| = \max\Big(|1-\eta\lambda_{\min}|,\ |1-\eta\lambda_{\max}|\Big)\]

Only the extremes matter. As \(\lambda\) ranges over \([\lambda_{\min},\lambda_{\max}]\), the function \(|1-\eta\lambda|\) is piecewise linear in \(\lambda\), so its maximum over the interval is attained at an endpoint. We must make the worse of those two as small as possible.

\[\text{for small }\eta:\ 1-\eta\lambda_{\min} \text{ dominates}; \quad \text{for large }\eta:\ \eta\lambda_{\max}-1 \text{ dominates}\]

Two competing terms. A small \(\eta\) leaves the flattest direction barely moving (contraction near 1); a large \(\eta\) makes the steepest direction overshoot. The first decreases in \(\eta\), the second increases. The best \(\eta\) balances them.

\[1-\eta\lambda_{\min} = \eta\lambda_{\max}-1 \quad\Longrightarrow\quad \eta^\star = \frac{2}{\lambda_{\min}+\lambda_{\max}}\]

Set them equal. A decreasing function and an increasing function are minimised in max at their crossing point — the standard minimax argument.

\[\text{rate} = 1-\eta^\star\lambda_{\min} = 1-\frac{2\lambda_{\min}}{\lambda_{\min}+\lambda_{\max}} = \frac{\lambda_{\max}-\lambda_{\min}}{\lambda_{\max}+\lambda_{\min}} = \frac{\kappa-1}{\kappa+1}\]

Substitute back. Dividing top and bottom by \(\lambda_{\min}\) gives the condition number. ∎

\[\text{steps to reach accuracy }\epsilon:\quad t \;\approx\; \frac{\kappa}{2}\ln\frac{1}{\epsilon}\]

Turn the rate into a count. We need \(\big(\frac{\kappa-1}{\kappa+1}\big)^t\le\epsilon\). Taking logs and using \(\ln\frac{\kappa-1}{\kappa+1}\approx-\frac{2}{\kappa}\) for large \(\kappa\) gives \(t\approx\frac{\kappa}{2}\ln\frac1\epsilon\). Linear in \(\kappa\). A problem with \(\kappa=10^4\) needs about ten thousand times as many steps as one with \(\kappa=1\), and no tuning of \(\eta\) can rescue it. Only changing the geometry — scaling features, preconditioning, momentum — helps.

Four optimisers on the same ill-conditioned bowlInteractive

gradient descent+ momentumAdamNewton (for reference)

log₁₀ condition number κ1.60

learning rate η (GD and momentum)0.080

momentum β0.90

steps30

Plain gradient descent zig-zags across the valley because the gradient points across it, not along it (Phase A Lecture 12). Momentum accumulates the consistent along-valley component while the oscillating across-valley components cancel. Adam rescales each coordinate by its own gradient history, which flattens the valley entirely. Newton is drawn for reference — it lands in one step, because it knows the curvature exactly.

2Stochastic gradient descent

Phase A Lecture 13 proved the two facts that matter: the mini-batch gradient is unbiased, \(\E[\vg_{\mathcal B}]=\vg_{\text{full}}\), and its noise scale is \(\sigma/\sqrt{B}\). Two consequences follow that shape every training recipe you will ever write.

Why constant-η SGD does not converge0 / 5

\[\vw_{t+1} = \vw_t - \eta\big(\vg_{\text{full}}(\vw_t) + \boldsymbol\xi_t\big), \qquad \E[\boldsymbol\xi_t]=\vzero,\ \Var(\boldsymbol\xi_t)=\tfrac{\sigma^2}{B}\]

Split the stochastic gradient into the true gradient plus zero-mean noise. Legitimate exactly because the mini-batch estimate is unbiased.

\[\text{near the optimum: } \vg_{\text{full}}(\vw)\approx \mA(\vw-\vw^\star) \quad\Longrightarrow\quad \ve_{t+1} = (\mI-\eta\mA)\ve_t - \eta\boldsymbol\xi_t\]

Linearise. Writing \(\ve_t=\vw_t-\vw^\star\), the error obeys a linear recursion driven by noise — a discrete Ornstein–Uhlenbeck process. Two forces: contraction towards zero, and a random kick each step.

\[\E\|\ve_{t+1}\|^2 = \E\|(\mI-\eta\mA)\ve_t\|^2 + \eta^2\frac{\sigma^2}{B}\]

Take expectations. The cross term vanishes because \(\boldsymbol\xi_t\) has mean zero and is independent of \(\ve_t\) (the batch at step \(t\) is drawn fresh). Variances add — Phase A Lecture 13.

\[\text{stationary point: } \E\|\ve\|^2_{\infty} \;\approx\; \frac{\eta\,\sigma^2}{2\lambda B}\]

Set \(\E\|\ve_{t+1}\|^2=\E\|\ve_t\|^2\) and solve (one dimension, \(\mA=\lambda\)): \((1-\eta\lambda)^2 E + \eta^2\sigma^2/B = E\) gives \(E(2\eta\lambda-\eta^2\lambda^2)=\eta^2\sigma^2/B\), so \(E=\frac{\eta\sigma^2}{\lambda B(2-\eta\lambda)}\approx\frac{\eta\sigma^2}{2\lambda B}\) for small \(\eta\).

\[\boxed{\text{noise floor} \;\propto\; \frac{\eta}{B}}\]

The result. SGD with constant \(\eta\) settles into a cloud around the optimum, not onto it, and the cloud's size is proportional to \(\eta/B\). It does not shrink with more steps. This is why every loss curve plateaus, why cutting the learning rate produces the characteristic sudden drop, and why the classical Robbins–Monro conditions \(\sum_t\eta_t=\infty\), \(\sum_t\eta_t^2<\infty\) are what guarantee actual convergence. ∎

The noise floor, and what a schedule does about itInteractive

losslearning ratenoise floor for the current η

0 = constant · 1 = step decay · 2 = cosine · 3 = warmup + cosine0

base learning rate η₀0.150

batch size B8

A 1-D quadratic with gradient noise, 600 steps, log loss on the vertical axis. With a constant rate the loss falls fast and then flatlines on the noise floor — extra steps buy nothing. Step decay cuts the rate and the loss immediately drops to the new, lower floor: that staircase is the most recognisable shape in deep learning. Cosine does the same thing smoothly. Raising the batch size lowers the floor too, at proportionally more compute per step.

3Momentum, and why it changes the exponent

Look again at the zig-zag in the first figure. Along the valley the gradient is small but consistent in direction; across the valley it is large but alternates in sign. If we average recent gradients, the consistent component accumulates and the alternating one cancels.

Momentum as an exponentially weighted average0 / 5

\[\vv_{t} = \beta\vv_{t-1} + \vg_t, \qquad \vw_{t+1} = \vw_t - \eta\,\vv_t\]

The update. \(\vv\) is a running accumulation of gradients, decayed by \(\beta\in[0,1)\) each step. \(\beta=0\) recovers plain gradient descent.

\[\vv_t = \sum_{k=0}^{t}\beta^{k}\,\vg_{t-k}\]

Unroll it. Momentum is an exponentially weighted sum of every past gradient, with the most recent weighted most heavily. Nothing mysterious is happening — it is a low-pass filter on the gradient sequence.

\[\text{if } \vg_t \equiv \vg \text{ (constant): } \quad \vv_\infty = \frac{\vg}{1-\beta}\]

The amplification factor. A persistent gradient direction gets amplified by \(\frac{1}{1-\beta}\): at \(\beta=0.9\) that is 10×, at \(\beta=0.99\) it is 100×. This is why you must reduce \(\eta\) when you increase \(\beta\) — the effective step size is \(\eta/(1-\beta)\).

\[\text{if } \vg_t \text{ alternates } +\vg,-\vg,\dots: \quad \vv_\infty = \frac{\vg}{1+\beta}\]

The suppression factor. An oscillating gradient gets damped to \(\frac{1}{1+\beta}\) — about \(0.53\times\) at \(\beta=0.9\). Amplify the signal by 10, damp the oscillation by half: the ratio between the two improves by a factor of roughly 20. That is the mechanism, entirely.

\[\text{optimal } \beta = \left(\frac{\sqrt\kappa-1}{\sqrt\kappa+1}\right)^2, \qquad \text{rate} = \frac{\sqrt\kappa-1}{\sqrt\kappa+1}\]

The payoff (heavy-ball analysis, stated). The contraction factor becomes \(\frac{\sqrt\kappa-1}{\sqrt\kappa+1}\) instead of \(\frac{\kappa-1}{\kappa+1}\), so the iteration count falls from \(O(\kappa)\) to \(O(\sqrt\kappa)\). At \(\kappa=10^4\) that is 10 000 steps versus 100 — a qualitative change, not a constant factor. The full derivation requires analysing a two-term linear recursion with complex roots; I am marking it deferred rather than skipped.

4AdaGrad → RMSProp → Adam, one change at a time

Momentum fixes the direction. The other problem is that different coordinates need different step sizes — which is exactly the preconditioning idea from Phase A Lecture 11, where Newton's \(\mH^{-1}\vg\) divided each eigen-coordinate by its curvature. Since we cannot afford the Hessian, we approximate its diagonal from the gradients we have already seen.

Building Adam0 / 5

\[\textbf{AdaGrad:}\quad s_{t,i} = \sum_{k\le t} g_{k,i}^2, \qquad w_{t+1,i} = w_{t,i} - \frac{\eta}{\sqrt{s_{t,i}}+\epsilon}\,g_{t,i}\]

Per-coordinate scaling by accumulated gradient magnitude. Coordinates that have seen large gradients get small steps; rarely-active coordinates get large ones. Excellent for sparse features (a rare word's embedding still moves). The fatal flaw: \(s_t\) only grows, so the effective learning rate decays monotonically to zero and learning stops, whether or not you have arrived.

\[\textbf{RMSProp:}\quad s_{t,i} = \beta_2 s_{t-1,i} + (1-\beta_2)g_{t,i}^2\]

Replace the sum with an exponential moving average. Now \(s_t\) tracks the recent squared gradient rather than all of history, so it can decrease again and learning never dies. \(\beta_2=0.999\) means an effective window of about 1000 steps. This single change is what made adaptive methods work for neural networks.

\[\textbf{Adam:}\quad m_{t,i} = \beta_1 m_{t-1,i}+(1-\beta_1)g_{t,i}, \qquad s_{t,i} = \beta_2 s_{t-1,i}+(1-\beta_2)g_{t,i}^2\]

Add momentum to RMSProp. \(m_t\) is the first moment (direction), \(s_t\) the second (scale). Adam = momentum + per-coordinate preconditioning. That is the whole idea; everything else is bookkeeping.

\[\hat m_t = \frac{m_t}{1-\beta_1^{t}}, \qquad \hat s_t = \frac{s_t}{1-\beta_2^{t}}\]

Bias correction — and it is genuinely necessary. Both averages start at zero, so early on they are biased towards zero. Assume the gradients are stationary with mean \(g\): unrolling gives \(\E[m_t]=(1-\beta_1^t)g\), so dividing by \((1-\beta_1^t)\) makes it unbiased. Without this, the very first steps of Adam are tiny — at \(t=1\), \(m_1=(1-\beta_1)g = 0.1g\) and \(\sqrt{s_1}=\sqrt{0.001}\,|g|\), so the raw ratio would be about \(3.16\), wildly wrong. Bias correction makes step 1 have magnitude \(\approx\eta\), as intended.

\[w_{t+1,i} = w_{t,i} - \eta\,\frac{\hat m_{t,i}}{\sqrt{\hat s_{t,i}}+\epsilon}\]

The Adam update. Note that \(\frac{\hat m}{\sqrt{\hat s}}\) is dimensionless and roughly bounded by 1 in magnitude, so \(\eta\) sets an approximate maximum step size in parameter units — which is why Adam's default \(\eta=10^{-3}\) transfers across problems so much better than SGD's does. Defaults \(\beta_1=0.9,\ \beta_2=0.999,\ \epsilon=10^{-8}\).

Common trap

Adam is the sensible default, not a universally superior method. Three things worth knowing:

SGD with momentum often generalises better on vision tasks, and the gap is a real and partially unexplained empirical finding, not folklore. Adam converges faster in training loss and sometimes to a worse test loss.
Adam's weight decay is not L2 regularisation. Adding \(\lambda\vw\) to the gradient gets divided by \(\sqrt{\hat s}\) along with everything else, so the effective decay differs per coordinate — which is not what you asked for. AdamW decouples them by subtracting \(\eta\lambda\vw\) directly, and it is the reason AdamW replaced Adam as the default for transformers.
Adam can fail to converge on some convex problems; the original convergence proof was wrong, and AMSGrad was proposed as a fix. In practice it rarely matters, but "Adam always works" is not a theorem.

Check yourself

Your training loss decreases smoothly for 2000 steps, then plateaus completely. You cut the learning rate by 10× and it immediately drops again to a new plateau. What is happening?

Correct. Correct, and this is the derivation above showing up in a real loss curve. The iterate had already reached the stationary distribution around the local optimum; further steps just resampled the cloud rather than shrinking it. Reducing \(\eta\) shrank the cloud.

What this tells you to do: the plateau is not "converged", it is "converged at this learning rate". A schedule that decays \(\eta\) is not an optional refinement — without it the run genuinely cannot reach the optimum. Increasing the batch size lowers the floor the same way, which is why progressive batch-size growth is an alternative to learning-rate decay.

Distinguishing it from other plateaus: a noise-floor plateau has visibly noisy loss values around a flat mean and responds instantly to an \(\eta\) cut. A saddle-point plateau has smooth, near-constant loss and small gradient norm, and does not respond to cutting \(\eta\) — if anything, cutting \(\eta\) makes it worse. Log the gradient norm and you can tell them apart.

Not quite. Overfitting would show as training loss continuing down while validation loss rises. Here training loss itself dropped on cutting \(\eta\), which is an optimisation phenomenon, not a generalisation one. (It is entirely possible for both to be happening; but the described symptom is the noise floor.)

Not quite. A smaller learning rate cannot make gradients larger — the gradient at a point does not depend on \(\eta\) at all. Vanishing gradients would also produce the opposite response: with tiny gradients, cutting \(\eta\) makes progress slower, not faster.

Five optimisers, one problem, honest comparisonPython · NumPy

import numpy as np
rng = np.random.default_rng(0)

# ---- a badly conditioned least-squares problem ----
n, d = 2000, 50
A = rng.normal(size=(d, d))
Q, _ = np.linalg.qr(A)
scales = np.logspace(0, -2, d)                 # condition number 100
X = rng.normal(size=(n, d)) @ (Q * scales) @ Q.T
w_true = rng.normal(size=d)
y = X @ w_true + 0.5*rng.normal(size=n)

H = 2 * X.T @ X / n
ev = np.linalg.eigvalsh(H)
kappa = ev.max()/ev.min()
print(f"lambda_max {ev.max():.4f}   lambda_min {ev.min():.6f}   kappa {kappa:.1f}")
print(f"stability limit eta < 2/lambda_max = {2/ev.max():.5f}")
print(f"optimal fixed eta   = 2/(lmax+lmin) = {2/(ev.max()+ev.min()):.5f}")
print(f"predicted GD rate   = (k-1)/(k+1)   = {(kappa-1)/(kappa+1):.6f} per step")
print()

w_star = np.linalg.lstsq(X, y, rcond=None)[0]
loss = lambda w: float(np.mean((X@w - y)**2))
L_star = loss(w_star)

def full_grad(w): return 2 * X.T @ (X@w - y) / n
def batch_grad(w, idx): return 2 * X[idx].T @ (X[idx]@w - y[idx]) / len(idx)

def run(name, step, T=400, stochastic=False, B=32, seed=1):
    r = np.random.default_rng(seed)
    w = np.zeros(d); state = {}
    for t in range(1, T+1):
        g = batch_grad(w, r.integers(0, n, B)) if stochastic else full_grad(w)
        w = step(w, g, t, state)
        if not np.isfinite(w).all(): return name, np.inf, np.inf
    return name, loss(w) - L_star, float(np.linalg.norm(w - w_star))

eta = 2/(ev.max()+ev.min())

def gd(w, g, t, s):        return w - eta*g
def mom(w, g, t, s):
    s['v'] = 0.9*s.get('v', 0) + g
    return w - eta*0.1*s['v']            # scale eta by (1-beta) to keep it stable
def nesterov(w, g, t, s):
    v_prev = s.get('v', np.zeros(d))
    s['v'] = 0.9*v_prev + g
    return w - eta*0.1*(0.9*s['v'] + g)
def rmsprop(w, g, t, s):
    s['s'] = 0.999*s.get('s', 0) + 0.001*g*g
    return w - 0.02*g/(np.sqrt(s['s']/(1-0.999**t)) + 1e-8)
def adam(w, g, t, s):
    s['m'] = 0.9*s.get('m', 0) + 0.1*g
    s['s'] = 0.999*s.get('s', 0) + 0.001*g*g
    mh = s['m']/(1-0.9**t); sh = s['s']/(1-0.999**t)
    return w - 0.02*mh/(np.sqrt(sh) + 1e-8)

print(f"{'optimiser':22s} {'excess loss':>14s} {'||w - w*||':>14s}")
print("-"*52)
for nm, st in [("GD", gd), ("GD + momentum", mom), ("Nesterov", nesterov),
               ("RMSProp", rmsprop), ("Adam", adam)]:
    name, ex, dist = run(nm, st)
    print(f"{name:22s} {ex:14.3e} {dist:14.3e}")
print()
for nm, st, B in [("SGD B=8", gd, 8), ("SGD B=64", gd, 64), ("SGD B=512", gd, 512)]:
    name, ex, dist = run(nm, st, stochastic=True, B=B)
    print(f"{name:22s} {ex:14.3e} {dist:14.3e}")
print("  -> noise floor falls roughly like 1/B, as derived")
print()

# ---- feature scaling is free preconditioning ----
Xs = (X - X.mean(0)) / X.std(0)
Hs = 2 * Xs.T @ Xs / n
print(f"kappa raw    {kappa:10.1f}")
print(f"kappa scaled {np.linalg.cond(Hs):10.1f}   <- the single cheapest optimisation improvement available")
print()

# ---- the eta sweep: where does it break? ----
print("eta sweep on full-batch GD, 400 steps:")
for mult in (0.1, 0.5, 1.0, 1.5, 1.9, 2.1):
    e = mult * 2/ev.max()
    w = np.zeros(d); ok = True
    for _ in range(400):
        w = w - e*full_grad(w)
        if not np.isfinite(w).all() or np.abs(w).max() > 1e10: ok = False; break
    tag = f"{loss(w)-L_star:.3e}" if ok else "DIVERGED"
    print(f"  eta = {mult:4.1f} x (2/lmax) = {e:.6f}  ->  {tag}")

Practical engineering note

A short list of things that actually go wrong, and what they look like.

Loss is nan from step one. Learning rate far too high, or a numerical issue in the loss itself — check for log(0) and use the fused stable forms from Phase A Lecture 2.
Loss explodes after a while. Learning rate slightly too high, or exploding gradients. Clip by global norm: \(\vg \leftarrow \vg\cdot\min(1, c/\|\vg\|)\).
Loss decreases far too slowly. Either \(\eta\) too small or \(\kappa\) too large. Check feature scaling first — it is free and often worth orders of magnitude.
Loss is flat from the start. Gradient is genuinely near zero: dead ReLUs, saturated sigmoids, or a bug where the gradient never reaches the parameters. Print \(\|\vg\|\) before blaming the optimiser.
Warmup. Adam's \(\hat s\) is a poor estimate for the first few dozen steps, so a large learning rate then can wreck the run. Linearly ramping \(\eta\) from 0 over a few hundred steps costs nothing and prevents it. This is standard for transformers for exactly this reason.

5Exercises

★

A quadratic loss has Hessian eigenvalues \(\{4, 0.04\}\).

What is the largest stable learning rate?
What is the optimal fixed learning rate, and the resulting per-step contraction?
Roughly how many steps to reduce the error by \(10^{6}\)?
What does momentum with the optimal \(\beta\) change these to?

1. \(\eta<2/\lambda_{\max}=2/4=0.5\).

2. \(\eta^\star = \frac{2}{\lambda_{\min}+\lambda_{\max}} = \frac{2}{4.04}=0.4950\). \(\kappa=4/0.04=100\), so the contraction is \(\frac{99}{101}=0.98020\) per step.

3. Need \(0.98020^t\le10^{-6}\), i.e. \(t\ge\frac{-6\ln10}{\ln0.98020}=\frac{-13.816}{-0.020} =689\) steps. (The approximation \(\frac{\kappa}{2}\ln\frac1\epsilon = 50\times13.82=691\) agrees.)

4. \(\sqrt\kappa=10\), so the optimal \(\beta=\left(\frac{10-1}{10+1}\right)^2 = \left(\frac{9}{11}\right)^2 = 0.6694\), and the rate becomes \(\frac{9}{11}=0.8182\).
Steps needed: \(t\ge\frac{-13.816}{\ln 0.8182}=\frac{-13.816}{-0.2007}=69\) — a tenfold reduction, matching \(\kappa\to\sqrt\kappa\).

Note the optimal \(\beta\) here (0.67) is well below the usual default of 0.9. The default corresponds to \(\sqrt\kappa\approx 39\), i.e. \(\kappa\approx1500\) — a reasonable guess for a neural network, and quite wrong for a well-conditioned problem. If momentum makes your training oscillate, \(\beta\) is too high for your \(\kappa\).

★★

Prove that the mini-batch gradient with batch size \(B\), sampled uniformly with replacement, has covariance \(\frac{1}{B}\mS\) where \(\mS\) is the covariance of the per-example gradients.

Then show that sampling without replacement gives \(\frac{1}{B}\cdot\frac{n-B}{n-1}\mS\), and interpret the extra factor. What happens at \(B=n\)?

With replacement. Let \(\vg_1,\dots,\vg_n\) be the per-example gradients, with mean \(\bar{\vg}\) and covariance \(\mS = \frac1n\sum_i(\vg_i-\bar{\vg})(\vg_i-\bar{\vg})^\top\). Draw indices \(I_1,\dots,I_B\) independently and uniformly. Then \(\vg_{\mathcal B}=\frac1B\sum_{j}\vg_{I_j}\) is an average of \(B\) i.i.d. draws each with covariance \(\mS\), so by Phase A Lecture 13's scaling rule

\[\Cov(\vg_{\mathcal B}) = \frac{1}{B^2}\cdot B\,\mS = \frac{\mS}{B}. \qquad\blacksquare\]

Without replacement. Now the draws are dependent — knowing \(I_1\) makes \(I_1\) less likely for \(I_2\). Expanding,

\[\Cov(\vg_{\mathcal B}) = \frac{1}{B^2}\Big[B\,\mS + B(B-1)\,\mathbf{C}\Big],\]

where \(\mathbf{C}=\Cov(\vg_{I_1},\vg_{I_2})\) for two distinct draws. For sampling without replacement from a finite population, \(\mathbf{C}=-\frac{\mS}{n-1}\) — the standard negative correlation, which holds because the full sum \(\sum_i \vg_i\) is fixed, so drawing an above-average gradient makes the next draw below-average on average. Substituting,

\[\Cov = \frac{\mS}{B} - \frac{(B-1)}{B(n-1)}\mS = \frac{\mS}{B}\cdot\frac{n-B}{n-1}. \qquad\blacksquare\]

Interpretation. The factor \(\frac{n-B}{n-1}\) is the finite population correction. Sampling without replacement is always at least as good, because you never waste a slot on a duplicate. The saving is negligible when \(B\ll n\) (the factor is \(\approx1\)) and total when \(B\) is a large fraction of \(n\).

At \(B=n\) the factor is exactly \(0\): sampling the whole dataset without replacement is the full-batch gradient, which has zero variance. Reassuringly, the formula knows this. Sampling with replacement at \(B=n\) still has variance \(\mS/n\), because a bootstrap sample is not the dataset.

Practical relevance: real training loops shuffle and partition into epochs, which is without-replacement sampling. For typical \(B/n\) around \(10^{-3}\) the distinction is irrelevant, and assuming independence (as the lecture's derivation does) is a good approximation.

★★

Show that Adam's bias correction is exactly what is needed.

Assume the gradients are i.i.d. with \(\E[g_t]=g\). Unroll \(m_t=\beta_1 m_{t-1}+(1-\beta_1)g_t\) with \(m_0=0\) and compute \(\E[m_t]\).
Show \(\E[m_t]/(1-\beta_1^t)=g\).
Compute the effective step size of an uncorrected Adam at \(t=1\) and explain why it is a problem.

1. Unrolling the recursion from \(m_0=0\):

\[m_t = (1-\beta_1)\sum_{k=1}^{t}\beta_1^{\,t-k}\,g_k.\]

Taking expectations and using \(\E[g_k]=g\) for all \(k\):

\[\E[m_t] = (1-\beta_1)\,g\sum_{k=1}^{t}\beta_1^{\,t-k} = (1-\beta_1)\,g\cdot\frac{1-\beta_1^{t}}{1-\beta_1} = \big(1-\beta_1^{t}\big)\,g,\]

using the finite geometric sum (Phase A Lecture 6).

2. Dividing: \(\E[\hat m_t] = \E[m_t]/(1-\beta_1^t) = g\) ✓. Exactly unbiased at every \(t\), not merely asymptotically. The same computation gives \(\E[s_t]=(1-\beta_2^t)\E[g^2]\).

3. At \(t=1\) with the defaults: \(m_1=(1-0.9)g=0.1g\) and \(s_1=(1-0.999)g^2=0.001g^2\), so \(\sqrt{s_1}=0.0316|g|\). The uncorrected update would be

\[\eta\,\frac{m_1}{\sqrt{s_1}} = \eta\,\frac{0.1\,g}{0.0316\,|g|} = 3.16\,\eta\,\operatorname{sign}(g),\]

a step 3.16 times larger than intended. With correction, \(\hat m_1 = g\) and \(\sqrt{\hat s_1}=|g|\), giving exactly \(\eta\).

Why it matters. The mismatch is worst precisely at the start, when the parameters are random and a 3× overshoot can push them somewhere they never recover from. The asymmetry between \(\beta_1=0.9\) and \(\beta_2=0.999\) is what creates it — the second moment warms up ten times more slowly than the first. It is also why warmup helps even with bias correction: correcting the bias in the mean does not fix the enormous variance of \(\hat s_t\) when it has only seen a handful of samples.

★★★

Coding. Reproduce the \(\kappa\) and \(\sqrt\kappa\) scaling laws empirically.

Build diagonal quadratics with \(\kappa\in\{10,100,1000,10^4\}\).
For each, run GD at the optimal \(\eta\) and momentum at the optimal \((\eta,\beta)\), counting steps to reach \(\|\ve\|\le10^{-6}\|\ve_0\|\).
Plot (or tabulate) steps against \(\kappa\) on a log-log scale and estimate the slope for each method.
Now add gradient noise and repeat. What happens to the advantage of momentum, and why?

import numpy as np
def steps_to(kappa, method, tol=1e-6, T=2_000_000, noise=0.0, seed=0):
    r = np.random.default_rng(seed)
    lam = np.array([1.0, 1.0/kappa])
    e0 = np.array([1.0, 1.0]); e = e0.copy(); v = np.zeros(2)
    if method == 'gd':
        eta, beta = 2/(lam.max()+lam.min()), 0.0
    else:
        sk = np.sqrt(kappa)
        beta = ((sk-1)/(sk+1))**2
        eta  = (1+np.sqrt(beta))**2 / lam.max() * 0.9    # heavy-ball choice
    for t in range(1, T+1):
        g = lam*e + (noise*r.normal(size=2) if noise else 0.0)
        v = beta*v + g
        e = e - eta*v
        if not np.isfinite(e).all(): return np.nan
        if np.linalg.norm(e) <= tol*np.linalg.norm(e0): return t
    return np.nan

print(f"{'kappa':>8} {'GD':>10} {'momentum':>10} {'ratio':>8} {'sqrt(k)':>9}")
res = []
for k in (10, 100, 1000, 10000):
    a, b = steps_to(k,'gd'), steps_to(k,'mom')
    res.append((k, a, b))
    print(f"{k:8d} {a:10.0f} {b:10.0f} {a/b:8.1f} {np.sqrt(k):9.1f}")

ks = np.log([r[0] for r in res])
for j, nm in ((1,'GD'), (2,'momentum')):
    ys = np.log([r[j] for r in res])
    slope = np.polyfit(ks, ys, 1)[0]
    print(f"log-log slope for {nm:9s}: {slope:.3f}")

print("\\nwith gradient noise = 0.01:")
for k in (10, 100, 1000):
    a, b = steps_to(k,'gd',noise=0.01), steps_to(k,'mom',noise=0.01)
    print(f"  kappa={k:6d}  GD {a}  momentum {b}")

2 & 3. The step counts should come out near \(690\kappa/100\) for GD and near \(70\sqrt\kappa/10\) for momentum. The log-log slopes come out at approximately 1.00 for GD and 0.50 for momentum — direct confirmation of \(O(\kappa)\) versus \(O(\sqrt\kappa)\). At \(\kappa=10^4\) the ratio is about \(100\times\), which is why momentum is on by default everywhere.

4. With noise, both methods stop improving once they hit their noise floor, and neither reaches \(10^{-6}\) at all — steps_to returns nan for the larger \(\kappa\). The advantage of momentum largely disappears in the noise-dominated regime.

The reason is that momentum's benefit is entirely about the deterministic contraction, and once the error is at the noise floor the deterministic part has already done its job. Worse, momentum amplifies the noise: the accumulated \(\vv\) is an average of \(\approx\frac{1}{1-\beta}\) noisy gradients, so its variance is reduced by that factor but the effective step is amplified by the same factor, and the net effect on the stationary variance is roughly \(\frac{\eta}{1-\beta}\) instead of \(\eta\).

The practical reading: momentum helps most early in training, when the deterministic gradient dominates and the landscape is ill-conditioned. Late in training, when you are on the noise floor, what helps is decaying \(\eta\) or growing \(B\). This is exactly the shape of standard recipes — high learning rate with momentum early, aggressive decay late — and it is not a coincidence.

Next: we have been assuming squared error without justification. Lecture 04 derives losses from noise models — one principle that produces squared error, cross-entropy, Poisson loss and the robust alternatives, and tells you which to use when.

Previous2 · Linear Regression, Properly Next4 · Where Loss Functions Come From