Optimization: Walking down a mountain blindfolded

In 1847 Augustin-Louis Cauchy presented to the Paris Academy of Sciences a three-page note entitled “General method for solving systems of simultaneous equations”. He was interested in astronomical calculations: determining the orbit of a celestial body from observations leads to systems of equations that cannot be solved exactly. His proposal was to turn the problem into minimizing a function that measures the error and to walk down it step by step, always in the direction of steepest slope. He had invented gradient descent.

The algorithm

Start from any parameters $θ_{0}$ and repeat:

θ_{t + 1} = θ_{t} - η \nabla ℒ (θ_{t}) .

The gradient, which backpropagation computes cheaply, points in the direction of steepest ascent, so we step the opposite way. The number $η > 0$ is the learning rate: the size of the step. It is the most important hyperparameter of any training run.

Click on the map to choose the starting point. The coloured bands are contour lines: the more intense the colour, the higher the loss. On the elongated bowl, try

η = 0.05

(slow progress),

0.18

(zigzag) and

0.21

(divergence): the limit is exactly

2 / λ_{max} = 0.2

. On the landscape with several minima, the valley you reach depends on where you start. Momentum

β

adds inertia: it crosses narrow valleys faster.

When does it work? Convexity and smoothness

Guaranteeing that the algorithm works takes two conditions on the landscape. A function is convex if the segment between any two points of its graph lies above it, that is, if it is bowl-shaped. It is $β$ -smooth if its gradient changes no faster than $β$ : $‖ \nabla f (x) - \nabla f (y) ‖ \leq β ‖ x - y ‖$ .

Theorem (minima of convex functions)

If $f$ is convex, every local minimum is a global minimum. If it is also differentiable, $\nabla f (x^{*}) = 0$ is enough for $x^{*}$ to be a global minimum.

Theorem (convergence of gradient descent)

If $f$ is convex and $β$ -smooth and $η = 1 / β$ is used, then after $k$ steps

f (θ_{k}) - f (θ^{*}) \leq \frac{β ‖ θ_{0} - θ^{*} ‖^{2}}{2 k} .

If it is also $μ$ -strongly convex, convergence is geometric: the error is multiplied at each step by at most $1 - μ / β$ .

The quadratic case, where everything is visible

Take $f (θ) = \frac{1}{2} θ^{⊤} A θ$ with $A$ symmetric and eigenvalues $0 < λ_{min} \leq \dots \leq λ_{max}$ . Since $\nabla f = A θ$ , the algorithm does $θ_{t + 1} = (I - η A) θ_{t}$ . In the eigenvector basis, each coordinate evolves on its own:

θ_{t}^{(i)} = (1 - η λ_{i})^{t} θ_{0}^{(i)} .

It converges if and only if $| 1 - η λ_{i} | < 1$ for every $i$ , that is, if $0 < η < 2 / λ_{max}$ . If $η λ_{i} > 1$ , that coordinate changes sign at every step (the zigzag). The slowest direction is that of $λ_{min}$ , and with the best $η$ the error shrinks per step by a factor $\frac{κ - 1}{κ + 1}$ , where $κ = λ_{max} / λ_{min}$ is the condition number. Narrow, elongated valleys (large $κ$ ) are slow.

Neural networks are not convex: their loss landscape is full of local minima, plateaus and saddle points. No classical theorem guarantees that gradient descent finds a good minimum there, and yet it does, again and again. Some partial explanations: in high dimension almost all critical points are saddle points rather than minima (Dauphin et al., 2014), and in heavily overparameterized networks most local minima are almost as good as the global one. It is one of the big open questions of the field.

Inertia: momentum and acceleration

In an elongated valley the gradient almost always points at the walls rather than towards the bottom. In 1964 Boris Polyak proposed the heavy-ball method: accumulate a velocity that averages previous gradients,

v_{t + 1} = β v_{t} - η \nabla ℒ (θ_{t}), θ_{t + 1} = θ_{t} + v_{t + 1} .

The oscillations cancel out and progress in the right direction builds up. In 1983 Yurii Nesterov found a variant that reaches error $O (1 / k^{2})$ instead of $O (1 / k)$ for smooth convex functions, and showed that this is the best possible:

