Diffusion Transformer（DiT）架构分析：从 U-Net 到 Transformer 的转变

引言
背景：为什么要放弃 U-Net？
- U-Net 的局限
- Transformer 的优势
DiT 架构详解
模型变体与缩放
- 分块大小的影响
应用 Classifier-Free Guidance
DiT 的后续影响
PyTorch 实现实践
- 自定义训练
总结

引言

2022 年底发表的《Scalable Diffusion Models with Transformers》（William Peebles & Saining Xie，2023 ICCV）成为了将图像生成的骨干网络从 U-Net 转向 Transformer 的起点。这篇论文提出的 DiT（Diffusion Transformer）此后成为了贯穿 OpenAI 的 SORA、Stability AI 的 Stable Diffusion 3，直至 2025 年 ICLR 发表的 Dynamic DiT 的核心架构。

背景：为什么要放弃 U-Net？

U-Net 的局限

传统的 DDPM、LDM（Latent Diffusion Model）都以 U-Net 作为骨干网络。U-Net 在图像生成中表现良好，但存在根本性的局限：

没有缩放定律：即使扩大模型规模，性能提升也无法预测
架构复杂：需要管理 skip connection 以及各种分辨率的 feature map
计算效率低：高分辨率下内存/计算成本会急剧增加
难以与其他模态集成：与文本、视频等的整合并不自然

Transformer 的优势

而 Transformer（ViT）则具备：

经过验证的缩放定律：性能提升与参数量成比例（已在 LLM 中得到证实）
架构简单：Self-attention + FFN 的重复堆叠
不限模态：文本、图像、视频都可以用相同结构处理
硬件优化：在 GPU/TPU 上高效并行处理

DiT 架构详解

整体流水线

DiT 运行在 Latent Diffusion 框架之上：

通过 VAE 编码器将图像转换到潜在空间（latent space）
将潜在表示分割为分块（与 ViT 相同）
通过DiT 块（Transformer）预测噪声
通过 VAE 解码器还原图像

分块嵌入

import torch
import torch.nn as nn

class PatchEmbed(nn.Module):
    """将潜在表示转换为分块序列"""
    def __init__(self, patch_size=2, in_channels=4, embed_dim=1152):
        super().__init__()
        self.proj = nn.Conv2d(
            in_channels, embed_dim,
            kernel_size=patch_size, stride=patch_size
        )

    def forward(self, x):
        # x: (B, C, H, W) -> (B, N, D)
        # 例如: (B, 4, 32, 32) -> (B, 256, 1152) (patch_size=2)
        x = self.proj(x)  # (B, D, H/p, W/p)
        x = x.flatten(2).transpose(1, 2)  # (B, N, D)
        return x

输入是将 256x256 图像通过 VAE 转换为 32x32 的潜在表示后，再以 patch_size=2 进行分割，得到 16x16=256 个 token。

条件化机制 — 四种变体

DiT 实验了四种注入时间步（t）与类别标签（c）的方式：

1. In-context Conditioning

class InContextConditioning(nn.Module):
    """将 t 和 c 作为附加 token 拼接到序列中"""
    def __init__(self, embed_dim):
        super().__init__()
        self.t_embed = TimestepEmbedder(embed_dim)
        self.c_embed = LabelEmbedder(1000, embed_dim)

    def forward(self, x, t, c):
        t_token = self.t_embed(t).unsqueeze(1)  # (B, 1, D)
        c_token = self.c_embed(c).unsqueeze(1)  # (B, 1, D)
        x = torch.cat([t_token, c_token, x], dim=1)  # (B, N+2, D)
        return x

2. Cross-Attention

class CrossAttentionBlock(nn.Module):
    """通过 cross-attention 注入 t 和 c"""
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(embed_dim, num_heads)
        self.cross_attn = nn.MultiheadAttention(embed_dim, num_heads)
        self.mlp = MLP(embed_dim)

    def forward(self, x, cond):
        x = x + self.self_attn(x, x, x)[0]
        x = x + self.cross_attn(x, cond, cond)[0]
        x = x + self.mlp(x)
        return x

3. Adaptive Layer Norm（adaLN）

class AdaLNBlock(nn.Module):
    """根据条件调整 LayerNorm 的 scale/shift"""
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False)
        self.attn = nn.MultiheadAttention(embed_dim, num_heads)
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(embed_dim, 2 * embed_dim)
        )

    def forward(self, x, cond):
        shift, scale = self.adaLN_modulation(cond).chunk(2, dim=-1)
        x = self.norm(x) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
        x = x + self.attn(x, x, x)[0]
        return x

