[論文レビュー] Function approximation by neural nets in the mean-field regime: Entropic regularization and controlled McKean-Vlasov dynamics

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

Stochastic Gradient Optimization Techniques被引用数 4

ひとこと要約

この論文は、$L^2$ 近似リスクとガウス事前分布からのKLダイバージェンスのバランスを取る自由エネルギー最小化問題として関数近似を定式化することにより、2層ニューラルネットワークにおけるラージトレーニングの平均場理論を確立する。エントロピー正則化を介した制御された McKean–Vlasov 動的システムを導入し、SGD が最適な分布的動的システムを定量的な $L^2$ 误差境界で近似できることを示している。

ABSTRACT

We consider the problem of function approximation by two-layer neural nets with random weights that are "nearly Gaussian" in the sense of Kullback-Leibler divergence. Our setting is the mean-field limit, where the finite population of neurons in the hidden layer is replaced by a continuous ensemble. We show that the problem can be phrased as global minimization of a free energy functional on the space of (finite-length) paths over probability measures on the weights. This functional trades off the $L^2$ approximation risk of the terminal measure against the KL divergence of the path with respect to an isotropic Brownian motion prior. We characterize the unique global minimizer and examine the dynamics in the space of probability measures over weights that can achieve it. In particular, we show that the optimal path-space measure corresponds to the Föllmer drift, the solution to a McKean-Vlasov optimal control problem closely related to the classic Schrödinger bridge problem. While the Föllmer drift cannot in general be obtained in closed form, thus limiting its potential algorithmic utility, we illustrate the viability of the mean-field Langevin diffusion as a finite-time approximation under various conditions on entropic regularization. Specifically, we show that it closely tracks the Föllmer drift when the regularization is such that the minimizing density is log-concave.

研究の動機と目的

• 2層ニューラルネットワークにおけるラージトレーニングの枠組み下での関数近似の平均場理論の構築。
• ガウス事前分布からのKLダイバージェンスと $L^2$ リスクを組み合わせた自由エネルギー関数の最小化として最適な重み分布を特定すること。
• エントロピー正則化を通じて最適な分布的動的システムを実現する制御された McKean–Vlasov ストキャスティック過程の構築。
• SGD が最適な平均場動的システムへのグリーディ近似としての $L^2$ 近似誤差に定量的な保証を与えること。
• 最適制御問題がシュレーディンガー・ブリッジ問題と Föllmer ドリフトにどのように関連するかを示し、平均場極限の動的解釈を可能にすること。

提案手法

• 自由エネルギー関数のグローバル最小化問題として平均場近似問題を定式化：$\mathcal{F}_{\tau,\beta}(\mu) = \mathbb{E}_{w \sim \mu}[\|f(w,\cdot) - f^*\|_{L^2}^2] + \beta \cdot D_{\mathrm{KL}}(\mu \| \gamma_\tau)$、ここで $\gamma_\tau$ は中心が原点のガウス事前分布である。
• 正則性条件の下で自由エネルギー関数のグローバル最小解の存在と一意性を証明し、$D_{\mathrm{KL}}(\mu^* \| \gamma_\tau)$ および $\mathcal{W}_2(\mu^*, \gamma_\tau)$ に対する明示的な境界を導出する。
• 終端コストに $\beta$ と $\tau$ を組み込んだ確率的最適制御問題を解く制御された McKean–Vlasov SDE を構築し、そのドリフトが $\varphi(w,t)$ であることを示す。
• 最適制御則が Föllmer ドリフトに一致することを特定し、これはエントロピー最適輸送（シュレーディンガー・ブリッジ）問題の解に相当する。
• 重み分布と価値関数の時間発展を記述する、前向き（Fokker–Planck）および後ろ向き（Hamilton–Jacobi–Bellman）の連立PDEを導出する。
• 適切に調整された学習率 $\eta$、ネットワーク幅 $N$、正則化パラメータ $\beta$ を用いた SGD が、高確率で最適な McKean–Vlasov 動的システムを、誤差 $O\left(\frac{\kappa}{N} + \kappa\sqrt{\frac{1}{N}\log\frac{N}{\delta}} + \kappa\beta + \kappa(1+\sqrt{\eta}\beta)(1+\beta)e^{\kappa\beta}\left(\beta + \sqrt{\tau d\log\frac{Nd}{\delta\eta}}\right)\right)$ で近似できることを示す。

実験結果

リサーチクエスチョン

• RQ1 2層ニューラルネットワークにおけるラージトレーニングの枠組みを、自由エネルギー最小化フレームワークとしてどのように形式化できるか？
• RQ2 近似誤差とガウス事前分布からの逸脱のトレードオフを最小化する最適な重み分布の構造は何か？
• RQ3 最適な分布的動的システムは、制御された McKean–Vlasov プロセスとして実現可能か？また、エントロピー最適輸送とはどのように関連するか？
• RQ4 確率的勾配降下法（SGD）は、ラージトレーニングの枠組みにおいて、最適な平均場動的システムをどの程度近似できるか？
• RQ5 有限標本および高確率的誤差境界として、SGD が最適な分布的動的システムを近似する際の誤差はどの程度か？

主な発見

• 自由エネルギー関数には一意なグローバル最小解 $\mu^*$ が存在し、KLダイバージェンス $D_{\mathrm{KL}}(\mu^* \| \gamma_\tau)$ と $L^2$ ワイエルシュトラス距離 $\mathcal{W}_2(\mu^*, \gamma_\tau)$ はそれぞれ $O(\beta)$ および $O(\sqrt{\tau D})$ で有界である。
• McKean–Vlasov 動的システムの最適制御則は Föllmer ドリフトであり、これはシュレーディンガー・ブリッジ問題の解に一致し、終端分布が最小解 $\mu^*$ に一致することを保証する。
• 制御された McKean–Vlasov 動的システムは、Fokker–Planck と Hamilton–Jacobi–Bellman の連立PDEによって記述され、価値関数は最小自由エネルギーを符号化している。
• $N$ 個のニューロンを用いた SGD は、時間 $t \in [0,1]$ の間、高確率 $1-\delta$ で最適な平均場動的システムを近似でき、$L^2$ 誤差境界が $O\left(\frac{\kappa}{N} + \kappa\sqrt{\frac{1}{N}\log\frac{N}{\delta}} + \kappa\beta + \kappa(1+\sqrt{\eta}\beta)(1+\beta)e^{\kappa\beta}\left(\beta + \sqrt{\tau d\log\frac{Nd}{\delta\eta}}\right)\right)$ に収束する。
• 近似誤差は $N$ および $\delta$ に対して有利にスケーリングされ、ガウス事前分布の分散 $\tau$ と正則化パラメータ $\beta$ の調整により小さくできる。
• 本分析により、ラージトレーニング、エントロピー正則化、McKean–Vlasov 動的システムの間の厳密な接続が確立され、ニューラル接線カーネルと平均場極限に関する先行研究が拡張されている。