A problem no learning rate solves

STATEMENT

Exhibit a quadratic on which every learning rate fails — not one badly chosen rate, but the whole stable range — and identify what is responsible. Then show that remembering the previous step fixes it, and by how much.

GIVEN

$\mat{H} = \mathrm{diag}(100,\ 0.01)$, so $\kappa = 10^4$. Start at $\vec{\theta}_0 = (1, 1)$; the target is an error norm below $10^{-3}$.

FIND

The stability limit; the step count for every $\eta$ across the stable range; the count at the optimal $\eta^*$; and the same for heavy-ball momentum at its own optimum, with the speed-up.

STRATEGY

The two directions decouple, as in I.7.B02, so the difficulty can be stated before anything is run: the steep direction demands $\eta < 0.02$, and at any such $\eta$ the flat direction’s multiplier is $1 - 0.01\eta \ge 0.9998$. One rate cannot serve both. Sweeping confirms there is no exception rather than merely no obvious one.

SOLUTION

Step 1 — the bind, stated in advance. Stability requires $\eta < 2/\lambda_{\max} = 0.02$. At the largest permissible rate the flat direction is multiplied by

so it loses at most two parts in ten thousand per step, however $\eta$ is chosen. That is the counterexample in one line; the sweep is the demonstration.

Step 2 — the sweep.

Nothing in the stable range does better than about $34\,700$ steps for a thousandfold reduction on a two-dimensional problem with an exactly known Hessian.

Step 3 — the most instructive row. At $\eta = 0.01$ the steep multiplier is exactly zero: $1 - 0.01(100) = 0$. That direction is solved perfectly, in one step, and the run still takes $69\,075$ steps — twice as many as the rates that solve it imperfectly. Perfect handling of the direction that constrains stability buys nothing, because the cost is entirely in the other direction. This is why “tune the learning rate” is not a complete answer to slow training.

Step 4 — the optimum, and its transient. $\eta^* = 0.0199980$ gives $\rho = 0.99980002$ and $36\,272$ steps — slightly worse than $\eta = 0.0199$ at $34\,709$. As in I.7.B01, $\eta^*$ minimises the asymptotic rate while the run is measured over a finite budget. The gap is small here because $\kappa$ is large enough that the asymptote arrives early.

Step 5 — momentum. Heavy ball keeps a running average of gradients, so the recurrence is second-order and its rate depends on $\sqrt{\kappa}$. At the classical optimum

with $\rho = (\sqrt{\kappa} - 1)/(\sqrt{\kappa} + 1) = 0.98019802$ and 708 steps, against $36\,272$: a speed-up of $51.2\times$.

Step 6 — why it works. Momentum does not make the steep direction stable at a larger rate; note that $\eta = 0.0392$ is above the plain-descent threshold of $0.02$. It works because the steep direction’s gradient alternates in sign when overshooting, so consecutive contributions to the moving average partly cancel, while the flat direction’s gradient keeps the same sign and accumulates. The average suppresses oscillation and amplifies drift, which is exactly the asymmetry the problem needs.

Check — sanity

The measured momentum rate matches the theory. $\rho = 0.98019802 = 99/101 = (\sqrt{\kappa} - 1)/(\sqrt{\kappa} + 1)$ for $\sqrt{\kappa} = 100$, exactly. And it equals the plain-descent rate for a problem of condition number 100 — momentum turns $\kappa = 10^4$ into an effective $10^2$, which is what “the square root” means concretely.

The speed-up is $51\times$, not the $100\times$ the rates suggest. Both methods are measured to a fixed target from a fixed start, so both pay a transient the asymptotic rates ignore, and momentum’s is proportionally larger because it needs several steps before its average carries information. Quoting $\sqrt{\kappa}$ as the speed-up would be quoting a limit as a measurement.

Divergence appears immediately above $0.02$. $\eta = 0.02001$ gives a steep multiplier of $1.001$ and diverges; $\eta = 0.0199$ gives $0.99$ and converges. Five parts in ten thousand separate them, which is the sharp edge of I.7.B01 in a stiffer problem.