在这里插入图片描述 在这里插入图片描述 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 PruningCode: https://github.com/EvelynZhang-epiclab/SiTo/tree/main
背景 随着生成式模型，尤其扩散模型（Diffusion Models）的飞速发展，它们在图像生成、艺术创作、超分辨率等领域展现了巨大的潜力。然而，这些模型通常伴随着极高的计算成本，导致在边缘设备和实时应用中的部署面临巨大挑战。针对这一问题， SiTo（Similarity-based Token Pruning）方法应运而生，它不仅大幅度加速了扩散模型的推理速度，还在不牺牲生成质量的情况下，显著降低了计算资源的消耗。
介绍 SiTo：加速扩散模型的创新解决方案
基础令牌 (Base Token) 概念，自适应地剪去冗余的令牌，从而在保证高质量生成结果的同时实现显著加速。
SiTo的三大核心创新 (I) 最大相似性 ：基础令牌被选为一组与所有其他令牌具有最高相似性的令牌。此外，在选择基础令牌后，进一步选择与基础令牌相似度最高的令牌作为剪枝令牌。这个选择基础令牌和剪枝令牌的策略确保了高相似性，即基础令牌和剪枝令牌之间的差异较小，从而进一步最小化从基础令牌恢复剪枝令牌时的误差。
(II) 均匀空间分布 ：经典的图像分析研究表明，图像同一区域的图块携带相似的信息，这意味着在空间维度上相邻的令牌可能具有相似的表示，因此，用空间相邻的令牌来恢复剪枝令牌比使用远离的令牌更合适。因此，并不是直接从图像中的所有令牌中选择基础令牌，而是在图像的每个局部区域中选择一个基础令牌，以保证基础令牌在图像的不同空间位置上均匀分布。此外，由于基础令牌不会被剪枝，这种方案还确保了每个区域至少有一个令牌没有被剪枝，从而避免了令牌剪枝引入的误差过于集中。
(III) 带有随机性的选择 ：前两条原则有效地最小化了单个去噪步骤中令牌剪枝
带来的恢复误差。然而，扩散模型的采样过程包含多个去噪步骤，相邻的步骤具有相似的令牌表示。因此，基础令牌和剪枝令牌的选择在相邻时间步中可能非常相似，甚至是相同的。此外，如下图(b)所示，由于剪枝令牌是通过直接复制其最相似的基础令牌来恢复的，这些剪枝令牌往往会在所有后续的去噪步骤中保持较高的相似性，因此它们很可能在几乎所有的时间步中都被剪枝。这种极度不平衡的令牌剪枝可能会导致生成质量的显著下降。为了解决这个问题，提出在不同令牌的相似度上添加高斯噪声，在基础令牌选择过程中引入随机性。如下图(a)所示，这种方法将相邻两个时间步中剪去相同令牌的比例从97%降低到72%，有效避免了不平衡剪枝问题。
在这里插入图片描述 SiTo的显著优势
无损加速，内存压缩 SiTo的应用效果显著，尤其在COCO30K和ImageNet数据集上，展示了强大的加速效果和内存压缩能力。例如，在Stable Diffusion v1.5中，应用SiTo方法后，加速达1.9倍，同时实现了2.7倍的内存压缩，并且FID值降低了1.33，不仅加速了推理过程，还提升了生成质量。 硬件友好，低计算开销 SiTo方法采用的操作非常低成本，不会增加额外的内存占用，且计算复杂度极低，非常适合在GPU等硬件平台上并行计算。这使得SiTo在边缘设备和实时互动应用中的应用具有广泛的前景，尤其是对于需要快速响应的应用场景，SiTo无疑是一种理想的加速方案。训练无关，数据无依赖 与传统的加速方法不同，SiTo完全不依赖训练数据或校准数据，它不需要额外的训练步骤和调整，因此具有极强的通用性和即插即用性。这一特点使得SiTo可以在不同的数据集、不同的模型以及不同的采样设置下自由应用，极大地拓展了其实际应用的范围。方法 在这里插入图片描述 SiTo的方法流程：
基础令牌选择 ：计算所有令牌之间的余弦相似度 。对于每个令牌，将它与所有其他令牌的相似度求和，得到该令牌的相似度得分（SimScore）
高斯噪声 ，以引入随机性，防止在不同时间步中选择相同的基础令牌和剪枝令牌。最后，选择在图像区域中具有最高噪声相似度得分（Noise SimScore） 的令牌作为基础令牌 。剪枝令牌选择 ：选择与基础令牌最相似的令牌作为剪枝令牌 。剪枝令牌恢复 ：将未剪枝的令牌输入到神经网络层，然后通过从最相似的基础令牌中复制来恢复剪枝令牌。结果 视觉效果：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倍的加速
11.67 ，在加速和图像质量方面始终优于ToMeSD。
消融 剪枝令牌选择方法的消融实验
对以下六种基础令牌选择方法进行了消融实验：
向SimScore添加高斯噪声后，选择 如下表所示，实验得出了两个关键结论：
在这里插入图片描述 时间分布均匀性： 基础令牌应该在不同的去噪时间步之间保持均匀分布。策略I ，即每个时间步始终选择相同的令牌，表现最差。策略VI 优于策略V ，因为在策略VI 中引入的随机噪声使得不同时间步的基础令牌选择有所变化。类似的推理解释了为什么策略II 优于策略III 。此外，还对策略V 和策略VI
策略VI 生成的图像草地纹理更为精细，花朵的数量也有所增加。在这里插入图片描述 空间分布均匀性： 基础令牌还应保持空间上的均匀分布。策略IV 、策略
V 和策略VI 的表现优于策略II 和策略III ，这表明在局部图块内选择基础令牌要比全局选择效果更好。这可能是因为全局选择会导致基础令牌集中在某些区域，从而在其他区域造成密集剪枝，导致信息损失难以恢复。
剪枝 vs. 合并 为了验证的剪枝操作，将剪枝操作替换为类似于ToMeSD的平均合并策略。如图5a所示，剪枝操作始终能够获得更低的FID得分。较高的令牌剪枝比率迫使基础令牌重建多个令牌，导致显著的误差和与真实分布的偏差，从而降低了图像质量。
不同图块大小的影响 在选择基础令牌的过程中，将特征划分为大小为
在这里插入图片描述 使用较少采样时间步的性能评估 评估了在较少采样时间步下，ToMeSD和SiTo的表现。如下图所示，SiTo在不同的时间步数下，始终在图像质量和加速效果上超过ToMeSD，这表明SiTo与DDIM等扩散加速方法是正交的。
在这里插入图片描述 SiTo在各模块中的应用
SiTo可以应用于所有模块，包括自注意力（Self-Attention）、交叉注意力（Cross-Attention）和前馈网络（Feed-Forward Network）。探讨了SiTo在各模块中的应用，以优化加速和图像质量之间的平衡。下表的结果显示，将SiTo应用于交叉注意力（CA）和前馈网络（MLP）仅带来了最小的加速收益，同时影响了图像生成质量。
在这里插入图片描述 此外，尽管SiTo和ToMeSD可以应用于更深的UNet块，下图显示，应用于更深的Unet层会导致生成质量的下降。
在这里插入图片描述
