Tim Lau

各層の最適化は、その構造や対称性に合わせて設計すべきであり、この観点からMuonを含むSpectral Optimizerを統一的に整理できる研究を紹介する。 T. Lau, W. Su, “Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers” まず通常の線形層を考える。 y = Wx ここでx, yは列ベクトル、Wは行列である。AdamでWを更新する場合、行列の各成分を独立したパラメータとして扱い、それぞれについてモーメンタムや勾配の二乗平均を計算し、更新幅を決定する。 この場合、更新はどの基底で行列を表したかに依存してしまう。 例えば、線形層の入力と出力の基底を直交変換（回転）して表したとしても、線形層そのものは同じ操作を表してほしい。 入力側を x -> Qx 出力側を y -> Py と回転させると、重み行列は W -> PWQ^T と書き換えられる。 実際、 (PWQ^T)(Qx) = PWx = Py である。 このとき勾配も同様に G -> PGQ^T と変換される。 このとき、最適化における更新写像をUとすると、 U(PGQ^T) = P U(G) Q^T を満たすのが自然である。 つまり、先に基底変換してから最適化を適用した結果と、最適化を適用してから基底変換した結果」が一致してほしい。これが同変性（equivariance）である。 Adamのように各座標に独立な非線形変換を適用するoptimizerは、一般にはこの性質を満たさない。そのため、同じ線形変換を表していても、基底を変えると最適化軌道が変わってしまう。 一方、Muonは勾配またはモーメンタムを行列として扱う。これらをまとめてDと書く。 Dを特異値分解して D = UΣV^T とすると、理想化したMuonの更新はpolar factor polar(D) = UV^T で表される。 すなわち、左特異ベクトルUと右特異ベクトルVは保ち、特異値だけを変換する。 このように、特異ベクトルを保存し、特異値だけを変換する写像をSpectral Operatorと呼ぶ。 一般にSpectral Operatorは U(D) = U Diag(Φ(σ(D))) V^T と書ける。ここでσ(D)はDの特異値を並べたベクトルであり、Φは特異値に対する変換である。 例えば勾配降下法では Φ(σ) = σ である。一方、理想的なpolar updateでは、非ゼロの特異値をすべて1に置き換えるため、 Φ(σ_i) = 1 とみなせる。 このpolar factorは polar(PDQ^T) = P polar(D) Q^T を満たす。このためMuonは、入力・出力の直交基底変換に対する同変性を保つ。 論文ではさらに、左右からの直交基底変換に対して同変かつ連続な更新写像がSpectral Operatorになることも示している。 この線形層以外の層は異なる形の対称性を持つ。 例えば、Embeddingを、行方向に各トークンの埋め込みベクトルが並ぶ行列として考える。 列方向のhidden featureについては基底を回転させてもよい。一方、行方向では任意の基底回転によって複数のトークンを混ぜることには意味がない。許されるのは、トークンの番号を入れ替える置換のみである。 したがって、Embeddingは、行方向には置換、列方向には直交変換に対する対称性を持つ。 これは通常の線形層が持つ左右両側の直交対称性よりも弱いため、例えば各行を独立に正規化するRowNormも対称性を壊さず使うことができる。 LM headもEmbeddingと似ており、語彙方向の置換とhidden featureの直交変換に対する対称性を持つ。 さらにLM headでは直後にSoftmaxを使うため、すべての語彙logitに同じ値を加えても出力確率が変化しない、という不変性がある。 そのため、語彙方向の平均成分は出力分布に影響しない冗長な方向であり、この平均成分を取り除いてから更新できる。 SwiGLUでは、中間層のニューロンを入れ替える置換に対して対称性がある。そのため、gate/up projectionでは中間ニューロンに対応する行方向の対称性、down projectionでは列方向の対称性に基づく更新が自然になる。 MoE Routerについても、expertの順番を入れ替える置換不変性と、Softmaxに由来する不変性が存在する。そのため、expert方向の平均成分を除去した上で、RowNormやLeft-Spectral updateを適用することができる。 このように、各層が持つ対称性をまず調べ、その対称性を保つように最適化法を設計する、というのが本論文の基本的な提案である。 コメント === 各パラメータの更新について、対称性という観点から最適化法を導出することを提案した興味深い研究である。 Muonなどによって、パラメータを要素ごとではなく行列として扱う方向性が注目されるようになったが、この論文ではさらに一歩進み、その行列がネットワーク内でどのような意味を持ち、どのような対称性を備えているかまで考慮して最適化法を設計することを提案している。 これまで、Muonについて実験的に示されている重要な利点の一つとして、大きなバッチサイズでも性能劣化が比較的小さく、大規模学習での学習効率を改善できる可能性があることが挙げられる。 また、低い特異値に対応する方向を相対的に強く更新することから、Long-tailな知識や低頻度方向の学習に有利なのではないかという仮説も考えられている。 こうしたMuonの性質を、単に特異値を均すからという観点だけでなく、モデルが持つ対称性と最適化の同変性から見る方向は興味深い。 こうした考え方は、パラメータの座標表示ではなく内在的な幾何に基づいて最適化するべきという、自然勾配法の流れともつながっているが、今回は各パラメーターブロックが持つ具体的な対称性から最適化を設計する方向を進めたものとみることができる。