论文阅读：Scalable Diffusion Models with Transformers

本文主要是介绍论文阅读：Scalable Diffusion Models with Transformers，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

Scalable Diffusion Models with Transformers

论文链接

介绍

传统的扩散模型基于一个U-Net骨架，这篇文章提出了一种新的扩散模型结构，将U-Net替换为一个transformer，并将这种结构称为Diffusion Transformers (DiTs)。他们还发现，transformer的规模越大（通过Gflops衡量），生成的图片的质量越好（FID越低）。
如图2所示，DiT的规模越大，图片生成的质量越好（左图），和当前流行的扩散模型相比，DiT的计算效率也表现优异。
ImageNet generation with Diffusion Transformers (DiTs)

方法（Diffusion Transformers）

预备知识

Diffusion formulation：扩散模型Diffusion Model（DM）在训练过程中，首先向图片中添加噪声，然后预测噪声来从图片中将噪声去除。这样，在推理过程中，首先初始化一个高斯噪声图片，然后去除预测的噪声，即可得到生成的图片。
Classifier-free guidance：条件扩散模型引入了额外信息 $c$ （比如，类别）作为输入。而classifier-free guidance可以引导生成的图片 $x$ 是类别 $c$ 的概率 $l o g (c ∣ x)$ 最大。
Latent diffusion models：扩散模型在像素空间上训练和推理的计算开销过大，Latent Diffusion Model（LDM）将像素空间替换为VAE编码得到的潜在空间 $z = E (x)$ ，可以提高计算效率。本文提出的DiT沿用了LDM中的潜在空间，但是在预测潜在空间特征的模型上，将LDM中的U-Net替换为了纯Transformer骨架。

Diffusion Transformer Design Space

Diffusion Transformers (DiTs)是基于Vision Transformer (ViT)的模型，它的大体结构如图3所示，从左图可以看到，输入的噪音特征被分解为不同批，然后被若干个DiT块处理；右边的三张图展示了DiT块的详细结构，分别是三种不同的变体。
The Diffusion Transformer (DiT) architecture
下面对DiT的各层进行分析：
Patchify. 从图3中可以看到，DiT的第一个层是Patchify，其将输入转化为 $T$ 个token序列。在这之后，作者使用标准ViT中基于频率的位置嵌入处理前面的token序列。而token序列的数量是由一个超参数 $p$ 决定的， $p$ 减半导致 $T$ 翻四倍，并且导致整个transformer的GFlops至少翻四倍，如图4所示。
Input specifications for DiT
DiT block design. 在patchfiy层之后，几个transformer块处理输入token以及一些额外的条件信息，比如，类标签 $c$ 和时间步数 $t$ 。作者尝试了4种不同的ViT变体：

In-context conditioning：这种变体直接将时间步数 $t$ 和类标签 $c$ 作为额外的token添加到输入token序列后面，类似于ViT的cls tokens，因此也可以直接使用标准的ViT块。这种方式引入的Gflops可以忽略不计。
Cross-attention block：这种变体将条件信息拼接为一个长度为2的序列，独立于图片输入序列。然后，在transformer块的self-attention层后添加了一个cross-attention层，类似于LDM，在cross-attention层将条件信息加入图片特征中。cross-attention方案增加的Gflops最多，大概15%。
Adaptive layer norm (adaLN) block：这种变体将transformer块中标准的layer norm layers替换为adaptive layer
norm (adaLN)，这一技术在GAN相关的模型中被广泛采用。不同于直接学习维度放缩和偏移因子 $\gamma$ 和 $\beta$ ，该方案回归 $t$ 和 $c$ 的嵌入的和得到这两个参数。在目前的三种方案中，该变体额外增加的Gflops最少。
adaLN-Zero block：先前的工作说明，ResNet中的恒等映射是有益处的。Diffusion U-Net在残差之前，零初始化了每个块中最后一个卷积层。作者采用了和Diffusion U-Net相同的方案。此外，除了回归 $\gamma$ 和 $\beta$ ，该方案还对DiT块中残差连接上的放缩因此 $\alpha$ 进行了回归。对于所有的 $\alpha$ ，作者初始化MLP以输出零向量，这使得DiT块为一个恒等函数。和adaLN方案一样，ada-Zero方案引入的Gflops也可以忽略不计。

Model Size. 作者设置了四种规模的DiT：DiT-S, DiT-B, DiT-L and DiT-XL，结构复杂度依次增大。
Transformer decoder. 在经过最后的DiT块之后，使用tranformer decoder将输入tokens转化为和输入同等性状的噪音预测。

综上，作者探索了DiT设计空间中的patch_size、transformer架构（4种，in-context，cross-attention， adaptive layer
norm and adaLN-Zero blocks）和model size（4种，DiT-S, DiT-B, DiT-L and DiT-XL）。

实验

实验设置

训练：在256 × 256和512 × 512 图片分辨率的ImageNet数据集上训练。超参数设置几乎和ADM一致。
Diffusion：和Stable DIffusion一样使用VAE编码图片和解码特征。
评估指标：主要使用Fr´echet Inception Distance (FID)，还使用了Inception Score [51], sFID [34] and Precision/Recall [32]
计算平台：在JAX [1]这个深度学习框架上实现了DiT，在TPU上训练模型。

实验结果

DiT block design. 四个不同的DiT块：in-context (119.4 Gflops), cross-attention (137.6 Gflops),
adaptive layer norm (adaLN, 118.6 Gflops) or adaLN-zero (118.6 Gflops)中， adaLN-zero (118.6 Gflops) 取得最低的FID。其中，adaLN-zero相较于adaptive layer norm的提升，说明了恒等映射的好处。（后续的实验除非特别说明都是在adaLN-zero上做的）

Comparing different conditioning strategies
Scaling model size and patch size. 模型size增大和patch zise减小，均会提高Gflops，降低FID。我们注意到，DiT-L 和DiT-XL的FID很接近，因为它们的Gflops也相对更接近。
Scaling the DiT model improves FID at all stages of training
DiT Gflops are critical to improving performance. 上面的图6再次说明了模型参数量的增大并不等同于DiT模型的图片质量提高，真正的关键是提高Gflops。比如，DiT S/2的表现和DiT B/4接近，因为小的batch size会增大Gflops，二者的Gflops接近，所以FID也接近。
Larger DiT models are more compute-efficient
小的DiT模型即便训练时间更长，相对于训练时间更短的大的DiT模型，其计算效率也是更差的。
这里，作者估计训练计算量的方式为model Gflops · batch size · training steps · 3。
Larger DiT models use large compute more effi-
ciently