复旦联合字节提出BlockDance！DiT加速新方法！保证生成质量同时加速50%！

🌐 社群导航

🔗 点击加入➡️【AIGC/LLM/MLLM/3D/自动驾驶】 技术交流群

🔗 点击加入➡️【2-3月CCF投稿】交流群

数源AI 最新论文解读系列

论文名：BlockDance: Reuse Structurally Similar Spatio-Temporal Features to Accelerate Diffusion Transformers

论文链接：https://arxiv.org/pdf/2503.15927

导读

由于扩散模型令人印象深刻的能力，它们已被公认为图像和视频生成任务的关键进展。最近，人们越来越有兴趣将扩散模型的架构从U - Net转向变压器（Transformer）架构，即DiT。这种改进的架构使这些模型不仅能够生成视觉上可信且艺术上引人注目的图像和视频，还能更好地遵循缩放定律。

简介

扩散模型已展现出令人印象深刻的生成能力，特别是近期借助变压器（Transformer）架构的进展，提升了视觉和艺术质量。然而，扩散变压器模型（Diffusion Transformers，DiTs）由于迭代去噪过程，仍然面临推理速度慢的挑战。为解决这一问题，我们提出了BlockDance，这是一种无需训练的方法，通过探索相邻时间步的特征相似性来加速DiTs。与之前缺乏针对不同尺度特征的定制复用策略的特征复用方法不同，BlockDance优先识别结构最相似的特征，即结构相似的时空（Structurally Similar Spatio - Temporal，STSS）特征。这些特征主要位于去噪后期变压器模型中注重结构的模块内。BlockDance缓存并复用这些高度相似的特征，以减少冗余计算，从而在最大程度上保证与原始模型生成结果一致性的同时加速DiTs。此外，考虑到生成内容的多样性和冗余特征的不同分布，我们引入了BlockDance - Ada，这是一种专为特定实例加速而设计的轻量级决策网络。BlockDance - Ada能够动态分配资源并提供更优的内容质量。BlockDance和BlockDance - Ada在各种生成任务和模型中均被证明是有效的，在保持生成质量的同时实现了25%至50%的加速。

方法与模型

1. 预备知识

扩散模式中的正向和反向过程 扩散模型逐渐向数据中添加噪声，然后学习逆转这个过程，从噪声中生成所需的无噪声数据。在本文中，我们关注[37]引入的在潜在空间中进行噪声添加和去噪的公式。在正向过程中，时间步时噪声潜在变量的后验概率有一个封闭形式：

其中和表示噪声方差调度。推理过程，即从噪声中生成数据的反向过程，是扩散模型框架的关键部分。一旦扩散模型训练完成，在反向过程中，传统采样器DDPM[14]会逐步对进行去噪，总共进行步。也可以使用像DDIM[43]这样更快的采样器，通过以下过程来加速采样过程：

在去噪过程中，模型在早期阶段生成图像的粗略结构，并在后期阶段通过添加纹理和详细信息逐步对其进行细化。

Transformer中的特征。去噪步骤的数量与DiT架构中的网络推理次数有关，该架构通常具有多个堆叠在一起的模块。每个模块根据前一个模块的输入依次计算其输出。靠近输入的浅层模块倾向于捕捉数据的全局结构和粗略轮廓。相比之下，靠近输出的深层模块会逐步细化特定细节，以实现既逼真又美观的输出。

2. DiTs中的特征相似度和冗余

DiTs的推理速度受到其固有的迭代推理性质的限制。本文旨在减少冗余计算以加速DiTs。

图3. PixArt - 模块的可视化。带有蓝色边框的特征图在相邻步骤中表现出高度相似性。

在重新审视各种扩散变压器（DiT）模型（包括DiT-XL/2 [32]、PixArt-α [4]和Open-Sora [57]）的去噪过程时，出现了两个关键发现：一、连续步骤之间存在显著的特征相似性，这表明去噪过程中存在冗余计算，如图2 (a)所示；二、这种高度相似性主要体现在变压器的浅层和中层块（0到20块之间），而深层块（21到27块之间）表现出更多的变化，如图2 (b)所示。我们将这种现象归因于以下事实：结构内容通常在初始步骤中生成，而纹理和细节则在后续步骤中生成。

