Reachability of learning tasks, by hand

Companion to the notes on Achille, Mbeng & Soatto, Dynamics and Reachability of Learning Tasks (arXiv:1810.02440). Back to the notes.

The paper models SGD as a noisy particle rolling on the loss surface, $\dot w = -\nabla U(w) + \sqrt{2D}\,n(t)$, where the noise level ("temperature") $D$ grows with the learning rate and shrinks with the batch size. It then asks how likely, and how quick, it is to get from the solution of one task to the solution of another. Each widget below takes one step of that argument and lets you watch it happen with actual particles, so the claims in the notes can be checked by eye.

1. What sets the waiting time: the barrier, or the end points?

Four hundred particles start in the left well at $w=-1$. The potential is $U(w)=h(w^2-1)^2+s\,(3w-w^3)/2$, with $h$ solved so the barrier above the left well is exactly the value you set. The tilt $s$ moves the right well up or down, so the end-to-end difference is $\Delta U = U(+1)-U(-1) = 2s$. Eq. (16) of the paper puts an end-to-end difference in the exponent of the waiting time. Kramers' law puts the barrier there. Move the tilt with the barrier fixed and watch which one the particles obey.

The potential (line) and where the particles are now (bars). Dashed: the barrier height.
Fraction of particles that have reached $w=+1$ at least once (dots), against $1-e^{-t/\tau}$ with the exact mean first-passage time $\tau$ (line).
Hold the barrier and slide the tilt from $+0.2$ to $-0.6$: $\Delta U$ changes by $1.6$, the exact $\tau$ barely moves, and so does the "Kramers" column. The "eq. (16)" column, $e^{\Delta U/D}$ relative to the untilted case, swings by orders of magnitude. The end-to-end term fixes only the ratio of forward to backward rates (detailed balance); the rate itself is set by the climb.

2. Sharp and flat minima: the $(D/2)\log b$ correction

Now two dimensions. The loss is $U(u,v)=a(u)+\tfrac12 b(u)\,v^2$: a double well along $u$, and a valley across it in $v$ whose stiffness $b(u)$ spikes to $1+$peak at the left well, so the left minimum is sharp and the right one is flat. The left well is also slightly deeper ($a(-1)=-0.05$, $a(+1)=+0.05$). Integrating out $v$ exactly gives the effective potential along the valley, $F(u)=a(u)+\tfrac D2\log b(u)$ plus a constant. The paper writes $U+D\log|H|$ (in one place $U-D\log|H|$). The particles decide.

Particles in the $(u,v)$ plane, over shading of $U$ (darker is lower). The left valley is pinched.
exact, $a+\frac D2\log b$paper, $a+D\log b$no curvature, $a$ only
Push $D$ above about $0.18$ (with peak $100$): the bottom of the sharp well turns into a bump of $F$, and the particles pile up on its two shoulders. That is the precise sense in which "sharp minima may not be minima at all" at high temperature. With the paper's coefficient the same thing would happen at half the temperature.

3. Figure 2, re-read from its own numbers

The paper prints two $8\times 8$ matrices: a static distance between tasks (left) and the epochs needed to fine-tune (centre). Rows are the target, columns the source. Click a row name to highlight that target. The question is whether the distance predicts which source is faster for a given target, or only which targets are slow.

Static distance $d(\text{source}\to\text{target})$. Nine entries are negative.
Epochs to fine-tune; blank = no convergence (and the diagonal).

4. Batch size and a noise floor

Figure 1 (centre) shows fewer steps to converge as the batch grows. Eq. (16), $\tau\propto e^{\Delta C/D}$ with $D=c/B$, cannot produce that curve with sensible limits. A different mechanism can: SGD on a quadratic loss in $k$ dimensions has a stationary loss of about $kD/2$ (equipartition), and a loss threshold close to that floor takes a long time to hit. Below is the exact mean time for the loss to fall from $50$ to the threshold, computed from the first-passage formula for $|w|$.

5. The structure function and its price, eqs. (2)–(3)

Linear regression with $N=100$ samples and 50 weight directions: 5 strong input directions carry the true weights, 45 weak ones carry only label noise. For each price $\beta$ the best "posterior" cloud $Q$ is known exactly, so every point below is exact. The blue curve is the structure function $S(t)$ of eq. (2): the best training loss with $t$ nats stored in the weights. Eq. (3) prices each nat at $\beta$: the optimum is where a line of slope $-\beta$ (dashed) touches $S$. Orange is the test loss of the same cloud.

train, $S(t)$testslope $-\beta$ (linear axis only)
Nats per weight direction at this $\beta$, strong directions on the left. Blue: the 5 that carry the signal.
Drag $\beta$ from right to left. At a high price the model stores almost nothing and both losses are large. Lowering it buys the signal directions first: the training and test losses fall together (structure). Near $\beta\approx0.005$ the test loss bottoms out at about 25 nats. Below that the extra nats go to the weak directions, fitting label noise: training loss keeps falling, test loss rises (nuisance). With random labels there is nothing to buy except the fact that the weights should be near zero.

All computations run in your browser: Euler–Maruyama for the particles (with the exact Ornstein–Uhlenbeck step across the valley in widget 2), trapezoidal quadrature for the first-passage times. The same numbers are printed by code/reachability.py.