本篇分享 AAAI 2025 论文 SiTo: Training-Free and Hardware-Friendly Acceleration for Diffusion Models via Similarity-based Token Pruning，上海交通大学提出 SiTo，通过基于相似性的令牌剪枝为扩散模型提供无需训练且硬件友好的加速。
论文链接：
https://www.researchgate.net/publication/387204421_Training-Free_and_Hardware-Friendly_Acceleration_for_Diffusion_Models_via_Similarity-based_Token_Pruning
项目链接：
https://github.com/EvelynZhang-epiclab/SiTo
介绍
1.1 SiTo：加速扩散模型的创新解决方案
先前的研究者们通常通过减少采样步数或压缩去噪网络等方式来降低扩散模型的计算开销。然而，这些方法往往在一定程度上牺牲了生成质量。SiTo 提出了一种全新的思路，它通过引入基础令牌（Base Token）概念，自适应地剪去冗余的令牌，从而在保证高质量生成结果的同时实现显著加速。
1.2 SiTo 的三大核心创新
（I）最大相似性：基础令牌与所有其他令牌具有最高相似性，剪枝令牌则选择与基础令牌相似度最高的令牌，从而最小化恢复误差。
（II）均匀空间分布：基础令牌在图像的不同局部区域均匀分布，相邻令牌的相似性较高，用空间相邻的令牌恢复剪枝令牌更为合适，避免了误差集中。
（III）带有随机性的选择：为避免相邻时间步中令牌的重复剪枝，加入高斯噪声使基础令牌选择带有随机性，减少了剪去相同令牌的比例，从而有效避免了不平衡剪枝导致的质量下降。
如下图（b）所示，由于剪枝令牌是通过直接复制其最相似的基础令牌来恢复的，这些剪枝令牌往往会在所有后续的去噪步骤中保持较高的相似性，因此它们很可能在几乎所有的时间步中都被剪枝。这种极度不平衡的令牌剪枝可能会导致生成质量的显著下降。
为了解决这个问题，提出在不同令牌的相似度上添加高斯噪声，在基础令牌选择过程中引入随机性。如下图（a）所示，这种方法将相邻两个时间步中剪去相同令牌的比例从 97% 降低到 72%，有效避免了不平衡剪枝问题。
1.3 SiTo 的显著优势
无损加速，内存压缩
SiTo 的应用效果显著，尤其在 COCO30K 和 ImageNet 数据集上，展示了强大的加速效果和内存压缩能力。例如，在 Stable Diffusion v1.5 中，应用 SiTo 方法后，加速达 1.9 倍，同时实现了 2.7 倍的内存压缩，并且 FID 值降低了 1.33，不仅加速了推理过程，还提升了生成质量。
硬件友好，低计算开销
SiTo 操作低成本、低内存占用，适合在 GPU 等硬件平台上并行计算，特别适用于边缘设备和实时应用。
训练无关，数据无依赖
SiTo 不依赖训练或校准数据，免去额外训练步骤，具有极强的通用性和即插即用性，可广泛应用于不同数据集、模型和采样设置。
方法
SiTo 的方法流程
基础令牌选择：
计算所有令牌之间的余弦相似度。对于每个令牌，将它与所有其他令牌的相似度求和，得到该令牌的相似度得分（SimScore）。接着，向相似度得分中添加高斯噪声，以引入随机性，防止在不同时间步中选择相同的基础令牌和剪枝令牌。最后，选择在图像区域中具有最高噪声相似度得分（Noise SimScore）的令牌作为基础令牌。
剪枝令牌选择：
选择与基础令牌最相似的令牌作为剪枝令牌。
剪枝令牌恢复：
将未剪枝的令牌输入到神经网络层，然后通过从最相似的基础令牌中复制来恢复剪枝令牌。
结果
视觉效果：SiTo 保留更多的图像细节，更加对齐文本提示，以及和原始模型生成的结果更吻合。
数值结果：在 ImageNet 和 COCO30K 数据集上的结果。
消融
4.1 剪枝令牌选择方法的消融实验
对以下六种基础令牌选择方法进行了消融实验：
始终选择 2x2 图块中的左上角令牌
在全局范围内随机选择 25% 的令牌
在全局范围内选择具有最高 SimScore 的令牌
在 2x2 图块中随机选择一个令牌
在 2x2 图块中选择具有最高 SimScore 的令牌
向 SimScore 添加高斯噪声后，选择 2x2 图块中具有最高得分的令牌
如下表所示，实验得出了两个关键结论：
时间分布均匀性
基础令牌应该在不同的去噪时间步之间保持均匀分布。策略 I，即每个时间步始终选择相同的令牌，表现最差。策略 VI 优于策略 V，因为在策略 VI 中引入的随机噪声使得不同时间步的基础令牌选择有所变化。类似的推理解释了为什么策略 II 优于策略 III。
此外，还对策略 V 和策略 VI 进行了视觉分析。如下图 a 所示，引入噪声可以防止某些令牌在多个时间步中被反复剪枝，从而使得图像呈现出更丰富的细节。
空间分布均匀性
基础令牌还应保持空间上的均匀分布。策略 IV、策略 V 和策略 VI 的表现优于策略 II 和策略 III，这表明在局部图块内选择基础令牌要比全局选择效果更好。
4.2 剪枝 vs. 合并
为了验证的剪枝操作，将剪枝操作替换为类似于 ToMeSD 的平均合并策略。如图 5a 所示，剪枝操作始终能够获得更低的 FID 得分。
4.3 不同图块大小的影响
如图 5b 所示，尽管较大的图块大小仅能带来微小的加速提升，但它们会导致 FID 值的显著增加。
4.4 使用较少采样时间步的性能评估
SiTo 在不同的时间步数下，始终在图像质量和加速效果上超过 ToMeSD，这表明 SiTo 与 DDIM 等扩散加速方法是正交的。
4.5 SiTo 在各模块中的应用
SiTo 可以应用于所有模块，包括自注意力（Self-Attention，SA）、交叉注意力（Cross-Attention，CA）和前馈网络（Feed-Forward Network，FFN）。下表的结果显示，将 SiTo 应用于 CA 和前馈网络 FFN 仅带来了很小的加速收益，同时影响了图像生成质量。
此外，尽管 SiTo 和 ToMeSD 可以应用于更深的 UNet 块，下图显示，应用这些方法会导致生成质量的下降。
总结
SiTo，一种用于高效扩散模型的令牌剪枝方法，旨在通过基于相似性的选择来减少剪枝误差。SiTo 不需要训练，硬件友好，不仅能显著降低内存和计算成本，还能提高生成质量。
它能够无缝地集成到现有的工作流程中，具有很好的跨模型和数据集的适应性，同时揭示了预训练扩散模型中的冗余，为更高效的模型设计提供了新的思路。
更多阅读
#投 稿 通 道#
让你的文字被更多人看到
如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。
总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。
📝 稿件基本要求：
• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注
• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题
• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道：
• 投稿邮箱：hr@paperweekly.site
• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿
△长按添加PaperWeekly小编
🔍
现在，在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧
·
·
·
扫一扫
784
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
