非欧几里得优化下的“稳定性边缘”：Muon 与 SignGD 的统一几何视角
Non-Euclidean Gradient Descent Operates at the Edge of Stability
总结
问题
方法
结果
要点
摘要
本文将深度学习中广泛存在的“稳定性边缘” (Edge of Stability, EoS) 现象从传统的欧几里得空间推广到了非欧几里得范数空间。通过引入基于方向平滑度 (Directional Smoothness) 的广义锐度 (Generalized Sharpness) 定义，证明了包括 SignGD, Block CD 和 Spectral GD (Muon) 在内的多种优化算法在训练后期均会维持在 2/η 的稳定性阈值附近。
TL;DR
在深度学习的黑盒中，有一个神秘的常数 2/η（η 为步长）。无论你是在做普通的梯度下降，还是在使用最前沿的 Muon 或 SignGD，模型的“锐度”最终都会像撞上一道无形的墙一样，在这条线上反复横跳。本文揭示了这一现象在非欧几里得范数下的普遍性，并提出了一套通用的度量工具。
背景定位：为何传统的锐度定义失效了？
经典的优化理论告诉我们，步长 η 必须小于 2/L 才能保证收敛（L 为平滑常数）。然而，Deep Learning 的真实情况是：Hessian 的最大特征值 λmax 经常在训练初期疯狂增长，直到跨过 2/η 的阈值。此时，损失函数开始剧烈震荡，但整体依然下降。这就是著名的 Edge of Stability (EoS)。
问题在于，λmax 是基于 L2 范数 定义的。当我们使用 SignGD（基于 范数）或者最近霸榜 LLM 训练的 Muon（基于谱范数）时，L2 锐度不再能准确描述算法的稳定性边缘。
核心直觉：方向平滑度 (Directional Smoothness)
作者提出，要理解非欧几里得下的 EoS，必须放弃全局平滑假设，转而关注 方向平滑度 (Directional Smoothness)。
简单来说，如果损失函数的变化 满足：
abla \mathcal{L}(w_t)\|_*^2$$ 那么，只要局部曲率 $D$ 满足 $D \le 2/\eta$，损失就会下降。这导出了一个关键结论：如果我们观察到损失开始震荡，那么对应的（广义）锐度必然在 2/η 附近盘旋。 ## 广义锐度的定义与计算 对于任何范数 $\|\cdot\|$，作者定义了广义锐度： $$S^{\|\cdot\|}(w) := \max_{\|d\| \le 1} d^ op abla^2 \mathcal{L}(w) d$$ - **L2 范数**：回归到传统的 $\lambda_{max}$。 - **$\ell_\infty$ 范数**：对应于寻找超立方体上的最大能量态（Ising 模型，NP-Hard）。 - **谱范数 (Muon)**：对应于矩阵层面的极向因子 (Polar Factor)。 由于许多范数下的最大化问题是 NP-Hard 的，作者引入了 **Frank-Wolfe (FW)** 算法来进行启发式逼近。  *上图展示了非欧几里得 GD 的基本更新逻辑，其核心在于对正则化线性化问题的求解。* ## 实验证据：EoS 是通用规律 作者在多种模型上验证了这一猜想。最直观的证据来自 **Spectral GD (Muon)** 的训练过程：  *图中清晰可见：随着训练进行，Directional Smoothness（绿线）和 Generalized Sharpness（红线）最终都收敛并稳定在 2/η（虚线）处。* 更有趣的是，在 **SignGD**（$\ell_\infty$ 范数）中，即使传统的 L2 锐度表现平平，如果你去看基于 $\ell_\infty$ 定义的广义锐度，它依然严格遵守 2/η 的“天条”。  ## 深度洞察：二次泰勒展开的“发散”印证 为了进一步验证 2/η 的物理含义，作者做了一个非常硬核的实验：在进入 EoS 阶段后，突然将真实的非凸损失函数替换为其局部**二次泰勒近似 (Quadratic Taylor Approximation)**。 - **EoS 之前**：在二次模型上继续优化，损失平稳下降。 - **EoS 之后**：在二次模型上继续优化，损失立即**指数级爆炸**。 这证明了 EoS 的本质：非线性项（如三阶导数）在某种程度上扮演了“刹车装置”的角色，防止了模型在跨越稳定性边界后彻底崩盘，实现了所谓的“自稳定 (Self-Stabilization)”。 ## 总结与局限 本文成功地将 EoS 现象从一种“实证观察”提升到了“几何特性”的高度。尽管在非欧几何下，证明“所有初值都会在 2/S 时发散”依然存在挑战（目前仅针对特定不变方向证明），但实验数据已经为研究者指明了方向。 **Takeaway**: 当你在调优 Muon 或 SignGD 的学习率时，请记住：此时起决定作用的不再是 Hessian 的特征值，而是那个隐藏在范数定义下的“广义锐度”。 --- *注：文中部分公式推导参考了原论文 Appendix E 关于二次函数的收敛性证明。*