=== 要約 ===

技術要約：低ランクアダプターによるMuonの近似（sMuon）

問題提起

Muonオプティマイザは、各レイヤーにおいて重みの更新を直交化し、特異値を効果的に均等化することで、大規模言語モデル（LLM）の事前学習において優れた性能を示してきました。この挙動は、スペクトルノルムの下での最急降下に対応しており、層をフラット化されたベクトルとしてではなく、線形作用素として捉え、ワーストケースの変化を制御します。しかし、Muonはパラメータ効率の良い微調整（PEFT）、特に低ランク適応（LoRA）においては滅多に使用されていません。

その核心となる数学的な障害は、Muonがレイヤー全体の重み行列に対する降下ステップを直交化する必要がある点にあります。LoRAの設定では、重みの更新が低ランクなパラメータ化 $\Delta W = BA$ （ここで $B \in \mathbb{R}^{d_1 \times r}$ および $A \in \mathbb{R}^{r \times d_2}$）に制約されているため、結果としての $\delta \Delta W$ が、Muonが要求するフルランクの直交化された勾配と完全に一致するような更新 $\delta B$ および $\delta A$ を一般的に見つけることは数学的に不可能です。

既存の回避策には以下が含まれます：

• Per-factor Muon: $B$ と $A$ を個別に直交化しますが、これはフルレイヤーの幾何学的構造を無視しています。
• LoRA-Muon: 三角不等式を用いて目的関数を緩和します。
• Riemannion: 低ランク多様体上のリーマン最適化として問題を定式化しますが、これはリトラクションステップにおいて高価な切断特異値分解（SVD）を必要とし、大幅な計算オーバーヘッドが生じます。

手法

著者らは、低ランク設定における緩和されたMuon目的関数の解を近似する手法として、sMuon（small Muon）を提案しています。このアプローチは、線形化と最小二乗最適化に基づき、最適な降下ステップを見つけ出します。

1. 線形化と最小二乗定式化

目標は、更新 $\delta B, \delta A$ によって導かれる更新 $\delta \Delta W \approx \delta B A + B \delta A$ が、理想的なMuonステップ $\delta \Delta W^* = -\eta \cdot \text{msign}(G_{\Delta W})$ に近似されるようにすることです。ここで $\text{msign}(X) = UV^\top$ はSVDの極因子（polar factor）です。

理想的なステップは一般にフルランクであり、低ランクの因子によって正確に実現することはできないため、著者らは二次項（$\delta B \delta A$）を無視して問題を線形化し、以下の最小二乗問題を解きます：
$$(\delta B^*, \delta A^*) = \arg \min_{\delta B, \delta A} \| \delta \Delta W^* - (\delta B A + B \delta A) \|_F^2$$
これにより、ムーア・ペンローズ擬似逆行列（$B^\dagger, A^\dagger$）を含む閉形式の解が得られます。重要な点は、因子 $(\delta B, \delta A)$ の観点からは最適な更新は一意ではないものの、フルレイヤーへの一次の変更である $\delta \Delta W$ は一意であるという点です。

2. 効率的な行列積のみによる計算

閉形式の解を直接実装する場合、大きな $d_1 \times d_2$ 行列を直交化する必要があり、計算コストが非常に高くなります。これを解決するために、著者らは大規模な行列の直交化を明示的に回避する効率的な定式化（アルゴリズム2）を導出しました。

• 低ランク勾配: フル勾配 $G_{\Delta W}$ の代わりに、最適化目的関数に対して一次同等である低ランク投影 $H = P_B G_{\Delta W} + (I-P_B)G_{\Delta W} P_A$ を利用します。
• 次元削減: 恒等式 $\text{msign}(H) = U \text{msign}(U^\top H V) V^\top$ を利用します。ここで $U$ と $V$ は、$H$ の列空間および行空間を張る正規直交基底です。$H$ のランクは最大 $2r$ であるため、これにより高価な $d_1 \times d_2$ の直交化を、$2r \times 2r$ の直交化へと削減できます。
• 実装: アルゴリズム全体は、行列乗算、行列の逆平方根、および小さな $r \times r$ または $2r \times 2r$ 行列の極分解のみに依存しています。これにより、Riemannionで必要とされるQR分解やSVD分解のオーバーヘッドを回避します。

3. アルゴリズムの詳細

sMuonオプティマイザは以下を組み込んでいます：

• 分割重み減衰（Split Weight Decay）: フルパラメータ学習のダイナミクスに合わせるため、勾配更新の前にスケーリング係数 $s = \sqrt{1-\lambda\eta}$ を用いて因子 $A$ と $B$ に重み減衰を適用します。
• モメンタム輸送（Momentum Transport）: パラメータ化の不変性（すなわち、特定の分解 $B, A$ に関わらずオプティマイザが一貫して動作すること）を確保するため、モメンタムバッファは毎ステップ現在の因子空間に投影されます。

主な貢献

1. sMuonの導出: スペクトルノルム正則化された最急降下目的関数の一次近似を取ることにより、LoRAアダプターを訓練するための閉形式の手法を導出しました。
2. 効率的な実装: 現代のハードウェアにとって計算効率が高く、高価な分解を回避する「行列積のみ」のアルゴリズムを提供します。これはPer-factor MuonやLoRA-Muonと同等の速度を提供しつつ、特にLoRAランクが大きい場合にはRiemannionよりも大幅に高速です。
3. 実証的評価: 4つのベースモデル（Muonで事前学習されたもの1つとAdamWで事前学習されたもの3つを含む）および11のベンチマーク（常識推論およびコーディング）にわたる広範な実験を行いました。

結果

論文では以下の知見を報告しています：

• 性能: sMuonはSFTベンチマークにおいて良好な性能を示しました。結果はモデルに依存しますが、Moonlight-16B-A3B（Muonで事前学習されたモデル）の微調整において、sMuonは11のタスク中6つのタスクで最高の精度を達成しました。全体として、AdamWおよびPer-factor Muonを上回っています。
• ReLoRA事前学習: ReLoRA（アダプターを定期的にマージおよび更新する手法）を用いた事前学習の設定において、sMuমはRiemannionと並んで最も低い最終検証損失を達成しましたが、速度は約30%高速でした。
• スペクトル特性: 重み更新スペクトルの実証分析により、sMuonはPer-factor MuonやLoRA-Muomと比較して、より一様なスペクトル（理想的なMuonの挙動に近い状態）を保持していることが示されました。
• 効率性: 表4は、s_MuonはAdamWよりはわずかに遅いものの、Riemannionよりも大幅に高速であることを示しています。ランク $r=64$ の場合、Riemannionはステップ時間に約33%のオーバーヘッドを追加しますが、sMuonの追加はわずか約0.8%です。

重要性と主張

本論文は、sMuonが、低ランクアダプターにMuon最適化を適用するための、原理に基づいたメモリ効率が良く、計算可能な方法を提供すると主張しています。LoRA因子を独立した行列として扱うのではなく、フルレイヤーの幾何学を尊重することで、sMuonは中程度の性能向上を実現しており、特にベースモデルがMuonで事前学習されている場合に顕著です。著者らは、性能を損なうことなく、Riemannian最適化の「高価な多様体制約」を回避できる点を強調しています。これにより、Muonベースの微調整が実用的なアプリケーションにおいてより利用しやすくなります。

著者らは、結果が「モデルおよび評価に依存する」ことを認めつつも、主要な貢献は、Muonの理論的利点とPEFTの実用的な制約との間の溝を埋める、クリーンで効率的な近似であると述べています。