DiT(Diffusion Transformer)を徹底解説 — 拡散モデルのU-NetをTransformerに置き換えた【ICCV 2023】

Sora(動画生成)やStable Diffusion 3 など、2024年以降の最先端の生成モデルには、ある共通の「背骨」が入っています。それが DiT(Diffusion Transformer) です。拡散モデルといえば長らく U-Net(畳み込みベースのノイズ除去ネットワーク)が定番でした。DiTはそこに真っ向から「Transformerにしたらどうなる?」と問い、U-Netを捨てて純粋なTransformerに置き換えたら、むしろSOTA(当時の最高品質)に到達してしまった——そんな論文です(Peebles & Xie, ICCV 2023)。

なぜこれを学ぶ価値があるのか。理由は2つあります。

  • 最新生成モデルの土台:Sora・Stable Diffusion 3・PixArt など、いま話題の生成モデルはほぼDiT系のアーキテクチャです。DiTを理解すれば、それらの設計思想がまとめて腑に落ちます。
  • 「スケーリング則」を画像生成に持ち込んだ:LLMで確立した「大きくすれば素直に良くなる」という性質を、DiTは画像生成でも成り立たせました。生成モデルを大規模化する方向性の理論的な裏付けになっています。

本記事では、論文の図(CC BY 4.0 ライセンスなので出典明記のうえ引用します)と日本語の自作図を組み合わせ、読めばDiTの全体像が分かることを目指します。

本記事の内容

  • DiTの位置づけ — 潜在拡散モデルのどこを置き換えたのか
  • 潜在のパッチ化(Transformerが食べられる形にする)
  • DiTブロックと adaLN-Zero による条件付けの仕組み
  • 条件付け戦略の比較(なぜadaLN-Zeroが勝つのか)
  • スケーリング則 — Gflopsを増やすほどFIDが下がる

前提知識

この記事を読む前に、以下を押さえておくとスムーズです。

DiTの位置づけ — 何を置き換えたのか

まず全体像です。DiTは拡散モデルを丸ごと作り直したわけではありません。潜在拡散モデル(Latent Diffusion)の「ノイズ除去ネットワーク」だけをTransformerに差し替えたものです。

DiTの位置づけ:画像→VAEで潜在へ→拡散でノイズ付加・DiTで除去→VAEデコードで画像、というパイプラインのノイズ除去部分をDiTが担う

パイプラインはこうです。まず画像(256×256×3)をVAEエンコーダで潜在表現(32×32×4)に圧縮します。拡散モデルはこの軽い潜在空間の上で動きます。学習時は潜在にノイズを加え、推論時はそのノイズを少しずつ除去してきれいな潜在を復元し、最後にVAEデコーダで画像に戻します。この「ノイズを除去する心臓部」が、従来はU-Net、DiTではTransformerになっている——これが唯一にして最大の変更点です。潜在空間で動かすので計算が軽く、高解像度にも対応しやすいという潜在拡散の利点はそのまま受け継いでいます。

そもそも拡散モデルが何をしているかを、1枚で確認しておきましょう。

拡散の前向き(ノイズを少しずつ足してデータを壊す)と逆向き(DiTが少しずつノイズを除いて生成する)の概念図

拡散モデルは2方向の過程からなります。前向き過程は学習データに少しずつガウスノイズを足していき、最終的に純粋なノイズにします。逆向き過程はその逆で、ノイズから出発して各ステップのノイズを予測・除去し、データを生成します。DiTが担うのは逆向きの「各ステップでどんなノイズが乗っているかを予測する」役割です。ここをうまく学習できれば、ノイズから画像を生成できます。

学習は驚くほどシンプルです。きれいな潜在 $\bm{z}_0$ に、ランダムな時刻 $t$ のノイズ $\bm{\epsilon}\sim\mathcal{N}(\bm{0},\bm{I})$ を規定の割合で混ぜてノイズ付き潜在 $\bm{z}_t$ を作り、DiTにそのノイズを当てさせます。損失は予測ノイズと真のノイズの単純な平均二乗誤差 $\|\bm{\epsilon}_\theta(\bm{z}_t, t, c) – \bm{\epsilon}\|^2$ です。DiTはこの $\bm{\epsilon}_\theta$(と分散)を出力するネットワークにすぎません。だからこそ、その中身をU-Netにしようと Transformer にしようと、枠組みは何も変わらない——これがDiTが成立する前提です。

では、Transformerは画像(潜在マップ)をそのままは食べられません。どうトークンにするのでしょうか。

潜在をパッチに切ってトークンにする

DiTはViT(Vision Transformer)と同じ発想で、潜在マップをパッチに分割してトークン列にします。

パッチ化の自作図:潜在マップをp×pの小区画に切り、各パッチを線形射影でトークンに変換し位置埋め込みを足す