4. adaLN-Zero（最终选择）

class DiTBlock(nn.Module):
    """adaLN-Zero：初始化时残差连接从零开始"""
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim, elementwise_affine=False)
        self.norm2 = nn.LayerNorm(embed_dim, elementwise_affine=False)
        self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, 4 * embed_dim),
            nn.GELU(approximate="tanh"),
            nn.Linear(4 * embed_dim, embed_dim),
        )
        # 6 个参数：gamma1, beta1, alpha1, gamma2, beta2, alpha2
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(embed_dim, 6 * embed_dim)
        )
        # 将 alpha 初始化为 0 -> 训练初期表现为 identity function
        nn.init.zeros_(self.adaLN_modulation[-1].weight)
        nn.init.zeros_(self.adaLN_modulation[-1].bias)

    def forward(self, x, cond):
        # cond: (B, D) - 时间步 + 类别嵌入之和
        shift_msa, scale_msa, gate_msa, \
        shift_mlp, scale_mlp, gate_mlp = \
            self.adaLN_modulation(cond).chunk(6, dim=-1)

        # Self-Attention with adaLN
        h = self.norm1(x)
        h = h * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)
        h = self.attn(h, h, h)[0]
        x = x + gate_msa.unsqueeze(1) * h  # gate 从 0 开始

        # FFN with adaLN
        h = self.norm2(x)
        h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
        h = self.mlp(h)
        x = x + gate_mlp.unsqueeze(1) * h  # gate 从 0 开始

        return x

adaLN-Zero 的核心：将 gate 参数初始化为 0，使每个 DiT 块在训练初期都表现为 identity function。这大幅提升了深层网络的训练稳定性。

完整的 DiT 模型

class DiT(nn.Module):
    def __init__(
        self,
        input_size=32,
        patch_size=2,
        in_channels=4,
        embed_dim=1152,
        depth=28,
        num_heads=16,
        num_classes=1000,
    ):
        super().__init__()
        self.patch_embed = PatchEmbed(patch_size, in_channels, embed_dim)
        num_patches = (input_size // patch_size) ** 2  # 256

        # 位置嵌入
        self.pos_embed = nn.Parameter(
            torch.zeros(1, num_patches, embed_dim)
        )

        # 时间步 & 类别嵌入
        self.t_embedder = TimestepEmbedder(embed_dim)
        self.y_embedder = LabelEmbedder(num_classes, embed_dim)

        # DiT 块
        self.blocks = nn.ModuleList([
            DiTBlock(embed_dim, num_heads) for _ in range(depth)
        ])

        # 最终层
        self.final_layer = FinalLayer(embed_dim, patch_size, in_channels * 2)

    def forward(self, x, t, y):
        # 分块嵌入 + 位置嵌入
        x = self.patch_embed(x) + self.pos_embed

        # 条件嵌入
        t_emb = self.t_embedder(t)  # (B, D)
        y_emb = self.y_embedder(y)  # (B, D)
        cond = t_emb + y_emb        # (B, D)

        # 通过 DiT 块
        for block in self.blocks:
            x = block(x, cond)

        # 预测噪声与方差
        x = self.final_layer(x, cond)  # (B, N, 2*C*p*p)
        return x

模型变体与缩放

模型	depth	embed_dim	heads	参数量	FID-50K
DiT-S/2	12	384	6	33M	68.4
DiT-B/2	12	768	12	130M	43.5
DiT-L/2	24	1024	16	458M	9.62
DiT-XL/2	28	1152	16	675M	2.27

核心发现：DiT 遵循与 LLM 相同的缩放定律。 模型规模越大，FID 就越能持续改善。

分块大小的影响

模型	patch_size	Token 数	Gflops	FID
DiT-XL/8	8	16	119	14.2
DiT-XL/4	4	64	525	5.11
DiT-XL/2	2	256	1121	2.27

分块越小（token 数越多），性能越好，但计算量也会大幅增加。

应用 Classifier-Free Guidance

def sample_with_cfg(model, z, class_labels, cfg_scale=4.0):
    """使用 Classifier-Free Guidance 进行采样"""
    # 同时进行条件与无条件预测
    z_combined = torch.cat([z, z], dim=0)
    y_combined = torch.cat([class_labels, torch.full_like(class_labels, 1000)])

    for t in reversed(range(1000)):
        t_batch = torch.full((z_combined.shape[0],), t, device=z.device)
        noise_pred = model(z_combined, t_batch, y_combined)

        # 应用 CFG
        cond_pred, uncond_pred = noise_pred.chunk(2)
        guided_pred = uncond_pred + cfg_scale * (cond_pred - uncond_pred)

        # DDPM step
        z = ddpm_step(z, guided_pred, t)

    return z

在 CFG scale=1.5 时，DiT-XL/2 取得了 FID 2.27，当时在 ImageNet 256x256 上创下了 SOTA。

DiT 的后续影响

SORA（OpenAI，2024）

将 DiT 扩展到视频领域的Spatial-Temporal DiT
3D 分块嵌入（空间 + 时间）
支持可变分辨率/时长

Stable Diffusion 3（Stability AI，2024）

DiT + Flow Matching
MM-DiT：在同一个 Transformer 中同时处理文本与图像
无需单独的 text encoder，直接进行 joint attention

Dynamic DiT（ICLR 2025）

通过 Early Exit 机制提升推理效率
噪声较大的早期步骤使用浅层，生成细节时使用深层
相比传统 DiT 减少 40% 推理时间

PyTorch 实现实践

# 克隆 DiT 官方代码
git clone https://github.com/facebookresearch/DiT.git
cd DiT

# 环境设置
pip install torch torchvision diffusers accelerate

# 下载预训练模型
python download.py DiT-XL/2

# 生成样本
python sample.py \
  --model DiT-XL/2 \
  --image-size 256 \
  --num-sampling-steps 250 \
  --seed 42 \
  --class-labels 207 360 387 974 88 979 417 279

自定义训练

from accelerate import Accelerator
from torch.utils.data import DataLoader

accelerator = Accelerator()

# 初始化模型
model = DiT(
    input_size=32,  # VAE latent size
    patch_size=2,
    in_channels=4,
    embed_dim=1152,
    depth=28,
    num_heads=16,
    num_classes=1000,
)

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=1e-4,
    weight_decay=0.0
)

