U-Netの終焉? Diffusion Transformer (DiT) アーキテクチャ徹底解剖 (Sora 2 & FLUX.2)

Diffusion Transformer (DiT)

U-Netの支配は終わった。2025年の生成AI革命を支えるアーキテクチャ、それがここにある。

長年にわたり、U-Netは拡散モデル(Diffusion Model)における不動の王者でした。Stable Diffusion 1.5から初期の動画モデルに至るまで、畳み込みネットワーク(CNN)がその中心に君臨していたのです。しかし、Sora 2FLUX.2の登場は、新たな現実を決定づけました。未来はDiffusion Transformer (DiT) のものです。

私たちは今、AIにおける「脳の移植手術」を目の当たりにしています。バックボーンを畳み込みニューラルネットワークからTransformerへと入れ替えることで、モデルは計算量(コンピュート)に応じて予測可能な形でスケーリングし、物理法則を理解し、これまでにないレベルで動画の時間的整合性を維持できるようになりました。

なぜU-Netは敗れ、DiTが勝利したのか

DiTを理解するには、まずそれが何を置き換えるのかを理解する必要があります。U-Netアーキテクチャ(SDXLやそれ以前のモデルで使用)は、畳み込み(Convolution) に依存しています。畳み込みは、局所的な詳細(エッジやテクスチャなど)の処理には非常に優れていますが、大域的なコンテキスト(Global Context) の把握には苦戦します。画像全体を一度に「見る」ことが苦手なため、生成されたオブジェクトの不整合や、動画における歪みが生じやすいのです。

Diffusion Transformer (DiT) は、GPT-4を支えるアーキテクチャ(Transformer)を視覚生成に応用したものです。単語の代わりに、画像や動画フレームの「パッチ(断片)」を予測します。

DiTの優位性:

  • スケーラビリティ: CNNとは異なり、Transformerはパラメータ数やデータを増やすことで、その性能が線形かつ予測可能に向上します(スケーリング則)。
  • 大域的アテンション(Global Attention): 画像内のすべてのパッチが、他のすべてのパッチと瞬時に「会話」できます。左端にある手が、右端にある腕の動きを正確に把握できるのはこのためです。
  • コンテキスト認識: 例えばFLUX.2は、240億パラメータの視覚言語モデル(Mistral-3)をDiTと結合しており、単なるキーワードのマッチングではなく、複雑なプロンプトや物理現象を理解することが可能です。

コード解説: シンプルなDiTブロック

Sora 2やFLUX.2の核心にあるのは「魔法」ではありません。それは、以下のブロックの積み重ねです。以下は、単一の DiTブロック のPyTorch実装です。ここで注目すべきは AdaLN (Adaptive Layer Norm) の使用です。これは生成プロセスにおいて、タイムステップやテキストプロンプトを注入する「操縦桿」の役割を果たします。

import torch
import torch.nn as nn

class DiTBlock(nn.Module):
    """
    単一のDiffusion Transformerブロック。
    
    Args:
        hidden_size (int): Transformerトークンの次元数。
        num_heads (int): アテンションヘッドの数。
        mlp_ratio (float): MLP内の隠れ層次元の倍率。
    """
    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.attn = nn.MultiheadAttention(hidden_size, num_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        
        # Feed Forward Network (MLP)
        mlp_hidden_dim = int(hidden_size * mlp_ratio)
        self.mlp = nn.Sequential(
            nn.Linear(hidden_size, mlp_hidden_dim),
            nn.GELU(),
            nn.Linear(mlp_hidden_dim, hidden_size),
        )
        
        # Adaptive Layer Norm (AdaLN) 変調
        # 条件付け(時間/ラベル)に基づいてシフト(gamma)とスケール(beta)を予測する
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
        )

    def forward(self, x, c):
        """
        x: 入力トークン [Batch, Sequence_Length, Hidden_Size]
        c: 条件付け埋め込み (Time + Text) [Batch, Hidden_Size]
        """
        # 1. 条件付けベクトルから変調パラメータを回帰予測
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
            self.adaLN_modulation(c).chunk(6, dim=1)
        )
        
        # 2. AdaLNを用いたSelf-Attentionブロック
        x_norm1 = (1 + scale_msa.unsqueeze(1)) * self.norm1(x) + shift_msa.unsqueeze(1)
        attn_output, _ = self.attn(x_norm1, x_norm1, x_norm1)
        x = x + gate_msa.unsqueeze(1) * attn_output
        
        # 3. AdaLNを用いたMLPブロック
        x_norm2 = (1 + scale_mlp.unsqueeze(1)) * self.norm2(x) + shift_mlp.unsqueeze(1)
        mlp_output = self.mlp(x_norm2)
        x = x + gate_mlp.unsqueeze(1) * mlp_output
        
        return x

