引言

我们正处于生成式视频的黄金时代。从 Sora 到 Open-Sora 再到 Latte,扩散 Transformer (Diffusion Transformers, DiTs) 已经解锁了通过简单的文本提示生成高保真、连贯视频的能力。然而,有一个巨大的瓶颈阻碍了这些工具在实时应用中的普及: 推理速度

在消费级硬件上生成仅仅一秒钟的视频可能需要惊人的长时间。这主要是由于扩散模型的顺序特性所致。为了生成一张图像或视频帧,模型必须在数十甚至数百个时间步中迭代地去除噪声。这是一个缓慢而有条不紊的过程,每一步都依赖于前一步的结果。

为了加速这一过程,研究人员将目光投向了缓存 (Caching) 。 其想法很简单: 如果模型在时间步 \(t\) 的输出与时间步 \(t+1\) 非常相似,为什么还要计算 \(t+1\) 呢?直接复用 \(t\) 的结果即可。以前的方法试图通过均匀缓存来实现这一点——例如,计算一步,跳过下一步,计算第三步,依此类推。

但这里有个问题: 并非所有时间步都是生而平等的。 在生成的某些阶段,视频内容变化剧烈;而在其他阶段,几乎没有变化。均匀缓存忽略了这一现实,导致在速度和视觉质量之间难以取得良好的平衡。

这就轮到 TeaCache (Timestep Embedding Aware Cache,时间步嵌入感知缓存) 出场了。在最近的一篇论文中,研究人员提出了一种新颖的、无需训练的方法,能够智能地决定何时进行缓存。TeaCache 不再是猜测或遵循固定的时间表,而是通过观察模型的输入来预测输出是否值得计算。结果如何?在质量损失几乎可以忽略不计的情况下实现了显著加速。

视频扩散模型的质量-延迟对比。TeaCache 在视觉质量和效率方面均显著优于 PAB。

如图 1 所示,TeaCache 突破了极限,与 PAB 等先前的最先进方法相比,它实现了更低的延迟 (更靠左) 同时保持了更高的质量 (更靠上) 。

背景: 扩散瓶颈

要理解为什么 TeaCache 是必要的,我们需要先看看视频扩散模型是如何工作的。

去噪过程

扩散模型通过逆转噪声过程来生成数据。在训练期间,我们通过逐渐添加高斯噪声来破坏图像。在推理 (生成) 期间,模型学习逐步去除这些噪声以恢复清晰的视频。

在数学上,前向过程在 \(T\) 个步骤中添加噪声:

方程 1: 添加噪声的前向扩散过程。

逆向过程,也就是当你点击“生成”时发生的过程,涉及模型从噪声输入中估计干净的数据分布:

方程 2: 重建数据的逆向扩散过程。

在现代视频模型中,这是由 扩散 Transformer (DiTs) 处理的。这些是巨大的神经网络,将视频帧作为令牌 (token) 序列进行处理。由于逆向过程是顺序的 (\(t=1000 \to t=999 \to \dots \to 0\)) ,你无法并行化这些步骤。如果你想要高质量,你需要很多步。如果你想要速度,你可以减少步数,但质量通常会受损。

缓存的概念

既然我们无法轻易并行化这些步骤,我们就尝试跳过计算。在许多情况下,模型在步骤 \(t\) 的内部特征或输出与步骤 \(t-1\) 极其相似。

像 PAB (Pyramid Attention Broadcast,金字塔注意力广播) 这样的现有方法使用均匀缓存 。 它们可能会说,“每隔一个奇数步缓存一次输出。”

TeaCache 与传统均匀缓存策略的对比。均匀缓存 (上) 是僵化的。TeaCache (下) 是动态的。

图 2 (上排) 展示了均匀缓存。模型在“计算” (实线) 和“复用” (虚线) 之间僵化地交替。这是次优的,因为输出之间的差异是波动的。有时你需要连续计算 5 步;有时你可以连续跳过 5 步。均匀缓存对这种动态变化视而不见。

核心方法: TeaCache

TeaCache 背后的研究人员提出了一个关键问题: 我们如何能在不实际运行繁重计算的情况下知道输出是否会改变?

如果我们运行模型来检查输出差异,那我们就已经浪费了本想节省的时间。我们需要一个廉价的“指示器”,提前告诉我们输出是否是多余的。

假设: 输入揭示一切

TeaCache 的核心洞察是输入与输出相关 。 如果 Transformer 块的输入在时间步 \(t\) 和 \(t-1\) 之间几乎没有变化,那么输出很可能也不会有太大变化。

但是扩散模型有多个输入。哪一个能反映真相呢?

  1. 文本嵌入 (Text Embedding) : 这代表你的提示词 (例如“一只猫在跑”) 。它在整个过程中保持不变。 对缓存无用。
  2. 噪声输入 (Noisy Input) : 这是正在被去噪的视频内容。它变化缓慢且包含图像数据。
  3. 时间步嵌入 (Timestep Embedding) : 代表当前时间 \(t\) 的向量。它每一步都在变化,但对图像内容一无所知。

研究人员分析了这些输入与实际模型输出的关系。他们关注的是“时间步嵌入调制的噪声输入 (Timestep Embedding Modulated Noisy Input) ”——本质上是结合了 Transformer 块内部时间信号的噪声图像。

输入差异与输出差异的可视化。绿线 (调制输入) 紧密跟随蓝线 (输出) 。

图 3 证实了他们的假设。请看蓝线 (模型输出差异) 和绿线 (调制噪声输入差异) 。

  • Open Sora (a) 中,输出差异先高,后降,再升 (呈“U”形) 。绿线跟随这一趋势。
  • Latte (b)OpenSora-Plan (c) 中,趋势不同,但绿线始终比红线 (原始噪声输入) 更好地反映了蓝线的行为。

