清华 & 华为 开源U-DiTS | 利用 U-Net 架构提升扩散 Transformer 的性能和效率！

扩散 Transformer （DiTs）将 Transformer 架构引入到潜在空间图像生成的扩散任务中。具有各向同性的架构，通过一系列 Transformer 块的链式组合，DiTs展示了竞争性的性能和良好的可扩展性；

但与此同时，DiTs放弃U-Net及其后续改进是值得重新思考的。为此，作者进行了一个简单的玩具实验，比较了一个采用U-Net架构的DiT与一个各向同性的DiT。结果表明，在U-Net的归纳偏置下，U-Net架构只获得了微小的优势，这表明U-Net风格的DiT中存在潜在的冗余。

受到U-Net Backbone 特征以低频为主的发现的启发，作者对自注意力中的 Query -键-值元组进行标记下采样，尽管计算量大幅减少，但仍带来了进一步的改进。

在基于下采样标记的自注意力基础上，本文提出了一系列U型DiTs（U-DiTs），并进行了广泛的实验来证明U-DiT模型的卓越性能。所提出的U-DiT仅用1/6的计算成本就能超越DiT-XL/2。

代码可在https://github.com/YuchuanTian/U-DiT获取。

1 Introduction

感谢注意力机制建立了长距离的空间依赖关系，Transformers 在各种视觉任务中，包括图像分类 、目标检测、分割  和图像恢复 ，都证明了其高效性。DiTs 将完整的Transformer架构引入到扩散模型中，这在小鼠图像空间和潜在空间生成任务上展示了卓越的性能和可扩展性。最近的后续工作通过将扩散Transformer的应用扩展到灵活分辨率的图像生成 、真实视频生成等领域，展示了其前景广阔。

有趣的是，DiTs 抛弃了在众多之前工作中普遍应用的 U-Net 架构，无论是像素空间还是潜在空间。DiTs 中的各向同性架构确实取得了成功，因为规模放大的DiT模型取得了最高性能。然而，DiTs 及其改进工作在潜在空间图像生成任务上抛弃了广泛应用的 U-Net 架构，引发了作者的好奇心，因为 U-Net 的归纳偏置被认为有助于去噪。因此，作者重新考虑在典型的 U-Net 架构上部署 DiTs。

为了实验 U-Net 与 DiT 的结合，作者首先提出了一个简单的 U-Net 风格的 DiT（DiT-UNet）并与类似大小的各向同性 DiT 进行比较。结果表明，在相似的计算成本下，DiT-UNet 仅与 DiT 相当。从这个玩具实验中推理，当简单地将 U-Net 和普通 Transformer 块组合时，并没有充分利用 U-Net 的归纳偏置。

因此，作者重新思考了 DiT-UNet 中的自注意力机制。潜在 U-Net 去噪器中的主干提供了主要包含低频成分的特征 [27]。这一发现暗示了主干特征中存在冗余：U-Net 扩散器中的注意力模块应该强调低频域。由于之前的理论赞扬了在扩散过程中下采样以过滤高频噪声，作者希望通过在自注意力之前对特征进行标记下采样，利用这一自然的低通滤波器。与之前只在关键值对上执行下采样的 Transformer 工作不同，作者彻底地下采样 Query -关键值-值元组，使得自注意力在降采样后的潜在标记之间执行。令人惊讶的是，当作者将带有降采样标记的自注意力整合到 DiT-UNet 中时，在潜在 U-Net 扩散器上获得了更好的结果，并且计算量显著减少。

基于这一发现，作者扩大了带有降采样自注意力的 U-Net，并提出了一系列最先进的 U 形扩散 Transformer（U-DiTs）。作者进行了许多实验，以验证作者的 U-DiT 模型相对于各向同性 DiTs 在性能和可扩展性方面的卓越表现。如图2 & 图2 所示，U-DiTs 可以大幅超越 DiTs。令人惊讶的是，所提出的 U-DiT 模型在 FLOPs 上比 DiT-XL/2 小 6 倍的情况下，还能表现得更好。

2 Preliminaries

