扩散模型的生成效果持续提升，但一个核心瓶颈也日益凸显——推理速度慢、计算成本高。

像FLUX、Wan这类主流的扩散生成模型，如今已能产出高分辨率、细节丰富且语义一致的图像与视频，表现确实令人印象深刻。不过在推理效率上，它们面临一个共同的挑战：多步去噪过程依赖大量计算，每一步都需要完整调用大型生成网络，导致算力开销居高不下。对于图像生成来说，这会显著拖慢交互式创作与在线部署的响应速度；而对于视频生成，模型还需同时兼顾空间细节与时间连续性，计算量更为庞大，对高效推理的需求也更为紧迫。

为应对这一难题，西湖大学AGI实验室提出了BudCache——一种专为固定计算预算设计的扩散模型step-level缓存方法。与那些依赖运行时阈值判断的缓存方案不同，BudCache不再在推理过程中临时决定计算多少步，而是先固定推理预算，再通过离线搜索找到该预算下的最优缓存策略。这样一来，既能严格把控推理成本，也能更好地保障最终生成质量。

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

扩散模型加速中的缓存技术：从重用到预算约束

在多步去噪过程中，相邻步骤之间通常存在大量可复用的信息。step-level缓存方法正是利用这一特性：选择关键的去噪步骤进行完整模型计算，而在其他步骤中复用之前保存的输出或中间特征，从而有效减少真实模型调用的次数。这种方法无需重新训练模型，也不改变原始参数，对于已训练好的FLUX、Wan等模型而言，缓存是一种即插即用的推理加速技术，直接作用于采样过程。

像TeaCache、MagCache等现有方法已充分证明了step-level缓存的有效性。然而，当前许多方法仍依赖启发式规则——在推理过程中，根据某种运行时信号判断当前步骤是否需要重新计算，或是否复用前一步的缓存结果。这种做法虽然能减少计算量，但存在两个关键问题。

1、计算量不稳定且不可控。启发式阈值方法的实际计算量通常由运行时信号触发，不同输入可能触发不同次数的完整模型计算，最终推理延迟难以稳定预测。而在真实部署场景中，开发者往往更希望提前指定一个明确的计算预算，而非被动接受不确定的开销。

2、决策局限于局部最优。某一步是否复用缓存，通常取决于当前步骤的差异度或误差信号，而非直接面向最终生成结果进行优化。换句话说，局部看似合理的缓存决策，未必能形成全局最优的生成轨迹。

BudCache的核心出发点正是重新定义这个问题：step-level缓存的关键，不仅在于判断“当前步骤能否缓存”，而是在固定计算预算下决定“哪些步骤最值得计算”。为了评估一个缓存策略的优劣，BudCache不再只依赖单步的局部误差，而是直接考察该策略对应的最终生成结果是否接近完整计算的结果——也就是说，看采用某一组缓存步骤后，最终生成的图像或视频，是否与完整计算产物高度一致。

从启发式缓存到固定预算搜索：迈向更可控的加速

BudCache将step-level缓存问题表述为一个预算约束下的离散优化问题。具体来说，首先固定模型实际前向计算的次数，也就是NFE预算。例如，在一个包含多个逻辑去噪步骤的采样过程中，可以预先设定只允许模型真实计算8次、9次或10次，其余步骤则复用缓存结果。这样一来，推理成本在部署前就已确定，不会随输入变化而产生不可控的延迟波动。

在此基础上，BudCache使用一个二进制掩码（mask）来表示缓存策略：对于每一个去噪步骤，如果对应位置为1，就执行完整模型计算；如果为0，则复用最近一次计算得到的缓存结果。由于掩码中1的数量被固定为预算B，BudCache能够严格保证实际模型调用次数。

然而，找到最优掩码并非易事。假设总共有K个逻辑去噪步骤，需要从中选出B个执行真实计算，可能的组合数量非常庞大。同时，扩散模型的去噪过程具有明显的序列依赖性：早期某一步的缓存决策可能影响后续整条采样轨迹，最终改变生成图像或视频的质量。因此，简单的局部贪心策略容易陷入局部最优，难以找到真正高质量的缓存方案。

为此，BudCache采用了“模拟退火+爬山算法”的混合搜索策略。第一阶段利用模拟退火进行全局探索：模拟退火允许搜索过程在早期接受一些暂时看起来更差的缓存策略，从而跳出局部最优，在更大的策略空间中寻找更优的候选。BudCache通过两类操作生成新策略：一种是Swap，交换一个计算步骤和一个缓存步骤的位置；另一种是Shift，将一个计算步骤移动到相邻位置。这两类操作都不改变总计算预算，确保搜索始终满足固定NFE约束。第二阶段使用爬山算法进行局部精修：当模拟退火找到较好的策略区域后，BudCache会检查当前策略附近的候选方案，并选择能让最终结果更接近完整计算结果的策略，持续迭代直到附近没有更好的候选为止。

这种混合搜索方式使BudCache能够在较小的离线成本下找到高质量的缓存策略。更重要的是，搜索过程仅在部署前进行一次；实际推理时直接使用已搜索好的固定缓存策略，无需在线搜索，也无需运行时阈值判断，因此不会引入任何额外的推理开销。

在缓存策略搜索完成后，BudCache还可以进一步进行cache-aware schedule alignment（缓存感知时间步校准）。这一步可以理解为一个可选的时间步校准模块：在固定缓存掩码的前提下，对采样时间步进行轻量级离线优化，使缓存采样器的最终输出更接近完整计算的结果。具体来说，完整计算的采样过程被视为教师模型（teacher），缓存后的采样过程被视为学生模型（student），BudCache优化学生模型使用的时间步分布，使其最终生成结果尽可能接近教师模型。优化完成后，新的时间步将与搜索得到的缓存策略一起固定下来，在推理阶段直接使用。需要强调的是，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缓存从“启发式阈值触发”推进到了“固定预算下的缓存策略搜索”。以往的缓存方法主要关注如何根据运行时信号判断当前步骤能否复用缓存，而BudCache关注的则是一个更适合实际部署的问题：当推理预算已经固定时，哪些步骤最值得进行真实计算？通过预算约束建模、离线缓存策略搜索以及时间步校准，BudCache能够在不重新训练模型、不改变原始模型参数、不增加在线推理开销的前提下，为FLUX和Wan等主流图像与视频生成模型提供更可控的加速效果。

总体来看，BudCache提供了一种更部署友好的扩散模型缓存加速方案：既能严格控制推理成本，又能在固定预算下更好地保持生成质量，为高效图像与视频生成开辟了新的技术路径。

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