Muon オプティマイザが Transformer の QK 重みで失敗する理由と、その解決策

From Muon to Gradient Clipping: Some Thoughts on QK Stability

Muon オプティマイザは、関数空間の視点から設計された新しい最適化手法で、Kimi K2 などの研究で有望視されています。しかし、Transformer の Query と Key の重み行列に直接適用すると、トレーニングが不安定になり、崩壊することさえあります。この記事では、その根本原因を理論的に分析し、Muon の原理に従った修正案を提案します。標準的な Muon はパラメータ空間ではなく関数空間で更新を制約しますが、QK の更新では、Q と K の積であるアテンションスコアの変化を直接制約する必要があることを示します。さらに、提案手法の実装上の課題と、実用的なヒューリスティックな解決策についても考察します。