AiPaper
论文状态:已完成

Compute Only 16 Tokens in One Timestep: Accelerating Diffusion Transformers with Cluster-Driven Feature Caching

发表:2025/09/12
原文链接PDF 下载
价格:0.10
价格:0.10
已有 5 人读过
本分析由 AI 生成,可能不完全准确,请以原文为准。

TL;DR 精炼摘要

扩散Transformer模型计算成本高。本文提出ClusCa,通过对Token进行空间聚类,每个簇仅计算一个代表性Token并传播其信息,将计算量减少超90%。无需训练即可应用于各种模型,实现高达4.96倍加速,同时保持甚至提升性能。

摘要

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.

思维导图

论文精读

中文精读

1. 论文基本信息 (Bibliographic Information)

  • 标题 (Title): Compute Only 16 Tokens in One Timestep: Accelerating Diffusion Transformers with Cluster-Driven Feature Caching (单步仅计算16个Token:利用簇驱动的特征缓存加速扩散Transformer)
  • 作者 (Authors): Zhixin Zheng, Xinyu Wang, Chang Zou, Shaobo Wang, and Linfeng Zhang. 作者主要来自上海交通大学 (Shanghai Jiao Tong University),并有电子科技大学 (University of Electronic Science and Technology of China) 和山东大学 (Shandong University) 的合作者。
  • 发表期刊/会议 (Journal/Conference): 论文提交至 33rd ACM International Conference on Multimedia (MM '25)。ACM Multimedia 是多媒体领域的顶级国际会议,具有很高的学术声誉和影响力。
  • 发表年份 (Publication Year): 2025 (根据论文信息)。
  • 摘要 (Abstract): 扩散Transformer模型在生成高质量图像和视频方面表现出色,但其迭代去噪过程带来了巨大的计算成本。现有的特征缓存技术通过利用时间维度上的相似性(即缓存并重用前一时间步的特征)来加速模型,但忽略了空间维度上的相似性。本文提出了一种名为 ClusCa (Cluster-Driven Feature Caching) 的新方法,作为对现有特征缓存技术的正交补充。ClusCa 在每个时间步对 token 进行空间聚类,然后只计算每个簇中的一个代表性 token,并将其信息传播给簇内其他 token。这种方法能将需要计算的 token 数量减少超过90%。在DiT、FLUX和HunyuanVideo等主流模型上的大量实验证明了该方法的有效性,并且它无需训练即可直接应用于任何扩散Transformer。例如,在FLUX模型上,ClusCa 实现了4.96倍的加速,同时ImageReward评分达到99.49%,甚至比原始模型还高0.51%。
  • 原文链接 (Source Link):
    • 原文链接: https://arxiv.org/abs/2509.10312
    • PDF 链接: http://arxiv.org/pdf/2509.10312v1
    • 发布状态:预印本 (Preprint)。ArXiv是一个发布未经同行评审的学术论文的平台。

2. 整体概括 (Executive Summary)

  • 研究背景与动机 (Background & Motivation - Why):

    • 核心问题: 扩散Transformer模型(如DiT)虽然生成效果拔群,但其核心的迭代去噪过程需要数百甚至上千次网络前向传播,导致推理速度极慢,计算成本高昂,难以部署在延迟敏感的应用中。
    • 现有挑战与空白 (Gap): 为了加速,现有研究主要有两个方向:减少采样步数和提升单步网络效率。在提升网络效率方面,特征缓存 (feature caching) 是一种流行的方法,它利用了去噪过程中相邻时间步 (temporal dimension) 之间特征的高度相似性。具体来说,就是缓存某个时间步计算出的特征,然后在接下来的几个时间步直接重用,从而跳过部分计算。然而,这些方法完全忽略了在单个时间步内部,不同空间位置的特征 (spatial dimension) 之间也存在的巨大冗余和相似性
    • 创新思路: 本文的切入点正是这个被忽略的空间维度相似性。作者假设,在一个时间步内,图像的不同区域(表示为 token)可以被聚类成若干组,每组内的 token 特征非常相似。因此,没有必要计算组内的每一个 token,只需计算一个代表,然后将更新后的信息传播给同组的其他成员即可。
  • 核心贡献/主要发现 (Main Contribution/Findings - What):

    • 提出了空间相似性分析: 通过实验系统性地揭示并验证了扩散Transformer中存在的空间冗余性。论文发现,在单个时间步内,聚类后同簇内的 token 不仅特征值高度相似,而且它们在去噪过程中的**变化趋势(运动模式)**也高度一致。此外,这种聚类结构在相邻时间步之间保持稳定。
    • 提出了 ClusCa 框架: 提出了一种创新的、结合了时空双重维度的特征重用范式——ClusCa。该方法通过对 token 进行聚类,仅计算极少数(例如每簇1个)代表性 token,然后将更新后的特征与历史缓存特征加权融合,从而更新所有 token。这极大地降低了单步计算量,同时有效缓解了传统缓存方法因缓存周期过长而导致的误差累积问题。
    • 实现了SOTA性能: 在多个主流图/视频生成模型(DiT, FLUX, HunyuanVideo)上取得了当前最优的效率-质量权衡。尤其在高加速比下,ClusCa 的生成质量远超其他方法,提供了一种在大幅提速的同时保持甚至提升生成质量的有效方案。

3. 预备知识与相关工作 (Prerequisite Knowledge & Related Work)

  • 基础概念 (Foundational Concepts):

    • 扩散模型 (Diffusion Models): 一类生成模型,通过两个过程工作。前向过程是逐步对一张清晰图像添加高斯噪声,直至其变为完全的随机噪声。反向过程(即生成过程)是训练一个神经网络,从纯噪声开始,逐步去除噪声,最终还原出一张清晰的图像。这个去噪过程是迭代的,每一步称为一个 timestep,总步数通常较多(如50-1000步),因此计算量大。
    • 扩散Transformer (Diffusion Transformer, DiT): 这是扩散模型中用于去噪的神经网络的一种架构。传统上常用U-Net架构,而DiT则完全基于Transformer。它将输入的带噪图片看作一系列的图像块 patch,每个 patch 作为一个 token 输入到Transformer中进行处理。DiT因其优异的可扩展性,已成为当前最先进的架构之一。
    • 特征缓存 (Feature Caching): 一种加速技术。基本思想是:在迭代计算中,如果相邻步骤的中间计算结果(即特征)变化不大,就可以把某一步的结果存起来(缓存),在接下来的一步或几步中直接使用,从而跳过这部分的计算。
    • K-Means聚类 (K-Means Clustering): 一种经典的无监督学习算法,用于将数据集中的样本点划分为 K 个簇 (cluster)。其目标是让同一个簇内的点尽可能相似(距离近),而不同簇的点尽可能不相似(距离远)。
  • 前人工作 (Previous Works):

    • 时间步优化 (Temporal Step Optimization): 这类工作旨在减少总的去噪步数,如 DDIM 提出确定性采样,DPM-Solver 使用更高阶的求解器加速收敛,Rectified Flow 通过拉直轨迹实现少步生成。它们从算法层面减少迭代次数。
    • 网络效率增强 (Network Efficiency Enhancement):
      • 架构优化:剪枝 (pruning) 去掉不重要的网络权重,量化 (quantization) 使用低精度计算,token压缩 (token compression) 减少序列长度。这些方法通常需要重新训练或微调模型。
      • 特征缓存 (Feature Caching): 这是与本文最相关的工作。
        • DeepCache: 首次在U-Net架构中提出缓存高层特征。
        • FORA, Δ-DiT: 将缓存思想应用于DiT。
        • ToCa, DuCa: 实现了更细粒度的 token 级别缓存,即在缓存周期内,只更新一部分“重要”的 token
        • TaylorSeer: 将“缓存-重用”范式升级为“缓存-预测”,使用泰勒级数展开来预测未来时间步的特征,以减少误差。
  • 差异化分析 (Differentiation):

    • 核心区别: 所有之前的特征缓存方法(如 ToCa, TaylorSeer)都只关注时间维度的冗余,即单个 tokentt 时刻和 t-1 时刻的相似性。
    • 本文创新: ClusCa 首次引入并利用了空间维度的冗余,即在同一个 tt 时刻token Atoken B 可能因为代表了图像中相似的区域(如大片天空)而具有相似的特征。ClusCa 将这两种冗余结合起来,形成了一个时空混合重用 (spatiotemporal reuse) 的新范式,这是其与所有先前工作的根本区别。

4. 方法论 (Methodology - Core Technology & Implementation Details)

ClusCa 的核心思想是利用 token 在空间上的聚集性,通过只计算少数代表来近似所有 token 的更新,从而实现加速。

  • 方法原理 (Methodology Principles):

    • 观察 1: 簇内 token 的特征与动态高度一致。 如下图所示,作者通过实验发现,将同一时间步的 token 进行聚类后,簇内的 token 间距离(左图)远小于全局 token 间的距离。更重要的是,在去噪过程中,同一个簇的 token 在特征空间中的演化轨迹(右图)也基本平行,这表明它们的动态变化模式是相似的。这为“用一个代表整个簇”提供了理论依据。

      Figure 1: Visualization of similarity in two dimensions Figure 1: Visualization of similarity in two dimensions Figure 2: (a) The distributions of distance between tokens in the same cluster (1/2/3) and all the tokens in the current timestep, showing that tokens in the same cluster have significantly lower distance. (b) PCA visualization on the evolution of tokens in different clusters. Points in the same line denote the same token in different timesteps. Tokens with the same color denote tokens in the same cluster.

    • 观察 2: 聚类结构在时间上保持稳定。 如下图所示,作者发现相邻时间步的聚类结果高度相似(通过 ARI 指标衡量,值接近1表示相似度高)。这意味着不需要在每个时间步都重新聚类,只需在一个缓存周期的开始聚类一次,大大降低了聚类带来的额外开销。

  • 方法步骤与流程 (Steps & Procedures): ClusCa 的工作流程可以分解为以下几个步骤,如下图所示:

    1. 划分缓存周期 (Cache Cycle): 将总的去噪时间步划分为若干个周期,每个周期包含 NN 个时间步。
    2. 全量计算与聚类 (Full Calculation & Clustering): 在每个周期的第一个时间步(例如 tt 时刻),对所有 token 执行完整的网络前向计算。计算完成后,对输出的 token 特征(通常是最后一层)使用 K-Means 算法将其划分为 KK 个簇。
    3. 部分计算 (Partial Calculation): 在该周期的后续 N-1 个时间步(例如 t1,t2,t-1, t-2, \dots),执行以下操作:
      • 代表选取: 从每个簇中随机选择一个 token 作为代表。总共选取 KK 个代表 token
      • 代表计算: 仅对这 KK 个代表 token 执行完整的网络计算,得到它们更新后的特征。
    4. 时空特征重用 (Spatiotemporal Feature Reuse): 对于所有未被计算的 token,它们的特征更新依赖于空间重用时间重用的结合:
      • 空间重用: 使用其所在簇的代表 token 刚刚计算出的新特征。
      • 时间重用: 使用它自己在上一个时间步缓存的旧特征。
      • 这两部分信息通过一个加权和的方式进行融合,以更新当前 token 的特征。
  • 数学公式与关键细节 (Mathematical Formulas & Key Details):

    • 对于未被计算的 token xix_i (即 iIComputei \notin \mathcal{I}_{\mathrm{Compute}}),其缓存 C(xi)C(x_i) 的更新公式为: C(xi)=γμ(i)+(1γ)C(xi) C ( x _ { i } ) = \gamma \cdot \mu _ { ( i ) } + ( 1 - \gamma ) C ( x _ { i } )

      • 符号解释:
        • C(xi)C(x_i): token xix_i 的特征缓存。
        • γ\gamma: 传播率 (propagation ratio),是一个超参数,用于平衡新旧信息。
        • μ(i)\mu_{(i)}: token xix_i 所在簇中所有被计算的代表 token 的平均特征。它代表了来自当前时间步的“新”空间信息。
        • (1γ)C(xi)(1-\gamma)C(x_i): token xix_i 自身的历史缓存信息,代表了来自前一时间步的“旧”时间信息。
    • μ(i)\mu_{(i)} 的计算方式为: μ(i)=jIComputeF(xj)jICompute[Ij=Ii] \mu _ { ( i ) } = \frac { \sum _ { j \in \mathcal { I _ { \mathrm { C o m p u t e } } } } \mathcal { F } ( x _ { j } ) } { \sum _ { j \in \mathcal { I _ { \mathrm { C o m p u t e } } } } [ \mathbf { I } _ { j } = \mathbf { I } _ { i } ] }

      • 符号解释:
        • ICompute\mathcal{I}_{\mathrm{Compute}}: 被选中计算的代表 token 的索引集合。
        • F(xj)\mathcal{F}(x_j): 对 token xjx_j 的网络计算结果。
        • [Ij=Ii][\mathbf{I}_j = \mathbf{I}_i]: 艾弗森括号,如果 token jjtoken ii 属于同一个簇,则值为1,否则为0。
        • 这个公式的本质就是计算 token ii 所在簇的所有代表 token(本文中每簇只有一个代表)的新特征的平均值。

5. 实验设置 (Experimental Setup)

  • 数据集 (Datasets):
    • ImageNet: 一个大规模图像分类数据集,包含1000个类别。用于类别条件图像生成 (class-conditional image generation) 任务的评估。
    • DrawBench: 一个包含200个挑战性文本提示 (prompt) 的基准,用于评估文生图 (text-to-image) 模型的性能。
    • VBench: 一个全面的视频生成评估基准,包含946个不同的提示,用于评估文生视频 (text-to-video) 模型的性能。
  • 评估指标 (Evaluation Metrics):
    • FID (Fréchet Inception Distance):
      1. 概念定义: FID 是衡量生成图像质量和多样性的黄金标准之一。它通过比较生成图像集和真实图像集在 Inception-V3 网络某一深层特征空间中的统计分布(均值和协方差)来计算距离。FID值越低,表示生成图像的质量和多样性越接近真实图像,效果越好。
      2. 数学公式: FID(x,g)=μxμg22+Tr(Σx+Σg2(ΣxΣg)1/2) \mathrm{FID}(x, g) = \left\| \mu_x - \mu_g \right\|_2^2 + \mathrm{Tr}\left(\Sigma_x + \Sigma_g - 2(\Sigma_x \Sigma_g)^{1/2}\right)
      3. 符号解释:
        • x, g: 分别代表真实图像和生成图像的特征向量集合。
        • μx,μg\mu_x, \mu_g: 真实图像和生成图像特征向量的均值。
        • Σx,Σg\Sigma_x, \Sigma_g: 真实图像和生成图像特征向量的协方差矩阵。
        • Tr()\mathrm{Tr}(\cdot): 矩阵的迹(主对角线元素之和)。
    • sFID (spatial FID):
      1. 概念定义: FID 的一种变体,更侧重于评估图像的空间结构和布局的合理性。计算方式与FID类似,但使用了不同的特征提取网络层。sFID值越低越好。
    • ImageReward:
      1. 概念定义: 一个基于大规模人类偏好数据训练的奖励模型。它能像人类一样对生成的图片质量、美学和文本对齐度进行打分。ImageReward分数越高,表示图像越符合人类偏好,效果越好。
    • CLIP Score:
      1. 概念定义: 用于衡量生成的图像与给定的文本描述之间的语义相似度。它利用预训练的 CLIP 模型分别提取图像和文本的特征,然后计算它们之间的余弦相似度。CLIP Score越高,表示图文匹配度越好。
    • VBench Score:
      1. 概念定义: 一个综合性的视频生成评估框架,它包含了对视频质量、时间一致性、动态性、图文对齐度等多个维度的子指标。VBench总分越高,表示视频生成综合能力越强。
  • 对比基线 (Baselines):
    • 减少DDIM步数: 最简单的加速方法,直接减少采样总步数。
    • FORA, ToCa, DuCa, TeaCache: 代表了现有的各种基于时间维度的特征缓存方法。
    • TaylorSeer: 代表了最先进的基于特征预测的加速方法。

6. 实验结果与分析

  • 核心结果分析 (Core Results Analysis):

    • 文生图 (FLUX模型, Table 1): ClusCa 表现出色。例如,在约 4.96倍 FLOPs 加速下,ClusCaImageReward 达到 0.9949,不仅远超同等加速比下的其他方法,甚至超过了未加速的原始模型。这表明该方法在大幅提速的同时,能有效保持甚至提升生成质量。

      方法 高效注意力 Latency(s) ↓ Speed ↑ FLOPs(T) ↓ Speed ↑ Image Reward ↑ CLIP Score ↑
      FLUX.1[dev]: 50 steps 25.82 1.00× 3719.50 1.00× 0.9898 19.761
      ... ... ... ... ... ... ... ...
      TaylorSeer (N=5, O=2) 7.46 3.46× 893.54 4.16× 0.9864 19.406
      ClusCa (N=6, O=1, K=16) 7.10 3.63× 748.48 4.96× 0.9762 19.533
      ClusCa (N=6, O=2, K=16) 7.12 3.62× 748.48 4.96× 0.9949 19.453
    • 文生视频 (HunyuanVideo模型, Table 2): ClusCa 同样领先。在 6.21倍 FLOPs 加速下,ClusCa 依然保持了 79.60% 的 VBench 高分,在所有对比方法中实现了最佳的加速-质量权衡。

      方法 高效注意力 Latency(s) ↓ Speed ↑ FLOPs(T) ↓ Speed ↑ VBench Score(%) ↑
      Original 334.96 1.00× 29773.0 1.00× 80.66
      ... ... ... ... ... ... ...
      TaylorSeer (N=6, O=1) 79.46 4.22× 5359.1 5.56× 79.78
      ClusCa (N=5, K=32) 87.35 3.83x 5968.1 4.99x 79.99
      ClusCa (N=6, K=32) 81.64 4.10× 5373.0 5.54× 79.96
      ClusCa (N=7, K=32) 74.88 4.47× 4796.2 6.21× 79.60
    • 类别条件生成 (DiT模型, Figure 3): 如下图所示,该图直观展示了 ClusCa 的优越性。横轴是加速比,纵轴是 FIDsFID(越低越好)。可以看出,随着加速比的增加,所有方法的生成质量都会下降(曲线向上)。然而,ClusCa(深蓝色菱形线)的曲线始终位于所有其他方法的下方,且上升趋势最为平缓。这有力地证明了,尤其是在高加速比下,ClusCa 能比其他方法保持显著更优的生成质量。

      Figure 3: Comparison between existing feature caching methods and ClusCa on DiT with FID and sFID (lower is better). Figure 3: Comparison between existing feature caching methods and ClusCa on DiT with FID and sFID (lower is better).

    • 特征轨迹可视化 (Figure 9): 作者通过PCA将特征演化轨迹降维可视化。结果显示,ClusCa(青色三角线)的轨迹与未加速的原始DiT模型(深蓝色圆点线,可视为“黄金标准”)的轨迹最为吻合,而其他方法的轨迹则偏离较远。这从机理上验证了 ClusCa 能更好地逼近原始模型的去噪过程,从而有效减少误差累积。

      Figure 9: PCA visualization on the trajectories of features in different timesteps. "DiT" here denotes the original model without acceleration, which can be considered as the ground-truth trajectory. Figure 9: PCA visualization on the trajectories of features in different timesteps. "DiT" here denotes the original model without acceleration, which can be considered as the ground-truth trajectory.

  • 消融实验/参数分析 (Ablation Studies / Parameter Analysis):

    • 传播率 γ\gamma 的影响 (Figure 10): 这个实验非常关键。作者固定缓存周期 NN,只改变传播率 γ\gamma。结果显示,FID(越低越好)随 γ\gamma 的增大呈现出先下降后上升的“U型”趋势。

      • γ=0\gamma=0 时,方法退化为只更新代表 token(类似 ToCa),效果不佳。

      • γ\gamma 适当增大时,空间信息被有效传播,质量得到提升,FID 降低。

      • γ\gamma 过大时,token 自身的时间信息被过度稀释,导致质量恶化,FID 升高。

      • 这证明了时空信息融合的必要性和有效性。图中还显示 ClusCa 在不同 NN 设置下都稳定优于 SOTA 方法 TaylorSeer

        Figure 10: ClusCa with different propagation ratios. Figure 10: ClusCa with different propagation ratios.

    • 聚类数量 K 的影响 (Table 5): 实验表明,增加聚类数 KK(即增加每步计算的代表 token 数量)可以提升生成质量(FID 降低),但会牺牲一部分加速比。作者选择 K=16K=16 是在性能和效率之间做的权衡。

7. 总结与思考 (Conclusion & Personal Thoughts)

  • 结论总结 (Conclusion Summary): 本文提出了一种名为 ClusCa 的新型扩散Transformer加速框架。与以往只关注时间冗余的特征缓存方法不同,ClusCa 创新性地同时利用了空间冗余时间冗余。通过在每个缓存周期开始时对 token 进行一次空间聚类,并在后续步骤中仅计算少数代表性 token,再通过独特的时空特征融合机制更新所有 tokenClusCa 实现了在极高加速比下仍能保持优异生成质量的SOTA性能。该方法无需训练、即插即用,为高效部署大规模生成模型提供了新的有效途径。

  • 局限性与未来工作 (Limitations & Future Work):

    • 局限性:
      1. 超参数敏感性: 方法引入了新的超参数,如聚类数 KK 和传播率 γγ,这些参数的最优值可能因模型、任务或加速比的不同而变化,需要进行一定的调优。
      2. 聚类开销: 虽然作者称聚类开销很小(通常 ≤ 5%),但在一些极端追求低延迟的场景下,这部分开销仍可能需要被考虑。
      3. 聚类策略: 当前采用的是简单的 K-Means 聚类和随机代表选取,可能存在更优的聚类或代表选取策略。
    • 未来工作:
      1. 自适应聚类: 研究如何根据图像内容或去噪阶段动态地调整聚类数量 KK
      2. 更智能的代表选取: 探索基于 token 重要性或其他度量的代表选取方法,而非随机选取。
      3. 与其他加速技术结合:ClusCa 与模型剪枝、量化等技术结合,探索实现极致加速的可能性。
  • 个人启发与批判 (Personal Insights & Critique):

    • 启发: 这篇论文最大的启发在于其思维的转换。当一个领域(如特征缓存加速)的研究似乎都集中在同一个维度(时间)时,跳出来从一个被忽略的维度(空间)入手,往往能打开全新的局面。这种“升维思考”的方法论非常有价值。ClusCa 将空间和时间信息优雅地结合,而不是简单地替代,体现了很好的设计哲学。

    • 批判性思考: 论文的方法非常巧妙且有效,但其核心假设——“同簇 token 动态一致”——在何种情况下可能失效,是一个值得探讨的问题。例如,在包含大量细碎、复杂纹理的图像中,或者在视频中物体发生剧烈形变或快速运动的阶段,聚类结构的稳定性可能会下降,从而影响加速效果。此外,虽然实验结果令人印象深刻,但如果能提供更多关于“坏 केस” (failure cases) 的分析,将有助于更全面地理解该方法的适用边界。


相似论文推荐

基于向量语义检索推荐的相关论文。

暂时没有找到相似论文。