Muon オプティマイザと Transformer における QK 更新の安定化への探究

1. はじめに：理論と実践の間での壁

深層学習においてオプティマイザはエンジンのような存在である。Adam などの標準的な手法に加え、Muon オプティマイザ という新たなアプローチが注目を集めている。これはケルラー・ジョルダンらによって提唱され、パラメータ空間ではなく機能空間（関数空間） の視点から更新ルールを設計したものである。

Kimi K2 などの最新研究で実装が進みつつあるが、重要な問題が浮上している。Muon をそのまま Transformer のクエリ・キー（Q, K）重みに適用すると、トレーニングが不安定になり、最悪の場合に収束しない（クラッシュする） 恐れがある。

本記事では以下の問いについて考察する：

• この現象の根本原因は何なのか？
• Muon の第一原理に従い、QK に適した改変された更新手法を提案することはできるのか？

これは理論的導出からエンジニアリングのボトルネックへ、そして実行可能なヒューリスティックな手法へと至る、巡り巡る旅路の記録である。

---

2. 理論的レビュー：なぜ標準的な Muon が QK に適合しないか

Muon の不整合を理解するには、Adam などのオプティマイザとの根本的な違いを知る必要がある。

Muon の導出と本質

最適化の目標は損失 $L$ を最小化することである。パラメータ $W$ の更新量 $\Delta W$ を決定する際、制約付き問題として形式化する。

• 伝統的オプティマイザ（SGD/Adam）： 更新の大きさをパラメータ空間で制限する。 $$ |\Delta W|_F \le \eta $$ これはパラメータ自身のユークリッド距離を制限するものである。

• Muon オプティマイザ： 更新の大きさを機能空間で制限する。 パラメータがどれだけ動くかよりも、モデルの関数（入力 $x$ に対する出力変化）がどれだけ変化するかに焦点を当てる。 $$ \sup_{|x|_2 = 1} |\Delta W x|_2 \le \eta $$ これはスペクトルノルム $|\Delta W|_2$ を制限することを意味する。

Muon の更新解は、行列の符号関数（matrix sign function,

msign

）を用いて表され： $$ \Delta W^* = -\eta , \operatorname{msign}(G) $$ ここで、特異値が等しくなるフルランクな更新となる。

QK 更新における幾何学的対立

Muon を Transformer の注意機構（Q, K 行列）に直接適用すると不安定化する理由には、2 つの要因がある。

1. 機能的結合（Functional coupling）
   • 注意力スコアは $QK^T$ という積項によって定義される。
   • $W_Q$ と $W_K$ それぞれを独立に制約しても、その積 $QK^T$ の変化が制御されるとは限らない。
2. フルランク更新挙動と MaxLogit の爆発
   • Muon は更新行列の特異値をすべて均等に扱う（フルランク）。
   • これに対し、パラメータ $W$ 自体や Adam の更新は低ランクである傾向がある。
   • 衝突の確率： Muon のフルランクエネルギーが、低ランクな既存のパラメータ構造と「衝突」しやすい。
   • 結果： この衝突によりパラメータの特異値が増大し、特に注意機構では $W_Q$ と $W_K$ の両方が増大することで内積計算時に爆発的成長が起きる（MaxLogit の爆発）。

---

3. 原理に基づく試み：Muon を QK 向けにカスタマイズ

単なるパラメータの独立制約は機能しない。重要なのは積項 $QK^T$ そのものの変化を制御することだ。

2.1 新しい最適化目的：パラメータから機能への転換

重み行列 $W_Q, W_K$ の直接更新ではなく、注意力スコア行列 $S = QK^T$ の変化 $\Delta S$ を制約する。

• 変化の線形近似： $$ \Delta S \approx X(\Delta W_Q W_K^T + W_Q \Delta W_K^T)X^T $$ （2 次項は無視）
• 目標制約： 最悪ケースの変化を制限する。 $$ |\Delta W_Q W_K^T + W_Q \Delta W_K^T|_2 \le \epsilon $$

2.2 完全な最適化問題

勾配降下法の原理より、損失変化 $\Delta L$ を最小化しつつ上記制約を満たす： $$ \min_{\Delta W_Q, \Delta W_K} \operatorname{Tr}(G_Q^T \Delta W_Q) + \operatorname{Tr}(G_K^T \Delta W_K) \quad \text{s.t.} \quad |\Delta W_Q W_K^T + W_Q \Delta W_K^T|_2 \le \epsilon $$

2.3 解への道筋と挫折

この問題の複雑さから、以下の段階的な試みがなされたが、すべて計算コストの壁にぶつかった。

試み A: 分離と擬逆行列の導入

結合項を三角不等式で分離し、独立な部分問題を解くアプローチ。

• 導出： 変数置換 $Z = \Delta W_Q W_K^T$ を用い、Moore-Penrose 擬逆行列 $(W_K^T)^+$ を導入して制約を解除した。
• 問題点： 計算に高価な特異値分解（SVD）と擬逆行列計算が必須となり、実装不可能である。

試み B: 近似による簡素化の失敗

擬逆行列を使わず、行列ノルムの乗法性を用いて単純化したアプローチ。

• 問題点： $W_K$ の大きさのみの情報しか考慮できず、$\Delta W_Q$ と $W_K$ の幾何学的方向情報の相互作用（射影効果）を捨象してしまっている。

