技術的サマリー：OrScale – 層ごとの信頼率スケーリングを用いた直交化最適化

1. 問題定義

Muon オプティマイザは、ニュートン・シュルツ反復（極因子の近似）を用いて行列値ステップの方向を固定する直交化更新を活用することで、大規模トレーニングにおける実用的なツールとして登場しました。しかし、Muon は重大な限界を有しています。すなわち、更新の方向を最適化する一方で、その大きさは単一の、本質的にグローバルな学習率によって制御されたままです。

このグローバルなスケーリングは最適ではありません。なぜなら、異なる層（アテンション射影、MLP のアップ/ダウン射影、埋め込みなど）は、はるかに異なる勾配分布と有効な曲率を示すからです。Moonlight 変種は、要素ごとの RMS を均等化するために静的な形状ごとの因子を導入しましたが、トレーニング中の層スケールの動的なドリフトに適応することはできません。

SGD および Adam においてステップサイズをグローバルな学習率から切り離すことに成功した層ごとの適応率スケーリング（LARS）および層ごとの適応モーメントスケーリング（LAMB）を Muon に移植しようとする以前の試みは失敗しました。本論文は、単純な移植が以下の 3 つの具体的な構造的な理由により破綻することを特定しています：

1. 退化した分母： 直交化された方向のフロベニウスノルム（$\|Q_\ell\|_F$）を信頼率の分母として使用すると、更新の大きさではなく行列の形状（ランク）によってほぼ一定の値が決まってしまうため、失敗します。
2. 単位の不一致とクリップ飽和： 生のモーメントノルム（$\|\tilde{M}_\ell\|_F$）を使用すると、勾配単位と重み単位の不一致が生じ、信頼率がほぼ普遍的に上限クリップ値（$r_{max}$）に達してしまい、適応が非効率的になります。
3. 切り離された重み減衰の暴走： 較正された分母と切り離された重み減衰を組み合わせると、正のフィードバックループが生じ、重みのノルムが制御不能に成長（10〜50 倍）し、トレーニングを不安定化させます。

2. 手法

著者らは、OrScale を提案しました。これは単一の設計原則に基づいたフレームワークです：層ごとの信頼率の分母は、オプティマイザがパラメータから減算する実際の更新方向のフロベニウスノルムを測定しなければならない。

中核となる原則

$D_{\ell,t}$ を、グローバルな学習率が適用される前に重み $W_{\ell,t}$ から減算されるパラメータ空間の実際の方向とします。信頼率は以下のように定義されます：
$$r_{\ell,t} = \frac{\|W_{\ell,t}\|_F}{\|D_{\ell,t}\|_F + \epsilon}$$
重要なのは、$D_{\ell,t}$ には重み減衰項と直交化された方向の両方が含まれなければならない点です。

2 つの変種

このフレームワークは、2 つの具体的な実装をもたらします：

1. OrScale（一般/ビジョンデフォルト）：

   • 更新方向： $D_{\ell,t} = \lambda W_{\ell,t} + Q_{\ell,t}$（$\lambda$ は重み減衰、$Q$ は直交化された方向）。
   • 形状因子： $s_\ell = 1$（静的スケーリングなし）。
   • 較正： なし。
   • 重み減衰： 信頼率スケーリングされた更新ステップに直接統合されます。

2. OrScale-LM（LLM デフォルト）：

   • 更新方向： $D_{\ell,t} = \lambda W_{\ell,t} + s_\ell Q_{\ell,t}$。
   • 形状因子： AdamW の要素ごとの RMS に一致させるため、Moonlight の静的因子 $s_\ell = 0.2\sqrt{\max(m_\ell, n_\ell)}$ を採用します。
   • 較正： 初期化時に 1 回計算される遅延な層ごとの較正定数（$c_{\text{denom}, \ell}$）を導入します。これにより、初期信頼率が正確に 1 になるようにアンカーされます（$r_{\ell,0}=1$）。
   • 較正の意義： これにより、OrScale-LM は、新しいハイパーパラメータ探索を必要とせずに、Moonlight（および延いては AdamW）用に調整された学習率と重み減衰のハイパーパラメータを継承できます。また、初期化における幅不変性を保証します。

アルゴリズム的詳細

• 結合された重み減衰： 重み減衰は、引き戻しとして別々に適用されるのではなく、信頼率の計算（$\lambda W$ が分母の一部）に統合されます。これにより、切り離された変種で観察された暴走的なフィードバックループが防止されます。
• 効率性： この手法は、追加のオーバーヘッドが極めて少ない（ウォールクロック時間の 1% 未満）です。行列層ごとに 1 つの追加スカラーと、ステップごとに 2 つのフロベニウスノルム計算のみを必要とし、これらは既存のニュートン・シュルツ反復によって支配されます。

