2025-11-23T17:13:17.428108

Token Pruning for Caching Better: 9 Times Acceleration on Stable Diffusion for Free

Zhang, Xiao, Tang et al.
Stable Diffusion has achieved remarkable success in the field of text-to-image generation, with its powerful generative capabilities and diverse generation results making a lasting impact. However, its iterative denoising introduces high computational costs and slows generation speed, limiting broader adoption. The community has made numerous efforts to reduce this computational burden, with methods like feature caching attracting attention due to their effectiveness and simplicity. Nonetheless, simply reusing features computed at previous timesteps causes the features across adjacent timesteps to become similar, reducing the dynamics of features over time and ultimately compromising the quality of generated images. In this paper, we introduce a dynamics-aware token pruning (DaTo) approach that addresses the limitations of feature caching. DaTo selectively prunes tokens with lower dynamics, allowing only high-dynamic tokens to participate in self-attention layers, thereby extending feature dynamics across timesteps. DaTo combines feature caching with token pruning in a training-free manner, achieving both temporal and token-wise information reuse. Applied to Stable Diffusion on the ImageNet, our approach delivered a 9$\times$ speedup while reducing FID by 0.33, indicating enhanced image quality. On the COCO-30k, we observed a 7$\times$ acceleration coupled with a notable FID reduction of 2.17.
academic

Token Pruning for Caching Better: 9 Times Acceleration on Stable Diffusion for Free

基本信息

  • 论文ID: 2501.00375
  • 标题: Token Pruning for Caching Better: 9× Acceleration on Stable Diffusion for Free
  • 作者: Evelyn Zhang, Bang Xiao, Jiayi Tang, Qianli Ma, Chang Zou, Xuefei Ning, Xuming Hu, Linfeng Zhang
  • 分类: cs.CV (Computer Vision), cs.LG (Machine Learning)
  • 发表时间: 2024年12月31日
  • 论文链接: https://arxiv.org/abs/2501.00375
  • 代码链接: github.com/EvelynZhang-epiclab/DaTo

摘要

Stable Diffusion在文本到图像生成领域取得了显著成功，但其迭代去噪机制带来了高计算成本和缓慢的生成速度。尽管特征缓存等方法因其有效性和简单性而受到关注，但简单地重用先前时间步计算的特征会导致相邻时间步的特征变得相似，降低特征随时间的动态性，最终影响生成图像的质量。本文提出了一种动态感知的令牌剪枝(DaTo)方法来解决特征缓存的局限性。DaTo选择性地剪枝具有较低动态性的令牌，仅允许高动态令牌参与自注意力层，从而在时间步间扩展特征动态性。在ImageNet上应用于Stable Diffusion时，该方法实现了9×加速，同时FID降低了0.33；在COCO-30k上观察到7×加速，FID显著降低了2.17。

研究背景与动机

问题背景

扩散模型在生成建模领域取得了重大进展，广泛应用于文本到图像生成、视频生成等任务。然而，扩散模型的迭代去噪机制导致了巨大的计算成本和缓慢的生成速度，限制了其更广泛的应用。

现有方法的局限性

当前加速扩散模型的方法主要包括：

  1. 减少采样步数：如DDIM等快速采样器
  2. 降低每步计算成本：包括知识蒸馏、结构剪枝、量化、令牌剪枝和特征缓存

其中，特征缓存因其有效性和简单性而广受欢迎，它存储先前时间步计算的特征并在后续时间步中重用。然而，特征重用强制不同时间步的特征具有相似值，降低了特征沿时间步的动态性，损害了原始扩散过程，从而降低了生成质量。

研究动机

论文通过实验观察发现，与原始Stable Diffusion相比，使用特征缓存的模型在相邻时间步间的特征差异显著降低。这引发了一个关键问题：是否可以在执行特征缓存的同时仍然保持正确的特征动态性？

核心贡献

  1. 提出了动态感知令牌剪枝(DaTo)方法：通过剪枝在不同时间步中动态性被特征缓存降低的令牌，并用具有大动态性的令牌恢复它们，避免了特征缓存导致的生成质量下降。
  2. 设计了进化搜索策略：提出通过进化方法搜索最优的特征缓存和令牌剪枝策略，充分释放DaTo的潜力。
  3. 实现了显著的性能提升：在Stable Diffusion和SDXL上的广泛实验表明，无需训练和额外数据，在Stable Diffusion上可获得高达9×的加速且生成质量无损失。

