DiT 扩散Transformer

DiT（Diffusion Transformer）把扩散模型中去噪网络从传统的 U-Net 换成了 Transformer，是近年来生成式 AI 架构最重要的一次范式转移。它继承自 Vision Transformer（ViT）的 patch 化思想，并与 Transformer 架构的核心机制深度结合。

2022 年 DiT 原论文发表时，扩散模型的主流骨干还是卷积 U-Net；到 2024—2025 年，行业格局已彻底改变——Stable Diffusion 3、Flux、Sora、HunyuanVideo、Wan 2.1 等几乎所有主流生成模型都采用了 DiT 或其变体作为骨干。促成这一转变的关键不仅是”Transformer 比 U-Net 更强”，更重要的是 DiT 展示了与大语言模型类似的 scaling law（标度律）：投入更多参数和算力，生成质量就持续、可预测地提升。这种”可扩展性”正是工业级大模型最看重的特质。

理解 DiT，也就理解了当今文生图、文生视频模型背后的骨干设计。

核心直觉

想象你有一幅大画，先把它裁成一张张同样大小的小卡片（patch），按从左到右、从上到下的顺序摊成一长串。每一张小卡片就是一个”故事段落”。

现在让一位 Transformer 编辑同时阅读所有段落。它的工作是：每读一轮，就把所有段落对照着看一遍，判断哪些地方有噪点、哪些地方颜色不协调，然后逐一修订。读的轮数越多，段落之间越能互相照应，最终把噪点擦干净、拼成一幅连贯的画面。

这就是 DiT 的精髓：把图像（或潜空间特征图）拆成 patch 序列，让 Transformer 的全局注意力来协调去噪。相比 U-Net 那种”先压缩再还原”的卷积管线，Transformer 的优势在于每个 patch 一开始就能”看到”整幅画的全局关系。

原理详解

从 U-Net 到 Transformer

经典扩散模型（如 Stable Diffusion 的 U-Net）使用卷积 U-Net 作为去噪骨干。U-Net 通过编码器下采样、解码器上采样、以及 skip connection 来聚合多尺度特征，在图像任务上效果出色，但其归纳偏置（局部性、平移不变性）也限制了模型在规模放大时的收益。

Peebles 和 Xie 在 2022 年的论文《Scalable Diffusion Models with Transformers》中提出：用纯 Transformer 替换 U-Net，并系统验证了 Transformer 在扩散模型上同样具备良好的可扩展性（scalability）。这篇工作后来被称为 DiT。

patch 化：图像变成 token 序列

DiT 的输入是扩散模型潜空间里的噪声特征图 z（由 VAE 编码器得到，通常空间分辨率较低，例如 32x32 或 64x64）。DiT 不直接处理二维网格，而是先把 z 切成不重叠的 patch：

把特征图按 p x p 的块切分（例如 p=2，即每 2x2 个像素归为一个 patch）。
每个 patch 展平后用一个线性投影映射成一个 token 向量。
所有 token 按光栅顺序排成序列，再加上位置编码（position embedding）告诉 Transformer 每个 patch 的空间位置。

若潜空间特征图大小为 I x I、通道数为 C、patch 大小为 p，则 token 数量为 (I/p) x (I/p)，每个 token 的维度由线性投影决定。patch 越大，序列越短、计算越省，但空间细节损失越多；patch 越小，序列越长、表达越精细，但计算量急剧上升。

与 ViT 的关系

DiT 的 patch 化几乎照搬了 ViT 的做法：ViT 把真实图像切 patch 喂给 Transformer 做分类，DiT 把潜空间噪声切 patch 喂给 Transformer 做去噪。区别在于：

ViT 通常有一个 class token 做分类输出；DiT 每个 token 都要还原成对应 patch 的噪声预测。
ViT 处理的是干净图像；DiT 处理的是加噪的潜空间特征，且每一步去噪都要知道当前时间步 t 和条件信息。
DiT 最后用一个线性层把 token 序列重新”拼回”与输入同形状的特征图，输出预测的噪声（或速度场），交给扩散采样器。

条件注入：AdaLN-Zero

扩散模型是条件去噪过程——每一步都需要知道当前时间步 t（噪声水平）以及类别标签或文本条件。DiT 没有采用把条件 token 拼接进序列的简单做法，而是设计了一种叫 AdaLN-Zero（Adaptive Layer Normalization with Zero Initialization） 的机制，将条件注入到每一个 Transformer 块里：

