2022/12/22 14:27

统治扩散模型的U-Net要被取代了，谢赛宁等引入Transformer提出DiT

来自 UC 伯克利的 William Peebles 以及纽约大学的谢赛宁撰文揭秘扩散模型中架构选择的意义，并为未来的生成模型研究提供经验基线。

近几年，在 Transformer 的推动下，机器学习正在经历复兴。过去五年中，用于自然语言处理、计算机视觉以及其他领域的神经架构在很大程度上已被 transformer 所占据。

不过还有许多图像级生成模型仍然不受这一趋势的影响，例如过去一年扩散模型在图像生成方面取得了惊人的成果，几乎所有这些模型都使用卷积 U-Net 作为主干。这有点令人惊讶！在过去的几年中，深度学习的大事件一直是跨领域的 Transformer 的主导地位。U-Net 或卷积是否有什么特别之处使它们在扩散模型中表现得如此出色？

将 U-Net 主干网络首次引入扩散模型的研究可追溯到 Ho 等人，这种设计模式继承了自回归生成模型 PixelCNN++，只是稍微进行了一些改动。而 PixelCNN++ 由卷积层组成，其包含许多的 ResNet 块。其与标准的 U-Net 相比，PixelCNN++ 附加的空间自注意力块成为 transformer 中的基本组件。不同于其他人的研究，Dhariwal 和 Nichol 等人消除了 U-Net 的几种架构选择，例如使用自适应归一化层为卷积层注入条件信息和通道计数。

本文中来自 UC 伯克利的 William Peebles 以及纽约大学的谢赛宁撰文《 Scalable Diffusion Models with Transformers 》，目标是揭开扩散模型中架构选择的意义，并为未来的生成模型研究提供经验基线。该研究表明，U-Net 归纳偏置对扩散模型的性能不是至关重要的，并且可以很容易地用标准设计（如 transformer）取代。

这一发现表明，扩散模型可以从架构统一趋势中受益，例如，扩散模型可以继承其他领域的最佳实践和训练方法，保留这些模型的可扩展性、鲁棒性和效率等有利特性。标准化架构也将为跨领域研究开辟新的可能性。

论文地址：https://arxiv.org/pdf/2212.09748.pdf
项目地址：https://github.com/facebookresearch/DiT
论文主页：https://www.wpeebles.com/DiT

该研究专注于一类新的基于 Transformer 的扩散模型：Diffusion Transformers（简称 DiTs）。DiTs 遵循 Vision Transformers (ViTs) 的最佳实践，有一些小但重要的调整。DiT 已被证明比传统的卷积网络（例如 ResNet ）具有更有效地扩展性。

具体而言，本文研究了 Transformer 在网络复杂度与样本质量方面的扩展行为。研究表明，通过在潜在扩散模型 (LDM) 框架下构建 DiT 设计空间并对其进行基准测试，其中扩散模型在 VAE 的潜在空间内进行训练，可以成功地用 transformer 替换 U-Net 主干。本文进一步表明 DiT 是扩散模型的可扩展架构：网络复杂性（由 Gflops 测量）与样本质量（由 FID 测量）之间存在很强的相关性。通过简单地扩展 DiT 并训练具有高容量主干（118.6 Gflops）的 LDM，可以在类条件 256 × 256 ImageNet 生成基准上实现 2.27 FID 的最新结果。

Diffusion Transformers

DiTs 是一种用于扩散模型的新架构，目标是尽可能忠实于标准 transformer 架构，以保留其可扩展性。DiT 保留了 ViT 的许多最佳实践，图 3 显示了完整 DiT 体系架构。

DiT 的输入为空间表示 z（对于 256 × 256 × 3 图像，z 的形状为 32 × 32 × 4）。DiT 的第一层是 patchify，该层通过将每个 patch 线性嵌入到输入中，以此将空间输入转换为一个 T token 序列。patchify 之后，本文将标准的基于 ViT 频率的位置嵌入应用于所有输入 token。

patchify 创建的 token T 的数量由 patch 大小超参数 p 决定。如图 4 所示，将 p 减半将使 T 翻四倍，因此至少能使 transformer Gflops 翻四倍。本文将 p = 2,4,8 添加到 DiT 设计空间。

DiT 块设计：在 patchify 之后，输入 token 由一系列 transformer 块处理。除了噪声图像输入之外，扩散模型有时还会处理额外的条件信息，例如噪声时间步长 t、类标签 c、自然语言等。本文探索了四种以不同方式处理条件输入的 transformer 块变体。这些设计对标准 ViT 块设计进行了微小但重要的修改。所有模块的设计如图 3 所示。

本文尝试了四种因模型深度和宽度而异的配置：DiT-S、DiT-B、DiT-L 和 DiT-XL。这些模型配置范围从 33M 到 675M 参数，Gflops 从 0.4 到 119 。

实验

研究者训练了四个最高 Gflop 的 DiT-XL/2 模型，每个模型使用不同的 block 设计 ——in-context（119.4Gflops）、cross-attention（137.6Gflops）、adaptive layer norm（adaLN，118.6Gflops）或 adaLN-zero（118.6Gflops）。然后在训练过程中测量 FID，图 5 为结果。

扩展模型大小和 patch 大小。图 2（左）给出了每个模型的 Gflops 和它们在 400K 训练迭代时的 FID 概况。可以发现，增加模型大小和减少 patch 大小会对扩散模型产生相当大的改进。

