SiTo: Training-Free and Hardware-Friendly Acceleration for Diffusion Models via Similarity-based Token Pruning

Paper: Training-Free and Hardware-Friendly Acceleration for Diffusion Models via Similarity-based Token Pruning
Code: https://github.com/EvelynZhang-epiclab/SiTo/tree/main

背景

随着生成式模型，尤其扩散模型（Diffusion Models）的飞速发展，它们在图像生成、艺术创作、超分辨率等领域展现了巨大的潜力。然而，这些模型通常伴随着极高的计算成本，导致在边缘设备和实时应用中的部署面临巨大挑战。针对这一问题， SiTo（Similarity-based Token Pruning）方法应运而生，它不仅大幅度加速了扩散模型的推理速度，还在不牺牲生成质量的情况下，显著降低了计算资源的消耗。

介绍

SiTo：加速扩散模型的创新解决方案
传统上，研究者们通过减少采样步数或压缩去噪网络等方式来降低扩散模型的计算开销。然而，这些方法往往在一定程度上牺牲了生成质量。SiTo提出了一种全新的思路，它通过引入基础令牌 (Base Token) 概念，自适应地剪去冗余的令牌，从而在保证高质量生成结果的同时实现显著加速。

SiTo的三大核心创新
(I) 最大相似性：基础令牌被选为一组与所有其他令牌具有最高相似性的令牌。此外，在选择基础令牌后，进一步选择与基础令牌相似度最高的令牌作为剪枝令牌。这个选择基础令牌和剪枝令牌的策略确保了高相似性，即基础令牌和剪枝令牌之间的差异较小，从而进一步最小化从基础令牌恢复剪枝令牌时的误差。

(II) 均匀空间分布：经典的图像分析研究表明，图像同一区域的图块携带相似的信息，这意味着在空间维度上相邻的令牌可能具有相似的表示，因此，用空间相邻的令牌来恢复剪枝令牌比使用远离的令牌更合适。因此，并不是直接从图像中的所有令牌中选择基础令牌，而是在图像的每个局部区域中选择一个基础令牌，以保证基础令牌在图像的不同空间位置上均匀分布。此外，由于基础令牌不会被剪枝，这种方案还确保了每个区域至少有一个令牌没有被剪枝，从而避免了令牌剪枝引入的误差过于集中。

(III) 带有随机性的选择：前两条原则有效地最小化了单个去噪步骤中令牌剪枝
带来的恢复误差。然而，扩散模型的采样过程包含多个去噪步骤，相邻的步骤具有相似的令牌表示。因此，基础令牌和剪枝令牌的选择在相邻时间步中可能非常相似，甚至是相同的。此外，如下图(b)所示，由于剪枝令牌是通过直接复制其最相似的基础令牌来恢复的，这些剪枝令牌往往会在所有后续的去噪步骤中保持较高的相似性，因此它们很可能在几乎所有的时间步中都被剪枝。这种极度不平衡的令牌剪枝可能会导致生成质量的显著下降。为了解决这个问题，提出在不同令牌的相似度上添加高斯噪声，在基础令牌选择过程中引入随机性。如下图(a)所示，这种方法将相邻两个时间步中剪去相同令牌的比例从97%降低到72%，有效避免了不平衡剪枝问题。

SiTo的显著优势

1. 无损加速，内存压缩
   SiTo的应用效果显著，尤其在COCO30K和ImageNet数据集上，展示了强大的加速效果和内存压缩能力。例如，在Stable Diffusion v1.5中，应用SiTo方法后，加速达1.9倍，同时实现了2.7倍的内存压缩，并且FID值降低了1.33，不仅加速了推理过程，还提升了生成质量。
2. 硬件友好，低计算开销
   SiTo方法采用的操作非常低成本，不会增加额外的内存占用，且计算复杂度极低，非常适合在GPU等硬件平台上并行计算。这使得SiTo在边缘设备和实时互动应用中的应用具有广泛的前景，尤其是对于需要快速响应的应用场景，SiTo无疑是一种理想的加速方案。
3. 训练无关，数据无依赖
   与传统的加速方法不同，SiTo完全不依赖训练数据或校准数据，它不需要额外的训练步骤和调整，因此具有极强的通用性和即插即用性。这一特点使得SiTo可以在不同的数据集、不同的模型以及不同的采样设置下自由应用，极大地拓展了其实际应用的范围。

方法

SiTo的方法流程：

