OTCache：扩散模型中基于最优传输的几何感知缓存

# OTCache：扩散模型中用于几何感知缓存的优化传输

来源：https://arxiv.org/html/2606.31026

11单位：中国联通数据科学与人工智能研究院  
22单位：中国联通数据智能部  
33单位：南京大学计算机软件新技术国家重点实验室  
33邮箱：\{gaohl51,zhaof50,wangk115,liansg\}@chinaunicom.cn  
†同等贡献。∗通讯作者。方赵†https://orcid.org/0009-0003-3465-3050 强辉https://orcid.org/0000-0002-3674-092X 付远石 赵韶安 李彦涛 谭超https://orcid.org/0009-0008-4957-6939 陆婷 尤鱼龙https://orcid.org/0009-0004-0168-9265 王凯∗https://orcid.org/0000-0002-1171-0281 连石国∗https://orcid.org/0000-0003-4308-7049

###### 摘要

我们提出 **OTCache**，一种无需训练的框架，通过缓存调度预测来加速扩散采样。现有的基于图的缓存方法通过优化最短路径目标来减少冗余计算，但依赖于加性独立性假设，该假设在低 NFE 条件下常常失效。为解决此问题，OTCache 将跨推理预算的缓存调度建模为策略空间中的平滑演变，受最优传输（OT）启发。该框架包含三个阶段：(1) 在保守预算下使用基于图的缓存方法获得高保真参考调度；(2) 在极端低预算设置下，通过 Optuna 优化并结合端到端感知目标执行轻量级锚点搜索；(3) 利用连续扭曲表示，通过参考策略与锚点策略之间的分位数插值来预测目标预算的调度。在 FLUX.1 [dev]、Qwen-Image 和 HunyuanVideo 上的实验表明，OTCache 分别实现了 4.5×、4.7× 和 3.66× 的加速，同时在生成保真度上持续优于最先进的缓存基线。这项工作通过最优传输启发的调度建模，为加速扩散模型提供了新视角。  
代码：https://github.com/UnicomAI/OTCache

## 1 引言

流匹配（FM）[lipman2022flow, albergo2022building] 已成为现代生成建模的基石，在图像 [flux2024]、视频 [kong2024hunyuanvideo, Open-Sora] 和多模态任务 [hung2024tangofluxsuperfastfaithful] 的高保真合成中取得了显著突破。通过瞬时速度场建模连续传输路径，FM 为采样提供了数学上严谨且有效的范例。然而，在 FLUX.1 [flux2024] 和 HunyuanVideo [kong2024hunyuanvideo] 等商业规模模型中，基于 Transformer 的去噪器参数规模巨大，加上漫长的迭代采样过程，导致了惊人的计算开销和内存占用。这种计算密集型特性造成了巨大的部署障碍，严重限制了其在交互式或资源受限场景中的适用性。

为缓解这一负担，蒸馏 [salimans2022progressive, sauer2024adversarial]、剪枝 [han2015deep] 和量化 [li2023q] 等成熟的加速策略已被广泛探索。然而，这些方法大多带来沉重的“采用成本”——通常需要在大规模数据集上进行密集的再训练、复杂的架构修改或复杂的工程流程，这限制了它们快速部署的灵活性。相比之下，缓存技术 [ma2023deepcache, wimbauer2023cache] 作为一种引人注目的无需训练替代方案而兴起。

最近如 MeanCache [gao2026meancache] 等进展，将缓存调度形式化为一个带约束的最短路径问题，允许在固定 NFE 预算下进行高效、无需训练的规划。然而，该范式依赖于*加性独立性*假设，将端到端退化近似为局部边误差的总和。我们发现这种替代指标在低 NFE 条件下变得愈发不可靠。如图 1A 所示，当 NFE 从 20 降至 8 时，MeanCache 与基于搜索的最优值（Optuna）之间仍存在明显差距，表明 LPIPS 仍有较大改进空间。

图 1：FLUX.1 [dev] 上的性能提升与结构洞察。(A) 加性替代指标的失败：在超低 NFE 条件下，MeanCache 偏离最优边界，反映了局部误差聚合的局限性。OTCache 恢复了这一保真度损失，显著缩小了与搜索最优值的差距。(B) 最优路径的结构规律性：不同 NFE（预算）下的最优调度表现出强结构关系而非独立模式，提示存在共享的底层策略轨迹，OTCache 通过几何插值加以利用。

