[論文レビュー] A mean-field theory of lazy training in two-layer neural nets: entropic regularization and controlled McKean-Vlasov dynamics.

Belinda Tzen, Maxim Raginsky|arXiv (Cornell University)|Feb 5, 2020

Stochastic Gradient Optimization Techniques参考文献 27被引用数 11

ひとこと要約

本稿は、重みの確率測度上の自由エネルギー最小化問題として学習問題を定式化することで、2層ニューラルネットワークにおけるラージトレーニングの平均場理論を構築する。この定式化により、近似誤差とガウス事前分布からのKLダイバージェンスのバランスをとる。制御されたMcKean–Vlasovダイナミクスを導入し、エントロピー正則化最適輸送（シュレーディンガー・ブリッジ）問題を解く。また、ラージトレーニング領域における確率的勾配降下法（SGD）が、定量的な$L^2$誤差バインディングを伴って、このダイナミクスをグリーディに近似することを示す。

ABSTRACT

We consider the problem of universal approximation of functions by two-layer neural nets with random weights that are nearly in the sense of Kullback-Leibler divergence. This problem is motivated by recent works on lazy training, where the weight updates generated by stochastic gradient descent do not move appreciably from the i.i.d. Gaussian initialization. We first consider the mean-field limit, where the finite population of neurons in the hidden layer is replaced by a continual ensemble, and show that our problem can be phrased as global minimization of a free-energy functional on the space of probability measures over the weights. This functional trades off the $L^2$ approximation risk against the KL divergence with respect to a centered Gaussian prior. We characterize the unique global minimizer and then construct a controlled nonlinear dynamics in the space of probability measures over weights that solves a McKean--Vlasov optimal control problem. This control problem is closely related to the Schrodinger bridge (or entropic optimal transport) problem, and its value is proportional to the minimum of the free energy. Finally, we show that SGD in the lazy training regime (which can be ensured by jointly tuning the variance of the Gaussian prior and the entropic regularization parameter) serves as a greedy approximation to the optimal McKean--Vlasov distributional dynamics and provide quantitative guarantees on the $L^2$ approximation error.

研究の動機と目的

• 2層ニューラルネットワークにおけるラージトレーニングを平均場極限によって形式化し、離散的ニューロン集団を重み上の連続的確率測度に置き換える。
• $L^2$近似リスクと中心ガウス事前分布からのKLダイバージェンスを組み合わせた自由エネルギー関数のグローバル最小化子として最適重み分布を特徴付ける。
• 重み空間における最適分布の時間発展を実現する制御されたMcKean–Vlasovダイナミクスを構築する。
• ラージトレーニング領域におけるSGDとこの最適ダイナミクスとの理論的関係を確立し、定量的な$L^2$近似誤差保証を提供する。

提案手法

• 重み上の確率測度の空間における自由エネルギー関数のグローバル最小化問題として、平均場極限における学習問題を定式化する。
• 自由エネルギー関数を、$L^2$近似リスクと中心ガウス事前分布からのKullback-Leiblerダイバージェンス項の和として定義する。
• 変分法とギブズ測度の性質を用いて、自由エネルギーの唯一のグローバル最小化子を特徴付ける。
• 最小化子に近づく重み分布の時間発展を実現する制御された非線形拡散過程（McKean–Vlasovダイナミクス）を構築し、最適制御問題として定式化する。
• 制御問題をシュレーディンガー・ブリッジ問題に結びつけ、その値が最小自由エネルギーに一致することを示す。
• ラージトレーニング領域におけるSGDが、この制御されたダイナミクスをグリーディに近似することを示し、その$ L^2 $誤差バインディングは事前分布の分散とエントロピー正則化パラメータに依存する。

実験結果

リサーチクエスチョン

• RQ1 2層ニューラルネットワークにおけるラージトレーニングの平均場極限下での最適重み分布は何か？
• RQ2 重み空間における重み分布の時間発展は、確率測度空間上での制御されたMcKean–Vlasov過程としてどのようにモデル化できるか？
• RQ3 最適制御問題とエントロピー正則化最適輸送（シュレーディンガー・ブリッジ）との関係は何か？
• RQ4 確率的勾配降下法（SGD）は、ラージトレーニング領域において、最適McKean–Vlasovダイナミクスをどの程度近似するか？
• RQ5 制御されたハイパーパrameterチューニング下で、SGDの定量的$L^2$近似誤差バインディングはどのように導出できるか？

主な発見

• 自由エネルギー関数の唯一のグローバル最小化子は、近似精度と事前分布への忠実度のバランスを取るギブズ測度である。
• 最適重み分布は、シュレーディンガー・ブリッジ問題を解く制御されたMcKean–Vlasovダイナミクスによって達成される。
• 最適制御問題の値は最小自由エネルギーに一致し、最適解の変分的特徴付けを提供する。
• ラージトレーニング領域におけるSGDは、最適McKean–Vlasovダイナミクスのグリーディ近似として機能する。
• SGDの$L^2$近似誤差は定量的にバインディングされ、ガウス事前分布の分散とエントロピー正則化パラメータのチューニングによって制御可能である。