基于Transformer的可扩展扩散模型

我们探索了一类基于Transformer架构的新型扩散模型。我们训练了图像的潜在扩散模型，将以往广泛使用的U-Net主干网络替换为作用于潜在块（latent patches）上的Transformer。通过前向传播复杂度（以Gflops衡量）的视角，我们分析了所提出的扩散Transformer（Diffusion Transformers, DiTs）的可扩展性。研究发现，具有更高Gflops的DiTs——无论是通过增加Transformer的深度或宽度，还是增加输入token的数量——均表现出更优的性能，FID指标持续降低。除了具备良好的可扩展性外，我们最大的DiT-XL/2模型在类别条件下的ImageNet 512×512和256×256基准测试中均超越了所有先前的扩散模型，尤其在256×256尺度上取得了当前最优的FID值2.27。