💻 computer science

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

本文介绍了无需训练的异构步数分配（HSA）推理算法，该算法通过一种新颖的键值缓存同步机制和缓存的欧拉更新，在保持全局上下文与质量的同时，将较少的去噪步数动态分配给低速度令牌，从而加速用于视频生成的扩散变换器。

原作者： Ernie Chu, Vishal M. Patel

发布于 2026-05-12

📖 1 分钟阅读☕ 轻松阅读

原作者： Ernie Chu, Vishal M. Patel

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象你正在执导一部规模宏大、预算高昂的电影。在这部电影中，每一个角色、每一棵背景树、甚至地板上的每一粒尘埃，都是演员。在标准的扩散 Transformer（即生成视频的 AI 模型）中，导演对待所有这些“演员”的方式完全相同。

如果剧本要求进行 40 轮排练（去噪步骤）以确保场景完美，导演就会强迫背景墙壁排练 40 次，静止的树木排练 40 次，跑过屏幕的英雄也排练 40 次。

问题出在哪里？墙壁和树木几乎不动。它们不需要 40 次排练；5 次就足够了。但 AI 却浪费大量时间和能量让它们照样排练，而那位（快速移动且形态变化的）英雄也获得了同样的 40 轮。这就像给一位超级巨星演员和一个纸板人偶支付完全相同的小时工资，以换取完全相同的工作量。

解决方案：异构步数分配（HSA）

本文作者 Ernie Chu 和 Vishal Patel 提出了一种更聪明的片场运作方式，称为异构步数分配（HSA）。

将 HSA 想象成一位聪明的导演，他观察演员并说道：

“你，背景墙壁？你既无聊又静止。你只需要5 次排练。”
“你，奔跑的英雄？你动态且复杂。你需要全部40 次排练。”
“你，摇曳的树木？你需要20 次排练。”

这节省了海量时间，因为 AI 不再将能量浪费在无聊的部分上。

他们解决的两个棘手问题

你可能会想：“好吧，但如果墙壁在 5 轮后停止排练，它怎么知道第 40 轮时英雄在做什么？它们需要互相看见，才能保持在同一场景中。”

本文提出了两个巧妙的技巧来解决这个问题：

1. “幽灵记忆”技巧（KV 缓存同步）
在 AI 中，演员需要“看见”彼此以保持同步（这称为注意力）。通常，如果一个演员停止排练，它就会从房间中消失，其他演员就看不见它了。

解决方案：AI 保留墙壁和树木的“幽灵记忆”（缓存）。即使它们不再主动排练，它们的“幽灵”仍留在房间里。活跃的演员（英雄）仍然可以观察不活跃演员的“幽灵”，从而知道它们的位置。
结果：英雄可以高强度排练，同时确切知道墙壁的位置，而墙壁实际上无需进行任何工作。

2. “时间旅行”技巧（缓存欧拉更新）
当墙壁休息时，它的位置会发生什么变化？它只是冻结不动吗？

解决方案：AI 记住墙壁上次移动的时间以及它的移动速度。它使用一个简单的数学公式，将墙壁的位置“快进”到被跳过的时间段。这就像在说：“墙壁上次每秒移动 1 英寸；我只需计算出它现在的位置，而无需让它实际移动。”
结果：墙壁在后台瞬间完成更新，AI 无需为其运行复杂的模拟。

结果：更快、更便宜、效果一样好

研究人员在强大的视频生成 AI 模型（如 Wan-2 和 LTX-2）上测试了这种方法。以下是他们的发现：

速度：他们可以将制作视频的时间缩短50% 甚至 75%（即仅需原始时间的 25%）。
质量：即使削减了这么多时间，视频看起来与慢速、全速生成的版本几乎完全相同。“英雄”依然清晰锐利，“墙壁”依然坚实稳固。
对比：之前的方法试图通过一次性为所有人跳过整轮排练来加速。当尝试过快时，这些方法会导致视频崩溃（就像建筑物倒塌一样）。HSA 保持了视频的稳定性，因为它只跳过了那些不需要工作的部分。

一句话总结

这篇论文的核心在于不要以同样的方式对待每一个像素。通过认识到视频的某些部分是无聊且静止的，而另一些部分则是激动人心且动态变化的，AI 可以停止在无聊的部分上浪费时间。它利用“幽灵记忆”和“数学捷径”来保持一切同步，使我们能够在极短的时间内生成高质量的视频。

这就好比一家工厂：一种做法是用相同的精度加工汽车上的每一颗螺丝；另一种做法是将 90% 的时间花在引擎上，而仅用 10% 的时间处理固定座椅套的螺丝。汽车依然完美运行，但建造速度快得多。

技术摘要：用于高效视频生成的异构步数分配（HSA）

1. 问题陈述

扩散 Transformer（DiTs）已在视频生成领域确立了最先进的性能，但其计算成本极其高昂。标准的推理协议对序列中的每个时空令牌（token）应用相同数量的去噪步数（ $T$ ）。这种单一整体方法忽视了视觉数据固有的不对称性：背景区域、静态块和粗略结构在感知上更简单，所需的去噪步数少于细节丰富的前景物体或高运动区域。此外，与静态组件相比，人类视觉对高度动态组件中的质量退化不那么敏感。因此，对所有令牌统一应用完整的去噪预算在计算上是浪费的，且未能与感知优先级保持一致。

现有的加速方法，如步级特征缓存（例如 TeaCache），通常做出全局决策：在任何给定的迭代中，要么计算所有令牌，要么所有令牌重用缓存输出。这些方法未能利用空间或令牌级别的异构性，将序列视为单一整体块，而非具有不同动态需求令牌的集合。