视觉 Transformer  (Vision Transformers) ViTs 通过将输入分割成块并将图像视为一系列块 Token ，将 Transformer 架构引入视觉任务，并在大规模图像分类任务上证明了其有效性。虽然ViTs采用等距架构，但后续一些关于视觉 Transformer 的工作提出了类似金字塔的分层架构，逐渐下采样特征。金字塔架构在分类和其他下游任务中已被证明非常有效。

视觉 Transformer 也是去噪模型的主流 Backbone 网络。IPT [6] 为去噪和其他低级任务引入了等距 Transformer  Backbone 网络。一些后续工作遵循等距约定，但其他去噪工作转向U-Net Backbone 网络作为其设计。U-ViT  和 DiT 的开创性工作将完整的 Transformer  Backbone 网络引入扩散作为去噪器。

**扩散 Transformer 中的最新进展 (Recent Advancements in Diffusion Transformers)**。在DiTs之后，一些研究探讨了扩散 Transformer 的训练和扩散策略。其他工作聚焦于DiT Backbone 网络的设计。DiffiT 引入了一种新的条件融合方法；FiT 和 VisionLLaMA 通过引入包括RoPE2D 和 SwishGLU在内的LLM技巧来加强DiT。这些基于 Transformer 的扩散工作一致赞同在潜在空间上采用等距架构，即在整个扩散模型中，潜在特征空间不下采样。DiT 的作者甚至认为U-Net的归纳偏置“并不关键”。

**扩散中的U-Net (U-Nets for Diffusion)**。从经典工作来看，U-Net 的设计理念在扩散中得到了普遍接受。具体来说，Stable Diffusion 在压缩的潜在空间上使用基于U-Net的去噪器进行高分辨率图像合成，这在许多生成任务中取得了巨大成功。一些之前在扩散 Transformer 上的尝试也在像素空间生成任务上采用U-Net；但奇怪的是，它们在潜在空间扩散上转向了类似等距的DiT结构。尽管它在像素空间扩散中很受欢迎，但U-Net架构在最近面向潜在空间扩散的 Transformer 导向工作中并未被广泛接受。

受到这一点的启发，作者致力于研究基于 Transformer  Backbone 的U-Net在潜在空间扩散中的潜力。

3 Investigating U-Net DiTs in Latent

正如所回顾的，U-Net架构在扩散应用中得到了广泛采用；对U-Net去噪器的理论评估也揭示了它们的优势，因为下采样U-Net阶段的转换可以过滤掉主导高频的噪声[35]。因此，对各单位向同性架构的潜在扩散 Transformer 的空前放弃是反直觉的。作者通过一个玩具实验重新思考和阐明基于 Transformer  Backbone 的U-Net去噪器在潜在扩散中的潜力。

图3：从DiT到提出的U-DiT的演变。 左边（a）：原始的DiT，使用各单位向同性架构。中间（b）：DiT-UNet，这是一个简单的U-Net风格的DiT。作者在玩具实验中尝试将DiT和U-Net简单结合。右边（c）：提出的U-DiT。作者建议对输入特征进行下采样以进行自注意力。下采样操作可以显著提高DiT-UNet，同时大幅减少计算量。

典型的U-Net风格的DiT。 首先，作者提出了一种简单的基于 Transformer  Backbone 的U-Net去噪器，名为DiT-UNet，通过将DiT块嵌入到典型的U-Net架构中。按照先前的U-Net设计，DiT-UNet由具有相等数量阶段的编码器和解码器组成。当编码器通过下采样图像作为阶段级量处理输入图像时，解码器从最压缩的阶段将编码图像上采样到输入大小。在每个编码器阶段转换时，空间下采样2倍，同时特征维度也翻倍。在每次阶段转换都提供跳跃连接。跳过的特征与来自前一个解码器阶段的上采样输出连接并融合，为由于特征下采样而带来的解码器信息损失补充信息。考虑到较小、狭窄的潜在空间（对于256×256大小的生成，为32×32），作者总共指定3个阶段，即特征下采样两次然后恢复到原始大小。

