解释小批量训练为何更平坦,揭示噪声如何让模型自动变稳定。
SGD at the Edge of Stability: The Stochastic Sharpness Gap
- 用随机自稳定机制解释梯度噪声如何抑制尖锐度
- 推导出尖锐度差距的闭式公式,与批次大小相关
- 适用于研究训练稳定性、优化器设计的读者
全批量梯度下降(GD)在步长η下,海森矩阵最大特征值——尖锐度S(θ)——会上升至2/η并保持稳定,称为稳定边缘现象。已有研究指出这是由损失函数三阶结构驱动的自稳定机制,并隐含地遵循约束S(θ)≤2/η的投影梯度下降(PGD)。对于小批量随机梯度下降(SGD),尖锐度稳定在2/η以下,且批次越小,差距越大,但此前缺乏理论解释。本文提出随机自稳定框架,揭示梯度噪声通过注入沿主海森特征向量的振荡动力学方差,增强了减少尖锐度的三次项作用力,使平衡点低于2/η。基于类似方法,我们定义相对于移动投影梯度轨迹的随机预测动态,并证明了随机耦合定理,控制了SGD与预测之间的偏差。推导出尖锐度差距的闭式表达式:ΔS = ηβσ²ᵘ/(4α),其中α为渐进尖锐化率,β为自稳定强度,σ²ᵘ为梯度噪声在主特征向量上的投影方差。该公式表明更小批次导致更平坦解,且当批次等于全数据集时退化为GD。
原文摘要 · Abstract (English)
When training neural networks with full-batch gradient descent (GD) and step size $η$, the largest eigenvalue of the Hessian -- the sharpness $S(\boldsymbolθ)$ -- rises to $2/η$ and hovers there, a phenomenon termed the Edge of Stability (EoS). \citet{damian2023selfstab} showed that this behavior is explained by a self-stabilization mechanism driven by third-order structure of the loss, and that GD implicitly follows projected gradient descent (PGD) on the constraint $ S(\boldsymbolθ)\leq 2/η$. For mini-batch stochastic gradient descent (SGD), the sharpness stabilizes below $2/η$, with the gap widening as the batch size decreases; yet no theoretical explanation exists for this suppression. We introduce stochastic self-stabilization, extending the self-stabilization framework to SGD. Our key insight is that gradient noise injects variance into the oscillatory dynamics along the top Hessian eigenvector, strengthening the cubic sharpness-reducing force and shifting the equilibrium below $2/η$. Following the approach of \citet{damian2023selfstab}, we define stochastic predicted dynamics relative to a moving projected gradient descent trajectory and prove a stochastic coupling theorem that bounds the deviation of SGD from these predictions. We derive a closed-form equilibrium sharpness gap: $ΔS = ηβσ_{\boldsymbol{u}}^{2}/(4α)$, where $α$ is the progressive sharpening rate, $β$ is the self-stabilization strength, and $σ_{ \boldsymbol{u}}^{2}$ is the gradient noise variance projected onto the top eigenvector. This formula predicts that smaller batch sizes yield flatter solutions and recovers GD when the batch equals the full dataset.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。