MuonオプティマイザはAnimaやSDXLのLoRA学習に使えるのか、実装前に調べた

AnimaやWAI-IllustriousのキャラLoRAを焼き続けているが、オプティマイザを自分で選んだことは一度もなく、ツール既定のまま回してきた。
一方LLM側では、Kimi K2やDeepSeek V4の事前学習でMuonというオプティマイザの採用が続いている。
これを画像LoRA学習で使えないか、実装前の下調べとして仕組みと適用条件を確認した。

Muonは更新行列を直交化するオプティマイザ

AdamWはパラメータを1要素ずつ扱う。
勾配の1次・2次モーメントを要素ごとに持ち、要素ごとに更新の歩幅を調整する。
行列のどの位置にある要素か、周りとどう並んでいるかは更新に影響しない。

Muon（MomentUm Orthogonalized by Newton-Schulz）は、重み行列への更新を1枚の行列としてまるごと扱う。
モメンタムに累積した勾配行列を、Newton–Schulz反復という行列積だけの演算で直交行列に近づけてから、重みに適用する。

flowchart TD
    G[ミニバッチ勾配 G] --> M[モメンタムへ累積]
    M --> NS[Newton–Schulz反復で直交化<br/>更新行列を UVᵀ に近づける]
    NS --> S[層の形状に応じてスケール]
    S --> W[重み行列を更新]

直交化には、更新行列の特異値を均す作用がある。
Muonの作者は、素の更新行列が少数の支配的な方向へ偏りやすいことを観測していて、直交化はこの偏りを打ち消す。
弱い方向にも歩幅を配ることが改善につながる、という機序の説明は、作者自身も現時点では仮説の扱いにしている。

もともとは2024年にKeller JordanがNanoGPTの学習速度競争用に出した手法で、Moonshot AIがMuon is Scalable for LLM Trainingで大規模化を検証した。
その後Kimi K2が派生のMuonClipで15.5Tトークンを学習し、Kimi K3はPer-Head Muonを採用、DeepSeek V4も事前学習の大部分のモジュールでMuonを使っている。

ただしMuonが主に対象とするのは隠れ層の2次元重み行列で、4次元の畳み込みも2次元へ並べ直せば適用はできる。
埋め込みの重みは2次元だが元の実装では経験的にAdamW側とされていて、LayerNorm/RMSNormのスケールやバイアスのような1次元パラメータもAdamWに残す。
つまり最初から、一部の層をAdamWに残す振り分けが前提の設計になっている。

Muonと並んで比較によく出てくるのがSOAPだ。
Shampoo系の手法で、勾配行列 $G$ から行側 $GG^\top$ と列側 $G^\top G$ の統計を持ち、その固有基底の上でAdamを回す。
前処理に使う情報はMuonより多いが、そのぶんオプティマイザが持つ状態のメモリも増える。

	AdamW	Muon	SOAP
更新の単位	要素ごと	行列ごと	固有基底上で要素ごと
対象パラメータ	全部	主に隠れ層の行列	主に行列
追加で持つ状態	1次・2次モーメント	モメンタムのみ	固有基底 + モーメント類
ステップ当たりの追加計算	なし	Newton–Schulz反復	固有分解の定期更新

状態がモメンタムだけで済むぶん、オプティマイザのメモリはMuonの方がAdamWより軽い。

AnimaとSDXLでかけられる層が違う

AnimaをSDXL系のつもりで考え始めたが、これは前提から間違っていて、初検証の記事に自分で書いた通りベースはNVIDIAのCosmos-Predict2 2B、SDXLとは別系統だ。
Cosmos-Predict2はセルフアテンション、クロスアテンション、フィードフォワードを重ねたDiT（Diffusion Transformer）で、時刻条件はadaptive LayerNormで入る。

Cosmos-Predict2のDiTブロックには、クロスアテンションや時刻条件のadaptive LayerNormなど、LLMには出てこない層も入っている。
それでもQ/K/V/出力の射影やフィードフォワードといった主要な重みは2次元の線形層で、Muonがそのまま適用できる形をしている。
埋め込みや正規化層やバイアスをAdamWへ残すLLMでの振り分けも、初期案としてそのまま流用する。

一方、WAI-IllustriousなどSDXL系はU-NetにConv2Dが混ざる。
畳み込みの重みは $[C_{out}, C_{in}, k_h, k_w]$ の4次元テンソルなので、Muonをかけるなら $[C_{out},\ C_{in} k_h k_w]$ の2次元に並べ直して直交化し、元の形へ戻す。
Muonの元実装も、ConvNetでは最初の畳み込みを除いてこの並べ直しを使っている。

ただ、2次元に並べ直せることと、その行列の特異値を均すのが良い更新であることは別問題だ。
カーネルの縦横2軸と入力チャネル軸を1本にまとめた行列では、直交化が何を均しているのかが変わる。
SDXL系ではまずアテンションと線形層だけMuonにして、畳み込みはAdamWのまま残す。

整理すると振り分けの初期案はこんな感じ。どの行も実測で確かめる前の設計案だ。

