[論文レビュー] 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最適化の制限段階におけるパラメータ空間における異常拡散の背後にあるメカニズムを特定すること。
• ハイパーパramータ、勾配ノイズ構造、ヘッセ行列が長期間にわたるダイナミクスをどのように統合的に形作るかを明らかにすること。
• 速度依存のダイナミクスと位相空間の振動を捉えることのできる有限バッチSGDの連続時間モデルを確立すること。
• ImageNetでトレーニングされたResNet-18において理論的予測を検証し、統計物理学と深層学習のダイナミクスを結びつけること。

提案手法

• 有限の学習率とバッチサイズのSGDをモデル化するため、連続時間の不釣り合いランジュバン方程式を定式化する。
• フォッカー＝プランク方程式を用いて、線形回帰におけるパラメータと速度のダイナミクスの正確な解析解を導出する。
• 元のトレーニング損失とは異なる、速度をパラメータ空間で暗黙的に正則化する修正された損失関数を導入する。
• 位相空間における確率流を分析し、パラメータ軌道の持続的で周期的な振動行動を説明する。
• 勾配ノイズ構造、ヘッセ曲率、およびそれらがもたらす非自明な拡散指数との相互作用を特徴付ける。
• ImageNetでトレーニングされたResNet-18の位相空間ダイナミクスを分析することで、理論的予測を検証する。

実験結果

リサーチクエスチョン

• RQ1 トレーニング損失が安定化した後も、ニューラルネットワークのパラメータが持続的に非収束的に移動する主な要因は何ですか？
• RQ2 勾配ノイズ、ヘッセ曲率、最適化ハイパーパramータの相互作用が、長期間にわたるパラメータダイナミクスをどのように形作るか？
• RQ3 速度正則化と確率流が、パラメータ-速度位相空間における周期的振動を生成する役割を果たすのはどのようなものか？
• RQ4 パラメータ空間における異常拡散指数は、修正された損失関数とフォッカー＝プランクダイナミクスによってどの程度説明可能か？
• RQ5 線形モデルで導出された理論枠組みは、ResNet-18のような深く非線形なネットワークのダイナミクスを説明するために拡張可能か？

主な発見

• 収束後も、深層ニューラルネットワークはパラメータ空間において異常拡散を示し、移動距離が更新回数のべき乗則に従う。
• 拡散指数は非自明であり、ヘッセ行列、勾配ノイズ構造、最適化ハイパーパラメータの相互作用によって決定される。
• ダイナミクスの主な駆動要因は元のトレーニング損失ではなく、パラメータ空間における速度を暗黙的に正則化する修正された損失関数である。
• 位相空間における確率流が持続的な振動を引き起こし、系が単純な平衡状態に到達することを防いでいる。
• 不釣り合いランジュバンダイナミクスを伴う線形回帰モデルからの理論的予測は、ImageNetでトレーニングされたResNet-18において定量的に検証された。
• 制限段階のダイナミクスは、速度正則化と非平衡的確率流のバランスによって支配されており、統計物理学の原則に基づいている。