清华 & 卡梅隆 & 上交大 打破计算瓶颈，DiTFastAttn 方法优化扩散 Transformer 的图像与视频生成 ！

扩散 Transformer （DiT）在图像和视频生成方面表现出色，但由于自注意力的二次复杂度，面临着计算上的挑战。作者提出了DiTFastAttn，一种新颖的后训练压缩方法，以减轻DiT的计算瓶颈。 作者在DiT推理过程中的注意力计算中识别出三种关键冗余： 1.空间冗余，许多注意力头聚焦于局部信息； 2.时间冗余，相邻步骤的注意力输出之间具有高相似性； 3.条件冗余，条件和无条件推理表现出显著相似性。 为了解决这些冗余，作者提出了三种技术： 1.带残差缓存的窗口注意力以减少空间冗余； 2.时间相似性降低以利用步骤间的相似性； 3.条件冗余消除以在条件生成过程中跳过冗余计算。 为了证明DiTFastAttn的有效性，作者将其应用于DiT、PixArt-Sigma进行图像生成任务，以及OpenSora进行视频生成任务。 评估结果显示，对于图像生成，作者的方法最多减少了88%的FLOPs，并在高分辨率生成时实现了最高1.6倍的加速。

1 Introduction

扩散 Transformer （DiT）最近在图像生成和视频生成中越来越受欢迎。然而，DiT的一个主要挑战是它们的计算需求量很大，特别是在生成高分辨率内容时特别明显。一方面，传统的 Transformer 架构，由于其自注意力机制，对输入标记长度L具有的复杂度。这种二次复杂度随着图像和视频分辨率的提高导致计算成本显著增加。如图1所示，随着图像分辨率的增加，注意力计算在推理过程中成为主要的计算瓶颈。具体来说，如果一个的图像被标记为16k个标记（Chen等人，2024），即使在像Nvidia A100这样的高端GPU上，注意力计算也需要几秒钟。另一方面，由于多个去噪步骤和分类器自由引导（CFG）技术，扩散推理过程需要大量的神经网络推理。

之前加速注意力机制的努力，例如局部注意力、Swin Transformer（Liu等人，2021）和组 Query 注意力（GQA）（Ainslie等人，2023），主要集中在对注意力机制或网络架构的设计上。尽管这些方法在降低计算成本方面有效，但它们需要较大的重新训练成本。由于训练DiT需要大量的数据和计算资源，因此需要后训练压缩方法。

在这项工作中，作者识别了DiT推理中注意力计算的三种冗余类型，并提出了一种后训练模型压缩方法DiTFastAttn来处理这些冗余：

(1) _空间维度的冗余_。许多注意力头主要捕捉局部空间信息，对于远距离标记的注意力值接近零。为了减少冗余，作者选择对某些层使用窗口注意力而不是全注意力。然而，直接丢弃窗口外的所有注意力计算会导致性能显著下降。为了在无需训练的情况下保持性能，作者提出缓存一步中全注意力和窗口注意力输出之间的残差，并在随后的几步中重用这个残差。作者将这种技术称为带残差共享的窗口注意力（WA-RS）。

(2) _相邻步骤之间注意力输出的相似性_。相同注意力头在相邻步骤之间的注意力输出可能非常相似。作者提出了跨时间步的注意力共享（AST）技术，利用这种步进相似性来加速注意力计算。

(3) _条件推理和无条件推理之间注意力输出的相似性_。作者观察到在CFG中，对于某些头和时间步，条件推理和无条件推理的注意力输出表现出显著的相似性（SSIM ）。因此，作者提出了跨CFG的注意力共享（ASC）技术，以在条件生成过程中跳过冗余计算。

作者进行了大量实验来评估使用多个DiT模型的DiTFastAttn，包括用于图像生成的DiT-XL（Peebles和Xie，2023）和PixArt-Sigma（Chen等人，2024），以及用于视频生成的Open-Sora（Open-Sora，2024）。作者的研究结果表明，DiTFastAttn一致地减少了计算成本。特别是，分辨率越高，计算和延迟的节省越多。例如，在使用PixArt-Sigma生成20482048图像时，DiTFastAttn可以在注意力计算上减少36%至88%，并将速度提升至1.6倍。

2 Related Work

Diffusion Models

扩散模型 因其与GANs（Creswell等人，2018）相比具有更优越的生成性能而受到广泛关注。早期的扩散模型是基于U-Net架构实现的。为了实现更好的可扩展性，DiT 采用了 Transformer 架构而不是U-Net。扩散 Transformer 被应用于图像和视频生成的领域。PixArt-Sigma（Chen等人，2024）展示了扩散 Transformer 生成高达4K分辨率图像的能力。OpenSora 展示了扩散 Transformer 生成视频的能力。

Vision Transformer Compression

