《基于 Transformer 的可扩展扩散模型》（DiT）

要解决的问题

扩散模型（diffusion model）生成图像的办法，是先训练一个网络学会把逐步加进去的噪声还原掉；采样时从纯噪声出发，一步一步去噪，最后得到图像。这类系统大多依赖 U-Net 主干：一种卷积网络，先把输入的空间分辨率压低，再放大回去，并用跳跃连接把浅层的细节送到深层。Transformer 的思路不同，它把输入当作一串 token，用注意力让 token 之间交换信息。这篇论文要问的是：U-Net 是不是必需的？如果换成 transformer，模型能不能获得语言模型那样的性质——规模变大、算力变多时，效果平滑而可预测地变好。^[1]

怎么做的

论文保留了扩散的训练配方和潜空间（latent space）训练方式，也就是让模型去噪的是压缩后的图像表示而不是原始像素，只把主干换掉。transformer 直接在潜变量的切块（patch）上工作：把潜表示切成网格状的小块，每块嵌入成一个 token，再交给 transformer 块依次处理。^[1] 网络其余部分，也就是含噪输入的接入与噪声预测，保持不变，这样比较才干净。

衡量规模用的是前向传播复杂度，单位 Gflops，即一次前向传播需要多少次浮点运算（以十亿次计）。作者通过三个可调的参数改变规模：transformer 的深度、宽度，以及输入 token 的数量，后者由切块的大小决定。^[1]

结果

论文给出两件事。第一，扩大规模确实有效。前向复杂度更高、也就是 Gflops 更大的扩散 transformer（Diffusion Transformer，DiT）——通过加深、加宽或增加输入 token 数量得到——FID 一致地更低。^[1] FID 即弗雷歇起始距离（Fréchet Inception Distance），用来比较生成图像与真实图像的分布差距，越低越好。这种关系不是某个特定规模上的偶然跳变，而是在作者训练的各个模型上连续成立的趋势，这正是把模型大小当作一个旋钮来调时所希望看到的。

第二，最大的模型超过此前的工作。最大的 DiT-XL/2 模型在类别条件的 ImageNet 512×512 和 256×256 基准上，表现优于此前所有扩散模型，在 256×256 上取得 2.27 的当前最好 FID。^[1]

算力代价也给得足够具体：这一系列中最耗算力的 DiT-XL/2，在全局批大小为 256 时，于 TPU v3-256 pod 上训练速度约为每秒 5.7 次迭代。^[1]

为什么重要

一个扩散模型其实捆绑了几个相对独立的决定：噪声如何调度、优化什么目标、以及用哪个网络去预测噪声。这篇论文把最后一项单独拿出来，说明一个通用、被研究得很透的序列架构可以充当主干，而且生成质量能随算力以可测量的方式改善。这样一来，「去噪网络该用哪种架构」就从设计口味问题，变成可以用 Gflops 与 FID 的对应关系来回答的问题。它也意味着扩散方向的研究可以直接借用围绕 transformer 积累起来的工具与扩规模的经验，而不必依赖图像生成专用的网络结构。论文给出的训练吞吐数字，则为其他团队估算自己的训练开销提供了一个具体参照。^[1]

^[1]: arXiv，《Scalable Diffusion Models with Transformers (DiT)》。