更根本的是，图 1B 揭示了一种被忽视的结构规律性：最优调度在不同预算下的间隔分布（一阶差分）并非无关；相反，其时序结构随着 NFE 预算的变化而平滑演变。这一观察促使视角转变：我们不再为每个预算独立求解离散最短路径问题，而是建模预算条件化的调度演化，并利用跨预算的结构关系。

基于这一洞察，我们提出 **OTCache**，一个无需训练的框架，通过最优传输插值来预测低 NFE 下的调度。通过用可靠的高预算参考和低预算锚点进行锚定预测，并在 Wasserstein 空间中在两者之间插值，OTCache 生成鲁棒的调度，在低 NFE 加速下保持稳定。本文的主要贡献总结如下：

- **重新思考低 NFE 下的缓存调度**：我们指出现有基于图调度器的两个基本局限：(i) 在非线性误差传播下，加性最短路径目标导致的替代指标错位；(ii) 忽略最优调度之间结构关系的独立预算假设。
- **预算条件化的调度演化**：我们引入一个三阶段加速框架，将缓存调度视为概率测度，并将其随 NFE 的演化建模为策略空间中的平滑轨迹。利用最优传输的几何结构，通过两个可靠端点之间的分位数插值获得目标预算调度。
- **卓越性能**：在 FLUX.1、Qwen-Image 和 HunyuanVideo 上的实验表明，OTCache 分别实现了 4.5×、4.7× 和 3.66× 的加速。特别是在 Qwen-Image 上，OTCache 的 LPIPS 达到 0.171。

## 2 相关工作

### 2.1 扩散模型加速

生成模型 [song2020ddim, song2020score] 在多种模态上的显著成功，很大程度上归功于采样速度的持续进步。早期工作主要集中在通过原则性的数值 SDE/ODE 求解器 [song2020score, jolicoeur2021gotta, chen2025optimizing] 来优化迭代去噪过程，例如 DDIM [song2020ddim]、EDM [karras2022elucidating] 和 DPM-Solver [lu2022dpm]，旨在用更少的离散化步骤保持高合成质量。为进一步压缩推理轨迹，知识蒸馏 [hinton2015distilling] 被广泛探索，将多步去噪映射为紧凑的少步甚至单步机制，例如渐进式蒸馏 [salimans2022progressive] 和一致性模型 [song2023consistency, sauer2024adversarial, wang2025target]。作为补充，正交策略——包括量化 [li2023q, shang2023post]、剪枝 [ma2023llm] 以及系统级并行化框架 [zhao2024dsp, chen2024asyncdiff]——已被研究以提升原始硬件吞吐量。最近，流匹配 [lipman2022flow, albergo2022building] 的出现引入了一种新范式，学习确定性速度场，本身具备用最少采样步骤实现高保真生成的潜力。然而，大多数现有方法仍需大量计算、大规模数据或复杂工程，限制了其在资源受限场景中的实际应用。

### 2.2 扩散模型中的缓存

作为一种无需训练的加速范式，缓存策略 [wimbauer2023cache, ma2024learning] 通过重用中间表示来绕过冗余计算而日益受到关注。早期方法如 DeepCache [ma2023deepcache] 引入了针对 UNet 主干网络的特定于架构的特征重用，而 T-GATE [zhang2024cross] 和 Δ-DiT [chen2024delta] 则将这些思想扩展到基于 Transformer 的架构 [peebles2023dit]。对于大规模视频生成，PAB [zhao2024real] 和 TeaCache [DBLP:journals/corr/abs-2411-19108] 利用时间相关性和误差阈值来触发特征重用。近期研究进一步探索了细粒度缓存标准，包括视频扩散 Transformer 的自适应重用 [kahatapitiya2025adaptive]、基于性能分析的缓存选择 [ma2025model]、频率感知缓存 [liu2025freqca]、投机性特征缓存 [liu2025speca] 以及基于搜索的策略发现 [aggarwal2025evolutionary]。该领域也朝着基于图的方向发展；LeMiCa [gao2025lemica] 将视频合成抽象为有向无环图（DAG）进行全局调度，而 MeanCache [gao2026meancache] 则从 MeanFlow [geng2025mean] 汲取灵感，将缓存从瞬时速度视角重新表述为平均速度视角，稳定了采样轨迹。尽管有这些进展，现有方法通常优化固定预算下的调度或依赖局部替代指标。如何在极高加速条件下，即低 NFE 设置中获得准确的*最优*缓存策略，仍是一个开放问题。

## 3 方法

### 3.1 预备知识

#### 3.1.1 流匹配

