[论文解读] Dynamic of Stochastic Gradient Descent with State-Dependent Noise

Qi Meng, Shiqi Gong+3 more|arXiv (Cornell University)|Jun 24, 2020

Stochastic Gradient Optimization Techniques参考文献 48被引用 8

一句话总结

本文提出一种带状态相关扩散的幂律随机微分方程来建模随机梯度下降（SGD），表明由此产生的平稳分布具有重尾特性，与实际观测结果一致。该文证明从尖锐极小值逃逸的时间为多项式时间，远快于恒定扩散模型中的指数时间，从而解释了SGD对泛化性能更优的平坦极小值的偏好。

ABSTRACT

Stochastic gradient descent (SGD) and its variants are mainstream methods to train deep neural networks. Since neural networks are non-convex, more and more works study the dynamic behavior of SGD and the impact to its generalization, especially the escaping efficiency from local minima. However, these works take the over-simplified assumption that the covariance of the noise in SGD is (or can be upper bounded by) constant, although it is actually state-dependent. In this work, we conduct a formal study on the dynamic behavior of SGD with state-dependent noise. Specifically, we show that the covariance of the noise of SGD in the local region of the local minima is a quadratic function of the state. Thus, we propose a novel power-law dynamic with state-dependent diffusion to approximate the dynamic of SGD. We prove that, power-law dynamic can escape from sharp minima exponentially faster than flat minima, while the previous dynamics can only escape sharp minima polynomially faster than flat minima. Our experiments well verified our theoretical results. Inspired by our theory, we propose to add additional state-dependent noise into (large-batch) SGD to further improve its generalization ability. Experiments verify that our method is effective.

研究动机与目标

• 为解决现有SGD动力学中假设梯度噪声与状态无关的局限性，该局限性无法解释实际中观测到的重尾参数分布。
• 构建一个理论基础扎实的动力学模型，以捕捉SGD噪声在局部极小值附近的真正状态相关特性。
• 分析该新动力学从局部极小值逃逸的效率，并将其与泛化性能关联。
• 提供一个PAC-Bayes泛化界，表明相较于恒定扩散动力学，该模型实现了更优的泛化性能。
• 通过在深度神经网络上的实验验证理论结果，表明参数分布与所提出的幂律模型拟合良好。

提出的方法

• 推导出局部极小值附近二次盆地中梯度噪声的协方差，将其表示为模型参数的二次函数，从而为状态相关扩散提供依据。
• 提出一种扩散系数与状态平方成正比的随机微分方程（SDE），称为幂律动力学。
• 证明该幂律动力学的平稳分布为具有尾指数κ的幂律κ分布，与实际中观测到的重尾参数分布一致。
• 应用随机扰动理论，计算在幂律动力学下从局部极小值逃逸的平均时间。
• 为幂律动力学推导PAC-Bayes泛化界，表明其泛化性能优于恒定扩散模型。
• 通过将训练得到的深度网络（如LeNet-5、ResNet-18、VGG-16等）的参数分布拟合到幂律κ分布，并与SGD的逃逸动力学进行比较，验证该模型。

实验结果

研究问题

• RQ1 SGD中的状态相关噪声如何影响模型参数的平稳分布？
• RQ2 具有状态相关扩散的幂律动力学能否在不假设无限方差的前提下，解释实际观测到的训练参数的重尾分布？
• RQ3 在幂律动力学下，从局部极小值逃逸的时间与恒定扩散动力学相比如何？
• RQ4 在所提出的动力学中，逃逸效率与泛化性能之间存在何种关系？
• RQ5 该幂律动力学在多大程度上能复现实际SGD在训练深度神经网络时的行为？

主要发现

• 所提出的幂律动力学的平稳分布为具有尾指数κ的重尾幂律κ分布，与训练深度网络中观测到的参数分布一致。
• 在幂律动力学下，从局部极小值逃逸的平均时间与势垒高度呈多项式关系，远快于恒定扩散动力学中的指数关系。
• 该幂律动力学比平坦极小值更高效地逃离尖锐极小值，表明SGD通过此机制避免尖锐极小值。
• 与恒定扩散动力学相比，幂律动力学的PAC-Bayes泛化界更紧，表明其具有更优的泛化性能。
• 实验表明，SGD训练模型（在MNIST、CIFAR-10、ImageNet上）的参数分布可被幂律κ分布良好拟合。
• 逃逸动力学的比较显示，幂律动力学比恒定扩散动力学更贴近SGD的实际行为。