CMuon、DiTの融合重みを分割してから直交化——675MモデルでAdamW比2倍の学習速度

拡散Transformer（DiT）では、行列モメンタムにNewton–Schulz直交化を適用して学習効率を高めるため、AdamWの代わりにMuonを採用する動きが始まっている。新たな研究「CMuon」は、Muonをそのまま適用した場合、モデル実装に由来する隠れた問題が残ると指摘する。計算効率を高めるため、DiTでは通常、Q、K、Vの射影に加え、AdaLNのスケール、シフト、ゲーティングの各パラメータを、より大きな重みテンソルに融合する。標準のMuonはテンソル全体に対して1回の直交化を共有する。これは、すべての部分行列の勾配統計から同一のプリコンディショナーを構築することに相当し、本来は機能的に独立しているパラメータ更新を互いに結合してしまう。

CMuonは直交化の前に、パラメータ本来の機能境界に沿って複数のブロックへ再分割し、各ブロックのモメンタム行列を個別に処理する。この変更はDiTのフォワードアーキテクチャを変えず、追加の表現アライメント損失も必要としない。論文によれば、計算オーバーヘッドは無視できるという。研究チームはImageNetの256×256画像を用い、バッチサイズ1,024で675Mパラメータのモデルを学習した。200 epoch後にFID 1.18を達成し、同一実験の標準Muonを上回ったほか、AdamWと比べて収束が約2倍高速になったと報告している。重要な違いは、学習初期に急速に改善した後で停滞するのではなく、優位性が学習後半まで持続した点にある。

エンジニアリング上、この結果は、行列ベースのオプティマイザーを使用する際、フレームワークにおける重みの融合が単なるメモリ配置上の詳細ではなく、最適化の幾何学も変えることを示している。実装ではQKVやAdaLNなどのテンソルの意味的な境界を保持し、分割方向がモデルのバージョンと一致していることを確認する必要がある。現時点のエビデンスは、特定のDiT、データセット、学習レシピに限られており、研究チームも完全なコードをまだ公開していない。また、FIDだけではテキストとの整合性や高解像度画像の生成品質を評価できない。今後は、数十億パラメータ規模のtext-to-imageモデルや動画モデルでも結果を再現できるか、さらにブロック分割後の通信コストとkernelのオーバーヘッドを引き続き無視できるかを検証する必要がある。