Optimization

The problem and the classical rates#

Minimise f(theta) over theta in R**n, with access to grad_f. Convergence guarantees depend entirely on which assumptions hold.

kappa = L / mu is the condition number. For the quadratic in the Developer block, L = 500 and mu = 1, so kappa = 500. Plain descent contracts the flat direction by 1 - eta * mu per step; with eta capped at 2/L, that is 1 - 2/kappa. Reaching a fixed accuracy therefore needs O(kappa) steps, while Nesterov needs O(sqrt(kappa)). The measured gap in that experiment is exactly this, and it widens as the landscape stretches.

The stability threshold eta < 2/L comes from the descent lemma; the derivation is in calculus.

Nesterov's lower bound (1983) says no first-order method beats O(1/t**2) on this class. Acceleration is therefore optimal, not an incremental gain.

Momentum#

Polyak's heavy ball (1964) and Nesterov's accelerated gradient (1983) differ in where the gradient is evaluated:

heavy ball:  v_{t+1} = beta v_t + grad_f(x_t)          x_{t+1} = x_t - eta v_{t+1}
Nesterov:    v_{t+1} = beta v_t + grad_f(x_t - eta beta v_t)
             x_{t+1} = x_t - eta v_{t+1}

Nesterov evaluates the gradient at the look-ahead point, which damps overshoot and yields the accelerated rate on convex problems. Heavy ball achieves the same rate on quadratics. On general smooth strongly convex functions it can fail to converge at all (Lessard, Recht and Packard, 2016).

The effective step along a consistently-signed direction is amplified by 1 / (1 - beta), so beta = 0.9 multiplies it by roughly 10. That is why momentum and learning rate must be retuned together.

Adaptive methods#

Adam (Kingma and Ba, 2015):

m_t = b1 m_{t-1} + (1 - b1) g_t                first moment, an average of recent gradients
v_t = b2 v_{t-1} + (1 - b2) g_t**2             second moment, elementwise
m_hat = m_t / (1 - b1**t)                      bias correction
v_hat = v_t / (1 - b2**t)
theta_t = theta_{t-1} - eta * m_hat / (sqrt(v_hat) + eps)

Bias correction is not cosmetic. With b2 = 0.999 and v_0 = 0, the uncorrected v_t is low by a factor of 1 - 0.999**t, which is about 0.001 at t = 1. Ignoring it produces enormous or vanishing early steps depending on which correction is dropped.

Three results are worth carrying:

• Adam can fail to converge on convex problems. Reddi, Kale and Kumar (2018) constructed an online convex counterexample where the exponential moving average of v forgets a rare large gradient. AMSGrad, which uses a running maximum of v, fixes it in theory and is rarely used in practice.
• Decoupled weight decay. Loshchilov and Hutter (2019) showed that adding lambda * theta to the gradient inside Adam gives per-parameter-scaled decay, which is not L2 regularisation. AdamW applies theta <- theta - eta * lambda * theta separately. This is now the default for transformer training.
• eps placement matters. Adding it inside the square root instead of outside changes behaviour for small gradients, and implementations differ. Check the source before porting a hyperparameter between frameworks.

Adam is scale-invariant per coordinate. That is why it tolerates badly conditioned and heterogeneous parameter groups, such as embeddings and layer norms alongside dense weights.

On well-conditioned vision problems, well-tuned SGD with momentum still frequently generalises better (Wilson et al., 2017).

Batch size and gradient noise#

For a mini-batch of size B drawn i.i.d., the gradient estimator is unbiased with covariance Sigma / B. The relative error therefore falls as 1/sqrt(B), which is the trend measured in the Developer block.

McCandlish et al. (2018) formalised the trade-off with the gradient noise scale B_crit = tr(Sigma) / (grad.T @ H @ grad), approximately tr(Sigma) / ||grad||**2. Below B_crit, doubling the batch roughly halves the number of steps needed. Above it, extra batch size buys almost nothing and wastes compute. B_crit grows during training as the gradient shrinks, which is the basis for batch-size ramps.

Goyal et al. (2017) established the linear scaling rule: multiply the learning rate by k when multiplying the batch by k. With a gradient warmup, they trained ImageNet at batch 8192 without accuracy loss.

The rule breaks down at very large batch sizes. LARS and LAMB (You et al., 2017, 2020) use layerwise trust ratios to push further.

The relevant scaling under SGD is the ratio eta / B, which controls the size of the stationary noise ball. Two runs with the same ratio behave similarly, which is a useful invariant when changing hardware.

