DiffiT | 英伟达提出用于图像生成的扩散ViT架构DiffiT，达成新SOTA！

本文首发于AIWalker，欢迎关注。

https://arxiv.org/abs/2312.02139
https://github.com/NVlabs/DiffiT

扩散模型以其强大的表达能力和高样本质量在许多领域得到了新的应用。对于样本生成，这些模型依赖于通过迭代去噪生成图像的去噪神经网络。然而，去噪网络架构的作用并没有得到很好的研究，大多数工作都依赖于卷积残差U-Nets。
本文研究了视觉transformer在基于扩散的生成学习中的有效性。本文提出一种新模型，称为扩散视觉transformer (DiffiT)，由一个具有U形编码器和解码器的混合分层架构组成。本文提出一种新的依赖时间的自注意力模块，允许注意力层以有效的方式自适应其在去噪过程的不同阶段的行为。
此外，本文还提出了LatentDiffiT，由transformer模型和所提出的自注意力层组成，用于高分辨率图像生成。结果表明，DiffiT在生成高保真图像方面惊人地有效，并在各种类条件和无条件合成任务中实现了最先进的(SOTA)基准。在潜空间中，DiffiT在ImageNet256数据集上取得了新的SOTA FID分数1.73。

本文方案

上图为本文所提DiffIT整体架构示意图，很明显，核心是所提DiffiT ResBlock，故我们仅对该核心模块进行简要介绍。

上图为DiffiT模块示意图，可以描述如下：
$${\hat{x}}_s=TMSA(LN(x_s),x_t)+x_s$$
$$x_s=MLP(LN(\widehat{(}x{)}_s))+\widehat{(}x{)}_s$$

关于TMSA，在每一层，我们的transformer块接收{xs}，一组标记在其输入中空间上排列在2D网格上。它还接收xt，这是表示时间步长的时间标记。通过将位置时间嵌入提供给具有swish激活的小型MLP来获得时间标记[19]。这次令牌被传递到我们的去噪网络中的所有层。本文提出了时间依赖的多头自注意力，通过在共享空间中投影特征和时间标记嵌入来捕捉长程空间和时间依赖性。具体来说，共享空间中的时间依赖查询q、键k和值v是
通过空间和时间嵌入的线性投影xs和xt via计算

key、query和value都是时间和空间token的线性函数，它们可以针对不同的时间步长自适应地修改注意力的行为。

图像空间

DiffIT结构 DiffiT使用对称的u型编码器-解码器架构，其中收缩和扩展路径在每个分辨率上通过跳跃连接相互连接。编码器或解码器路径的每个分辨率由L连续的DiffiT块组成，其中包含所提出的时间依赖的自注意力模块。在每条路径的开始，对于编码器和解码器，使用一个卷积层来匹配特征图的数量。此外，卷积上采样或下采样层还用于每个分辨率之间的过渡。我们推测，这些卷积层的使用嵌入了可以进一步提高性能的归纳图像偏差。在本节的其余部
分，我们讨论了DiffiT Transformer块和提出的时间依赖自注意力机制。在构建u型去噪架构时，使用所提出的Transformer块作为残差单元。

DiffiT ResBlock 通过将所提出的DiffiT Transformer块
与额外的卷积层相结合，定义最终的残差单元:

潜空间

潜扩散模型被证明可以有效地生成高质量的大分辨率图像。在图4中，我们展示了隐DiffiT模型的架构。我们首先使用预训练的变分自编码器网络对图像进行编码，然后将特征图转换为不重叠的块并投影到新的嵌入空间。与DiT模型[52]类似，我们使用视觉transformer，没有上采样或下采样层，作为潜空间中的去噪网络。此外，还利用三通道无分类器指导来提高生成样本的质量。架构的最后一层是一个简单的线性层，用于对输出进行解码。

本文实验

DiffiT在CIFAR-
10数据集上取得了最先进的图像生成FID分数1.95，
超 过 了EDM [34]和LSGM [69]等 最 先 进 的 扩 散 模
型 。 与 最 近 的 两 个 基 于vit的 扩 散 模 型 相 比 ， 所
提出的DiffiT在CIFAR-10数据集上的FID分数明显优
于U-ViT [7]和GenViT [76]模 型 。 此 外 ， 就FID分 数
而言，DiffiT在VP和VE训练配置方面明显优于EDM
[34]和ddpm++ [66]模型。在图5中，我们对FFHQ-64数
据集上生成的图像进行了说明。

• 在ImageNet-256数据集中，
  潜在DiffiT模型在FID分数方面优于MDT-G [21]、DiTXL/2-G [52]和StyleGAN-XL [61]等竞争方法，并设
  置了新的SOTA FID分数为1.73。在IS和sFID等其他
  指 标 方 面 ， 潜DiffiT模 型 表 现 出 了 有 竞 争 力 的 性
  能 ， 从 而 表 明 了 所 提 出 的 时 间 依 赖 自 注 意 力 的
  有 效 性 。
• 在ImageNet-512数 据 集 中 ， 隐DiffiT模 型
  在FID和Inception分数(IS)方面明显优于DiT-XL/2-G。
  尽管StyleGAN-XL [61]在FID和IS方面显示了更好的
  性能，但众所周知，基于gan的模型存在多样性低
  的问题，这些问题无法被FID分数捕获。这些问题
  反映在StyleGAN-XL在准确率和召回率方面的次优
  性能上。
• 此外，在图6中，我们展示了在ImageNet-
  256和ImageNet-512数据集上生成的未策划图像的可视
  化。潜DiffiT模型能够在不同的类别中生成各种高质量
  的图像。