技術概要：ScheduleFree+ 大規模言語モデルへの学習率不要かつスケジュール不要な学習の拡張

1. 問題提起

従来の機械学習の最適化は、収束を確保するために学習率スケジュール（例：線形減衰、コサイン、ウォームアップ - 安定 - 減衰）に大きく依存している。これらは効果的であるが、大幅なチューニングを必要とし、しばしば不安定または予測不可能な学習曲線をもたらす。最近の「スケジュール不要（Schedule-Free）」学習手法は、スケジュールを積極的な反復平均化に置き換えることで小規模ベンチマークにおいて有望な結果を示したが、歴史的には大規模言語モデル（LLM）のトレーニングへの拡張に失敗してきた。具体的には、従来の Schedule-Free 実装は以下の課題に直面していた：

• 大バッチサイズにおける発散： バッチサイズが小規模の限界を超えて増加すると、パフォーマンスが著しく低下した。
• 勾配ノルムのドリフト： この手法は機能するために異常に小さな重み減衰値を必要とし、その結果、重みノルムが不安定になり、長期トレーニングにおいてパフォーマンスが低下した。
• 長期実行における最適でない収束： この手法は、1000 トークン/パラメータなどの長期トレーニング期間における平均化の恩恵を十分に活用できなかった。

本論文は、LLM の事前学習において、最先端（SOTA）のスケジュールベース手法のパフォーマンスに匹敵、あるいはそれを超えるスケジュール不要学習の拡張という課題に取り組み、学習率のチューニングを必要としない安定した「いつでも（anytime）」トレーニング手順の実現を目指している。

2. 手法：ScheduleFree+

著者らは、LLM に対して堅牢にするために、基本の Schedule-Free アルゴリズムにいくつかの具体的な修正を組み込んだ複合手法**ScheduleFree+**を提案する。中核となる構成要素は以下の通りである：

A. 内部モーメントの再導入

元の Schedule-Free の研究はメモリ節約のために内部モーメント（例：Adam の $\beta_1$）を除去したが、著者らはこれが大バッチサイズにおいて有害であることを発見した。

• 対策： 基本オプティマイザに内部モーメント（例：$\beta_1 = 0.75$ または $0.9$）を再導入する。
• 効果： これにより大バッチサイズに対する堅牢性が回復し、発散することなくより高い学習率を可能にする。これにより、標準的な Schedule-Free AdamW で観察された「スケーリングの崖」に対処できる。

B. 逆勾配ノルム重み付け（Polyak ステップサイズ）

標準的な Schedule-Free 学習は、増加する勾配ノルムを補償するために小さな重み減衰を必要とし、これが不安定さを引き起こしていた。著者らは、Polyak ステップサイズに基づいた理論的な補正を導き出した。

• メカニズム： ステップサイズ $\gamma_t$ は、勾配の L1 ノルムに反比例してスケーリングされる（$\gamma_t \propto 1/\|g_t\|_1$）。
• 実装： 勾配 L1 ノルムの指数移動平均（EMA）を用いて分母を推定し、凸回帰 bound から導出された最適な重み付けを近似する。
• 結果： この補正は、正規化レイヤーによって引き起こされる実効学習率のドリフトを相殺し、標準的かつ大きな重み減衰値（例：0.05–0.5）の使用を可能にする。これにより重みノルムが安定し、長期実行を不安定にする「重みノルムの縮小」現象を防ぐ。

C. Polyak ステップによる学習率不要最適化

逆勾配重み付けに基づき、著者らは修正された Polyak ステップサイズを用いた完全な学習率不要バリアントを実装する：
$$\gamma_t = \frac{f(y_t) - f^* + \beta \langle \nabla f(y_t), z_t - x_t \rangle}{\sqrt{\pi/2} \| \nabla f(y_t) \|_1}$$