为了适应跨多尺度阶段的各个特征维度的时间嵌入和条件嵌入，作者为各个阶段使用独立的嵌入器。此外，作者避免将潜在空间分割为块，因为U-Net架构本身下采样潜在空间，无需进一步的空间压缩。通过玩具实验，作者将提出的U-Net风格的DiT与采用各单位向同性架构的原始DiT进行比较。为了使模型与DiT设计保持一致，作者在每个阶段重复使用简单的DiT块。每个DiT块包括一个自注意力模块作为标记混合器和一个两层的前馈网络作为通道混合器。作者通过训练U-Net风格的DiT 40万次迭代来进行实验，并与大小相当的DiT-S/4进行比较。所有训练超参数保持不变。结果显示，U-Net风格的DiT只比原始各单位向同性DiT有有限的优势。U-Net的归纳偏置没有得到充分利用。

通过标记下采样改进U-Net风格的DiT。 为了更好地将注意力引入到扩散U-Net中的 Transformer ，作者回顾了U-Net Backbone 作为扩散去噪器的作用。最近关于潜在扩散模型[27]的工作对U-Net Backbone 中间特征进行了频率分析，并得出结论认为能量集中在低频域。这一频率域的发现暗示了 Backbone 中的潜在冗余：U-Net Backbone 应该从全局角度强调粗糙的物体，而不是高频细节。

作者采用带有降采样标记的注意力机制。降采样操作是一种自然的低通滤波器，可以丢弃高频成分。在扩散场景下，已经研究过了降采样的低通特性，结论是降采样有助于扩散过程中的去噪器，因为它自动“丢弃了那些被噪声主导的高频子空间”[35]。因此，作者选择对标记进行降采样。

实际上，对降采样标记的注意力并不是新方法。先前的关于视觉 Transformer [15; 38]的工作提出了降低键值对计算成本的方法。最近在无需训练的扩散加速[28]上的研究也在Stable Diffusion模型上应用了键值降采样。但这些工作保持了 Query 的数量，因此降采样操作并未完全执行。此外，这些降采样措施通常涉及张量大小的减少，可能导致信息的重大损失。

与这些工作不同，作者为DiT-UNets提出了一种简单而激进的标记降采样方法：作者同时降采样 Query 、键和值以进行扩散友好的自注意力，但同时保持整体张量大小以避免信息损失。具体步骤如下：特征图输入首先通过降采样器（降采样器的设计在4.2节中详细介绍）转换为四个 降采样特征。然后，将降采样特征映射到 ， ， 进行自注意力。在每个降采样特征内执行自注意力。在注意力操作之后，降采样标记在空间上合并为一个单元以恢复原始的标记数量。值得注意的是，在整个过程中保持了特征维度不变。与U-Net降采样不同，在降采样过程中，作者并未减少或增加特征中的元素数量。相反，作者以并行方式将四个降采样标记送入自注意力。

带有降采样标记的自注意力确实有助于DiT-UNets在潜在扩散任务上的表现。如表格1所示，将降采样自注意力替换为全尺度自注意力，在Frechet Inception Distance (FID)指标上带来了轻微的提升，尽管FLOPs显著减少。

复杂性分析。除了性能优势之外，作者注意到，与全尺度自注意力相比，降采样自注意力可以节省最多1/3的总计算成本。作者对自注意力机制进行了简短的复杂度分析，以解释节省成本的原因。

表格1：关于U-Net风格的DiTs的玩具实验。简单的DiT-UNet比各向同性的DiT-S/4稍好一些；但有趣的是，当作者为自注意力应用标记降采样时，DiT-UNet在成本较低的情况下表现更佳。

给定一个大小为 、维度为 的输入特征，作者用 表示映射后的 Query -键-值元组。自注意力机制的复杂性分析如下：在提出的在降采样标记上的自注意力中，四组降采样的 Query -键-值元组 分别执行自注意力。尽管每个自注意力操作的复杂度仅为完整尺度自注意力的1/16，但降采样自注意力的总复杂度是完整尺度自注意力的1/4。通过标记降采样，自注意力机制的计算成本节省了3/4。

简而言之，作者从简单的实验中展示了，通过为自注意力降采样标记，DiT-UNet 的冗余性得到了减少。

4 Scaling the Model Up

基于作者在玩具实验中的发现，作者提出了一系列U型DiTs（U-DiT），通过应用降采样自注意力（在第3节中提出）并将U-Net风格的DiT放大。

