U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers

Diffusion Transformers（DiTs）将变压器架构引入潜在空间图像生成的扩散任务中。DiTs具有各向同性的架构，通过一系列变压器块的链接，展示出竞争性能和良好的可扩展性；但同时，DiTs放弃了U-Net，并对其进行了改进，这值得重新思考。为此，我们通过比较U-Net架构的DiT和各向同性的DiT进行了简单的玩具实验。结果表明，U-Net架构仅在U-Net归纳偏差方面获得了轻微优势，表明U-Net风格的DiT中存在潜在的冗余。受到U-Net骨干特征低频占主导的发现的启发，我们对自我注意力的查询-键-值元组进行令牌下采样，并带来了进一步的改进，尽管计算量大大降低。基于令牌下采样的自我注意力，我们在论文中提出了一系列U形DiTs（U-DiTs），并进行了广泛的实验，以展示U-DiT模型的出色性能。所提出的U-DiT可以在仅有DiT-XL/2计算成本的1/6的情况下胜过它。代码可在https://github.com/YuchuanTian/U-DiT上获得。

论文试图通过比较U-Net架构和等向性架构的DiT模型，提出一种新的U-DiT模型解决图像生成任务中的低频特征问题，并优化计算效率。

通过对U-Net架构和等向性架构的DiT模型进行比较，发现U-Net架构在具有U-Net归纳偏差的情况下仅略有优势，因此提出U-DiT模型解决低频特征问题，并通过下采样优化计算效率。

论文提出了U-DiT模型，能够在低计算代价的情况下超过DiT-XL/2的性能；通过下采样优化了自注意力机制，提出了一系列U-DiT模型，并在多个数据集上进行了实验。

最近的相关研究包括：DiT模型在图像生成任务中的应用，Transformer架构的改进，U-Net架构在图像生成任务中的应用等。