パラメータ	Anima（DiT）	SDXL（U-Net）
アテンション Q/K/V/出力	Muon	Muon
フィードフォワード / MLP	Muon	Muon
時刻・条件の射影	Muon候補	Muon候補
Conv2D	ほぼ登場しない	AdamW
埋め込み	AdamW	AdamW
正規化層・バイアス	AdamW	AdamW

拡散モデルとViTでの実測報告

拡散系の目的関数でMuonがどうなるかは、Optimization Benchmark for Diffusion Models on Dynamical Systemsに同条件比較がある。
対象は力学系の軌道を学習する小規模な拡散・フローモデルで、1024エポック回した最終lossはMuonとSOAPがAdamWより18%低かった。
一方で1エポックにかかる時間はMuonが約1.45倍、SOAPが約1.72倍で、最適学習率はAdamWのおよそ2倍と報告されている。
Stable Diffusion級の画像生成で品質が18%上がる話とは別物だが、拡散の損失でもMuon系が下げ余地を持つというデータにはなる。

画像認識側ではMuon in Vision Transformersが、ViTの分類でMuonがAdamWを上回ったと報告している。
ただしデータ拡張を盛った設定ほど差が開くという但し書きがあり、効果は学習レシピとの組み合わせに依存している。

LoRAの二因子分解との相性

自分の用途はLoRA学習なので、Muonとの相性が一番はっきりしないのはここだ。

LoRAは重み本体 $W$ を凍結し、低ランク分解 $\Delta W = BA$ の小さい2枚だけを学習する。
$A$ と $B$ はそれぞれ2次元行列だから、個別にMuonをかけること自体はできる。
ただし生成結果に反映されるのは積の $\Delta W$ で、2枚を別々に直交化した更新が積として見ても良い更新になるのかは自明ではない。

LoRA meets Riemannionは、 $BA = B'A'$ となる分解の取り方が無数にあるというLoRAの不定性に注目し、固定ランク行列の多様体上で直接最適化するRiemannionを提案した。
Muonをリーマン幾何の側へ一般化した位置づけで、LLMと拡散モデルの両方で素のLoRAより収束と最終性能が改善したと報告している。
拡散側の実験はStable Diffusion 2に数枚の参照画像で特定の被写体を覚えさせるもので、LoRAとRiemannionをランク4/8/16で比較している。400ステップの定量比較はランク4、ランク別の可視化は600ステップという規模だ。
参照画像が数枚、ランク4/8/16、特定の被写体を覚えさせるという構成は、キャラLoRAの設定に重なる部分が多い。

Uniform Spectral Growth and Convergence of Muon in LoRA-Style Matrix Factorizationは理論寄りの結果で、 $A$ と $B$ に別々にMuonをかけても、積 $\Delta W$ の特異値がスペクトル全体でほぼ均等に成長することを簡約した設定で示した。
さっき挙げた、積として良い更新になる保証がないという懸念への部分的な回答になっている。

それでも、数十枚の画像と数千ステップで焼くAnima系キャラLoRAそのものを検証した報告は見つからなかった。
そこは自分で回して確かめるしかない。

差し替えで何が変わりうるか

LoRA学習の1ステップはモデル本体の順伝播・逆伝播が支配的なはずで、オプティマイザを差し替えても学習時間の短縮は見込めない。
Newton–Schulzのぶんステップ当たりの計算はむしろ増えるし、さっきの拡散ベンチでも1エポックにかかる時間はMuon側が長かった。

同じステップ数を回したときに、より低い損失に落ち着くことを期待している。
18%という数字も、同じ26.6Kステップを回し終えた時点の最終loss同士を比べた差だ。
それがランク4〜128、バッチ1〜4、数千ステップという画像キャラLoRA特有の条件でも残るのかは、文献からは読めなかった。

実装の段取り

学習経路はAnimaLoraToolkitを想定している。
4月の検証記事ではsd-scriptsベースと書いたが、現行の公開ソースは独自のanima_train.py（FHfanshu/Anima_Trainerが基礎）になっていて、設定YAMLにオプティマイザ指定の項目もない。
sd-scripts本体には任意のモジュールパスのクラスをオプティマイザとして読み込む仕組みがあるものの、この経路はToolkitからは使えない。なのでToolkitのオプティマイザ生成箇所へ直接手を入れる。

差し替え、ペアの対応付け、直交化を一度に入れると、失敗したときにどれが原因か分けられない。順に1つずつ確かめる。

まずToolkitのオプティマイザ生成箇所を読み、中身はAdamWのままの独自クラスへ差し替えて、経路が通っているかだけ確認する。
次にlora_downとlora_upをペアとして扱えるようにする。同じ形のLoRA行列は層をまたいで大量にあるから、形状や登録順の対応付けでは別の層と取り違えかねない。LoRA注入時の層名かペア情報を明示的に持たせる。
最後に更新の直交化を入れる。最初はNewton–Schulzではなく、SVDで極分解の因子 $UV^\top$ を直接計算して正しさだけ確かめ、動いてから置き換える。

最小実験は画像数枚、ランク4か8、100ステップ前後で足りる。
品質はまだ確かめない。損失が発散せず最後まで走るか、保存したLoRAがComfyUIで読めるか、重みが初期値から変化しているかだけを確認する。
そこまで通ったら、AdamWと同条件で数百ステップ回してloss曲線と所要時間を並べる。