为了证实这一点，我们使用主成分分析（PCA）对PixArt-α的块特征进行可视化，如图3所示。在去噪的初始阶段，网络主要专注于生成结构内容，如人体姿势和其他基本形态。随着去噪过程的推进，网络的浅层和中层块仍然专注于生成低频结构内容，而深层块则将重点转向生成更复杂的高频纹理信息，如景深中的云朵和人群。因此，在结构建立之后，图3中蓝色框突出显示的特征图在相邻步骤中表现出高度一致性。我们将这种计算定义为冗余计算，它与变压器浅层和中层块所关注的低级结构有关。基于这些观察结果，我们认为分配大量计算资源来重新生成这些相似特征收益甚微，但会导致更高的计算成本。因此，我们的目标是设计一种策略，利用这些高度相似的特征来减少冗余计算并加速去噪过程。

3. 免训练加速方法

我们提出了BlockDance，这是一种简单而有效的方法，通过利用去噪过程中步骤之间的特征相似性来加速扩散变压器（DiTs）。通过策略性地缓存高度相似的结构特征并在后续步骤中重用它们，我们减少了冗余计算。

图4. BlockDance概述。重用步骤基于缓存步骤中的结构特征生成x_{t - 1}，节省前n个块的计算以加速推理。

具体来说，我们将去噪步骤设计为两种类型：缓存步骤和重用步骤，如图4所示。在连续的时间步中，缓存步骤首先基于x_t进行标准的网络前向传播，输出x_{t - 1}，并保存第n个块的特征F_n。对于接下来的时间步——重用步骤——我们不进行完整的网络前向计算；相反，我们进行部分推理。更具体地说，我们将缓存步骤中缓存的特征F_n用作重用步骤中第n个块的输入。因此，由于变压器块的顺序推理特性，重用步骤中前n个块的计算可以节省下来，只有比第n个块更深的块需要重新计算。

为此，确定最佳的块索引以及应集中进行重用的去噪过程阶段至关重要。基于图2和图3的见解，我们将索引设置为20，并将重用集中在去噪过程的后T'步，即结构稳定之后。这些设置使得特征重用能够解耦，并专门重用结构相似的时空特征。因此，我们将前T - T'个去噪步骤设置为缓存步骤，并将剩余的T'个去噪步骤均匀地分成几个组，每个组包含k个步骤。每个组的第一步被指定为缓存步骤，而后续的k - 1个步骤为重用步骤，以加速推理。随着新组的到来，新的缓存步骤更新缓存的特征，然后在该组内的重用步骤中使用这些特征。这个过程会重复进行，直到去噪过程结束。较大的k表示更高的重用频率。我们将这种缓存和重用策略称为BlockDance-k，它在免训练范式下运行，可以在保持生成内容质量的同时有效加速多种类型的扩散变压器（DiTs）。

4. 特定实例加速方法

图6. 去噪过程中生成图像特征的余弦相似度。由红色框突出显示的高相似度步骤特征随着结构复杂性的增加而减少。

然而，如图6所示，生成的内容呈现出不同的特征相似度分布。我们可视化了每个去噪步骤中块索引的特征与其他步骤的特征的余弦相似度矩阵。我们发现相似特征的分布与生成内容的结构复杂性有关。在图6中，随着结构复杂性从左到右增加，适合重用的相似特征数量减少。为了改进BlockDance策略，我们引入了BlockDance-Ada，这是一个轻量级框架，用于学习特定实例的缓存和重用策略。

图5. BlockDance-Ada概述。给定中间潜在特征和提示嵌入，重用决策网络学习每个样本的结构复杂性，并得出相应的重用策略。这些策略决定了DiTs中的每个后续步骤是缓存步骤还是重用步骤。奖励函数平衡了图像质量和速度之间的权衡。

BlockDance-Ada利用强化学习来处理不可微的二元决策，即顺序门控，以动态确定去噪过程中的每个步骤是缓存步骤还是重用步骤。如图5所示，对于总共个去噪步骤，给定潜在特征和提示的文本嵌入，我们最初执行步正常计算，即缓存步骤，以获得中间潜在特征。状态空间定义为和，决策模型中的动作被定义为确定后续步中的每一步是缓存步骤还是重用步骤。更正式地说，由参数化的决策网络学习特征相似度的分布，然后将其映射到向量：

这里，中的每个元素被归一化到范围内，表示执行缓存步骤的可能性。我们定义了一个具有维伯努利分布的重用策略：

