Schedule-Free Optimizer（スケジュールフリーオプティマイザ）
Meta（旧 Facebook）の研究チームが 2024 年に提案した、学習率スケジュールを必要としない最適化アルゴリズム。従来の Cosine Decay や Linear Decay といった事前設計のスケジュールなしに、AdamW + Cosine Decay と同等以上の性能を達成する。ハイパーパラメータチューニングの工数を大幅に削減できる実用的な利点がある。
Schedule-Free Optimizerとは
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（平均系列）: z の重み付き移動平均。推論・評価に使用する「安定用」パラメータ
訓練時の勾配計算には z_t と x_t の補間点 y_t を使用する：
y_t = (1 − β) · z_t + β · x_t
ここで β はモーメンタム係数に対応する。この補間により、探索の自由度（z の役割）と安定性（x の役割）を両立する。
Schedule-Free AdamW のアルゴリズム
- 補間点計算: y_t = (1 − β1) · z_t + β1 · x_t
- 勾配計算: g_t = ∇L(y_t)
- z の更新（AdamW ステップ）: z_{t+1} = AdamW_step(z_t, g_t)
- x の更新（平均化）: x_{t+1} = (1 − 1/(t+1)) · x_t + 1/(t+1) · z_{t+1}
- 推論時は x_t を使用
理論的背景
Schedule-Free の理論的基盤は、Polyak-Ruppert 平均化と Primal Averaging の研究に遡る。定数学習率のオプティマイザの軌跡を適切に平均化すると、学習率を減衰させた場合と同等の収束率が得られることが理論的に証明されている。
直感的には、一定の学習率で「探索」を続けるオプティマイザ（z 系列）は損失ランドスケープの最適解周辺を振動するが、その軌跡の平均（x 系列）は最適解に収束する。学習率スケジュールは本質的にこの平均化を暗黙的に行っていたと解釈でき、Schedule-Free はこれを明示的に分離したものである。
AdamW + Cosine Decayとの比較
論文で報告された各種ベンチマークでの比較：
| タスク | 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 で不要になるハイパーパラメータ：
- 総ステップ数（事前決定不要）
- Warmup ステップ数（Warmup 自体は引き続き推奨されるが、メインのスケジュールは不要）
- スケジュール形状（Cosine/Linear/WSD の選択が不要）
- 最終学習率（Cosine Decay の min_lr 設定が不要）
実装と利用方法
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 で管理できるため、追加メモリは最小限に抑えられている。
FAQ
Q1. Schedule-FreeはWarmupも不要にするのか？
Warmup 自体は引き続き推奨される。Schedule-Free が不要にするのはメインの学習率減衰スケジュール（Cosine Decay 等）であり、初期の数百〜数千ステップでの Warmup は学習の安定化に効果がある。ただし、Warmup の期間は Cosine Decay ほど厳密に設定する必要はなく、多少のずれが性能に大きく影響しない。
Q2. 既存のAdamW訓練コードからの移行コストは？
非常に低い。schedulefree ライブラリをインストールし、torch.optim.AdamW を schedulefree.AdamWScheduleFree に置き換え、評価時に optimizer.eval() / optimizer.train() を呼ぶだけである。学習率スケジューラのコード（lr_scheduler）は削除できる。既存のハイパーパラメータ（lr, betas, weight_decay）はそのまま流用可能。
Q3. Schedule-Freeは全ての場面でCosine Decayより優れるか？
ほとんどの標準的なベンチマークで同等以上の結果が報告されているが、非常に長い訓練（数百万ステップ）での検証はまだ限定的である。また、WSD（Warmup-Stable-Decay）スケジュールの方が中間チェックポイントの性能で優れるケースも報告されている。ただし、チューニングコストの削減を考慮すると、多くの実務シナリオで Schedule-Free が最も効率的な選択肢である。