Scalable Diffusion Models with Transformers (Peebles, 2022)

3行要約

💡 Transformerアーキテクチャに基づく新しい拡散モデルであるDiffusion Transformer (DiT) を提案し、従来のU-NetバックボーンをTransformerに置き換えました。
📈 DiTは、計算量（Gflops）の増加に伴いFIDが着実に低下するなど、優れたスケーリング特性を示し、モデルサイズや入力トークン数を増やすことで性能が向上します。
✨ 最大のDiTモデルであるDiT-XL/2は、クラス条件付きImageNet 256x256および512x512ベンチマークにおいて、既存のすべての拡散モデルを上回る最先端のFIDを達成しました。

Q:この論文の概要を教えてください

A:

本研究は、Transformerアーキテクチャをバックボーンとして採用した新しいクラスの拡散モデル、Diffusion Transformers（DiT）を提案します。画像生成タスクにおいて、従来の拡散モデルで一般的に使用されてきたU-Netバックボーンを、潜在空間のパッチ上で動作するTransformerに置き換えます。

DiTモデルは、Variational Autoencoder (VAE) の潜在空間で拡散モデルを訓練するLatent Diffusion Models (LDM) フレームワークに基づいています。入力画像 $x$ はまずVAEのEncoder $E$ により潜在表現 $z = E(x)$ に圧縮され、DiTはこの潜在表現 $z$ に加えられたノイズを予測するように訓練されます。ノイズ予測の目的関数は、元のノイズ $\epsilon_t$ とモデルの予測 $\epsilon_\theta(x_t)$ との二乗誤差の期待値として表されます：

$ L_{\text{simple}}(\theta) = \mathbb{E}{t, x_0, \epsilon_t} \left[ ||\epsilon\theta(x_t) - \epsilon_t||_2^2 \right] $

ここで、x_t = \sqrt{\bar{\alpha}t}x_0 + \sqrt{1 - \bar{\alpha}t}\epsilon_t です。また、モデルは共分散 \Sigma\theta も学習しますが、主要な訓練はノイズ予測に焦点を当てます。推論時には、生成プロセスをクラスラベル c などの条件で誘導するために、Classifier-free guidanceが使用されます。これは、条件付きモデル \epsilon\theta(x_t, c) と無条件モデル \epsilon_\theta(x_t, \emptyset) の予測を
組み合わせて行われます：
$ \hat{\epsilon}\theta(x_t, c) = \epsilon\theta(x_t, \emptyset) + s \cdot (\epsilon_\theta(x_t, c) - \epsilon_\theta(x_t, \emptyset)) $

ここで $s > 1$ はガイダンススケールです。

DiTのアーキテクチャはVision Transformer (ViT) に倣っています。まず、入力であるノイズ付き潜在表現（例えば、256x256x3画像に対応する32x32x4）をパッチに分割し、各パッチを線形埋め込みにより次元 $d$ のトークンに変換する「Patchify」層を通します。パッチサイズ $p \times p$ により、入力トークン列の長さ $T$ は $(I/p)^2$ となります（入力潜在空間が $I \times I \times C$ の形状の場合）。このトークン列に位置エンベディングが加えられます。続いて、複数のDiTブロックでトークン列を処理します。

本研究では、条件付け（ノイズタイムステップ $t$ とクラスラベル $c$）を組み込むためのTransformerブロック設計の4つのバリアントを探求しました。その中でもAdaptive Layer Norm (adaLN) を用いた設計が最も効果的であり、特にadaLN-Zeroブロックが優れていました。adaLN-Zeroブロックは、標準的なLayer NormをAdaLNに置き換えるとともに、残差接続の直前に適用されるスケーリングパラメータ $\alpha$ を追加で回帰的に予測します。これらのパラメータ $\gamma, \beta, \alpha$ は $t$ と $c$ の埋め込みの合計から生成されます。$\alpha$ を出力する層をゼロ初期化することで、DiTブロックがIdentity functionとして開始するように工夫されています。