# 使用例
# batch_size=4, seq_len=256 (パッチ数), dim=1152
x = torch.randn(4, 256, 1152) 
c = torch.randn(4, 1152) # 条件付けベクトル (時間 + テキスト)
model = DiTBlock(hidden_size=1152, num_heads=16)
output = model(x, c)
print(f"Output Shape: {output.shape}")

アーキテクチャの可視化

flowchart TD
    subgraph "前処理 (Preprocessing)"
    A["入力画像/動画"] --> B["VAE エンコーダー"]
    B --> C["潜在表現 (Latent)"]
    C --> D["パッチ化 (Patchify)"]
    end

    subgraph "Diffusion Transformer (DiT)"
    D --> E["線形埋め込み + 位置エンコーディング"]
    E --> F["DiT ブロック 1 (Attn + MLP)"]
    F --> G["DiT ブロック 2 (Attn + MLP)"]
    G --> H["... DiT ブロック N ..."]
    H --> I["パッチ結合 (Depatchify)"]
    end

    subgraph "条件付け (Conditioning)"
    J["テキストプロンプト"] --> K["T5/CLIP/Mistral エンコーダー"]
    L["タイムステップ (ノイズレベル)"] --> M["MLP"]
    K --> N["条件付けベクトル (c)"]
    M --> N
    N --"AdaLN経由で注入"--> F
    N --"AdaLN経由で注入"--> G
    end

    subgraph "出力 (Output)"
    I --> O["予測されたノイズ / 潜在変数"]
    O --> P["VAE デコーダー"]
    P --> Q["最終生成メディア"]
    end
    
    style A fill:#333,stroke:#fff,color:#fff
    style Q fill:#333,stroke:#fff,color:#fff
    style F fill:#000,stroke:#0f0,color:#fff
    style G fill:#000,stroke:#0f0,color:#fff

ステップ・バイ・ステップ: FLUX.2をローカルで動かす

Sora 2はクローズドソースであるため、高性能なDiTモデルに入門するにはFLUX.2が最適です。ただし、これは320億パラメータという「怪物」であるため、十分なVRAM(少なくとも24GB)を用意するか、8-bit量子化を使用する必要があります。

  1. 環境の準備:
    FLUX.2アーキテクチャをサポートする最新の diffusers ライブラリが必要です。

    pip install -U diffusers transformers accelerate sentencepiece
    
  2. Hugging Faceでの認証:
    black-forest-labs/FLUX.2-dev へのアクセスには、Hugging Face上でのライセンス同意が必要です。

    huggingface-cli login
    # HFトークンを入力
    
  3. パイプラインの実行 (FP8量子化):
    コンシューマー向けのRTX 4090などで動作させるため、8-bitでロードします。

    import torch
    from diffusers import FluxPipeline
    
    # FP8最適化を有効にしてモデルをロード
    pipe = FluxPipeline.from_pretrained(
        "black-forest-labs/FLUX.2-dev", 
        torch_dtype=torch.bfloat16
    )
    
    # VRAM節約のためにCPUオフロードを有効化
    pipe.enable_model_cpu_offload() 
    
    prompt = "A cinematic shot of a futuristic datacenter, glowing blue lights, 8k resolution, photorealistic"
    
    image = pipe(
        prompt,
        height=1024,
        width=1024,
        guidance_scale=3.5,
        num_inference_steps=50,
        max_sequence_length=512
    ).images[0]
    
    image.save("flux2_output.png")
    
  4. DiTのためのプロンプト戦略:
    • 説明的に書く: 初期のSDモデルで必要だった「呪文の羅列(例: masterpiece, best quality)」とは異なり、FLUX.2のようなDiTモデルは強力なLLMエンコーダー(Mistral/T5)を使用しています。自然言語を理解できるため、タグではなく文章で書いてください。
    • 物理/照明への言及: DiTは光の伝播(Light Transport)のレンダリングに優れています。「窓から差し込むボリューメトリックライト(volumetric lighting)」のように、照明条件を明示的に記述すると効果的です。

Diffusion Transformerへの移行は、単なるアップデートではありません。それは今後5年間の生成AIにおける標準規格です。Sora 2のアプリを使う場合でも、FLUX.2をローカルで動かす場合でも、あなたはピクセルや畳み込みではなく、パッチとアテンションで「思考」するモデルと対話しているのです。このアーキテクチャを習得することは、デジタルクリエイションの未来を習得することに他なりません。