Blog Tags GitHub
Back to blog
diffusion-model video-generation kv-cache

TeaCache: Timestep Embedding Aware Cache for Video Diffusion Model

基于时间步嵌入感知的缓存策略加速视频扩散模型推理

TeaCache: Timestep Embedding Aware Cache for Video Diffusion Model

一、论文概述

项目内容
标题Timestep Embedding Tells: It’s Time to Cache for Video Diffusion Model
作者Feng Liu, Shiwei Zhang, Xiaofeng Wang, Yujie Wei, Haonan Qiu, Yuzhong Zhao, Yingya Zhang, Qixiang Ye, Fang Wan
机构UCAS, Alibaba Group, Fudan University, NTU
论文arXiv:2411.19108
代码Project Page
发布2024年11月28日
许可arXiv 非独占分发许可

二、核心思想

问题定义

视频扩散模型 (DiT) 因去噪过程的顺序性导致推理速度慢。现有缓存方法在均匀选择的时间步缓存模型输出，忽略了模型输出差异在不同时间步的非均匀性，导致推理效率和视觉质量之间难以平衡。

解决方案概述

TeaCache (Timestep Embedding Aware Cache) 是一种无需训练的缓存策略：

  • 核心观察: 模型输入与输出之间存在强相关性，且输入计算成本极低
  • 时间步嵌入调制: 使用时间步嵌入调制的噪声输入来近似模型输出差异
  • 多项式拟合重缩放: 通过多项式拟合修正输入差异与输出差异之间的缩放偏差
  • 选择性缓存: 基于累积相对 L1 距离决定是否缓存

核心性能

指标数值
最大加速比4.41× (Open-Sora-Plan)
质量损失-0.07% VBench 分数
训练需求无需训练
兼容性Open-Sora, Latte, OpenSora-Plan

三、技术架构

整体框架图

TeaCache vs 均匀缓存

Figure 2: TeaCache 与传统均匀缓存策略对比。TeaCache 能够选择性缓存信息量大的中间模型输出，在保持性能的同时加速 DiT 模型。

核心公式

扩散模型前向过程

xt=αtxt−1+1−αtztfort=1,…,T\mathbf{x}_{t}=\sqrt{\alpha_{t}}\mathbf{x}_{t-1}+\sqrt{1-\alpha_{t}}\mathbf{z}_{t} \quad \text{for} \quad t=1,\ldots,T

扩散模型反向过程

pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),Σθ(xt,t))p_{\theta}(\mathbf{x}_{t-1}\mid\mathbf{x}_{t})=\mathcal{N}(\mathbf{x}_{t-1};\mu_{\theta}(\mathbf{x}_{t},t),\Sigma_{\theta}(\mathbf{x}_{t},t))

时间步嵌入

Tt=MLP(sinusoidal(t))fort=1,…,T\mathbf{T}_{t}=MLP(sinusoidal(t)) \quad \text{for} \quad t=1,\ldots,T

相对 L1 距离 (输出差异度量)

L1rel(O,t)=∥Ot−Ot+1∥1∥Ot+1∥1\text{L1}_{\text{rel}}(\mathbf{O},t)=\frac{\|\mathbf{O}_{t}-\mathbf{O}_{t+1}\|_{1}}{\|\mathbf{O}_{t+1}\|_{1}}

大值表示 Ot\mathbf{O}_t 相对于 Ot+1\mathbf{O}_{t+1} 信息量大，应缓存；小值表示两者相似，可复用。

朴素缓存策略

∑t=tatb−1L1rel(F,t)≤δ<∑t=tatbL1rel(F,t)\sum_{t=t_{a}}^{t_{b}-1}\text{L1}_{\text{rel}}(\mathbf{F},t)\leq\delta<\sum_{t=t_{a}}^{t_{b}}\text{L1}_{\text{rel}}(\mathbf{F},t)

多项式拟合重缩放

y=f(x)=a0+a1x+a2x2+⋯+anxny=f(x)=a_{0}+a_{1}x+a_{2}x^{2}+\cdots+a_{n}x^{n}

其中 yy 是估计的模型输出差异，xx 是时间步嵌入调制噪声输入的差异。

最终缓存指标

∑t=tatb−1f(L1rel(F,t))≤δ<∑t=tatbf(L1rel(F,t))\sum_{t=t_{a}}^{t_{b}-1}f(\text{L1}_{\text{rel}}(\mathbf{F},t))\leq\delta<\sum_{t=t_{a}}^{t_{b}}f(\text{L1}_{\text{rel}}(\mathbf{F},t))

模型组件

组件说明关键参数
时间步嵌入正弦嵌入 + MLP标量 tt → 向量 Tt\mathbf{T}_t
调制层时间步嵌入调制噪声输入控制输入输出幅度
差异估计器相对 L1 距离计算标量差异
多项式拟合numpy poly1d4 阶多项式
缓存决策器累积阈值比较阈值 δ\delta

缓存流程

  1. 校准阶段: 在第一个样本上计算输入差异和输出差异，拟合多项式
  2. 推理阶段: 计算时间步嵌入调制噪声输入的相对 L1 距离
  3. 重缩放: 使用拟合的多项式修正差异估计
  4. 决策: 累积差异超过阈值时计算新输出，否则复用缓存

四、核心创新

创新点说明理论/实验依据
时间步嵌入感知缓存利用时间步嵌入调制的噪声输入估计输出差异输入输出相关性分析 (Fig. 3)
多项式拟合重缩放修正输入差异与输出差异的缩放偏差4 阶多项式拟合效果最佳 (Table 3)
选择性缓存非均匀策略，信息量大时计算，小时复用优于均匀缓存 (Table 1)
无需训练完全兼容 DiT 扩散模型校准阶段开销可忽略

