[論文レビュー] The large learning rate phase of deep learning: the catapult mechanism

Aitor Lewkowycz, Yasaman Bahri+3 more|arXiv (Cornell University)|Mar 4, 2020

Stochastic Gradient Optimization Techniques参考文献 37被引用数 60

ひとこと要約

本論文は、勾配降下法における3つの学習率フェーズ（lazy, catapult, divergent）を導入し、catapultダイナミクスがより平坦な極小点へと導く解ける有限幅モデルを提示し、実際の深層ネットワークにおける予測を検証し、最適な性能は大きな学習率のcatapultフェーズにしばしば存在するという実証を示す。

ABSTRACT

The choice of initial learning rate can have a profound effect on the performance of deep networks. We present a class of neural networks with solvable training dynamics, and confirm their predictions empirically in practical deep learning settings. The networks exhibit sharply distinct behaviors at small and large learning rates. The two regimes are separated by a phase transition. In the small learning rate phase, training can be understood using the existing theory of infinitely wide neural networks. At large learning rates the model captures qualitatively distinct phenomena, including the convergence of gradient descent dynamics to flatter minima. One key prediction of our model is a narrow range of large, stable learning rates. We find good agreement between our model's predictions and training dynamics in realistic deep learning settings. Furthermore, we find that the optimal performance in such settings is often found in the large learning rate phase. We believe our results shed light on characteristics of models trained at different learning rates. In particular, they fill a gap between existing wide neural network theory, and the nonlinear, large learning rate, training dynamics relevant to practice.

研究の動機と目的

• 初期学習率が深層ネットワークの訓練と一般化に著しく影響することを動機づけ、特徴づける。
• 三つの異なる学習率レジームを予測する扱いやすい有限幅モデルを導入する。
• アーキテクチャを横断する現実的な深層ネットワークでフェーズ予測を検証し、理論と実践を橋渡しする。
• 最適な性能がしばしば大きな学習率（catapult）フェーズで生じることを示す。
• 平坦性と一般化のダイナミクスをSGDノイズから切り離し、学習率主導の効果に焦点を当てる。

提案手法

• 大きいが有限幅の1層隠れ線形ネットワークとMSE損失に対する厳密な勾配降下更新を導出する。
• NTKの最大固有値を曲率の代理として用い、3つの学習率レジームを識別・分析する。
• 入力をd次元、訓練サンプル数をmとする完全モデルへ解析を拡張し、類似の更新ダイナミクスを導出する。
• 全結合、畳み込み、残差ネットワークに渡る実証実験を行い、フェーズ予測を検証する。
• アーキテクチャ依存定数c_actを用いて実践的な最大学習率を推定し、実験に基づくReLUは約12。

実験結果

リサーチクエスチョン

• RQ1 広いが有限のネットワークにおいて、初期学習率を変化させたときに勾配降下のダイナミカルフェーズは何か？
• RQ2 学習率は訓練中のカーネル曲率、特にNTKのトップ固有値にどのような影響を与えるか？
• RQ3 大きな学習率を安定化させてより平坦な極小点へ収束させ、これが一般化にどう影響するか？
• RQ4 理論的なフェーズ予測は現実的なアーキテクチャとSGD設定で成り立つか？
• RQ5 アーキテクチャ、非線形性と最大安定学習率との実証的な関係は何か？

主な発見

• 3つの学習率フェーズが存在する：lazy（eta < 2/λ_0）、catapult（2/λ_0 < eta < eta_max）、divergent（eta > eta_max）。
• catapultフェーズでは、初期損失が急速な曲率の低下とともに上昇し、その後lazy相より平坦な極小点へ収束する。
• 最大安定学習率は概ね eta_max = c_act / lambda_0、c_actは非線形性に依存する（理論上≈4、実践ではReLUで約12）。
• CNN、ResNet、全結合ネットワークにおける実証結果はフェーズ境界と一致し、catapultフェーズでピーク性能を示す。
• 最適な性能はしばしば大きな学習率のcatapultフェーズで発生し、アーキテクチャや訓練予算を問わず一致している。
• catapultの後、モデルの挙動はほぼ一定のカーネルを伴う線形ダイナミクスに近づき、線形風のダイナミクスの回復を示唆している。