Loss surfaces and gradient descent

Before this lesson: partial derivatives and the gradient (calculus 02); vectors and dot products (linear algebra 01). Eigenvalues appear only as a preview. Time: about 55 minutes, plus the notebook.

Why this matters #

Every optimizer used to train neural networks is gradient descent with extra parts: SGD uses a noisy gradient from a minibatch (lesson 05), momentum adds memory of past steps (lesson 04), and Adam rescales each coordinate (lesson 06). So the basic behaviour of gradient descent explains a lot of what you see in practice. A loss that explodes in the first steps usually means a learning rate above the stability limit. A loss that falls quickly and then crawls usually means a narrow valley (an ill-conditioned problem, lesson 02). Standardizing input features before training linear models works because it makes the valley rounder.

1. Predict first #

2. The picture #

Loss surfaces and level sets #

An optimization problem asks for a point $x^\star$ that makes a function $f: \R^d \to \R$ as small as possible:

$f$ is the objective. In machine learning it is the loss $L(\theta)$, the average error of a model with parameters $\theta$ on the training data. We write $f$ and $x$ in general statements and $L$ and $\theta$ when we talk about training.

With two variables you can draw $f$ as a surface over the plane: the loss surface. A contour plot shows the same information from above. Each curve is a level set $\{x : f(x) = c\}$, the set of points with the same height $c$. Where the curves are close together the surface is steep; where they are far apart it is flat.

• A minimizer $x^\star$ satisfies $f(x^\star) \le f(x)$ for all $x$. The value $f^\star = f(x^\star)$ is the optimal value. This is a global minimum.
• A local minimum is a point that is lowest only in a small neighbourhood. A function can have several local minima with different heights.
• A stationary point (or critical point) has $\nabla f(x) = 0$. Minima, maxima and saddle points are stationary points. At a saddle point the surface goes up in some directions and down in others.

The gradient points uphill, across the contours #

The gradient $\nabla f(x)$ is the vector of partial derivatives. For a unit vector $u$, the directional derivative $\nabla f(x)^\top u$ is the slope of $f$ when you walk from $x$ in the direction $u$. Since $\nabla f(x)^\top u = \|\nabla f(x)\| \cos\theta$, where $\theta$ is the angle between $u$ and $\nabla f(x)$:

• the slope is largest, $\|\nabla f(x)\|$, when $u$ points along $\nabla f(x)$ ($\theta = 0$);
• the slope is most negative, $-\|\nabla f(x)\|$, when $u$ points along $-\nabla f(x)$;
• the slope is zero when $u$ is perpendicular to $\nabla f(x)$. These are the directions along the level set, where the height does not change.

Gradient descent #

Gradient descent (GD) starts at a point $x_0$ and repeats

The positive number $\eta$ is the step size, called the learning rate in machine learning. Each step moves downhill, perpendicular to the current contour line. The step is long where the surface is steep (large gradient) and short where it is flat.

One variable: four regimes #

Take the simplest loss, a parabola $f(x) = \tfrac{\lambda}{2}x^2$ with curvature $\lambda > 0$. Then $f'(x) = \lambda x$ and one step gives

and so, after $k$ steps,

Everything depends on the single number $1 - \eta\lambda$:

So gradient descent on this parabola converges exactly when $0 < \eta < 2/\lambda$. A step that is too large does not fail at once. First it overshoots and oscillates (for $1/\lambda < \eta < 2/\lambda$); it diverges only after the overshoot becomes larger than the distance it started from.

Two variables: one step size for two curvatures #

Now take $f(x) = \tfrac12(x_1^2 + 10\,x_2^2)$. Its contours are ellipses, ten times steeper in the $x_2$ direction: a narrow valley along the $x_1$ axis. The gradient is $\nabla f(x) = (x_1,\ 10\,x_2)$, so gradient descent works on each coordinate separately:

These are two copies of the 1D problem, with $\lambda = 1$ and $\lambda = 10$, but one $\eta$ must serve both:

