Diffusion Transformer (DiT)

Diffusion Transformer (DiT)は、ノイズ除去ネットワークに畳み込み型のU-Netではなくトランスフォーマーを用いる拡散モデルです。ノイズのある画像の潜在表現をパッチに分割し、Vision Transformer (ViT)が画像を処理するのと同様に、自己注意を使ってトークンの系列として処理し、各ノイズ除去ステップで取り除くノイズを予測します。William PeeblesとSaining Xieは、Scalable Diffusion Models with Transformers (ICCV 2023)でこのモデルを発表しました。同論文ではDiT-XL/2がクラス条件付きImageNet 256×256でFréchet Inception Distance (FID) 2.27を達成し、潜在拡散モデルによる従来の最高値3.60を上回りました。また、論文はDiTの計算量を増やすと画像品質が着実に向上することを示しました。その後、Stable Diffusion 3、FLUX.1、OpenAIのSoraなど、ほとんどの大規模な画像・動画生成モデルで、U-Netに代わってトランスフォーマーのノイズ除去ネットワークが使われています。

Diffusion Transformerの仕組み#

DiTは、Stable Diffusionが使用する事前学習済み変分オートエンコーダー (VAE)の潜在空間で動作します。このVAEのエンコーダーは画像を8分の1にダウンサンプリングするため、256×256×3の画像は32×32×4の潜在表現になります。パッチ化レイヤーは、ノイズのある潜在表現をp×pのパッチに分割し、それぞれを線形変換してトークンにします。パッチサイズが2、4、8の場合、トークン数はそれぞれ256、64、16です。正弦・余弦位置埋め込みで各トークンの位置を示し、トランスフォーマーブロックのスタックで系列を処理します。最後の線形レイヤーは、各トークンを予測ノイズと分散を含むパッチに戻し、潜在表現と同じ形状に再構成します。拡散タイムステップとクラスラベル、またはテキスト画像モデルの場合はテキストプロンプトが、各ブロックを条件付けます。生成では、純粋なノイズから最後のステップまで予測を繰り返し、VAEデコーダーがクリーンな潜在表現をピクセルに変換します。

論文では、各ブロックに条件を与える4つの方法を比較しました。系列への追加トークン、クロスアテンション、適応レイヤーノルム (adaLN)、adaLN-Zeroです。adaLN-Zeroは最も低いFIDを達成し、計算効率も最も優れていました。小さなネットワークがタイムステップ埋め込みとクラス埋め込みの和を読み取り、各レイヤーノルムの後に適用するシフトとスケール、および各残差分岐のスケールを調整するゲートを回帰します。ゲートはゼロから始まるため、すべてのブロックは恒等関数として学習を開始します。このPyTorchのコードは、潜在表現をパッチ化し、1つのブロックのアテンション分岐を適用します。

import torch
from torch import nn

torch.manual_seed(0)
# A 256x256 image becomes a 4x32x32 latent; 2x2 patches turn it into 256 tokens of width 384 (DiT-S/2)
latent = torch.randn(1, 4, 32, 32)
tokens = nn.Conv2d(4, 384, kernel_size=2, stride=2)(latent).flatten(2).transpose(1, 2)

# adaLN-Zero: the timestep and class embedding regresses a shift, scale, and gate for the attention branch
condition = torch.randn(1, 384)
modulation = nn.Linear(384, 3 * 384)
nn.init.zeros_(modulation.weight)
nn.init.zeros_(modulation.bias)
shift, scale, gate = modulation(condition).unsqueeze(1).chunk(3, dim=-1)

# Self-attention branch of a DiT block, modulated by the condition and scaled by the gate
attention = nn.MultiheadAttention(384, num_heads=6, batch_first=True)
h = nn.functional.layer_norm(tokens, (384,)) * (1 + scale) + shift
out = tokens + gate * attention(h, h, h)[0]
print(f"tokens {tuple(tokens.shape)}, block is identity at initialization: {torch.equal(out, tokens)}")

tokens (1, 256, 384), block is identity at initialization: Trueと出力されます。幅384の256個のトークンは最小モデルのDiT-S/2に相当し、ゼロゲートにより学習が進むまでは入力が変更されずに通過します。完全なDiTブロックでは、MLP分岐用に2つ目のシフト、スケール、ゲートも回帰します。

トークン数と計算量によるスケーリング#

論文では、12層・幅384のDiT-Sから28層・幅1,152のDiT-XLまで、4つのモデルサイズを学習し、それぞれにパッチサイズ8、4、2を用いました。全12モデルにおいて、計算量がより深いネットワーク、より幅広いネットワーク、あるいはより多くのトークンのどれに由来するかにかかわらず、FIDはGflops単位のフォワードパス計算量に対応していました。パッチサイズを半分にすると、パラメーター数はほとんど増えずにトークン数と計算量が4倍になります。DiT-XL/2は、8,400万パラメーターのVAEを除いて6億7,500万パラメーターを持ち、256×256でフォワードパスあたり119 Gflopsを使用します。512×512では、64×64×4の潜在表現から1,024トークンを処理し、525 GflopsでFID 3.04を達成します。自己注意は各トークンを他のすべてのトークンと比較するため、高解像度や長時間の動画では計算コストが二次的に増加します。Ultralytics YOLO26は、640×640画像で400位置を持つ1/32スケールのステージだけに注意機構を配置し、それ以外では畳み込みを維持してリアルタイム検出を実現し、このコストを回避します。

DiT、U-Net、ViTの比較#

• U-Netノイズ除去ネットワーク: DDPMやStable Diffusion 1.x、2.xのU-Netは、複数の解像度で動作する畳み込み型のエンコーダー・デコーダーです。同じ解像度のレベル間にはスキップ接続があり、低解像度では注意機構が追加されます。DiTは1つのトークン解像度を維持し、すべてのブロックで全体に注意を向けます。また、画像固有の設計上の選択が少ないため、大規模言語モデルのトランスフォーマーと同様にスケールします。
• Vision Transformer: ViTは、分類などのタスクに向けて画像をエンコードします。DiTはViTのパッチ埋め込みとトランスフォーマーブロックを再利用し、adaLN-Zeroでタイムステップ条件付けを加え、すべてのパッチについて予測を出力します。
• 潜在拡散とフローマッチング: 潜在拡散モデルはVAEの潜在空間でノイズ除去を行う全体的な枠組みであり、拡散またはフローマッチングは学習目的です。DiTはその枠組みの中のノイズ除去ネットワークで、現在の生成モデルではDiTバックボーンとフローマッチングの目的関数が組み合わされています。

テキスト画像・動画DiT#

• マルチモーダルDiT (MMDiT): Stable Diffusion 3は、画像トークンとテキストトークンにモダリティごとに別々の重みを適用し、単一の注意機構で結合します。そのため、プロンプトと画像の間で双方向に情報が流れます。これにより、タイポグラフィとプロンプトへの追従性が向上しました。Stable Diffusion 3.5 Largeでは、この設計が約80億パラメーターまでスケールしています。
• FLUX.1: Black Forest LabsのFLUX.1は、テキスト画像生成を行う120億パラメーターのRectified Flowトランスフォーマーです。
• 動画: OpenAIのSora技術レポートでは、動画を潜在空間に圧縮して時空間パッチに分割する拡散トランスフォーマーが説明されています。これにより、長さ、解像度、アスペクト比の異なるクリップを、長さの異なるトークン系列として表現できます。動画生成では、Metaの最大規模のMovie Genモデルは300億パラメーターのトランスフォーマーで、7万3,000動画トークンのコンテキストは、毎秒16フレームで16秒間の動画に相当します。