AAAI2025|免訓練低成本！上海交大提出用於高效擴散模型的令牌剪枝方法

本篇分享 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 剪枝令牌選擇方法的消融實驗

對以下六種基礎令牌選擇方法進行了消融實驗：

始終選擇 2×2 圖塊中的左上角令牌
在全域性範圍內隨機選擇 25% 的令牌
在全域性範圍內選擇具有最高 SimScore 的令牌
在 2×2 圖塊中隨機選擇一個令牌
在 2×2 圖塊中選擇具有最高 SimScore 的令牌
向 SimScore 新增高斯噪聲後，選擇 2×2 圖塊中具有最高得分的令牌

如下表所示，實驗得出了兩個關鍵結論：

時間分佈均勻性

基礎令牌應該在不同的去噪時間步之間保持均勻分佈。策略 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 僅帶來了很小的加速收益，同時影響了影像生成質量。