arXiv:2606.05326math.OCcs.AI2026-06被引 1

提出自由能模型,解析大学习率下神经网络的震荡训练动态

Gradient descent at the Edge of Stability: free energy model and kinetic description of the two-layer network

论文配图:Gradient descent at the Edge of Stability: free energy model and kinetic description of the two-layer network
图 1 · 摘自论文原文
  • 构建连续时间有效模型,跟踪平均轨迹与快速振荡的协方差
  • 发现有效自由能可预测训练中的波动峰值,准确捕捉振荡包络
  • 推导均值场动力学方程,揭示其为广义梯度流,适合研究不稳训练

我们研究了在边缘稳定性区间内梯度下降的动力学行为,该区域学习率足够大,导致损失和尖锐度产生持续振荡。本文提出一个连续时间有效模型,追踪平均轨迹及其快速振荡的时间平均协方差。分析表明,在此类不稳定区域中,应监控一种结合原始风险函数与曲率相关“熵项”的有效自由能。该模型即使在振荡包络动态与平均权重演化时间尺度相近时,仍可准确追踪振荡峰值,适用于某些神经网络架构的训练过程。针对宽两层神经网络在非衰减振荡下的稳定优化,我们推导出均值场极限,得到描述权重及其波动联合分布的新动力学方程。该方程可解释为宏观自由能的Wasserstein-2梯度流。最后,我们在矩阵分解和深度学习任务(CIFAR-10)上提供了数值证据,验证了模型对振荡包络的捕捉能力及有效自由能的预测性能。

原文摘要 · Abstract (English)

We study the dynamics of gradient descent in the Edge of Stability regime, where the learning rate is large enough to induce persistent oscillations in the loss and the sharpness. We propose a continuous-time effective model that tracks the evolution of the average trajectory coupled with the time-averaged covariance of its fast oscillations. Our analysis reveals that the natural quantity to monitor in such unstable regimes is an effective free energy, which combines the original risk functional with a curvature-related "entropic" term. Our model allows us to track the envelope of the oscillations even in situations where its dynamics evolve on similar timescales as the averaged weights. Otherwise stated, we can track the spikes that occur during the training of some neural network architectures. For wide two-layer neural networks optimized under stable non-vanishing oscillations, we derive a mean-field limit that results in a novel kinetic equation describing the joint distribution of weights and their fluctuations. We show that this equation can be interpreted as a Wasserstein-2 gradient flow of a macroscopic free energy. Finally, we provide numerical evidence on matrix factorization and deep learning tasks (CIFAR-10) to demonstrate the model's accuracy in capturing the envelope of the oscillations and the predictive power of the effective free energy.

梯度下降自由能边缘稳定性均值场

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。