Mean Mode Screamingとは？1000層DiTを安定化するMV-Split残差手法

拡散モデルの400〜1000層化で生じるトークン均質化崩壊「Mean Mode Screaming（MMS）」を初めて体系的に特定・解明した研究
平均成分と分散成分を独立ゲインで制御するMV-Split残差がLayerScaleを超え、400層でFID 2.60・1000層でFID 2.68を達成
世界初の1000層Diffusion Transformer安定学習を実証し、モデル重みとコードをオープンソースで公開

研究の背景

Diffusion Transformer（DiT）は高品質な画像生成を実現するアーキテクチャとして注目を集めています。一般的なTransformerと同様に、層数を増やすほど表現能力が高まることが期待されますが、実際には400層を超えたあたりから深刻な学習崩壊が頻発するという壁がありました。

これまでの研究では、この崩壊を「勾配爆発」として処理し、LayerScaleのような残差スケーリング手法で対処することが一般的でした。しかし崩壊の根本的なメカニズムは解明されておらず、既存の安定化手法を適用しても1000層規模では学習が成立しない状態が続いていました。本研究はその根本原因を特定することから出発しています。

MMS現象の発見と解明

著者らが特定した崩壊メカニズムは「Mean Mode Screaming（MMS）」と命名されています。MMSとは、深いDiTの学習中に突発的に発生する崩壊事象で、すべてのトークンの表現が互いに似通った均質な状態へ収束し、各トークン固有の特徴変動が失われる現象です。

図1: 400層DiTにおけるMMS崩壊事象の実測推移。垂直破線が崩壊発生ステップを示す。グローバル勾配ノルムの急増、平均一貫性勾配成分の集中増幅、Q/K勾配の4桁急落、残差分岐の開放と平均/分散エネルギー比の上昇、深い層でのトークンコサイン類似度の均質化を順に示す

MMS の崩壊は2段階で進行します。第1段階の後方トリガーでは、学習中に突如として「平均一貫性勾配成分」が急増します。これはすべてのトークンに共通した方向の勾配が一斉に積み上がる現象で、Attention機構（入力の重要な部分に注目する仕組み）のQ/K勾配が約4桁急落する一方、残差ライターの勾配は残り続けます。

続く第2段階の前方ロックインでは、残差分岐が大きく開いてネットワーク全体が平均主導の状態に固定されます。深い層でのAttentionは各トークン固有の特徴変動を補充する能力を失い、最終的にすべてのトークンのコサイン類似度が1に近づいて表現の均質化が完成します。

この崩壊の数学的根拠として、「Softmax零空間効果」が重要な役割を果たします。バリューベクトルが均質化するとAttentionのロジット勾配が構造的に消滅してQ/Kの学習が止まる一方、残差ライターの勾配だけが残るという非対称性が生まれます。さらにトークン間の整合度が高まるにつれて平均一貫性勾配がトークン数Tに比例してO(T)スケールで増大するという増幅則も理論的に導出されています。

MV-Split残差の仕組み

MMSの原因が「平均成分と中心成分が同一のゲインで制御されていること」にあることが判明したため、著者らはこの2つを分離して独立に制御する「Mean-Variance Split（MV-Split）残差」を提案しています。

従来のLayerScaleは残差全体を1つのスカラーゲインで等方的にスケーリングします。これに対してMV-Splitは、残差を「中心成分パス」と「平均成分パス」の2系統に分割し、それぞれ独立したゲインβとαで制御します。中心成分パスは各トークン固有の特徴変動を担い、平均成分パスはトークン全体に共通するグローバルな状態をリーキーな積分器（情報をゆっくりと蓄積する仕組み）として扱います。

図2: 残差ライター勾配の2モード分解。ステップごとの深さ方向中央値として平均一貫性勾配（Gmean）と中心勾配（Gctr）を比較。LayerScaleは両成分を等しく圧縮するが、MV-SplitはGmeanを抑制しつつGctrを高い安定したバンドに維持する

後方の効果として、2つのパスが独立したゲインを持つため、平均一貫性勾配が急増しようとしても中心成分の勾配バンドへの影響が遮断されます。これによりMMSの連鎖的な崩壊プロセスを根本から断ち切ることができます。LayerScaleが両成分を同時に圧縮してしまうのとは対照的な設計です。

また本研究では興味深い副次的な知見も示されています。トークンの平均成分は崩壊を引き起こす危険な方向である一方、訓練済みモデルではグローバルなタイムステップ情報のキャリアとしても機能しているという二面性です。MV-Splitはこの有用な信号を保ちつつ、崩壊への増幅を抑制するという精巧な設計になっています。

実験結果

ImageNet 256×256のクラス条件付き画像生成タスクで、400層DiTを用いた80,000ステップの比較実験が行われました。FID（Fréchet Inception Distance）は値が小さいほど生成品質が高く、IS（Inception Score）は値が大きいほど高品質であることを示す指標です。

手法	@20k FID	@20k IS	@50k FID	@50k IS
ベースライン（未安定化）	崩壊	—	崩壊	—
LayerScale（400層）	6.50	96.6	2.90	165.5
MV-Split（400層）	3.64	139.9	2.60	185.5
MV-Split（1000層）	—	—	2.68	—

図3: 品質指標とオプティマイザー安定性（80,000ステップ、ImageNet 256×256）。上段がFID-50KとInception Score、下段がグローバル勾配ノルム。MV-Splitは400層・1000層ともに安定した勾配バンドを維持し、LayerScaleを超える品質を達成

400層での比較において、MV-Splitは20,000ステップ時点でLayerScaleに対してFIDを2.86ポイント改善し、学習効率でも優位を示しています。勾配ノルムの推移を見ると、MV-SplitはLayerScaleよりも高い安定した勾配バンドを維持しており、より活発な学習が継続していることが確認できます。

さらに1000層DiTへのスケールアップ実験では、FID-50K 2.68という高品質な結果が得られました。拡散モデルの学習効率を改善する研究は他にも進んでいますが、学習安定性の根本原因から対処して極端な深さへのスケーリングを実現したアプローチは本研究が先駆的です。

図4: 1000層MV-Split DiTによるクラス条件付き画像生成サンプル（ImageNet 256×256）

まとめ

本研究は、拡散モデルの超深層化を妨げてきた「Mean Mode Screaming」という崩壊現象を初めて体系的に特定し、その発生メカニズムを詳細に解明しました。そしてMV-Split残差という解法を導き出し、400層・1000層という極限的な深さのDiTで安定学習を実証しています。

層数をスケールアップすることへの技術的障壁が下がれば、パラメータ効率やモデルの表現能力についての理解が深まる可能性があります。一方で、1000層モデルの学習コストやメモリ消費は依然として大きな課題であり、実用的なスケールへの展開にはさらなる工夫が必要です。コードと1000層モデルの重みはGitHubおよびHugging Faceでオープンソース公開されており、発展研究への貢献が期待されます。