设置。作者采用了DiT的训练设置。采用了与潜在扩散模型[25]相同的VAE（即sd-vae-ft-ema）和AdamW优化器。训练超参数保持不变，包括全局批量大小256，学习率 ，权重衰减0，全局种子0。训练是在ImageNet 2012[10]的训练集上进行的。

除了在第3节中介绍的玩具实验中的降采样自注意力之外，作者进一步对U-DiTs进行了一系列修改，包括余弦相似度注意力[20; 18]，RoPE2D[30; 22; 8]，深度卷积FFN[34; 3; 38]和重参化[12; 31]。每一项修改的贡献在 Sec. 6 中进行了定量评估。

U-DiT at Larger Scales

表2：**在ImageNet 256×256生成任务上比较U-DiTs与DiTs。在超级标记 的实验中，根据DiT的官方代码进行复制。作者比较了训练400K次迭代的模型，使用了DiT的标准训练超参数。U-DiTs的性能非常出色：U-DiT-B仅用1/6的推理FLOPs就能超越DiT-XL/2；U-DiT-L可以比DiT-XL/2低10个FIDs。从表2可以得出结论，所有的U-DiT模型都以相当大的优势超越了它们各向同性的对应模型。特别是，U-DiT-S和U-DiT-B能够比尺寸相近的DiTs低约30个FIDs；U-DiT-L可以比DiT-XL/2低约10个FIDs。令人震惊的是，U-DiT-B仅用1/6的计算成本就能超越DiT-XL/2。为了更好地展示作者方法的优势，作者还包含了U-DiTs在FID-50K对FLOPs图中的性能（图1）。除了DiTs和U-DiTs，作者还包括了其他最先进的方法：SiT [23]，它为DiTs提出了一个插值框架，以及SiT-LLaMA [8]，它结合了最先进的DiT主干网络VisionLLaMA和SiT。在图中，U-DiTs与其他 Baseline 的优势非常明显。结果显示了所提出的U-DiT模型非凡的可扩展性。

与DiTs及其改进的比较。 为了验证所提出的U-DiT模型在简单玩具实验之外的有效性，作者将它们扩大规模，并与更大尺寸的DiTs [24]进行比较。为了公平比较，作者使用了与DiT相同的训练超参数集；所有模型都训练了400K次迭代。在ImageNet 256×256上的结果如表2所示，作者分别将U-DiTs扩展到约6e9、约20e9、约80e9 FLOPs，并与类似计算成本的DiTs进行比较。在无分类器引导的生成场景中，U-DiTs同样表现优异。在表3中，作者在 的情况下比较了U-DiTs与DiTs。为了公平比较，作者在相同的设置下训练了U-DiTs和DiTs，共400K次迭代。

扩展训练步骤。 作者通过将训练步骤扩展到100万次来挖掘U-DiTs的潜力。图2进一步证明，U-DiTs在所有训练步骤中的优势是一致的。随着训练步骤逐渐增加到100万次，U-DiTs的性能在提高（表4）。作者可视化了这个过程，其中图像质量逐渐变得更好（图4）。值得注意的是，仅用600K训练步骤的U-DiT-L就能在没有分类器自由引导的情况下超越训练了7M步骤的DiT-XL/2。如图5额外展示的那样，在仅仅100万次迭代下，U-DiT模型就能有条件地生成真实的图像。

Ablations

图4：随着训练的进行，生成样本质量的提升。 作者从在ImageNet 256 256上训练了不同迭代次数的U-DiT模型中进行采样。更多的训练确实能提高生成质量。最佳在屏幕上查看。

表3：使用无分类器引导的生成性能。 作者在400K训练步骤时测量了U-DiTs和DiTs的性能， 。带有超级标记 的实验是根据DiT的官方代码复制的。U-DiTs在条件生成方面也表现出色。

下采样器的设计。 所提出的U-DiT中的下采样操作将完整的特征转换为多个空间下采样的特征。基于以往的智慧，作者得出这样一个结论：先前的工作要么直接执行像素重排，要么在像素重排之前应用卷积层。作者认为直接将像素重排作为下采样过于死板，而在计算成本方面，应用卷积几乎是负担不起的。具体来说，普通卷积由于涉及密集的通道连接而代价高昂：使用基于卷积的下采样器可能使计算成本翻倍。作为妥协，作者改用深度卷积。作者还添加了一个短路，直接连接这个深度卷积，这对于提高性能至关重要。这个短路为模型增加的计算成本可以忽略不计，实际上，在推理阶段可以通过重参化技巧将其移除。结果在表5中展示。