32×32×4の潜在を、$p \times p$ の小区画(パッチ)に区切り、各パッチを線形射影で $d$ 次元のトークンに変換します。位置埋め込みを足せば、Transformerが処理できるトークン列の完成です。論文の図でこの入力仕様を見てみましょう。

DiTの入力仕様:パッチサイズpを小さくするほどトークン数Tが増える(出典の論文Fig.4)

出典: Peebles & Xie, “Scalable Diffusion Models with Transformers” (ICCV 2023), Fig.4 (CC BY 4.0)

ここで重要なのがパッチサイズ $p$ の効き方です。$p$ を小さくする(例:$p=2$)と、潜在を細かく刻むのでトークン数が増え、Transformerの計算量(Gflops)が増えます。後で見るように、この計算量こそが品質を決める鍵になります。パッチを小さくすると、より細かい構造を捉えられて生成品質が上がるのです。

トークン列ができたら、いよいよDiTブロックで処理します。ここに、拡散モデルならではの「条件付け」の工夫が入ります。

DiTブロックと adaLN-Zero による条件付け

拡散モデルのノイズ除去ネットワークは、ただトークンを処理するだけでは足りません。「いまどのノイズレベル(時刻 $t$)か」「どのクラス $c$ の画像を作りたいか」という条件を必ず受け取る必要があります。DiTはこの条件の渡し方を3種類試し、最も良いものを採用しました。論文のアーキテクチャ図がこれです。

DiTアーキテクチャ:左が全体のLatent Diffusion Transformer、右が3種のDiTブロック(adaLN-Zero/Cross-Attention/In-Context)の詳細(出典の論文Fig.3)

出典: Peebles & Xie, “Scalable Diffusion Models with Transformers” (ICCV 2023), Fig.3 (CC BY 4.0)

図の左が全体構成、右が3種類のDiTブロックです。3つの方式を整理すると次の通りです。

  • In-Context:条件 $t, c$ を追加のトークンとして入力列の末尾に並べ、普通の自己注意で混ぜる(GPTのプロンプトのような発想)。追加計算はわずかだが性能は伸び悩む。
  • Cross-Attention:条件を別系列として用意し、各ブロックでCross-Attentionを挟んで参照する。表現力は高いが計算量が約15%増える。
  • adaLN-Zero:条件から正規化のスケール・シフト・ゲートを生成して各層を変調する(後述)。最も軽量かつ高性能。

最も性能が良かったのが一番左の adaLN-Zero ブロックでした。この仕組みを日本語で噛み砕きます。

adaLN-Zeroの自作図:条件t,cをMLPに通してscale γ・shift β・gate αを生成し、Layer Normの出力を変調、ゲートαを0初期化して恒等写像から学習を始める

adaLN(adaptive Layer Norm)は、条件 $t, c$ から小さなMLPでスケール $\gamma$・シフト $\beta$ を作り、Layer Normの出力を $\gamma \odot (\cdot) + \beta$ と変調します。つまり「いまの時刻・クラスに応じて、各層の正規化の効き具合を調整する」わけです。さらにDiTは gate $\alpha$(ゲート)も条件から生成し、各残差ブロックの出力に掛けてから足します。

ここで効いてくるのが 「-Zero」 の部分です。このゲート $\alpha$ をゼロで初期化します。すると学習開始時は各ブロックの寄与が $\alpha=0$ で消え、ネットワークは恒等写像として振る舞います。深いネットワークを「何もしない状態」から学習し始められるので、訓練が非常に安定するのです。これはResNetやLSTMで見た「最初は素通りさせる」という安定化の発想と同じ精神です。

なぜadaLN-Zeroがそんなに良いのか。論文は条件付け戦略を定量比較しています。

条件付け戦略の比較 — なぜadaLN-Zeroが勝つのか

条件付け戦略の比較:adaLN-Zeroがin-context・cross-attention・adaLNよりも低いFIDを達成(出典の論文Fig.5)

出典: Peebles & Xie, “Scalable Diffusion Models with Transformers” (ICCV 2023), Fig.5 (CC BY 4.0)

この図は、同じ計算量のもとで4つの条件付け方式(条件をトークンとして並べる in-contextcross-attention、通常の adaLN、そして adaLN-Zero)を比べたものです。縦軸はFID(小さいほど高品質)。adaLN-Zeroが明確に最も低いFIDを達成しています。条件をトークンに足すだけのin-contextや、別途cross-attentionを挟む方式よりも、「正規化の変調+ゼロ初期化ゲート」という軽量な仕掛けのほうが効率的に条件を伝えられる、という発見です。追加の計算をほとんど増やさずに済むのも利点です。

条件付けが決まれば、あとは「どれだけ大きくするか」です。ここでDiTの最大の主張——スケーリング則——が登場します。

スケーリング則 — Gflopsを増やすほどFIDが下がる

