[论文解读] 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被引用 4

一句话总结

该论文揭示，由于损失函数的修改以及由优化超参数和梯度噪声引起的非平衡概率流，随机梯度下降（SGD）在性能收敛后仍会通过参数空间中的异常扩散持续驱动深度神经网络。通过连续时间下的欠阻尼朗之万模型，作者在线性回归中推导出精确动力学，并在ImageNet上的ResNet-18中验证了预测结果，显示理论与实验在异常扩散指数和瞬时速度方面具有精确的定量一致性。

ABSTRACT

In this work we explore the limiting dynamics of deep neural networks trained with stochastic gradient descent (SGD). As observed previously, 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的连续时间模型，将其表示为欠阻尼朗之万方程，以捕捉有限学习率和小批量大小的影响。
• 在线性回归中解析推导参数动力学的平稳分布，并证明其对应于一个打破细致平衡的修改后损失函数。
• 在真实世界模型（如ImageNet上的ResNet-18）上验证理论预测，展示与经验异常扩散的定量一致性。

提出的方法

• 推导一个连续时间下的欠阻尼朗之万方程，以建模具有有限学习率和由小批量大小引起的噪声的SGD，将优化过程视为随机动力系统。
• 在线性回归中使用福克-普朗克方程分析动力学，表明平稳分布是修改后损失函数上的吉布斯测度，而非原始训练损失。
• 证明修改后的损失打破了细致平衡，导致非零概率流，从而在相空间（参数与速度）中引发振荡运动。
• 在数据协方差矩阵的特征基下，将相空间动力学表示为阻尼谐振子的叠加，从而在解析上求解各阶矩。
• 利用福克-普朗克方程推导平稳分布，并证明系统表现出非平衡稳态行为，具有持续运动。
• 在ImageNet上微调的ResNet-18模型上验证理论预测，追踪参数轨迹并拟合幂律以提取异常扩散指数。

实验结果

研究问题

• RQ1 在训练收敛后，是什么驱动了神经网络参数在参数空间中的持续非平衡运动？
• RQ2 优化超参数——学习率、小批量大小和动量——如何定量影响极限动力学中的异常扩散指数和瞬时速度？
• RQ3 海森矩阵和梯度噪声结构在塑造SGD相空间动力学中起什么作用？
• RQ4 参数动力学的平稳分布是否对应于原始训练损失，还是由一个修改后的隐式正则化项所控制？
• RQ5 基于欠阻尼朗之万动力学的理论框架能否定量预测真实模型（如ImageNet上的ResNet-18）中的经验异常扩散行为？

主要发现

• SGD的极限动力学并非由原始训练损失决定，而是由一个隐式正则化速度的修改后损失函数控制，导致非平衡稳态。
• 相空间中非零的概率流——源于细致平衡的破坏——驱动振荡性、非扩散性运动，解释了收敛后持续的参数移动。
• 在线性回归中，相空间动力学可在数据协方差矩阵的特征基下被精确解析为阻尼谐振子的叠加。
• 在ImageNet上训练的ResNet-18中，异常扩散指数和瞬时速度被理论模型精确预测，与实测结果完全一致。
• 尽管原始损失景观高度各向异性，ResNet-18的极限动力学在参数空间中表现出各向同性行为，与模型预测一致。
• 拟合的异常扩散指数随轨迹长度增加，表明其为幂律之和，且主导项随时间逐渐显现，与理论框架一致。