[論文レビュー] Where Does Warm-Up Come From? Adaptive Scheduling for Norm-Constrained Optimizers

Artem Riabinin, Andrey Veprikov+3 more|arXiv (Cornell University)|Feb 5, 2026

Stochastic Gradient Optimization Techniques被引用数 0

ひとこと要約

この論文は、ノルム制約付き最適化器（例：Muon、Lion、normSGD）に対する理論的枠組みと実用的な適応型ウォームアップスケジューラを提供します。局所的曲率が最適性欠如に比例して成長する一般化された滑らかさ仮定を導入し、ウォームアップが自然に発生することを証明し、適応型ウォームアップが追加のハイパーパラメータ探索を行わずにLLMプリトレーニングを改善することを示します。

ABSTRACT

We study adaptive learning rate scheduling for norm-constrained optimizers (e.g., Muon and Lion). We introduce a generalized smoothness assumption under which local curvature decreases with the suboptimality gap and empirically verify that this behavior holds along optimization trajectories. Under this assumption, we establish convergence guarantees under an appropriate choice of learning rate, for which warm-up followed by decay arises naturally from the proof rather than being imposed heuristically. Building on this theory, we develop a practical learning rate scheduler that relies only on standard hyperparameters and adapts the warm-up duration automatically at the beginning of training. We evaluate this method on large language model pretraining with LLaMA architectures and show that our adaptive warm-up selection consistently outperforms or at least matches the best manually tuned warm-up schedules across all considered setups, without additional hyperparameter search. Our source code is available at https://github.com/brain-lab-research/llm-baselines/tree/warmup

研究の動機と目的

• ノルム制約付き最適化器のウォームアップを、経験的なヒューリスティックを超えて動機づけ・正当化する。
• 局所的曲率と最適性欠如を結びつける一般化滑らかさモデルを導入する。
• このモデルの下で、ウォームアップ-その後の減衰学習率で収束を証明する。
• 標準的なハイパーパラメータに依存する実用的な適応型ウォームアップスケジューラを開発する。
• 大規模言語モデルのプリトレーニングでスケジューラを検証し、競争力のある性能を示す。

提案手法

• LMOベースの更新式 x^{t+1}=x^{t}+ eta^{t} LMO(g^{t}) を定義し、2次近似による損失と関連付ける。
• Assumption 2: ( ho, K0, K1, Kρ)-滑らかさを、曲率がサブオプティマリティ f(x)-f^{*} に依存する形で提案する。
• 学習率 eta^{t}= (illed{Δ^{t}})/(D·K(x^{t})) のとき、欠陥度 Δ^{t} が減少し、K(x^{t}) も減少する（定理1）。
• ウェイト減衰を拡張：x^{t+1}=(1−λη^{t})x^{t}+η^{t} LMO(g^{t}) として収束を証明する（定理2）。
• 確率的設定へ拡張し、スケールされた勾配正規化と補間の仮定を用いた（定理3）。
• 実用的な適応型ウォームアップスケジューラを、Δに対して η(Δ)=Δ/(K0+K1Δ+K2Δ^2) という3パラメータモデルに適合させ、Δ′ でデカイやすへ切替える（アルゴリズム5）。

実験結果

リサーチクエスチョン

• RQ1 LMOベースの最適化器の学習率ウォームアップは、理論的に正当化できるか（純粋なヒューリスティックに留まらないか）？
• RQ2 トレーニング開始時にマニュアルな調整を要することなく、ウォームアップ期間を自動適応できるか？
• RQ3 サブ最適性に依存する滑らかさモデルは、ウォームアップと減衰を最適化の自然なダイナミクスとして説明できるか？
• RQ4 適応型ウォームアップスケジューラは、ハイパーパラメータ探索なしに大規模LLMプリトレーニングで良好に機能するか？

主な発見

• 局所的曲率がサブオプティマリティギャップに依存する一般化滑らかさモデルが、最適化の経路に沿って経験的に支持される。
• このモデルの下では、LMOベースの最適化器に対する収束証明からウォームアップとデカイが自然に出現する。
• 標準的なハイパーパラメータを使用し、訓練開始時に K0, K1, K2 と Δ′ を推定する実用的な適応型ウォームアップスケジューラが導出される。
• 適応型ウォームアップは、LLamaプリトレーニングにおいてMuon、Lion、normSGD の最適化スケジュールの中で、追加のハイパーパラメータ探索を要せず、手動で調整した最高のウォームアップと同等かそれを上回る性能を達成する。
• このアプローチは、モデルサイズやバッチ設定を問わず堅牢であり、特に小さなバッチ設定で顕著な恩恵をもたらす。