DiT Scalable Diffusion Models with Transformers_diffusion with transformer

DiT: Scalable Diffusion Models with Transformers

TL; DR：将 Diffusion 中常用的 UNet 模型替换为 Transformer 模型，并探究了为 Transformer 模型注入条件的结构设计和 scaling 的能力。实验显示，随着模型计算复杂度的上升，生成质量也稳步提升，Transformer 在 Diffusion 中的 scaling 能力得到了验证。而 adaLN 相较于 in-context、cross-attention 等结构更适合为 Transformer 模型注入条件。

---

Diffusion 仅要求其去噪网络是一个输入输出等尺寸的图到图模型，基于 CNN 的 UNet 是一个很自然的选择。但近些年来，凭借着更强的全局理解能力和 scaling 能力，Transformer 在视觉领域也大放异彩。本文中作者使用 ViT 替换掉 UNet，在 Diffusion 中也验证了 Transformer 的 scaling 能力。并试验了 in-context、cross-attention、adaLN 等不同的条件注入方式。

Diffusion Transformer及其条件注入结构设计

patchify

我们知道 Transformer 处理的是一维的序列，而在（图片） Diffusion 中，我们处理的是二维的图片或者特征图，这里需要将特征图展平为一维的序列。为了指示patch 之间的位置关系，DiT 使用了 sin-cos 的位置编码，这个位置编码是不可学习的。

DiT block design

在 patchify 之后，就可以把各个 patch token 送到 Transformer 中进行处理，此时就已经可以做无条件生成了。如果要做有条件生成，还需要将条件编码并输入到模型中。在常用的 Classfier-free Guidance 中，就是要将条件 embedding 和非条件 embedding 输入到模型中。在 LDM 中，是通过 cross-attention 的方式输入到 UNet 模型中的。那么在 Transformer 结构中，如何把条件注入到模型中呢？

DiT 探究了四种对 Transformer 的条件注入方法：

1. In-context conditioning：将两个 embedding 作为两个 special tokens 拼接到图像块 token 后，类似 ViT 中的 cls token，实现起来比较简单，基本没有额外的计算量。
2. Cross-attention：将两个 embedding 拼接起来，然后在 transformer block 中插入一个 cross attention，将 embedding 作为 cross attention的 K 和 V；这也是，该方法需要引入的额外计算量最大，约增加 15%。
3. Adaptive layer norm (adaLN)：adaLN 在 GAN 这类生成模型中的应用非常防范。将常规的 LN 替换为 adaLN，回归 scale 和 shift 两个参数，这种方式也基本不增加计算量。
4. adaLN-zero：即采用零初始化，将 adaLN 的线性层参数初始化为零，网络初始化时 transformer block 的残差模块就是一个 identity 函数。除了回归scale 和 shift，还在每个残差模块结束之前回归一个 scale。

实验对比这四种结构的性能，adaLN-zero 最优。

Transformer Decoder

在隐层进行 Diffusion 去噪之后，需要输出的还是特征图，再送到 vae decoder 中解码出真实图片。但 Transformer 的输出是一维序列，所以这里还需要一个 Transformer Decoder 来将一维序列变换为与输入同尺寸的特征图。这里直接用了一个标准的线性解码器。

Scaling 性能

之所以要用 Transformer 替换掉 UNet，很重要的一个原因就是 Transformer 的 scaling 能力更强，给更多的参数和计算量，就能有更好的性能。

关于 scaling 能力最关键的实验如下图所示。可以看到，生成质量的提升能跟住 log 尺度的计算量增加，scaling 能力相当不错了。

总结

DiT 的思路和做法都是比较直接的，就是用 Transformer 替换掉 UNet，并探索了其 scaling 能力和条件注入的结构。

在 Diffusion 成功之后，以 DiT 为代表的，有很多用 Transformer 替换 UNet 的工作，但都没有受到很大的关注。究其原因，应该是 UNet 已经做的足够好，计算开销也低，使用 Transformer scaling 上去，虽然生成质量有所提升，但也没那么显著。

然而，在最近 OpenAI 的 sora 炸裂登场之后，借助 Transformer 实现了任意长度任意尺寸的视频生成，最大长度更是达到了惊人的 60 秒，对比一众在四五秒挣扎的现有模型来说，真算得上是降维打击了。我们或许应该好好想想如何更好地利用起 Transformer 的特点了。