NUM~ Numerical methods Algorithm

Gradient descent

Also known as: Gradient descent · steepest descent · heavy ball method · heavy ball · momentum · SGD

A point rolls down the mountains above Almaty, every step against the gradient of the real heights. This is how the minimum is found for functions with no formula for the answer, from a sum of squared misses to a neural network's loss. How to choose the step, why the ball needs momentum and why on real terrain it gets stuck in pits.

Intermediate 25 minutes optimization gradient iterative approximation geodata science and physics classic interviews

In front of you are the mountains south of Almaty, thirty-odd kilometres from the edge of the city to the main ridge; the heights come from NASA’s radar survey from the Space Shuttle Endeavour in 2000. Tap the map anywhere, and a point will roll from there. It does not see the whole map: at each step it knows only the slope under its feet and steps where the slope goes down most steeply. That is gradient descent, the way to find the minimum of a function when there is no formula for the answer.

The first point starts from Big Almaty Peak, at a height of 3620 metres, and does not reach the city: at 1271 metres, in a valley, it begins to jump on the spot. Try other steps. A cautious step gets stuck even higher, at 1654 metres, a huge one jumps chaotically over ridges, and a point with momentum gets down almost to the plain, to 749 metres. Press “Rain”: of a hundred and fifty drops, a dozen and a half leave the map, and the rest get stuck. Why, and what is done about it, is below.

The idea

Let a function f depend on several numbers: here on the two coordinates of a point, for a neural network on billions of weights. Its gradient ∇f is the vector of partial derivatives: how fast f grows if you move a little along each coordinate. It points in the direction of steepest ascent, and its length equals the steepness.

If you don’t know where the minimum is but can measure the slope under your feet, step against the gradient: x ← x − η∇f(x). A small step in that direction always decreases the function, and repeating the steps, the point slides down into a pit.

Why it decreases: for a small step the function is almost linear, f(x − η∇f) ≈ f(x) − η·|∇f|², and the subtracted term is positive as long as the gradient is not zero. The number η is the step length; in the scene it is “metres of path per unit of slope”, so on a slope of steepness 0.5 with η = 300 the point travels 150 metres. The step length changes by itself with the steepness: on a steep slope the point rushes, on a gentle one it crawls, and where it is flat it stops.

Descent finds not the lowest minimum but the nearest one, a point where the gradient has become zero. On a smooth convex function, like the sum of squared misses in linear regression, these are the same thing. In the mountains they are not.

Where the gradient comes from

The heights in the scene are a table: 201 × 201 points every 163–167 metres. Between grid nodes the height is filled in from the four neighbours: along the north side of the cell one straight line, along the south side another, and between them a third. Such a function is called bilinear, and its derivatives are written in one line. The listing below does exactly this, and the scene computes the same way.

When a function is given by a formula, its derivatives are derived by the rules of calculus. A neural network’s formula is a long chain of simple operations, and the gradient over all the weights at once is found by backpropagation: a pass along the chain from end to start costs only a few times more than computing the function itself. But computing the gradient head-on, shifting each variable in turn, costs d + 1 function evaluations per step: for two coordinates that is nothing, for a billion weights it is out of reach.

The step

The whole tuning of descent is one number η, and it decides everything. The listing lets the point go from three places in four ways and prints where it stopped:

Big Almaty Peak, 3620 m
  step   30                 stuck at 1654 m after 2391 steps
  step  300                 stuck at 1271 m after 729 steps
  step 3000                 stuck at 802 m after 430 steps
  step  300, momentum 0.9   stuck at 749 m after 547 steps

A small step is careful: the point honestly goes down, but on the gentle floor of a valley it takes steps of a few metres, and any dip the size of a grid cell stops it. A big step flies over small dips, but a huge one no longer has anything in common with descent: kilometres long, it jumps over ridges and sometimes lands higher than it stood. That it ended up lower is luck, not a sign that the method is good.

There is also an exact limit. Near a minimum any smooth function looks like a bowl, and a bowl has a largest curvature L, the sharpness of its bend in the steepest direction. If η > 2/L, each step overshoots the bottom and lands on the opposite slope higher than it was: the descent blows up. The proof for a bowl with its axes along the coordinates is in “Mathematics, the Queen of the Sciences”, and for a rotated bowl in the first problem of the article.

The ravine

Descent has the worst time in a ravine: the walls are steep, the floor gentle. A step safe for the walls is too small on the floor, and the point rushes from wall to wall, slowly sliding along. In the scene this shows in the valleys: the path becomes a saw. How bad it is is told by the ratio of the largest curvature to the smallest, the condition number κ = L/μ. On the function (x² + κy²)/2 with the best constant step 2/(1 + κ), descent needs about κ steps for each reduction of the error by a factor of e² ≈ 7.4, and to reach a precision of 10⁻⁶ with κ = 100, 709 steps. Stretch the ravine a hundred times and there are a hundred times as many steps.

The heavy ball

A real ball in a ravine behaves more cleverly than a point: it speeds up along the floor, and its own inertia damps the rushing across. In 1964 Boris Polyak proposed giving descent the same thing, a velocity that remembers past steps:

$$v \leftarrow \beta v - \eta\,\nabla f(x), \qquad x \leftarrow x + v.$$

This is the heavy ball method, in machine learning called momentum. With β = 0 it is ordinary descent. With β = 0.9, where the gradient points the same way for many steps in a row, the velocity builds up almost 1/(1 − β) = 10 times; where it changes sign every other step, across the ravine, neighbouring pushes cancel each other.

For a quadratic function Polyak found the best parameters: η = 4/(√L + √μ)², β = ((√L − √μ)/(√L + √μ))². With them the number of steps grows not as κ but as √κ:

 kappa   plain  momentum  ratio
    10      71        27    2.6
   100     709        95    7.5
  1000    7082       321   22.1
 10000   70811      1074   65.9