图 6（顶部）展示了 FID 是如何随着模型大小的增加和 patch 大小保持不变而变化的。在四种设置中，通过使 Transformer 更深、更宽，训练的所有阶段都获得了 FID 的明显提升。同样，图 6（底部）展示了 patch 大小减少和模型大小保持不变时的 FID。研究者再次观察到，在整个训练过程中，通过简单地扩大 DiT 处理的 token 数量，并保持参数的大致固定，FID 会得到相当大的改善。

图 8 中展示了 FID-50K 在 400K 训练步数下与模型 Gflops 的对比：

SOTA 扩散模型 256×256 ImageNet。在对扩展分析之后，研究者继续训练最高 Gflop 模型 DiT-XL/2，步数为 7M。图 1 展示了该模型的样本，并与类别条件生成 SOTA 模型进行比较，表 2 中展示了结果。

当使用无分类器指导时，DiT-XL/2 优于之前所有的扩散模型，将之前由 LDM 实现的 3.60 的最佳 FID-50K 降至 2.27。如图 2（右）所示，相对于 LDM-4（103.6 Gflops）这样的潜在空间 U-Net 模型来说，DiT-XL/2（118.6 Gflops）计算效率高得多，也比 ADM（1120 Gflops）或 ADM-U（742 Gflops）这样的像素空间 U-Net 模型效率高很多。

表 3 展示了与 SOTA 方法的比较。XL/2 在这一分辨率下再次胜过之前的所有扩散模型，将 ADM 之前取得的 3.85 的最佳 FID 提高到 3.04。

更多研究细节，可参考原论文。

理论扩散模型

相关数据

深度学习技术

深度学习（deep learning）是机器学习的分支，是一种试图使用包含复杂结构或由多重非线性变换构成的多个处理层对数据进行高层抽象的算法。深度学习是机器学习中一种基于对数据进行表征学习的算法，至今已有数种深度学习框架，如卷积神经网络和深度置信网络和递归神经网络等已被应用在计算机视觉、语音识别、自然语言处理、音频识别与生物信息学等领域并获取了极好的效果。

来源：LeCun, Y., Bengio, Y., & Hinton, G. (2015). Deep learning. nature, 521(7553), 436.

机器学习技术

机器学习是人工智能的一个分支，是一门多领域交叉学科，涉及概率论、统计学、逼近论、凸分析、计算复杂性理论等多门学科。机器学习理论主要是设计和分析一些让计算机可以自动“学习”的算法。因为学习算法中涉及了大量的统计学理论，机器学习与推断统计学联系尤为密切，也被称为统计学习理论。算法设计方面，机器学习理论关注可以实现的，行之有效的学习算法。

来源：Mitchell, T. (1997). Machine Learning. McGraw Hill.

自注意力技术

自注意力（Self-attention），有时也称为内部注意力，它是一种涉及单序列不同位置的注意力机制，并能计算序列的表征。自注意力在多种任务中都有非常成功的应用，例如阅读理解、摘要概括、文字蕴含和语句表征等。自注意力这种在序列内部执行 Attention 的方法可以视为搜索序列内部的隐藏关系，这种内部关系对于翻译以及序列任务的性能非常重要。

基准技术

一种简单的模型或启发法，用作比较模型效果时的参考点。基准有助于模型开发者针对特定问题量化最低预期效果。

来源：Google ML Glossary

参数技术

在数学和统计学裡，参数（英语：parameter）是使用通用变量来建立函数和变量之间关系（当这种关系很难用方程来阐述时）的一个数量。

来源：维基百科

超参数技术

在机器学习中，超参数是在学习过程开始之前设置其值的参数。相反，其他参数的值是通过训练得出的。不同的模型训练算法需要不同的超参数，一些简单的算法（如普通最小二乘回归）不需要。给定这些超参数，训练算法从数据中学习参数。相同种类的机器学习模型可能需要不同的超参数来适应不同的数据模式，并且必须对其进行调整以便模型能够最优地解决机器学习问题。在实际应用中一般需要对超参数进行优化，以找到一个超参数元组（tuple），由这些超参数元组形成一个最优化模型，该模型可以将在给定的独立数据上预定义的损失函数最小化。

来源：Wikipedia

计算机视觉技术

计算机视觉（CV）是指机器感知环境的能力。这一技术类别中的经典任务有图像形成、图像处理、图像提取和图像的三维推理。目标识别和面部识别也是很重要的研究领域。

来源：机器之心

图像生成技术

图像生成（合成）是从现有数据集生成新图像的任务。

来源：paperswithcode

自然语言处理技术

自然语言处理（英语：natural language processing，缩写作 NLP）是人工智能和语言学领域的分支学科。此领域探讨如何处理及运用自然语言；自然语言认知则是指让电脑“懂”人类的语言。自然语言生成系统把计算机数据转化为自然语言。自然语言理解系统把自然语言转化为计算机程序更易于处理的形式。

来源：维基百科

生成模型技术

在概率统计理论中, 生成模型是指能够随机生成观测数据的模型，尤其是在给定某些隐含参数的条件下。它给观测值和标注数据序列指定一个联合概率分布。在机器学习中，生成模型可以用来直接对数据建模（例如根据某个变量的概率密度函数进行数据采样），也可以用来建立变量间的条件概率分布。

来源：维基百科