• The steep coordinate $x_2$ needs $\eta < 2/10 = 0.2$, or it diverges.
• With $\eta$ just below $0.2$, the factor for $x_2$ is close to $-1$: $x_2$ jumps from one wall of the valley to the other. This is the zig-zag.
• With such a small $\eta$, the factor for $x_1$ is $1 - \eta \ge 0.8$: progress along the valley floor is slow. This is the crawl.

The best compromise makes the two factors equally bad: $|1 - \eta| = |1 - 10\eta|$ gives $\eta = 2/11$, and then both factors have size $9/11 \approx 0.82$. Even with this best fixed step, the error shrinks by only $9/11$ per step: about 69 steps to reduce it by a factor $10^6$, against 1 step when both curvatures are equal (a round bowl with $\eta = 1/\lambda$).

The ratio of the largest to the smallest curvature, here $10/1$, is the condition number $\kappa$. Lesson 02 defines it for any function and shows that it controls the speed of gradient descent.

A first nonconvex surface #

Real losses are not quadratic. The double well $f(x) = (x_1^2 - 1)^2 + \tfrac12 x_2^2$ has two minimizers, $(-1, 0)$ and $(1, 0)$, and a saddle point at $(0, 0)$ between them. Gradient descent goes downhill from where it starts, so the start point decides which minimum it reaches. A start with $x_1 > 0$ ends at $(1, 0)$; a start with $x_1 < 0$ ends at $(-1, 0)$. A start exactly on the line $x_1 = 0$ slides down to the saddle point and stops there, because the gradient has no $x_1$-component on that line. This start is a very special case; any tiny change of $x_1$ moves the path to one of the minima. Lesson 11 studies saddle points in high dimension.

Try it in the trajectory figure above: choose double well and drag the start across the line $x_1 = 0$.

3. Definitions and main results #

Why $-\nabla f$ is the steepest descent direction #

Theorem 1 (steepest ascent and descent). Let $f$ be differentiable at $x$ with $\nabla f(x) \ne 0$. Among all unit vectors $u$, the directional derivative $\nabla f(x)^\top u$ is largest for $u = \nabla f(x)/\|\nabla f(x)\|$ and smallest for $u = -\nabla f(x)/\|\nabla f(x)\|$.

Proof.

1. The directional derivative in direction $u$ is $\frac{d}{dt} f(x + tu)\big|_{t=0} = \nabla f(x)^\top u$ (chain rule, calculus 02).
2. By the Cauchy–Schwarz inequality, $|\nabla f(x)^\top u| \le \|\nabla f(x)\|\,\|u\| = \|\nabla f(x)\|$.
3. So $-\|\nabla f(x)\| \le \nabla f(x)^\top u \le \|\nabla f(x)\|$. Equality on the right holds only when $u$ is a positive multiple of $\nabla f(x)$, and on the left only when $u$ is a negative multiple. $\blacksquare$

Theorem 2 (the gradient is perpendicular to the level set). Let $\gamma(t)$ be a differentiable curve inside the level set $\{f = c\}$, with $\gamma(0) = x$. Then $\nabla f(x)^\top \gamma'(0) = 0$.

Proof. $f(\gamma(t)) = c$ for all $t$. Differentiate with the chain rule: $0 = \frac{d}{dt} f(\gamma(t)) = \nabla f(\gamma(t))^\top \gamma'(t)$. Set $t = 0$. $\blacksquare$

So the gradient is perpendicular to every tangent direction of the level set. In 2D the level set is a curve, and the gradient is the normal of that curve.

A small enough step always goes down #

Theorem 3 (descent for small steps). Let $f$ be differentiable at $x$ with $\nabla f(x) \ne 0$. Then $f(x - \eta\nabla f(x)) < f(x)$ for every small enough $\eta > 0$.

Proof.

