通过空间聚类只计算每组中的一个令牌，大幅加速扩散Transformer。
Compute Only 16 Tokens in One Timestep: Accelerating Diffusion Transformers with Cluster-Driven Feature Caching
- 对每步的令牌进行空间聚类，仅计算每簇一个代表令牌。
- 可减少90%以上计算量，实现4.96倍加速，图像质量损失极小。
- 无需训练即可适配任意扩散Transformer，适合高效生成场景。
扩散Transformer在图像与视频生成中表现优异，但其迭代去噪过程带来巨大计算开销。现有特征缓存方法利用时间相似性，却忽略了空间相似性。本文提出空间聚类驱动的特征缓存（ClusCa），在每一步对令牌进行空间聚类，仅计算每个簇中的一个代表性令牌，并将其信息传播至该簇所有其他令牌，使计算量降低超过90%。在DiT、FLUX和HunyuanVideo上的大量实验表明，该方法在文生图与文生视频任务中均有效。例如，ClusCa在FLUX上实现4.96倍加速，图像奖励得分99.49%，优于原模型0.51%。该方法无需训练，可直接应用于任意扩散Transformer。代码已开源：https://github.com/Shenyi-Z/Cache4Diffusion。
原文摘要 · Abstract (English)
Diffusion transformers have gained significant attention in recent years for their ability to generate high-quality images and videos, yet still suffer from a huge computational cost due to their iterative denoising process. Recently, feature caching has been introduced to accelerate diffusion transformers by caching the feature computation in previous timesteps and reusing it in the following timesteps, which leverage the temporal similarity of diffusion models while ignoring the similarity in the spatial dimension. In this paper, we introduce Cluster-Driven Feature Caching (ClusCa) as an orthogonal and complementary perspective for previous feature caching. Specifically, ClusCa performs spatial clustering on tokens in each timestep, computes only one token in each cluster and propagates their information to all the other tokens, which is able to reduce the number of tokens by over 90%. Extensive experiments on DiT, FLUX and HunyuanVideo demonstrate its effectiveness in both text-to-image and text-to-video generation. Besides, it can be directly applied to any diffusion transformer without requirements for training. For instance, ClusCa achieves 4.96x acceleration on FLUX with an ImageReward of 99.49%, surpassing the original model by 0.51%. The code is available at https://github.com/Shenyi-Z/Cache4Diffusion.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。