选摘自 西湖大学AGI实验室

量子位 | 公众号 QbitAI

扩散模型生成得越来越好，但也越来越慢。

近年来，以FLUX、Wan为代表的扩散生成模型，在文本生成图像、文本生成视频上，已经能产出高分辨率、细节丰富又语义一致的内容。

但它们都有一个共同的麻烦——推理依赖多步去噪，每一步都要完整调用一次大型生成网络，算力成本居高不下。

对图像生成来说，这会直接拖慢交互式创作和在线部署的响应；对视频生成来说，模型还要同时兼顾空间细节与时间连续性，计算量更大，对高效推理的需求也更突出。

为应对这一问题，西湖大学AGI实验室提出了BudCache：一种面向固定计算预算的扩散模型step-level cache方法。

和已有那些基于启发式阈值的缓存方法不同，BudCache不再在运行时靠阈值临时决定到底算多少步，而是先固定推理预算，再离线搜出该预算下最优的缓存策略。

这样既能严格卡住推理成本，也能更好地保住最终生成质量。

论文一作是西湖大学AGI实验室研究助理雷明坤，指导老师是西湖大学助理教授张驰。

扩散模型加速中的cache技术

在多步去噪过程中，相邻步骤之间通常存在可以复用的信息。因此，step-level cache方法会选择一部分关键去噪步骤进行完整模型计算，而在其他步骤中复用此前保存的输出或中间特征，从而减少真实模型调用次数。

这种方法的优势在于不需要重新训练模型，也不需要改变原始模型参数。对于已经训练好的FLUX和Wan等生成模型而言，cache可以作为一种即插即用的推理加速技术，直接作用于采样过程。

已有方法如TeaCache、MagCache等已经证明了step-level cache的有效性。

但当前许多方法仍然主要依赖启发式规则进行缓存决策——在推理过程中根据某种运行时信号，判断当前步骤是否需要重新计算、或是否可以复用前一步的缓存结果。

这种做法虽然能减少计算，但仍存在两个关键问题。

1、计算量不稳定、不可控。 启发式阈值方法的实际计算量通常由运行时信号触发，不同输入可能触发不同数量的完整模型计算，因此最终推理延迟并不总是稳定可控。

而真实部署场景往往更希望提前指定一个明确的计算预算，而不是在运行时被动接受不确定的开销。

2、决策是局部的。 某一步是否复用缓存，往往取决于当前步骤的差异度或误差信号，而不是直接面向最终生成结果做优化。

换句话说，局部看似合理的缓存决策，并不一定能形成全局最优的生成轨迹。

BudCache的核心出发点正是重新定义这一问题：

step-level cache的关键不只是判断“当前步骤能不能缓存”，而是在固定计算预算下决定“哪些步骤最值得计算”。

为了判断一个缓存策略好不好，BudCache不再只依赖单步的局部误差信号，而是直接考察该策略对应的最终生成结果是否接近完整计算的结果——

也就是直接看采用某一组缓存步骤后，最后生成出来的图像或视频，是否接近完整计算的产物。

从启发式缓存到固定预算搜索

BudCache将step-level cache表述为一个预算约束下的离散优化问题。

具体来说，BudCache先固定模型实际前向计算的次数，也就是NFE预算。

例如在一个包含多个逻辑去噪步骤的采样过程中，可以预先设定只允许模型真实计算8次、9次或10次，其余步骤复用缓存结果。这样一来，推理成本在部署前就已确定，不会随输入变化产生不可控的延迟波动。

在此基础上，BudCache用一个二进制mask表示缓存策略：

对每一个去噪步骤，如果对应位置为1，就执行完整模型计算；如果为0，则复用最近一次计算得到的缓存结果。由于mask中1的数量被固定为预算B，BudCache可以严格保证实际模型调用次数。

但如何找到最优mask并不简单。假设总共有K个逻辑去噪步骤，需要从中选B个执行真实计算，那么可能的组合数量非常大。

同时，扩散模型的去噪过程具有明显的序列依赖：早期某一步的缓存决策可能影响后续整条采样轨迹，最终改变生成图像或视频的质量。

因此，简单的局部贪心策略容易陷入局部最优，难以找到真正高质量的缓存方案。

为此，BudCache采用了“模拟退火+爬山算法”的混合搜索策略。

第一阶段用模拟退火做全局探索。 模拟退火允许搜索过程在早期接受一些暂时看起来更差的缓存策略，从而跳出局部最优，在更大的策略空间里寻找更好的候选。

BudCache通过两类操作生成新策略：一种是Swap，交换一个计算步骤和一个缓存步骤的位置；另一种是Shift，把一个计算步骤移动到相邻位置。这两类操作都不改变总计算预算，因此搜索始终满足固定NFE约束。

