Interactive research note

Mirror Descent Framework for Edge of Stability

This page turns the paper's CIFAR-10 experiments into a small interactive guide. The central question is whether mirror descent methods, once measured in their own Bregman geometry, train at the same edge-of-stability boundary that full-batch gradient descent reaches.

Interactive format inspired by the Central Flows companion site: centralflows.github.io/part1.

01

Why the edge exists

On a quadratic loss, gradient descent is stable only when the learning rate times the relevant curvature stays below 2. Deep networks do something stranger: sharpness can rise to that boundary, oscillate, and still make long-run progress.

\[ \text{GD edge:}\qquad \eta\,\lambda_{\max}(H) \approx 2 \]
02

Mirror descent changes the ruler

Mirror descent replaces Euclidean distance with a geometry induced by a mirror map. In the paper's experiments, this means the relevant curvature is the top eigenvalue of an effective Hessian, where the optimizer's local metric rescales the landscape.

\[ \text{MD diagnostic:}\qquad \eta\,\lambda_{\max}\!\left(G(\theta)^{-1/2}H(\theta)G(\theta)^{-1/2}\right) \]
03

What the experiments show

The full-batch GD run is the cleanest EoS reproduction. Several stochastic, adaptive, and mirror-descent variants sharpen early or spike past the edge, but they often settle below the normalized edge rather than hovering there.

That makes the main result deliberately cautious: non-GD mirror descent routines may not train at EoS as often as full-batch GD on the same MLP and CIFAR-10 setup.

Toy model

Stability threshold

Geometry

Bregman ruler

Experiment 2

Learning-rate sensitivity

Reading the results

Full-batch GD reaches the boundary.

The normalized EoS ratio for full-batch GD rises toward and then remains around the edge, ending near 1.45 after 3000 epochs.

Mirror descent spikes are not the same as hovering.

Entropy, p-norm, cosh, and Tsallis mirror descent runs all show early sharpness spikes, but their late-run normalized ratios sit well below 1 in these saved runs.

Adam has its own scale.

Following the adaptive-EoS literature, Adam is normalized by a momentum-adjusted threshold of 38 when beta1 is 0.9, rather than the GD threshold of 2.

Sources and context

This explainer uses the saved metrics and paper text from the EOS mirror descent project. Background framing follows the original EoS paper by Cohen, Kaur, Li, Kolter, and Talwalkar (2021), the adaptive-EoS extension for Adam by Cohen et al. (2022), and the interactive style and central-flow perspective of Cohen, Damian, Talwalkar, Kolter, and Lee (2025).