1. Write $g = \nabla f(x)$ and $\varphi(\eta) = f(x - \eta g)$. Then $\varphi(0) = f(x)$.
2. By the chain rule, $\varphi'(0) = \nabla f(x)^\top(-g) = -\|g\|^2 < 0$.
3. A function with negative derivative at 0 is smaller than $\varphi(0)$ just to the right of 0: by the definition of the derivative, $\big(\varphi(\eta) - \varphi(0)\big)/\eta \to -\|g\|^2$, so the quotient is negative for all small $\eta > 0$. $\blacksquare$

The theorem says "small enough", but not how small. For a quadratic we can compute the exact range; lesson 02 gives a range for any $L$-smooth function: $0 < \eta < 2/L$.

Fixed points are stationary points. If gradient descent converges, $x_k \to \bar x$, and $\nabla f$ is continuous, then taking the limit in $x_{k+1} = x_k - \eta\nabla f(x_k)$ gives $\bar x = \bar x - \eta\nabla f(\bar x)$, so $\nabla f(\bar x) = 0$. The limit is a stationary point: a minimum, or (rarely) a saddle point as in the double well.

Gradient descent on quadratics #

Theorem 4 (1D quadratic). For $f(x) = \tfrac{\lambda}{2}x^2$ with $\lambda > 0$, gradient descent gives $x_k = (1 - \eta\lambda)^k x_0$. It converges to 0 from every start if and only if $0 < \eta < 2/\lambda$.

Proof. The formula follows from $x_{k+1} = (1 - \eta\lambda)x_k$ by induction. A geometric sequence $r^k x_0$ with $x_0 \ne 0$ goes to 0 if and only if $|r| < 1$. Here $|1 - \eta\lambda| < 1$ means $-1 < 1 - \eta\lambda < 1$, that is, $0 < \eta\lambda < 2$. $\blacksquare$

Theorem 5 (separable 2D quadratic). Let $f(x) = \tfrac12(\lambda_1 x_1^2 + \lambda_2 x_2^2)$ with $0 < \lambda_1 \le \lambda_2$. Gradient descent converges from every start if and only if $\eta < 2/\lambda_2$. The fixed step that makes the slower of the two factors $|1 - \eta\lambda_i|$ as small as possible is $\eta = 2/(\lambda_1 + \lambda_2)$, and then both factors have size $(\lambda_2 - \lambda_1)/(\lambda_2 + \lambda_1)$.

Proof.

1. The gradient is $(\lambda_1 x_1, \lambda_2 x_2)$, so each coordinate follows Theorem 4 with its own $\lambda_i$. Both converge if and only if $\eta < 2/\lambda_1$ and $\eta < 2/\lambda_2$. Since $\lambda_2 \ge \lambda_1$, the second condition is the binding one.
2. The speed is set by the slower coordinate, $\rho(\eta) = \max\big(|1 - \eta\lambda_1|,\ |1 - \eta\lambda_2|\big)$.
3. For $\eta < 1/\lambda_2$ both factors fall as $\eta$ grows, and for $\eta > 1/\lambda_1$ both rise. So the best $\eta$ lies in $[1/\lambda_2, 1/\lambda_1]$. There $1 - \eta\lambda_1$ falls and $\eta\lambda_2 - 1$ rises, and their maximum is smallest where they are equal: $1 - \eta\lambda_1 = \eta\lambda_2 - 1$, so $\eta = 2/(\lambda_1 + \lambda_2)$.
4. Insert this $\eta$: $1 - 2\lambda_1/(\lambda_1 + \lambda_2) = (\lambda_2 - \lambda_1)/(\lambda_2 + \lambda_1)$. $\blacksquare$

For $\lambda_1 = 1$, $\lambda_2 = 10$: the limit is $\eta < 0.2$, the best step is $2/11$, and the factor is $9/11$. Lesson 03 proves the same for any quadratic $\tfrac12 x^\top A x - b^\top x$, with $\lambda_1, \lambda_2$ replaced by the smallest and largest eigenvalues of $A$.

Stopping rules #

Gradient descent never stops by itself. Common rules:

