DiT：Transformer 架构如何重塑扩散模型

图1：DiT-XL/2 在 256×256 ImageNet 上生成的样本，FID 达到 2.27

一个困扰研究者多年的问题

2022 年底之前，扩散模型（Diffusion Models）的天下属于 U-Net。

从 DDPM 到 Stable Diffusion，几乎所有主流扩散模型都使用 U-Net 作为噪声预测的核心网络。这种卷积神经网络结构成熟、训练稳定，但当研究人员试图扩大模型规模时，却撞上了一堵隐形的墙——参数量增大到一定程度后，模型性能的提升开始趋于平缓，效率与表达能力之间的权衡变得越来越棘手。

就在这时，两位来自 UC 伯克利和纽约大学的研究者 William Peebles 与 Saining Xie 给出了一个反直觉的答案：把 U-Net 换成 Transformer。

他们的论文《Scalable Diffusion Models with Transformers》于 2022 年 12 月发布，很快成为扩散模型领域的里程碑。DiT 证明了一个关键结论——当用 Transformer 替代 U-Net，并让它以"图像块"（patch）为基本单位处理潜在空间中的数据时，模型的扩展特性变得异常优雅：更大的模型、更长的训练步数，几乎总能带来稳定的 FID 下降。这正是 U-Net 难以实现的"扩展法则"（Scaling Law）。

架构解析：DiT 做了什么不同

DiT 的核心创新在于将标准的图像扩散流程与 Transformer 架构深度融合。这一融合并非简单的模块替换，而是从输入表示到网络结构的系统性重新设计。

潜在扩散范式：DiT 并非直接在像素空间操作图像，而是沿用了 Latent Diffusion Model 的范式——先将图像通过 VAE（变分自编码器）压缩到低维潜在空间，然后在潜在空间中进行扩散和去噪。这一设计显著降低了计算开销，使得 Transformer 能够在可接受的计算成本下处理高分辨率图像。

Patchify 机制：与 ViT（Vision Transformer）将图像切分为固定大小的块类似，DiT 将 VAE 编码后的潜在表示进一步切分为小块（patch）。例如，256×256 的图像经 VAE 8 倍下采样后变成 32×32 的潜在表示，再按 2×2 或其他尺寸划分为 token 序列，送入 Transformer 处理。这种表示方式让 Transformer 的自注意力机制能够高效建模全局图像结构。

条件注入：DiT 巧妙地将扩散过程中的时间步（timestep）和类别标签（class label）通过自适应层归一化（Adaptive Layer Norm, AdaLN）机制注入 Transformer 块。不同于在 token 序列中拼接条件向量，AdaLN 直接调制每个层的特征统计量，引入的额外参数量极少，却能实现与条件向量拼接相当甚至更好的条件控制效果。

模型规模层次：DiT 提供了从 DiT-S 到 DiT-XL 多个规模的模型变体，核心区别在于 Transformer 的层数和隐藏维度宽度。DiT-XL/2（XL 规模，patch 大小为 2）是论文中的旗舰模型，在 ImageNet 256×256 上达到了 2.27 的 FID 分数，刷新了该基准的最优记录。

图2：DiT 在不同类别条件下的生成样本展示

技术栈与代码质量

DiT 的代码库体现了顶级学术研究的工程水准。

PyTorch 原生实现：项目完全基于 PyTorch 实现，不依赖任何特殊的自定义算子。所有模型定义集中在 models.py 中，代码结构清晰，注释详尽。扩散过程的核心逻辑位于 diffusion/ 目录，包含标准的高斯扩散实现和重采样策略。

模块化架构：models.py 中完整实现了 DiT 的各层组件——TimestepEmbedder（时间步嵌入）、LabelEmbedder（类别嵌入）、DiT_Block（核心 Transformer 块）和 FinalLayer（输出层）。各模块职责单一，可独立测试和复用。train.py 使用 PyTorch DDP 实现多卡分布式训练。

依赖精简：核心依赖仅包括 PyTorch、torchvision、timm（提供 ViT 基础组件）和 diffusers（HuggingFace 的标准化扩散模型库）。训练脚本约 11,000 字，推理脚本约 3,000 行，代码体量适中，易于阅读和修改。

开放的预训练生态：DiT 模型已被直接集成进 HuggingFace diffusers 库，这意味着任何熟悉 diffusers API 的开发者都可以用两三行代码加载 DiT-XL/2 进行推理。这种"研究代码即产品"的实践大幅降低了应用门槛。

上手体验

DiT 的部署体验对研究者友好，但有一定门槛。

推理阶段：运行预训练模型推理极为简单。克隆仓库、创建 conda 环境后，只需一条命令即可生成图像：

python sample.py --image-size 512 --seed 1

预训练权重（约 1.4GB）会自动从 Meta 的公共文件服务器下载。HuggingFace Spaces 和 Google Colab 提供了完全免部署的在线体验。

训练阶段：使用 ImageNet 训练 DiT 需要专业硬件配置。论文中使用 8 块 A100 GPU 训练 DiT-XL/2 到 40 万步。torchrun 多节点分布式启动配置略复杂但文档说明清晰。

硬件需求：DiT-XL/2 推理需要约 12GB VRAM，生成单张 512×512 图像在 A100 上约需数十秒到一分钟。训练则需要多卡 A100 环境。

局限与挑战

DiT 并非没有短板。首先是许可证限制——DiT 采用 CC-BY-NC（仅限非商业用途）许可证，限制了商业应用空间。

其次，DiT 目前仅支持 ImageNet 的 1000 类条件生成，不支持文本到图像的语义控制，与同期的商业方案存在明显差距。

最后，推理效率仍是 Transformer 扩散模型的痛点。相比经过大量优化的 U-Net，DiT 的自注意力机制在高分辨率下计算量增长显著。

行业意义

DiT 的发布对 AI 视觉领域产生了深远影响。

从技术演进看，DiT 证明了 Transformer 的扩展法则同样适用于扩散模型，直接推动了后续一系列大型 Transformer 扩散模型的研究热潮。DiT 的架构思路被广泛借鉴，HuggingFace diffusers 库中如今已有数十个基于 DiT 衍生的模型。

从工程实践看，DiT 的代码组织方式为后续研究项目树立了标杆——清晰的模块划分、完善的分布式训练支持、自动下载的预训练权重，这些实践后来成为 AI 研究开源项目的标准配置。

DiT 的故事说明了一个朴素但深刻的道理：有时，问题的答案不在现有架构的修补中，而在一个看似离经叛道的基础架构替换里。

---

项目主页：https://www.wpeebles.com/DiT | 论文：arXiv:2212.09748 | HuggingFace Space：wpeebles/DiT