model, optimizer, dataloader = accelerator.prepare(
    model, optimizer, dataloader
)

# 训练循环
for epoch in range(epochs):
    for batch in dataloader:
        images, labels = batch
        # VAE 编码
        with torch.no_grad():
            latents = vae.encode(images).latent_dist.sample() * 0.18215

        # 随机时间步
        t = torch.randint(0, 1000, (latents.shape[0],), device=latents.device)

        # 添加噪声
        noise = torch.randn_like(latents)
        noisy_latents = scheduler.add_noise(latents, noise, t)

        # 预测噪声
        noise_pred = model(noisy_latents, t, labels)
        loss = F.mse_loss(noise_pred, noise)

        accelerator.backward(loss)
        optimizer.step()
        optimizer.zero_grad()

总结

DiT 在图像生成领域引领了架构范式的转变：

将 U-Net 的复杂结构替换为纯 Transformer
证明了与 LLM 相同的缩放定律同样适用于图像生成
通过 adaLN-Zero 实现了深层 Transformer 的稳定训练
成为 SORA、SD3 等后续模型的基础架构

✅ 测验：DiT 架构理解检查（8 题）

Q1. DiT 选择使用 Transformer 而非 U-Net 的主要原因是什么？

经过验证的缩放定律、简单的架构，以及易于与其他模态集成。

Q2. 在 DiT 中，图像是如何转换为 token 序列的？

先通过 VAE 转换为潜在表示（latent），再按分块进行分割并做线性投影。

Q3. adaLN-Zero 中的「Zero」指的是什么？

将 gate 参数初始化为 0，使每个块在训练初期都表现为 identity function。

Q4. DiT 实验的四种条件化方式中，最终选择了哪一种？

adaLN-Zero。在 In-context、Cross-Attention、adaLN、adaLN-Zero 中，它取得了最好的 FID。

Q5. 分块越小会带来怎样的权衡？

token 数增加，图像质量（FID）随之提升，但计算量（Gflops）会按平方比例增长。

Q6. Classifier-Free Guidance 在 DiT 中是如何实现的？

训练时对类别标签进行随机丢弃，推理时则用条件预测与无条件预测的加权和进行采样。

Q7. DiT-XL/2 在 ImageNet 256x256 上的 FID-50K 分数是多少？

2.27，当时创下了 class-conditional 生成的 SOTA。

Q8. Dynamic DiT（ICLR 2025）的核心思路是什么？

通过 Early Exit 机制，根据噪声水平动态调整所使用的层数，从而将推理时间减少 40%。