注意力机制的计算开销受到了广泛关注。FlashAttention 将输入标记划分为更小的块，以最小化冗余的内存访问并优化延迟。一些研究强调了注意力计算的二次复杂度，并通过在网络的各个阶段过滤 或合并（Lu等人，2023；Huang等人，2023；Wu等人，2023）标记来提高效率。DynamicViT（Rao等人，2021）使用预测网络动态过滤标记。自适应稀疏ViT 通过同时考虑注意力值和特征L2范数来过滤标记。Lu等人（2023）训练了一个带有分割标签的网络，以指导在内容相似的区域中标记的合并操作。Huang等人（2023）在降采样标记后进行注意力计算，然后进行上采样以恢复空间分辨率。Wu等人（2023）证明，深层更适合过滤标记，而浅层更适合合并标记。

Local Attention

多项研究探讨了局部注意力模式的应用，其中每个标记在固定窗口大小内关注一组相邻标记，旨在减轻处理长序列的计算负担。局部窗口化注意力的概念最初由Beltagy等人引入于Longformer中，提出了一种与序列长度成线性关系扩展的注意力机制。Bigbird 扩展了这一思想，通过结合窗口注意力、随机注意力和全局注意力机制，在降低计算成本的同时保持长距离依赖。在计算机视觉领域，Swin Transformer 采用类似方法，将注意力计算限制在非重叠的局部窗口内，并使用在不同层间移动的窗口来有效地捕捉全局上下文。Twins Transformer 和FasterViT 采用基于窗口的注意力来提高计算效率，并通过不同的模块设计，如全局子采样注意力和分层注意力，有效地利用全局上下文。在作者的工作中，作者采用固定大小的窗口注意力来加速预训练的Diffusion Transformer模型，并引入了一种名为“带残差共享的窗口注意力”的新技术，以保持图像标记之间的长距离依赖。

Attention Sharing

GQA（Ainslie等人，2023年）将 Query 头分为组。每个 Query 保留自己的参数，而每组共享一个键和一个值，从而减少内存使用并提高效率。PSVIT 指出，ViT中不同层之间的注意力图具有显著相似性，并建议跨层共享注意力图以减少冗余计算。Deepcache 证明了在U-Net框架扩散模型中，高层特征在不同时间步之间是相似的。Deepcache Proposal 重用U-Net的高层特征并跳过中间层的计算以加速去噪过程。TGATE（Zhang等人，2024年）表明，文本条件扩散模型的跨注意力输出在几个去噪时间步后收敛到一个固定点。TGATE在输出收敛后缓存这个输出，并在剩余的去噪步骤中保持固定，以减少计算成本。在DiTFastAttn中，作者展示了注意力输出在CFG（条件流图）和步骤上的相似性。作者还考虑了在不同步骤上不同层之间相似性的差异，以在CFG和步骤上共享注意力输出。

3 Method

Overview

在本节中，作者展示了在采用 Transformer 的扩散模型推理过程中的冗余性。在去噪过程中，作者识别出三种类型的冗余，如图2所示：（1）空间维度上的冗余。（2）注意力输出中相邻步骤之间的相似性。（3）注意力输出中条件推理与无条件推理之间的相似性。为了解决这些冗余，作者提出了三种压缩技术，如图2所示：（1）在3.2节中，作者引入了带残差共享的窗口注意力以减少空间冗余。（2）在3.3节中，作者引入了跨时间步的注意力共享来利用步骤间的相似性，从而提高模型效率。（3）在3.4节中，作者引入了跨CFG的注意力共享，通过利用条件生成与无条件生成之间的相似性来减少冗余。在3.5节中，作者介绍了一种简单的贪心方法来决定压缩计划，即选择每个层和步骤适当的压缩技术。

Window Attention with Residual Sharing (WA-RS)

作者可以观察到，在许多预训练的DiTs的 Transformer 层中，注意力的空间局部性。如图3(a)所示，注意力的值在注意力矩阵的对角线区域内集中在一个窗口内。因此，对于某些层来说，用固定大小的窗口注意力替换全注意力可以在推理过程中保留注意力矩阵中的大部分值。通过仅在指定窗口内计算注意力值，可以大大降低注意力的计算成本。

然而，一些标记仍然关注一组空间上距离较远的标记。丢弃这些依赖关系会负面影响模型性能。仅使用窗口注意力来缓解这个问题需要一个大窗口尺寸来捕获这些依赖关系。因此，这种方法在计算成本上的降低非常小，从而阻碍了加速努力。

为窗口注意力缓存并重用残差。 为了解决上述问题，作者研究了使用窗口注意力所造成的信息损失。如图3(a)所示，在全注意力和窗口注意力输出之间的残差在各个步骤中变化很小，与窗口注意力的输出不同。这一观察促使作者在一个步骤中缓存窗口注意力和全注意力的残差，并在后续步骤中重用它。

