SGD at the Edge of Stability: The Stochastic Sharpness Gap

当使用全批量梯度下降（GD）并以学习率 $η$ 训练神经网络时，损失函数 Hessian 矩阵的最大特征值——即“尖锐度”$S(\boldsymbolθ)$——会上升至 $2/η$ 并在该值附近持续振荡，这一现象被称为“稳定性边缘”（Edge of Stability, EoS）。 \citet{damian2023selfstab} 指出，该行为可由损失函数三阶结构所驱动的一种自稳定机制予以解释；并且，GD 实质上是在约束条件 $S(\boldsymbolθ) \leq 2/η$ 下隐式执行投影梯度下降（Projected Gradient Descent, PGD）。 而对于小批量随机梯度下降（SGD），尖锐度则稳定在低于 $2/η$ 的水平，且随着批量大小减小，该稳定值与 $2/η$ 之间的差距进一步扩大；然而，目前尚无理论能解释这种尖锐度受抑现象。 本文提出“随机自稳定”（stochastic self-stabilization）机制，将原有自稳定框架拓展至 SGD 场景。我们的核心洞见在于：梯度噪声为沿 Hessian 矩阵主特征向量方向的振荡动力学注入了方差，从而增强了三次项主导的、降低尖锐度的作用力，并使系统平衡点下移至 $2/η$ 以下。 沿用 \citet{damian2023selfstab} 的分析路径，我们相对于一条动态演化的投影梯度下降轨迹，定义了“随机预测动力学”（stochastic predicted dynamics），并证明了一个随机耦合定理（stochastic coupling theorem），用于严格界定 SGD 轨迹与其预测轨迹之间的偏差。 我们进一步推导出尖锐度平衡态缺口（equilibrium sharpness gap）的闭式表达式： $$ ΔS = \frac{ηβσ_{\boldsymbol{u}}^{2}}{4α}, $$ 其中，$α$ 表示渐进式尖锐化速率（progressive sharpening rate），$β$ 表示自稳定强度（self-stabilization strength），而 $σ_{\boldsymbol{u}}^{2}$ 则是梯度噪声在 Hessian 主特征向量方向上的投影方差。 该公式表明：批量越小，所得解越平坦；当批量等于整个数据集时，该公式自然退化为全批量 GD 情形。

论文试图解释为何在mini-batch SGD中，神经网络训练时的Hessian最大特征值（sharpness）稳定在低于2/η的水平，且随批大小减小而进一步降低——这一‘sharpness抑制’现象缺乏理论解释，而此前Damian et al. (2023) 的Edge of Stability理论仅适用于全批量GD。

提出‘随机自稳定’（stochastic self-stabilization）新机制：梯度噪声在主Hessian特征方向引入方差，增强三阶loss结构所驱动的sharpness衰减力，从而将动态平衡点下拉至2/η以下；并严格建模为带噪声的投影梯度流，在移动PGD轨迹上建立随机耦合定理，导出sharpness gap的闭式解ΔS = ηβσ_u²/(4α)。

首次为SGD下的Edge of Stability抑制现象提供可验证、参数化的理论解释；推导出sharpness gap与学习率η、噪声方差σ_u²、批大小（通过σ_u²隐式控制）的定量关系；无需额外假设（如PL条件或强凸性），适用于标准深度网络；实验验证gap随batch size减小而增大，与理论一致；代码已开源（见GitHub: stochastic-self-stab）；未来可拓展至自适应优化器、非各向同性噪声建模及泛化界连接。

Damian et al., 'Self-Stabilization Properties of Gradient Descent', NeurIPS 2023；Cohen et al., 'The Break-Even Point on Optimization Trajectories of Deep Neural Networks', ICLR 2021；Jin et al., 'How Does Batch Size Affect Generalization in SGD?', arXiv 2022；Wu et al., 'On the Dynamics of Stochastic Gradient Descent with Variable Batch Sizes', ICML 2023；Zhu et al., 'Anisotropic Noise Helps SGD to Escape Sharp Minima', NeurIPS 2022