[论文解读] Stochastic Modified Flows, Mean-Field Limits and Dynamics of Stochastic Gradient Descent

Benjamin Gess, Sebastian Kassing+1 more|arXiv (Cornell University)|Feb 14, 2023

Markov Chains and Monte Carlo Methods被引用 4

一句话总结

本文提出了随机修正流（SMFs），这是一种新型的随机微分方程（SDE），在小学习率条件下对随机梯度下降（SGD）进行建模，其扩散系数具有规则性，避免了先前随机修正方程（SMEs）中存在的不规则性问题。此外，本文进一步建立了依赖于分布的SMFs，作为无限宽、小学习率尺度下SGD的波动性均场极限，能够匹配多点统计特性，并实现对过参数化网络的严格分析。

ABSTRACT

We propose new limiting dynamics for stochastic gradient descent in the small learning rate regime called stochastic modified flows. These SDEs are driven by a cylindrical Brownian motion and improve the so-called stochastic modified equations by having regular diffusion coefficients and by matching the multi-point statistics. As a second contribution, we introduce distribution dependent stochastic modified flows which we prove to describe the fluctuating limiting dynamics of stochastic gradient descent in the small learning rate - infinite width scaling regime.

研究动机与目标

• 解决经典随机修正方程（SMEs）中扩散系数不规则的问题，该问题阻碍了在过参数化深度学习中的分析。
• 开发一种能匹配SGD多点统计特性的极限SDE模型，改进先前的连续极限方法。
• 在无限宽、小学习率的设定下，利用依赖于分布的随机修正流，建立SGD的波动性均场极限。
• 提供一个统一的框架，用于分析SGD的动力学行为，且对网络架构和数据依赖性变化具有鲁棒性。

提出的方法

• 提出随机修正流（SMFs）作为由圆柱形布朗运动驱动的SDE，其扩散系数直接来源于单个损失函数梯度，而非退化的协方差矩阵Σ。
• 将SMF SDE定义为：dXₜ^η(x) = -∇(R(Xₜ^η(x)) + η/4 |∇R(Xₜ^η(x))|²) dt + √η ∫_Θ G(Xₜ^η(x),θ) W(dθ,dt)，其中G(x,θ) = ∇R̃(x,θ) - ∇R(x)。
• 采用鞅问题的表述方法，证明SMFs在统计行为上与SMEs一致，但具有更好的正则性。
• 引入依赖于分布的SMFs，用于在无限宽极限下建模SGD的经验测度动力学，其系数依赖于解的分布律。
• 应用Wasserstein-2距离和混沌传播论证，比较SMF解的经验测度与有限SGD动力学之间的差异。
• 结合SDE理论、Malliavin微积分和混沌传播方法，证明了在期望和分布意义下的收敛速率。

实验结果

研究问题

• RQ1 能否为SGD构造一种基于SDE的新极限动力学模型，避免经典随机修正方程中扩散系数的不规则性？
• RQ2 所提出的随机修正流模型是否保留了SGD的多点统计特性，从而实现更精确的连续逼近？
• RQ3 在无限宽、小学习率的设定下，SGD的波动性均场极限是什么？如何通过依赖于分布的SDE来描述它？
• RQ4 能否以分布和矩误差的形式量化有限SGD动力学向所提出的极限SDE的收敛性？

主要发现

• 所提出的随机修正流（SMF）SDE具有由单个梯度差值导出的规则扩散系数，避免了对协方差矩阵Σ的病态平方根运算。
• SMF与SGD的多点统计特性一致，并满足与经典SME相同的鞅问题，确保了统计一致性。
• 对于任意f ∈ C⁴_b(ℝᵈ)，SMF与SGD动力学下f的期望值在有限时间范围内相差O(η²)，建立了二阶精度。
• 在无限宽、小学习率的设定下，SGD的经验测度收敛到一个依赖于分布的SMF，该SMF捕捉了波动性均场动力学。
• 经验测度向极限SMF的收敛速率为O(η²)（在期望下），并额外包含依赖于初始测度矩和样本数M的误差界。
• 在适当的矩和正则性条件下，即使初始测度由有限经验测度近似，收敛速率仍为O(η²)，且显式依赖于M和η。