图3(b)说明了WA-RS的计算方法：在每一步中，对于每个窗口注意力层，作者计算窗口注意力并将从前一步缓存的残差添加到输出中。作者将共享残差值的步骤集合表示为，步骤的全注意力表示为，步骤的窗口注意力表示为。对于集合中的第一步，WA-RS的计算如下：

对于集合中的后续步骤，WA-RS的计算如下：

Attention Sharing across Timesteps (AST)

去噪过程中扩散模型的序列性质是推理速度的一个主要瓶颈。在这里，作者比较了去噪过程中不同步骤的注意力输出。作者发现对于某些层，特定步骤的注意力输出与相邻步骤的输出显示出显著的相似性。图4(a)展示了不同步骤间注意力输出的余弦相似度。作者可以得出两个主要观察结果：(1)注意力输出之间存在明显的时间相似性；(2)这种相似性在步骤和层之间有所不同。

为了利用这种相似性来降低计算成本，作者提出了AST技术。具体来说，对于一组注意力输出彼此相似的步骤，作者缓存最早步骤的注意力输出并重新使用，从而跳过后续步骤的计算。

Attention Sharing across CFG (ASC)

无分类器引导（CFG）被广泛用于条件生成。在条件生成的推理过程中的每一步，CFG执行两次神经网络推理：一次带条件输入，一次不带。与无条件生成相比，这使计算成本翻倍。如图4(b)所示，对于许多层和步骤，条件和无条件神经网络评估中的注意力输出之间的高度相似。基于这一观察，作者提出了一种ASC技术，该技术在无条件神经网络评估中重用条件神经网络评估的注意力输出。

Method for Deciding the Compression Plan

上述技术，包括WA-RS、AST和ASC，可以在保持性能的同时有效降低计算成本。如图3和图4所示，不同层在不同时间步具有不同程度的冗余。因此，恰当决定压缩计划至关重要，即在每个时间步对每个层应用哪些技术。

作者开发了一个简单的贪心方法，从策略列表 [AST, WA-RS + ASC, WA-RS, ASC] 中为每个时间步和每个层选择合适的策略（技术的组合）。如算法1所示，作者逐步逐层确定策略。对于每个时间步和转换层，作者应用四种压缩策略中的每一种，并计算当前步骤模型输出在有压缩和无压缩情况下的损失，。然后，作者选择在损失低于阈值 的条件下具有最高计算减少比例的策略，其中 是层索引， 是模型中的层数。如果四种策略中没有一种满足阈值，作者则不对该层在该步骤应用压缩。

算法1 决定压缩计划的方法

4 Experiments

Settings

作者评估了DiTFastAttn在三种常用的扩散 Transformer 上的表现：用于图像生成任务的DiT（Peebles & Xie，2023）和Pixart-Sigma，以及用于视频生成任务的Open-Sora。为了展示与快速采样方法的兼容性，作者在DiT和Pixart-Sigma上基于20步DPM-Solver构建作者的方法，在Open-Sora上基于100步IDDPM。

作者使用ImageNet作为计算质量度量的评估数据集，以及MS-COCO 2017（Lin等人，2014）标题作为Pixart-Sigma模型图像生成的文本提示。对于ImageNet，作者生成了5k张图像来评估生成质量。遵循以前的研究，作者采用FID（Heusel等人，2017）、IS（Salimans等人，2016）和CLIP分数（Hessel等人，2021）作为评估指标。作者在单个Nvidia A100 40G GPU上测量了每个样本的延迟。

作者使用的平均相对绝对误差，并实验了不同的阈值，间隔为0.025。作者将这些阈值设置分别表示为D1（=0.025）、D2（=0.05），...，D6（=0.15）。作者将WA-RS的窗口大小设置为 Token 大小的1/8。

Results on Image Generation

评估指标结果及FLOPs数量。 DiTFastAttn被应用于预训练的DiT-2-XL-512、PixArt-Sigma-1024和PixArt-Sigma-2K模型上。图5展示了这些模型在ImageNet-5k验证数据集上的评估结果。对于DiT-2-XL-512和PixArt-Sigma-1024模型，配置D1、D2和D3几乎在IS和FID指标上与原始模型性能相匹配。设置D4、D5和D6的IS和CLIP分数略有下降，作为达到更高压缩率的折衷。从D1到D6，DiT-2-XL-512的注意力计算分别减少到88.82%、76.65%、67.28%、58.77%、50.82%和44.28%；对于PixArt-Sigma-1024，分别减少到90.49%、74.37%、62.50%、50.67%、42.03%和36.79%。PixArt-Sigma-2K模型在D5之前保持了与原始模型接近的性能，只有在D6时性能明显下降，此时注意力计算减少到88.01%、70.98%、56.65%、47.03%、39.60%和35.54%。对比这三个模型的压缩效果和评估指标可以看出，随着图像分辨率的增加，DiTFastAttn不仅实现了更大的压缩，同时也能更好地保持模型的生成性能。

