Non-Euclidean Gradient Descent Operates at the Edge of Stability

In plain words

Gradient descent improves a model by taking repeated steps downhill. With a fixed step size, a known pattern appears in deep learning: the loss first falls while the landscape around the current point gets steeper, until the steepness reaches a ceiling set by the step size. From there the loss no longer falls smoothly. It jumps around while still falling over the long run. Researchers call this the edge of stability.

Most of that work assumed the usual way of measuring distance, where a step is simply opposite to the gradient. Many newer optimizers measure distance differently, for example by moving every weight by the same amount, by changing one block of weights at a time, or by using a matrix-aware geometry. This paper defines steepness for any such choice of distance. It then shows that, for several of these optimizers, the new steepness measure climbs to the same ceiling and stays near it. The authors also offer a small theory for simple bowl-shaped losses, where the same ceiling separates stable steps from divergent ones. The result is a diagnostic that can be read during training.

TL;DR

• Extends the edge-of-stability analysis from plain gradient descent to gradient descent under any norm, covering l-infinity (sign-like), block coordinate and spectral updates, and their normalized versions §1 §3 §4.
• Uses directional smoothness, the average curvature between two consecutive iterates, to define a generalized sharpness; for the Euclidean norm it reduces to the largest Hessian eigenvalue §2 §2.1 §3.
• Author-reported: on MLP, CNN and Transformer runs the generalized sharpness rises and hovers near, or slightly above, 2/eta, while the l2 sharpness stays well below it §3 fig. 2 fig. 4.
• Theory covers quadratics only: stable below the threshold, and divergent from some starting points above it §5 app. J.
• No official code was found.

Problem

Edge-of-stability behaviour is documented for vanilla gradient descent, preconditioned or adaptive methods, and sharpness-aware minimization §1 §1.1. The authors note that most prior analysis covers a narrow family of algorithms, and that it is unclear whether the same picture holds for methods such as Muon, Scion and SignGD §1.1. For the non-Euclidean methods studied here, the paper finds that the l2 sharpness stays well below the threshold while the generalized sharpness reaches it §3, so a norm-aware definition is needed §2.

Method

Setting: full-batch loss L(w), step size eta, and a norm ||.|| on parameters with dual norm ||.||* and dual gradient (g)* §1. The dual gradient is a maximiser of <g, y> over the unit ball. The divergence result refers to a "valid dual-gradient selection", so a choice may be needed when it is not unique §5.

Non-Euclidean gradient descent. The update minimises a linearised loss plus a squared-norm penalty. Written out eq. 1, with the dual norm and dual gradient defined eq. 2:

# Non-Euclidean GD step  [eq. 1] [eq. 2]
# Input:  w (d,), gradient g (d,), step size eta, norm
# Output: new w (d,)

def nonEuclidean_step(w, g, eta, norm):
    # Dual norm of the gradient: largest <g, y> over ||y|| <= 1  [eq. 2]
    dual_size = dual_norm(g, norm)
    # A maximiser y of <g, y> over ||y|| = 1  [eq. 2]
    dual_dir = dual_gradient(g, norm)
    # Step length eta * ||g||_*, pointing along the dual gradient  [eq. 1]
    return w - eta * dual_size * dual_dir

Normalized variant. Dropping the dual-norm factor gives the normalized method, w_{t+1} = w_t - eta (g)_* eq. 3 §1. For the l-infinity norm, the first form is l-infinity descent and the second is SignGD; for the spectral norm, the second is normalized spectral GD §1.

Directional smoothness. For a step from w to y, the quantity D is the ratio of the loss gap above the linear model at w, to half the squared step length in the chosen norm eq. 5. It is the tightest constant for which the quadratic upper bound holds with equality §2. It can be negative §2.

# Directional smoothness of one step  [eq. 5]
# Input:  loss L, weights w, next weights w_next, gradient g_w at w
# Output: scalar D

def directional_smoothness(L, w, w_next, g_w, norm):
    # Loss gap above the linear model at w, evaluated at w_next  [eq. 5]
    gap = L(w_next) - L(w) - dot(g_w, w_next - w)
    # Divide by half the squared step length in the chosen norm  [eq. 5]
    step = w_next - w
    return gap / (0.5 * norm(step) ** 2)

