Diffusion Transformer (DiT)

扩散 Transformer（DiT）是一种扩散模型，其去噪网络采用Transformer，而不是卷积式 U-Net。它会将带噪图像潜变量切分成多个图块，像视觉 Transformer（ViT）处理图像那样，以自注意力处理 token 序列，并在每个去噪步骤预测要移除的噪声。William Peebles 和 Saining Xie 在《Scalable Diffusion Models with Transformers》（ICCV 2023）中提出了这一方法，其中 DiT-XL/2 在类别条件式 ImageNet 256×256 上的弗雷歇 Inception 距离（FID）达到 2.27，优于此前潜空间扩散模型创下的 3.60。论文表明，随着 DiT 的计算量增加，图像质量会持续提升；此后，Transformer 去噪器已在大多数大型图像和视频生成器中取代 U-Net，包括 Stable Diffusion 3、FLUX.1 和 OpenAI 的 Sora。

扩散 Transformer 的工作原理#

DiT 在潜空间中运行，该潜空间属于 Stable Diffusion 使用的预训练变分自编码器（VAE）；其编码器会将图像下采样 8 倍，因此 256×256×3 的图像会变成 32×32×4 的潜变量。patchify 层会将带噪潜变量拆分为 p×p 图块，并将每个图块线性嵌入为一个 token；图块大小为 2、4 和 8 时，分别会得到 256、64 和 16 个 token。正弦余弦位置嵌入会标记每个 token 的来源位置，多个 Transformer 块会处理该序列，最后的线性层则将每个 token 转换回包含预测噪声和方差的图块，再将其重新组合成潜变量的形状。扩散时间步和类别标签，或文生图模型中的文本提示词，都会对每个块进行条件控制。生成过程会从纯噪声开始重复预测，直到最后一步；随后，VAE 解码器会将干净的潜变量转换为像素：

论文比较了四种将条件输入每个块的方法：序列中的额外 token、交叉注意力、自适应层归一化（adaLN）和 adaLN-Zero。adaLN-Zero 的 FID 最低，计算效率也最高。一个小型网络会读取时间步嵌入和类别嵌入的和，并回归一个偏移量和缩放量，将它们应用于每个层归一化之后；此外，它还会生成一个门控值，用于缩放每个残差分支。门控值从零开始，因此每个块开始训练时都是恒等函数。这段 PyTorch代码会将潜变量切分为图块，并应用其中一个块的注意力分支：

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 个 token 与最小模型 DiT-S/2 相符；零门控会让输入保持原样，直到训练对其进行调整。完整的 DiT 块还会为其 MLP 分支回归第二组偏移量、缩放量和门控值。

Token 数量与计算量的扩展#

论文训练了四种不同规模的模型，从 DiT-S（12 层，宽度 384）到 DiT-XL（28 层，宽度 1,152），每种模型都使用 8、4 和 2 三种图块大小。在全部 12 个模型中，FID 都与前向计算量（Gflops）相关，无论计算量来自更深或更宽的网络，还是更多的 token。将图块大小减半会使 token 数量和计算量都增加到四倍，但几乎不会增加参数。DiT-XL/2 有 6.75 亿个参数（不包括 8400 万参数的 VAE），在 256×256 分辨率下每次前向计算需要 119 Gflops。在 512×512 分辨率下，它会处理来自 64×64×4 潜变量的 1,024 个 token，计算量为 525 Gflops，FID 达到 3.04。由于自注意力会比较每个 token 与其他所有 token，因此更高分辨率和更长的视频都会使其成本按平方增长。Ultralytics YOLO26通过仅在 1/32 尺度阶段应用注意力来避免这类成本；对于 640×640 图像，该阶段有 400 个位置，其余部分则都使用卷积，以实现实时检测。

DiT、U-Net 与 ViT 的对比#

• U-Net 去噪器：DDPM 和 Stable Diffusion 1.x、2.x 中的 U-Net 是卷积式编码器-解码器，会在多个分辨率下工作，并在对应层级之间使用跳跃连接，同时在较低分辨率下加入注意力。DiT 保持单一 token 分辨率，在每个块中进行全局注意力，而且图像专属的设计选择更少，因此它可以像大型语言模型中的 Transformer 一样扩展。
• 视觉 Transformer：ViT 会为分类等任务编码图像。DiT 复用了 ViT 的图块嵌入和 Transformer 块，通过 adaLN-Zero 添加时间步条件，并为每个图块输出预测结果。
• 潜空间扩散与流匹配：潜空间扩散模型是指在 VAE 潜空间中执行去噪的整体框架，而扩散或流匹配则是训练目标。DiT 是该框架中的去噪网络；当前的生成器会将 DiT 主干网络与流匹配目标结合使用。

文生图与视频 DiT#

• 多模态 DiT（MMDiT）：Stable Diffusion 3 会通过针对每种模态使用单独权重来处理图像 token 和文本 token，并在单次注意力运算中将它们结合起来，使提示词和图像能够双向传递信息。这一设计改进了排版效果和提示词遵循能力，而 Stable Diffusion 3.5 Large 将其扩展到了约 80 亿个参数。
• FLUX.1：Black Forest Labs 的 FLUX.1是一种拥有 120 亿个参数的整流流 Transformer，用于文生图生成。
• 视频：OpenAI 的 Sora 技术报告介绍了一种扩散 Transformer，它会将视频压缩到潜空间并拆分为时空图块，因此不同长度、分辨率和宽高比的片段都会变成长度各异的 token 序列。在视频生成方面，Meta 最大的 Movie Gen 模型是一种拥有 300 亿个参数的 Transformer，其上下文中的 73,000 个视频 token 相当于 16 帧/秒的视频中的 16 秒内容。