【论文研读】U-DiTs：在U型扩散Transformer中引入下采样Token，以更低计算成本超越DiT-XL/2

速览

论文概述

本文研究了在潜在空间图像生成任务中使用U型扩散Transformer（U-DiT）模型。传统的Diffusion Transformers（DiTs）采用了各向同性的Transformer架构，在图像生成任务中表现出了强大的性能和可扩展性。然而，DiTs放弃了在许多先前工作中广泛应用的U-Net架构。本文首先通过简单的实验发现，U-Net架构的DiT（DiT-UNet）相比各向同性的DiT仅表现出微小的优势，表明U-Net架构中存在潜在的冗余。受U-Net骨干特征以低频成分为主这一发现的启发，本文提出了对自注意力机制中的查询-键-值元组进行下采样，从而在减少计算量的同时进一步提高了性能。基于下采样自注意力的方法，本文提出了一系列U型DiT模型（U-DiTs），并通过广泛的实验证明了U-DiT模型的卓越性能和可扩展性。

主要贡献

提出了U型扩散Transformer（U-DiT）模型：通过对自注意力机制中的查询-键-值元组进行下采样，本文设计了一系列U型扩散Transformer模型（U-DiTs），这些模型在保持高性能的同时显著减少了计算量。
实验验证了U-DiT模型的有效性：通过大量实验，本文证明了U-DiT模型在潜在空间图像生成任务中的卓越性能和可扩展性。特别是在与DiT及其改进版本相比时，U-DiT模型能够以更低的计算成本实现更好的性能。
深入分析了下采样自注意力的效果：本文通过对U-DiT模型中下采样自注意力机制的分析，揭示了其在减少计算冗余和提高模型性能方面的优势。此外，还通过实验验证了不同下采样方法和组件对U-DiT模型性能的影响。

论文及作者相关信息

论文标题: U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers
作者姓名：Yuchuan Tian, Zhijun Tu, Hanting Chen, Jie Hu, Chao Xu
作者单位：北京大学, 华为诺亚方舟实验室
原文地址：https://arxiv.org/pdf/2405.02730
项目地址：https://github.com/YuchuanTian/U-DiT

论文内容解读

背景与挑战

研究背景

随着深度学习技术的飞速发展，特别是在计算机视觉和图像生成领域，Transformer架构逐渐崭露头角。Transformer通过其独特的注意力机制，能够有效捕捉图像中的长距离空间依赖关系，从而在多种视觉任务中表现出卓越的性能。从图像分类、目标检测到图像分割和图像修复，Transformer都展现出了其强大的应用潜力。然而，在图像生成任务，尤其是潜在空间图像生成任务中，Transformer的应用仍面临诸多挑战。尽管近年来提出的Diffusion Transformers（DiTs）通过引入Transformer架构到扩散任务中，取得了令人瞩目的性能和可扩展性，但现有的DiTs及其改进版本大多摒弃了传统的U-Net架构，转而采用各向同性的标准Transformer架构。这种转变虽然带来了性能上的提升，但也引发了研究者们对U-Net架构在降噪方面独特优势的重新思考。因此，本研究致力于探索在潜在空间图像生成任务中，将Transformer架构与U-Net架构相结合的潜力，以期在保持高性能的同时，进一步挖掘U-Net在降噪方面的独特优势。

方法介绍

U-DiT 方法概述

在这篇论文中，作者提出了一种名为 U-DiT（U-shaped Diffusion Transformers）的方法，旨在优化基于 Transformer 架构的扩散模型在潜在空间图像生成任务中的表现。该方法的核心在于将 U-Net 架构与 Transformer 架构相结合，并通过下采样 tokens 的方式优化自注意力机制，从而在减少计算量的同时提升模型性能。

核心关键词

U-Net 架构：一种常用于图像分割和图像生成任务的卷积神经网络架构，其特点是具有编码器-解码器结构，通过下采样和上采样操作逐步提取和恢复图像特征。
Transformer 架构：一种基于自注意力机制的神经网络架构，广泛应用于自然语言处理和计算机视觉任务中。
Tokens 下采样：在自注意力机制中，对 query、key 和 value 进行下采样操作，以减少计算量并优化模型性能。

方法亮点

结合 U-Net 与 Transformer 架构：传统扩散模型通常采用 U-Net 架构，而近期的研究表明，基于 Transformer 的扩散模型（如 DiT）在性能和可扩展性方面表现出色。本文通过将两者结合，旨在充分利用 U-Net 的归纳偏置和 Transformer 的强大建模能力。
Tokens 下采样优化自注意力机制：通过下采样 tokens，可以在保持模型性能的同时显著减少计算量。这一优化策略基于 U-Net 骨干特征的低频主导特性，即高频成分主要包含噪声，而下采样操作可以自然地滤除这些高频噪声。

方法步骤

1. 构建 U-Net 风格的 DiT 模型（DiT-UNet）

