Sora爆火引热议！深扒其背后的Transformer论文!

这两天，OpenAI首款文生视频大模型Sora以黑马之姿占据AI领域话题中心，有趣的是，OpenAI工程师William Peebles此前以Sora为基础发布的Transformer论文，却曾因“缺乏新颖性”而被计算机视觉学术顶会之一拒收。我们来了解一下这篇论文吧！

显然Transformers已经广泛应用于包括NLP、CV在内的机器学习的各个领域。然而，很多图片level的生成模型还坚持使用卷积神经网络，比如扩散模型采用的就是U-Net的主干网络架构。2023年3月2日，谢赛宁和William Peebles在arXiv上发布了一篇题为“Scalable Diffusion Models with Transformers”的论文，文章提出使用Transformers替换扩散模型中U-Net主干网络，分析发现，这种Diffusion Transformers（DiTs）不仅速度更快（更高的Gflops），而且在ImageNet 512×512和256×256的类别条件图片生成任务上，取得了更好的效果，256×256上实现了SOTA的FID指标。

关键词

扩散变换器（DiTs） 生成模型 变换器架构 图像生成 SOTA。

正文内容

实验设置

• 模型训练：作者在ImageNet数据集上训练了类条件的潜在Diffusion Transformer (DiT) 模型，使用的图像分辨率为256×256和512×512。使用AdamW优化器进行训练，初始化最终的线性层为零，其余使用ViT的标准权重初始化技术。

• 训练超参数：维持恒定的学习率为1×10^-4，不使用权重衰减，批处理大小为256。唯一使用的数据增强是水平翻转。实验中没有发现需要学习率预热或正则化来训练DiT达到高性能。

DiT设计空间

• 输入处理：采用空间表示的方式（来自VAE的噪声潜在空间），将图像划分为序列化的patches，然后通过Transformer模型进行处理。

• 条件信息处理：实验了不同的Transformer块设计，包括在上下文中直接添加条件信息的tokens、交叉注意力块、自适应层规范化（adaLN）块以及adaLN-Zero块，后者在训练过程中展现了最佳性能。

性能评估




• 评估指标：使用Frechet Inception Distance (FID) 作为主要评估生成模型图像质量的标准度量。此外，还报告了Inception得分、sFID和精度/召回率作为辅助指标。

• 计算资源：所有模型均使用JAX实现，并在TPU-v3 Pods上训练。DiT-XL/2模型是最计算密集的模型，在TPU v3-256 Pod上以大约5.7次迭代/秒的速度进行训练。

实验结果

• 模型性能：在多个指标上，包括FID、Inception得分等，DiT模型表现出与传统U-Net模型相比，在图像生成任务上的优越性能。实验结果显示，随着模型尺寸的增加和patch尺寸的减小，DiT模型的性能得到了显著提高。

• 计算效率：通过对不同的DiT配置进行实验，作者发现增加模型的Gflops是提高性能的关键。尽管在保持模型大小不变的情况下减小patch大小不会改变模型的总参数数量，但会显著提高模型的计算需求，从而改善生成的图像质量。

这项研究证明了变换器架构可以有效地应用于扩散模型中，为图像生成任务提供了一种新的方法。通过将变换器（Transformer）技术与扩散模型相结合，研究开辟了一条提高生成模型性能和效率的新途径。这不仅在理论上扩展了变换器在不同机器学习领域的应用范围，也在实践中提供了一种高效生成高质量图像的新策略。未来的研究方向可能包括进一步优化变换器和扩散模型的结合方式，提升模型在更多图像生成任务上的性能，以及探索变换器在其他类型的生成模型中的应用潜力。