拡散トランスフォーマ（DiT）

動機：拡散の骨格はU-Netでなければならないのか

拡散モデルの標準的な実装は、長らくノイズ予測ネットワークに U-Net を使ってきました。畳み込みで局所特徴を抽出し、ダウン／アップサンプリングで多解像度を作り、要所にセルフアテンションを挟む——潜在拡散モデル のStable Diffusionもこの系譜です。しかしU-Netの畳み込み・スキップ接続・多解像度構造は、拡散に本質的に必要な要素なのかという問いが残っていました。

DiT（Diffusion Transformer, 2022）の答えは明快です。骨格をまるごとトランスフォーマに置き換えてもよい。Vision Transformer が「画像をパッチ列にすれば畳み込みなしで分類できる」と示したのと同じ発想を、拡散のノイズ予測器に持ち込みます。畳み込みの帰納バイアスを捨て、代わりにトランスフォーマの素直なスケーラビリティを得る、というトレードオフです。

パッチ化：潜在をトークン列に変換する

DiTはピクセルではなくVAE潜在の上で動きます。例えば 256×256 の画像はVAEで 32×32×4 程度の潜在テンソルへ圧縮されます。この潜在を、ViTと同じ要領で p×p のパッチに区切り、各パッチを線形射影して1トークンにします（patchify）。

# パッチ化：潜在テンソル → トークン列
z        : (H, W, C)                       # VAE潜在（例 32×32×4）
patches  : (N, p*p*C)   N = (H/p)*(W/p)    # 重なりなしに分割・平坦化
tokens   : patches @ W_E                   # W_E: 学習可能な線形射影
# tokens: (N, d)  ← N 個のパッチトークン。位置埋め込みを加算

パッチサイズ p（DiTでは 2, 4, 8）が系列長 N を決めます。セルフアテンションの計算量はトークン数に対し O(N^2) なので、p を小さくすると N が増え、表現は細かくなる代わりにGFLOPsが跳ね上がります。ここはViTと同じ勘所です。位置情報はアテンションが順序を区別しないため、パッチトークンに位置埋め込み（正弦波など）を加算して与えます。

adaLN：条件を「正規化のスケール・シフト」として注入する

DiTの技術的な核心は、時刻やクラスといった条件をどうトランスフォーマに効かせるかです。テキスト条件を鍵・値として差し込む 潜在拡散 のクロスアテンションとは別の、より軽量な経路を採ります。

出発点は正規化層です。LayerNorm は入力を標準化したうえで、学習可能なスケール γ とシフト β で復元します（y = γ ⊙ x̂ + β）。adaLN（adaptive LayerNorm） はこの γ／β を固定パラメータにするのをやめ、条件ベクトルから小さなMLPで回帰します。

# adaLN：条件 c から正規化のスケール・シフトを作る
c      = embed(t) + embed(y)               # 時刻・クラスなどを埋め込んで合算
γ, β   = MLP(c)                            # 条件からブロックごとに回帰
x_norm = γ ⊙ LayerNorm(x) + β             # 標準化後に条件依存のアフィン変換

DiTはさらに標準ブロックへ拡張し、正規化のスケール・シフトに加えて残差経路のゲートまで条件から作ります。1つのトランスフォーマブロック（セルフアテンション部とMLP部）に対し、{γ1, β1, α1, γ2, β2, α2} の6組を条件 c から一括で回帰し、それぞれの前段正規化と残差の重みに割り当てます。条件はネットワーク全体に薄く広く効き、追加の計算はブロックあたりごく小さなMLP1つで済みます。

adaLN-Zero：残差ゲートを0初期化して恒等から始める

DiTが最良の結果を出したのは、adaLNの改良版 adaLN-Zero です。着想は残差ネットワークの安定化技法（各ブロックを恒等写像から学習開始する）を条件付けに持ち込むことにあります。

具体的には、残差経路のゲート α（アテンション出力とMLP出力を足し込む重み）を回帰するMLPの最終層の重みとバイアスを0で初期化します。すると学習開始時は α = 0、つまり各ブロックの出力は残差の恒等成分だけになります。

