[論文レビュー] Training (Overparametrized) Neural Networks in Near-Linear Time

Jan van den Brand, Binghui Peng+2 more|arXiv (Cornell University)|Jun 20, 2020

Stochastic Gradient Optimization Techniques参考文献 89被引用数 23

ひとこと要約

本稿では、ガウス＝ニュートン法の新規再定式化を用いて、過パラメータ化されたReLUニューラルネットワークのトレーニングのための近線形時間アルゴリズムを提示する。この再定式化はℓ²正則化回帰問題として扱われ、Fast-Johnson-Lindenstrauss（Fast-JL）次元削減を用いてグラム行列を前処理することで高速化される。1反復あたりの計算コストはeO(mnd + n³)、全実行時間はeO((mnd + n³) log(1/ǫ))で、訓練損失をǫまで低下させる。これは従来の2次最適化手法と比較して2乗の高速化を達成し、深層学習最適化における高度なランダム化線形代数の実用可能性を示している。

ABSTRACT

The slow convergence rate and pathological curvature issues of first-order gradient methods for training deep neural networks, initiated an ongoing effort for developing faster $\mathit{second}$-$\mathit{order}$ optimization algorithms beyond SGD, without compromising the generalization error. Despite their remarkable convergence rate ($\mathit{independent}$ of the training batch size $n$), second-order algorithms incur a daunting slowdown in the $\mathit{cost}$ $\mathit{per}$ $\mathit{iteration}$ (inverting the Hessian matrix of the loss function), which renders them impractical. Very recently, this computational overhead was mitigated by the works of [ZMG19,CGH+19}, yielding an $O(mn^2)$-time second-order algorithm for training two-layer overparametrized neural networks of polynomial width $m$. We show how to speed up the algorithm of [CGH+19], achieving an $ ilde{O}(mn)$-time backpropagation algorithm for training (mildly overparametrized) ReLU networks, which is near-linear in the dimension ($mn$) of the full gradient (Jacobian) matrix. The centerpiece of our algorithm is to reformulate the Gauss-Newton iteration as an $\ell_2$-regression problem, and then use a Fast-JL type dimension reduction to $\mathit{precondition}$ the underlying Gram matrix in time independent of $M$, allowing to find a sufficiently good approximate solution via $\mathit{first}$-$\mathit{order}$ conjugate gradient. Our result provides a proof-of-concept that advanced machinery from randomized linear algebra -- which led to recent breakthroughs in $\mathit{convex}$ $\mathit{optimization}$ (ERM, LPs, Regression) -- can be carried over to the realm of deep learning as well.

研究の動機と目的

• 深層学習における2次最適化の高い計算コスト、特に過パラメータ化ネットワークに対する従来のガウス＝ニュートン法の1反復あたりO(mn²)のコストを軽減すること。
• 滑らかでない活性化関数を用いるネットワークに、2次最適化を拡張すること。これは、滑らかでない活性化関数を用いるネットワークの方がより複雑かつ現実的であるため。
• 通常の2次最適化手法でHessian行列の逆行列計算がボトルネックとなっている問題を克服し、全勾配の次元（mn）に関して近線形時間でのトレーニングを達成すること。
• これまで凸最適化で用いられてきた高度なランダム化線形代数技術が、非凸な深層学習の文脈にも効果的に適用可能であることを示すこと。

提案手法

• 共役勾配法を用いて効率的に解けるように、ジャコビアン行列上でℓ²正則化回帰問題にガウス＝ニュートン更新を再定式化する。
• Fast-Johnson-Lindenstrauss（Fast-JL）型の次元削減を用いて、グラム行列J_t J_t^Tを前処理し、サイズを縮小しながら解の品質を保持する。
• 1次最適化の共役勾配法を用いて、元の行列サイズMに依存しない時間で前処理済み回帰問題を解く。
• ニューラル接線カーネル（NTK）理論を活用して、過パラメータ化ネットワークの局所線形化を正当化し、カーネル回帰に還元可能であることを示す。
• 集中不等式を用いてジャコビアン近似誤差と回帰部分問題の解の誤差を抑え、収束を保証する。
• 1反復あたりの計算コストがeO(mnd + n³)であるように、回帰ソルバーをバックプロパゲーションフレームワークに統合。このコストはジャコビアン評価と回帰問題の解法が支配的である。

実験結果

リサーチクエスチョン

• RQ1 過パラメータ化されたReLUネットワークに対する2次最適化を、1反復あたりのコストをO(mn²)未満に抑えることで実用化可能か？
• RQ2 Fast-JL次元削減を用いてガウス＝ニュートンのヘッセ行列近似を前処理することで、収束保証を維持できるか？
• RQ3 2層ReLUネットワークに対して、全勾配の次元（mn）に関して近線形時間でのトレーニングが達成可能か？
• RQ4 凸最適化で成功したランダム化線形代数の手法が、非凸な深層学習トレーニングに適応可能か？
• RQ5 提案手法は、SGDと比較して著しく高速な収束を達成する一方で、一般化性能を維持できるか？

主な発見

• 本手法は、[CGH+19]のO(mn²)法と比較して2乗の高速化を達成し、1反復あたりのコストをeO(mnd + n³)に低減した。
• 本手法は、過パラメータ化の程度がやや強い条件下（m = Ω(max{λ⁻⁴n⁴, λ⁻²n²d log(n/δ)})) において、全勾配次元（mn）に関して近線形時間で動作する、ReLUネットワークに対する最初の2次最適化アルゴリズムである。
• 本手法は1反復あたり∥ft+1 − y∥² ≤ ½∥ft − y∥²を満たし、目標損失への線形収束を保証する。
• 訓練損失をǫ未満に低下させるまでの全実行時間はeO((mnd + n³) log(1/ǫ))であり、高速行列乗算を用いることでeO((mnd + n^ω) log(1/ǫ))にまで短縮可能である。
• 本手法は重み更新量を有界に保つため、ネットワークの重みが初期値から著しく逸脱することを防ぎ、一般化性能の維持に寄与する。
• 本手法は凸最適化にも成功裏に適用され、γ-strongly convexかつβ-smoothな関数（HessianがL-Lipschitz）に対して、ニュートン法の実行時間をeO((nd log(κ) + d³) log(1/ǫ))に改善した。