A wedge landscape you can walk through

Companion to the notes on Fort & Jastrzebski, Large Scale Structure of Neural Network Loss Landscapes (NeurIPS 2019). Back to the notes.

The paper's toy: the low-loss set is a union of flat pieces (wedges) and the loss is the distance to the nearest one. With $D=3$ weights and wedges of dimension $n=2$, the wedges are the three coordinate planes, and the loss of a point is $\min(|x|,|y|,|z|)$. Put two optima on different planes, and the straight line between them leaves every plane. A path that moves within a hyperplane normal to the line and slides to a wedge at each step (the paper's tunnel algorithm) does not.

Presets:

plane xy plane xz plane yz straight line tunnel

Drag to rotate. The tunnel is built the way Section 4.2 does it: 60 segments of the line, at each a plane normal to the line, and projected gradient descent on the loss inside that plane.

A angle radius

B angle radius

What to look for

Same plane: both optima lie on one wedge, so the straight line stays on it and has zero loss: no barrier, and averaging the two weight vectors is safe. This is the paper's account of why Stochastic Weight Averaging works at a low learning rate.
Different planes: the midpoint has two nonzero coordinates it should not (it is off both planes), so it has positive loss. The tunnel avoids this by walking through the intersection of the planes, here the axis they share, and the loss stays at zero at its optimised points (small bumps appear between neighbouring points).
Look at the tunnel in 3-D: it goes toward the shared axis, along the first wedge, then out along the second. The paper's Figure 3 measures the same: the deviation from the line points in one direction for the first half and a roughly orthogonal one for the second.
In 100 dimensions the wedges are far less empty: the notes' simulation shows the same wall and tunnel with $D=100$.