Nonconvex landscapes: saddle points, sharpness and the edge of stability

Before this lesson: lesson 02 (the descent lemma), lesson 03 (the stability limit $\eta < 2/\lambda_{\max}$), lesson 05 (SGD noise), lesson 07 (warmup) and lesson 10 (Newton's method). From other tracks: power iteration (linear algebra 05), Lanczos (linear algebra 14) and Hessian-vector products by automatic differentiation (calculus 08). Time: about 85 minutes, plus the notebook. This is a difficult lesson: take the figures one at a time.

---

Why this matters #

Lessons 01–06 mostly assumed a convex loss with one curvature bound $L$. A neural-network loss is nonconvex, and its curvature changes by orders of magnitude during a run. This lesson answers three practical questions. What can gradient descent promise? A small gradient, not a minimum. Do saddle points stop training? Rarely, but they can slow it down. Which learning rates are stable, and why does training work even when the theory says it should diverge? The answer is the sharpness and the edge of stability. Sharpness is also measured in practice: to debug unstable runs of large models, to explain warmup (lesson 07), and in methods such as SAM that look for flat minima.

1. Predict first #

2. The picture #

Stationary points and the Hessian. A stationary point (or critical point) of $f$ is a point $x$ with $\nabla f(x) = 0$. Near it, $f(x + h) \approx f(x) + \tfrac12 h^\top H h$, where $H = \nabla^2 f(x)$ is the Hessian, the matrix of second derivatives. Write $H = Q\Lambda Q^\top$ with orthonormal eigenvectors $q_1, \dots, q_d$ and eigenvalues $\lambda_1, \dots, \lambda_d$. Along $q_i$ the function curves up if $\lambda_i > 0$ and down if $\lambda_i < 0$. So the signs of the eigenvalues give the shape:

Eigenvalues of $H$	Shape near the point	Name
all $> 0$	bowl	strict local minimum
all $< 0$	upside-down bowl	strict local maximum
some $> 0$, some $< 0$	up in some directions, down in others	saddle point
some $= 0$	the quadratic term is flat in some direction	degenerate: the test cannot decide

Interactive figure: loading…

Gradient descent near a saddle. Take the simplest saddle, $f(x, y) = \tfrac12(x^2 - y^2)$. The gradient is $(x, -y)$, so gradient descent (GD) with step $\eta$ gives

The $x$-coordinate shrinks: $x = 0$ is the stable direction. The $y$-coordinate grows by the factor $1 + \eta$ at every step, however small it is at the start. Only a start with $y_0 = 0$ exactly converges to the saddle. A start with $y_0 = \delta \ne 0$ first moves toward the saddle along $x$, then leaves along $y$.

Interactive figure: loading…

3. Definitions and main results #

Stationary points and the second-order test #

Definitions. Let $f: \mathbb R^d \to \mathbb R$ be twice continuously differentiable. A stationary point has $\nabla f(x) = 0$. A strict saddle is a stationary point with $\lambda_{\min}(\nabla^2 f(x)) < 0$. (With this definition, a strict local maximum is also a strict saddle; the important property is that some direction goes down.)

Theorem (second-order conditions).

1. Necessary. If $x^\star$ is a local minimum, then $\nabla f(x^\star) = 0$ and $\nabla^2 f(x^\star) \succeq 0$.
2. Sufficient. If $\nabla f(x^\star) = 0$ and $\nabla^2 f(x^\star) \succ 0$, then $x^\star$ is a strict local minimum.

Proof. By Taylor's theorem, $f(x^\star + h) = f(x^\star) + \nabla f(x^\star)^\top h + \tfrac12 h^\top H h + o(\|h\|^2)$ with $H = \nabla^2 f(x^\star)$.

1. If $\nabla f(x^\star) \ne 0$, then $h = -t\nabla f(x^\star)$ with small $t > 0$ lowers $f$. So $\nabla f(x^\star) = 0$. If $q^\top H q < 0$ for some unit $q$, then $f(x^\star + tq) = f(x^\star) + \tfrac{t^2}{2} q^\top H q + o(t^2) < f(x^\star)$ for small $t$. So $H \succeq 0$.
2. With $\nabla f(x^\star) = 0$: $f(x^\star + h) - f(x^\star) \ge \tfrac12\lambda_{\min}(H)\|h\|^2 + o(\|h\|^2)$, and $\lambda_{\min}(H) > 0$. For small $h \ne 0$ the first term wins, so the difference is positive. $\blacksquare$

Degenerate cases. When $H \succeq 0$ is singular, the test does not decide. In one variable at $x = 0$: $x^4$ has a minimum, $-x^4$ a maximum and $x^3$ neither, and all three have $f''(0) = 0$. The monkey saddle $x^3 - 3xy^2$ has $H = 0$ and three directions that go down.

What gradient descent guarantees on a nonconvex function #

Theorem (gradient descent on an $L$-smooth function). Let $f$ be $L$-smooth and bounded below by $f_{\inf}$. Gradient descent with step $1/L$, $x_{i+1} = x_i - \tfrac1L\nabla f(x_i)$, satisfies

Proof.

1. The descent lemma (lesson 02) needs no convexity: $f(y) \le f(x) + \nabla f(x)^\top(y - x) + \tfrac L2\|y - x\|^2$.
2. Put $y = x_{i+1} = x_i - \nabla f(x_i)/L$. Then $y - x_i = -\nabla f(x_i)/L$, and $f(x_{i+1}) \le f(x_i) - \tfrac1L\|\nabla f(x_i)\|^2 + \tfrac{1}{2L}\|\nabla f(x_i)\|^2 = f(x_i) - \tfrac{1}{2L}\|\nabla f(x_i)\|^2$.
3. Sum over $i = 0, \dots, k-1$. The values of $f$ cancel in pairs (the sum telescopes): $\tfrac{1}{2L}\sum_{i<k}\|\nabla f(x_i)\|^2 \le f(x_0) - f(x_k) \le f(x_0) - f_{\inf}$.
4. The smallest of $k$ numbers is at most their mean. $\blacksquare$

What it says and what it does not say.

• It is about the gradient, at the best iterate so far. It says nothing about $f(x_k) - f^\star$, and nothing about whether the point is a minimum or a saddle.
• $O(1/k)$ for $\|\nabla f\|^2$ means $O(1/\varepsilon^2)$ iterations to reach $\|\nabla f\| \le \varepsilon$.
• The loss decreases at every step (step 2), so GD can never climb out of a basin.

Why strict saddles rarely trap gradient descent #

Linear picture. Near a stationary point $x_s$ with Hessian $H$, $\nabla f(x) \approx H(x - x_s)$. So the error $e_k = x_k - x_s$ follows $e_{k+1} \approx (I - \eta H)\,e_k$. In the eigenbasis of $H$, the component along $q_i$ is multiplied by $1 - \eta\lambda_i$ at every step.

• For $\lambda_i > 0$ (and $\eta < 2/\lambda_i$): $|1 - \eta\lambda_i| < 1$, the component shrinks.
• For $\lambda_i < 0$: $1 - \eta\lambda_i = 1 + \eta|\lambda_i| > 1$, the component grows.

Only starts with exactly zero component along every direction with $\lambda_i < 0$ converge to the saddle. In the linear picture, this is a subspace of lower dimension; for the real $f$ it is a curved set of lower dimension, the stable manifold. A random start lands on it with probability zero.

Escape time. A start at distance $\delta$ from the stable manifold, along the direction with $\lambda_{\min} < 0$, reaches distance $r$ after about

For $\eta = 0.1$, $\lambda_{\min} = -1$, $\delta = 10^{-6}$ and $r = 1$: $k \approx 13.82/0.0953 = 145$ steps.

What has been proved.

• Lee, Simchowitz, Jordan and Recht (2016). Let $f$ be twice continuously differentiable with an $L$-Lipschitz gradient, and let $\eta < 1/L$. Then the set of starting points from which GD converges to a strict saddle has measure zero. So from a random start, if GD converges and every saddle of $f$ is strict, the limit is a local minimum. The idea: the GD map $x \mapsto x - \eta\nabla f(x)$ has the invertible Jacobian $I - \eta\nabla^2 f$, and at a strict saddle it has an eigenvalue $1 + \eta|\lambda_{\min}| > 1$; the stable manifold theorem does the rest.
• The escape can be slow. Du et al. (2017) built functions where GD needs exponential time to pass a sequence of saddles. Plateaus near saddles can slow training even when they do not stop it.
• Noise helps. Jin, Ge, Netrapalli, Kakade and Jordan (2017): perturbed GD, which adds a small random kick when the gradient is small, reaches an approximate second-order stationary point ($\|\nabla f\| \le \varepsilon$ and $\lambda_{\min}(\nabla^2 f) \ge -\sqrt{\rho\varepsilon}$, for a $\rho$-Lipschitz Hessian) in $\tilde O\big(L(f(x_0) - f^\star)/\varepsilon^2\big)$ iterations, with only a polylogarithmic factor in the dimension. The noise of SGD plays a similar role.
• In high dimensions, saddles are the typical obstacle. Dauphin et al. (2014) argued, with theory from random functions and experiments on networks, that most stationary points with a high loss are saddles, not bad local minima, and that they are often surrounded by plateaus.

Newton's method is attracted to saddles. Newton's method (lesson 10) solves $\nabla f = 0$; it does not look at the sign of the curvature. On the quadratic model, $e_{k+1} = e_k - H^{-1}He_k = 0$: one step goes exactly to the stationary point, whether it is a minimum or a saddle. Saddle-free Newton (Dauphin et al.) uses $|H| = Q|\Lambda|Q^\top$ instead: along a direction with $\lambda_i < 0$ it moves away from the saddle.

Why the escape time depends on log(1/δ)

In the linear picture, the component along $q_{\min}$ after $k$ steps is $\delta(1 + \eta|\lambda_{\min}|)^k$. Set it equal to $r$ and take logarithms: $\log\delta + k\log(1 + \eta|\lambda_{\min}|) = \log r$. So the number of steps grows like $\log(1/\delta)$. A start $10^{6}$ times closer to the stable manifold costs only $\log(10^6)/\log(1 + \eta|\lambda_{\min}|)$ extra steps, not $10^6$ times more. With a small $|\lambda_{\min}|$ (a flat saddle) the factor $1/\log(1 + \eta|\lambda_{\min}|) \approx 1/(\eta|\lambda_{\min}|)$ is large: this is a plateau.

Measuring curvature: Hessian-vector products and power iteration #

For a model with $d = 10^8$ parameters, the Hessian has $10^{16}$ entries: it cannot be stored. But most methods only need its product with a vector.

Hessian-vector product (HVP). For a vector $v$,

that is, the derivative of $\nabla f(x + tv)$ with respect to $t$ at $t = 0$: the rate at which the gradient changes when you move along $v$. Three ways to compute it (calculus 08 has the details):

• reverse-over-reverse (Pearlmutter 1994): $Hv = \nabla_x\big(\nabla f(x)^\top v\big)$, a second backward pass through the gradient computation;
• forward-over-reverse: a forward-mode derivative (JVP) of the gradient function in direction $v$;
• central differences: $\big(\nabla f(x + hv) - \nabla f(x - hv)\big)/(2h)$, two gradients. The error is $O(h^2)$ from truncation plus $O(\varepsilon_{\text{mach}}/h)$ from rounding, so the best $h$ is about $\varepsilon_{\text{mach}}^{1/3}$ times the scale of $x$ (about $6\cdot10^{-6}$ in float64), with a relative error near $10^{-10}$. In float32 use automatic differentiation.

The two exact methods cost a small constant times one gradient (often quoted as 2 to 5) and $O(d)$ extra memory. No $d \times d$ matrix appears.

Power iteration (linear algebra 05) needs only HVPs: start from a random unit vector $v_0$ and repeat

where the estimate $r_k$ is the Rayleigh quotient. If $|\lambda_1| > |\lambda_2| \ge \dots$, the direction error shrinks like $|\lambda_2/\lambda_1|^k$, and for a symmetric $H$ the eigenvalue error like $|\lambda_2/\lambda_1|^{2k}$. Caution: it finds the eigenvalue of largest absolute value. Near a saddle with a large negative eigenvalue, that is the negative one. The Lanczos method (linear algebra 14; scipy.sparse.linalg.eigsh with which="LA") needs fewer products and can target the largest algebraic eigenvalue.

Interactive figure: loading…

Newton-CG uses the same products. The Newton step solves $Hp = -\nabla f$. Conjugate gradients needs one HVP per iteration, so a truncated Newton step also never forms $H$. It must stop when it meets negative curvature ($p^\top Hp \le 0$), because CG assumes $H \succ 0$.

Sharpness and the edge of stability #

Definition. The sharpness of the training loss $L$ at parameters $\theta$ is $S(\theta) = \lambda_{\max}\big(\nabla^2 L(\theta)\big)$, the largest eigenvalue of the Hessian.

Proposition (stability on a quadratic). Let $L(\theta) = \tfrac12(\theta - \theta^\star)^\top H(\theta - \theta^\star)$ with $H \succ 0$. Gradient descent with step $\eta$ converges from every start if and only if $S = \lambda_{\max}(H) < 2/\eta$.

Proof. The error $e_k = \theta_k - \theta^\star$ follows $e_{k+1} = (I - \eta H)e_k$ exactly. Along the eigenvector $q_i$ the component is multiplied by $1 - \eta\lambda_i$. All components shrink if and only if $|1 - \eta\lambda_i| < 1$ for every $i$, that is $0 < \eta\lambda_i < 2$. Since all $\lambda_i > 0$, this is $\eta\lambda_{\max} < 2$. If $\eta\lambda_{\max} > 2$, the component along $q_1$ grows by $|1 - \eta\lambda_{\max}| > 1$ per step, with a sign flip at each step: the iterate jumps from side to side with growing size. $\blacksquare$

The same argument, applied to the linear picture at a minimum $\theta^\star$ of a nonquadratic loss, says: gradient descent can stay at a minimum only if its sharpness is at most $2/\eta$. If $S(\theta^\star) > 2/\eta$, the minimum is unstable: tiny errors along $q_1$ grow. So the step size decides which minima GD can converge to. (The worked example in section 4 shows this on a function with a whole curve of minima.)

What real training does (Cohen et al. 2021). Cohen, Kaur, Li, Kolter and Talwalkar trained many networks with full-batch GD and measured the sharpness. They observed, across architectures and tasks:

1. Progressive sharpening: while $S < 2/\eta$, the sharpness tends to rise during training.
2. Edge of stability: once $S$ reaches $2/\eta$, it stops rising and hovers at or just above $2/\eta$.
3. The loss goes up and down over short times, but keeps decreasing over long times.

The classical theory of lesson 02 needs $\eta < 2/L$ for a monotone decrease. Here the local curvature sits at the limit and training still makes progress. The step size does not need to adapt to the curvature: the curvature adapts to the step size.

Interactive figure: loading…

Adam at the edge of stability. For Adam, the relevant quantity is the preconditioned sharpness $\lambda_{\max}(P^{-1}H)$, where $P$ is Adam's diagonal preconditioner. Cohen et al. (2022) observed that it rises to and hovers near $38/\eta$ for $\beta_1 = 0.9$, the stability limit of momentum, $(2 + 2\beta_1)/((1 - \beta_1)\eta)$. With minibatch SGD the sharpness does not settle at one value, but larger steps and smaller batches still lead to flatter regions (Cohen et al. 2021, section 6).

Why the sharpness stops at 2/η: self-stabilization (informal)

Damian, Nichani and Lee (2023) gave a mechanism. When $S$ rises above $2/\eta$, the iterate starts to oscillate along the top eigenvector $q_1$ with growing size. The loss is not exactly quadratic, and its third derivative couples this oscillation to the gradient of $S$: the oscillation pushes $\theta$ in a direction that lowers the sharpness. When $S$ is back below $2/\eta$, the oscillation shrinks again. The net effect is that GD follows, approximately, the path of a constrained method that keeps $S \le 2/\eta$. This explains the hovering in the figure and the short up-and-down steps of the loss.

Sharp and flat minima #

A minimum is sharp if the loss rises fast around it (large Hessian eigenvalues) and flat if it rises slowly. An old idea is that flat minima generalize better: a flat minimum is less sensitive to the small differences between the training and the test data. The evidence is mixed, and how you measure sharpness matters.

• Keskar et al. (2017). Large-batch training tends to find sharper minimizers that generalize worse; small-batch training, with more gradient noise, finds flatter ones. (They measured sharpness as the largest loss increase in a small box around the minimum.)
• Dinh et al. (2017): sharpness depends on the parametrization. For a ReLU network, $\operatorname{relu}(\alpha z) = \alpha\operatorname{relu}(z)$ for $\alpha > 0$. So the change $(W_1, W_2) \mapsto (\alpha W_1, W_2/\alpha)$ does not change the function $x \mapsto W_2\operatorname{relu}(W_1x)$, but it changes the Hessian and can make a minimum arbitrarily sharp. The same function, so the same test error, with any sharpness you like. Sharpness alone cannot explain generalization.
• Li et al. (2018): filter normalization. Plots of the loss along random directions have the same scale problem. They rescale each filter of the random direction to the norm of the matching filter of the weights. With this normalization, flatter-looking minima went with better test error in their experiments, and skip connections gave smoother landscapes.
• Later studies found both positive and negative results. Large studies of generalization measures (Jiang et al. 2020) ranked some sharpness-based measures among the better ones, while Andriushchenko et al. (2023) found weak or even negative correlation for modern architectures and training setups. Treat "flat means better" as a useful heuristic, not a law.

Sharpness-aware minimization (SAM; Foret et al. 2021). SAM minimizes the worst loss in a ball of radius $\rho$ around the weights:

At a minimum $w^\star$ with Hessian $H \succeq 0$, the second-order Taylor expansion gives $L^{\text{SAM}}_\rho(w^\star) \approx L(w^\star) + \tfrac12\rho^2\lambda_{\max}(H)$: the SAM objective adds a penalty proportional to the sharpness.

The SAM update, step by step.

1. Linearize the inner problem: $L(w + \varepsilon) \approx L(w) + \varepsilon^\top g$ with $g = \nabla L(w)$.
2. Maximize the linear term over the ball: $\varepsilon^\top g \le \|\varepsilon\|\,\|g\| \le \rho\|g\|$ (Cauchy–Schwarz), with equality for $\hat\varepsilon = \rho\,g/\|g\|$. This is a step of length $\rho$ uphill.
3. Take the gradient at the perturbed point, $g_{\text{SAM}} = \nabla L(w + \hat\varepsilon)$, ignoring how $\hat\varepsilon$ depends on $w$.
4. Update with any optimizer, for example $w \leftarrow w - \eta\,g_{\text{SAM}}$.

Each step costs two forward-backward passes. A typical radius in image classification is $\rho = 0.05$.

Interactive figure: loading…

Answers to section 1. (1) Typically it rises (progressive sharpening) to about $2/\eta = 200$ and then hovers there (the edge of stability), while the loss goes up and down over short times but decreases over long times. (2) GD does not reach the saddle: $|y_k| = 10^{-6}\cdot 1.1^k$ grows, and $|y| \ge 1$ after 145 steps. (3) Yes: power iteration or Lanczos on Hessian-vector products. One iteration is one HVP, a small constant times the cost of one gradient, with $O(d)$ memory. Tens of iterations are usually enough for the top eigenvalue.

4. Worked example #

(a) A nonconvex function in two variables. $f(x, y) = \cos x + \tfrac14 y^2$. Then $\nabla f = (-\sin x,\ y/2)$ and $\nabla^2 f = \operatorname{diag}(-\cos x,\ 1/2)$. The Hessian norm is at most 1, so $L = 1$; and $f \ge -1$, so $f_{\inf} = -1$.

• Stationary points: $\sin x = 0$ and $y = 0$, so $x = m\pi$ for integers $m$.
• At $x = 0, \pm2\pi, \dots$: $\cos x = 1$, $\nabla^2 f = \operatorname{diag}(-1, 1/2)$: one negative and one positive eigenvalue, so strict saddles, with $f = 1$.
• At $x = \pm\pi, \pm3\pi, \dots$: $\cos x = -1$, $\nabla^2 f = \operatorname{diag}(1, 1/2) \succ 0$: strict local minima, with $f = -1$.

(b) The guarantee and the real run. GD with step $1/L = 1$ from $(0.5, 1)$, where $f = \cos 0.5 + 0.25 = 1.128$.

• The theorem: $\min_{i<k}\|\nabla f(x_i)\|^2 \le 2\cdot 1\cdot(1.128 + 1)/k = 4.26/k$. For $k = 8$: $0.53$.
• The run: $x_{k+1} = x_k + \sin x_k$ and $y_{k+1} = y_k/2$. So $x$ goes $0.5 \to 0.979 \to 1.810 \to 2.781 \to 3.134 \to 3.1416$, and $y$ halves at each step. Among the first 8 iterates the smallest squared gradient norm is $1.5\cdot10^{-5}$.

The bound is true but loose: it must hold for every $L$-smooth function, including bad ones.

(c) Near the saddle. From $(0, 1)$ exactly, $x$ stays 0 (since $\sin 0 = 0$), and GD converges to the saddle $(0, 0)$: this start is on the stable manifold. From $(10^{-8}, 0)$: near the saddle $x_{k+1} = x_k + \sin x_k \approx 2x_k$, the factor $1 + \eta|\lambda_{\min}| = 2$. So $|x|$ first exceeds 1 at step 27 ($\log_2 10^8 \approx 26.6$), and $x$ is within $10^{-6}$ of $\pi$ at step 31.

(d) Power iteration and the stability limit. $f(x) = \tfrac12 x^\top Ax$ with $A = \begin{pmatrix} 3 & 1 \\ 1 & 3\end{pmatrix}$ (eigenvalues 4 and 2), so $Hv = Av$. From $v_0 = (1, 0)$:

• $Av_0 = (3, 1)$, $v_1 = (3, 1)/\sqrt{10}$, and $r_1 = v_1^\top Av_1 = (3\cdot 10 + 1\cdot 6)/10 = 3.6$ (here $Av_1 = (10, 6)/\sqrt{10}$).
• Next $v_2 = (5, 3)/\sqrt{34}$ and $r_2 = 132/34 \approx 3.882$; then $3.969$, $3.992$, $3.998$.
• The errors $0.4,\ 0.118,\ 0.031,\ 0.0078,\ 0.0020$ shrink by $(\lambda_2/\lambda_1)^2 = 1/4$ per step.

The sharpness is 4, so GD needs $\eta < 2/4 = 0.5$. With $\eta = 0.45$ the factors are $|1 - 0.45\cdot 4| = 0.8$ and $|1 - 0.45\cdot 2| = 0.1$: convergence. With $\eta = 0.55$ the first factor is $|1 - 2.2| = 1.2$: from $(1, 0)$, after 50 steps $x \approx (4550, 4550)$, along the top eigenvector $(1, 1)/\sqrt2$.

(e) The step size selects the minimum. $f(u, v) = \tfrac12(uv - 1)^2$ is the smallest "two-layer linear network": the model is the product $uv$. Its minimizers form the hyperbola $uv = 1$.

• The gradient is $\nabla f = (uv - 1)\,(v, u)$. On the hyperbola, $uv - 1 = 0$, so the Hessian is $\nabla^2 f = \begin{pmatrix} v^2 & 1 \\ 1 & u^2\end{pmatrix}$. Its trace is $u^2 + v^2$ and its determinant is $u^2v^2 - 1 = 0$, so the eigenvalues are 0 (along the hyperbola) and $u^2 + v^2$.
• So the minimum $(1, 1)$ has sharpness 2, and $(10, 0.1)$ has sharpness $100.01$: the same function $uv = 1$, 50 times sharper. This is Dinh et al.'s rescaling $(u, v) \mapsto (\alpha u, v/\alpha)$ with two parameters.

Run GD from $(3, 0.1)$, where $f = 0.245$:

• $\eta = 0.05$ ($2/\eta = 40$): the loss decreases at every step, and GD stops at a nearby minimum with sharpness $9.19 < 40$.
• $\eta = 0.25$ ($2/\eta = 8$): the nearby minima have sharpness above 8, so GD cannot stay there. The loss rises for four steps ($0.245 \to 0.39 \to 0.63 \to 0.72 \to 0.77$), $v$ jumps between about $-0.09$ and $0.77$, and GD settles at the flatter minimum $(2.432, 0.411)$, with sharpness $6.08 < 8$.

GD can only stay where the sharpness is at most $2/\eta$. This is the rule behind the edge of stability.

5. In code #

NumPy checks the numbers of section 4. The notebook has these cells and the exercises.

import numpy as np

# (a) classify the stationary points of f(x, y) = cos x + y^2/4 by the Hessian
hess = lambda x, y: np.array([[-np.cos(x), 0.0], [0.0, 0.5]])
for x in (0.0, np.pi):
    print(round(x, 4), np.linalg.eigvalsh(hess(x, 0.0)))
# 0.0 [-1.   0.5]      -> strict saddle
# 3.1416 [0.5 1. ]     -> strict local minimum

# (b) escape from the saddle of f = (x^2 - y^2)/2: steps until |y| >= 1
eta, y, k = 0.1, 1e-6, 0
while abs(y) < 1:
    y, k = (1 + eta) * y, k + 1
print(k, round(np.log(1e6) / np.log(1 + eta), 2))   # 145 144.95

# (c) the nonconvex GD bound on f = cos x + y^2/4 (L = 1, f_inf = -1), step 1/L = 1
p = np.array([0.5, 1.0])
grad = lambda p: np.array([-np.sin(p[0]), p[1] / 2])
f0, sq = np.cos(0.5) + 0.25, []
for _ in range(8):
    sq.append(grad(p) @ grad(p))
    p = p - grad(p)
print(f"min |grad|^2 = {min(sq):.1e}, bound 2L(f0 - f_inf)/k = {2 * (f0 + 1) / 8:.2f}")
# min |grad|^2 = 1.5e-05, bound 2L(f0 - f_inf)/k = 0.53

# (d) sharpness of Rosenbrock at its minimizer: power iteration on finite-difference HVPs
def rosen_grad(z):
    x, y = z
    return np.array([-2 * (1 - x) - 400 * x * (y - x**2), 200 * (y - x**2)])
def hvp(z, v, h=1e-5):
    return (rosen_grad(z + h * v) - rosen_grad(z - h * v)) / (2 * h)
z, v = np.array([1.0, 1.0]), np.array([1.0, 0.0])
for _ in range(5):
    w = hvp(z, v)
    lam, v = v @ w, w / np.linalg.norm(w)
print(round(lam, 3), np.linalg.eigvalsh([[802.0, -400.0], [-400.0, 200.0]]).round(3).tolist())
# 1001.601 [0.399, 1001.601]

# (e) the step size selects the minimum: GD on f(u, v) = (uv - 1)^2 / 2 from (3, 0.1)
for lr in (0.05, 0.25):
    u, w = 3.0, 0.1
    for _ in range(2000):
        r = u * w - 1
        u, w = u - lr * r * w, w - lr * r * u
    print(f"lr {lr}: uv = {u * w:.6f}, sharpness u^2 + v^2 = {u * u + w * w:.2f}, 2/lr = {2 / lr:.0f}")
# lr 0.05: uv = 1.000000, sharpness u^2 + v^2 = 9.19, 2/lr = 40
# lr 0.25: uv = 1.000000, sharpness u^2 + v^2 = 6.08, 2/lr = 8

In PyTorch, an exact HVP is torch.func.jvp(torch.func.grad(loss_fn), (params,), (v,))[1] (forward-over-reverse), and scipy.sparse.linalg.eigsh on a LinearOperator built from it gives the top eigenvalue with Lanczos.

6. Where it shows up #

• Training large models. When a run becomes unstable, people estimate the sharpness on a fixed batch with Lanczos on HVPs and compare it with $2/\eta$ (for Adam: the preconditioned sharpness with $38/\eta$). Gilmer et al. (2022) linked training instabilities to the sharpness exceeding this limit, and found that good choices, such as the learning-rate warmup of lesson 07, help training avoid or leave regions of high curvature.
• Optimizer design. SAM and its variants are used in image classification; for example, SAM let vision transformers train well on ImageNet without large-scale pretraining (Chen, Hsieh and Gong 2022). Loss-landscape plots with filter normalization are a common diagnostic.
• ML for science: saddle points as the goal. In chemistry and materials science, a reaction goes from one minimum of the potential energy surface to another through a transition state, an index-1 saddle point (exactly one negative Hessian eigenvalue). Methods such as the nudged elastic band and the dimer method search for these saddles on purpose. The dimer method finds the lowest-curvature direction from gradients only, in the same spirit as a finite-difference HVP. Machine-learned interatomic potentials make these searches much cheaper.
• Physics-informed neural networks. Their loss landscapes are often badly conditioned because of the differential operators in the loss. Rathore et al. (2024) studied this with Hessian spectra and recommended Adam followed by L-BFGS or a Newton-type method.

7. Common mistakes #

• "GD converged, so it found a minimum." The theorem gives a small gradient. The Hessian decides what the point is.
• "Saddles trap gradient descent." Strict saddles attract only a measure-zero set of starts. The real cost is the time spent near them (plateaus), which grows with $\log(1/\delta)$ and $1/|\lambda_{\min}|$.
• Using Newton's method on a nonconvex loss without changes. It is attracted to saddles. Use damping, a trust region, or a positive semidefinite model (Gauss–Newton, lesson 10).
• Power iteration near a saddle. It returns the eigenvalue of largest absolute value, possibly a large negative one. Use Lanczos with which="LA", or shift $H$.
• Building the Hessian to get its top eigenvalue. For a big model, use HVPs; never form the $d \times d$ matrix.
• The wrong finite-difference step. $h = 10^{-8}$ suits a one-sided difference; a central difference needs about $\varepsilon_{\text{mach}}^{1/3}$ times the scale.
• Calling a short rise of the loss a divergence. At the edge of stability the loss goes up and down but decreases over long times. A real divergence keeps growing.
• Comparing sharpness across parametrizations or across models of different scale, or reading it from loss plots without filter normalization.

8. Check yourself #

1. Find and classify the stationary points of $f(x, y) = x^3 - 3x + y^2$. Can GD from a random start converge to the saddle? Where else can it go?
   Answer

   $\nabla f = (3x^2 - 3,\ 2y) = 0$ gives $(\pm1, 0)$, and $\nabla^2 f = \operatorname{diag}(6x, 2)$. At $(1, 0)$: $\operatorname{diag}(6, 2) \succ 0$, a strict local minimum with $f = -2$. At $(-1, 0)$: $\operatorname{diag}(-6, 2)$, a strict saddle with $f = 2$. GD with a small step reaches the saddle only from a measure-zero set of starts. But $f$ is unbounded below: for $x < -1$ the partial derivative $3x^2 - 3 > 0$ pushes $x$ to $-\infty$. The theorem of Lee et al. only says "if GD converges, the limit is a minimizer".
2. Full-batch GD with $\eta = 0.02$ has measured sharpness 104, and the loss went up in the last three steps. Is the run diverging?
   Answer

   Not necessarily. $2/\eta = 100$, and at the edge of stability the sharpness hovers at or just above $2/\eta$ while the loss goes up and down over short times but decreases over long times. Look at the long-time trend; a real divergence keeps growing.
3. For $f(u, v) = \tfrac12(uv - 1)^2$, at which of the minima $(1, 1)$ and $(2, 0.5)$ can GD with $\eta = 0.5$ stay? Which one generalizes better?
   Answer

   The sharpness $u^2 + v^2$ is 2 at $(1, 1)$ and $4.25$ at $(2, 0.5)$. GD needs it at most $2/\eta = 4$, so it can stay only at $(1, 1)$. Both compute the same function $uv = 1$, so they generalize equally well: sharpness alone does not decide generalization.
4. A Hessian has eigenvalues $2$ and $-3$. What does power iteration return, and how do you get $\lambda_{\max} = 2$ instead?
   Answer

   It returns $-3$, the eigenvalue of largest absolute value. Use Lanczos with `which="LA"`, or run power iteration on $H + sI$ with a shift $s > 3$ (for example $s = 5$: the eigenvalues become 7 and 2, power iteration finds 7, and $7 - 5 = 2$).
5. You estimate the top Hessian eigenvalue of a $10^9$-parameter model with 30 Lanczos iterations. Roughly what does it cost, compared with one training step?
   Answer

   Each iteration is one HVP, a small constant (about 2 to 5) times the cost of one gradient. So about 60 to 150 gradient computations on the chosen batch, plus $O(d)$ memory per stored vector. That is the cost of tens to a hundred training steps, and no $10^9 \times 10^9$ matrix.
6. At a minimum $w^\star$ with $L(w^\star) = 0.1$ and Hessian eigenvalues 40 and 1, estimate the SAM objective for $\rho = 0.05$.
   Answer

   $L^{\text{SAM}}_\rho(w^\star) \approx L(w^\star) + \tfrac12\rho^2\lambda_{\max} = 0.1 + \tfrac12\cdot 0.0025\cdot 40 = 0.15$. The worst point in the ball lies along the sharpest direction.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. What is a strict saddle?

   Show answer

2. What does GD with step $1/L$ guarantee on an $L$-smooth nonconvex $f$?

   Show answer

3. Why do strict saddles rarely trap gradient descent?

   Show answer

4. How many GD steps to escape a saddle from distance $\delta$?

   Show answer

5. Why is Newton's method attracted to saddles?

   Show answer

6. Write the Hessian-vector product as a derivative, and give its cost.

   Show answer

7. When is GD on a quadratic stable, in terms of the sharpness?

   Show answer

8. What are the edge-of-stability observations (Cohen et al. 2021)?

   Show answer

9. What is Dinh et al.'s objection to "flat minima generalize better"?

   Show answer

10. Write the SAM objective and its ascent step.

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Do saddle points trap SGD?" — Not in general. GD from a random start reaches a strict saddle with probability zero (Lee et al.), and noise escapes in polynomial time (Jin et al.). Saddles and the plateaus near them can still slow training.
• "What does gradient descent guarantee on a nonconvex loss?" — With step $1/L$, the best squared gradient norm is at most $2L(f(x_0) - f_{\inf})/k$: a first-order stationary point, $O(1/\varepsilon^2)$ iterations; prove it with the descent lemma and a telescoping sum.
• "How do you get the top Hessian eigenvalue of a 1B-parameter model?" — Lanczos or power iteration on HVPs from autodiff, on a fixed batch; a few gradients per iteration, tens of iterations.
• "What is the edge of stability?" — Full-batch GD drives the sharpness up to $2/\eta$ and trains there with a non-monotone but decreasing loss; the step size selects how sharp a region GD can stay in.
• "Do flat minima generalize better?" — Often correlated, with the right normalization (Li et al.), but sharpness changes under reparametrization (Dinh et al.). SAM optimizes for flatness at twice the cost per step.

10. Practice #

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

1. Measure the escape time of GD from a saddle and compare it with $\log(1/\delta)/\log(1 + \eta|\lambda_{\min}|)$.
2. Write a central-difference HVP and test it on a quadratic and on a small NumPy network.
3. Write power iteration on HVPs and compare its result with np.linalg.eigvalsh of the full Hessian.
4. Train the small network with full-batch GD for two step sizes and track the sharpness against $2/\eta$.
5. Write one SAM step and compare the SAM objective at a sharp and a flat minimum.

11. Watch and read #

The Problem of Local Optima (C2W3L10)Andrew Ng · DeepLearning.AI · 5 min

Why bad local minima are rare in high dimensions and why plateaus near saddles are the real problem.

The Misconception that Almost Stopped AI [How Models Learn Part 1]Welch Labs · 23 min

A visual story of loss landscapes in high dimensions and of the old worry about bad local minima.

Spotlight Talk: How to Escape Saddle Points EfficientlyPraneeth Netrapalli · Simons Institute · 16 min

Perturbed gradient descent and its almost dimension-free escape guarantee (Jin et al. 2017).

11/02/2021 -- Jeremy Cohen (CMU)CMU AI Seminar · 58 min

The first author presents "Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability".

• Cohen et al., Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability (ICLR 2021).
• Cohen et al., Adaptive Gradient Methods at the Edge of Stability (2022).
• Damian, Nichani and Lee, Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability (ICLR 2023).
• Lee, Simchowitz, Jordan and Recht, Gradient Descent Converges to Minimizers (COLT 2016).
• Jin et al., How to Escape Saddle Points Efficiently (ICML 2017).
• Dauphin et al., Identifying and attacking the saddle point problem in high-dimensional non-convex optimization (NeurIPS 2014).
• Keskar et al., On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima (ICLR 2017).
• Dinh et al., Sharp Minima Can Generalize For Deep Nets (ICML 2017).
• Li et al., Visualizing the Loss Landscape of Neural Nets (NeurIPS 2018).
• Foret et al., Sharpness-Aware Minimization for Efficiently Improving Generalization (ICLR 2021).
• Andriushchenko et al., A Modern Look at the Relationship between Sharpness and Generalization (ICML 2023).