1. Gradient norm: stop when $\|\nabla f(x_k)\| \le \varepsilon$.
2. Small change: stop when $|f(x_{k}) - f(x_{k+1})| \le \varepsilon\,|f(x_k)|$ (relative change) or when $\|x_{k+1} - x_k\|$ is small.
3. Budget: stop after a maximum number of steps (or time). In deep learning this is the usual rule, combined with a validation loss.

Why the gradient norm is the most direct test. A minimizer of a differentiable function must have $\nabla f = 0$, so $\|\nabla f(x_k)\|$ measures how far $x_k$ is from satisfying this condition. It does not need $f^\star$, which you usually do not know. A small change of $f$ or $x$ only says that progress is slow, and progress can be slow far from the minimum: with a tiny step, or on a flat plateau, or along the floor of a narrow valley. For example, on $f(x) = \tfrac12 x^2$ with $\eta = 0.001$ from $x_0 = 10$, one step changes $f$ by only 0.2%. A relative-change rule with $\varepsilon = 10^{-2}$ stops at once, while the gradient norm is 10.

The gradient norm has a limit too: on a very flat function it can be small far from the minimizer. Lesson 02 shows that for a $\mu$-strongly convex $f$, $f(x) - f^\star \le \|\nabla f(x)\|^2/(2\mu)$, so a small gradient does guarantee a small gap when $\mu$ is not tiny. Always also keep a maximum number of steps.

Answers to section 1. (1) $\eta\lambda = 0.4 \cdot 4 = 1.6$, so $x_k = (-0.6)^k$: it converges, and the sign changes every step. (2) No. At $\eta = 0.2$ the factor for $x_2$ is $1 - 10 \cdot 0.2 = -1$: $x_2$ jumps between $1$ and $-1$ forever, and gradient descent no longer converges. (3) Neither. The gradient is perpendicular to the contour line. In a narrow valley this means it points mostly across the valley, up the wall, and not toward the minimum. The descent direction $-\nabla f$ points mostly down to the valley floor.

4. Worked example #

Run gradient descent by hand on $f(x) = \tfrac12(x_1^2 + 10\,x_2^2)$ from $x_0 = (2, 1)$. The minimizer is $x^\star = (0, 0)$ with $f^\star = 0$, and $f(x_0) = \tfrac12(4 + 10) = 7$. The gradient is $\nabla f(x) = (x_1,\ 10\,x_2)$, so $\nabla f(x_0) = (2, 10)$ and $\|\nabla f(x_0)\| = \sqrt{104} \approx 10.2$.

(a) Step $\eta = 0.1$ ($= 1/L$).

The steep coordinate is solved in one step ($1 - 10 \cdot 0.1 = 0$). The flat coordinate shrinks by $0.9$ per step: $x_2 = (1.62, 0)$, $f(x_2) = 1.3122$. After that, only the slow factor $0.9$ is left.

(b) Step $\eta = 0.18$ (close to the limit 0.2).

Then $f(x_1) = \tfrac12(1.64^2 + 10 \cdot 0.8^2) = \tfrac12(2.6896 + 6.4) = 4.5448$. Next, $\nabla f(x_1) = (1.64, -8)$ and

$x_2$ changed sign twice (factor $1 - 1.8 = -0.8$): the zig-zag. $x_1$ shrinks by $0.82$ per step.

(c) Step $\eta = 0.21$ (above the limit). $x_1 = (2 - 0.42,\ 1 - 2.1) = (1.58,\ -1.1)$ and $f(x_1) = \tfrac12(2.4964 + 12.1) = 7.2982 > 7$. The loss went up. The factor for $x_2$ is $1 - 2.1 = -1.1$, so $|x_2|$ grows by 10% each step: $x_2 = (1.2482,\ 1.21)$, $f(x_2) = 8.0995$.

How many steps? The error shrinks at least by the slower factor $\rho = \max(|1 - \eta|, |1 - 10\eta|)$ per step, so a reduction by $10^6$ needs about $\ln(10^6)/\ln(1/\rho)$ steps:

