该研究专注于一类新的基于Transformer的扩散模型:Diffusion Transformers(简称 DiTs)。DiTs遵循Vision Transformers (ViTs)的最佳实践,有一些小但重要的调整。DiT已被证明比传统的卷积网络(例如ResNet)具有更有效地扩展性。
本文研究了Transformer 在网络复杂度与样本质量方面的扩展行为。研究表明,通过在潜在扩散模型 (LDM) 框架下构建 DiT设计空间并对其进行基准测试,其中扩散模型在VAE的潜在空间内进行训练,可以成功地用transformer替换 U-Net主干。
本文进一步表明DiT是扩散模型的可扩展架构:网络复杂性(由 Gflops 测量)与样本质量(由FID测量)之间存在很强的相关性。通过简单地扩展DiT并训练具有高容量主干(118.6 Gflops)的LDM,可以在类条件256×256 ImageNet生成基准上实现2.27 FID的最新结果。
1、Diffusion Transformers
DiTs 是一种用于扩散模型的新架构,目标是尽可能忠实于标准 transformer 架构,以保留其可扩展性。DiT 保留了 ViT 的许多最佳实践,图 3 显示了完整 DiT 体系架构。
DiT 的输入为空间表示 z(对于 256 × 256 × 3 图像,z 的形状为 32 × 32 × 4)。DiT 的第一层是 patchify,该层通过将每个 patch 线性嵌入到输入中,以此将空间输入转换为一个 T token 序列。patchify 之后,本文将标准的基于 ViT 频率的位置嵌入应用于所有输入 token。
patchify 创建的 token T 的数量由 patch 大小超参数 p 决定。如图 4 所示,将 p 减半将使 T 翻四倍,因此至少能使 transformer Gflops 翻四倍。本文将 p = 2,4,8 添加到 DiT 设计空间。
DiT 块设计:在 patchify 之后,输入 token 由一系列 transformer 块处理。除了噪声图像输入之外,扩散模型有时还会处理额外的条件信息,例如噪声时间步长 t、类标签 c、自然语言等。本文探索了四种以不同方式处理条件输入的 transformer 块变体。这些设计对标准 ViT 块设计进行了微小但重要的修改。所有模块的设计如图 3 所示。
本文尝试了四种因模型深度和宽度而异的配置:DiT-S、DiT-B、DiT-L 和 DiT-XL。这些模型配置范围从 33M 到 675M 参数,Gflops 从 0.4 到 119 。
2、实验
研究者训练了四个最高 Gflop 的 DiT-XL/2 模型,每个模型使用不同的 block 设计 ——in-context(119.4Gflops)、cross-attention(137.6Gflops)、adaptive layer norm(adaLN 118.6Gflops)或 adaLN-zero(118.6Gflops)。然后在训练过程中测量 FID,图 5 为结果。
扩展模型大小和 patch 大小。图 2(左)给出了每个模型的 Gflops 和它们在 400K 训练迭代时的 FID 概况。可以发现,增加模型大小和减少 patch 大小会对扩散模型产生相当大的改进。
图 6(顶部)展示了 FID 是如何随着模型大小的增加和 patch 大小保持不变而变化的。在四种设置中,通过使 Transformer 更深、更宽,训练的所有阶段都获得了 FID 的明显提升。同样,图 6(底部)展示了 patch 大小减少和模型大小保持不变时的 FID。研究者再次观察到,在整个训练过程中,通过简单地扩大 DiT 处理的 token 数量,并保持参数的大致固定,FID 会得到相当大的改善。
图 8 中展示了 FID-50K 在 400K 训练步数下与模型 Gflops 的对比:
SOTA 扩散模型 256×256 ImageNet。在对扩展分析之后,研究者继续训练最高 Gflop 模型 DiT-XL/2,步数为 7M。图 1 展示了该模型的样本,并与类别条件生成 SOTA 模型进行比较,表 2 中展示了结果。
当使用无分类器指导时,DiT-XL/2 优于之前所有的扩散模型,将之前由 LDM 实现的 3.60 的最佳 FID-50K 降至 2.27。如图 2(右)所示,相对于 LDM-4(103.6 Gflops)这样的潜在空间 U-Net 模型来说,DiT-XL/2(118.6 Gflops)计算效率高得多,也比 ADM(1120 Gflops)或 ADM-U(742 Gflops)这样的像素空间 U-Net 模型效率高很多。
表 3 展示了与 SOTA 方法的比较。XL/2 在这一分辨率下再次胜过之前的所有扩散模型,将 ADM 之前取得的 3.85 的最佳 FID 提高到 3.04。