[论文解读] Rethinking the limiting dynamics of SGD: modified loss, phase space oscillations, and anomalous diffusion

Daniel Kunin, Javier Sagastuy-Breña+5 more|arXiv (Cornell University)|Jul 19, 2021

Stochastic Gradient Optimization Techniques参考文献 34被引用 4

一句话总结

该论文揭示，随机梯度下降（SGD）在训练损失收敛后，仍通过参数空间中的异常扩散持续驱动深度神经网络，且行进距离与更新次数的幂律关系成比例。通过连续时间下的惯性朗之万模型和线性回归中的福克-普朗克分析，识别出由梯度噪声和海森矩阵决定的修正损失函数与概率流，这些是相空间振荡和非平凡扩散指数的关键驱动力。

ABSTRACT

In this work we explore the limiting dynamics of deep neural networks trained with stochastic gradient descent (SGD). We find empirically that long after performance has converged, networks continue to move through parameter space by a process of anomalous diffusion in which distance travelled grows as a power law in the number of gradient updates with a nontrivial exponent. We reveal an intricate interaction between the hyperparameters of optimization, the structure in the gradient noise, and the Hessian matrix at the end of training that explains this anomalous diffusion. To build this understanding, we first derive a continuous-time model for SGD with finite learning rates and batch sizes as an underdamped Langevin equation. We study this equation in the setting of linear regression, where we can derive exact, analytic expressions for the phase space dynamics of the parameters and their instantaneous velocities from initialization to stationarity. Using the Fokker-Planck equation, we show that the key ingredient driving these dynamics is not the original training loss, but rather the combination of a modified loss, which implicitly regularizes the velocity, and probability currents, which cause oscillations in phase space. We identify qualitative and quantitative predictions of this theory in the dynamics of a ResNet-18 model trained on ImageNet. Through the lens of statistical physics, we uncover a mechanistic origin for the anomalous limiting dynamics of deep neural networks trained with SGD.

研究动机与目标

• 理解训练损失稳定后，神经网络参数中持续存在的非平衡动力学。
• 识别SGD优化极限阶段中参数空间内异常扩散的机制。
• 揭示超参数、梯度噪声结构与海森矩阵如何共同塑造长期动力学。
• 建立一个捕捉速度依赖动力学与相空间振荡的有限小批量SGD连续时间模型。
• 在ImageNet上训练的ResNet-18中验证理论预测，建立统计物理与深度学习动力学之间的联系。

提出的方法

• 构建连续时间惯性朗之万方程，以建模有限学习率与小批量大小的SGD。
• 利用福克-普朗克方程，在线性回归中推导出参数与速度动力学的精确解析解。
• 提出一种修正损失函数，其在参数空间中隐式正则化速度，与原始训练损失不同。
• 通过相空间中的概率流分析，解释参数轨迹中持续振荡行为的成因。
• 表征梯度噪声结构、海森曲率与由此产生的非平凡扩散指数之间的相互作用。
• 通过分析在ImageNet上训练的ResNet-18的相空间动力学，验证理论预测。

实验结果

研究问题

• RQ1 在训练损失稳定后，是什么驱动了神经网络参数的持续非收敛运动？
• RQ2 梯度噪声、海森曲率与优化超参数之间的相互作用，如何塑造长期参数动力学？
• RQ3 速度正则化与概率流在参数-速度相空间中生成振荡行为方面起什么作用？
• RQ4 在多大程度上，参数空间中的异常扩散指数可由修正损失与福克-普朗克动力学解释？
• RQ5 在线性模型中推导出的理论框架，能否扩展以解释ResNet-18等深度非线性网络的动力学？

主要发现

• 收敛后，深度神经网络在参数空间中仍表现出异常扩散，行进距离与更新次数的幂律关系成比例。
• 扩散指数非平凡，由海森矩阵、梯度噪声结构与优化超参数之间的相互作用决定。
• 动力学的关键驱动力并非原始训练损失，而是隐式正则化参数空间中速度的修正损失。
• 相空间中的概率流引发持续振荡，阻止系统达到简单平衡状态。
• 基于惯性朗之万动力学的线性回归模型的理论预测，在ImageNet上训练的ResNet-18中得到定量验证。
• 极限动力学由速度正则化与非平衡流之间的平衡所主导，其根源在于统计物理原理。