同样算力质量更高,西湖大学把扩散缓存做成了“离线搜索题”
来源:互联网
时间:2026-06-28 13:24:57
扩散模型的生成效果越来越好,但一个老问题也越来越突出——
。
像
、
这类主流的扩散生成模型,现在能产出高分辨率、细节丰富又语义一致的图像和视频,水平确实不错。但它们有个共同的烦恼:推理依赖于多步去噪,每一步都要完整调用一次大型生成网络,算力成本居高不下。对图像生成来说,这会拖慢交互式创作和在线部署的响应速度;对视频生成来说,模型还要同时兼顾空间细节和时间连续性,计算量更大,对高效推理的需求也更迫切。
为了应对这个难题,
提出了
——一种面向固定计算预算的扩散模型step-level cache方法。和那些依赖运行时阈值判断的缓存方法不同,BudCache不再在推理过程中临时决定要算多少步,而是
先固定推理预算,再离线搜出这个预算下最优的缓存策略
。这样既能严格卡住推理成本,也能更好地保住最终生成质量。
论文一作是西湖大学AGI实验室研究助理
,指导老师是西湖大学助理教授
。

扩散模型加速中的cache技术
在多步去噪过程中,相邻步骤之间通常存在可以复用的信息。step-level cache方法正是利用这一点:选择一部分关键去噪步骤进行完整模型计算,在其他步骤中复用此前保存的输出或中间特征,从而减少真实模型调用次数。这种方法不需要重新训练模型,也不改变原始参数,对已经训练好的FLUX和Wan等模型来说,cache是一种
的推理加速技术,直接作用于采样过程。
像
、
这些已有方法已经证明了step-level cache的有效性。但当前不少方法仍然依赖启发式规则——在推理过程中根据某种运行时信号,判断当前步骤是否需要重新计算,或是否复用前一步的缓存结果。这种做法虽然能减少计算,但存在两个关键问题。
启发式阈值方法的实际计算量通常由运行时信号触发,不同输入可能触发不同数量的完整模型计算,最终推理延迟并不总是稳定可控。而真实部署场景往往更希望提前指定一个明确的计算预算,而不是被动接受不确定的开销。
某一步是否复用缓存,往往取决于当前步骤的差异度或误差信号,而不是直接面向最终生成结果做优化。换句话说,局部看似合理的缓存决策,并不一定能形成全局最优的生成轨迹。
BudCache的核心出发点正是重新定义这个问题:step-level cache的关键不只是判断“当前步骤能不能缓存”,而是在固定计算预算下决定“
”。为了判断一个缓存策略的好坏,BudCache不再只依赖单步的局部误差,而是直接考察该策略对应的最终生成结果是否接近完整计算的结果——也就是看采用某一组缓存步骤后,最后生成出来的图像或视频,是否接近完整计算的产物。
从启发式缓存到固定预算搜索
BudCache将step-level cache表述为一个
。具体来说,先固定模型实际前向计算的次数,也就是
。例如,在一个包含多个逻辑去噪步骤的采样过程中,可以预先设定只允许模型真实计算8次、9次或10次,其余步骤复用缓存结果。这样一来,推理成本在部署前就已确定,不会随输入变化产生不可控的延迟波动。
在此基础上,BudCache用一个
表示缓存策略:对每一个去噪步骤,如果对应位置为1,就执行完整模型计算;如果为0,则复用最近一次计算得到的缓存结果。由于mask中1的数量被固定为预算B,BudCache可以严格保证实际模型调用次数。
但找到最优mask并不简单。假设总共有K个逻辑去噪步骤,需要从中选B个执行真实计算,可能的组合数量非常大。同时,扩散模型的去噪过程具有明显的序列依赖:早期某一步的缓存决策可能影响后续整条采样轨迹,最终改变生成图像或视频的质量。因此,简单的局部贪心策略容易陷入局部最优,难以找到真正高质量的缓存方案。
为此,BudCache采用了“
”的混合搜索策略。第一阶段用模拟退火做全局探索:模拟退火允许搜索过程在早期接受一些暂时看起来更差的缓存策略,从而跳出局部最优,在更大的策略空间里寻找更好的候选。BudCache通过两类操作生成新策略:一种是
,交换一个计算步骤和一个缓存步骤的位置;另一种是
,把一个计算步骤移动到相邻位置。这两类操作都不改变总计算预算,搜索始终满足固定NFE约束。第二阶段用爬山算法做局部精修:当模拟退火找到较好的策略区域后,BudCache会检查当前策略附近的候选方案,并选择能让最终结果更接近完整计算结果的策略,持续进行直到附近没有更好的候选为止。
这种混合搜索方式让BudCache能在较小的离线成本下找到高质量缓存策略。更重要的是,搜索过程只在部署前进行一次;真正推理时直接使用已经搜好的固定缓存策略,
,因此不会引入额外的推理开销。
在缓存策略搜索完成后,BudCache还可以进一步进行
cache-aware schedule alignment
。这一步可以理解为一个可选的时间步校准模块:在固定缓存mask的前提下,对采样时间步做轻量级离线优化,让缓存采样器的最终输出更接近完整计算的结果。具体来说,完整计算的采样过程被视为
,缓存后的采样过程被视为
,BudCache优化student使用的时间步分布,使其最终生成结果尽可能接近teacher。优化完成后,新的时间步会和搜索得到的缓存策略一起固定下来,在推理阶段直接使用。需要强调的是,cache-aware schedule alignment并不改变原始生成模型,也不增加推理时的模型调用次数,只是在固定缓存策略的基础上让采样过程更好地适应缓存带来的轨迹变化。

主流图像与视频生成模型上的实验结果
研究团队首先在当前主流图像生成模型
上评估了BudCache,并在
上与多种代表性缓存方法做了比较,包括TeaCache、MagCache、LeMiCa、DiCache、ERTACache和TaylorSeer。结果显示,在相同推理预算下,BudCache能更好地保持生成质量。

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

视频生成方面,研究团队进一步在
上做了实验。视频生成通常需要处理更长的视觉序列和更复杂的时空关系,因此对推理加速提出了更高要求。实验结果表明,BudCache在视频生成中同样有效:在Wan2.1-T2V-1.3B上,完整推理需要约
,TeaCache约
,而BudCache只需约
。同时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