• 安定性： 分子はノイズの多い現在の反復値ではなく、関数値ギャップの推定値（$f(z_t) - f^*$）を使用し、分母は平滑化された L1 ノルムを使用する。これにより、確率的 Polyak ステップで一般的に見られる分母の消失問題を回避する。
• 適応性： このアプローチは、バッチサイズに応じて学習率を自動的にスケーリングし（AdamW における $\sqrt{b}$ スケーリングに従う）、手動の学習率グリッドサーチを不要にする。

D. ウォームスタートとアニーリング戦略

• C-ウォームアップ： 平均化バッファが初期の急速な収束段階を希薄化しないよう、著者らは平均化更新（$c_t = 1$）を初期期間（例：学習率ウォームアップ期間の 2 倍）遅らせる。
• ベータアニーリング： 現在の反復値 $z_t$ と平均 $x_t$ を混合する補間パラメータ $\beta$ を、トレーニング実行中に低い値（例：0.9）から高い値（例：0.965）へアニーリングする。これにより、初期の高速な収束を捉えつつ、長期的な損失削減を最適化する。
• 重み付き平均（$r=1$）： 長期実行では、平均化重みをデフォルトの $r=0$ ではなく $r=1$（線形に増加する重み）に設定する。これにより、強凸でない問題における収束が加速される。

3. 主要な結果

著者らは、1 億 2000 万から 10 億パラメータの LM に対して、パラメータあたり最大 1000 トークンでトレーニングされた大規模スケーリング実験を実施した。

• 長期トレーニング（1000 TPP）： この規模において、ScheduleFree+ はウォームアップ - 安定 - 減衰（WSD）および線形減衰スケジュールの両方を大幅に上回る。
  • 効率性： 線形減衰スケジュールと同じ最終損失を、トレーニング時間の 31% 短縮（目標損失に到達するために必要なトークンの 45% 削減）で達成する。
  • 安定性： この手法は、理論的な $1/\sqrt{t}$ 減衰率に従う滑らかで予測可能な損失曲線を生み出し、初期トレーニングデータから最終損失を正確に予測可能にする。
• 短期トレーニング（20–100 TPP）： ScheduleFree+ は、すべてのモデルサイズにおいて WSD スケジュールを上回る。線形減衰に対しては、定常状態に達するまでの時間が必要なため、非常に短い期間（20 TPP）では同等かわずかに劣るが、一般的には競争力のあるパフォーマンスを示す。
• バッチサイズスケーリング： 内部モーメントにより、ScheduleFree+ は 800 万のバッチサイズまで堅牢にスケーリングし、線形減衰を伴う標準的な AdamW のパフォーマンスに匹敵する。
• 重みノルムの安定性： Polyak ステップサイズと大きな重み減衰を使用することで、この手法は安定した重みノルムを維持し、修正されていない Schedule-Free 手法が要求する小さな重み減衰値に関連する発散や量子化の問題を回避する。

4. 意義と主張

本論文は、ScheduleFree+ が事前学習中のモデル平均化とチェックポイントマージの理論的基盤を提供し、オンライン最適化とオフラインモデルスープの間のギャップを埋めると主張する。

• いつでもトレーニング（Anytime Training）： 最終的な損失向上を実現するためにトレーニングの最後に特定の減衰フェーズを必要とする WSD スケジュールとは異なり、ScheduleFree+ は真の「いつでも」機能を提供する。モデルは、専用の減衰フェーズなしに、任意の停止時点で使用可能かつ最適である。
• チューニングの排除： この手法は学習率のグリッドサーチを不要とし、代わりに Polyak ステップサイズと逆勾配重み付けに依存して、バッチサイズやトレーニング期間に自動的に適応する。
• 予測可能性： 著者らは、Schedule-Free 学習の損失曲線は経験的に安定しているだけでなく、理論的根拠に基づいており、非凸の LM 設定であっても凸最適化理論が予測する $1/\sqrt{t}$ 収束率に密接に従うと論じている。

要約すると、本論文は、内部モーメント、逆勾配重み付け、適応的アニーリングという特定の修正を加えることで、スケジュール不要学習を拡張し、大規模言語モデルのトレーニングにおける従来の学習率スケジュールに代わる、優れており、安定し、チューニング不要な代替手段となり得ることを実証している。