流匹配 [lipman2022flow, albergo2022building] 和整流流 [liu2022flow] 通过构造噪声分布 π₁ 和数据分布 π₀ 之间的连续传输路径，引入了基于扩散的生成建模的新范式。通过线性插值定义概率路径：

x_t = (1 - t) x_0 + t x₁, t ∈ [0,1], (1)

这些方法旨在学习状态空间中的直线轨迹。由于在去噪（生成）过程中干净数据 x_0 未知，使用学习到的速度场 v_θ(x_t, t) 来近似方向 (x_0 - x₁)，从而构建神经 ODE 模型：

d x̂_t = v_θ(x_t, t) dt (2)

数值求解器将该 ODE 离散化为 N 个步骤以恢复数据分布。总计算成本主要由函数评估次数（NFE）决定，因此高效的逐步调度对于推理加速至关重要。

#### 3.1.2 基于图的缓存调度

特征缓存已成为一种主导的无需训练范式，通过跨相邻时间步重用中间状态来加速推理。最新最先进的方法 [gao2025lemica, gao2026meancache] 将其形式化为有向多重图 G = (V, E) 上的带约束最短路径问题。在该形式化中，每个节点 v ∈ V 对应一个离散时间步，边 e = (t → s) 表示一个缓存转移，其中 s 处的速度场使用从 t 缓存的特征估计。保真度损失由边权重量化：

L(t → s) = ||v(x_s, s) - v̂(x_s, s; x_t)||_p, (3)

其中 v̂ 表示缓存估计器。全局调度问题旨在找到最小化累积误差的最优路径 π^⋆：

π^⋆ = argmin_{π ∈ P} ∑_{e ∈ π} L(e)^γ, s.t. |π| ≤ B, (4)

其中 P 是可行路径集合，B 是计算预算。

### 3.2 重新思考：超越加性最短路径替代指标

##### 局限性

基于图的调度器将缓存形式化为带约束最短路径问题（公式 4）。在此形式化中，任意两个时间步 (t, s) 之间的缓存成本首先作为边权重 L(e) 独立估计，然后通过最小化这些边成本的总和获得最优调度。这隐含地假设了*加性独立性*：一个调度的端到端退化可以通过聚合独立估计的边损失来近似。

这种近似在高 NFE 条件下通常是可靠的，因为缓存间隔短，每个决策对后续步骤的影响有限。然而，在低 NFE 条件下，缓存操作的数量和缓存状态之间的间隔都增加。因此，早期缓存决策的影响沿着轨迹传播得更远，并与后面的决策相互作用，违反了加性独立性假设。结果，最短路径替代指标在估计真实路径退化时变得越来越不准确。与此观察一致，图 1A 显示 MeanCache 与最优边界之间的性能差距随着 NFE 的减小而增大。

##### 观察

尽管存在这一局限性，我们观察到清晰的结构规律性：不同 NFE 预算下的近最优调度*并非独立*。如图 1B 所示，最优调度的时间间隔模式随预算变化而平滑演变。这表明不同 NFE 下的最优调度遵循结构化演化，而非形成无关解。

##### 核心洞察

我们假设不同 NFE 预算下的最优调度对应于策略空间中共享的*理想轨迹*在不同分辨率下的观测。尽管观测分辨率（NFE）变化，底层轨迹仍保持固定。在此视角下，跨预算的调度在策略流形上形成平滑演化，我们将其解释为*策略测地线*。这一视角启发我们通过策略空间中的几何插值来预测调度，我们在下一节中使用最优传输实现这一想法。

图 2：OTCache 概览。阶段 1：使用基于图的缓存方法获得可靠的高 NFE（预算）调度作为参考策略。阶段 2：执行黑盒搜索，找到近最优的低 NFE 调度作为锚点策略。阶段 3：受最优传输启发，通过 PCHIP 将两个端点调度转换为连续扭曲曲线，并应用分位数插值预测目标预算调度 π_B。

### 3.3 OTCache

受第 3.2 节观察的启发，我们提出 OTCache，一个无需训练的三阶段框架，通过缓存调度预测加速扩散采样，如图 2 所示。

#### 3.3.1 阶段 1：参考

经验上，我们发现基于图的缓存方法（例如 MeanCache）在低 NFE 条件下可能无法恢复最优调度。然而，在高 NFE 条件下，这些方法通常能生成高度可靠的调度，因为加性独立性假设成立。因此，我们保守地选择一个相对高的 NFE 预算 B_high，并使用基于图的调度器 G 获得一个参考调度 π_ref = G(B_high)。该参考调度充当结构上的锚点。