模型结构：DiT-UNet 模型采用 U-Net 架构，包括编码器和解码器两部分，每部分包含多个阶段。在每个阶段，编码器通过下采样操作逐步减少特征图的空间尺寸，同时增加特征维度；解码器则通过上采样操作逐步恢复特征图的空间尺寸。
Skip 连接：在每个阶段之间，通过 skip 连接将编码器的特征图与解码器的对应阶段特征图进行融合，以补充解码器在上采样过程中丢失的信息。

2. 对 DiT-UNet 模型进行 Tokens 下采样优化

下采样操作：在下采样 tokens 的过程中，将输入特征图转换为四个空间尺寸减半的特征图，然后分别将这四个特征图映射为 query、key 和 value，并在每个下采样后的特征图上分别执行自注意力操作。
特征融合：在自注意力操作完成后，将下采样后的 tokens 在空间上进行融合，以恢复原始数量的 tokens。整个过程中，特征维度保持不变。

设计初衷与有效性分析

设计初衷：通过结合 U-Net 架构和 Transformer 架构，并利用下采样 tokens 优化自注意力机制，旨在构建一种既具有强大建模能力又具有良好可扩展性的扩散模型。
有效性分析：

图中是对比U-DiTs与DiTs及其改进效果。在40万步训练后，本文作者们绘制了FID-50K与去噪器GFLOPs（采用对数刻度）的关系图。U-DiTs展现出了优于其同类模型的性能。
下采样 tokens 策略能够显著减少计算量，同时保持或提升模型性能。这是因为 U-Net 骨干特征的低频主导特性使得高频成分主要包含噪声，而下采样操作可以自然地滤除这些高频噪声。如图所示，U-DiTs 在与 DiTs 的对比中，展现出了更优的性能表现，进一步验证了该方法的有效性。

3. 扩展模型规模并引入其他优化策略

扩展模型规模：在验证了下采样 tokens 策略的有效性后，作者进一步扩展了 U-DiT 模型的规模，并与更大规模的 DiT 模型进行了比较。

图中展示了不同规模下U-DiTs与DiTs的性能对比。随着训练步骤的增加，U-DiTs的表现始终优于DiTs。标记的大小定性地反映了模型的计算成本。
这一对比结果进一步证明了 U-DiT 方法在扩展模型规模时的优越性能。
模型配置：

图中展示了U-DiTs架构在不同模型规模下的配置情况。其中，“通道”指的是第一层初始输出通道的数量，“编码器-解码器”则表示编码器与解码器模块中的变压器块数量。
U-DiTs 架构的配置情况对于理解其性能扩展和计算成本具有重要意义。通过调整通道数量和编码器-解码器模块中的变压器块数量，U-DiTs 能够在不同规模下实现性能与计算成本的平衡。
引入其他优化策略：除了下采样 tokens 策略外，作者还引入了余弦相似度注意力、RoPE2D 位置编码、深度可分离卷积 FFN 和重参数化等优化策略，以进一步提升 U-DiT 模型的性能。

性能对比与生成质量

性能对比：

上表展示了在ImageNet $512\times512$ 图像生成任务上，U-DiTs与DiTs的性能对比。标有星号（∗）的实验是根据DiT的官方代码进行复现的。本文作者们比较了在标准训练超参数下，训练40万次迭代后的模型表现。
在更高分辨率的图像生成任务上，U-DiTs 与 DiTs 的性能对比进一步验证了 U-DiT 方法的优越性。通过标准训练超参数下的复现实验，U-DiTs 在 $512\\times512$ 图像生成任务上展现出了更优的性能。

在这里插入图片描述
图中是U-DiT-L在100万次迭代后生成的样本。令人惊叹的是，U-DiT仅需100万步训练便能达到如此逼真的视觉质量。建议在屏幕上观看以获得最佳效果。

上述图片展示了 U-DiT-L 模型在 100 万次迭代后生成的样本，其视觉质量令人惊叹。这一图像不仅为读者提供了对 U-DiT 方法生成能力的直观感受，还有助于理解该方法在实际应用中的潜在效果。此外，图片中提及的 U-DiT-L 模型与本文讨论的 U-DiT 方法在潜在空间图像生成任务中的处理对象相符，进一步增强了文章内容的连贯性和可读性。

在这里插入图片描述
上图展示了随着训练的进行，生成样本的质量逐步提升。本文作者们从在ImageNet $256\times256$ 数据集上训练了不同迭代次数的U-DiT模型中采样。更多的训练确实能增强生成质量。最佳观赏效果需通过屏幕查看。

通过上图，可以直观地看到 U-DiT 方法在 ImageNet 数据集上随着训练迭代次数增加，生成样本质量的逐步提升。这一图像进一步验证了 U-DiT 方法的有效性，并为读者提供了对训练过程中生成质量变化的直观感受。

