ウォームアップは、機械学習および深層学習で使用されるトレーニング手法であり、一連の初期トレーニングステップにわたって学習率を小さな初期値から目標値へ徐々に増加させます。この手法は、ニューラルネットワーク、特に大規模言語モデルやトランスフォーマーなどの大規模モデルのトレーニング初期段階を安定させるために設計されています。開始時に大きな学習率を使用することを避けることで、ウォームアップは不安定な勾配を防ぎ、モデルがより良い最終解に収束するのを助けます。
この手法は、コサインアニーリングや線形減衰などの後続の減衰スケジュールと組み合わせて使用されることが多く、完全な学習率スケジュールを構成します。ウォームアップは現在、OpenAI、Anthropic、Google DeepMindによって開発されたものなど、多くの最先端モデルのトレーニングにおける標準的な要素となっています。これは通常、特定の数のウォームアップステップまたは学習率の増加方法を定義するウォームアップスケジュールを通じて実装されます。
動機と問題
トレーニングの初期段階では、モデルの重みは、XavierやHe初期化などの方法を使用してランダムに初期化されます。この時点で、特に深いネットワークでは、勾配が大きくノイズが多い可能性があります。高い学習率が即座に適用されると、オプティマイザーは大きく不規則なステップを踏む可能性があり、損失の発散やモデルが悪い局所最小値に落ち着く原因となります。この問題は、残差ネットワークやトランスフォーマーなど、層が多いモデルでは悪化し、勾配の大きさが層間で大きく変動する可能性があります。
ウォームアップは、モデルがまず損失ランドスケープ内の勾配がより安定して情報豊かな領域に移動できるようにすることで、この問題に対処します。小さな初期学習率は正則化の一形態として機能し、初期の重み構成を破壊する可能性のある急激な更新を防ぎます。これは、活性化の統計が安定する時間を必要とするレイヤー正規化やバッチ正規化などの技術を使用するモデルにとって特に重要です。
歴史的背景
ウォームアップの概念は新しいものではありませんが、ディープラーニングの台頭とともに顕著になりました。2010年代の深層ニューラルネットワークのトレーニングに関する初期の研究、例えばトロント大学やスタンフォードAIラボからのものは、学習率を徐々に増加させる利点に言及していました。しかし、2017年のGoogleの研究者（Jakob UszkoreitやLukasz Kaiserを含む）による論文「Attention Is All You Need」でトランスフォーマーアーキテクチャが導入された後、広く採用される慣行となりました。その論文は、トランスフォーマーモデルのトレーニングに4000ステップのウォームアップフェーズを明示的に推奨し、これが一般的なデフォルトとなりました。
それ以来、ウォームアップは多くのトレーニングフレームワークに組み込まれ、Hugging FaceのTransformersやPyTorchの学習率スケジューラーなどのライブラリではデフォルト設定となっています。モデルのスケーリングに伴い、大規模モデルはハイパーパラメータの選択により敏感であるため、その採用は成長しています。
ウォームアップの仕組み
ウォームアップは通常、ウォームアップステップ数（\(W\)と表記）によって定義されます。これらのステップ中、学習率は小さな値（多くの場合ゼロに近い）から目標学習率\(\eta_{max}\)まで線形または二次的に増加します。一般的な線形スケジュールは次のとおりです：
\[ \eta_t = \eta_{max} \cdot \frac{t}{W} \]
ここで、\(t = 1, 2, \ldots, W\)です。ウォームアップ期間後、学習率はコサイン減衰、指数減衰、ステップワイズ削減などの減衰スケジュールに従います。一部のスケジュールでは、特定のアーキテクチャに有益な平方根や三次の増加などの非線形ウォームアップを使用します。
\(W\)の選択は、モデルサイズとデータセットに依存します。小規模モデルでは、ウォームアップは数百ステップしか続かない場合がありますが、大規模言語モデルでは、数千から数万ステップに及ぶことがあります。例えば、数千億のパラメータを持つGPT-3のようなモデルのトレーニングでは、不安定性を避けるために数千ステップのウォームアップが必要でした。
利点と効果
ウォームアップはいくつかの具体的な利点を提供します。第一に、トレーニングの安定性を向上させ、損失の発散や勾配のNaN値のリスクを低減します。第二に、モデルが損失ランドスケープのより好ましい領域を探索できるため、最終的なパフォーマンスが向上することがよくあります。第三に、ウォームアップフェーズ後の収束を加速できる高いピーク学習率の使用を可能にします。
実証研究では、ウォームアップがAdamや他の適応オプティマイザーと組み合わせた場合に特に効果的であることが示されています。Adamは勾配の履歴に基づいてパラメータごとの学習率を維持しますが、初期ステップではこれらの推定値がノイズが多いです。ウォームアップは、オプティマイザーの分散推定が安定するまでステップサイズを制限することでこれを軽減します。
さらに、ウォームアップは勾配クリッピングや重み初期化などの他の技術と相互作用します。例えば、残差ネットワークを使用する場合、ウォームアップは残差接続が初期エラーを増幅しないようにするのに役立ちます。
バリエーションと拡張
ウォームアップのいくつかのバリエーションが開発されています。一般的なバリアントの1つは「段階的ウォームアップ」で、学習率が固定されたステップ数にわたって線形に増加します。もう1つは「指数ウォームアップ」で、指数関数的に学習率を増加させますが、これはあまり一般的ではありません。一部のスケジュールでは、「ウォームアップ後の減衰」アプローチを使用し、ウォームアップ後にコサイン減衰でゼロに減少させます。これは元のトランスフォーマー論文で「Noam」スケジュールとして普及しました。
もう1つの拡張は「確率的ウォームアップ」で、ウォームアップステップが各トレーニング実行ごとにランダムにサンプリングされ、ハイパーパラメータ正則化の一形態として機能します。場合によっては、ウォームアップはトランスフォーマーの埋め込み層など特定の層やパラメータにのみ適用され、他の場合では完全なウォームアップが使用されます。
ウォームアップはまた、モデルがまず簡単な例でトレーニングされてから難しい例に進むカリキュラム学習とも関連しています。両方の手法は初期トレーニングを安定させることを目的としていますが、異なる軸で動作します：ウォームアップは学習率を調整し、カリキュラム学習はデータ分布を調整します。
実装の実践
実際には、ウォームアップを実装するには、初期学習率（多くの場合ピークレートの小さな割合、例えば0.1や0.01に設定）、ピーク学習率、ウォームアップステップ数の3つのハイパーパラメータを設定する必要があります。ピーク学習率は通常ハイパーパラメータチューニングを通じて決定され、ウォームアップステップはしばしば総トレーニングステップのパーセンテージ（1%から10%など）に設定されます。
例えば、機械翻訳用のトランスフォーマーモデルのトレーニングでは、一般的な構成はピーク学習率0.0001と4000ウォームアップステップです。対照的に、数兆トークンを持つ大規模言語モデルでは、ウォームアップは総計500,000ステップのうち2000ステップ続く可能性があります。
PyTorchやTensorFlowなどのフレームワークは、ウォームアップをサポートする組み込みスケジューラーを提供しています。例えば、PyTorchのget_cosine_schedule_with_warmup関数は、線形ウォームアップとその後のコサイン減衰を備えたスケジュールを作成します。同様に、Hugging Faceのget_schedulerはウォームアップ比率を指定できます。
他の技術との関係
ウォームアップは、しばしば他の安定化技術と併用されます。勾配クリッピングは勾配の大きさを制限し、ウォームアップなしでも爆発を防ぐことができますが、組み合わせにより堅牢性が向上します。レイヤー正規化やバッチ正規化は活性化を安定させるのに役立ちますが、初期の大きな勾配の問題に直接対処するわけではありません。
ウォームアップはまた、ドロップアウトや他の正則化手法を補完します。低い学習率から始めることで、モデルは初期段階でノイズに過剰適合する可能性が低くなり、ドロップアウトが後でより効果的になります。
さらに、ウォームアップは、画像セグメンテーション用のU-Netやさまざまなシーケンス・ツー・シーケンスモデルを含む、多くの現代アーキテクチャのトレーニングレシピの重要なコンポーネントです。また、事前トレーニングされたモデルのファインチューニングでも使用され、短いウォームアップ（例えば100ステップ）が破滅的な忘却なしにモデルを新しいタスクに適応させるのに役立ちます。
課題と限界
その利点にもかかわらず、ウォームアップはチューニングが必要な追加のハイパーパラメータを導入します。ウォームアップステップが少なすぎると不安定性につながる可能性があり、多すぎると収束が遅くなる可能性があります。最適なウォームアップ期間はモデルとデータセットによって異なり、普遍的なルールはありません。
一部の研究では、固定の低学習率の使用や勾配統計に基づいてレートを調整する適応的手法の採用など、ウォームアップの代替案が探求されています。しかし、これらのアプローチは実際にはウォームアップを完全に置き換えておらず、ウォームアップはシンプルで効果的なベースラインとして残っています。
もう1つの限界は、ウォームアップが主に確率的勾配降下法とその変種向けに設計されていることです。二次法などの他の最適化アルゴリズムでは、利点がそれほど顕著ではない可能性がありますが、一部のケースではまだ使用されています。
将来の方向性
モデルが成長し続けるにつれて、ウォームアップは活発な研究分野のままです。最近の研究では、検証損失を使用してウォームアップから減衰への遷移をトリガーするなど、ウォームアップ期間を設定する自動化方法が調査されています。また、ウォームアップの理論的基盤を理解することへの関心もあり、一部の研究では損失ランドスケープの曲率との関連が示されています。
生成AIと大規模トレーニングの文脈では、ウォームアップは標準的な慣行として残る可能性があります。OpenAI、Anthropic、Google DeepMindなどの企業はすべて、ウォームアップフェーズを含むトレーニング詳細を公開しており、多くのオープンソースのトレーニングスクリプトでデフォルトとなっています。新しいアーキテクチャが登場するにつれて、ウォームアップは適応され洗練され続けるでしょう。
要約すると、ウォームアップはニューラルネットワークのトレーニングの安定性とパフォーマンスを向上させるシンプルでありながら強力な手法です。学習率を徐々に増加させることで、モデルが最適化の初期の混沌としたフェーズをナビゲートできるようにし、最終的な結果の向上につながります。この分野での広範な採用は、現代の機械学習におけるその重要性を強調しています。