[論文レビュー] Optimal learning rate schedules in high-dimensional non-convex optimization problems

Stéphane d’Ascoli, Maria Refinetti+1 more|arXiv (Cornell University)|Feb 9, 2022

Stochastic Gradient Optimization Techniques被引用数 4

ひとこと要約

本稿は、時間に依存する学習率 $\eta(t) = \eta_0 t^{-\beta}$ を用いたランジュヴィンダイナミクスを用いて、高次元非凸最適化における最適な学習率スケジューリングの最初の解析的研究を提供する。純粋に非凸な地形では $\beta < 1$ が鞍点を回避することで最適化を加速するが、信号復元問題では、探索段階で定数 $\eta$ を使用し、収束段階で $\beta = 1$ の減衰をとる二段階戦略が最適な性能を達成する。

ABSTRACT

Learning rate schedules are ubiquitously used to speed up and improve optimisation. Many different policies have been introduced on an empirical basis, and theoretical analyses have been developed for convex settings. However, in many realistic problems the loss-landscape is high-dimensional and non convex -- a case for which results are scarce. In this paper we present a first analytical study of the role of learning rate scheduling in this setting, focusing on Langevin optimization with a learning rate decaying as $η(t)=t^{-β}$. We begin by considering models where the loss is a Gaussian random function on the $N$-dimensional sphere ($N ightarrow \infty$), featuring an extensive number of critical points. We find that to speed up optimization without getting stuck in saddles, one must choose a decay rate $β&lt;1$, contrary to convex setups where $β=1$ is generally optimal. We then add to the problem a signal to be recovered. In this setting, the dynamics decompose into two phases: an \emph{exploration} phase where the dynamics navigates through rough parts of the landscape, followed by a \emph{convergence} phase where the signal is detected and the dynamics enter a convex basin. In this case, it is optimal to keep a large learning rate during the exploration phase to escape the non-convex region as quickly as possible, then use the convex criterion $β=1$ to converge rapidly to the solution. Finally, we demonstrate that our conclusions hold in a common regression task involving neural networks.

研究の動機と目的

• 現代の機械学習で一般的な高次元非凸最適化における学習率スケジューリングの理論的理解の欠如に対処すること。
• 時間に依存する学習率減衰スケジュール $\eta(t) = \eta_0 t^{-\beta}$ が非凸地形における最適化ダイナミクスに与える影響を分析すること。
• 純粋に非凸および信号復元設定の両方において、損失を最小化する最適な減衰指数 $\beta$ を特定すること。
• SGDを用いたニューラルネットワーク回帰タスクを通じて、実用的なディープラーニング設定で解析的結果を検証すること。

提案手法

• 球面制約と時間依存学習率 $\eta(t) = \eta_0 t^{-\beta}$ を用いた最適化問題を定式化し、温度 $T$ を用いて SGD のノイズをモデル化する。
• ガウス確率的損失関数を伴う $N$ 次元球面上での $N \to \infty$ の極限において、球面シェリントン・キルパトリック（SK）モデル（$p=2$）とスパイクド行列テンソル（SMT）モデル（$p>2$）の二つのモデルを分析する。
• 統計物理学的手法、特にカック＝ライス法とダイナミカルリプラー理論を用いて相関関数と応答関数を計算し、有効ダイナミクスを導出する。
• 球面制約を強制するラグランジュ乗数 $z_\infty$ の漸近的値を導出し、非凸領域における閾値損失を決定する。
• 初期の探索段階（粗い非凸領域）と、信号が存在する際に生じる凸領域への収束段階という二段階のダイナミクス的挙動を同定する。
• ミニバッチ SGD（$B=1$）を用いたティーチャー・スタディオのニューラルネットワーク回帰タスクで結果を検証し、最適スケジューリング戦略を実証的に確認する。

実験結果

リサーチクエスチョン

• RQ1 純粋に非凸で高次元な地形では、最適な学習率減衰指数 $\beta$ は何か？ これは最適化の加速に寄与する。
• RQ2 復元可能な信号が存在する場合、非凸最適化における最適スケジューリング戦略はどのように変化するか？
• RQ3 実際のディープラーニング設定（ニューラルネットワークを用いた）で、最適 $\beta$ の解析的予測は検証可能か？
• RQ4 信号復元問題において、探索から収束へのダイナミクス的相転移を示すのは何であるか？
• RQ5 ノイズ（温度 $T$）と学習率減衰の相互作用は、収束速度と最終損失にどのように影響するか？

主な発見

• 純粋に非凸な地形（例：$p=2$ のSKモデル）では、最適な減衰率は $\beta = 1/2$ であり、凸設定で用いられる $\beta = 1$ よりも著しく遅い。
• 高次の非凸性（$p > 2$）では、最適な $\beta$ は1未満であることが判明し、これは鞍点を効率的に脱出するためにはより遅い減衰が必要であることを示唆する。
• 信号が存在する場合、最適戦略は二段階スケジューリングである：探索段階では定数学習率を用い、非凸領域からの迅速な脱出を図り、収束段階では $\beta = 1$ の減衰をとる。
• 探索から収束へのクロスオーバー時間 $t_{\text{cross}}$ は、定数学習率下で損失がプラトー化する「特化」遷移の終了時刻と一致する。
• $B=1$ のティーチャー・スタディオのニューラルネットワーク回帰タスクにおいて、学習率を早めに減衰させると損失が高くなるが、$t_{\text{cross}}$ の後に遅らせることで、損失がゼロに収束することが可能になる。
• 非凸領域における解析的閾値損失は、低温では温度 $T$ に対して線形に増加し、導出された $z_\infty$ と $\ell_{\text{th}}$ の式の妥当性を確認する。