先把时间步 t（一个标量）通过一个多层感知机（MLP）编码成 embedding。
把类别条件（或文本 cross-attention 的结果）同样编码，与时间步 embedding 拼接或相加，得到一个总的条件向量 c。
在每个 Transformer 块里，用条件向量 c 通过一个线性层回归出六组调制参数：scale 和 shift 各用于 attention 前后的层归一化、以及一个 gate。
对层归一化后的特征做逐通道仿射变换：归一化输出乘 scale 再加 shift，即 $h = \text{scale} \cdot \text{LN}(h) + \text{shift}$ 。
关键的”Zero”：这个 gate 参数在初始化时被设为全零，使得每个 Transformer 块在训练初期相当于一个恒等映射（输入直接输出），从而极大稳定训练。

AdaLN-Zero 让每一层都能根据”当前是第几个去噪步、当前要满足什么条件”来动态调整自己的行为，这是 DiT 相比普通 ViT 在扩散任务上能稳定收敛的重要原因。

可扩展性：FID 随 FLOPs 平滑下降

DiT 论文最核心的实证贡献，是展示了 模型越大、计算量越大，生成质量越好，而且关系非常平滑。作者用 GFLOPs（每秒十亿次浮点运算）衡量单次前向计算量，定义了从 DiT-S（Small）、DiT-B（Base）、DiT-L（Large）到 DiT-XL（Extra Large）一系列模型，参数量从约 33M 到约 675M。

实验在 ImageNet 类别条件生成上测量 FID（Fréchet Inception Distance，越低越好）。结论是：

固定模型大小，增大 patch 数（即减小 patch 尺寸 p）会降低 FID。
固定 patch 配置，从 DiT-S 放大到 DiT-XL，FID 持续下降。
把所有配置点画在”FLOPs vs FID”的图上，呈现出平滑的幂律下降趋势——这正是可扩展性（scaling behavior）的标志，与大语言模型上观察到的 scaling law 异曲同工。

这一发现意义重大：它说明扩散模型的去噪网络不必依赖卷积的归纳偏置，纯 Transformer 架构只要持续堆参数和算力，就能持续变好。这直接催生了后来一系列基于 DiT 的大规模生成模型。

扩散目标函数

DiT 本身只是去噪骨干，它外围仍遵循标准扩散模型的训练目标。给定干净数据 x_0，前向加噪过程可近似为：

x_t = \sqrt{\bar{\alpha}_t} \cdot x_0 + \sqrt{1 - \bar{\alpha}_t} \cdot \text{noise}

其中 $\bar{\alpha}_t$ 是累积噪声调度系数。网络以 $x_t$ 和时间步 $t$ 为输入，预测所加的噪声 $\text{noise}$ ，损失为预测噪声与真实噪声的均方误差。采样时用各类采样器逐步去噪，最终通过 VAE 解码器还原成像素图像。关于条件引导的细节，可参考 CFG 引导。

原理图解

动手实践

下面用一个极简的 DiT 块示意如何把潜空间特征图 patch 化、注入时间步条件、并做一次前向：

import torch
import torch.nn as nn

class TinyDiTBlock(nn.Module):  # 极简 DiT 块示意
    def __init__(self, dim, num_heads):
        super().__init__()
        self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True)
        self.mlp = nn.Sequential(nn.Linear(dim, dim*4), nn.GELU(), nn.Linear(dim*4, dim))
        # AdaLN-Zero: 由条件向量回归出 scale/shift/gate
        self.ada = nn.Linear(dim, dim*6)
        nn.init.zeros_(self.ada.weight)  # Zero 初始化, 训练初期为恒等映射
        nn.init.zeros_(self.ada.bias)

    def forward(self, x, c):
        s1, b1, g1, s2, b2, g2 = self.ada(c).chunk(6, dim=-1)  # 六组调制参数
        h = (1 + s1.unsqueeze(1)) * x + b1.unsqueeze(1)  # 注意力前的 AdaLN
        a, _ = self.attn(h, h, h); h = x + g1.unsqueeze(1) * a  # 门控残差
        h2 = (1 + s2.unsqueeze(1)) * h + b2.unsqueeze(1)  # MLP 前的 AdaLN
        return h + g2.unsqueeze(1) * self.mlp(h2)  # 门控残差

