Not All Tokens Need 40 Steps: Heterogeneous Step Allocation in Diffusion Transformers for Efficient Video Generation

📄 arXiv: 2605.06892v1 📥 PDF

作者: Ernie Chu, Vishal M. Patel

分类: cs.CV

发布日期: 2026-05-07

备注: Project page: https://ernestchu.github.io/hsa

🔗 代码/项目: PROJECT_PAGE


💡 一句话要点

提出异构步长分配(HSA)算法,通过动态调整去噪步长实现高效视频生成

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱八:物理动画 (Physics-based Animation)

关键词: 扩散模型 视频生成 推理加速 Transformer优化 动态计算 流匹配

📋 核心要点

  1. 现有DiTs模型对所有时空Token采用统一的去噪步长,导致大量计算资源浪费在冗余的背景或静态区域上。
  2. 提出异构步长分配(HSA)算法,根据Token的运动动态分配计算预算,并结合KV缓存同步机制实现高效推理。
  3. 实验证明HSA在Wan-2和LTX-2模型上显著提升了推理效率,在大幅压缩运行时间的同时保持了高质量的生成效果。

📝 摘要(中文)

扩散Transformer(DiTs)在视频生成领域表现卓越,但其推理成本极高,原因在于标准推理对序列中所有时空Token均施加相同的去噪步长。受人类视觉对冗余运动忽略机制的启发,本文提出了异构步长分配(HSA),这是一种无需训练的推理算法,根据时空Token的运动速度动态分配去噪预算。为解决序列长度不匹配问题,HSA引入了KV缓存同步机制,使活跃Token能关注全局序列,同时跳过非活跃Token。此外,通过推导缓存欧拉更新,可在不增加模型评估的情况下推进跳过Token的潜在状态。在Wan-2和LTX-2模型上的实验表明,HSA在大幅加速(如50%和25%运行时间)下,显著优于现有的缓存方法和流匹配基线,在无需离线分析的情况下实现了更优的质量-运行时间帕累托前沿。

🔬 方法详解

问题定义:视频生成模型中,所有时空Token被同等对待,导致计算资源在静态或低运动区域被过度消耗。现有方法缺乏针对不同Token重要性的差异化计算分配策略。

核心思路:基于“并非所有Token都需要相同去噪步长”的假设,利用Token的运动速度作为重要性度量,动态分配计算资源,从而在保证全局上下文一致性的前提下降低推理开销。

技术框架:HSA算法包含三个核心模块:基于速度动态的步长分配器、用于处理序列长度不匹配的KV缓存同步机制,以及用于更新跳过Token状态的缓存欧拉更新算子。

关键创新:引入了无需训练的推理优化策略,通过KV缓存同步机制,在跳过非活跃Token的同时,确保活跃Token仍能获取完整的全局上下文信息,避免了传统剪枝方法带来的上下文丢失问题。

关键设计:利用缓存欧拉更新(Cached Euler Update)在单次操作中推进跳过Token的潜在状态,无需额外的模型前向传播,从而在保持结构完整性的同时实现显著的计算加速。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

HSA在Wan-2和LTX-2模型上进行了广泛评估,在50%和25%的运行时间预算下,显著超越了现有的缓存方法和Vanilla Flow Matching基线。该方法无需昂贵的离线分析,在保持生成质量和结构完整性的同时,成功构建了更优的质量-运行时间帕累托前沿,展现了极强的鲁棒性。

🎯 应用场景

该技术主要应用于大规模视频生成模型(如T2V和I2V)的推理加速。在计算资源受限的边缘设备或实时视频生成场景中,HSA能够显著降低推理延迟,提升生成效率,对于推动高质量视频生成模型的商业化落地和实时交互应用具有重要价值。

📄 摘要(原文)

Diffusion Transformers (DiTs) have achieved state-of-the-art video generation quality, but they incur immense computational cost because standard inference applies the same number of denoising steps uniformly to every token in the sequence. It is well known that human vision ignores vast amounts of redundant motion. Why, then, do our densest models treat every spatiotemporal token with equal priority? In this paper, we introduce Heterogeneous Step Allocation (HSA), a training-free inference algorithm that assigns varying step budgets to different spatiotemporal tokens based on their velocity dynamics. To resolve the resulting sequence-length mismatch without sacrificing global context, HSA introduces a KV-cache synchronization mechanism that allows active tokens to attend to the full sequence while entirely bypassing inactive tokens. Furthermore, we derive a cached Euler update that advances the latent states of skipped tokens in a single operation without additional model evaluations. We evaluate HSA on the Wan-2 and LTX-2 models for both text-to-video (T2V) and image-to-video (I2V) generation. Our results demonstrate that HSA significantly outperforms previous state-of-the-art caching methods and the vanilla Flow Matching baseline, especially at aggressive acceleration regimes (e.g., 50% and 25% runtimes). Crucially, HSA achieves a superior quality-runtime Pareto frontier without the need for expensive offline profiling, robustly preserving structural integrity and generation quality even under tight computational budgets. Project page: https://ernestchu.github.io/hsa