技術要約: LoRA-Muon: 低ランク多様体上のスペクトル最急降下法

1. 問題提起

低ランク適応（LoRA）は、計算量（FLOPs）とメモリ要件を削減する効率性の高さから、大規模ディープラーニングモデルのファインチューニングに広く採用されています。しかし、AdamWのような標準的なオプティマイザをLoRAの因子（$A$ および $B$）に対して独立に適用すると、全結合（dense）トレーニングと比較して、以下のような重大なチューニングの困難が生じます：

• 初期化への敏感さ: パフォーマンスが初期値の選択に強く依存する。
• ハイパーパラメータの転移性の低さ: 最適な学習率が、ランク、幅、深さ、または因子のスケーリング規則を越えてうまく転移しない。
• ダイナミクスの乖離: 合成された重み（$W = AB^T$）が同一の状態から開始した場合でも、因子のスケールが離れてドリフトすると、学習ダイナミクスが乖離することがある。
• サブオプティマルな性能: ネイティブな低ランク事前学習は、追加の安定化なしには、全結合のベースラインに匹敵する性能を得るのが困難である。

核心となる課題は、全結合オプティマイザ（特にMuonおよびShampooファミリー）の幾何学的特性を維持しつつ、ランク、幅、深さ、および因子の再スケーリング間で学習率が最適に転移することを保証する、低ランク多様体上のオプティマイザを設計することです。

2. 手法

2.1 幾何学的基礎

本論文は、オプティマイザの設計を、あるオプティマイザがどのノルムの下で最急降下を行うかを特定するという幾何学的な問題として再定義しています。

• Muonオプティマイザ: アンビエント行列空間 $\mathbb{R}^{m \times n}$ におけるスペクトルノルム（$\|\cdot\|_{2\to2}$）の下で最急降下を行います。更新方向は行列サイン関数 $-\text{msign}(G)$ によって与えられます。
• 低ランク多様体: LoRAは、重みを $\mathcal{M}_r = \{W \in \mathbb{R}^{m \times n} : \text{rank}(W) = r\}$ という多様体に制限します。著者らは、因子 $A$ と $B$ を独立した座標として最適化するのではなく、この多様体の接空間 $T_W \mathcal{M}_r$ 上で直接最急降下問題を解くことを提案しています。

2.2 LoRA-Muonの導出

著者らは、スペクトルノルムの下での低ランク接空間における制約付き線形最小化問題を解くことにより、LoRA-Muonを導出しています。

1. 接空間の分解: $W$ への更新は、$\Delta W = \Delta A B^T + A \Delta B^T$ に分解されます。
2. 信頼領域の分割: 問題を扱いやすくするため、信頼領域の予算 $\eta$ を2つの接成分に均等に分割します。これは、訓練中に2つの成分が急速に整列し、三角不等式がほぼタイトになるという経験的な観察に基づいています。
3. 射影形式: 単位不変性を用いて、問題は投影された勾配に対する線形最小化オラクル（LMO）を見つける問題へと簡略化されます。理想的な更新は以下の通りです：
   $$\Delta W = \frac{\eta}{2} \text{LMO}(\nabla W f \cdot P_B) + \frac{\eta}{2} \text{LMO}(P_A \cdot \nabla W f)$$
   ここで、$P_A$ と $P_B$ は $A$ および $B$ の列空間への射影器です。
4. スペクトルの特化: スペクトルノルムの場合、LMOは $-\text{msign}(\cdot)$ です。結果として得られる因子更新（式14–15）は、現在のグラム行列（$S_A = A^TA, S_B = B^TB$）によって因子勾配をホワイトニングし、スペクトルサイン関数を適用することを含みます：
   $$\Delta A^* = -\frac{\eta}{2} \text{msign}(\nabla_A f \cdot S_B^{-1/2}) S_B^{-1/2}$$
   $$\Delta B^* = -\frac{\eta}{2} \text{msign}(\nabla_B f \cdot S_A^{-1/2}) S_A^{-1/2}$$

2.3 分割重み減衰（Split Weight Decay）

標準的な因子ごとの重み減衰は、全結合の更新ダイナミクスと低ランクの更新ダイナミクスの間の不一致を引き起こします。著者らは、減衰係数 $s = \sqrt{1 - \lambda \eta}$ が非対称に適用される分割重み減衰ルールを導出しました：
$$A_{t+1} = s A_t + \frac{1}{s} \Delta A^*_t, \quad B_{t+1} = s B_t + \frac{1}{s} \Delta B^*_t$$
これにより、合成された重みの更新が所望の全結合ダイナミクスと一致し、ハイパーパラメータの転移性が保持されます。