DiTモデルのスケーラビリティは、Forward Passの計算量（Gflops）を主要な指標として分析されました。Transformerの深さや幅を増加させるモデルサイズのスケーリング（DiT-S, DiT-B, DiT-L, DiT-XL）と、パッチサイズ $p$ を小さくして入力トークン数を増加させるスケーリングの両方が探求されました。実験の結果、モデルGflopsが高いDiTモデルほど、FID（Fréchet Inception Distance）が継続的に低下することが示され、モデルGflopsとFIDの間に強い負の相関（約 -0.93）があることが確認されました。この傾向はFID以外の評価指標（Inception Score, Precision, Recall）でも見られました。より大きなDiTモデルは、総訓練計算量（モデルGflops × バッチサイズ × 訓練ステップ数）に対して、より効率的に性能を向上させることが示されています。

最大モデルであるDiT-XL/2は、ImageNetベンチマークにおいて、従来のU-Netベースの拡散モデルやGANなどの先行研究を凌駕するState-of-the-Art性能を達成しました。クラス条件付きImageNet 256x256生成では、FID 2.27（Classifier-free guidance scale 1.5使用）を達成し、従来の最高性能であったLDM-4のFID 3.60を大幅に改善しました。同様に、512x512生成ではFID 3.04（同scale使用）を達成し、ADMのFID 3.85を上回りました。DiT-XL/2は、これらのSOTA結果を、ピクセル空間で動作するADMのようなU-Netモデルと比較して大幅に少ないGflopsで達成しており、高い計算効率を示しています。

また、本研究では、訓練後にサンプリングステップ数を増加させることによる計算量増加が、モデル自体の計算量不足を補うことができないことを示しました。モデルGflopsが大きいDiTは、サンプリング計算量が少なくても、モデルGflopsが小さいDiTがより多くのサンプリングステップを使用した場合よりも優れた性能を示します。

結論として、本研究は、拡散モデルのバックボーンとしてTransformerが有効であることを示し、その優れたスケーリング特性が画像生成タスクにおいても性能向上に直結することを確認しました。DiTは、従来のU-Netモデルを凌駕する性能と計算効率を両立し、将来の画像生成研究における有力なアーキテクチャとなる可能性を示唆しています。今後の展望として、DiTをさらに大規模にスケーリングすることや、Text-to-Imageモデルへの応用が挙げられています。

Q:提案手法の処理の流れを順序だてて説明してください

A:

DiTは、Latent Diffusion Models (LDMs) [48] のフレームワーク内で、従来のU-NetバックボーンをTransformerに置き換えたモデルです。

処理は、VAEの潜在空間で行われます。主な流れは以下の通りです。

入力の準備 (Noised Latent + Conditioning):
- まず、オリジナルの画像はVAEのEncoder (E) によって潜在表現 (z) に圧縮されます。論文では、256x256x3の画像が32x32x4の潜在表現に変換される例が示されています。
- この潜在表現 (z) に、Forward Diffusion Processに従ってノイズが加えられ、「ノイズ付き潜在表現 (Noised Latent)」($z_t$) が生成されます。これがDiTへの主要な入力となります。
- さらに、Denoising Diffusion Probabilistic Models (DDPMs) では、ノイズの量を示す「タイムステップ (Timestep)」(t) や、生成したい画像の「クラスラベル (Label)」(y) などの条件情報も入力として与えられます。
Patchify:
- 入力された「ノイズ付き潜在表現」($z_t$) は、Vision Transformer (ViT) [10] と同様に、小さな空間的な「パッチ (Patch)」に分割されます (Figure 4参照)。論文では、パッチサイズ (p) として2, 4, 8などが実験されています。
- 各パッチは、線形変換によって、Transformerが処理できる固定次元 (d) のベクトル（「入力トークン (Input Token)」）に埋め込まれます。パッチサイズが小さいほど、トークン数は増えます。
- 生成されたトークンシーケンスには、パッチの元の位置情報を加えるために、位置埋め込み (Positional Embeddings) が追加されます。
DiT Blockによる処理:
- このトークンシーケンスは、積み重ねられた複数の「DiT Block」によって順次処理されます (Figure 3参照)。各DiT Blockは基本的にTransformerの構造に基づいています。
- 重要なのは、ここでタイムステップ (t) とクラスラベル (y) といった条件情報がモデルに注入される点です。論文ではいくつかの方法を実験し、特に「adaLN-Zero Block」が最も効果的であることを示しています (Figure 5参照)。
- adaLN-Zero Blockでは、Multi-Head Self-Attention や Pointwise Feedforward などの標準的なTransformer層に加え、Adaptive Layer Normalization (adaLN) が使用されます。adaLNは、条件情報（tとyの埋め込みの合計）からLayer Normalizationのスケール ($\gamma$) やシフト ($\beta$) パラメータ、および残差接続の前のスケーリングパラメータ ($\alpha$) を動的に生成し、適用します。これにより、各トークンの正規化やスケーリングが条件に依存するようになります。
Transformer Decoder:
- 最後のDiT Blockを通過した後のトークンシーケンスは、デコーダーによって元の空間表現と同じ形状の出力に変換されます。
- 最終的なLayer Normalizationが適用された後、線形層によって各トークンがデコードされます。
- これらのデコードされた結果を元のパッチ配置に並べ替えることで、モデルの最終的な出力であるノイズ ($\epsilon_\theta$) の予測や、対角共分散 ($\Sigma_\theta$) の予測が得られます。
出力:
- モデルの最終的な出力は、ノイズ付き潜在表現 ($z_t$) に加えられたノイズ ($\epsilon_t$) を予測した ($\epsilon_\theta$) や、学習された逆過程の共分散 ($\Sigma_\theta$) となります。

