Non-Euclidean Gradient Descent Operates at the Edge of Stability

Summary

This paper generalizes the Edge of Stability (EoS) phenomenon from vanilla and preconditioned gradient descent to a broad family of non-Euclidean gradient methods indexed by an arbitrary norm — including $\ell_\infty$-descent (SignGD), Block CD, and Spectral GD (the update underlying Muon and Scion). The conceptual move is to route the argument through Mishkin et al.'s directional smoothness $D^{\|\cdot\|}$, which admits an exact loss-change identity: whenever $\|\nabla L\|_\ast > 0$, the loss decreases iff $D^{\|\cdot\|} < 2/\eta$. A second-order expansion of $D^{\|\cdot\|}$ then motivates a norm-adapted generalized sharpness $S^{\|\cdot\|} := \max_{\|d\|\le 1} d^\top \nabla^2 L(w) d$ that reduces to $\lambda_{\max}(\nabla^2 L)$ for the Euclidean norm and to the preconditioned-Hessian sharpness for preconditioned $\ell_2$ norms, unifying earlier EoS characterizations. For norms without a closed form, $S^{\|\cdot\|}$ is estimated via multi-restart Frank-Wolfe with a projection back to the unit sphere. Empirically, across MLPs, CNNs, and Transformers, $S^{\|\cdot\|}$ progressively sharpens and then hovers near — sometimes slightly above — $2/\eta$, whereas the standard $\ell_2$ sharpness does not track the threshold under non-Euclidean updates (Figures 12 and 18 make this contrast most sharply on ResNet20/VGG11). A local quadratic analysis (Theorems 5.1–5.2) shows convergence below $2/S$ and constructs a diverging trajectory above $2/S$ from a specific invariant-direction initialization. The paper also flags a genuinely new phenomenon — a pre-EoS oscillatory regime where $D^{\|\cdot\|}$ climbs before $S^{\|\cdot\|}$ crosses the threshold — that does not appear for Euclidean GD. The main conceptual limitation is the gap between the exact $D^{\|\cdot\|}$-level identity and the empirically observed $S^{\|\cdot\|}$-level plateau: passing from $D$ to $S$ involves a Taylor step and an upper-bounding max over directions, so the sharpness result is a one-sided diagnostic rather than an equivalence. The framework is nevertheless a clean and useful lens for optimizer geometry.

Strengths

• Conceptual unification. The move to route EoS through directional smoothness and then define a norm-adapted sharpness gives a single object $S^{\|\cdot\|}$ that recovers Cohen et al.'s $\lambda_{\max}$ result for GD and the preconditioned-Hessian sharpness of Cohen et al. (2022) as special cases, while extending coherently to $\ell_\infty$-descent, Block CD, and Spectral GD.
• Timely coverage of modern optimizers. By covering the spectral norm, the paper directly addresses geometry-aware EoS for the update underlying Muon and Scion, which have become practically important and were not previously covered by the EoS literature.
• Clean loss-change identity. The two-line derivation from (5) through (7) is genuinely illuminating: it makes precise, for the first time in the non-Euclidean setting, the sense in which the algorithm's own step-size dictates a threshold on directional smoothness whenever the dual gradient is nonzero.
• Honest empirical scope. The paper reports the phenomenon across three quite different data/architecture pairs (MLP/CNN on CIFAR-10, Transformer on Tiny Shakespeare, plus ResNet20/VGG11 on full CIFAR-10) and repeatedly notes when the generalized sharpness sits 'slightly above' the threshold and offers a plausible multi-eigenvalue mechanism (Appendix C).
• Frank-Wolfe machinery documented as a heuristic. The paper is candid that $S^{\|\cdot\|}$ is NP-hard in general and that FW is used with random restarts as a heuristic estimator, and includes sensitivity plots (Figures 11, 13, 15, 16) that quantify how much the restart count and PolarExpress step budget matter.
• New pre-EoS oscillatory regime. Figures 7–8 identify a phenomenon — sustained network-output oscillations under $\ell_\infty$-descent while $S^{\|\cdot\|}$ is still below $2/\eta$ — that does not exist for Euclidean GD and that the paper honestly flags as an open direction.

Weaknesses

