技術要約：ホライゾンフリーな学習率スケジュールの活用による、Anytime（いつでも利用可能）な事前学習

問題提起

大規模言語モデル（LLM）は、総学習ホライゾン（処理される総トークン数）を事前に知ることができない、継続的またはオープンエンドな設定での学習が増えています。それにもかかわらず、標準的な事前学習のレシピは、ホライゾン依存型の学習率スケジュール、特にコサイン減衰に依存しています。コサイン減衰は、減衰曲線を定義するために学習の総長 $N$ を事前に知る必要があるため、データが継続的に到着したり、トレーニング予算が動的に延長されたりするシナリオには適していません。

既存の「ニアリー・エニータイム（ほぼいつでも利用可能）」な代替案である Warmup-Stable-Decay (WSD) スケジュールは、この依存性を軽減してはいるものの、依然として追加のオーバーヘッド（特定のインターバルでのチェックポイント保存）や、学習の延長時における再チューニングや再減衰を必要とし、これが計算コストを生み出し、「エニータイム」の特性を損なわせます。核心となる課題は、以下の条件を満たす学習率スケジュールを特定することです：

1. 学習ホライゾン $N$ に依存しないこと。
2. 事前に $N$ の知識がない状態で、任意の中間チェックポイント $T < N$ において競争力のあるパフォーマンス（「コサイン・エンベロープ」を追跡できること）を実現すること。

メソドロジー

理論的分析

著者らは、ガウスデータを用いた過剰パラメータ化された線形回帰を用いた理論的分析を提供しています。これはLLMの事前学習のプロキシ（代理）として機能します。

• セットアップ: 本研究では、二次損失関数上でイテレート平均（重み平均）を用いた確率的勾座降下法（SGD）を検討しています。データの共分散行列 $H$ とターゲットベクトル $w^*$ は、容量指数 $a$（固有値の減衰 $\lambda_k \sim k^{-a}$）とソース指数 $b$（ターゲット座標の減衰 $(w^*_k)^2 \sim k^{-b}$）によって特徴付けられるべき乗則分布に従います。
• 主要な洞察: 著者らはリスクのバイアス・バリアンス分解を分析しています。彼らは、定数学習率と平均化の組み合わせが特定のレジーム（$b \le a$）では最適となり得る一方で、$b > a$ の場合にはホライゾン依存のスケーリングが必要になることを示しています。
• 提案されるスケジュール: 本論文では、$\eta_t = \eta / t^\gamma$ という形式の多項式学習率スケジュールとイテレート平均（具体的には指数移動平均、EMA）を組み合わせることで、$N$ を知ることなくミニマックス最適収束率を達成できることを導き出しています。
  • $1 < a < b \le 2a$ のレジームでは、最適な減衰指数は $\gamma^* = 1 - a/b$ と導出されます。
  • 特筆すべき特殊なケースとして $\gamma = 1/2$（すなわち $1/\sqrt{t}$）があり、これは特定のスペクトル条件下において、対数因子と定数を除いてエニータイム最適であることが示されています。
• 重み平均: 分析は、重み平均がこれらのレートを達成するために中心的であることを強調しています。平均化なしでは、多項式減衰を用いた標準的なSGDは、線形回帰の設定において最適なレートに到達できないことが多いです。

実証的評価

理論的な処方箋は、OLMoコードベースを用いたLLMの事前学習を用いてストレス・テストされています。

• モデル: 150M および 300M パラメータの2つのモデルを訓練しました。
• スケール: 150Mモデルに対してはChinchilla最適トークン数の32倍、300Mモデルに対しては16倍までの学習を実施しました。
• ベースライン:
  • コサイン・エンベロープ: すべての2の累乗のホライゾン（1$\times$から32$\times$）に対して独立してチューニングされたコサイン・スケジュールを構築し、各時点での最良の検証損失を取ることで作成された「ゴールドスタンダード」のベースライン。
  • WSD: 目標ホライゾンの90%でチェックポイントを保存し、最後の10%で減衰させる Warmup-Stable-Decay スケジュール。
• 候補:
  • 定数LR + EMA: 固定学習率と、時間に比例して成長するEMAウィンドウ（$f=25$）。
  • $1/\sqrt{t}$ + EMA: $\eta_t = \eta \sqrt{\alpha/(t+\alpha)}$ という多項式減衰スケジュールとEMA。