DiTのスケーリング:モデルを大きく(バブル=Gflops大)するほど、またパッチを小さくするほどFIDが下がる(出典の論文Fig.2)

出典: Peebles & Xie, “Scalable Diffusion Models with Transformers” (ICCV 2023), Fig.2 (CC BY 4.0)

DiTの核心的な発見はこれです。Transformerの計算量(Gflops)を増やすほど、生成品質(FID)が予測可能な形で改善する。バブルの大きさが計算量を表し、大きいモデル・小さいパッチほど右下(低FID=高品質)に向かいます。モデル幅・深さを増やす、あるいはパッチを細かくする——どの方向でも計算量を増やせばFIDが下がる、という素直な関係が成り立ちました。

これは画像生成にとって大きな意味を持ちます。U-Netでは「大きくすれば良くなる」という関係が必ずしも明瞭ではありませんでした。DiTは、LLMで確立したスケーリング則を画像生成でも成立させたのです。「計算資源を投じれば品質が上がる」と分かっていれば、安心して大規模化に投資できます。

なぜTransformerだとこれが効くのでしょうか。

U-NetとDiTの比較の自作図:U-Netは畳み込み中心でスケール則が不明瞭、DiTは自己注意で大域を見て素直なスケール則を持ち、Sora・SD3の土台になった

U-Netは畳み込みの強い帰納バイアス(局所性)を持ち、解像度ごとに専用の設計が要るなど、スケールの仕方が一様ではありませんでした。一方Transformerは自己注意で大域的な関係を捉え、ブロックを積むだけで素直に大きくできます。この「均一に大きくできる」性質が、スケーリング則のきれいさを生みます。実際、最大モデル DiT-XL/2 はImageNet 256×256のクラス条件付き生成でFID 2.27 を達成し、当時のSOTAを更新しました。

最後に、ここまで何度も出てきた評価指標FIDを直感で押さえておきましょう。

評価指標FIDの直感

FIDの直感図:生成画像と実画像の特徴分布の近さを測る。良い生成は実画像分布に重なり(FID小)、悪い生成は離れる(FID大)

FID(Fréchet Inception Distance)は、生成画像と実画像を特徴抽出器に通し、その特徴分布がどれだけ近いかを測る指標です。良い生成(緑)は実画像(灰)の分布に重なり、FIDが小さくなります。悪い生成(赤)は分布が離れ、FIDが大きくなります。1枚ごとの正解と比べるのではなく「生成物の集団が実データの集団とどれだけ似た広がりを持つか」を見るのがポイントです。

結果と、その後の影響

DiTの最大モデル DiT-XL/2(パッチサイズ2の最大構成)は、ImageNet 256×256 のクラス条件付き生成で FID 2.27 を達成し、それまでの拡散モデル(ADMなどU-Netベース)を上回って当時のSOTAになりました。512×512でも同様に高品質な生成を示しています。重要なのは、この結果が「巧妙な専用設計」ではなく「素直にTransformerを大きくしただけ」で得られた点です。スケーリング則が成り立つので、計算資源を増やせばさらに伸びる余地があることも示されました。

そしてDiTの本当のインパクトは、その後の生成モデルの設計を決定づけたことにあります。OpenAIの動画生成 Sora は「DiTを動画に拡張したもの」と説明されていますし、Stable Diffusion 3 は DiT に整流フロー(Rectified Flow)とマルチモーダル注意を組み合わせた MM-DiT を採用しています。軽量な text-to-image モデルの PixArt 系も DiT ベースです。「拡散のバックボーンはTransformer」という今の常識は、この論文が作りました。U-Netからの世代交代を象徴する一本と言えます。

まとめ

DiT(Diffusion Transformer)を、論文の図を引用しながら解説しました。

  • 置き換えたのは1か所:潜在拡散モデルのノイズ除去ネットワークを、U-NetからTransformerへ。
  • パッチ化:潜在マップを $p\times p$ パッチに切ってトークン化。パッチを小さくするほど計算量が増え品質が上がる。
  • adaLN-Zero:条件(時刻・クラス)から scale/shift/gate を生成して各層を変調。ゲートをゼロ初期化して恒等写像から学習し、安定かつ高性能。
  • スケーリング則:Gflopsを増やすほどFIDが素直に下がる。画像生成にもスケール則を持ち込んだ。
  • 影響:DiT-XL/2でFID 2.27のSOTA。Sora・Stable Diffusion 3 などの土台になった。

「拡散モデルのバックボーンをTransformer化し、スケール則に乗せる」——このシンプルな発想が、現在の生成AIの主流アーキテクチャを決定づけました。

次に読むと理解が深まる記事:

参考文献:William Peebles, Saining Xie, “Scalable Diffusion Models with Transformers”, ICCV 2023 (arXiv:2212.09748, CC BY 4.0)。本記事中の論文図(Fig.2/3/4/5)は同論文よりCC BY 4.0ライセンスのもと引用しました。