Schedules#

Warmup exists because early curvature estimates are unreliable. For Adam specifically, v_hat has high variance in the first steps; Liu et al. (2020) analysed this and proposed RAdam as a principled alternative to a hand-set warmup. Empirically, linear warmup over 500 to 4000 steps remains the standard for transformers. Its absence is a leading cause of early divergence.

Decay shrinks the noise ball as training proceeds. Cosine decay (Loshchilov and Hutter, 2017) is the common default. Fixed-budget schedules must be planned in advance, which motivated later work on schedule-free and warmup-stable-decay approaches.

Smith and Le (2018) argued that two moves are near-equivalent: decaying the learning rate, and increasing the batch size. Both reduce the same noise term. The choice between them is about hardware utilisation.

What actually stops progress#

Local minima are not the problem. Dauphin et al. (2014) argued from random matrix theory that critical points in high dimensions are overwhelmingly saddles. Requiring all n Hessian eigenvalues to be positive is exponentially unlikely. Ge et al. (2015) showed noisy SGD escapes strict saddles — those with at least one strictly negative eigenvalue — in polynomial time.

Plateaus and conditioning are the problem. Long regions of small gradient dominate wall-clock time, and the condition number sets the step count.

Not every minimum is equally good. Keskar et al. (2017) reported that large-batch training finds sharper minima that generalise worse, though Dinh et al. (2017) showed sharpness as usually defined is not reparameterisation-invariant, so the argument needs care. Sharpness-Aware Minimisation (Foret et al., 2021) optimises the worst case within a neighbourhood. It improves generalisation, at roughly double the compute.

Cohen et al. (2021) documented the edge of stability. Full-batch gradient descent drives the largest Hessian eigenvalue to about 2/eta, then oscillates there. The loss keeps falling, non-monotonically. Classical smooth-optimisation theory does not describe this regime, and it is the regime real training occupies.

References#

• Nesterov, Y. Lectures on Convex Optimization, 2nd ed., Springer, 2018.
• Polyak, B. T. "Some methods of speeding up the convergence of iteration methods." USSR Comp. Math. 4(5), 1964.
• Robbins, H., Monro, S. "A Stochastic Approximation Method." Ann. Math. Statist. 22(3), 400–407, 1951.
• Kingma, D. P., Ba, J. "Adam: A Method for Stochastic Optimization." ICLR, 2015. arXiv:1412.6980
• Reddi, S. J., Kale, S., Kumar, S. "On the Convergence of Adam and Beyond." ICLR, 2018. arXiv:1904.09237
• Loshchilov, I., Hutter, F. "Decoupled Weight Decay Regularization." ICLR, 2019. arXiv:1711.05101
• Loshchilov, I., Hutter, F. "SGDR: Stochastic Gradient Descent with Warm Restarts." ICLR, 2017. arXiv:1608.03983
• Goyal, P., et al. "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour." arXiv:1706.02677, 2017.
• McCandlish, S., Kaplan, J., Amodei, D., et al. "An Empirical Model of Large-Batch Training." arXiv:1812.06162, 2018.
• Smith, S. L., Le, Q. V. "Don't Decay the Learning Rate, Increase the Batch Size." ICLR, 2018. arXiv:1711.00489
• Dauphin, Y., et al. "Identifying and attacking the saddle point problem." NeurIPS, 2014. arXiv:1406.2572
• Keskar, N. S., et al. "On Large-Batch Training for Deep Learning." ICLR, 2017. arXiv:1609.04836
• Foret, P., Kleiner, A., Mobahi, H., Neyshabur, B. "Sharpness-Aware Minimization." ICLR, 2021. arXiv:2010.01412
• Wilson, A. C., et al. "The Marginal Value of Adaptive Gradient Methods in Machine Learning." NeurIPS, 2017. arXiv:1705.08292
• Lessard, L., Recht, B., Packard, A. "Analysis and Design of Optimization Algorithms via Integral Quadratic Constraints." SIAM J. Optim. 26(1), 2016.
• Cohen, J. M., et al. "Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability." ICLR, 2021. arXiv:2103.00065
• Liu, L., et al. "On the Variance of the Adaptive Learning Rate and Beyond" (RAdam). ICLR, 2020. arXiv:1908.03265

What to learn next#

• Gradient descent — the algorithm as implemented.
• Backpropagation — where each g_t comes from.
• Experiment tracking — recording the sweeps this lesson forces you to run.