【Efficient AIGC】SiTo: Similarity-based Token Pruning (AAAI-2025)

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)所示，由于剪枝令牌是通过直接复制其最相似的基础令牌来恢复的，这些剪枝令牌往往会在所有后续的去噪步骤中保持较高的相似性，因此它们很可能在几乎所有的时间步中都被剪枝。这种极度不平衡的令牌剪枝可能会导致生成质量的显著下降。为了解决这个问题，提出在不同令牌的相似度上添加高斯噪声，在基础令牌选择过程中引入随机性。如