一个时间步内仅计算 16 个 Token:基于聚类驱动特征缓存的扩散 Transformer 加速
计算机视觉与模式识别
2025-09-15 v1
摘要
扩散 Transformer 因其生成高质量图像和视频的能力而在近年来受到广泛关注,但由于其迭代去噪过程,仍面临巨大的计算开销。最近,特征缓存被引入以加速扩散 Transformer,其通过缓存先前时间步的特征计算并在后续时间步中复用,这利用了扩散模型的时间相似性,却忽略了空间维度上的相似性。在本文中,我们引入了聚类驱动特征缓存作为先前特征缓存的一种正交且互补的视角。具体而言,ClusCa 在每个时间步对 token 进行空间聚类,在每个簇中仅计算一个 token 并将其信息传播至所有其他 token,从而能够将 token 数量减少 90% 以上。在 DiT、FLUX 和 HunyuanVideo 上的大量实验证明了其在文本到图像和文本到视频生成中的有效性。此外,它可以直接应用于任何扩散 Transformer 而无需训练。例如,ClusCa 在 FLUX 上实现了 4.96 倍的加速,ImageReward 得分为 99.49%,超过原始模型 0.51%。代码可在 https://github.com/Shenyi-Z/Cache4Diffusion 获取。
引用
@article{arxiv.2509.10312,
title = {Compute Only 16 Tokens in One Timestep: Accelerating Diffusion Transformers with Cluster-Driven Feature Caching},
author = {Zhixin Zheng and Xinyu Wang and Chang Zou and Shaobo Wang and Linfeng Zhang},
journal= {arXiv preprint arXiv:2509.10312},
year = {2025}
}
备注
11 pages, 11 figures; Accepted by ACM MM2025; Mainly focus on feature caching for diffusion transformers acceleration