• Identity is at the $D$-level, not the $S$-level. The rigorous statement is the biconditional $\Delta L < 0 \iff D^{\|\cdot\|} < 2/\eta$ for the tightest directional smoothness. Passing to $S^{\|\cdot\|}$ requires a Taylor expansion with a Lagrange remainder (Eq. 8), an upper-bounding max over unit-norm directions (Eq. 9), and a 'Hessian almost constant along the chord' assumption. The paper's contributions and abstract phrasing sometimes read as though the plateau at $2/\eta$ is a proven property of $S^{\|\cdot\|}$; it is not.
• Sign-crossing vs. magnitude proximity. The identity implies that on ascent steps $D^{\|\cdot\|} > 2/\eta$ and on descent steps $D^{\|\cdot\|} < 2/\eta$ — i.e., $D^{\|\cdot\|}$ must cross the threshold — but does not imply that it hovers near $2/\eta$ in magnitude. The main text's phrasing 'must oscillate around $2/\eta$' invites the stronger reading.
• Theorem 5.2 is much weaker than a converse to 5.1. Divergence for $\eta > 2/S$ is established only from $w_0 \in \mathrm{span}(\hat d)$ and under a favorable dual-gradient selection. For non-smooth norms (sign / $\ell_\infty$, spectral norm with degenerate singular values) the dual map is set-valued, so a differently valid selection can immediately push the iterate off $\mathrm{span}(\hat d)$ and destroy divergence. The dual-selection dependence is not discussed in the surrounding prose.
• Empirical claims are single-run and unquantified. 'Hovers near' and 'slightly above' are the paper's headline observations, but no tolerance is reported, and each figure shows a single trajectory. Given that the FW estimate of $S^{\|\cdot\|}$ is itself heuristic and known (from the paper's own ablations) to be biased low with few restarts, the reader has no way to tell whether 'slightly above' is a robust feature or an estimation artifact.
• Appendix J typos and notation clashes. Theorem J.5 is stated for '$w_0 = 0$' (the fixed point that provably does not diverge; the intended condition is $w_0 \neq 0$) and its bound uses an undefined $\mu$, while the theorem's own definition of $S^{\|\cdot\|}$ is never used. Its proof also mislabels the gradient dual norm $\|H w_t\|_\ast$ as a 'parameter norm'. The Frank-Wolfe step-size formula appears mis-transcribed (rendered as 'k+22' where the standard is $2/(k+2)$).
• Frank-Wolfe convergence guarantee vs. non-convex sphere. The cited stationary-point guarantee for FW requires a convex feasible set, but the underlying problem is over the non-convex unit sphere; the paper 'relaxes' but never states exactly which convex set FW operates on, and then reinstates the sphere via a boundary projection. The reader cannot reconcile the projection with the cited guarantee.
• Minor consistency issues. 'Block GD' and 'Block CD' are used interchangeably in Appendix F.1 while the main text says only 'Block CD'; the sentence 'Increasing the number of restarts to 15…' is placed in the $\ell_\infty$ subsection whose restart grid is $\{1,5,10\}$; Figure 17's caption calls the real-objective curve both 'gray' and 'blue'; the sentence 'In opposite, the dynamics on the quadratic model…' is grammatically broken; Figure 5's caption uses a doubled subscript $\|\nabla L(w_t)\|_{\ast\ast}$ where $\|\cdot\|_\ast$ is intended.

Reproducibility & code

The evidence.json bundle (paper-code mode) contains no actual released source: the replication/codebase directory is empty, so I can only assess what the manuscript documents against what a reader would need to redo the experiments. The paper says implementations are 'based on open source code from Cohen et al. (2021)' but does not link to a public repository for this project.

• Architecture details promised but missing. The main text explicitly defers 'architecture and implementation details' to the appendix, but Appendix D contains only losses, Frank-Wolfe defaults, and PolarExpress step count. The MLP width/depth, CNN channel/kernel choices, and the Transformer's layers/heads/embedding are absent for every figure. This blocks reproduction of every empirical result.
• Seeds and initialization absent. The word 'seed' does not appear anywhere in paper.md, and no weight-initialization scheme, no CIFAR-10-5k subset construction, and no dataset preprocessing are described. Single-seed trajectories underpin every 'hovers near $2/\eta$' claim, so the missing randomness controls matter for the headline conclusions.
• Tiny Shakespeare pipeline undocumented. The tokenizer, sequence length, and batch construction for the Transformer runs (Figures 2, 4, 5) are not stated; this is the only non-vision setting and is used to argue architecture-independence.
• Normalization layers unresolved for ResNet20 / VGG11. Canonical implementations of these networks include BatchNorm; the manuscript never states whether BN layers are present, frozen, or absorbed into the spectral / block update. Since Figures 12 and 18 are the main evidence that $S^{\|\cdot\|}$ tracks the threshold on 'real' architectures where $\ell_2$ sharpness does not, the treatment of normalization is load-bearing.
• Convolutional matrix unfolding for Spectral GD unspecified. The polar-factor update assumes each layer is a matrix, but no reshape convention is given for conv kernels ((C_out, C_in·k·k) vs. alternative unfoldings), and whether biases are updated Euclidean-style or included in the spectral block is not stated.
• RMSprop hyperparameters missing. The Appendix H claim that the adaptive-EoS characterization breaks down as $\beta_2 \to 0$ depends on RMSprop's learning rate, $\varepsilon$, and second-moment initialization — none of which are documented — so the observed departure from $2/\eta$ cannot be cleanly attributed to $\beta_2$ alone.
• Block partition not fully described. For Block CD the paper says blocks correspond to 'a weight matrix or bias vector of a layer', but the treatment of biases, per-filter grouping in conv layers, and skip connections is not spelled out; this matters because Figure 3's sudden drops in $D^{\|\cdot\|}$ are attributed to inter-block transitions.