2. 方法论：异构步数分配（HSA）

作者提出了异构步数分配（HSA），这是一种无需训练、模型无关的推理算法，它在预训练的 DiT 中为令牌分配特定的去噪轨迹。其核心思想是根据速度动态将令牌序列划分为组，为演化较慢的令牌分配较少的步数，同时为基准组保持完整的时间表。

2.1 令牌划分与调度

$N$ 个令牌的集合被划分为 $K$ 个不相交的组（ $G_1, \dots, G_K$ ）。每个组 $G_k$ 被分配一个步数预算 $s_k$ ，该预算必须是总步数 $T$ 的约数。

基准组（ $G_1$ ）： 分配完整预算 $s_1 = T$ 。这些令牌在每次迭代中均处于激活状态。
缩减组： 分配较小的预算（例如 $T/2, T/4$ ）。这些令牌仅在全局迭代 $i$ 满足 $i \pmod{T/s_k} = 0$ 时处于激活状态。
动态选择： 本文评估了四种将令牌分配给组的策略。基于每个令牌的速度预测 L1 相对变化来分配令牌的动态令牌选择策略（用于识别背景等演化缓慢的区域）始终表现出最佳性能。

2.2 KV 缓存同步

异构调度的主要挑战在于，在任何给定的挂钟迭代时刻，具有不同步数的令牌处于不同的噪声水平（ $\sigma$ ）。自注意力机制要求所有令牌相互关注，从而产生同步不匹配。HSA 通过KV 缓存同步机制解决这一问题：

激活令牌： 仅激活令牌子集（ $A_i$ ）执行通过 QKV 投影的前向传播。
缓存更新： 将激活令牌的新鲜键（ $K$ ）和值（ $V$ ）投影写入每层缓存，覆盖过时的条目。
全上下文注意力： 激活令牌针对完整的缓存 $K, V$ 矩阵计算查询（其中包含激活令牌的新鲜条目和被跳过令牌的过时但有效的条目）。
复杂度降低： 这将自注意力复杂度从 $O(N^2)$ 降低到 $O(|A_i| \cdot N)$ ，其中 $|A_i|$ 是第 $i$ 步的激活令牌数量。被跳过的令牌在非激活迭代中不产生查询计算、交叉注意力或前馈操作。

2.3 缓存欧拉更新

为了在不进行额外模型评估的情况下推进被跳过令牌的潜在状态，HSA 采用缓存欧拉更新。

每个令牌存储其最近一次激活步预测的速度（ $\hat{v}_n$ ）。
在每次全局迭代中，所有令牌（激活和被跳过的）均通过相同的增量欧拉步 $(\sigma_{i+1} - \sigma_i)$ 进行推进。
激活令牌使用其新计算的速度；被跳过的令牌重用其缓存的速度。
该更新被表述为对所有 $N$ 个令牌的单个张量操作，避免了分支并保持了 GPU 效率。

2.4 相位感知缓存窗口

认识到去噪的早期和晚期阶段分别对全局结构和精细细节至关重要，HSA 将缓存限制在轨迹的中间部分。前 $m$ 步和后 $m$ 步对所有令牌执行完整的、无缓存的前向传播，确保结构完整性和细节细化不会因过时的近似而受损。

3. 主要贡献

HSA 算法： 一种完全无需训练、即插即用的推理范式，根据速度动态为时空令牌组分配异构去噪轨迹。
KV 缓存同步： 一种机制，在保留激活令牌完整 $N$ 令牌注意力上下文的同时，完全绕过非激活令牌，将每迭代注意力成本降低至 $O(|A_i| \cdot N)$ 。
缓存欧拉更新： 一种简单的张量操作，利用缓存速度推进被跳过的令牌，在不进行额外模型评估的情况下保持生成质量。
综合评估： 对四种令牌分组策略和四种预算预设的分析，证明了其在不同推理预算下的鲁棒性。

4. 实验结果

作者在 Wan-2.1-1.3B 和 Wan-2.2-A14B（文本到视频和图像到视频）以及 LTX-2（音频 - 视频）上评估了 HSA。

质量 - 运行时间权衡： 与均匀流匹配（FM）和 TeaCache 等最先进的缓存方法相比，HSA 实现了更优的帕累托前沿。
激进加速： 在激进加速区域（50% 和 25% 运行时间），优势最为明显。在 25% 运行时间下，HSA 在 VBench 基准测试中显著优于基线（总分：79.87%，而 FM 为 75.68%，TeaCache 为 78.33%）。
维度鲁棒性： 虽然基线方法在严格预算下遭受“灾难性维度崩溃”（在运动平滑度或时间一致性等特定 VBench 维度上严重退化），但 HSA 在所有 16 个评估维度上都能稳健地跟踪全预算参考。
图像到视频（I2V）： HSA 成功保持了图像条件对齐和视觉保真度，而减少步数（ $T=10$ ）的均匀 FM 基线则表现出严重的结构崩溃。
无需离线分析： 与需要昂贵离线分析来估计模型输出波动的 TeaCache 不同，HSA 无需任何离线调整或微调即可运行。

5. 意义与主张

本文声称，HSA 通过将计算资源与视觉数据的感知和动态复杂性对齐，解决了 DiT 推理中的根本性低效问题。通过单独处理令牌而非全局处理，HSA 在实现显著加速的同时，避免了通常与激进加速相关的质量下降。

作者强调，HSA 是一个灵活框架，可以通过各种令牌分组策略进行实例化。虽然动态选择目前产生最强结果，但该框架预期随着更复杂的分组策略而进一步改进。这项工作为利用时空数据固有异构性以实现高保真视频生成民主化的新一类推理算法奠定了基础。