Muonから勾配クリッピングまで: QK安定性に関する考察

前置き：エレガントな理論と厄介な問題深層学習のツールボックスにおいて、オプティマイザーはエンジンです。Adamとその派生形を調整することに慣れている人がほとんどですが、全く異なる哲学に基づいた別のオプティマイザーが、研究者の注目を徐々に集めています。それがMuonオプティマイザーです。Keller JordanとJeremy Bernsteinが、彼らの研究[1]で最初に提案しました。その深遠さは、もはや馴染みのあるパラメータ空間にとどまらないという点にあります。代わりに、関数空間の視点から出発し、そこから更新ルールを構築します。このアイデアは最近、Kimi K2[7]のような研究で進歩し、成功裏に実践されており、大きな可能性を示しています。しかし、K2に関する公開資料は、困難な問題も指摘しています。すなわち、オリジナルのMuonオプティマイザーをTransformerのQuery (Q) と Key (K) の重み行列に直接適用すると、学習が非常に不安定になり、崩壊することさえあります。理論のエレガンスは、実践において壁にぶつかります。これは自然に2つの疑問を提起します。この現象の根本原因は何でしょうか？Muonの第一原理に従って、QKに適した修正された更新アイデアを提案できるでしょうか？この投稿は、その道を考え抜こうとする私の試みを記録したものです。それは曲がりくねった旅です。厳密な理論的導出から、エンジニアリングのボトルネック、そして最終的には実行可能なヒューリスティックなスキームまで。それは洗練された答えというよりは、理論、エンジニアリング、直感を横断する思考の記録です。1. 理論レビュー：標準的なMuonがQKに適合しない理由この不一致を理解するために、まずMuonの本来の意図と、それがAdamのようなオプティマイザーと根本的にどのように異なるかを理解する必要があります。Muonの導出と本質すべては、最適化の最も基本的な目的から始まります。重み行列Wについて、損失Lに対する勾配Gを計算したとします。損失を可能な限り減らす更新ΔWを見つけたいとします。数学的には、一次テイラー展開は次のように述べています。 ΔL ≈ ⟨G, ΔW⟩ = Tr(GᵀΔW)。 Lを可能な限り速く減少させるためには、ΔWがGの反対方向に可能な限り多く向くようにしたいです。これは制約付き最適化問題として形式化できます。 min ΔW Tr(GᵀΔW)　ただし、ΔWのサイズ制約付き。 明らかに、制約がなければ、ΔWのエントリは無限大になる可能性があり、それは無意味です。すべてのオプティマイザーは、更新ステップの「サイズ」を制限する何らかの制約を課さなければなりません。違いは、このサイズがどのように定義されるかです。パラメータ空間における従来の制約は、SGD、Adam、および類似のオプティマイザーから知っているものです。これはΔW自体に直接課されます。最も一般的な選択はフロベニウスノルムです。 ‖ΔW‖_F ≤ η。 その幾何学的な意味は直感的です。パラメータ空間でのユークリッド距離の移動は大きすぎてはなりません。Muonの制約は、関数空間におけるもので、全く異なります。Muonは、パラメータWがどれだけ移動したかを主に気にかけるべきではないと言います。モデルの関数がどれだけ変化するかを気にかけるべきです。線形層の場合、関数は入力xを出力Wxにマッピングします。したがって、パラメータ更新の関数的効果は、出力の変化ΔWxです。Muonの核心的な考え方は、この関数的変化を制限すべきであり、その制限はすべての可能な入力に対して有効であるべきだということです。xの大きさの影響を取り除くために、単位ノルムの入力のみを考慮します。したがって、Muonの制約は次のようになります。 sup ‖x‖₂=1 ‖ΔW x‖₂ ≤ η。 これは、長さ1の任意の入力ベクトルに対して、更新行列ΔWによる変換の後、出力ベクトルの長さがηを超えてはならないことを意味します。行列理論を学んだことがある人なら誰でも、この上限がスペクトルノルム‖ΔW‖₂の定義そのものであり、行列の最大特異値に等しいことを認識するでしょう。これでMuonの最適化問題は明確になりました。 min ΔW Tr(GᵀΔW)　ただし、‖ΔW‖₂ ≤ η。 スペクトルノルム制約下でのこの最適化問題は、美しい解析解を持ちます。それは特性行列符号関数です。 ΔW* = -η・msign(G)。 msign関数はGの特異ベクトルを保持しますが、そのすべての特異値を1に変更します。これは方向制御の極端な形態です。更新「エネルギー」をすべての特異方向全体に均等に分配します。QK更新における幾何学的な衝突Muonのメカニズムを念頭に置いて、QK更新での失敗を再検討できます。Su Jianlinが彼のブログ記事[2]で分析しているように、問題はアテンションの双線形性にあります。比較的浅い直感は、関数的な結合です。QとKの関数はペアとして現れます。それらはQKᵀの積を通してのみ意味のあるアテンションスコアを生成します。W_QとW_Kの関数的変化を独立に制約しても、それらの結合された積の関数的変化が制御されることは保証されません。Su Jianlinの仮説[2]に続く、より深い幾何学的な説明は、Muonのフルランク更新の挙動に関連しています。実践で見られる「MaxLogit爆発」は、しばしばW_QまたはW_Kのスペクトルノルムが爆発の兆候を示していることを意味します。したがって、問題は次のようになります。なぜMuonはスペクトルノルムを爆発させる可能性が高いのでしょうか？スペクトルノルムは最大特異値なので、これは次のようになります。なぜMuonは特異値を増加させる傾向があるのでしょうか？MuonとAdamの違いは何でしょうか？Muonの更新はmsignによって処理されるため、更新のすべての特異値は等しくなります。この意味で、その実効ランクはフルです。対照的に、一般的な行列は通常、異なるサイズの特異値を持ち、最初の数個によって支配されます。実効ランクの観点からは、低ランクです。Adam更新についても同様の低ランク仮説を立てることができます。この仮説は新しいものではありません。高次のmuPは、Adam更新の低ランク性について同様の仮説を立てています。数式で表すと、パラメータWのSVDをUΣVᵀ、Muon更新のSVDをUmΣmVmᵀ、Adam更新のSVDをUaΣaVaᵀとします。更新された重みは模式的に次のように書くことができます。 W_new = W + ΔW。 もしWの特異ベクトルペアが更新の特異ベクトルペアに近い場合、対応する特異値は直接加算される傾向があり、その方向でのWの特異値が増加します。Muonの更新はフルランクであるため、そのエネルギーはすべての特異方向全体に均等に分散されます。したがって、低ランクのAdam更新よりもWの特異ベクトルと「衝突」する可能性がはるかに高くなります。Muonは自然にパラメータの特異値を増加させる可能性が高くなります。この効果はアテンションで増幅されます。アテンションスコアの核心は、QKᵀの双線形形式です。もしMuonがW_QとW_Kの両方のスペクトルノルムを同時に増加させている場合、最終的な内積が計算されるときに増加は乗算されます。爆発のリスクは急激に上昇します。「悪貨は良貨を駆逐する」というフィードバックループが形成され、最終的に学習が崩壊する可能性があります。この時点で、これほどエレガントに設計されたオプティマイザーがTransformer QK重みには不向きになる理論的な説明が得られました。根本的な問題は、2つの独立した変更を制約するのではなく、積項の変更を直接制約する必要があるということです。この洞察は、次のステップの出発点となります。2. 原則に基づいた試み：QKのためのMuonのカスタマイズ最初のセクションで、重要な洞察が得られました。W_QとW_Kへの更新を個別に制約するだけではうまくいかないということです。アテンションでは、この2つは関数的に結合されています。最終的に重要なのはそれらの積です。独立した制約は、2つのベクトルのノルムをほとんど変化させないようにするかもしれませんが、それらの方向が整列することで、それらの内積、つまりアテンションスコアが爆発する可能性があります。したがって、Muonの哲学に真に従う原則に基づいた修正は、根本的な問題から始める必要があります。すなわち、直接制約する