方法详解

任务定义

本文的任务是在保持图像生成质量的前提下，显著加速Stable Diffusion模型的推理过程。输入为文本提示，输出为对应的高质量图像，约束条件是无需重新训练模型。

模型架构

1. 动态感知令牌剪枝(DaTo)

基础令牌选择：

  • 时间噪声差异分数：对于第t个时间步，计算前两个相邻时间步输出的绝对差异： 
    DiffScore = (1/C) * Σ|f_up_0(x_{t+2})_c - f_up_0(x_{t+1})_c|
  • 基于补丁的令牌选择：将图像划分为s×s的非重叠补丁，在每个补丁中选择具有最高DiffScore的令牌作为基础令牌。

CFG对齐： 为了处理分类器自由引导(CFG)，将条件生成的基础令牌位置复制到无条件生成中：

X_base,i,j[k] = X_base,i,j[k - B/2], k ∈ {B/2, B/2+1, ..., B-1}

剪枝令牌选择： 基于余弦相似度选择与基础令牌最相似的K个令牌进行剪枝：

X_prune = arg topK max Cosine_Similarity(X_i, X_j)

剪枝令牌恢复： 将剪枝的令牌通过直接复制其最相似的基础令牌来恢复。

2. 时间步感知特征缓存

搜索空间剪枝：

  • 缓存深度d限制为{0, 1, 1/2}
  • 剪枝比率r限制为{0.3, 0.4, 0.5, 0.6, 0.7}

进化搜索算法： 使用NSGA-II多目标优化算法，优化目标包括：

  • 推理延迟
  • 生成质量(FID)

搜索过程包括选择、交叉、变异等标准进化操作，最终获得最优的步骤感知策略F(t)。

技术创新点

  1. 动态性恢复机制：通过选择性剪枝低动态令牌并用高动态令牌恢复，成功恢复了被特征缓存破坏的特征动态性分布。
  2. 统一的缓存-剪枝框架：将特征缓存和令牌剪枝结合在一个训练无关的框架中，实现了时间和令牌级别的信息重用。
  3. 自适应策略搜索：针对不同时间步的不同冗余特性，提出自动搜索最优缓存深度和剪枝比率的方法。

实验设置

数据集

  • ImageNet-1k：生成2000张512×512图像(每类2张)
  • COCO-30k：生成30000张图像(每个标题1张)
  • MS COCO验证集：用于SDXL评估，生成5k张1024×1024图像

评价指标

  • FID (Fréchet Inception Distance)：衡量生成质量
  • CLIP Score：评估文本-图像对齐
  • Inception Score：图像质量评估
  • 延迟和加速比：效率评估

对比方法

  • DDIM/DPM：快速采样器
  • ToMeSD：令牌合并方法
  • DeepCache：特征缓存方法
  • DeepCache & ToMeSD：朴素组合方法

实现细节

  • 使用NSGA-II进化算法，种群大小20，运行100代
  • CFG scale: 7.5 (SD v1.5), 9.0 (SD v2), 7.0 (SDXL)
  • 采样步数：50步PLMS
  • 在单个4090 GPU上测试

实验结果

主要结果

Stable Diffusion v1.5 (ImageNet)：

  • 配置e1：9.01×加速，FID从27.64降至27.31
  • 在所有配置下均优于对比方法

Stable Diffusion v2 (ImageNet)：

  • 配置e2：7.25×加速，FID为28.20
  • 相比原始模型FID从29.8降至28.20

COCO-30k数据集：

  • SD v1.5：7×加速，FID从12.15降至9.98(降低2.17)
  • SD v2：7.25×加速，FID从13.68降至13.88

SDXL (MS COCO)：

  • 2.32×加速，FID从24.25降至23.10
  • 显著优于DeepCache(1.75×)和DeepCache&ToMeSD(1.78×)

消融实验