其中是基于的动作，表示第步是缓存步骤，中的零元素表示重用步骤。在训练期间，是通过从相应策略中采样产生的，在测试时使用贪心方法。通过这种方法，DiTs根据重用策略生成潜在特征，然后解码器将潜在特征解码为像素级图像。

基于此，我们设计了一个奖励函数，以激励在保证质量的同时最大化计算节省。该奖励函数由两个主要部分组成：图像质量奖励和计算奖励，用于平衡生成质量和推理速度。对于图像质量奖励，我们使用质量奖励模型[52]，根据视觉美学和对提示的遵循程度对生成的图像进行评分，即。计算奖励定义为归一化的重用步骤数，由以下公式给出：

最后，整体奖励函数形式化为，其中调节图像质量的重要性。此时，可以优化决策网络以最大化预期奖励：。我们使用策略梯度方法[45]来学习的参数，预期梯度推导如下：

我们使用小批量样本计算预期梯度，并将公式6近似为：

其中是小批量中的样本数量。然后使用Adam [16]优化器将梯度反向传播以训练。通过这个训练过程，决策网络能够感知特定实例的缓存和重用策略，从而实现高效的动态推理。

实验与结果

1. 实验细节

1.1. 模型、数据集和评估指标

我们对类别条件图像生成、文本到图像生成和文本到视频生成进行了评估。对于类别条件图像生成，我们使用DiT - [32]，通过DDIM采样器[43]在ImageNet [6]数据集上为每个类别生成张图像，引导系数为4.0。对于文本到图像生成，我们使用PixArt - ，通过DPMSolver采样器[24]在COCO2017 [23]的验证集上生成张图像，引导系数为4.5。对于文本到视频生成，我们使用Open - Sora 1.0 [57]，通过DDIM采样器在MSR - VTT [51]的2990测试集上生成分辨率为的16帧视频，引导系数为7.0。我们遵循先前的工作来评估这些任务，并报告文本到图像生成的IQS分数[52]和Pickscore [17]。我们通过生成延迟（即每张图像/视频的生成时间）在A100 GPU上测量推理速度。

1.2. 实现细节

对于BlockDance，PixArt - 中的缓存和重用步骤主要发生在去噪过程的和之间，而在DiT - XL/2和Open - Sora中，它们主要发生在去噪过程的和之间。使用这三种模型生成每个内容时，缓存特征的大小分别为18MB、4.5MB和72MB。默认块索引设置为20。对于BlockDance - Ada，我们将决策网络设计为一个轻量级架构，由三个Transformer块和一个多层感知器组成。决策网络的参数总量为0.08B。我们将设置为去噪步骤总数的。奖励函数中的参数设置为2。对于PixArt - ，我们从其训练数据集中选取10000个样本子集，以16为批量大小对步骤选择网络进行100个轮次的训练。我们使用学习率为的Adam优化器。

2. 主要结果

2.1. 免训练范式实验

加速PixArt - 的文本到图像生成。如表1所示，在 COCO2017验证集上的结果证明了BlockDance的有效性。我们将ToMe [3]和DeepCache [26]扩展应用到PixArt - 作为基线方法。对于ToMe，我们通过合并操作移除25%的标记来降低计算成本。对于DeepCache，我们在整个去噪过程中每隔2步重用特征，具体来说，重用PixArt - 中28个块里前14个块的输出。使用时，BlockDance将PixArt - 加速了，且在视觉美感和遵循提示方面，图像质量均无显著下降。不同的速度 - 质量权衡可以通过进行调节。

与ToMe相比，无论重用频率  如何，BlockDance始终明显优于ToMe。这可归因于DiTs（去噪变压器，Denoising Transformer）的架构比基于U - Net的架构更依赖注意力机制，因此在DiTs中持续使用令牌合并会加剧质量下降。与DeepCache相比，BlockDance通过关注高相似性特征，在速度相当的情况下，在所有图像指标上都取得了更好的性能。我们特别减少了去噪后期阶段的冗余结构计算，避免了不相似特征的重用，并将图像质量损失降至最低。然而，DeepCache在整个去噪过程中都重用特征，并且没有专门针对高度相似的特征进行重用。这导致重用集中包含了不相似的特征，从而导致结构失真和提示对齐度下降。与通过减少交叉注意力计算中的冗余来加速的TGATE [55] 相比，BlockDance支持不包含交叉注意力的DiTs，如SD3 [8] 和Flux [18]。此外，实验结果表明，在相同的加速效果下，BlockDance在各种指标上都优于TGATE。与通过一致性蒸馏训练获得的PixArt - LCM [5] 相比，虽然BlockDance需要更多的推理时间，但在无需额外训练的情况下，在多个指标上实现了更高的生成质量。值得注意的是，由于我们有针对性的重用策略，BlockDance生成的图像与基础模型的一致性更高，

