[論文レビュー] 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

ひとこと要約

本稿では、正則な拡散係数を備えたSDEの新クラス、確率的修正流れ（SMF）を導入する。これは、学習率が小さい領域における確率的勾配降下法（SGD）をモデル化するもので、従来の確率的修正方程式（SME）が抱える不規則性の問題を回避する。さらに、無限幅・小学習率スケーリング領域において、SMFを用いて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.

研究の動機と目的

• 過パラメータ化された深層学習における解析を妨げる、従来の確率的修正方程式（SME）における拡散係数の不規則性を解消すること。
• SGDの多点統計と一致する制限SDEモデルを構築し、従来の連続極限を改善すること。
• 無限幅・小学習率領域におけるSGDのフラクチュエーティング平均場極限を、分布に依存する確率的修正流れを用いて確立すること。
• アーキテクチャ的およびデータ依存的変動に対して頑健な、SGDのダイナミクスを統一的に分析する枠組みを提供すること。

提案手法

• 個々の損失勾配から直接得られる拡散係数を備えた、円形ブラウン運動によって駆動されるSDEとして、確率的修正流れ（SMF）を提案する。これは、退化した共分散行列Σの平方根を必要としない。
• SMFのSDEを、dXₜ^η(x) = -∇(R(Xₜ^η(x)) + η/4 |∇R(Xₜ^η(x))|²) dt + √η ∫_Θ G(Xₜ^η(x),θ) W(dθ,dt) として定義する。ここで、G(x,θ) = ∇R̃(x,θ) - ∇R(x) である。
• マーティングル問題の定式化を用いて、SMFがSMEと同一の統計的挙動を再現するが、より優れた正則性を有することを示す。
• 解の分布に依存するSMFを導入し、無限幅極限における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(η²)の誤差内に収束し、2次精度が確立される。
• 無限幅・小学習率領域において、SGDの経験的測度は分布に依存するSMFに収束し、フラクチュエーティング平均場ダイナミクスを捉える。
• 経験的測度から極限SMFへの収束速度は、期待値においてO(η²)であり、初期測度のモーメントおよび標本サイズMに依存する追加の誤差バウンドが存在する。
• 適切なモーメントおよび正則性条件の下で、初期測度が有限の経験的測度によって近似されている場合でも、収束速度はO(η²)である。Mおよびηに明示的な依存関係が得られる。