2.4 ゲージ不変性（Gauge Invariance）

LoRA-Muonの重要な特性は、ゲージ不変性です。合成された重み $W$ への更新は、任意の可逆変換 $(A, B) \to (AR, BR^{-T})$ に対して不変です。これは、更新が射影器 $P_A$ および $P_B$ にのみ依存しているために達成されます（これら自体もゲージ不変です）。

3. 主な貢献

1. LoRA-Muonの導出: 単位不変ノルム（スペクトルノルムに特化）の下で低ランク多様体上の最急降下問題を解くことによって導出された、新しいオプティマイザ・ファミリー。
2. 分割重み減衰ルール: 低ランクの訓練ダイナミクスを全結合の訓練ダイナミクスに適合させ、ハイパーパラメータの転移を可能にする特定の更新ルール。
3. ゲージ不変性の証明: LoRA-Muonの更新が完全なゲージ群作用に対して不変であることを示す理論的証明。これにより、オプティマイザが任意の因子座標ではなく、誘導された重みの更新に依存することが保証される。
4. Spectronとの比較: Spectronはゲージ不変ではなく、任意の因子の再スケーリングに対して敏感であることを示したこと。
5. LoRA-RITEとの関係: LoRA-RITEの簡略化されたQR座標コアが、代数的にLoRA-Muonのスペクトル更新と同一であることを証明。ただし、LoRA-MuonはQR分解や二次モーメントの保存に伴う計算オーバーヘッドを回避している。
6. 実証的検証: LoRA-Muonが、計算量が一致した全結合および低ランクのベースラインに対する効果的なプロキシとして機能することを示す証拠。

4. 実験結果

実験は、TinyShakespeareデータセットを用いた文字レベルの言語モデリングにおいて、計算量を一致させた全結合および低ランクのベースラインを用いて行われました。

• 学習率の転移: 全結合のMuonの最適学習率（$lr = 0.1$）は、ランク（2, 4, 8, 16, 32）、幅、深さ、およびスカラー因子の再スケーリングのスイープを通じて、LoRA-Muonへ完璧に転移しました。
  • ランク2において、LoRA-Muonは全結合のベストテストされた学習率を再現しました。
  • ランク32において、Loora-Muonはシード平均スイープにおいて、全結合ベースライン（$1.789 \pm 0.002$）よりも低い平均検証損失（$1.776 \pm 0.002$）を達成しました。
• ゲージ再バランシング: 因子の初期化において大きなスカラー不均衡（スケール9）を与えた場合、再バランシングステップを適用してもLoRA-Muonへの影響は無視できる程度であり（損失の変化 $\approx -5.1 \times 10^{-5}$）、ゲージ不変性が確認されました。対照的に、Spectronのパフォーマンスは大きく変化し（$2.0 \times 10^{-2}$）、因子のスケーリングに対する感受性が確認されました。
• 効率性: ランク2のプロキシは、全結合モデルの訓練可能パラメータのわずか約4.3%しか使用せず、同じ計算制約の下で約9.2倍大きなオプティマイザ・ステップ・予算をサポートしています。

5. 意義と主張

本論文は、LoRA-MuonがLoRAのチューニングにおける困難に対する理論的に裏付けられた解決策を提供すると主張しています。LoRAを因子化された座標系としてではなく、低ランク多様体上の最適化問題として扱うことで、著者らは以下を実現しました：

• 堅牢性: オプティマイザは、任意の因子の選択（ゲージ対称性）に対して不変です。
• 転移性: 最適な学習率（および、スケーリング則を通じた他のハイパーパラメータ）が、ランク、幅、深さ、および因子のスケールを越えてシームレスに転移します。
• 効率性: QR分解や二次モーメントの保存といった高価な処理を行うことなく、スペクトル更新を計算でき、アクセラレータフレンドリーです。

著者らは、LoRA-Muonを全結合のMuonやShampooファミリーの優れた低ランク・プロキシとして位置付けています。彼らは、低ランクのパラメータ化がオプティマイザの幾何学を決定すべきではなく、代わりに更新は低ランク行列多様体自体の幾何学によって定義されるべきであると強調しています。

限界: 著者らは、現在の実証的な主張はTinyShakespeare規模のモデルに基づいており、大規模な事前学習やダウンストリームのファインチューニングのベンチマークではまだ検証されていないことを述べています。本研究は、完全なハイパーパラメータ・スイープではなく、学習率の転移に焦点を当てています。