在这里插入图片描述
上表是U-DiT组件的消融研究。除了第3节中的玩具示例外，本文作者们还通过比较在相同浮点运算量（FLOPs）下的U-DiT与精简版DiT-UNet，进一步验证了下采样的有效性。结果显示，下采样可为DiT-UNet带来约18个FID的改进。在U-DiT架构基础上的进一步修改，每项可额外提升2至5个FID。

通过上表的消融研究，本文作者们进一步确认了下采样策略在U-DiT架构中的关键作用，以及该策略对模型性能的显著提升。这一发现不仅加深了对U-DiT方法的理解，也为未来在扩散模型优化方面的研究提供了有价值的参考。

实验与结果

实验设计

U-Net风格的DiT设计

为了验证U-Net架构在DiT中的潜力，作者首先设计了一个简单的U-Net风格的DiT（DiT-UNet）。该模型由编码器和解码器组成，每个编码器或解码器阶段包含相同数量的DiT块。在每个编码阶段，特征图通过2倍下采样，特征维度翻倍，同时利用跳跃连接将编码阶段的输出与解码阶段的输出进行融合，以补充解码过程中因下采样丢失的信息。作者选择了三个编码/解码阶段，以适应紧凑的潜在空间（对于256x256大小的图像，潜在空间为32x32）。

为了比较U-Net风格的DiT与原始DiT的性能，作者将DiT-UNet与计算成本相似的DiT-S/4进行了对比。训练迭代次数为400K，所有训练超参数保持不变。

在这里插入图片描述

上图展示了在U-Net风格DiT上的玩具实验。原始的DiT-UNet表现略优于各向同性DiT $S /4$ ；但有趣的是，当对自注意力应用标记下采样时，DiT-UNet以更低的成本实现了更好的性能。

在这里插入图片描述
上图展示了从DiT到所提U-DiT的演进。左图(a)：原始DiT，采用各向同性架构。中图(b)：DiT-UNet，一种朴素的U-Net风格DiT。在玩具实验中，尝试将DiT与U-Net进行简单结合。右图©：所提U-DiT。本文作者提出对输入特征进行下采样以用于自注意力机制。该下采样操作能显著优化DiT-UNet，大幅削减计算量。

改进后的U-Net风格DiT

在初步实验中发现，U-Net风格的DiT相较于原始DiT仅有微弱的优势。为了进一步提高性能，作者提出了对注意力机制进行改进。他们观察到，在潜在空间扩散模型中，U-Net骨干网络提供的特征主要以低频分量为主。基于这一发现，作者提出了对查询（Query）、键（Key）和值（Value）三元组进行下采样的方法，以过滤掉高频噪声，同时减少计算量。

下采样操作将输入特征图转换为四个2倍下采样的特征图，然后将这些下采样后的特征图映射为Q、K、V进行自注意力计算。自注意力在每个下采样后的特征图内独立进行，之后将下采样后的特征图合并以恢复原始数量的特征图。这种方法不仅减少了计算量，还保留了原始特征图的维度。

实验结果

初步实验结果

在初步实验中，作者发现U-Net风格的DiT（DiT-UNet）相较于原始DiT（DiT-S/4）在性能上仅有微弱的优势。这表明，简单地将U-Net与DiT块结合并没有充分利用U-Net的归纳偏置。

改进后的U-Net风格DiT性能

通过引入下采样注意力机制，U-Net风格的DiT在性能上有了显著提升。尽管计算量显著减少，但FID（Fréchet Inception Distance）指标却有所降低，表明图像生成质量得到了提高。具体来说，下采样注意力机制在保持计算成本相近的情况下，使FID降低了约18个点。

更大规模的模型比较

为了验证所提出方法的有效性，作者将U-DiT模型扩展到更大规模，并与更大规模的DiT模型进行了比较。实验结果显示，U-DiT模型在性能上显著优于DiT模型。具体来说，U-DiT-B模型在计算成本仅为DiT-XL/2模型的1/6的情况下，仍然能够取得更好的FID指标。U-DiT-L模型甚至能够在FID上比DiT-XL/2模型低约10个点。

此外，作者还进行了条件生成和无条件生成的实验，结果显示U-DiT模型在分类器指导下的生成性能也优于DiT模型。在扩展训练步数到100万步后，U-DiT模型的性能仍然持续提升，表明其具有良好的可扩展性。

消融实验

作者还进行了消融实验，以评估不同组件对U-DiT模型性能的贡献。结果显示，下采样注意力机制对性能的提升最为显著，而其他改进（如余弦相似度注意力、RoPE2D位置编码、深度可分离卷积FFN等）也进一步提高了模型的性能。即使在没有这些额外改进的情况下，原始的U-DiT模型仍然能够在性能上超越DiT-XL/2模型。

在这里插入图片描述
图中展示的是ImageNet数据集内的一张图片，尺寸为256x256像素。