Landscapes

Dive into Deep Learning · §9.1

What makes deep-net optimization hard
risk vs. empirical risk · where gradients vanish · curvature and noise · the edge of stability

An optimizer is three decisions

The chapter’s frame

  1. A descent direction — which way is “down”? Depends on which norm measures the step. Euclidean → the gradient. Other norms → other algorithms (payoff: Muon).
  2. A step size over time — how far to trust the local slope, and how that trust changes over a run (schedules, warmup).
  3. A way of living with noise — every gradient is a minibatch estimate; batch size and averaging set the noise level.

Every method in this chapter, GD through Muon, is one way of making these three decisions.

Optimization vs. learning

Optimization minimizes the empirical risk (training loss). Learning wants low risk (expected loss on the population). The optimizer only ever sees the former:

%matplotlib inline
from d2l import torch as d2l
import numpy as np
import torch
def f(x):
    return x * d2l.cos(np.pi * x)

def g(x):
    return f(x) + 0.2 * d2l.cos(5 * np.pi * x)

The two minima sit in different places — and no optimizer can fix that:

def annotate(text, xy, xytext):
    d2l.plt.gca().annotate(text, xy=xy, xytext=xytext,
                           arrowprops=dict(arrowstyle='->'))

x = d2l.arange(0.5, 1.5, 0.01)
d2l.set_figsize((4.5, 2.5))
d2l.plot(x, [f(x), g(x)], 'x', 'risk')
annotate('min of\nempirical risk', (1.0, -1.2), (0.5, -1.1))
annotate('min of risk', (1.1, -1.05), (0.95, -0.5))

Local minima

f(x) = x \cos(\pi x) has a local minimum that is not global. Near it, the gradient goes to zero — the signal cannot tell the two apart:

x = d2l.arange(-1.0, 2.0, 0.01)
d2l.plot(x, [f(x), ], 'x', 'f(x)')
annotate('local minimum', (-0.3, -0.25), (-0.77, -1.0))
annotate('global minimum', (1.1, -0.95), (0.6, 0.8))

Noise can knock the iterate out of a shallow basin — minibatch variance supplies exactly that.

Saddle points

1D: f(x) = x^3 has f'(0) = 0, yet no minimum:

x = d2l.arange(-2.0, 2.0, 0.01)
d2l.plot(x, [x**3], 'x', 'f(x)')
annotate('saddle point', (0, -0.2), (-0.52, -5.0))

High-dim: a zero-gradient point is a minimum only if all Hessian eigenvalues are positive — with mixed signs it is a saddle. At 10^6 parameters, essentially every critical point is a saddle:

x, y = d2l.meshgrid(
    d2l.linspace(-1.0, 1.0, 101), d2l.linspace(-1.0, 1.0, 101))
z = x**2 - y**2

ax = d2l.plt.figure().add_subplot(111, projection='3d')
ax.plot_wireframe(x, y, z, **{'rstride': 10, 'cstride': 10})
ax.plot([0], [0], [0], 'rx')
ticks = [-1, 0, 1]
d2l.plt.xticks(ticks)
d2l.plt.yticks(ticks)
ax.set_zticks(ticks)
d2l.plt.xlabel('x')
d2l.plt.ylabel('y');

Vanishing gradients

No critical point needed: f(x) = \tanh(x) at x = 4 has f'(4) \approx 0.0013. The surface is just flat where we stand:

x = d2l.arange(-2.0, 5.0, 0.01)
d2l.plot(x, [d2l.tanh(x)], 'x', 'f(x)')
annotate('vanishing gradient', (4, 1), (2, 0.0))

ReLU and good initialization fixed this at the model level — not the optimizer’s job.

The first villain: curvature

f(\mathbf{x}) = 0.1 x_1^2 + 2 x_2^2: curvatures 0.2 and 4, one learning rate. Steep direction caps \eta < 0.5; flat direction then keeps > 90\% of its value per step:

def f_valley(x1, x2):  # Second derivatives 0.2 and 4
    return 0.1 * x1 ** 2 + 2 * x2 ** 2

def gd_valley(x1, x2, s1, s2):
    eta = 0.45  # Just under the stability ceiling of 0.5
    return (x1 - eta * 0.2 * x1, x2 - eta * 4 * x2, 0, 0)

d2l.show_trace_2d(f_valley, d2l.train_2d(gd_valley, steps=30))

epoch 30, x1: -0.295265, x2: -0.002476

Zig-zag across, crawl along. Condition number \kappa = \lambda_{\max}/\lambda_{\min} = 20; iterations scale linearly with \kappa. Real networks: \kappa in the thousands.

Momentum → \sqrt{\kappa}. Adam → per-coordinate rescaling. Muon → per-matrix rescaling. Most of the chapter fights this picture.

The edge of stability

Classical advice: measure sharpness \lambda_{\max}, pick \eta < 2/\lambda_{\max}.

Measured reality (Cohen et al., 2021): causality runs backwards — training raises sharpness (“progressive sharpening”) until it reaches \approx 2/\eta, then hovers there. Loss keeps falling, non-monotonically, in the “forbidden” regime.

  • The ceiling is an attractor, not a fence: pick \eta, the network adapts its curvature to it.
  • Training does not live in the tidy descent regime the proofs analyze.
  • One reason warmup and schedules matter (§ Schedules); measured on a 25-parameter net in the math appendix.

The second villain: noise

The gradient is a minibatch estimate: unbiased, variance \propto 1/b (measured on a real network in the SGD section).

  • Constant \eta → no convergence: a noise ball of squared radius \propto \eta. Hence decaying learning rates and schedules.
  • Batch size = a second dial, with hardware consequences (Minibatches) and diminishing returns at scale (Batch Size).
  • Momentum’s second job: averaging noise over time.
  • Not purely a tax — noise kicks the iterate out of saddles and shallow minima; deep barriers stay expensive.

What convexity still buys

Deep losses are not convex — permutation symmetry alone gives every minimum d! separated copies; convex minima form one connected set.

What survives:

  • Language and baselines: condition number, rates, noise ball — all theorems in the convex world. A method that fails on a quadratic has no business near a transformer.
  • Local honesty: near a good minimum the loss is approximately a quadratic bowl — which is why the valley cartoon predicts late-training behavior (and why weight averaging works).

Full treatment: the convexity chapter of the math appendix.

Recap

  • Minimizing training loss ≠ minimizing test loss; that gap belongs to regularization, not the optimizer.
  • Classical hazards: local minima, saddles (dominant in high dim), vanishing gradients.
  • Practical hazards: curvature (condition number \kappa) and noise (minibatch variance) — the chapter’s two villains.
  • Modern twist: training equilibrates at the edge of stability.
  • The toolkit ahead = three decisions: direction, step size over time, living with noise.