Recommended Changes

Essential

• Sharpen the $D$-vs-$S$ distinction. In the abstract, contributions, and Section 2.1, explicitly state that the biconditional in (7) is for directional smoothness $D^{\|\cdot\|}$, and that the corresponding statement for $S^{\|\cdot\|}$ is an empirical, upper-bound-based diagnostic that follows from the Taylor step (Eq. 8) and the max-over-directions relaxation (Eq. 9). Reword contribution 2 to say $D^{\|\cdot\|}$ must cross $2/\eta$, not 'oscillate around' it.
• State the dual-selection hedge in Theorem 5.2. In the prose immediately following Theorem 5.2, name the dual-gradient selection dependence (in addition to the initialization dependence already discussed), and give an example — e.g. degenerate spectral norm singular values or ties in SignGD — where a different valid selection destroys the constructed divergence.
• Fix Appendix J. In Theorem J.5: change 'every initial point $w_0 = 0$' to '$w_0 \neq 0$', define $\mu$ in the theorem statement (or drop the unused $S^{\|\cdot\|}$ line), and correct 'the parameter norm $\|H w_t\|_\ast$' to name the gradient dual norm. Verify and re-typeset the Frank-Wolfe step size to $\gamma_k = 2/(k+2)$ in Lemmas I.4 and I.9.
• Report per-seed variance and a quantitative 'closeness'. Add per-seed trajectories (or at least mean $\pm$ range envelopes across $\ge 3$ seeds) for the main figures (1–5, 12, 18) and report the empirical distribution of $S^{\|\cdot\|}/(2/\eta)$ in the post-sharpening window. This is the missing quantification of the paper's headline 'hovers near, sometimes slightly above' claim.
• Add a training-details table. In Appendix D, give MLP widths/depths, CNN channel counts and kernel sizes, Transformer configuration (layers/heads/embedding), Tiny Shakespeare tokenizer and sequence length, seed(s), weight-init distribution, CIFAR-10-5k subset construction, and any preprocessing. Cover every figure.
• Document ResNet20/VGG11 normalization layer handling and the convolutional matrix unfolding used by Spectral GD. State whether BatchNorm layers are included, frozen, or removed; state the reshape convention (e.g. (C_out, C_in·k·k)) for conv kernels under the polar-factor update; state whether biases participate in the block or spectral update or default to Euclidean.
• Release the code. Point to a public repository (or at least commit to release upon acceptance), and clarify which changes are on top of the Cohen et al. (2021) codebase, in particular the FW-with-restarts routine, the layer-wise polar update, and the Block CD scheduler.

Suggested

• Clarify the Frank-Wolfe relaxation. In Section 2.1 / Appendix A, state explicitly the convex set FW operates over (presumably the unit ball, then a boundary projection) and reconcile the boundary projection step with the cited Lacoste-Julien (2016) stationary-point guarantee.
• Bring the invariance argument into the main text. Add one sentence hinting at why $\hat d$ is a 'nonlinear eigenvector' of the dual map $d \mapsto (Hd)_\ast$ (via the KKT conditions of the max problem), so Section 5's central mechanism is visible without reading Appendix J.
• Reconcile the 'restarts not necessary' claim with the Section 3 default of $M=5$. Either weaken the claim in G.2 to a quantitative tolerance ('$M=1$ agrees with $M=10$ within $x\%$ on Spectral GD') or drop the normative recommendation, since Figures 4/5 were produced with $M=5$.
• Unify the sensitivity restart grids across Figures 11, 13, 15 (currently $\{1,5,10\}$ vs. $\{1,7,15\}$) so the cross-method sensitivity comparison in G.2 is apples-to-apples.
• Rename or clarify 'Block GD' vs 'Block CD' throughout Appendix F.1 (main text and figures already use Block CD). Fix the 'Increasing the number of restarts to 15' sentence in Appendix E.2 (whose restart grid is $\{1,5,10\}$), the color-name clash for the real objective in Figure 17's caption, the doubled subscript $\|\cdot\|_{\ast\ast}$ in Figure 5's caption, and the garbled 'In opposite' sentence in Appendix G.4.
• Explicitly note the sign of $D^{\|\cdot\|}$. Add a one-line remark near Definition 2.1 that the tightest $D^{\|\cdot\|}$ can be negative on descent steps and that 'progressive sharpening' refers to crossing behaviour, not monotone positive growth.
• Document RMSprop hyperparameters for the Appendix H sweep ($\eta$, $\varepsilon$, initial second-moment values, any warm-up), since the AEoS-breakdown conclusion at small $\beta_2$ depends on them.