3. 主要な貢献

• アルゴリズム： 層ごとの信頼率スケーリングを Muon に成功裏に統合する、統一されたパラメータ軽量のレシピ（アルゴリズム 1）。著者らは、設計空間におけるユニークな「実用可能なコーナー」（実際の更新方向の分母 + 結合された WD）を特定しました。
• 失敗分析： Muon の信頼率設計空間の厳密なマッピング（表 1）により、3 つの失敗モード（退化した分母、クリップ飽和、重み成長の暴走）を分析的かつ経験的に特徴付け、それらがなぜ失敗し、提案された解決策がなぜ成功するのかを実証しました。
• 理論：
  • 収束： 核ノルムによる $O(1/\sqrt{T})$ の非凸収束保証を証明しました。
  • 適応的ゲイン： 測定可能な層の異質性のもとで、層適応的な降下定数 $\kappa_{\text{eff}} > 1$ を確立し、Muon の均質な定数よりも厳密に改善することを示しました。
  • 分離： クリップ飽和の条件下において、MuonTrust/MuScale との厳密な分離を証明しました。
  • 較正： OrScale-LM が初期化において $\mu$P 風の学習率転送を維持することを示しました。
• 経験的検証： ビジョンおよび言語モデルタスクにおける包括的な評価。

4. 実験結果

ビジョン（CIFAR-10 / DavidNet）

• 性能： OrScale は 94.05% $\pm$ 0.08 の検証トップ 1 精度を達成し、5 つのベースラインの中で 1 位となりました。
• 改善： これは、標準的な Muon（93.70%）に対して +0.35 ポイント、Muon+Moonlight（93.75%）に対して +0.30 ポイントの改善を表します。
• 安定性： OrScale は、層適応的ゲインの理論的予測と一致し、Muon に比べて広い安定した学習率ウィンドウを示しました。
• アブレーション： 失敗モードのアブレーション研究により、特定された失敗モードに対応する変種（例：$\|Q_\ell\|_F$ または生のモーメントの使用）は、予測された症状（例：99.5〜100% のクリップ飽和または重みノルムの暴走）を示し、OrScale の性能に匹敵しなかったことが確認されました。

大規模言語モデル（FineWeb-Edu 事前学習）

• スケーリング： 4 つのモデルスケール（1 億 2500 万〜11 億パラメータ）で評価されました。
• 比較： OrScale-LM は、4 つのスケールのうち 3 つ（1 億 2500 万、5 億 4500 万、11 億）で Muon+Moonlight を上回り、すべてのスケールで AdamW を打ち破りました。
• 損失低減： 11 億パラメータのスケールにおいて、OrScale-LM は最終検証クロスエントロピーを 2.6251 まで達成しました。これは Muon+Moonlight の 2.6360、AdamW の 2.7304 と比較して優れています。
• 効率性： この優位性は計算範囲全体で維持され、OrScale-LM は LLM にとって計算が最も重要となる領域において明確な勝利を示しました。

5. 意義と主張

本論文は、OrScale が Muon のスペクトル幾何学と LARS/LAMB の層ごとの適応性の間のギャップを成功裏に埋めたと主張しています。その主な意義は以下の点にあります：

1. 大きさ問題の解決： 構造的な失敗モードに対して堅牢な、Muon 向けの最初の時間変化する層ごとの大きさ制御を提供します。
2. 実用的な転送性： OrScale-LM 変種と遅延較正を通じて、実務者が AdamW や Moonlight 用に調整された学習率を、再調整なしで OrScale-LM に直接転送することを可能にし、採用の障壁を下げます。
3. 理論的妥当性： 以前のハイブリッド試みがなぜ失敗したのかについての厳密な理論的根拠を提供し、提案された手法が異質性のもとで厳密に優れた収束定数をもたらすことを証明します。
4. 経験的検証： これらの理論的改善が、小規模なビジョンタスクと大規模な LLM 事前学習の両方で具体的な利益に変換されることを実証し、OrScale を行列値層のための最先端オプティマイザとして確立しました。

著者らは、LLM の結果については慎重な態度を維持しており、11 億パラメータのスケールでの優位性は、より小さなスケールでは単一のシードのノイズと同等であると指摘しています。したがって、この結果は、すべての領域における決定的な統計的優位性というよりは、手法の有効性を示す方向的な証拠として見るべきです。しかし、スケール全体にわたる一貫した改善と、失敗モードの成功裏なアブレーションは、提案された設計原則を強く支持しています。