Theorem (Nesterov's lower bound)

No method that only uses gradients can guarantee, for every convex $β$ -smooth function (in high enough dimension), an error smaller than $\frac{3 β ‖ θ_{0} - θ^{*} ‖^{2}}{32 (k + 1)^{2}}$ after $k$ steps. Nesterov's accelerated method achieves that order.

Stochastic gradient descent

The loss of an LLM is an average over trillions of training words: $ℒ (θ) = \frac{1}{N} \sum_{i = 1}^{N} ℓ_{i} (θ)$ . Computing the exact gradient at every step would mean going through all the data. The alternative is to estimate it with a random mini-batch $B$ of a few examples:

g_{t} = \frac{1}{| B |} \sum_{i \in B} \nabla ℓ_{i} (θ_{t}), 𝔼 [g_{t}] = \nabla ℒ (θ_{t}) .

The estimate is noisy but unbiased: on average it points the right way. Herbert Robbins and Sutton Monro proved in 1951 that this kind of procedure converges if the learning rate decreases at the right pace.

Theorem (Robbins–Monro conditions, 1951)

Under regularity conditions, stochastic approximation converges if the rates $η_{t}$ satisfy

\sum_{t = 1}^{\infty} η_{t} = \infty and \sum_{t = 1}^{\infty} η_{t}^{2} < \infty,

for example with $η_{t} = c / t$ . The first condition ensures you can get anywhere, and the second that the noise eventually dies out.

Fitting a line

\hat{y} = w x + b

with stochastic gradient descent. Left: the data and the current line. Right: the path of the parameters

(w, b)

over the contour lines of the loss. With the full batch the path is smooth; with batches of one example it is erratic, but each step costs 40 times less.

That is why small batches are always used in practice: many cheap, noisy steps beat a few expensive, exact ones. The noise also seems to help escape “sharp” minima and find “flatter” ones, which tend to generalize better.

Adam: a learning rate for every parameter

The parameters of a network are not alike: some receive huge gradients and others tiny ones. In 2014 Diederik Kingma and Jimmy Ba proposed Adam, which combines momentum with a normalization by the recent magnitude of each gradient:

m_{t} = β_{1} m_{t - 1} + (1 - β_{1}) g_{t}, v_{t} = β_{2} v_{t - 1} + (1 - β_{2}) g_{t}^{2}, θ_{t + 1} = θ_{t} - η \frac{{\hat{m}}_{t}}{\sqrt{{\hat{v}}_{t}} + ϵ} .

Here ${\hat{m}}_{t}$ and ${\hat{v}}_{t}$ are the averages corrected for their initial bias. Its variant with decoupled weight decay, AdamW (2017), is the optimizer practically every current LLM is trained with, almost always with a rate that rises during an initial warm-up and then decays, in the spirit of what Robbins and Monro asked for.

We now have a general procedure for fitting parameters. The next question is which family of functions to parameterize, and here the oldest idea in AI appears: imitating neurons.

References

A.-L. Cauchy (1847). “Méthode générale pour la résolution des systèmes d'équations simultanées”. Comptes Rendus de l'Académie des Sciences, 25.
H. Robbins and S. Monro (1951). “A Stochastic Approximation Method”. Annals of Mathematical Statistics, 22(3).
B. T. Polyak (1964). “Some methods of speeding up the convergence of iteration methods”. USSR Computational Mathematics and Mathematical Physics, 4(5).
Y. Nesterov (1983). “A method for solving the convex programming problem with convergence rate O(1/k²)”. Doklady Akademii Nauk SSSR, 269.
Y. Dauphin et al. (2014). “Identifying and attacking the saddle point problem in high-dimensional non-convex optimization”. NeurIPS.
D. P. Kingma and J. Ba (2015). “Adam: A Method for Stochastic Optimization”. ICLR.
I. Loshchilov and F. Hutter (2019). “Decoupled Weight Decay Regularization”. ICLR.
S. Bubeck (2015). “Convex Optimization: Algorithms and Complexity”. Foundations and Trends in Machine Learning, 8(3–4).

§The algorithm

§When does it work? Convexity and smoothness

§Inertia: momentum and acceleration

§Stochastic gradient descent

§Adam: a learning rate for every parameter