At κ = 10,000 it is 1074 steps instead of 70,811. Notice that the ball’s step η is greater than 2/L: without momentum descent would blow up with such a step, but with it, it does not.

Polyak’s guarantee is for a quadratic function and for a neighbourhood of the minimum. Far from it, the ball with the same parameters may fail to converge: in 2014 Laurent Lessard, Benjamin Recht and Andrew Packard built a strongly convex function of one variable with κ = 25 on which it does not converge to the minimum. A method that accelerates with a guarantee on any smooth convex function was proposed by Yurii Nesterov in 1983; it differs in that the gradient is taken not at the current point but where momentum is already carrying the point.

Pits

Go back to the mountains and press “Rain”. The drops get stuck in pits, places where the gradient points uphill on all sides. On real terrain there are hundreds of them: the basin of Big Almaty Lake, the unevenness of valley floors, the roofs and trees of the city, which the radar recorded together with the ground. Water in such a pit would rise until it overflowed the edge and would run on. Descent cannot do that: it sees only the slope under its feet, and in a pit the slope leads back on all sides.

Hence all the usual remedies. Run descent many times from different places and take the best result. Add momentum: the ball runs through small dips at speed, but it can also run through the one you need. Add noise, more on that below. And there are problems with no pits at all: a convex function, like a sum of squared misses, has a single minimum, and descent finds it from any point. In training neural networks the function is not convex, but according to the observations of Yann Dauphin and co-authors (2014), in spaces of millions of variables points with zero gradient more often turn out to be saddles, down in some directions and up in others, than pits, and what mostly hinders descent is the flat plateaus around such saddles.

Code

The first listing is descent over the real terrain. The heights are on the server where the code runs, on the same grid as in the scene. The function height_and_slope is the bilinear height and its gradient, descend is the descent itself with momentum (without it when β = 0).

When to stop? For a smooth function, when the gradient is almost zero. For terrain made of cells this does not work: in a pit a point with a big step jumps from wall to wall forever, and the gradient under its feet does not shrink. So the listing compares the mean position over the last hundred steps with the mean position two hundred steps earlier. If they almost coincide, the point is marking time. Slowly sliding along the floor of a ravine does not look like this: its mean position moves.

The second listing is the ravine (x² + κy²)/2 from the widget above: ordinary descent with the best constant step against Polyak’s ball, and what happens with a step greater than 2/L.

What it costs

A step costs one gradient evaluation: for a formula with d variables usually O(d) operations, for a neural network roughly two or three passes over it. Memory is O(d): the point, the gradient and, for the ball, the velocity.

The number of steps depends on the function. For a smooth function with one bowl-shaped minimum (strongly convex) and a constant step, the error decreases geometrically with a ratio of about 1 − 2/κ, and reaching a precision ε takes O(κ·log(1/ε)) steps. Polyak’s ball on a quadratic function and Nesterov’s method on any smooth strongly convex one give O(√κ·log(1/ε)). If the function is convex but the bowl degenerates, the guarantee is weaker: the error decreases as 1/k for ordinary descent and as 1/k² for Nesterov’s method. For non-convex functions only one thing is guaranteed, that descent will reach a point where the gradient is small; whether that will be a bottom, a saddle or a pit halfway down, the theory does not promise.

Hence the main practical advice, which is not about the step but about the data: make the ravine rounder. In problems like regression the columns are shifted to zero and their scales evened out for this. The first problem of the article shows how much this changes κ on real data: for a line through the mean annual temperatures of Almaty, from 2.6·10¹⁰ to 602.

Noise instead of precision

A neural network’s loss function is a sum of misses over millions of examples, and an honest gradient requires going over all of them. Instead, a random batch of a few dozen examples is taken and the gradient is computed over it: it is inaccurate but thousands of times cheaper. This is stochastic gradient descent. Its roots are in the 1951 work of Herbert Robbins and Sutton Monro on stochastic approximation: it proves that noise does not stop you reaching the goal if the step is decreased so that the sum of the steps is infinite and the sum of their squares finite.

Noise can be useful in itself: it knocks the point off a plateau and pushes it out of shallow pits, just like the huge step in the scene, only more gently. The method popular today, Adam (Diederik Kingma and Jimmy Ba, 2014), adds one more idea to momentum: choose the step length separately for each variable, by how strongly and how often its gradient has changed. A ravine stretched along one of the axes is almost levelled out by such a method.

History

The method was proposed by Augustin-Louis Cauchy in a three-page note of 1847 in the “Comptes rendus” of the Academy of Sciences. Cauchy needed to compute the orbits of celestial bodies: six unknown orbital elements from a system of equations of motion. Cauchy reduced the system to finding the minimum of a sum of squared residuals and proposed going towards it in small steps against the gradient. The continuation with details promised in the note never came out.

Momentum was added by Boris Polyak in 1964, and acceleration with a guarantee was found by Yurii Nesterov in 1983. And when in the 1980s people learned to compute a neural network’s gradient quickly by backpropagation, descent became the main way of training models: today neural networks are trained almost exclusively by its variants.

Problems

In the first problem you have to say in advance, without taking a single step, which step is safe: the ravine is rotated, and the limit from “Mathematics, the Queen of the Sciences” does not apply. In the second, descent runs over a function with kinks, a sum of distances to places, and the step chooses itself.

The longest safe step A pick-up point

Data: Relief of the mountains over Almaty: NASA SRTM v3 (1″) via Open Topo Data, public domain Almaty places: OpenStreetMap and Nominatim, ODbL Almaty weather: the ERA5 reanalysis (Copernicus) via Open-Meteo, CC BY 4.0