# 模拟: 32x32 潜空间, patch=2 -> 256 个 token
z = torch.randn(2, 32*32, 256)          # (batch, token 数, 维度)
c = torch.randn(2, 256)                 # 时间步+类别 embedding
blk = TinyDiTBlock(256, num_heads=8)
out = blk(z, c)
print(out.shape)  # → torch.Size([2, 1024, 256]) 与输入同形

实践要点

patch 大小是关键超参：p=2 通常在质量和速度之间取得较好平衡；p=1 质量更高但序列长 4 倍、显存暴涨。
AdaLN-Zero 几乎是标配：若把条件改成简单的 cross-attention 拼接，训练初期容易不稳；零初始化的门控让深层 DiT 也能收敛。
潜空间扩散是主流：直接在像素空间做 DiT 计算量过大，绝大多数实现（Sora、PixArt、Stable Diffusion 3）都在 VAE 压缩后的潜空间操作。
Scaling law 指导资源分配：DiT 的平滑可扩展性意味着，与其精心设计网络结构，不如把预算花在加大模型和提升数据质量上。
文本条件常用 cross-attention：类别条件适合 AdaLN 注入；但开放域文本提示通常用额外的 cross-attention 层与文本 token 交互，再与时间步 embedding 一起喂给 AdaLN。
采样成本主要在去噪网络：DiT 的前向比同等 U-Net 略重，配合高效采样器（如 DPM-Solver）能显著减少采样步数。
训练分辨率与采样分辨率可分离：借助 patch 化和位置编码的灵活性，DiT 较容易迁移到不同分辨率或长宽比（这是 Sora 能生成任意比例视频的基础）。

典型应用

Sora（OpenAI）：文生视频模型，基于 DiT 架构，在时空 patch 上做扩散去噪，能生成长达一分钟、物理一致性较强的视频片段。
PixArt-α / PixArt-Σ（华为、蚂蚁等）：高效的文生图 DiT 系列，通过分解训练阶段（先类别条件、再文本条件）大幅降低训练成本，画质逼近甚至超越更大规模的模型。
Stable Diffusion 3（Stability AI）：第三代 Stable Diffusion 将骨干从 U-Net 切换为 MMDiT（Multimodal DiT），文本与图像 token 在统一 Transformer 中交互。
SiT（Scalable Interpolant Transformers）：在 DiT 基础上改用可插值（interpolant）的扩散框架，探索更灵活的前向/反向过程，进一步提升可扩展性。
Flux（Black Forest Labs）：开源的高质量文生图 DiT 模型，社区广泛使用，常作为 ComfyUI 工作流的骨干。

典型类库与工具

类库	语言	说明
DiT 官方仓库	Python	Peebles & Xie 的原版实现，含 DiT-S/B/L/XL 及 ImageNet 预训练权重
diffusers	Python	Hugging Face 出品，内置多个 DiT 变体（PixArt、Stable Diffusion 3、Flux 等），推理与微调一体
xDiT	Python	针对 DiT 类模型的多 GPU 并行推理引擎，大幅加速大模型生成
PixArt 仓库	Python	华为/蚂蚁的 PixArt-α/Σ 官方实现，含高效训练策略代码
ComfyUI	Python	节点式生成工作流，原生支持 Flux、SD3 等 DiT 系模型

术语表

术语	英文	解释
扩散 Transformer	Diffusion Transformer (DiT)	以 Transformer 作为去噪骨干的扩散模型架构
patch 化	Patchify / Patch Embedding	将特征图切成不重叠的小块并线性投影成 token
自适应层归一化	Adaptive Layer Normalization (AdaLN)	用条件向量回归出 scale/shift 参数来调制层归一化输出
零初始化	Zero Initialization	AdaLN 的门控参数初始化为零，使初始状态为恒等映射
可扩展性	Scalability	模型随参数量/计算量增大而性能持续提升的性质
标度律	Scaling Law	性能与计算量之间的幂律关系
潜空间	Latent Space	VAE 压缩后的低维表示空间，扩散在此进行
位置编码	Position Embedding	给每个 token 附加空间位置信息的向量
类别条件生成	Class-Conditional Generation	以类别标签为条件的图像生成
多模态 DiT	Multimodal DiT (MMDiT)	文本与图像 token 在统一 Transformer 中交互的 DiT 变体