“U-Net 已死，Transformer 成为扩散模型新 SOTA 了！”
就在 ChatGPT 占尽 AI 圈风头时，纽约大学谢赛宁的图像生成模型新论文横空出世，收获一众同行惊讶的声音。
论文创意性地将 Transformer 与扩散模型融合，在计算效率和生成效果上均超越了基于 U-Net 的经典模型 ADM 和 LDM，打破了 U-Net 统治扩散模型的“普遍认知”。
网友给这对新组合命名也是脑洞大开：
All we need is U-Transformer
希望他们没有错过 Transffusion 这个名字。
要知道，这几年虽然 Transformer 占尽风头，但 U-Net 在扩散模型领域仍然一枝独秀 ——
无论是“前任王者”DALL・E2 还是“新晋生成 AI”Stable Diffusion，都没有使用 Transformer 作为图像生成架构。
如今新研究表明，U-Net 并非不可用 Transformer 替代。
“U-Net 并非不可替代”
论文提出的新架构名叫 Diffusion Transformers（DiTs）。
架构保留了很多 ViT 的特性，其中整体架构如图左（包含多个 DiT 模块），具体的 DiT 模块组成如图右：
更右边的两个灰色框的模块，则是 DiT 架构的“变体”。主要是探讨在条件输入下，不同的架构是否能对信息进行更好的处理，包括交叉注意力等。
最终结果表明，还是层归一化（Layer Normalization）更好用，这里最终选用了 Adaptive Layer Normalization（自适应层归一化）的方法。
对于这篇论文研究的目的，作者表示希望探讨扩散模型中不同架构选择的重要性，以及也是给将来生成模型的评估做一个评判标准。
先说结果 —— 作者认为，U-Net 的归纳偏置（inductive bias），对于扩散模型性能提升不是必须的。
与之相反，他们能“轻松地”（readily）被 Transformer 的标准架构取代。
有网友发现，DALL・E 和 DALL・E2 似乎都有用到 Transformer。
这篇论文和它们的差异究竟在哪里？
事实上，DALL・E 虽然是 Transformer，但并非扩散模型，本质是基于 VQVAE 架构实现的；
至于 DALL・E2 和 Stable Diffusion，虽然都分别将 Transformer 用在了 CLIP 和文本编码器上，但关键的图像生成用的还是 U-Net。
不过，DiT 还不是一个文本生成图像模型 —— 目前只能基于训练标签生成对应的新图像。
虽然生成的图片还带着股“ImageNet 风”，不过英伟达 AI 科学家 Jim Fan 认为，将它改造成想要的风格和加上文本生成功能，都不是难点。
如果将标签输入调整成其他向量、乃至于文本嵌入，就能很快地将 DiT 改造成一个文生图模型：
Stable-DiT 马上就要来了！
所以 DiTs 在生成效果和运算速率上，相比其他图像生成模型究竟如何？
在 ImageNet 基准上取得 SOTA
为了验证 DiTs 的最终效果，研究者将 DiTs 沿“模型大小”和“输入标记数量”两个轴进行了缩放。
具体来说，他们尝试了四种不同模型深度和宽度的配置：DiT-S、DiT-B、DiT-L 和 DiT-XL，在此基础上又分别训练了 3 个潜块大小为 8、4 和 2 的模型，总共是 12 个模型。
从 FID 测量结果可以看出，就像其他领域一样，增加模型大小和减少输入标记数量可以大大提高 DiT 的性能。
FID 是计算真实图像和生成图像的特征向量之间距离的一种度量，越小越好。
换句话说，较大的 DiTs 模型相对于较小的模型是计算效率高的，而且较大的模型比较小的模型需要更少的训练计算来达到给定的 FID。
其中，Gflop 最高的模型是 DiT-XL / 2，它使用最大的 XL 配置，patch 大小为 2，当训练时间足够长时，DiT-XL / 2 就是里面的最佳模型。
于是在接下来，研究人员就专注于 DiT-XL / 2，他们在 ImageNet 上训练了两个版本的 DiT-XL / 2，分辨率分别为 256×256 和 512×512，步骤分别为 7M 和 3M。
当使用无分类器指导时，DiT-XL / 2 比之前的扩散模型数据都要更好，取得 SOTA 效果：
在 256×256 分辨率下，DiT-XL / 2 将之前由 LDM 实现的最佳 FID-50K 从 3.60 降至了 2.27。
并且与基线相比，DiTs 模型本身的计算效率也很高：
DiT-XL / 2 的计算效率为 119 Gflops，相比而言 LDM-4 是 103 Gflops，ADM-U 则是 742 Gflops。
同样，在 512×512 分辨率下，DiT-XL / 2 也将 ADM-U 之前获得的最佳 FID 3.85 降至了 3.04。
不过此时 ADM-U 的计算效率是 2813 Gflops，而 XL / 2 只有 525 Gflops。
研究作者
本篇论文作者为 UC 伯克利的 William Peebles 和纽约大学的谢赛宁。
William Peebles，目前是 UC 伯克利的四年级博士生，本科毕业于麻省理工学院。研究方向是深度学习和人工智能，重点是深度生成模型。
之前曾在 Meta、Adobe、英伟达实习过，这篇论文就是在 Meta 实习期间完成。
谢赛宁，纽约大学计算机科学系助理教授，之前曾是 Meta FAIR 研究员，本科就读于上海交通大学 ACM 班，博士毕业于 UC 圣迭戈分校。
谢赛宁读博士时曾在 FAIR 实习，期间与何恺明合作完成 ResNeXt，是该论文的一作，之前何恺明一作论文 MAE 他也有参与。
当然，对于这次 Transformer 的表现，也有研究者们表示“U-Net 不服”。
例如三星 AI Lab 科学家 Alexia Jolicoeur-Martineau 就表示：
U-Net 仍然充满生机，我相信只需要经过细小调整，有人能将它做得比 Transformer 更好。
看来，图像生成领域很快又要掀起新的“较量风暴”了。
论文地址：
https://arxiv.org/abs/2212.09748v1
参考链接：
-
[1]https://twitter.com/ethanCaballero/status/1605621603135471616
-
[2]https://www.wpeebles.com/DiT
-
[3]https://paperswithcode.com/paper/scalable-diffusion-models-with-transformers#code
本文来自微信公众号：量子位 （ID：QbitAI），作者：羿阁 萧箫