訓練時には、この予測 ($\epsilon_\theta$, $\Sigma_\theta$) と正解の値（実際に加えたノイズ $\epsilon_t$ など）との誤差（損失関数 Lsimple や L(θ)）を最小化するように学習が進められます。サンプリング時には、初期ノイズから始めて、この学習された逆過程モデル（$p_\theta(z_{t-1}|z_t, c)$）を用いて段階的にノイズを取り除き、最終的な潜在表現 (z) を生成し、それをVAEのDecoder (D) で画像に戻します。

DiTのコアアイデアは、画像生成のためのバックボーンとして、U-Netの畳み込み構造に依存する代わりに、Transformerのスケーリング特性を活用することにあります。特に、トークン数を増やすこと（パッチサイズを小さくすること）やTransformerの深さ/幅を増やすことが、FIDの向上に強く相関することが示されています (Figure 8参照)。

この説明が、DiTの処理の流れを理解する助けとなれば幸いです。何かさらに詳しく知りたい点があれば、遠慮なく質問してください。

Q:関連手法との比較について詳しく説明してください

A:

DiTは、Diffusion Modelのバックボーンとして、広く用いられてきたU-NetアーキテクチャをTransformerに置き換えるという、比較的シンプルな変更を提案していますが、これがスケーリング特性と性能に大きな影響を与えることを示しています。

主な比較対象は以下の手法です。