1. 基础令牌选择：
   计算所有令牌之间的余弦相似度。对于每个令牌，将它与所有其他令牌的相似度求和，得到该令牌的相似度得分（SimScore）。接着，向相似度得分中添加高斯噪声，以引入随机性，防止在不同时间步中选择相同的基础令牌和剪枝令牌。最后，选择在图像区域中具有最高噪声相似度得分（Noise SimScore）的令牌作为基础令牌。
2. 剪枝令牌选择：
   选择与基础令牌最相似的令牌作为剪枝令牌。
3. 剪枝令牌恢复：
   将未剪枝的令牌输入到神经网络层，然后通过从最相似的基础令牌中复制来恢复剪枝令牌。

结果

• 视觉效果：SiTo 保留更多的图像细节，更加对齐文本提示，以及和原始模型生成的结果更吻合
  
• 数值结果

（1）在ImageNet数据集上的评估
将ToMeSD和SiTo应用于SD v1.5和SD v2。表1显示，SiTo在所有加速比设置下相比ToMeSD具有更低的FID得分、更高的加速比和更低的内存使用。在剪枝比率为0.7的情况下，SiTo在硬件上实现了1.9倍的加速和2.70倍的内存压缩，同时FID得分减少了1.33。

（2）在COCO30K数据集上的评估
为了进一步验证SiTo的有效性，在COCO30K数据集上进行了实验。如表1所示，SiTo在SD v1.5上实现了1.75倍的加速，并取得了11.17的FID得分；在SD v2上，SiTo实现了1.65倍的加速，FID得分为11.67，在加速和图像质量方面始终优于ToMeSD。

消融

---

剪枝令牌选择方法的消融实验

对以下六种基础令牌选择方法进行了消融实验：

1. 始终选择$2\times 2$图块中的左上角令牌
2. 在全局范围内随机选择25%的令牌
3. 在全局范围内选择具有最高SimScore的令牌
4. 在$2\times 2$图块中随机选择一个令牌
5. 在$2\times 2$图块中选择具有最高SimScore的令牌
6. 向SimScore添加高斯噪声后，选择$2\times 2$图块中具有最高分数的令牌（SiTo的默认设置）

如下表所示，实验得出了两个关键结论：

1. 时间分布均匀性：
   基础令牌应该在不同的去噪时间步之间保持均匀分布。策略I，即每个时间步始终选择相同的令牌，表现最差。策略VI优于策略V，因为在策略VI中引入的随机噪声使得不同时间步的基础令牌选择有所变化。类似的推理解释了为什么策略II优于策略III。
   此外，还对策略V和策略VI进行了视觉分析。分析结果表明，引入噪声可以防止某些令牌在多个时间步中被反复剪枝，从而使得图像呈现出更丰富的细节。例如，下图a中显示，使用策略VI生成的图像草地纹理更为精细，花朵的数量也有所增加。
   
2. 空间分布均匀性：
   基础令牌还应保持空间上的均匀分布。策略IV、策略
   V和策略VI的表现优于策略II和策略III，这表明在局部图块内选择基础令牌要比全局选择效果更好。这可能是因为全局选择会导致基础令牌集中在某些区域，从而在其他区域造成密集剪枝，导致信息损失难以恢复。

---

剪枝 vs. 合并
为了验证的剪枝操作，将剪枝操作替换为类似于ToMeSD的平均合并策略。如图5a所示，剪枝操作始终能够获得更低的FID得分。较高的令牌剪枝比率迫使基础令牌重建多个令牌，导致显著的误差和与真实分布的偏差，从而降低了图像质量。

---

不同图块大小的影响
在选择基础令牌的过程中，将特征划分为大小为 $s\times s$的图块，在每个图块内选择一个令牌作为基础令牌。随着图块大小的增大，基础令牌的比例相应减少。如图5b所示，尽管较大的图块大小仅能带来微小的加速提升，但它们会导致FID值的显著增加。

---

使用较少采样时间步的性能评估
评估了在较少采样时间步下，ToMeSD和SiTo的表现。如下图所示，SiTo在不同的时间步数下，始终在图像质量和加速效果上超过ToMeSD，这表明SiTo与DDIM等扩散加速方法是正交的。

---

SiTo在各模块中的应用

SiTo可以应用于所有模块，包括自注意力（Self-Attention）、交叉注意力（Cross-Attention）和前馈网络（Feed-Forward Network）。探讨了SiTo在各模块中的应用，以优化加速和图像质量之间的平衡。下表的结果显示，将SiTo应用于交叉注意力（CA）和前馈网络（MLP）仅带来了最小的加速收益，同时影响了图像生成质量。

此外，尽管SiTo和ToMeSD可以应用于更深的UNet块，下图显示，应用这些方法会导致生成质量的下降。