Meta(旧 Facebook)の研究チームが 2024 年に提案した、学習率スケジュールを必要としない最適化アルゴリズム。従来の Cosine Decay や Linear Decay といった事前設計のスケジュールなしに、AdamW + Cosine Decay と同等以上の性能を達成する。ハイパーパラメータチューニングの工数を大幅に削減できる実用的な利点がある。
Schedule-Free Optimizer は、2024 年に Meta の Aaron Defazio と Konstantin Mishchenko が論文「The Road Less Scheduled」で提案した最適化手法のフレームワークである。従来の深層学習訓練では、学習率を徐々に減衰させるスケジュール(Cosine Decay, Linear Decay 等)が性能向上に不可欠とされてきたが、Schedule-Free はスケジュールを使わずに同等以上の性能を実現する。
学習率スケジュールの設計と調整は LLM 訓練における最も時間のかかるハイパーパラメータチューニングの一つである:
| 問題 | 具体例 | 影響 |
|---|---|---|
| 総ステップ数の事前決定 | Cosine Decay では終了ステップを最初に決める必要 | 訓練途中での延長・短縮が困難 |
| Warmup 期間の設計 | 短すぎると発散、長すぎると収束遅延 | 試行錯誤が必要 |
| スケジュール形状の選択 | Cosine, Linear, WSD 等の選択 | タスク依存で最適解が異なる |
| 中間チェックポイントの性能 | Cosine Decay は終了付近でのみ最高性能 | 途中時点での評価が不安定 |
Schedule-Free Optimizer はこれらの問題を根本的に解決する。
Schedule-Free の核心は「パラメータ平均化(iterate averaging)」と「補間(interpolation)」の組み合わせにある。
従来のオプティマイザでは、現在のパラメータ θ_t がそのまま推論にも使われる。Schedule-Free では 2 つのパラメータ系列を維持する:
訓練時の勾配計算には z_t と x_t の補間点 y_t を使用する:
y_t = (1 − β) · z_t + β · x_t
ここで β はモーメンタム係数に対応する。この補間により、探索の自由度(z の役割)と安定性(x の役割)を両立する。
Schedule-Free の理論的基盤は、Polyak-Ruppert 平均化と Primal Averaging の研究に遡る。定数学習率のオプティマイザの軌跡を適切に平均化すると、学習率を減衰させた場合と同等の収束率が得られることが理論的に証明されている。
直感的には、一定の学習率で「探索」を続けるオプティマイザ(z 系列)は損失ランドスケープの最適解周辺を振動するが、その軌跡の平均(x 系列)は最適解に収束する。学習率スケジュールは本質的にこの平均化を暗黙的に行っていたと解釈でき、Schedule-Free はこれを明示的に分離したものである。
論文で報告された各種ベンチマークでの比較:
| タスク | AdamW + Cosine | Schedule-Free AdamW | 結果 |
|---|---|---|---|
| GPT-2 124M 事前学習 | ベースライン | 同等〜やや上回る | 勝ち |
| ImageNet ViT-B/16 | ベースライン | 同等 | 引き分け |
| CIFAR-100 ResNet-50 | ベースライン | やや上回る | 勝ち |
| IWSLT'14 翻訳 | ベースライン | 同等 | 引き分け |
| MLPerf ベンチマーク | ベースライン | 8 タスク中 7 勝 | 勝ち |
特筆すべきは、Schedule-Free が学習率スケジュールのチューニングなしにこれらの結果を達成している点である。Cosine Decay 側は各タスクに最適化されたスケジュールを使用しているため、実質的なチューニングコストを含めると Schedule-Free の優位性はさらに大きい。
| シナリオ | Cosine Decay | Schedule-Free |
|---|---|---|
| 訓練延長 | スケジュール再設計が必要 | そのまま継続可能 |
| 訓練短縮 | 性能が最適でない地点で終了 | 任意の時点で停止可能 |
| 中間チェックポイント | 性能が不安定 | x 系列は常に安定 |
| 転移学習 | スケジュール再設定が必要 | ベース学習率のみ調整 |
Schedule-Free で不要になるハイパーパラメータ:
Schedule-Free は公式の PyTorch ライブラリとして公開されている:
import schedulefree
optimizer = schedulefree.AdamWScheduleFree(
model.parameters(),
lr=3e-4,
betas=(0.9, 0.95),
weight_decay=0.1
)
# 訓練時
optimizer.train()
for batch in dataloader:
loss = model(batch)
loss.backward()
optimizer.step()
optimizer.zero_grad()
# 評価時(x 系列に切り替え)
optimizer.eval()
evaluate(model)
重要な注意点として、評価時には optimizer.eval() を呼び出して x 系列(平均化パラメータ)に切り替える必要がある。訓練時のパラメータ(y 系列)で評価すると、z 系列の振動の影響で性能が低く出る。
Schedule-Free は PyTorch の FSDP(Fully Sharded Data Parallel)および DeepSpeed と互換性がある。パラメータの平均化操作(x の更新)は各 GPU のローカルシャード上で独立に行えるため、追加の通信オーバーヘッドは発生しない。
メモリ使用量は AdamW と比較して z と x の 2 つのパラメータコピーが必要だが、実装上は x を in-place で管理できるため、追加メモリは最小限に抑えられている。
Warmup 自体は引き続き推奨される。Schedule-Free が不要にするのはメインの学習率減衰スケジュール(Cosine Decay 等)であり、初期の数百〜数千ステップでの Warmup は学習の安定化に効果がある。ただし、Warmup の期間は Cosine Decay ほど厳密に設定する必要はなく、多少のずれが性能に大きく影響しない。
非常に低い。schedulefree ライブラリをインストールし、torch.optim.AdamW を schedulefree.AdamWScheduleFree に置き換え、評価時に optimizer.eval() / optimizer.train() を呼ぶだけである。学習率スケジューラのコード(lr_scheduler)は削除できる。既存のハイパーパラメータ(lr, betas, weight_decay)はそのまま流用可能。
ほとんどの標準的なベンチマークで同等以上の結果が報告されているが、非常に長い訓練(数百万ステップ)での検証はまだ限定的である。また、WSD(Warmup-Stable-Decay)スケジュールの方が中間チェックポイントの性能で優れるケースも報告されている。ただし、チューニングコストの削減を考慮すると、多くの実務シナリオで Schedule-Free が最も効率的な選択肢である。