U-Netをバックボーンとする従来のDiffusion Models:
- ADM (Dhariwal and Nichol, 2021) [9]: Denoising Diffusion Probabilistic Models (DDPMs) [19] を改良し、画像生成タスクでGANsを凌駕する性能を示した seminal work です。高解像度画像生成のために、カスケーディング拡散モデルや classifier guidance を導入しました。ADMは、ResNetブロックを主体とした畳み込みU-Netをバックボーンとして採用しています。
- LDM (Rombach et al., 2022) [48]: 画像を高解像度のピクセル空間でなく、VAEの潜在空間で拡散モデルを学習することで、計算効率を大幅に向上させた手法です。LDMもまた、畳み込みU-Netアーキテクチャをバックボーンとして使用しています（ただし、条件付けのためにcross-attention層も追加しています）。
- DiTとの比較:
  - アーキテクチャの根本的な違い: 最も重要な違いは、バックボーンとして畳み込みベースのU-Netを使用するか、Transformerを使用するかという点です。論文では、DiTはU-Netの持つ畳み込み的帰納バイアス（locality, translation equivariance）がDiffusion Modelの性能に必須ではないことを示し、Transformerでも高精度な生成が可能であることを実証しました (Page 2)。
  - スケーリング特性: 論文の主要な発見の一つは、DiTがTransformerの優れたスケーリング特性を継承していることです。U-Netベースのモデルもスケーリング研究が行われていますが [9]、DiTはモデルサイズ（Transformerの深さ/幅）だけでなく、入力トークン数（パッチサイズ）をスケーリングすることでも性能が向上することを明確に示しています (Figure 6)。特に、パッチサイズを小さくしてトークン数を増やすと、パラメータ数はほとんど変わらずGflopsが増加し、これが性能向上に強く寄与することを示しました (Figure 8)。これは、Gflopsがモデル品質を決定するクリティカルな要素であり、TransformerはU-NetよりもGflopsを性能に効率的に変換できる可能性を示唆しています。
  - 計算効率: DiT-XL/2は、同じ性能レベルを達成する際に、ADMやLDMといった既存のU-Netベースモデルよりも計算効率が良いことが示されています (Figure 2, Table 6)。特にピクセル空間で動作するADMと比較すると、潜在空間で動作するDiT-XL/2は大幅にGflopsを削減しています。LDMとは潜在空間で比較され、DiT-XL/2 (118.6 Gflops) がLDM-4-G (103.6 Gflops) を上回るFIDを達成しています (Table 2)。
  - 条件付けメカニズム: 論文では、Transformerブロック内での条件付け（タイムステップやクラスラベル）の方法を詳細に実験しています (Page 4-5)。In-context conditioning, Cross-attention, adaLN, adaLN-Zeroを比較し、adaLN-Zeroが最も優れた性能を示すことを発見しました (Figure 5)。これは、Transformerアーキテクチャに合わせた効果的な条件付け手法が重要であることを示しています。従来のU-Netモデルもadaptive normalization [40] などを使用していますが、Transformerの構造に合わせた設計が性能に影響を与えています。
  - 性能: DiT-XL/2モデルは、ImageNet 256x256および512x512のクラス条件付き画像生成タスクにおいて、既存のすべてのDiffusion Model（ADM, LDM, CDM [20] など）を上回るState-of-the-ArtのFIDを達成しました (Table 2, 3)。
他のTransformerベースの生成モデル:
- Autoregressive Models (e.g., PixelRNN/CNN [58, 52], Image Transformer [38], VQ-VAE + Transformer [11, 47, 62]): これらのモデルは、画像をピクセル単位または離散コードブックのトークン単位で自己回帰的に生成します。Transformerはこれらのモデルでも広く使用され、優れたスケーリング特性が報告されています [17, 62]。
- Masked Generative Models (e.g., MaskGIT [4, 14]): 入力の一部をマスクし、それを予測するように学習するモデルで、Transformerが使用されます。
- Diffusion Models on embeddings (e.g., DALL·E 2 [41, 46]): 画像ピクセルではなく、CLIP [41] の画像埋め込みのような潜在表現に対してDiffusion Modelを適用し、その後デコーダーで画像に戻す手法です。DALL·E 2は拡散モデルの部分にTransformerを使用しています。
- DiTとの比較:
  - 生成パラダイム: DiTはDiffusion Modelのフレームワーク内で動作します。これは自己回帰モデルやGANとは異なる生成プロセスを持ちます。
  - Transformerの応用: DiTはこれらのモデルと同様にTransformerの強力な表現学習能力とスケーリング能力を活用しています。しかし、DiTは（DALL·E 2の一部を除いて）画像自体の空間的な潜在表現（VAEの出力）に対してTransformerを適用している点で異なります。ViTのように画像をパッチに分割し、そのパッチシーケンスを直接Transformerで処理するというアプローチは、画像生成のためのTransformerバックボーンとしての新しい方向性を示しています。
  - スケーリング研究: 自己回帰モデルで示されてきたTransformerのスケーリング法則 [26] が、Diffusion Modelのアーキテクチャにも適用できることを、DiTは実証的に示しました (Figure 9, 13)。これは、生成モデルの様々なクラスにおいて、Transformerのアーキテクチャとスケーリング特性が共通して有効である可能性を示唆しています。