DiffScore的有效性： 在不同缓存设置和剪枝比率下，使用DiffScore consistently提高了FID分数，证明了时间噪声差异分数的有效性。

CFG对齐的影响： 随着剪枝比率的增加，CFG对齐配置带来的收益逐渐增大，在高剪枝比率(0.7)下FID改善从13个点到30个点不等。

案例分析

视觉比较结果显示DaTo在多个方面表现优异：

  • 内容保真度：与原始图像内容高度相似
  • 细节保留：在高细节场景中保持精细纹理
  • 风格适应：在图像到图像任务中平衡内容保存与风格准确性
  • 提示对齐：准确生成复杂文本提示中的所有元素

实验发现

  1. 特征动态性恢复：DaTo成功将特征差异分布恢复到接近原始Stable Diffusion的水平
  2. 稀疏编码效应：适度的令牌剪枝和特征缓存可以通过专注于关键特征来改善模型性能
  3. 策略泛化性：在SD v1.5上搜索的策略在SDXL和其他数据集上表现良好

相关工作

高效扩散模型

  • 采样步数减少：DDIM、一致性模型等
  • 网络压缩：量化、剪枝、蒸馏等
  • 架构优化：U-Net改进、Transformer优化等

令牌减少策略

  • 学习方法：DynamicViT、A-ViT等使用辅助模型排序和剪枝
  • 启发式方法：Token Pooling、Token Merging等训练无关方法
  • 扩散模型应用：ToMeSD、AT-EDM等针对生成任务的适配

缓存机制

  • U-Net缓存：DeepCache利用时间冗余缓存特征
  • DiT缓存：Δ-DiT针对Diffusion Transformer的缓存策略
  • 优化挑战：平衡效率提升与生成质量保持

结论与讨论

主要结论

  1. DaTo成功解决了特征缓存导致的特征动态性损失问题
  2. 通过进化搜索获得的自适应策略显著优于固定配置
  3. 该方法在多个模型和数据集上实现了显著的加速和质量提升

局限性

  1. 搜索成本：虽然≤20 GPU小时可接受，但仍需额外计算资源
  2. 硬件依赖：性能提升可能因硬件配置而异
  3. 极端设置下的限制：过高的剪枝比率或过低的缓存更新频率会损害性能

未来方向

  1. 自适应策略学习：开发更智能的自适应缓存和剪枝策略
  2. 其他架构适配：扩展到更多扩散模型架构
  3. 理论分析：深入理解稀疏编码原理在扩散模型中的作用

深度评价

优点

  1. 创新性强：首次系统性地解决特征缓存的动态性损失问题
  2. 方法实用：训练无关，易于部署和集成
  3. 实验充分：在多个模型、数据集上进行了全面评估
  4. 理论支撑：提供了稀疏编码的理论解释
  5. 开源友好：提供了完整的代码实现

不足

  1. 理论分析不够深入：对为什么该方法能改善FID的理论解释相对简单
  2. 搜索算法依赖：需要进化搜索来找到最优策略，增加了使用复杂度
  3. 评估指标单一：主要依赖FID评估，缺乏更多样化的质量指标
  4. 用户研究缺失：没有进行人类评估来验证生成质量

影响力

  1. 学术价值：为扩散模型加速提供了新的思路和方法
  2. 实用价值：可直接应用于现有Stable Diffusion模型
  3. 可复现性：提供了详细的实现细节和开源代码
  4. 启发性：为令牌级优化在生成模型中的应用提供了范例

适用场景

  1. 资源受限环境：移动设备、边缘计算场景
  2. 实时应用：需要快速生成的交互式应用
  3. 批量生成：大规模图像生成任务
  4. 研究原型：需要快速迭代的研究项目

参考文献

论文引用了46篇相关文献，涵盖了扩散模型、令牌减少、缓存机制等多个相关领域的重要工作，为本研究提供了坚实的理论基础和对比基准。

总体评价：这是一篇高质量的计算机视觉论文，在扩散模型加速这一重要问题上提出了创新性解决方案。方法设计巧妙，实验评估全面，实用价值突出。尽管在理论分析深度上有所不足，但其实际贡献和影响力值得认可。