技術的要約：返却されるモデルのための学習

問題提起

現代の言語モデル（LM）のパイプラインは、最終的なイテレート（反復値）そのものではなく、通常は指数移動平均（EMA）を用いた平均化されたモデルを返却します。イテレートの平均化は最適化を安定させ、汎化性能を向上させることが知られていますが、根本的な疑問が残っています。最終的な出力がイテレートの平均であるならば、この特定の推定器の性能を最適化するために、学習プロセス自体をどのように修正すべきか？ という点です。

先行研究では、平均化によってより積極的な学習率が可能になるという観察を通じて、この問題に暗黙的に対処してきました。しかし、本論文は直接的なアプローチを取ります。すなわち、返却されるイテレート平均の推定誤差を最小化するように、学習の軌跡を明示的に制御する最適化アルゴリズムの設計を目指します。

手法

理論的定式化

著者らは、この問題を連続時間確率二次モデルにおける最適制御問題として定式化しています。

• ダイナミクス: パラメータの軌跡 $\theta_t$ は、二次極小値付近でのノイズを含む勾配降下法を表す確率微分方程式（Ornstein–Uhlenbeck過程）に従い、制御入力 $u_t$ によって拡張されます。
• 目的: 制御介入の大きさに対する二次ペナルティ（正則化パラメータ $\lambda$）を課した上で、返却される時間平均推定器 $\hat{\theta}_{EMA, T}$ と真の最適値 $\mu^*$ との間の平均二乗誤差を最小化する制御戦略 $u_t$ を見つけることです。
• 導出: 線形二次ガウス（LQG）制御理論と確実性等価性の原理を用いて、最適なコントローラーを導出します。厳密な解は、軌跡から与えられる最適値の事後平均に依存しますが、これは大規模なLMにおいては計算不可能です。

実用的近似：PACE

解決策を実装可能にするため、著者らは一連のヒューリスティックな近似を導入しています。

1. 推定器の置換: 未知の最適値の事後平均を、現在のイテレート平均（$\hat{\theta}_{EMA, t}$）で置き換えます。
2. 時間の近似: 現在の時間 $t$ が総学習時間 $T$ よりも十分に小さい（$t \ll T$）と仮定し、時間依存の係数を簡略化します。
3. 離散化と前処理: 連続的な制御を離散化し、未知のヘッセ行列をAdamWで使用される対角前処理（二次のモーメント推定値 $v_t$）で置き換えます。

これにより、PACE (Pullback Averaging Control for Efficient Optimization) が得られます。これはAdamWの軽量なラッパーであり、ライブウェイトをEMAウェイトへと引き戻す「プルバック（引き戻し）」項を加えるものです。
$$\theta_{k+1} = \theta_k - \eta g_k + C(\hat{\theta}_{EMA, k} - \theta_k)$$
ここで、$C$ はAdamの前処理行列から導出される対角行列であり、プルバックの強さは、更新が標準的なステップとプルバックの凸結合に留まるようにクリップされます。

主な貢献

1. 最適制御の定式化: 本論文は、最終的なイテレートではなく、イテレート平均の誤差を最小化するために特別に設計された最適化アルゴリズムの、初の原理的な導出を提供します。
2. 理論的保証:
   • 凸設定: 著者らは、様式化されたバージョンのPACEが、偏りのない確率的勾配を持つ任意の凸損失関数に対して、EMAの強度に依存する定数倍の範囲内で、標準的な $O(1/\sqrt{T})$ のレートで収束することを証明しています。これにより、本手法が標準的なSGDよりも著しく劣ることは決してないことが保証されます。
   • 二次設定: 導出に使用された二次設定において、制御されたダイナミクスは、制御されていないダイナミクス（標準的なEMA付きSGD）と比較して、イテレート平均推定器の限界二乗誤差を厳密に改善します。この改善は、特定の二次インスタンスにおいて任意に大きくなります。
3. 実証的検証: 本手法は以下のモデルおよびデータセットで評価されました。
   • ファインチューニング: SmolLM2-1.7B, Qwen3-1.7B, Gemma3-1B モデル（smol-smoltalk データセット）。
   • 事前学習: GPT-2 (124M)（Chinchilla最適トークン予算における FineWeb）。
   • ベースライン: PACEは、バニラAdamW、EMA評価付きAdamW、およびSchedule-Freeアルゴリズムと比較されました。

結果

• パフォーマンス: PACEは、テストされたすべてのモデル（1B–2Bパラメータ）およびデータセットにおいて、バニラAdamWおよびEMA評価付きAdamWの両方を厳密に上回りました。
• 堅牢性: この改善は、幅広い学習率、減衰スケジュール（一定、コサイン、Warmup-Stable-Decay）、およびハイパーパラメータ（プルバック強度 $c$、EMA指数 $\kappa$）に対して堅牢です。
• Schedule-Freeとの比較: PACEは、学習率の減衰を不要にするためにイテレート平均を使用するSchedule-Freeアルゴリズム（Defazio et al., 2024）と同等の性能を示します。ただし、PACEは異なる更新規則を用いています。
• 学習率の減衰: PACEは一定の学習率でも良好に動作しており、学習率の減衰が通常必要とされる領域においても、学習を安定化できることを示唆しています。トークン予算の比較において、PACEの一定学習率の軌跡は、WSDブランチと同等に、あるいはそれ以上に、損失関数のランドスケープの「リバーフロア（川底）」を効果的に追跡します。

意義と限界

意義:
本論文は、最終的なイテレートではなく、返却されるモデル（平均）を最適化することが、具体的な性能向上につながることを示しています。学習の軌跡を制御問題として扱うことで、PACEは複雑なスケジュールのチューニングを必要とせずに、学習を安定させ、汎化性能を向上させるための理論的根拠に基づいた手法を提供します。著者らは、PACEが一般的な凸設定において標準的な手法よりもわずかに遅い理論的収束保証を持つものの、二次設定においては標準的な手法と比較して改善された理論的収束を示すことを指摘しています。

限界:

• 理論的仮定: 導出は、加法的ノイズを伴う簡略化された二次モデルに依存しています。一般の凸関数に対する収束は証明されていますが、モメンタムや適応的前処理（AdamWの中核）を理論的解析に厳密に取り入れることは、依然として未解決の課題です。
• スケール: 実験は、ファインチューニングでは最大2Bパラメータ、事前学習では124Mまでのモデルに限定されています。より大規模なスケール（例：70B以上）での有効性は検証されていません。
• メモリオーバーヘッド: PACEは、学習中にモデルウェイトの追加コピー（EMA状態）を保持する必要があるため、無視できないメモリオーバーヘッドが生じますが、実用上は大きなバッチサイズによってこれが緩和されます。
• シードへの感受性: 計算リソースの制約により、主要な実行には単一のシードが使用されましたが、マルチシードによる堅牢性調査により、手法の順序の安定性が確認されています。

著者らは、PACEが、平均化されたモデルが最終製品となる現代のLMの実際のデプロイ戦略と、最適化アルゴリズムを一致させるための有望な一歩であると結論付けています。