GANs (Generative Adversarial Networks):
- BigGAN-deep [2], StyleGAN-XL [53]: Diffusion Modelが登場するまで、多くの画像生成ベンチマークでState-of-the-Artを達成していたモデルクラスです。
- DiTとの比較:
  - 生成手法: GANはGeneratorとDiscriminatorの敵対的な学習によって画像を生成します。Diffusion ModelであるDiTとは、学習・生成プロセスが全く異なります。
  - 性能: 近年、Diffusion Modelは画像品質においてGANsを凌駕する例が増えており [9]、DiTもその流れを汲んでいます。DiT-XL/2は、ImageNet 256x256および512x512において、StyleGAN-XLを含む既存のGANsを上回るFIDを達成しています (Table 2, 3)。

要約すると、DiTはDiffusion ModelにおけるU-NetバックボーンをTransformerに置き換えることで、Transformerの持つ優れたスケーリング特性をDiffusion Modelに持ち込み、計算効率を維持しながらState-of-the-Artの画像生成性能を実現した手法です。従来のU-NetベースのDiffusion Modelに対しては、アーキテクチャの柔軟性とスケーリング効率で優位性を示し、他のTransformerベース生成モデルとは異なる応用（画像空間潜在表現への直接適用）を行い、GANsに対しては最新のDiffusion Modelとして性能面で優位に立っています。

特にGflopsと性能の強い相関を示し、計算量を増やせば増やすほど性能が向上するというスケーリング法則が、Diffusion ModelとTransformerの組み合わせでも成り立つことを明確にした点が、この研究の重要な貢献と言えます。

この説明が、関連手法との比較についてより深く理解する助けとなれば幸いです。さらに具体的な点について知りたい場合は、お気軽に追加の質問をしてください。

Q:提案手法の限界と今後の課題はなんですか？

A:

提案手法DiT（Diffusion Transformer）の限界と今後の課題について、論文の内容や一般的な視点から考察し、詳しく説明します。

論文で示された成果は非常に印象的ですが、どのような手法にも限界と今後の探求すべき方向性があります。

DiTの限界

潜在空間への依存: DiTはLatent Diffusion Models (LDMs) [48] のフレームワークに基づいているため、生成品質は使用するVariational Autoencoder (VAE) の性能に大きく依存します。VAEのEncoderが情報を失ったり、Decoderがアーティファクトを生成したりする場合、その影響はDiTの最終的な画像品質にも現れます。論文中でも、異なるVAEデコーダーでの性能比較が行われています (Table 5)。VAEのボトルネックを解消するか、またはDiT自体をピクセル空間で効率的に動作させる方法を模索する必要があります。
計算コスト: 論文では計算効率の改善が強調されていますが、それでも高解像度・高品質な画像を生成するための計算コストは依然として高いです。DiT-XL/2は、512x512 ImageNet生成において524.6 Gflopsを使用します (Table 4)。Transformerはシーケンス長（ここではパッチ数）に対して計算量が二次的に増加する傾向があるため、より高解像度の画像（例：1024x1024以上）を生成する際には、計算量が爆発的に増加する可能性があります。パッチ数を効率的に扱うための工夫（例：Sparse Attention [7], Window Attentionなど）が必要になるかもしれません。
学習の長さ: State-of-the-Artの性能を達成するために、DiT-XL/2モデルは非常に長い期間学習されています（256x256 ImageNetで7Mステップ、512x512 ImageNetで3Mステップ）(Table 4)。これは計算資源と時間を大量に必要とします。学習効率をさらに向上させる方法が望まれます。
アーキテクチャの複雑さ: U-Netは skip connections を用いて異なる解像度の特徴を効率的に結合する構造を持っていますが、ViTベースのDiTは基本的にフラットなシーケンス処理を行います。空間的な構造を明示的に扱うU-Netの帰納バイアスがないことが、データ効率や特定のタスク（例：詳細な空間操作が必要なタスク）において不利に働く可能性も考えられます。ただし、論文ではImageNet生成タスクにおいてこれが性能のボトルネックになっていないことを示唆していますが、他の画像タスク（例：セマンティックセグメンテーションなど、DiTは直接解いていませんが）では異なるかもしれません。
サンプリング速度: Diffusion Model一般の課題として、サンプリングには多数のステップが必要であり、GANなどの他の生成モデルと比較して時間がかかります。論文ではサンプリングステップ数を増やしてもモデル容量不足を補えないことが示されており (Figure 10)、モデル自体を大きくすることの重要性が強調されています。サンプリング速度を高速化する手法（例：DDIM [55], score-based SDE solvers [56] の改良や、より少ないステップで高品質なサンプリングを可能にするモデル設計）は、DiTにとっても引き続き重要な課題です。

