Optimisation Theory

Classical optimisation theory says you should not be able to train a neural network. The loss is wildly non-convex, has exponentially many critical points, and gradient descent has no convergence guarantees. Yet SGD on a 100B-parameter transformer reliably finds solutions that generalise. Understanding why this works - and where the classical intuitions break - is what separates a practitioner who can train models from one who can debug them.

Convex vs non-convex landscapes

A function f is convex if for all x, y and lambda in [0, 1]:

f(lambda x + (1 - lambda) y) <= lambda f(x) + (1 - lambda) f(y)

Equivalently, every local minimum is a global minimum and the Hessian is positive semi-definite everywhere. Convex optimisation has a complete theory: gradient descent converges at rate O(1/t), accelerated methods at O(1/t^2), the optimum is unique. Linear regression, logistic regression, SVMs are all convex.

Neural networks are aggressively non-convex. The loss landscape has:

• Permutation symmetries. Swap any two hidden units; the loss is unchanged. A 1000-unit layer has 1000! equivalent global minima.
• Scaling symmetries. Multiply one layer's weights by c, divide the next by c (for ReLU networks). Same function, different parameters.
• Saddle points at every scale. Far more saddles than minima, especially in high dimension.

Despite this, SGD finds solutions that perform well on held-out data. Why?

Why SGD works on non-convex landscapes

Three empirical observations, none fully explained theoretically:

1. Loss-surface findings. Visualisations of large-network loss landscapes (Goodfellow et al, Li et al) show that the path from initialisation to the trained solution is essentially monotonic, no high-loss ridges to cross. The geometry near the solution is wide and flat in most directions.
2. Local minima are mostly equivalent. Choromanska et al (2015) and others argued, with caveats, that for large enough networks the local minima found by SGD all sit at similar loss values. There is no "vastly better" minimum being missed.
3. The lottery-ticket framing. Frankle and Carbin (2018) showed that dense networks contain sparse subnetworks ("winning tickets") that, when trained from the original initialisation, match or exceed the dense network's performance. The implication: overparameterisation gives SGD many independent paths to a good solution.

The working intuition: in very high dimension, a "bad" critical point would require all eigenvalues of the Hessian to be positive (a true minimum that is worse than what SGD finds). The probability of this drops sharply as dimension grows.

Saddle points dominate

Dauphin et al (2014) showed that critical points in high-dimensional non-convex landscapes are overwhelmingly saddles, not local minima. Reasoning: at a random critical point, each eigenvalue of the Hessian is independently roughly equally likely to be positive or negative. The probability that all n are the same sign drops as 2^{-n}.

Consequences for training:

• Loss plateaus during training often reflect slow escape from saddles, not actual local minima.
• Stochastic noise from minibatching helps escape saddles - the noise occasionally kicks the optimiser off the ridge into a descending direction.
• Pure second-order methods (Newton) are attracted to saddles, because they converge to any stationary point. This is one practical argument against Newton in deep learning.

Momentum as a damped oscillator

The momentum update:

v_t = mu * v_{t-1} + grad
x_t = x_{t-1} - lr * v_t

Read as a discretisation of a continuous-time differential equation, this is a damped harmonic oscillator. The parameter mu controls damping; lr controls the time step. Gabriel Goh's Distill article "Why Momentum Really Works" walks through the analysis: on a quadratic with condition number kappa, plain gradient descent converges at rate (kappa - 1) / (kappa + 1) per step, while optimal momentum gets (sqrt(kappa) - 1) / (sqrt(kappa) + 1). For kappa = 10000, momentum is 100x faster per step.

The mental model:

• mu near 0: no memory, behaves like plain SGD.
• mu near 1: under-damped, oscillates around minima with slow decay.
• mu = 0.9 (the universal default): roughly critically damped on typical curvatures, smooths gradient noise without wild oscillation.

Adam as adaptive RMSProp + momentum

Adam tracks two running averages per parameter:

m_t = beta1 * m_{t-1} + (1 - beta1) * grad        # momentum-like
v_t = beta2 * v_{t-1} + (1 - beta2) * grad^2      # variance-like
update = m_t / (sqrt(v_t) + eps)

Two effects:

1. Per-parameter learning rate scaling. Parameters with large historical gradient magnitudes get smaller effective steps. This is what makes Adam robust across model architectures - the LR you choose globally is rescaled per parameter to the appropriate magnitude.
2. Momentum smoothing. Same accelerative effect as classical momentum.

Adam can be derived as a diagonal approximation to natural gradient under specific assumptions. The diagonal approximation is what makes it tractable; it is also what makes it suboptimal compared to (intractable) full natural gradient. AdamW (decoupled weight decay) is the actually-used variant - see the optimiser note for details.

Learning-rate schedules as implicit regularisation

The schedule (cosine, linear, step) often matters more than the optimiser. The dominant modern recipe:

1. Warmup (1-5% of steps) to avoid huge first-step updates from Adam bias correction.
2. Cosine decay to ~10% of peak over the rest.

There is mounting evidence that the schedule itself is a form of regularisation. Smith et al and others observed that decaying the learning rate has a similar effect to increasing batch size, both of which reduce gradient noise. Reduced noise at the end of training lets the optimiser settle into a wider, flatter minimum, which generalises better.

Implications:

• Stopping cosine decay early ("warmup-stable-decay" schedules used in some Llama variants) does measurably worse than running it to the end.
• Re-warming after a partial run is hard to do without losing some generalisation.
• The same final LR with different schedule shapes can give noticeably different held-out loss.

When the intuition fails

• Non-convex theory does not predict success. Convergence proofs for SGD on non-convex losses give bounds of the form "approximate stationary point in O(1/eps^4) steps." Useless in practice; SGD trains a 70B model to convergence in ~1e5 steps.
• Sharpness vs flatness is contested. "Flat minima generalise better" is a working heuristic but counterexamples exist (Dinh et al 2017). The relationship between geometry and generalisation is genuinely not settled.
• Batch size effects are non-linear. Doubling batch size and halving LR (or scaling LR by sqrt(2)) is the textbook rule. It breaks above some critical batch size that varies by model and dataset. Past that point, you waste compute for no convergence speedup.
• Learning rate transfer across model sizes. A learning rate tuned at 1B parameters does not transfer to 70B without correction. Mup (Yang et al) and related re-parameterisations try to make optimal hyperparameters scale-invariant, with partial success.

Further reading

• Why Momentum Really Works - Gabriel Goh on Distill; the canonical visual explanation of momentum.
• Identifying and attacking the saddle point problem in high-dimensional non-convex optimization - Dauphin et al 2014.
• Adam: A Method for Stochastic Optimization - Kingma and Ba.
• Entropy-SGD: Biasing Gradient Descent Into Wide Valleys - Chaudhari et al on flat minima and generalisation.
• Deep Learning Book - Chapter 8: Optimization - the standard ML-focused treatment.