加速DiT/XL - 2进行类别条件生成。 ImageNet图像的结果如表2所示。我们将ToMe和DeepCache扩展到DiT/XL - 2作为基线。在时，BlockDance不仅能在保持图像质量的同时将DiT/XL - 2加速37.4%，而且在与基础模型生成的图像保持更高一致性的情况下，性能优于DeepCache。通过增加（相关参数），BlockDance的加速比可高达57.5%，并且图像质量始终优于ToMe。

加速Open - Sora的文本到视频生成。BlockDance在加速视频生成任务方面也很有效。在MSR - VTT上的加速结果如表3所示。在时，BlockDance在保持视频质量（包括视觉质量和时间一致性）的同时，将Open - Sora的速度提高了。增加可以实现不同的质量 - 速度权衡。相比之下，DeepCache存在明显的质量下降问题，如FVD等指标的恶化就证明了这一点。这是因为DeepCache重用了低相似度的特征，比如去噪早期的结构信息。

定性结果。如图7所示，我们进一步对我们的方法进行了定性分析。ToMe合并相邻的相似标记以节省自注意力计算，但这种方法对依赖Transformer架构的模型不太友好，会导致生成低质量且带有“块状伪影”的图像。虽然DeepCache和TGATE实现了约的加速，但在一些复杂情况下，它们可能会导致与原始图像存在明显的结构差异，并出现伪影和语义损失。PixArt - LCM通过额外的一致性蒸馏训练加速了PixArt - ，实现了显著的加速，但在视觉美感和遵循提示方面有明显下降。相比之下，BlockDance在无需额外训练成本的情况下实现了25.4%的加速，在结构和细节上与原始图像保持了高度一致。

图7. 定性结果。BlockDance具有高保真度，并与原始图像保持一致。

2.2. BlockDance - Ada评估

PixArt - 的动态推理。表4详细列出了BlockDance - Ada的性能。通过基于特定实例策略为每个样本动态分配计算资源，BlockDance - Ada有效减少了冗余计算，在实现接近BlockDance 的加速的同时，提供了更优的图像质量。与BlockDance 相比，BlockDance - Ada在质量相近的情况下提供了更大的加速优势。

3. 讨论

3.1. 消融研究

重用频率的影响。如图8所示，我们展示了生成图像如何随着的增加而演变。随着生成时间的减少，图像的主要主体保持一致，但细节的保真度逐渐降低，这与表1中的见解一致。的不同值为各种速度 - 质量权衡提供了灵活的选择。

在不同去噪阶段使用BlockDance的效果。如图9所示，我们研究了在去噪过程的不同阶段应用BlockDance的影响：初始阶段和后期阶段。BlockDance主要重用结构特征；因此，在关注结构的初始阶段应用它可能会导致结构变化或伪影（如图9 (a)中的红色框所示），因为此时结构尚未稳定。相反，在后期阶段，结构信息已经稳定，重点转向纹理细节，重用结构特征可以在质量损失最小的情况下加速推理，如图9 (b)所示。

在不同深度重用块的影响。如图10所示，我们研究了仅重用浅层和中层块与在Transformer中也重用深层块的影响。由于

图10. 在不同深度重用块的影响。

深层块中特征的相似度较低，重用它们会导致与细节相关的计算丢失，从而导致纹理细节下降，如图10 (b)中的红色框所示。相反，重用相似度较高的浅层和中层块（这些块关注结构信息），质量下降最小。

结论

在本文中，我们介绍了BlockDance，这是一种新颖的免训练加速方法，通过在结构稳定后缓存和重用结构级特征，即结构相似的时空特征，BlockDance在质量损失最小的情况下显著加速了DiTs，并与基础模型保持高度一致。此外，我们还介绍了BlockDance - Ada，它通过根据特定实例的重用策略动态分配计算资源来增强BlockDance。