[論文レビュー] Adam Improves Muon: Adaptive Moment Estimation with Orthogonalized Momentum

Minxin Zhang, Yuxuan Liu+1 more|arXiv (Cornell University)|Feb 19, 2026

Muon and positron interactions and applications被引用数 0

ひとこと要約

この論文は NAMO および NAMO-D という新しい最適化アルゴリズムを導入し、直交化されたモーメントと適応的モーメント推定を組み合わせ、GPT-2 の事前学習における AdamW および Muon ベースラインより理論的収束性と経験的性能向上を示します。

ABSTRACT

Efficient stochastic optimization typically integrates an update direction that performs well in the deterministic regime with a mechanism adapting to stochastic perturbations. While Adam uses adaptive moment estimates to promote stability, Muon utilizes the weight layers' matrix structure via orthogonalized momentum, showing superior performance in large language model training. We propose a new optimizer and a diagonal extension, NAMO and NAMO-D, providing the first principled integration of orthogonalized momentum with norm-based Adam-type noise adaptation. NAMO scales orthogonalized momentum using a single adaptive stepsize, preserving orthogonality while improving upon Muon at negligible additional cost. NAMO-D instead right-multiplies orthogonalized momentum by a diagonal matrix with clamped entries. This design enables neuron-wise noise adaptation and aligns with the common near block-diagonal Hessian structure. Under standard assumptions, we establish optimal convergence rates for both algorithms in the deterministic setting and show that, in the stochastic setting, their convergence guarantees adapt to the noise level of stochastic gradients. Experiments on pretraining GPT-2 models demonstrate improved performance of both NAMO and NAMO-D compared to the AdamW and Muon baselines, with NAMO-D achieving further gains over NAMO via an additional clamping hyperparameter that balances the competing goals of maintaining a well-conditioned update direction and leveraging fine-grained noise adaptation.

研究の動機と目的

• Deterministic な領域で効果的な更新方向を確保しつつ、確率的勾配ノイズに適応する必要性を動機づける。
• NAMO を直交化モーメントのノルムベースの適応スケーリングとして提案し、安定性と性能を向上させる。
• NAMO-D を対角拡張として導入し、更新条件を保持しつつニューロン単位のノイズ適応を可能にする。
• 決定論的および確率的設定の双方において、両アルゴリズムの理論的収束保証を提供する。
• GPT-2 の事前学習において、AdamW および Muon のベースラインと比較して実験的改善を示す。

提案手法

• NAMO を更新における直交化モーメントを乗ずるノルムベースの適応スカラーとして定義する。 Θ_t = Θ_{t-1} - η α_t O_t ただし α_t = (||M_t|| / (sqrt(v_t) + ε))。
• NAMO-D を直交化モーメントを対角行列 D_t で右掛けしてニューロン単位の適応ステップ幅を実現する。
• 勾配の偏りのある一階および二階モーメント推定量 (M_t および v_t) とそのノルムを計算し、バイアス修正を行って hat M_t と hat v_t を得る。
• Ortho(M_t) によってモーメントを直交化し、更新方向 O_t を得る。
• NAMO および NAMO-D の決定論的設定と確率的設定の標準的な滑らかさと分散境界仮定の下での収束速度を確立する。
• 実験で用いられる直交化の実用的近似（例：Newton–Schulz 法の反復）を説明する。

実験結果

リサーチクエスチョン

• RQ1 直交化されたモーメントを Adam 型の分散適応と principled に組み合わせることは可能か。
• RQ2 NAMO および NAMO-D は AdamW および Muon と同等かそれ以上の証明可能な収束保証を提供するか。
• RQ3 これらの手法は GPT-2 のような大規模言語モデルの事前学習タスクで実用的な性能向上をもたらすか。
• RQ4 ニューロン単位の適応（NAMO-D）による近傍のブロック対角 Hessian 構造に対して、スカラー適応（NAMO）より有意な利点があるか。

主な発見

• NAMO および NAMO-D は GPT-2 事前学習において AdamW および Muon ベースラインより性能が向上する。
• NAMO は追加コストなしでスケールされた直交化モーメントの単一の適応ステップ幅を提供する。
• NAMO-D は対角スケーリングとフィルタリング機構によりニューロン単位のノイズ適応を細かく制御し、更新を良好に条件付けする。
• 理論的分析により、決定論的設定で最適収束速度を、確率的設定でノイズ適応の保証を両アルゴリズムに対して示す。
• NAMO-D は追加のクラッピングハイパーパラメータにより条件付けとノイズ適応のバランスを取ることで NAMO よりさらなる利得をもたらす。