ICLR 2025 | 无需训练的Token级 DiT加速方法

前言本文分享 ICLR 2025 论文 ToCa: Accelerating Diffusion Transformers with Token-wise Feature Caching，提出的 ToCa 模型通过 token 粒度的缓存方法，实现了图像和视频生成模型上无需训练的两倍以上的加速。

欢迎关注公众号CV技术指南，专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。

本文转载自我爱计算机视觉

仅用于学术分享，若侵权请联系删除

CV方向的准研究生们，未来三年如何度过？

招聘高光谱图像、语义分割、diffusion等方向论文指导老师

论文：https://arxiv.org/abs/2410.05317
Github：https://github.com/Shenyi-Z/ToCa
联系方式：shenyizou@outlook.com
通讯作者：zhanglinfeng@sjtu.edu.cn

摘要 Abstract

Diffusion Transformer在图像和视频生成中展现了显著的效果，但代价是巨大的计算成本。

为了解决这一问题，特征缓存方法被引入，用于通过缓存前几个时间步的特征并在后续时间步中复用它们来加速扩散Transformer。然而，之前的缓存方法忽略了不同的token对特征缓存表现出不同的敏感性，而对某些token的特征缓存可能导致生成质量整体上高达10倍的破坏，相较于其他token。

在本文中，我们提出了基于token的特征缓存方法，允许我们自适应地选择最适合进行缓存的token，并进一步使我们能够为不同类型和深度的神经网络层应用不同的缓存比率。

通过在PixArt-α、OpenSora和DiT, 以及FLUX上的广泛实验，我们证明了在图像和视频生成中无需训练即可实现我们方法的有效性。例如，在OpenSora和PixArt-α上分别实现了2.36倍和1.93倍的接近无损的生成加速。

背景 Backgrounds

扩散模型(Diffusion Models)在图像生成、视频生成等多种生成任务中展现了出色的性能。近年来，以FLUX, Sora, 可灵等模型为代表的 Diffusion Transformers 通过扩展参数量和计算规模进一步推动了视觉生成领域的发展。

然而，Diffusion Transformers面临的一个重大挑战在于其高计算成本，这导致推理速度缓慢，从而阻碍了其在实时场景中的实际应用。为了解决这一问题，研究者们提出了一系列加速方法，主要集中在减少采样步数和加速去噪网络模型。

近期，基于特征缓存来实现去噪模型加速的方法由于其优秀的无损加速性能，以及无需训练的优良性能，受到工业界的广泛关注。

上海交通大学张林峰团队进一步注意到一个自然而有趣的现象：不同计算层，以及同计算层的不同 Token 对于缓存误差的适应性不同，同样的缓存误差在不同位置对模型影响最高可以达到数十，百倍的差异，因此有必要进一步将模型加速的粒度由特征级进一步到 token 级，并考虑了如何衡量视觉生成模型中 token 的重要性，以实现重要 token 的筛选保留。

核心贡献 TL,DR

ToCa 首次在DiT加速中中引入 token 级的缓存复用策略，并首次从误差积累与传播的角度分析特征缓存方法。
ToCa 提出 4 种从不同角度出发，适用于不同情形的 token selection策略：(a)基于Self-Attention Map来评估 token对其它token的影响; (b)基于Cross-Attention Map评估文生图/视频任务中 image token 对 text token的关注分布，以加强控制能力; (c) 基于该 token在先前去噪步中的被连续缓存复用的次数设计增益策略，鼓励token在时间步上被更均匀地计算，避免局部误差积累过大，破坏全局图像; (d)将各个token的重要性得分基于空间分布进行加权，鼓励被计算的token在空间上分布更均匀。
ToCa 被应用于多种最新模型上开展实验，证明了其相比现有方法更加优秀，包含文生图模型PixArt-alpha，FLUX-dev和FLUX-schnell，文生视频模型 OpenSora，以及基于ImageNet类标签生成图像的DiT模型。

研究动机

如图1所示，不同token在相邻两步间进行特征缓存引入的误差值的差异高达几十上百倍；图2说明不同token上引入同样大小的误差，这最初幅度相同的误差在模型推理过程经过积累和传播，对模型的输出的影响差异也极大。

因此，有必要考虑token级别的特征缓存-复用策略，使得模型的计算更集中在关键被需要的token上。

方法

计算流程

ToCa的缓存-复用流程如图3(a)所示：

Cache 初始化首先推理一个完整的时间步，将各层的特征放入cache中以便使用。
重要性得分计算在使用ToCa的时间步上，对于每一层：先计算各个token的重要性得分，将最低的部分token 标记为cache状态（例如图示中ID为1和3的token），不传入网络层进行计算。
部分计算对于被传入的token(2,4,5),执行正常的计算,得到它们的输出。
Cache更新从cache中调出存储的 token 1,3的输出，并将计算得到的新的token 2,4,5输出更到cache中。

通常这样的一个循环长度为2~4个时间步，即1步充分计算后续搭配1至3个ToCa step。此外，ToCa还基于不同层的重要性，设计了随着层深度上升而衰减的计算比例，详情请参考论文。