# adaLN-Zero ブロック（概念）
γ1,β1,α1, γ2,β2,α2 = MLP_zero(c)   # MLP_zero の最終層は 0 初期化 → 初期 α ≈ 0
h = x + α1 ⊙ Attn( γ1 ⊙ LN(x) + β1 )   # 初期は α1=0 → h = x（恒等）
y = h + α2 ⊙ MLP ( γ2 ⊙ LN(h) + β2 )   # 初期は α2=0 → y = h（恒等）

DiT全体が最初は入力潜在をほぼそのまま通す恒等写像に近い初期状態から出発するため、深いトランスフォーマでも勾配が暴れにくく、学習が安定して速く収束します。論文の比較では、条件注入方式を クロスアテンション < 単純なadaLN < adaLN-Zero の順で品質（FID）が改善し、adaLN-Zeroが標準採用になりました（不等号の向きは「右ほど良い」の意）。

出力ヘッドと学習目標：拡散はそのまま

トランスフォーマを通ったトークン列は、最後に線形層でパッチ形状へ戻し（unpatchify）、潜在と同じ空間解像度のテンソルに再構成します。DiTの出力ヘッドは各パッチについて2つを予測します——予測ノイズと、対角共分散（分散） です。前者は標準の拡散損失、後者は分散を学習する定式化（learned variance）に使う成分で、ここは拡散モデル側の設計であってトランスフォーマ化とは独立です。

# 出力の再構成と学習目標
per_token = Linear(tokens)          # 各パッチについて (ε_pred, Σ_pred)
ε_θ, Σ_θ  = unpatchify(per_token)   # パッチ列 → 潜在解像度のテンソルへ
# 学習目標は従来どおり： L = E[ ‖ ε − ε_θ(z_t, t, c) ‖² ]（+ 分散項）

つまりDiTは、拡散モデル の「ノイズを予測して二乗誤差を最小化する」という骨太の目標を一切変えていません。骨格をトランスフォーマにしただけで、拡散の理論・サンプリング手続き・ガイダンス（分類器フリーガイダンス）はそのまま流用できます。

スケーリング特性：計算量を増やすほど素直に良くなる

DiTが決定的に重要視されたのは、きれいなスケーリング特性を示したからです。トランスフォーマの規模（層数・幅・ヘッド数）とパッチサイズを変えてモデルの計算量（GFLOPs）を振ると、GFLOPsを増やすほど生成品質（FIDの低下）が単調かつ予測可能に改善しました。

重要なのは、品質を支配するのが**パラメータ数そのものよりモデルの計算量（GFLOPs）**だと示された点です。パッチを細かくしてNを増やすのはパラメータをほとんど増やさずGFLOPsを上げますが、これでも品質は改善しました。この「規模で殴れば素直に報われる」関係は、言語モデルの スケーリング則 が拡散の骨格にも成り立つことを意味し、生成モデルを大規模化する明確な指針を与えました。

生成基盤への影響：画像から動画への共通骨格

DiTの意義は、単一モデルの性能以上に共通土台としての波及にあります。骨格がトランスフォーマなら、系列長を伸ばすだけで扱う対象を素直に拡張できます。

とりわけ動画生成では、画像を時空間パッチ（spacetime patch） に切ってトークン化すれば、空間と時間を区別せず1本の系列として扱えます。U-Netの固定的な空間構造では時間軸の追加が設計上の負担になりますが、トランスフォーマは系列長の問題に還元できるため拡張が自然です。Soraが「拡散トランスフォーマを時空間パッチに適用する」設計を採ったのは、まさにこのスケーラビリティと拡張の素直さを狙ったものです。

まとめ：帰納バイアスを捨て、スケーラビリティを買う

DiTの核心は、Vision Transformer が分類で示した「畳み込みを捨ててデータと計算で買う」というトレードオフを、拡散のノイズ予測器に持ち込んだ点にあります。パッチ化で潜在を系列に変え、adaLN-Zeroで条件を軽く安定に効かせ、あとはトランスフォーマの素直なスケーラビリティに任せる。この設計が、画像・動画・マルチモーダルを貫く生成基盤の共通骨格を成立させました。土台となる拡散そのものは 拡散モデル と 潜在拡散モデル を、大規模化の一般則は スケーリング則 を合わせて読むと、DiTの立ち位置が一本につながります。