[論文レビュー] 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）が、性能収束後も、最適化ハイパーパrameterと勾配ノイズに起因する非平衡確率流を伴い、損失関数を変更することで、パラメータ空間における異常拡散を引き続き駆動していることを明らかにする。連続時間の不釣れランジュバンモデルを用いて、線形回帰における正確なダイナミクスを導出し、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.

研究の動機と目的

• トレーニング収束後も続く、非平衡的パラメータ空間ダイナミクスが異常拡散として現れる理由を説明すること。
• 最適化ハイパーパrameter、勾配ノイズ構造、収束時のヘッセ行列との関連を通じて、これらのダイナミクスの機構的起源を同定すること。
• 有限学習率とバッチサイズに起因するノイズを捉える連続時間の不釣れランジュバン方程式としてのSGDモデルを構築すること。
• 線形回帰におけるパラメータダイナミクスの定常分布を解析的に導出し、元のトレーニング損失ではなく、詳細なバランスが破れた修正損失関数に従うことを示すこと。
• 実世界のモデル（例：ImageNet上のResNet-18）における理論的予測を検証し、実験的異常拡散と定量的に一致することを示すこと。

提案手法

• 有限学習率とバッチサイズに起因するノイズを考慮した連続時間の不釣れランジュバン方程式を導出し、最適化プロセスを確率的力学系としてモデル化すること。
• 線形回帰におけるダイナミクスをフォッカー・プランク方程式を用いて解析し、定常分布が元のトレーニング損失ではなく、修正された損失関数上のギブス測度であることを示すこと。
• 修正された損失関数が詳細なバランスを破ることを証明し、位相空間（パラメータと速度）における非ゼロの確率流が、振動的で拡散的でない運動を駆動することを示すこと。
• データ共分散行列の固有基底において、位相空間ダイナミクスを減衰する調和振動子の和として表現し、モーメントの正確な解析的解を得ること。
• フォッカー・プランク方程式を用いて定常分布を導出し、系が非平衡定常状態を示し、持続的な運動を示すことを示すこと。
• ImageNetで微調整されたResNet-18モデルに対して理論的予測を検証し、パラメータ軌道を追跡し、全距離のべき乗則フィッティングにより異常拡散指数を抽出すること。

実験結果

リサーチクエスチョン

• RQ1 トレーニング収束後もパラメータ空間で持続的かつ非平衡的な運動が生じる要因は何か？
• RQ2 最適化ハイパーパrameter（学習率、バッチサイズ、モーメンタム）が、限界ダイナミクスにおける異常拡散指数および瞬間速度に定量的にどのように影響を与えるか？
• RQ3 ヘッセ行列と勾配ノイズ構造は、SGDの位相空間ダイナミクスにどのような役割を果たすか？
• RQ4 パラメータダイナミクスの定常分布は元のトレーニング損失に対応するのか、それとも、詳細なバランスが破れた暗黙の正則化項によって支配されるのか？
• RQ5 不釣れランジュバンダイナミクスに基づく理論的枠組みは、ImageNet上のResNet-18のような実モデルにおける実験的異常拡散行動を定量的に予測できるか？

主な発見

• SGDの限界的ダイナミクスは、元のトレーニング損失ではなく、速度を暗黙的に正則化する修正損失関数に支配されており、非平衡定常状態を生じる。
• 詳細なバランスが破れたことによる非ゼロの位相空間確率流が、振動的で拡散的でない運動を駆動しており、収束後もパラメータが動く理由を説明する。
• 線形回帰では、データ共分散行列の固有基底において、位相空間ダイナミクスが減衰する調和振動子の和として解析的に解ける。
• ImageNetでトレーニングされたResNet-18における異常拡散指数および瞬間速度は、理論モデルによって正確に予測され、実測値と完全に一致した。
• 極めて非等方的な元の損失関数の形状であっても、ResNet-18の限界的ダイナミクスはパラメータ空間で等方的であることが示され、モデルの予測と整合的である。
• フィッティングされた異常拡散指数は、軌道長に従って増加し、時間とともに支配的項が顕在化するべき乗則の和であると考えられ、理論的枠組みと整合的である。