五、实验结果

质量-延迟对比

质量延迟对比

Figure 1: 视频扩散模型质量-延迟对比。TeaCache 在视觉质量和效率上显著优于 PAB。

主要结果 (Latte, 16帧, 512×512)

方法FLOPs (P) ↓加速比 ↑延迟 (s) ↓VBench ↑
Latte (T=50)3.361×26.9077.40%
Δ-DiT3.361.02×-52.00%
T-GATE2.991.13×-75.42%
PAB-slow2.701.21×22.1676.32%
PAB-fast2.521.34×19.9873.13%
TeaCache-slow1.861.86×14.4677.40%
TeaCache-fast1.522.21×12.1677.22%

输入输出相关性分析

输入输出差异可视化

Figure 3: Open Sora、Latte 和 OpenSora-Plan 的连续时间步输入差异和输出差异可视化。时间步嵌入和时间步嵌入调制噪声输入与模型输出有强相关性。

多项式拟合效果

多项式拟合可视化

Figure 5: 输入差异与输出差异的相关性可视化。原始数据点偏离线性相关，多项式拟合减少偏差。

不同模型加速效果

模型原始延迟TeaCache 延迟加速比VBench 变化
Open-Sora--2.50×-0.78%
Latte26.90s12.16s2.21×-0.18%
Open-Sora-Plan--4.41×-0.07%

缓存 vs 减少时间步

缓存 vs 减少时间步

Figure 6: 缓存机制 vs 减少时间步。减少推理时间步会导致视觉质量下降，而 TeaCache 保持视觉质量。

不同分辨率和视频长度

不同分辨率效率

Figure 8: TeaCache 在不同视频长度和分辨率下的推理效率。

消融实验

缓存指标消融 (Table 2)

指标VBench ↑LPIPS ↓SSIM ↑PSNR ↑
Timestep embedding76.30%0.19040.721819.44
Timestep embedding modulated input77.22%0.18310.724819.49

多项式拟合消融 (Table 3)

阶数VBench ↑LPIPS ↓SSIM ↑PSNR ↑
无拟合76.50%0.18790.719619.45
1 阶76.80%0.18520.722119.47
2 阶77.06%0.18410.723519.48
3 阶77.18%0.18350.724419.49
4 阶77.22%0.18310.724819.49

多 GPU 扩展 (Table 4)

设置加速比VBench ↑
单 GPU2.21×77.22%
2 GPU (DSP)3.84×77.21%
4 GPU (DSP)6.49×77.18%

六、与现有方法对比

方面均匀缓存 (DeepCache)PABFasterCacheTeaCache
缓存策略均匀间隔按注意力块特性CFG 冗余优化时间步嵌入感知
训练需求无需训练无需训练无需训练无需训练
自适应性无部分部分完全自适应
最大加速~2×1.34×~2×4.41×
质量损失中等较大较小极小
兼容性UNetDiTDiTDiT

七、相关工作

相关工作与本文关系
DeepCacheUNet 缓存启发，TeaCache 扩展到 DiT
PABDiT 缓存基线，TeaCache 显著优于
Δ-DiTDiT 缓存基线，质量损失大
FasterCacheCFG 冗余优化，互补方法
AdaCache自适应缓存，但需要内容复杂度评估
DiT (Peebles & Xie)基础架构，TeaCache 完全兼容

八、总结

核心贡献

  1. TeaCache: 首个基于时间步嵌入感知的训练无关缓存策略
  2. 两阶段估计: 粗估计 (时间步嵌入调制) + 精炼 (多项式拟合)
  3. 4.41× 加速: 在 Open-Sora-Plan 上实现 4.41× 加速，质量损失仅 -0.07%
  4. 广泛兼容: 适用于 Open-Sora、Latte、OpenSora-Plan 等多种 DiT 模型

技术影响

  • 视频生成加速: 为 DiT 视频扩散模型提供高效推理方案
  • 无需训练: 零额外训练成本，即插即用
  • 自适应缓存: 根据时间步动态调整缓存策略
  • 多 GPU 扩展: 与 DSP 结合实现 6.49× 加速

局限性

  • 需要校准阶段 (第一个样本)
  • 多项式拟合阶数需要手动选择
  • 仅在视频生成任务上验证
  • 未与蒸馏方法结合

九、参考资源

十、关键公式速查

公式说明
Tt=MLP(sinusoidal(t))\mathbf{T}_{t}=MLP(sinusoidal(t))时间步嵌入
L1rel(O,t)=∥Ot−Ot+1∥1∥Ot+1∥1\text{L1}_{\text{rel}}(\mathbf{O},t)=\frac{\|\mathbf{O}_{t}-\mathbf{O}_{t+1}\|_{1}}{\|\mathbf{O}_{t+1}\|_{1}}相对 L1 距离
y=f(x)=a0+a1x+a2x2+⋯+anxny=f(x)=a_{0}+a_{1}x+a_{2}x^{2}+\cdots+a_{n}x^{n}多项式拟合
∑f(L1rel(F,t))≤δ\sum f(\text{L1}_{\text{rel}}(\mathbf{F},t))\leq\delta缓存决策

十一、关键图片索引

图片说明文件名
Figure 1质量-延迟对比latency.png
Figure 2TeaCache vs 均匀缓存框架x2.png
Figure 3输入输出差异可视化opensora-diff.png, latte-diff.png, opensora-plan-diff.png
Figure 4DiT 扩散模块结构x1.png
Figure 5多项式拟合效果opensora-fit.png, latte-fit.png, opensora-plan-fit.png
Figure 6缓存 vs 减少时间步x3.png
Figure 7定性对比x4.png
Figure 8不同分辨率效率480p-48f.png, 480p-192f.png, 360p-240f.png, 720p-48f.png