試み C: SVD を回避する反復法

擬逆行列を計算するために必ずしも SVD を使わない、Newton-Schulz 反復法などの検討。

• 問題点： 低精度浮動小数点数（bfloat16 など）では数値的不安定性が高まる。キャッシュやウォームスタートによる工夫はあるが、理論的な完全性を損なう代償が高い。

結論として、理論的には完璧だが実装にほぼ不可能な更新式が導出されただけで終わった。

---

4. 転換点：双線形注意の再考と低ランク近似

複雑な解析解を追求する限り、壁は高くなる。そこで視点を変える必要がある。

思考実験：統合された双線形行列

注意力スコア $S = XBX^T$ とし、相互作用行列 $B = W_Q W_K^T$ を統一して扱う。

• この場合、Muonn の制約 $|\Delta B|_2 \le \epsilon$ がそのまま適用できる。
• 解は単純な標準 Muon 更新となる：$\Delta B^* = -\epsilon \operatorname{msign}(G_B)$。

現実のコストとパラメータ爆発

この「統合行列 $B$」アプローチが採用されない理由は効率性にある。

• パラメータ数：
  • 標準（独立）: $2dh$
  • 統合：$d^2$
  • （例：$d=4096, h=128$ の場合、統合は約 16 倍のパラメータが必要）。
• 計算量： $XBX^T$ の計算も極めて重い。

結論：Muon-LoRA への架け橋

標準的なマルチヘッド注意（MHA）は、フルランク双線形注意の低ランク近似である。 モデルアーキテクチャが自動的に低ランク制約 $B \approx W_Q W_K^T$ を強制している。

この視点から新たな可能性が開ける：

1. Muon-LoRA： 大規模事前学習済みモデル $W_0$ のファインチューニングにおいて、更新分 $\Delta W = AB$ （低ランク）として扱う。Muonn の機能空間のアイデアをそのまま適用できる。
2. 解析解の見直し： 複雑な擬逆行列や射影が必要なのは、低ランク構造を保ちつつ最適化しようとしたからである。低ランク空間内で直接単純な近似を検索すればよい。

---

5. インタールードと次のステップ

これまでの理論的探求は計算コストが過大であることが分かった。理論的な完璧さを追求するのではなく、エンジニアリング的実現可能性を優先すべきだ。

5.1 経験的近似へ：MaxLogit の活用

理論的に「すべての入力に対する最悪ケース」を制約することは高コストだが、トレーニング中の不安定化の直結原因は「現在のバッチでの MaxLogit の爆発」である。

• 代替案： 理論的なスペクトルノルム推定として、経験的な最大ログリット（MaxLogit） を使用する。 $$ M = \max_{i,j} (\log \text{softmax}(S)_{ij}) $$
• 補正戦略（MuonClip）：
  1. 現在のバッチで計算した $M$ と閾値 $\tau$ を比較。
  2. $M > \tau$ の場合、スケーリング因子 $s = \tau/M$ を計算。
  3. 事後補正として重みを正規化する（例：$W_Q \leftarrow s^\alpha W_Q$）。

5.2 手法の比較と哲学

特徴	グラディエントクリッピング	MuonClip (Kimi AI)
介入タイミング	バックパス後、更新前	更新後、フォワード計算へ干渉
メカニズム	極端な勾配値を単純に切り捨てる	スケーリングによる温度調整（出力分布の平滑化）
作用範囲	ローカルな変化量制限	グローバルな観測量制限

5.3 課題と展望

• 推定精度： 有限サンプルから最悪ケースを推定するが、巨大なバッチサイズ（数百万トークンペア）により十分信頼性が高い。
• 未知数： クリッピング閾値やバランス係数などの超パラメータ感度、ステートフルオプティマイザ（Adam）との相互作用などについて、さらに実験検証が必要である。

---

結論

Muon オプティマイザの機能空間への着想は素晴らしいが、Transformer の QK 結合に対する直接適用には理論と実践のギャップが存在する。

1. 完全な解析解： 擬逆行列や高次元射影を必要とし、実装不可能である。
2. 現実的な解決策： MuonClip のような、経験的な MaxLogit を用いた事後補正が、理論的理想に近い安定化効果を与える。

低ランク近似の視点（LoRA や MHA の構造）と組み合わせることで、パラメータ効率の良い新しいファインチューニングオプティマイザ（Muon-LoRA など）の開発が可能となる。理論の美しさとエンジニアリングの実用性のバランスを踏まえたアプローチが求められる。

---

参考文献

1. Jordan, K., Bernstein, J. et al. (2024). "Muon: An optimizer for hidden layers in neural networks."
2. Su Jianlin. (2025). "QK-Clip: Letting Muon Go Further on the Road to Scaleup."
3. Su Jianlin. (2025). "Muon Sequel: Why Did We Choose to Try Muon?"
4. Liu, J., Su, J. et al. (2025). "Muon is Scalable for LLM Training." arXiv preprint.
5. Jordan, K. et al. (2024). "modded-nanogpt: Speedrunning the NanoGPT baseline." GitHub.
6. Jordan, K. et al. (2024). "Muon Optimizer." GitHub.
7. Moonshot AI. (2025). "Kimi K2: Open Agentic Intelligence." Project page.