DiT Scalable Diffusion Models with Transformers

最新推荐文章于 2025-05-06 20:00:05 发布

Adenialzz

最新推荐文章于 2025-05-06 20:00:05 发布

阅读量2.5k

点赞数 22

文章标签：人工智能语言模型计算机视觉

本文链接：https://blog.csdn.net/weixin_44966641/article/details/136170071

版权

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 的条件注入方法：

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

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

在这里插入图片描述