DiTの今後の課題と研究方向

さらなるスケーリング: 論文で示されたように、DiTはGflopsと性能に強い相関があります (Figure 8)。これは、現在のDiT-XL/2よりもさらに大規模なモデルや、より小さなパッチサイズ（より多くのトークン）で学習することで、さらなる性能向上が期待できることを示唆しています。Transformerのスケーリング法則 [14, 22] に従い、計算リソースの増加に伴ってどの程度性能が向上し続けるのかを探求することは重要な方向性です。
ピクセル空間DiT: 現在のDiTは潜在空間で動作していますが、計算効率の課題を克服できれば、高解像度ピクセル空間で直接動作するDiTも考えられます。これにより、VAEのボトルネックを回避し、より高品質な画像を生成できる可能性があります。計算量を削減する技術（例：効率的なAttentionメカニズム [7, 24], 階層的な構造の導入など）と組み合わせることで実現可能かもしれません。
テキスト条件付けモデルへの応用: 論文の結論でも述べられているように、DiTはDALL·E 2 [46] やStable Diffusion [48] のようなテキスト-画像生成モデルのバックボーンとして応用できる可能性があります。現在のStable DiffusionはLDMフレームワークとU-Netバックボーンを使用しており、テキスト条件付けにはcross-attentionを用いています。DiTをバックボーンに置き換えることで、テキスト-画像生成においてもスケーリングによる性能向上や計算効率の改善が期待されます。論文の実験ではクラス条件付けのみですが、テキストエンコーダーからの出力を効果的にDiTに条件付けする方法（cross-attention [60] や adaLN [40] の適用など）を開発・検証する必要があります。
多様な画像生成タスクへの適用: ImageNetのクラス条件付き生成だけでなく、インペインティング、アウトペインティング、画像編集、超解像度など、他の画像生成・編集タスクへのDiTの適用を探求すること。これらのタスクでは、入力画像の一部や追加のマスク情報などを条件として扱う必要があります。
より効率的なアーキテクチャの探求: 現在のDiTはViTの標準的な構造を踏襲していますが、Diffusion Modelの特性に合わせたTransformerの変種を開発することも考えられます。例えば、異なるタイムステップで必要な情報量や空間的な関連性が異なることを考慮した動的な計算割り当て [24] や、階層的な処理を組み合わせるなどが考えられます。
ゼロ初期化 (Zero Initialization) のさらなる理解と応用: adaLN-Zeroブロックは、各ブロックを初期状態でIdentity関数に近づけるゼロ初期化が性能向上に寄与することを示しました (Figure 5)。これはResNet [13] などでも有効な手法ですが、Diffusion ModelやTransformerにおけるこのメカニズムをさらに深く理解し、他のモデル設計やタスクに応用すること。
サンプリング効率と品質の両立: サンプリング速度を高速化しつつ、DiTが達成した高品質を維持する研究。教師なしのサンプル生成だけでなく、インタラクティブな生成やリアルタイム応用のためには、サンプリング速度の向上が不可欠です。

DiTは、Diffusion Modelのアーキテクチャ研究において、U-Net一辺倒だった状況に一石を投じ、Transformerの強力な能力を有効活用できることを示した重要な一歩です。しかし、前述のような限界や未解決の課題も多く存在します。今後の研究は、これらの課題を克服し、DiTベースのモデルをさらに高性能かつ汎用的に、そして効率的にすることを目指すことになるでしょう。