Decrease criterion. Substituting one step into the bound gives the change in loss eq. 6. Whenever the gradient is nonzero, the loss falls exactly when D is at most 2/eta eq. 7 §2. With normalization, the same criterion reads D <= 2 ||g||_* / eta eq. 21 eq. 22 §4.

Generalized sharpness. For a twice-differentiable loss and a nonzero step direction, D equals a Hessian quadratic form in the step direction, evaluated at an intermediate point eq. 8. Bounding it by the maximum over all directions and assuming the Hessian is nearly constant on the segment gives the generalized sharpness S(w) as a maximum over the unit ball eq. 9 eq. 10 eq. 11 §2.1. In the Euclidean case it is the largest Hessian eigenvalue §3.

Special cases with the closed forms used in the paper §3:

• Preconditioned l2 norm with matrix P_t: update w - eta P_t^{-1} grad eq. 12; sharpness is the largest eigenvalue of P_t^{-1/2} H P_t^{-1/2} eq. 13. This covers AdaGrad, RMSprop and Newton's method §3.
• l-infinity norm: update w - eta ||g||_1 sign(g) eq. 14; sharpness is the maximum of d^T H d over sign vectors d eq. 15. The paper notes this is NP-hard in general §3.
• Block l_{1,2} norm: only the block with the largest gradient norm moves, which is Block coordinate descent eq. 16 app. I. For positive semidefinite Hessians, sharpness is the largest block eigenvalue eq. 17 eq. 18 app. I.
• Spectral norm on layer matrices: each layer moves along the polar factor U V of its gradient, scaled by the sum of nuclear norms over all layers eq. 19 app. I. Sharpness is the maximum over matrix directions of the Hessian-vector product eq. 20 §3.

Estimating sharpness (Algorithm 1). For norms other than the Euclidean and block cases, S is approximated with Frank-Wolfe from several random starts, taking the maximum over restarts alg. 1 §2.1.

# Frank-Wolfe estimate of generalized sharpness  [alg. 1] [§2.1]
# Input:  hvp(d) = Hessian times d; project(d) onto the unit ball; M, K
# Output: estimate S of max over ||d|| = 1 of d^T H d

def sharpness_estimate(hvp, project, M, K):
    # Start from zero, so the estimate is never negative  [alg. 1]
    S = 0.0
    for restart in range(M):
        # Random start, projected onto the unit ball  [alg. 1]
        d = project(random_normal_like_params())
        for k in range(K):
            # Linear oracle: v maximises <hvp(d), v> over ||v|| <= 1  [alg. 1]
            v = linear_oracle(hvp(d))
            # Step size gamma_k = 2 / (2 + k)  [alg. 1]
            gamma = 2.0 / (2.0 + k)
            d = (1.0 - gamma) * d + gamma * v
        # Quadratic form at the final iterate  [alg. 1]
        S = max(S, dot(d, hvp(d)))
    return S

Mechanism on quadratics. For L(w) = 1/2 w^T H w with H positive definite and S the generalized sharpness, the authors prove two results §5:

• If eta < 2/S, the method converges linearly from any start. The contraction factor per step is 1 - 2 mu eta (1 - eta S / 2), with mu the smallest value of the quadratic form over the unit sphere app. J.
• If eta > 2/S, there is a start and a dual-gradient choice for which the method diverges. The proof uses a lemma that the maximising direction is invariant under the update, so along that direction the iterate follows w_t = (1 - eta S)^t w_0 eq. 23 §5.
• An appendix aside shows divergence for every nonzero start when eta > 2/mu, with the weaker mu in place of S app. J.

Hyperparameters

Training details, as stated in the paper:

Results

Author-reported. The paper has no numeric benchmark table. Setup: full-batch training on CIFAR-10 (5k subset and full set) and Tiny Shakespeare; sharpness values are Frank-Wolfe estimates §3 app. D. No baseline other than the l2 sharpness and the 2/eta threshold.

Paper vs code

No official code.

The paper names no project page, and its only pointer to other code is the implementation base: "Our implementation is based on open source code from" Cohen et al. (2021) app. D. That is another paper's code and is not recorded here.

We searched the following. The paper and its appendix contain no repository link. The GitHub repository search for "non-euclidean edge of stability" returned MachineLearning-Nerd/icml26-non-euclidean-edge-stability, described as an independent reproduction audit of this paper, with no license and an owner who is not an author. It is not used. Searches for "generalized sharpness non-euclidean", "directional smoothness edge of stability", "signgd sharpness edge stability", "edge of stability" with the author Islamov, with user gowerrobert, and with org flatironinstitute returned nothing.