第二阶段用爬山算法做局部精修。 当模拟退火找到较好的策略区域后，BudCache会检查当前策略附近的候选方案，并选择能让最终结果更接近完整计算结果的策略，持续进行直到附近没有更好的候选为止。

这种混合搜索方式让BudCache能在较小的离线成本下找到高质量缓存策略。

更重要的是，搜索过程只在部署前进行一次；真正推理时直接使用已经搜好的固定缓存策略，不需要在线搜索，也不需要运行时阈值判断，因此不会引入额外的推理开销。

在缓存策略搜索完成后，BudCache还可以进一步进行cache-aware schedule alignment。

这一步可以理解为一个可选的时间步校准模块：在固定缓存mask的前提下，对采样时间步做轻量级离线优化，让缓存采样器的最终输出更接近完整计算的结果。

具体而言，完整计算的采样过程被视为teacher，缓存后的采样过程被视为student，BudCache优化student使用的时间步分布，使其最终生成结果尽可能接近teacher。

优化完成后，新的时间步会和搜索得到的缓存策略一起固定下来，在推理阶段直接使用。

需要强调的是，cache-aware schedule alignment并不改变原始生成模型，也不增加推理时的模型调用次数，只是在固定缓存策略的基础上让采样过程更好地适应缓存带来的轨迹变化。

主流图像与视频生成模型上的实验结果

研究团队首先在当前主流图像生成模型FLUX.1-dev上评估了BudCache，并在DrawBench上与多种代表性缓存方法做了比较，包括TeaCache、MagCache、LeMiCa、DiCache、ERTACache和TaylorSeer。

结果显示，在相同推理预算下，BudCache能更好地保持生成质量。

从可视化结果来看，BudCache在文字生成、复杂结构和细节保持方面更加稳定。

例如在包含“Google Brain Toronto”或“CLIMATE CHANGE”等文字的图像中，部分启发式缓存方法会出现字符错误、拼写混乱或文字模糊，而BudCache能更好地保持文字内容的正确性。在机械钟表、雕像结构、物体相对位置等复杂场景中，BudCache也能更好地保留几何细节和语义一致性。

视频生成方面，研究团队进一步在Wan2.1-T2V-1.3B上做了实验。视频生成通常需要处理更长的视觉序列和更复杂的时空关系，因此对推理加速提出了更高要求。

实验结果表明，BudCache在视频生成中同样有效：

在Wan2.1-T2V-1.3B上，完整推理需要约189秒，TeaCache约100秒，而BudCache只需约82秒。同时BudCache在PSNR、SSIM和LPIPS等指标上均优于TeaCache，说明它不仅速度更快，也能更好地保持视频重建质量。

从可视化结果来看，BudCache能在加速条件下保持更稳定的视频细节，例如人物服装颜色、摩托车外观、画面主体与背景关系等都能被更好地保留。

这说明BudCache的固定预算搜索策略不仅适用于图像生成，也能扩展到视频生成模型。

泛化性实验：搜索策略能否迁移到更多设置？

除了在固定设置下比较性能，研究团队还进一步验证了BudCache的泛化能力。

一个关键问题是：离线搜索得到的缓存策略，是否只对某一个特定配置有效？如果每次更换设置都必须重新搜索，方法的实际使用成本就会提高。

实验结果显示，BudCache搜索得到的缓存策略具有较好的迁移能力。

研究团队将一个在源设置下搜索得到的缓存配置，直接迁移到不同solver、不同分辨率和不同CFG scale的设置中，并在FLUX.1-dev上进行评估。

结果表明，在这些不同推理条件下，BudCache仍然稳定优于TeaCache，说明搜索得到的缓存策略并没有过拟合到单一采样配置。

更多不同设置下的实验结果、消融分析与可视化对比，可参考原论文。

总结

BudCache的核心贡献，是把step-level cache从“启发式阈值触发”推进到“固定预算下的缓存策略搜索”。

以往的缓存方法主要关注如何根据运行时信号判断当前步骤能否复用缓存，而BudCache关注的是一个更适合部署的问题：

当推理预算已经固定时，哪些步骤最值得进行真实计算？

通过预算约束建模、离线缓存策略搜索和时间步校准，BudCache能在不重新训练模型、不改变原始模型参数、不增加在线推理开销的前提下，为FLUX和Wan等主流图像与视频生成模型提供更可控的加速效果。

总体来看，BudCache提供了一种更部署友好的扩散模型cache加速方案：

既能严格控制推理成本，也能在固定预算下更好地保持生成质量，为高效图像与视频生成提供了新的技术路径。

论文标题：Budget-Constrained Step-Level Diffusion Caching

项目地址：https://github.com/Westlake-AGI-Lab/BudCache

论文链接：https://arxiv.org/abs/2606.13496