The code in section 5 counts the steps exactly (it gets 1,364 for $\eta = 0.01$, because it starts from $\|x_0\| = \sqrt5$, not from a point on the slow axis).

5. In code #

NumPy checks the worked example and counts the steps.

import numpy as np

lam = np.array([1.0, 10.0])                 # f(x) = 1/2 (x1^2 + 10 x2^2), minimizer x* = 0
f = lambda x: 0.5 * np.sum(lam * x**2)
grad = lambda x: lam * x

def gd(x0, eta, steps):
    xs = [np.array(x0, dtype=float)]
    for _ in range(steps):
        xs.append(xs[-1] - eta * grad(xs[-1]))
    return np.array(xs)

x0 = [2.0, 1.0]
for eta in (0.1, 0.18, 0.21):
    xs = gd(x0, eta, 2)
    print(eta, xs[1].round(4), xs[2].round(4), [round(float(f(x)), 4) for x in xs])
# 0.1 [1.8 0. ] [1.62 0.  ] [7.0, 1.62, 1.3122]
# 0.18 [ 1.64 -0.8 ] [1.3448 0.64  ] [7.0, 4.5448, 2.9522]
# 0.21 [ 1.58 -1.1 ] [1.2482 1.21  ] [7.0, 7.2982, 8.0995]

def steps_to(eta, tol=1e-6, kmax=5000):
    """Steps until ||x_k|| <= tol * ||x_0||, or None if that never happens."""
    x = np.array(x0); n0 = np.linalg.norm(x)
    for k in range(kmax):
        if np.linalg.norm(x) <= tol * n0:
            return k
        if np.linalg.norm(x) > 1e6:          # diverged
            return None
        x = x - eta * grad(x)
    return None

print([steps_to(eta) for eta in (0.01, 0.1, 0.18, 2 / 11, 0.2, 0.21)])
# [1364, 131, 70, 69, None, None]
print(2 / lam.max(), abs(1 - 0.2 * lam))   # stability limit and the two factors at eta = 0.2
# 0.2 [0.8 1. ]

The notebook adds a gradient check by finite differences, the four 1D regimes, a numerical search for the largest stable step, and gradient descent on a real regression dataset.

6. Where it shows up #

• Training neural networks. Every training step of SGD, momentum or Adam is a gradient step on the loss, with a gradient from backpropagation (calculus 06). The learning rate is usually the most important hyperparameter to tune.
• Divergence and loss spikes. Near a minimum, the loss behaves like a quadratic whose largest curvature is the largest Hessian eigenvalue. A step above about $2/\lambda_{\max}$ makes the loss oscillate and grow, as in the worked example (c). Lessons 07 and 11 connect this to warmup and to the "edge of stability".
• Linear and logistic regression. Their losses are convex, so gradient descent reaches the global minimum. Unscaled features (one in metres, one in millimetres) make a narrow valley and slow gradient descent down; standardizing them is the standard fix (lesson 02).
• ML for science. Gradient-based optimization fits physical models to data: the parameters of an ODE model, the initial state in weather data assimilation (4D-Var), or the material model in full-waveform inversion. The gradient comes from automatic differentiation or an adjoint method, and the same step-size questions appear.

7. Common mistakes #

• Drawing the gradient along the valley. The gradient is perpendicular to the contour line. In a narrow valley it points mostly across the valley, not toward the minimum.
• "A larger learning rate is always faster." Only up to $2/\lambda_{\max}$. Close to the limit the steep direction zig-zags, and above it gradient descent diverges.
• "$\eta\lambda > 1$ means divergence." Between 1 and 2 the iterates overshoot but still converge. Divergence starts at $\eta\lambda > 2$.
• "The loss went up once, so the step is too large." On a quadratic with a stable step, $f$ decreases every step. But with noisy minibatch gradients (lesson 05) or on nonconvex losses, single increases are normal; look at the trend.
• Stopping on a small change of the loss alone. With a small step or on a plateau, the change is small far from the minimum. Check the gradient norm and keep a step budget.
• "Gradient descent finds the global minimum." Only for convex functions (lesson 02). On a nonconvex loss it finds a stationary point that depends on the start.