这意味着我们可以观察计算成本低廉的输入来预测计算成本高昂的输出

架构

让我们可视化这在模型中的发生位置。

带有 Transformer 的扩散模块。时间步嵌入调制输入。

如图 4 所示, 时间步嵌入 (黄色) 在噪声输入进入自注意力层之前对其进行调制。这个组合信号——“调制输入”——正是 TeaCache 所监控的对象。

度量标准

为了量化“差异”,论文使用了相对 L1 距离 。 如果 \(O_t\) 是时间 \(t\) 的输出,则差异为:

方程 4: 相对 L1 距离公式。

值越大意味着输出变化很大 (不要缓存!) 。值越小意味着输出稳定 (可以安全缓存!) 。

重缩放策略 (多项式拟合)

这里有一个陷阱。虽然输入的趋势与输出相匹配,但幅度可能会有偏差。它们的关系并不是完美的 1:1。

为了解决这个问题,TeaCache 使用了多项式拟合 。 他们取几个样本点并拟合一条曲线,将“输入差异”映射到预期的“输出差异”。

多项式拟合可视化。拟合 (橙色/绿色) 使输入差异更接近与输出差异的对角线。

图 5 展示了为什么这是必要的。原始数据点 (蓝色星号) 是分散的。通过应用多项式函数 \(f(x)\) (见下方方程) ,他们可以调整输入差异,使其成为输出差异的高精度预测指标。

方程 6: 多项式拟合函数。

TeaCache 算法

综上所述,TeaCache 通过累加估计的差异来工作。

  1. 计算当前步骤的输入差异。
  2. 使用多项式函数 \(f\) 对其进行重缩放。
  3. 将其加到一个运行总和中。
  4. 决策时间:
  • 如果总和低于阈值 \(\delta\),则复用缓存的输出。
  • 如果总和高于 \(\delta\),则运行模型计算新输出并重置总和。

该逻辑在以下不等式中形式化:

方程 7: TeaCache 指示器不等式。

这创建了一个动态的时间表 (如图 2 底部所示) ,模型仅在必要时才全力工作。

为什么不直接减少时间步?

你可能会问,“为什么不直接让模型运行更少的步数 (例如 50 步而不是 100 步) ?”

减少步数会导致更粗糙的噪声调度,这会显著降低质量。TeaCache 保持了细粒度的噪声调度,只是跳过了冗余的计算。

对比: 缓存 vs. 减少时间步。TeaCache 更好地保留了细节。

图 6 清晰地展示了这一点。减少步数 (中间栏) 会洗掉细节。TeaCache (右栏) 在速度更快的同时,保留了原始 30 步生成的纹理和光照。

实验与结果

研究人员在三个主要的开源视频模型上测试了 TeaCache: Open-Sora 1.2LatteOpenSora-Plan 。 他们将其与 T-GATE、\(\Delta\)-DiT 和 PAB 进行了比较。

定量性能

结果令人印象深刻。TeaCache 始终能获得更高的加速比和更好的 VBench (视频基准) 分数。

表 1: 定量评估。TeaCache 显示出卓越的加速和 VBench 分数。

观察表 1:

  • 速度: 在 OpenSora-Plan 上,TeaCache 实现了巨大的 4.41 倍加速 , 同时保持 VBench 分数为 80.32% (实际上与基线相同) 。更快的版本达到了 6.83 倍加速。
  • 质量: 几乎在每一行中,在相当的加速比下,TeaCache 的 PSNR、SSIM 和 VBench 分数都高于 PAB。

视觉质量

数字固然重要,但视频生成归根结底是关于视觉效果的。

视觉对比。TeaCache 的帧在相似或更快的速度下比 PAB 更清晰。

在图 7 中,我们可以看到并排对比。

  • Open-Sora (上) : 与原图相比,PAB 结果在女性头发的纹理上有所丢失。TeaCache 结果保留了这种清晰度,同时将推理时间从 44 秒减少到 28 秒。
  • Open-Sora-Plan (下) : TeaCache 将时间从 99 秒大幅减少到仅 22 秒——这对于实际使用来说是一个改变游戏规则的加速——且没有引入明显的伪影。

规模化效率

TeaCache 在不同分辨率和硬件规模下也表现良好。

推理效率图表。加速比在不同分辨率和 GPU 数量下保持一致。

图 8 显示,无论你是生成 480p 还是 720p,或者使用 1 个 GPU 还是 8 个 GPU,TeaCache 都能提供一致的推理时间减少 (橙色条明显低于蓝色的“Original”条) 。

重缩放的重要性

最后,研究人员验证了他们的多项式拟合策略。

表 3: 多项式拟合的消融实验。

表 3 显示,与使用原始数据相比,使用 4 阶多项式拟合提高了 VBench 分数。这证实了考虑输入和输出之间的非线性关系对于选择正确的缓存时机至关重要。

结论

视频生成发展迅速,但推理延迟仍然是一个主要障碍。 TeaCache 为这个问题提供了一个智能的、无需训练的解决方案。它不再盲目地跳过步骤,而是倾听模型的输入。通过跟踪“时间步嵌入调制的噪声输入”,TeaCache 能够准确识别模型何时在做冗余工作并跳过它。

其意义重大:

  1. 即插即用方案: 它不需要重新训练庞大的视频模型。
  2. 高保真度: 它保留了简单的步数减少所破坏的视觉细节。
  3. 速度: 凭借 1.5 倍到近 7 倍的加速,它使高质量视频生成更接近实时性能。

随着扩散模型在规模和复杂性上不断增长,像 TeaCache 这样的智能缓存策略对于让每个人都能高效地使用这些强大的工具将至关重要。