• プロトコル: 著者らは、単一の固定構成（一度選ばれたハイパーパラメータ）が、各ホライゾンに対して再チューニングを行うことなく、中間チェックポイント全体にわたってコサイン・エンベロープを追跡できるかどうかを評価しています。

主要な貢献

1. エニータイム・スケジュールの理論的導出: 本論文は、べき乗則データを用いた線形回帰において、$1/t^\gamma$ 形式の学習率スケジュールとイテレート平均を組み合わせることで、学習ホライゾンに依存せずにミニマックス最適レートを達成することを証明しています。具体的には、最適な $\gamma$ をソース指数と容量指数（$a$ および $b$）の関数として確立しています。
2. 重み平均の役割: 本研究は、重み平均（モデル・マージング）が単なる汎化のためのヒューリスティックではなく、過剰パラメータ化された設定において最適な収束レートを達成するための理論的な必然性であることを強調しています。
3. 大規模での実証的検証: 著者らは、単純なホライゾンフリーのスケジュール（Constant LR + EMA および $1/\sqrt{t}$ + EMA）が、大規模なトレーニング範囲（最大32$\times$ Chinchilla）にわたって、注意深くチューニングされたホライゾン依存のコサイン・エンベロープと同等の性能を示すことを実証しています。
4. 将来的な転用性: 本研究は、中間チェックポイント（例：4$\times$ Chinchilla）でチューニングされたハイパーパラメータが、再チューニングなしで、より長いトレーニング実行（最大32$\times$）へと正常に転用できることを示しており、オープンエンドな学習のための実用的なプロトコルを提供しています。

結果

• コサインとの比較: 全トレーニング範囲（1$\times$ から 32$\times$）にわたって、提案されたエニータイム・スケジュール（Constant + EMA および $1/\sqrt{t}$ + EMA）は、独立してチューニングされたコサイン・エンベロープに匹敵する検証損失を達成しています。その差は極めて小さく（多くの場合、検証損失ユニットで0.005以内）、最小限です。
• 単一軌跡の成功: 異なるホライゾンに対して再チューニングを必要とするコサイン減衰とは異なり、エニータイム・スケジュールの単一の実行は、効果的に「コサイン・エンベロープ」を追跡します。
• ラージバッチ・レジーム: 勾配ノイズが抑制されるラージバッチ実験（バッチサイズ 4096）では、定数学習率と平均化の組み合わせが、コサイン減衰および $1/\sqrt{t}$ スケジュールよりも優れた性能を示しました。これは、分散が低い場合、減衰は不要であり、より大きな定数レートがバイアス減少を加速させるという理論的予測と一致しています。
• WSDとの比較: WSD スケジュールは、コサインおよび提案されたエニータイム手法と同等の性能を発揮します。しかし、WSDは減衰フェーズを開始するために実行の90%でチェックポイントを保存する必要があるため、厳密な意味でのホライゾンフリーではありません。これに対し、提案手法はこのような介入を必要としません。

意義と主張

本論文は、重み平均を組み合わせた単純なホライゾンフリーのステップサイズが、大規模言語モデルの事前学習におけるコサイン学習率スケジュールの実用的かつ効果的な代替案となることを主張しています。

• 実用性: 提案手法は、総トレーニング予算を事前に知る必要性を排除し、真に継続的かつオープンエンドな学習を容易にします。
• 簡潔さ: 提案手法は、WSD（減衰開始のためのチェックポイント保存オーバーヘッド）よりも計算コストが低く、ホライゾン依存のコサイン・アニーリングよりも実装が容易です。
• 理論的根拠: 線形回帰レジームにおける厳密な理論的分析によって、経験的な成功が裏付けられており、これらのスケジュールの利点が単なる経験的なアーティファクトではなく、過剰パラメータ化された設定における根本的な最適化特性に由来することを示唆しています。

著者らは謙虚な姿勢を保ち、彼らの理論的境界が定数と対数因子を除いてのみミニマックスレートと一致すること、および現在の経験的な結論が特定のモデルサイズ（150M/300M）およびオプティマイザ（重み減衰のないAdamW）に限定されていることを述べています。彼らは、最適なチューニングポイント（例：4$\times$ Chinchilla）を自動的に選択する方法の理解や、これらの知見をより大きなスケールや異なるオプティマイザへ拡張することへの今後の課題を挙げています。