8. Check yourself #

1. $f(x) = 3x^2$. Which step sizes make gradient descent converge from every start? Which step reaches the minimum in one step?
   Answer

   $f'(x) = 6x$, so $\lambda = 6$ and $x_{k+1} = (1 - 6\eta)x_k$. It converges for $0 < \eta < 2/6 = 1/3$. The step $\eta = 1/6$ gives factor 0: one step.
2. Compute one gradient-descent step with $\eta = 0.5$ on $f(x) = x_1^2 + x_1 x_2 + x_2^2$ from $(1, 1)$.
   Answer

   $\nabla f = (2x_1 + x_2,\ x_1 + 2x_2) = (3, 3)$ at $(1, 1)$. The step gives $(1, 1) - 0.5\,(3, 3) = (-0.5,\ -0.5)$. The value goes from $3$ to $0.75$.
3. At the point $(1, 1)$ of the level set $x_1^2 + 4x_2^2 = 5$, give a vector perpendicular to the level set and a vector tangent to it.
   Answer

   With $f = x_1^2 + 4x_2^2$, $\nabla f(1, 1) = (2, 8)$ is perpendicular. A tangent vector is perpendicular to it, for example $(8, -2)$ or $(4, -1)$: $(2, 8)\cdot(4, -1) = 0$.
4. On $f(x) = \tfrac12(x_1^2 + 100\,x_2^2)$, what is the largest stable step, the best fixed step, and the factor per step with the best step?
   Answer

   Largest stable: $\eta < 2/100 = 0.02$. Best: $\eta = 2/101 \approx 0.0198$. Factor: $99/101 \approx 0.980$, so about $\ln(10^6)/\ln(101/99) \approx 691$ steps for a reduction by $10^6$.
5. Gradient descent converges to a point $\bar x$. Must $\bar x$ be a minimizer?
   Answer

   No. If $\nabla f$ is continuous, $\bar x$ is a stationary point, $\nabla f(\bar x) = 0$. It can be a saddle point, for example $(0, 0)$ on the double well when the start has $x_1 = 0$. For a convex $f$, every stationary point is a global minimizer (lesson 02).

9. Interview angle #

• "What happens if the learning rate is too large?" — The loss oscillates and then explodes. On a quadratic, the error along the steepest direction is multiplied by $1 - \eta\lambda_{\max}$, which is below $-1$ once $\eta > 2/\lambda_{\max}$.
• "Why does gradient descent zig-zag?" — The gradient is perpendicular to the contours, so in a narrow valley it points across the valley. The step that is safe for the steep wall is too small for the flat floor.
• "Why do we standardize features?" — It makes the contours of a linear model's loss rounder (a smaller condition number), so one learning rate fits all directions (lesson 02).
• "Does gradient descent find the global minimum?" — For a convex function, yes (with a suitable step). For a neural network it finds a stationary point, usually a good local minimum; the start and the step size matter.

10. Practice #

The notebook for this lesson runs in Colab. It has four exercises, each with a check cell and a folded solution:

1. Write the gradients of $\tfrac12 x^\top A x - b^\top x$ and of least squares, and check them with a finite-difference gradient check.
2. Write gradient descent that returns the whole path, and test the four 1D regimes.
3. Find the largest stable step numerically by bisection and compare it with $2/\lambda_{\max}$.
4. Gradient descent with a gradient-norm stopping rule on the diabetes regression data, with raw and with standardized features: count the steps.

11. Watch and read #

• Gabriel Goh, Why Momentum Really Works, Distill (2017): the first section, "First Steps: Gradient Descent", has an interactive version of the narrow valley.
• Stanford CS231n, Optimization: stochastic gradient descent: loss surfaces of linear classifiers and the gradient check.
• Boyd and Vandenberghe, Convex Optimization, free PDF: §9.2 (descent methods) and §9.3 (gradient descent, with the stopping rule on $\|\nabla f\|$).