2025-11-17T17:13:12.426721

Cached Adaptive Token Merging: Dynamic Token Reduction and Redundant Computation Elimination in Diffusion Model

Saghatchian, Moghadam, Nickabadi

Diffusion models have emerged as a promising approach for generating high-quality, high-dimensional images. Nevertheless, these models are hindered by their high computational cost and slow inference, partly due to the quadratic computational complexity of the self-attention mechanisms with respect to input size. Various approaches have been proposed to address this drawback. One such approach focuses on reducing the number of tokens fed into the self-attention, known as token merging (ToMe). In our method, which is called cached adaptive token merging(CA-ToMe), we calculate the similarity between tokens and then merge the r proportion of the most similar tokens. However, due to the repetitive patterns observed in adjacent steps and the variation in the frequency of similarities, we aim to enhance this approach by implementing an adaptive threshold for merging tokens and adding a caching mechanism that stores similar pairs across several adjacent steps. Empirical results demonstrate that our method operates as a training-free acceleration method, achieving a speedup factor of 1.24 in the denoising process while maintaining the same FID scores compared to existing approaches.

academic

Cached Adaptive Token Merging: Dynamic Token Reduction and Redundant Computation Elimination in Diffusion Model

基本信息

论文ID: 2501.00946
标题: Cached Adaptive Token Merging: Dynamic Token Reduction and Redundant Computation Elimination in Diffusion Model
作者: Omid Saghatchian, Atiyeh Gh. Moghadam, Ahmad Nickabadi (Amirkabir University of Technology)
分类: cs.CV (Computer Vision)
发表时间: 2025年1月1日 (arXiv预印本)
论文链接: https://arxiv.org/abs/2501.00946
代码链接: https://github.com/omidiu/ca_tome

摘要

扩散模型已成为生成高质量、高维图像的有前途方法。然而，这些模型受到高计算成本和慢推理速度的阻碍，部分原因是自注意力机制相对于输入大小的二次计算复杂度。本文提出了缓存自适应token合并(CA-ToMe)方法，通过计算token之间的相似性并合并相似度大于阈值参数t的token来解决这一问题。由于相邻步骤中观察到的重复模式和相似性频率的变化，该方法通过实现自适应阈值和添加缓存机制来增强token合并方法。实验结果表明，该方法作为免训练加速方法，在去噪过程中实现了1.24倍的加速，同时保持与现有方法相同的FID分数。

研究背景与动机

问题定义

扩散模型在图像生成任务中表现出色，但面临严重的计算效率问题：

高计算成本：自注意力机制的二次复杂度导致推理速度慢
串行去噪过程：无法并行化，每个去噪步骤都需要重复计算
冗余计算：相邻时间步之间存在大量重复计算

问题重要性

扩散模型的高延迟限制了其在需要快速推理的应用中的使用
计算成本高昂使得模型部署困难，特别是在资源受限的环境中
现有加速方法要么需要重新训练，要么在质量上有显著损失

现有方法局限性

减少采样步数的方法通常需要重新训练或使用复杂的数值求解器
Token剪枝方法会导致信息丢失和性能下降
**传统Token合并(ToMe)**使用固定合并率，无法适应不同时间步和层的相似性分布变化

研究动机

基于观察到的两个关键现象：

不同时间步和层中token相似性分布存在显著变化
相邻推理步骤之间的token对表现出高度相似性

核心贡献

提出自适应阈值机制：根据token相似性分布动态调整合并策略，替代固定合并率
设计缓存机制：利用相邻步骤间的相似性，缓存token对以减少重复计算
实现免训练加速：方法可直接应用于预训练模型，无需重新训练
达到更好的质量-速度权衡：相比基线ToMe方法，在保持图像质量的同时实现更快的推理速度

方法详解

任务定义

输入：扩散模型去噪过程中的token序列输出：经过自适应合并和缓存优化的加速推理过程约束：保持生成图像质量不显著下降

模型架构

1. 自适应Token合并机制

传统ToMe方法使用固定比例r进行token合并，而CA-ToMe引入相似性阈值t：

核心思想：

将图像分割为大小为sx × sy的步长区域
选择每个步长区域的左上角token作为目标token
计算源token与目标token间的余弦相似度
仅合并相似度超过阈值t的token对

优势分析：

场景A：当大多数token相似度较低时，固定合并率会强制合并不相似的token，导致信息丢失。自适应阈值确保只合并高相似度token
场景B：当大多数token高度相似时（如去噪初期），固定合并率限制了合并数量。自适应阈值允许合并更多token，提高效率

2. 缓存机制设计

基于Jaccard距离分析发现相邻步骤间token对的高相似性：

$JaccardDistance(A_n, A_{n+1}) = 1 - \frac{|A_n \cap A_{n+1}|}{|A_n \cup A_{n+1}|}$

其中An表示第n步的所有源-目标token对集合。

实现策略：

设置检查点(checkpoints)，仅在特定时间步计算相似性矩阵
在非检查点步骤重用之前计算的token对
显著减少相似性矩阵的重复计算开销

技术创新点

动态自适应性：根据相似性分布自动调整合并策略，避免固定参数的局限性
时间维度优化：利用时间步间的冗余性，通过缓存减少计算量
层级选择性应用：专门针对计算密集的U-Net顶层（D1和U1）应用优化
无需重训练：作为即插即用的加速方法，可直接应用于现有模型

实验设置

数据集

ImageNet-1k数据集：生成2000张512×512分辨率图像（每类2张，共1000类）
验证集：使用5000张ImageNet-1k验证图像计算FID分数
提示模板："A high-quality photograph of a classname."

评价指标

FID (Fréchet Inception Distance)：衡量生成图像质量的主要指标
推理时间：生成2000张图像的平均时间
PSNR：峰值信噪比，衡量像素级重建质量
SSIM：结构相似性指数，评估空间和结构一致性

对比方法

Baseline：原始Stable Diffusion v1.5
ToMe：传统token合并方法(r=50%)

实现细节

硬件：Tesla V100S GPU
扩散步数：50步PLMS采样
CFG scale：7.5
步长大小：固定为2×2
应用层：仅在U-Net的D1和U1层应用

实验结果

主要结果

模型	FID	平均时间(s)	加速比
Baseline	33.66	7.61±0.001	1.0×
ToMe	34.16	6.39±0.006	1.19×
CA-ToMe	34.05	6.09±0.001	1.24×

关键发现：

CA-ToMe实现了最快的推理速度(6.09s)
FID分数(34.05)优于ToMe(34.16)，接近baseline(33.66)
在速度和质量间达到了最佳平衡

阈值参数分析

阈值t	FID	平均时间(s)	PSNR	SSIM
0.4	35.28	6.07±0.007	27.90	0.191
0.5	35.46	6.07±0.004	27.909	0.208
0.6	35.56	6.10±0.005	27.908	0.218
0.7	34.30	6.23±0.002	27.910	0.234
0.8	33.80	6.58±0.004	27.904	0.239
0.9	33.42	6.92±0.003	27.907	0.238

观察结果：

阈值0.4-0.6范围内变化较小，因为大多数token相似度≥0.6
阈值0.7提供了最佳的质量-速度权衡
更高阈值提升质量但降低速度

缓存配置对比

配置	检查点设置	时间(s)	FID
CONFIG 1	0,1,2,3,5,10,15,25,35	6.18±0.02	36.14
CONFIG 2	0,10,11,12,15,20,25,30,35,45	6.13±0.001	34.33
CONFIG 3	0,8,11,13,20,25,30,35,45,46,47,48,49	6.09±0.001	34.05