[論文レビュー] Optimal Linear Decay Learning Rate Schedules and Further Refinements

Aaron Defazio, Ashok Cutkosky+2 more|arXiv (Cornell University)|Oct 11, 2023

Stochastic Gradient Optimization Techniques被引用数 5

ひとこと要約

この論文は、深層学習最適化のための理論的裏付けがあり、問題に適応する学習率スケジューリングフレームワークを提案する。過去の実行から得られた勾配ノルムの系列を精錬することで、従来はヒューリスティックにしか使われてこなかったウォームアップと線形減衰を併せ持つスケジューリングを導出する。10種類の多様な深層学習タスクにおいて、線形減衰がコサイン減衰や他のスケジューリングを上回ることを実証し、さらに精錬によって性能が向上することを示している。

ABSTRACT

Learning rate schedules used in practice bear little resemblance to those recommended by theory. We close much of this theory/practice gap, and as a consequence are able to derive new problem-adaptive learning rate schedules. Our main technical contribution is a refined analysis of learning rate schedules for a wide class of optimization algorithms (including SGD). When considering only worst-case analysis, our theory predicts that the optimal choice is the linear decay schedule where the step-size is set proportional to 1 - t/T, where t is the current iteration and T is the total number of steps. To go beyond this worst-case analysis, we use the observed gradient norms to derive schedules refined for any particular task. These refined schedules exhibit learning rate warm-up and rapid learning rate annealing near the end of training. Ours is the first systematic approach to automatically yield both of these properties. We perform the most comprehensive evaluation of learning rate schedules to date, evaluating across 10 diverse deep learning problems, a series of LLMs, and a suite of logistic regression problems. We validate that overall, the linear-decay schedule outperforms all commonly used default schedules including cosine annealing. Our adaptive schedule refinement method gives further improvements.

研究の動機と目的

• 深層学習における理論的学習率スケジューリングと実践的選択との間の継続的なギャップを解消すること。
• 最終反復収束の理論的分析と現実の最適化実践を調和させること。
• 観測された勾配統計を用いて、学習率スケジューリングを体系的かつ問題に適応させる方法を開発すること。
• 最悪ケースの仮定下で、SGDの最終反復に対して線形減衰が理論的に最適であることを示すこと。
• 理論的にウォームアップと線形減衰の挙動が自然に生じるか、それらは単なる経験的ヒューリスティックに過ぎないのかを解消すること。

提案手法

• 全期待勾配ノルムの系列を用いて、最終反復の関数値に対する新たな理論的バウンドを提案する。
• 各ステップサイズを、将来の勾配ノルムの逆二乗の積として計算する学習率スケジューリングを設計：$\eta_t = w_t \sum_{p=t+1}^T w_p $ ここで $ w_t = \hat{G}_t^{-2} $。
• ノイズを平滑化し、スケジューリング精錬の耐性を向上させるために、原始的な勾配ノルムにメディアンフィルタを適用する。
• 事前トレーニング走行で勾配ノルムの系列を収集し、その後のトレーニング用に精錬された適応的スケジューリングを生成する。
• SGDに限らず、劣化レジストが消える任意の最適化アルゴリズムへ一般化し、Adamや類似手法への応用を可能にする。
• 最終スケジューリングを正規化して、最大学習率が標準ベースラインと一致するように保証する。

Figure 1: Example gradient norm sequences (top row) and the resulting refined schedules given by Algorithm 1 (bottom row). Black dashed line at $y=1$ shown for reference. Percentage of runtime on the x-axis.

実験結果

リサーチクエスチョン

• RQ1 なぜ標準的な理論的学習率スケジューリングは、深層学習における実践的性能と一致しないのか？
• RQ2 最悪ケースの仮定下で、SGDの最終反復に対して理論的に最適な学習率スケジューリングは何か？
• RQ3 過去の走行から得た勾配ノルムの系列を用いて、一般化性能と収束性を向上させる適応的スケジューリングを導出できるか？
• RQ4 理論的にウォームアップと線形減衰の挙動が自然に生じるのか、それともそれらは単なる経験的ヒューリスティックに過ぎないのか？
• RQ5 提案された精錬手法は、視覚、自然言語処理、レコメンデーションシステムを含む多様な深層学習タスクにおいて、標準スケジューリング（例：コサイン減衰）と比較してどのように異なるか？

主な発見

• 従来は経験的にのみ使われていた線形減衰学習率スケジューリングは、最悪ケースの仮定下でSGDの最終反復に対して理論的に最適であることが判明した。
• 提案された精錬手法により、ウォームアップと急速な減衰を併せ持つスケジューリングが生成され、標準的な理論的スケジューリングには存在しない。
• 非適応的スケジューリングの中で、線形減衰は10の評価済み深層学習タスクすべてでコサイン減衰を同等または上回る性能を示した。
• 精錬されたスケジューリングは、ResNet-50、ViT、GPT、RoBERTa、DLRMモデルにおいて、標準スケジューリングよりも一貫して一般化性能を向上させた。
• 本手法は一般化可能である：Adamを含む、劣化レジストが消える任意の最適化アルゴリズムに適用可能であり、理論的裏付けのある適応的スケジューリングを可能にする。
• 実験結果により、観測された勾配ノルムを用いて精錬された理論的導出スケジューリングが、実際の性能向上を顕著に示すことが裏付けられた。

Figure 2: Gradient Norm sequences and the resulting Refined schedules, generated using an initial linear decay schedule with warmup for the initial run. Log scale views are inset for scale reference. Fitted polynomial decay schedules are overlayed in blue.