Pitfalls

• Sharpness is an estimate for most norms. For l-infinity, spectral and general block cases, S is approximated by Frank-Wolfe. With too few restarts the estimate underestimates the maximum, as the paper shows for l-infinity (M = 1) and block CD app. E app. F. Reproduce the estimate with the same M and K and check that it is stable in M before you read the 2/eta level off it.
• Algorithm 1 and its text disagree on the final projection. The algorithm computes u_K = Pi(d_K), the projection onto the unit ball, but the estimate it returns uses d_K, not u_K alg. 1. The text says the projection "always improved the estimation" §2.1. It is unclear which quantity produced the plotted curves.
• Algorithm 1 clips at zero. S starts at 0, so a negative curvature in every direction yields an estimate of 0, not a negative value alg. 1.
• Restart counts differ between text and appendix. The block-CD text in §3 says M = 10, but the appendix sweeps {1, 7, 15} app. F. The l-infinity text says "increasing the number of restarts to 15", but the sweep and caption use {1, 5, 10} app. E fig. E.
• Gradient scaling in the spectral lemma. The Frank-Wolfe lemma for the spectral norm uses the gradient 2 H[D] app. I. Algorithm 1 uses H d with no factor 2. The argmax is scale-invariant, so the iterates are the same, but the printed formulas differ.
• Block CD ties split the step. If several blocks share the largest gradient norm, the step is divided among them: each tied block moves by eta / |J| times its gradient app. I. A singleton J gives ordinary Block CD. The paper's Block CD experiments do not say how often ties occurred.
• Spectral update scale. Unnormalized spectral GD multiplies the polar factors by the sum of nuclear norms over all layers, not per layer eq. 19 app. I. Normalized spectral GD drops this factor §4. Implement the two variants separately; they are not interchangeable.
• l-infinity step size scale. The unnormalized l-infinity update scales with the l1 norm of the gradient eq. 14. The paper's l-infinity captions use step sizes near 1e-7 app. B app. C app. E. The paper does not discuss transferring these values to other models.
• Directional smoothness can be negative. The paper does not require D to be positive §2; handle negative values in any plot or threshold check.
• Divergence is shown only for specific starts. Theorem 2 gives divergence from one line of initial points §5. The paper says generic divergence is not proved and is an open question §5 §6. The appendix aside covers all nonzero starts, but with the weaker constant mu and no proof for general norms app. J.
• Pre-EoS oscillation. For l-infinity and spectral descent, directional smoothness starts to rise toward 2/eta, and the outputs oscillate, before S reaches 2/eta. Euclidean GD does not do this §5 app. B. A monitor that waits for S to cross 2/eta will see the instability late.
• Estimated costs. The defaults stated in app. D imply M x K = 5 x 50 = 250 Hessian-vector products per sharpness estimate, plus 5 Polar Express steps per linear oracle call for spectral norms. This is derived from the stated defaults, not measured by the paper app. D.
• Euclidean case in practice. Under the l2 norm the paper's measure is the largest Hessian eigenvalue §3. The paper does not state how that eigenvalue was computed.
• Unstated details. Transformer and MLP architectures, dataset subset construction beyond "5k", training length, batch-size details beyond "full batch", and hardware are not stated §3 app. D.
• Version. This card follows arXiv v3, updated 2026-06-28, with the LaTeX source fetched by wikillm. The paper states it is accepted to ICML 2026 as an oral presentation [title footnote].

Lineage

• Builds on: the Edge of Stability for gradient descent (Cohen et al., 2021) §1; directional smoothness (Mishkin et al., 2024) §2; preconditioned-GD edge of stability (Cohen et al., 2022) §1.1; SAM edge of stability (Long et al., 2024) §1.1.
• Compared against: the l2 sharpness (largest Hessian eigenvalue) §3. The preconditioned-Hessian sharpness for adaptive methods is matched by the paper's definition under a preconditioned norm, citing Cohen et al. (2025, central flows) §3.
• Relates to: Muon and Scion (spectral geometry), SignGD (l-infinity geometry) and Block CD, which the paper covers as special cases of its framework §1.
• Superseded by: not known from this card.