每个单独修改的贡献。 在这部分中，作者从普通的U-Net风格的DiT（DiT-UNet）开始，评估各个组件的贡献。首先，作者检查下采样自注意力的优势。回顾第3节的玩具实验结果，用下采样自注意力替换全尺度自注意力将导致FID的改进和FLOPs的1/3减少。为了评估通过模型性能的下采样改进，作者还设计了DiT-UNet的精简版本（即DIT-UNet（Slim））。DiT-UNet（Slim）作为一个全尺度自注意力的 Baseline ，其计算量与作者的U-DiT大约相同（ GFLOPs）。如表6的上半部分所示，通过将U-DiT与DiT-UNet（Slim）进行比较，结果表明，在DiT-UNet中下采样标记可以带来约18FIDs的性能提升。

表4：关于训练迭代次数的U-DiT-B和U-DiT-L模型的性能。 这两个模型在ImageNet 256 256上的无条件生成性能随着训练的进行而持续改进，其中U-DiT-L在600K步骤时的表现明显优于DiT-XL/2在700万步骤时的表现。

表5：关于下采样器选择的消融研究。 作者尝试了多种下采样器设计，结果证明，快捷方式和深度卷积的并行连接是最合适的。作者避免使用普通卷积（即Conv.+PS），因为通道混合是昂贵的：传统的基于卷积的下采样器可能使计算量翻倍。带有传统下采样器的U-DiT总共需要多达2.22G FLOPs。

接下来，作者检查进一步改进U-DiTs的其他修改（表6的下半部分）。Swin Transformer V2 [20]提出了自注意力的一个更强版本：不是直接相乘Q和K矩阵，而是使用 Query 和键之间的余弦相似性。作者将这个设计应用到作者的自注意力中，带来了约2.5FIDs的改进。RoPE [30]是一种强大的位置嵌入方法，已被广泛应用于大型语言模型。遵循最新的扩散转换工作[22; 8]，作者在自注意力之前将二维RoPE（RoPE2D）注入 Query 和键中。引入RoPE2D使性能提高了约2.5FIDs。一些最近的研究工作通过在两个线性映射之间插入深度卷积层来加强MLP [34; 3; 38]。由于这些工作中证明了该措施的有效性，作者将其借鉴到作者的U-DiT模型中，改进了约5FIDs。由于训练期间的重参化[12]可以提高模型性能，作者将该技巧应用到FFN [31]上，并带来了额外的约3.5FIDs改进。

总之，基于上述组件，所提出的U-DiTs可以大幅度优于普通的DiT-UNets和各向同性的DiTs。

5 Conclusion

表6：对U-DiT组件的消融研究。除了第3节中的玩具示例外，作者还通过比较具有相等FLOPs的简化版DiT-UNet来进一步验证降采样效果。结果显示，降采样可以为DiT-UNet带来约18的FIDs改善。在U-DiT架构上的进一步修改可以提高2到5个FIDs。

图5：U-DiT-L在100万次迭代生成的样本。令人惊讶的是，U-DiT在仅100万训练步数时就能达到真实的视觉质量。最佳在屏幕上观看。

在本文中，作者强调了在U-Net架构中对DiTs在潜在空间生成方面的重要性。尽管具有各向同性架构的DiTs已经证明了其强大的可扩展性和卓越的性能，但U-Net的归纳偏置效果被忽视了。因此，作者重新思考了U-Net风格的DiTs。作者首先研究了简单的DiT-UNet，这是U-Net和DiT块的直接组合，并尝试减少U-Net主干中的计算冗余。

受到之前扩散智慧的启发，作者提出对视觉标记进行降采样以进行自注意力，并取得了非凡的结果：尽管FLOPs大幅减少，但性能却进一步提高。从这个有趣的发现出发，作者将U-Net架构进行扩展，并提出了一系列U形DiT模型（U-DiTs）。作者进行了各种实验来证明作者的U-DiTs的卓越性能和可扩展性。

参考

[1].U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers.