0%
💡 왼쪽 원문을 읽으면서 오른쪽에 따라 써보세요. Tab 키로 힌트를 받을 수 있습니다.
引言
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%。
현재 단락 (1/230)
2022 年底发表的《**Scalable Diffusion Models with Transformers**》（William Peebles & Saining Xie，2023 ICCV）...
작성 글자: 0자원문 글자: 7,874자작성 단락: 0/230