DiTFastAttn生成结果的可视化。 图6展示了DiTFastAttn的图像生成样本。对于DiT-2-XL-512和PixArt-Sigma-1024模型，D1、D2和D3在视觉生成质量上与原始模型相似。尽管D4、D5和D6实现了更大的压缩，并在生成图像的细节上略有差异，但它们仍然产生了可用质量的图像。PixArt-Sigma-2K模型在D4之前保持了与原始模型相当图像质量，配置D5和D6仍然生成高质量的图像。这表明，即使在高分辨率下将注意力计算减少一半以上，压缩到36%，作者的压缩方法仍然保持了良好的生成质量。

Results on Video Generation

作者在OpenSora上应用DitFastAttn进行视频生成。结果如图10所示。具体来说，标记为D1至D6的配置计算减少分别为：4.29%，11.51%，19.55%，26.74%，32.62%，40.37%。在扩展分析中，附加结果在附录中提供。

#FLOPs减少与加速

图7：在240p分辨率下，使用OpenSora V1.1以16帧进行视频生成的比较。

在不同序列长度上DiTFastAttn的压缩结果。 作者基于FlashAttention-2（Dao, 2023）实现DiTFastAttn。表1显示了在Diffusion Transformers中，DiTFastAttn与原始注意力机制的FLOPs比例和延迟比例。ASC技术将注意力计算减少了50%，随着分辨率的增加，延迟减少略有增加。随着分辨率的增加，WA-RS可以将注意力计算从69%减少到23%，延迟减少从93%到70%。WA-RS和ASC技术是正交的；它们可以同时使用，而不会产生额外的开销。DiTFastAttn的整体延迟。 图8显示了在应用DiTFastAttn时，图像生成和注意力的延迟如何随着计算减少而变化。在DiT-XL-512模型上，在低压缩比设置中，图像生成和注意力的整体延迟略有增加。DiTFastAttn在推理过程中不引入开销。延迟增加是由于作者的 Kernel 实现不够好，导致与FlashAttention-2相比性能略有下降。在高压缩比设置中，与原始相比延迟降低，D6达到最低延迟：总体生成延迟约为原始的94.27%，总体注意力延迟为原始的82.32%。对于PixArt-Sigma-1024，随着FLOPs的减少，延迟继续降低。PixArt-Sigma-2K表现最佳，在D6时总体生成延迟为原始的76.32%，总体注意力延迟为64.27%。随着分辨率的增加，DiTFastAttn在降低总体注意力和图像生成的延迟方面表现更佳。值得注意的是，进一步优化作者的 Kernel 实现可能会导致更好的延迟降低。

Ablation Study

DiTFastAttn优于单一方法。如图9左侧所示，在相同的计算预算下，DiTFastAttn相较于单个技术保持了更高的质量指标。在单一技术中，AST展现出最佳的生成质量。然而，当超过2.2 FLOPs时，使用AST进一步压缩会显著降低输出质量，导致搜索算法终止。DiTFastAttn在进一步压缩的同时保持了更好的质量。

更高的步骤提升了DiTFastAttn的性能。如图9中间所示，作者比较了不同步骤下DiTFastAttn的性能。很明显，随着步骤的增加，DiTFastAttn可以在保持质量的同时压缩更多的计算。

残差缓存技术在维持性能方面至关重要。如图9右侧所示，带有残差共享的窗口注意力在相同压缩比下比窗口注意力保持了更好的生成性能。没有残差的情况下，窗口注意力会导致性能显著下降。

5 Conclusion

在本文中，作者引入了一种新颖的后训练压缩方法DiTFastAttention，以加速扩散模型。作者确定了三种类型的冗余：

(1) 空间维度上的冗余。

(2) 注意力输出中相邻步骤之间的相似性。

(3) 注意力输出中条件推理与无条件推理之间的相似性。

并且作者提出了相应的压缩技术：

(1) 带残差共享的窗口注意力，(2) 跨时间步的注意力共享，(3) 跨CFG的注意力共享。

实验表明，DiTFastAttention显著降低了注意力的成本并加快了计算速度。

局限性和未来工作。 首先，作者的方法是一种后训练压缩技术，因此无法利用训练来避免性能下降。其次，作者的贪心压缩方法选择简单，但可能找不到最优解。第三，作者的方法仅降低了注意力模块的成本。在未来的工作中，作者计划探索训练感知的压缩方法。作者还致力于将作者的方法扩展到其他模块。此外，进一步的 Kernel 级优化可能会为所提出的压缩技术解锁更大的速度提升。

参考

[1].DiTFastAttn: Attention Compression for Diffusion Transformer Models.