[論文レビュー] High-dimensional limit theorems for SGD: Effective dynamics and critical scaling

Gérard Ben Arous, Reza Gheissari+1 more|arXiv (Cornell University)|Jun 8, 2022

Markov Chains and Monte Carlo Methods被引用数 7

ひとこと要約

本稿は、定数ステップサイズを用いた確率的勾配降下法（SGD）の高次元極限における極限定理を確立し、SGDの要約統計量が、初期化およびステップサイズのスケーリングに応じて常微分方程式（ODE）または確率微分方程式（SDE）に支配される有効な力学に収束することを示している。重要なスケーリング領域が特定され、勾配フローが新たな項によって補正され、非自明な相図と正の確率で最適でない収束を生じるが、過パラメータ化によってこれを緩和できる。

ABSTRACT

We study the scaling limits of stochastic gradient descent (SGD) with constant step-size in the high-dimensional regime. We prove limit theorems for the trajectories of summary statistics (i.e., finite-dimensional functions) of SGD as the dimension goes to infinity. Our approach allows one to choose the summary statistics that are tracked, the initialization, and the step-size. It yields both ballistic (ODE) and diffusive (SDE) limits, with the limit depending dramatically on the former choices. We show a critical scaling regime for the step-size, below which the effective ballistic dynamics matches gradient flow for the population loss, but at which, a new correction term appears which changes the phase diagram. About the fixed points of this effective dynamics, the corresponding diffusive limits can be quite complex and even degenerate. We demonstrate our approach on popular examples including estimation for spiked matrix and tensor models and classification via two-layer networks for binary and XOR-type Gaussian mixture models. These examples exhibit surprising phenomena including multimodal timescales to convergence as well as convergence to sub-optimal solutions with probability bounded away from zero from random (e.g., Gaussian) initializations. At the same time, we demonstrate the benefit of overparametrization by showing that the latter probability goes to zero as the second layer width grows.

研究の動機と目的

• 定数学習率を用いた高次元設定におけるSGDのスケーリング極限を統一的に扱うフレームワークを構築すること。
• 初期化、ステップサイズのスケーリング、パラメータ領域の選択が要約統計量の有効な力学に与える影響を特定すること。
• 勾配フローが新たな項によって補正され、収束行動が変化する臨界スケーリング領域を同定すること。
• 過パラメータ化が最適でない解への収束に与える影響を分析すること。
• スパイク行列／テンソルモデルおよびガウス混合分類のための2層ネットワークにおいて理論を実証すること。

提案手法

• 次元が発散する極限において、次元が有限の要約統計量のSGDの極限定理を、やや弱い正則性条件の下で導出する。
• 初期化およびステップサイズのスケーリングに応じて、力学的段階ではODEへの収束、拡散的段階ではSDEへの収束を確立する。
• 要約統計量の軌道に対する関数的中心極限定理を用いて、有効な力学を導出する。
• ステップサイズが $ O(1/ ext{dimension}) $ のスケーリングとなる臨界スケーリング領域を特定し、ODE極限に補正項が現れることを示す。
• 損失、重みの振幅、真の値との相関など、特定の統計量の追跡が可能である。
• 高次元設定におけるガウス近似およびモーメントバウンドを用い、ReLUやシグモイド活性化関数に対しても分析を展開する。

実験結果

リサーチクエスチョン

• RQ1 定数ステップサイズを用いた高次元極限において、SGDの要約統計量の軌道はどのように収束するか？
• RQ2 有効な力学が力学的（ODE）か拡散的（SDE）かを決定するのは何か？
• RQ3 ステップサイズのスケーリングが有効な力学および収束行動に与える役割は何か？
• RQ4 非凸問題において、確率的初期化が最適でない解への収束に与える影響は何か？
• RQ5 過パラメータ化によって、最適でない解への収束確率を完全に消去できるか？

主な発見

• 要約統計量の有効な力学は、初期化およびステップサイズのスケーリングに応じてODEまたはSDEに収束するが、臨界スケーリング領域ではODE極限に新たな補正項が現れる。
• 臨界スケーリング領域では、ODE極限は母集団損失の勾配フローと一致するが、補正項が加わり、相図が変化し、最適でない固定点への収束が生じる可能性がある。
• XOR型のガウス混合モデルにおける2層ネットワークでは、確率的初期化のもとでSGDは、ゼロでない下限を持つ確率で最適でない解へ収束する。
• 2番目の層の過パラメータ化により、最適でない解への収束確率は低下し、2番目の層の幅が増加するにつれてゼロに近づく。
• 本手法はスパイク行列およびテンソルモデルへも適用可能であり、収束までのマルチタイムスケールと複雑な拡散的極限を明らかにする。
• 分析により、拡散的極限が退化することがあり、収束行動は要約統計量および初期化の選択に強く依存することが示された。