宝塔服务器面板，一键全能部署及管理，送你10850元礼包，点我领取

“U-Net 已死，Transformer 成为扩散模型新 SOTA 了！”

就在 ChatGPT 占尽 AI 圈风头时，纽约大学谢赛宁的图像生成模型新论文横空出世，收获一众同行惊讶的声音。

扩散模型和 Transformer 梦幻联动，一举拿下新 SOTA，MILA 博士：U-Net 已死-风君子博客

△MILA 在读 ML 博士生 Ethan Caballero

论文创意性地将 Transformer 与扩散模型融合，在计算效率和生成效果上均超越了基于 U-Net 的经典模型 ADM 和 LDM，打破了 U-Net 统治扩散模型的“普遍认知”。

扩散模型和 Transformer 梦幻联动，一举拿下新 SOTA，MILA 博士：U-Net 已死-风君子博客

网友给这对新组合命名也是脑洞大开：

All we need is U-Transformer

希望他们没有错过 Transffusion 这个名字。

扩散模型和 Transformer 梦幻联动，一举拿下新 SOTA，MILA 博士：U-Net 已死-风君子博客

要知道，这几年虽然 Transformer 占尽风头，但 U-Net 在扩散模型领域仍然一枝独秀 ——

无论是“前任王者”DALL・E2 还是“新晋生成 AI”Stable Diffusion，都没有使用 Transformer 作为图像生成架构。

扩散模型和 Transformer 梦幻联动，一举拿下新 SOTA，MILA 博士：U-Net 已死-风君子博客

△ 英伟达 AI 科学家 Jim Fan

如今新研究表明，U-Net 并非不可用 Transformer 替代。

“U-Net 并非不可替代”

论文提出的新架构名叫 Diffusion Transformers（DiTs）。

架构保留了很多 ViT 的特性，其中整体架构如图左（包含多个 DiT 模块），具体的 DiT 模块组成如图右：

扩散模型和 Transformer 梦幻联动，一举拿下新 SOTA，MILA 博士：U-Net 已死-风君子博客

更右边的两个灰色框的模块，则是 DiT 架构的“变体”。主要是探讨在条件输入下，不同的架构是否能对信息进行更好的处理，包括交叉注意力等。

最终结果表明，还是层归一化（Layer Normalization）更好用，这里最终选用了 Adaptive Layer Normalization（自适应层归一化）的方法。

对于这篇论文研究的目的，作者表示希望探讨扩散模型中不同架构选择的重要性，以及也是给将来生成模型的评估做一个评判标准。

先说结果 —— 作者认为，U-Net 的归纳偏置（inductive bias），对于扩散模型性能提升不是必须的。

与之相反，他们能“轻松地”（readily）被 Transformer 的标准架构取代。

扩散模型和 Transformer 梦幻联动，一举拿下新 SOTA，MILA 博士：U-Net 已死-风君子博客

有网友发现，DALL・E 和 DALL・E2 似乎都有用到 Transformer。

这篇论文和它们的差异究竟在哪里？

事实上，DALL・E 虽然是 Transformer，但并非扩散模型，本质是基于 VQVAE 架构实现的；

扩散模型和 Transformer 梦幻联动，一举拿下新 SOTA，MILA 博士：U-Net 已死-风君子博客

至于 DALL・E2 和 Stable Diffusion，虽然都分别将 Transformer 用在了 CLIP 和文本编码器上，但关键的图像生成用的还是 U-Net。

扩散模型和 Transformer 梦幻联动，一举拿下新 SOTA，MILA 博士：U-Net 已死-风君子博客

△ 经典 U-Net 架构

不过，DiT 还不是一个文本生成图像模型 —— 目前只能基于训练标签生成对应的新图像。

虽然生成的图片还带着股“ImageNet 风”，不过英伟达 AI 科学家 Jim Fan 认为，将它改造成想要的风格和加上文本生成功能，都不是难点。

如果将标签输入调整成其他向量、乃至于文本嵌入，就能很快地将 DiT 改造成一个文生图模型：

Stable-DiT 马上就要来了！

扩散模型和 Transformer 梦幻联动，一举拿下新 SOTA，MILA 博士：U-Net 已死-风君子博客

所以 DiTs 在生成效果和运算速率上，相比其他图像生成模型究竟如何？

在 ImageNet 基准上取得 SOTA

为了验证 DiTs 的最终效果，研究者将 DiTs 沿“模型大小”和“输入标记数量”两个轴进行了缩放。

具体来说，他们尝试了四种不同模型深度和宽度的配置：DiT-S、DiT-B、DiT-L 和 DiT-XL，在此基础上又分别训练了 3 个潜块大小为 8、4 和 2 的模型，总共是 12 个模型。

扩散模型和 Transformer 梦幻联动，一举拿下新 SOTA，MILA 博士：U-Net 已死-风君子博客

从 FID 测量结果可以看出，就像其他领域一样，增加模型大小和减少输入标记数量可以大大提高 DiT 的性能。

FID 是计算真实图像和生成图像的特征向量之间距离的一种度量，越小越好。

换句话说，较大的 DiTs 模型相对于较小的模型是计算效率高的，而且较大的模型比较小的模型需要更少的训练计算来达到给定的 FID。

其中，Gflop 最高的模型是 DiT-XL / 2，它使用最大的 XL 配置，patch 大小为 2，当训练时间足够长时，DiT-XL / 2 就是里面的最佳模型。

扩散模型和 Transformer 梦幻联动，一举拿下新 SOTA，MILA 博士：U-Net 已死-风君子博客

于是在接下来，研究人员就专注于 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。

扩散模型和 Transformer 梦幻联动，一举拿下新 SOTA，MILA 博士：U-Net 已死-风君子博客

同样，在 512×512 分辨率下，DiT-XL / 2 也将 ADM-U 之前获得的最佳 FID 3.85 降至了 3.04。

不过此时 ADM-U 的计算效率是 2813 Gflops，而 XL / 2 只有 525 Gflops。

扩散模型和 Transformer 梦幻联动，一举拿下新 SOTA，MILA 博士：U-Net 已死-风君子博客

研究作者

本篇论文作者为 UC 伯克利的 William Peebles 和纽约大学的谢赛宁。

扩散模型和 Transformer 梦幻联动，一举拿下新 SOTA，MILA 博士：U-Net 已死-风君子博客

William Peebles，目前是 UC 伯克利的四年级博士生，本科毕业于麻省理工学院。研究方向是深度学习和人工智能，重点是深度生成模型。

扩散模型和 Transformer 梦幻联动，一举拿下新 SOTA，MILA 博士：U-Net 已死-风君子博客

之前曾在 Meta、Adobe、英伟达实习过，这篇论文就是在 Meta 实习期间完成。

谢赛宁，纽约大学计算机科学系助理教授，之前曾是 Meta FAIR 研究员，本科就读于上海交通大学 ACM 班，博士毕业于 UC 圣迭戈分校。

谢赛宁读博士时曾在 FAIR 实习，期间与何恺明合作完成 ResNeXt，是该论文的一作，之前何恺明一作论文 MAE 他也有参与。

扩散模型和 Transformer 梦幻联动，一举拿下新 SOTA，MILA 博士：U-Net 已死-风君子博客

当然，对于这次 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），作者：羿阁 萧箫