2025-01-24 12:31 江苏
关注公众号，发现CV技术之美
本篇分享 AAAI-2025 论文SiTo: Training-Free and Hardware-Friendly Acceleration for Diffusion Models via Similarity-based Token Pruning，上海交通大学提出 SiTo，通过基于相似性的令牌剪枝为扩散模型提供无需训练且硬件友好的加速。
介绍
SiTo：加速扩散模型的创新解决方案
先前的研究者们通常通过减少采样步数或压缩去噪网络等方式来降低扩散模型的计算开销。然而，这些方法往往在一定程度上牺牲了生成质量。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的方法流程
基础令牌选择：
计算所有令牌之间的余弦相似度。对于每个令牌，将它与所有其他令牌的相似度求和，得到该令牌的相似度得分（SimScore）。接着，向相似度得分中添加高斯噪声，以引入随机性，防止在不同时间步中选择相同的基础令牌和剪枝令牌。最后，选择在图像区域中具有最高噪声相似度得分（Noise SimScore）的令牌作为基础令牌。
剪枝令牌选择：
选择与基础令牌最相似的令牌作为剪枝令牌。
剪枝令牌恢复：
将未剪枝的令牌输入到神经网络层，然后通过从最相似的基础令牌中复制来恢复剪枝令牌。
结果
视觉效果
SiTo 保留更多的图像细节，更加对齐文本提示，以及和原始模型生成的结果更吻合
数值结果
在ImageNet 和COCO30K数据集上的结果
消融
1.剪枝令牌选择方法的消融实验
对以下六种基础令牌选择方法进行了消融实验：
始终选择2x2图块中的左上角令牌
在全局范围内随机选择25%的令牌
在全局范围内选择具有最高SimScore的令牌
在2x2图块中随机选择一个令牌
在2x2图块中选择具有最高SimScore的令牌
向SimScore添加高斯噪声后，选择2x2图块中具有最高得分的令牌
如下表所示，实验得出了两个关键结论：
时间分布均匀性
基础令牌应该在不同的去噪时间步之间保持均匀分布。策略I，即每个时间步始终选择相同的令牌，表现最差。策略VI优于策略V，因为在策略VI中引入的随机噪声使得不同时间步的基础令牌选择有所变化。类似的推理解释了为什么策略II优于策略III。
此外，还对策略V和策略VI进行了视觉分析。如下图a所示，引入噪声可以防止某些令牌在多个时间步中被反复剪枝，从而使得图像呈现出更丰富的细节。
空间分布均匀性
基础令牌还应保持空间上的均匀分布。策略IV、策略V和策略VI的表现优于策略II和策略III，这表明在局部图块内选择基础令牌要比全局选择效果更好。
剪枝 vs. 合并
为了验证的剪枝操作，将剪枝操作替换为类似于ToMeSD的平均合并策略。如图5a所示，剪枝操作始终能够获得更低的FID得分。
不同图块大小的影响
如图5b所示，尽管较大的图块大小仅能带来微小的加速提升，但它们会导致FID值的显著增加。
使用较少采样时间步的性能评估
SiTo在不同的时间步数下，始终在图像质量和加速效果上超过ToMeSD，这表明SiTo与DDIM等扩散加速方法是正交的。
SiTo在各模块中的应用
SiTo可以应用于所有模块，包括自注意力（Self-Attention，SA）、交叉注意力（Cross-Attention,CA）和前馈网络（Feed-Forward Network,FFN）。下表的结果显示，将SiTo应用于CA和前馈网络FFN仅带来了很小的加速收益，同时影响了图像生成质量。
此外，尽管SiTo和ToMeSD可以应用于更深的UNet块，下图显示，应用这些方法会导致生成质量的下降。
总结
SiTo，一种用于高效扩散模型的令牌剪枝方法，旨在通过基于相似性的选择来减少剪枝误差。SiTo不需要训练，硬件友好，不仅能显著降低内存和计算成本，还能提高生成质量。
它能够无缝地集成到现有的工作流程中，具有很好的跨模型和数据集的适应性，同时揭示了预训练扩散模型中的冗余，为更高效的模型设计提供了新的思路。
张林峰（上海交通大学）投稿
最新 AI 进展报道
请联系：amos@52cv.net
END
欢迎加入「剪枝」交流群?备注：剪枝
