无需训练让扩散模型提速2倍，上交大提出Token级缓存方案｜ICLR‘25

OpenSora/PixArt都能提速

原标题：无需训练让扩散模型提速2倍，上交大提出Token级缓存方案｜ICLR‘25
文章来源：量子位
内容字数：5376字

上海交大团队提出ToCa：Token级缓存加速扩散Transformer模型

近日，上海交通大学等团队提出的Toca（Token-wise Caching）模型，通过token粒度的缓存方法，实现了图像和视频生成模型上无需训练的两倍以上的加速，相关论文已被ICLR 2025接收。该研究针对扩散Transformer模型（如FLUX、Sora等）计算成本高昂的问题，提出了一种创新的解决方案。

1. 背景：扩散Transformer模型的挑战

扩散模型在图像和视频生成领域表现出色，但基于扩散Transformer的模型计算成本巨大，推理速度缓慢，限制了其在实时应用中的发展。虽然已有特征缓存方法用于加速，但这些方法忽略了不同token对缓存的敏感性差异，某些token的缓存甚至可能导致生成质量大幅下降。

2. ToCa的核心贡献

ToCa首次在扩散Transformer模型加速中引入token级的缓存复用策略，并从误差积累与传播的角度分析特征缓存方法。其主要贡献包括：

1. 提出了基于token的特征缓存方法，自适应地选择最适合缓存的token，并根据不同网络层应用不同的缓存比率。
2. 设计了四种token选择策略，分别基于Self-Attention Map、Cross-Attention Map、token连续缓存次数和空间分布加权，以更有效地选择需要缓存的token。
3. 在PixArt-α、OpenSora、DiT和FLUX等多种模型上进行了广泛实验，证明了ToCa的有效性，在OpenSora和PixArt-α上分别实现了2.36倍和1.93倍的接近无损加速。

3. ToCa的工作机制

ToCa的缓存-复用流程包括三个步骤：

1. Cache初始化：推理一个完整的时间步，将各层的特征放入缓存。
2. 重要性得分计算：计算每个token的重要性得分，将重要性最低的token标记为缓存状态，不参与后续计算。
3. Cache更新：从缓存中调出被缓存token的输出，并将新计算的token输出更新到缓存中。这个循环通常持续2-4个时间步。

ToCa设计了四种从不同角度计算token重要性得分的方法，并在实际应用中进行加权求和。

4. 实验结果与分析

ToCa在图像生成模型（PixArt-α、FLUX、DiT）和视频生成模型（OpenSora）上的实验结果表明，其加速效果显著优于其他方法，并在生成质量上保持了较高的水平。例如，在OpenSora上实现了高达2.36倍的无损加速。实验还显示，ToCa在图像文本对齐和细节还原方面也具有优势。

5. 总结

ToCa作为首个基于token级缓存加速扩散Transformer模型的方法，具有良好的适配性和高效的加速效果，为扩散模型的实际应用提供了新的可能性。其无需训练的特性以及在多种模型和任务上的优异表现，使其成为一种具有前景的模型加速方案。

联系作者

文章来源：量子位
作者微信：
作者简介：追踪人工智能新趋势，关注科技行业新突破