[論文レビュー] Rigorous dynamical mean field theory for stochastic gradient descent methods

Cédric Gerbelot, Emanuele Troiani+3 more|arXiv (Cornell University)|Oct 12, 2022

Stochastic Gradient Optimization Techniques被引用数 6

ひとこと要約

本稿は、高次元設定における離散時間確率的勾配降下法（SGD）および関連する一次順序最適化手法について、厳密な動的平均場理論（DMFT）方程式を確立する。SGDの漸近的ダイナミクスが、一般の共分散行列を有するガウス分布データ上で正確に離散DMFT方程式によって記述されることを証明しており、自己整合的カーネルを通じて記憶効果を捉え、磁化およびノルムのダイナミクスによる一般化性能の正確な予測を可能にする。

ABSTRACT

We prove closed-form equations for the exact high-dimensional asymptotics of a family of first order gradient-based methods, learning an estimator (e.g. M-estimator, shallow neural network, ...) from observations on Gaussian data with empirical risk minimization. This includes widely used algorithms such as stochastic gradient descent (SGD) or Nesterov acceleration. The obtained equations match those resulting from the discretization of dynamical mean-field theory (DMFT) equations from statistical physics when applied to gradient flow. Our proof method allows us to give an explicit description of how memory kernels build up in the effective dynamics, and to include non-separable update functions, allowing datasets with non-identity covariance matrices. Finally, we provide numerical implementations of the equations for SGD with generic extensive batch-size and with constant learning rates.

研究の動機と目的

• 離散時間確率的勾配降下法（SGD）および関連最適化アルゴリズムの動的平均場理論（DMFT）方程式の厳密な導出を提供すること。
• 従来のDMFT結果を、i.i.d.なサブガウス分布データに限らず、任意の共分散行列を有する一般のガウス分布データへと拡張すること。
• 有効な確率的積分微分方程式における明示的な記憶カーネルを導出することで、ダイナミクスにおける記憶効果を扱うこと。
• 高次元極限における磁化およびノルムのダイナミクスを通じて、一般化性能（例えば、推定器と真値とのコサイン類似度など）の正確な予測を可能にすること。
• 広範なバッチサイズおよび一定の学習率を有するSGDに対して、離散DMFTフレームワークを数値的に検証すること。

提案手法

• 反復的ガウス条件付き推定と摂動理論を用いて、SGDにおける重み更新の無限次元極限を分析することで、離散時間DMFT方程式を導出する。
• 推定器の重みベクトルおよび残差誤差の有効ダイナミクスを記述する、自己整合的システムの確率的積分微分方程式を導入する。
• 損失関数の過去の勾配に関する関数的微分を用いて記憶カーネルを計算し、ダイナミクスにおける非マルコフ的効果を捉える。
• 参照軌道の一次摂動展開を用いて有効ダイナミクスを導出し、項が高次元極限でガウス過程に収束することを示す。
• 磁化（真値とのコサイン類似度）および重みノルムの閉形式方程式を導出し、一般化性能の予測を可能にする。
• 一定の学習率および広範なバッチサイズを有するSGDに対して、導出されたDMFT方程式を数値的に実装し、漸近的予測の妥当性を検証する。

実験結果

リサーチクエスチョン

• RQ1 非i.i.d.データ、特に非単位共分散行列を有する場合においても、離散時間SGDに対して厳密なDMFT方程式を導出可能か？
• RQ2 勾配の履歴に起因する記憶効果が、高次元におけるSGDの漸近的ダイナミクスにどのように現れるか？
• RQ3 離散DMFT方程式が、高次元極限におけるSGDの一般化性能をどの程度正確に記述できるか？
• RQ4 DMFTフレームワークは、i.i.d.サブガウス設計を超えて、非可分な更新関数および非i.i.d.データへと拡張可能か？
• RQ5 離散時間SGDのダイナミクスと連続時間勾配フローとの間の正確な関係は、DMFTの文脈でどのように特定されるか？

主な発見

• 本稿は、非i.i.d.ガウス分布データでさえも、離散時間DMFT方程式が高次元極限においてSGDダイナミクスの正確な漸近的記述を提供することを証明している。
• 記憶カーネルは損失関数の関数的微分から自然に導かれ、過去の勾配が現在の更新に与える影響を明示的に計算することで捉える。
• 一般化性能は、磁化 $ m^t = ext{corr}(m{w}^*, m{w}^t) $ によって完全に特徴付けられ、損失導関数およびノイズ共分散を含む自己整合的方程式に従って時間発展する。
• 有効ダイナミクスは、記憶カーネル $ R_g(t,t') $、共分散 $ C_g(t,t') $ を持つノイズ過程 $ u^t $、および自己正則化項 $ ilde{ u}^t $ を含む確率的積分微分方程式系によって記述される。
• DMFT方程式の数値的実装は、高次元設定における推定器のノルムおよび真値とのコサイン類似度の時間発展を正確に予測する。
• 本フレームワークは、SGD、ネステロフ加速、モーメンタムベースのアルゴリズムを含む、一次順序手法の広いクラスに、経験的リスク最小化の下で適用可能である。