用强化学习自动优化扩散模型的缓存策略,按预算控计算量并提升生成质量。
ReCache: Learning Budget-Aware Caching Schedules for Diffusion Models via REINFORCE

- 用REINFORCE算法学习按预算自适应的重计算调度策略。
- 在FLUX上减少5.04倍浮点运算,LPIPS降低31%;在Wan 2.1上提速2.6倍,质量提升显著。
- 无需标注数据,适配多种缓存机制,支持运行时动态调整预算。
现代扩散模型生成高质量图像与视频,但其迭代去噪过程推理成本高昂。特征缓存通过复用或预测相邻去噪步骤的中间激活来加速采样,利用反向轨迹中的计算冗余。本文聚焦缓存调度:选择哪些步骤需完全重计算。现有策略或为固定(如均匀),或基于每步误差启发式自适应;两者均将实际计算开销作为调参阈值的副作用,而非用户可直接指定的量。我们提出ReCache,反向设计:给定目标预算k,学习最大化生成质量的重计算调度,使计算量成为可直接控制的输入。ReCache通过策略梯度训练,避免对完整扩散推理进行反向传播,且无需标注数据。未缓存推理结果作为匹配目标,结合生成质量奖励。ReCache兼容任意缓存机制,包括特征复用与预测;每种机制仅需一个训练好的策略,即可在推理时跨预算自适应。实验表明,相比基线,ReCache持续领先:在FLUX上实现×5.04 FLOPs降低,LPIPS从0.456降至0.316(降31%);在Wan 2.1上实现约×2.6加速,LPIPS从0.480降至0.169(降65%),VBench得分提升7%(+5.6分,从70.4升至76.0)。代码已开源。
原文摘要 · Abstract (English)
Modern diffusion models generate high-quality images and videos, but their iterative denoising process makes inference expensive. Feature caching accelerates sampling by reusing or predicting intermediate activations across neighboring denoising steps, exploiting the redundancy of computations along the reverse trajectory. In this work, we focus on the caching schedule: selecting which denoising steps should be fully recomputed. Existing schedules are either fixed (e.g. uniform) or chosen adaptively from per-step error heuristics; in both cases, the actual compute cost is a side-effect of hand-tuned thresholds rather than a quantity the user can specify. We propose ReCache, which inverts this: given a target budget k, it learns the recomputation schedule that maximizes generation quality, turning compute into a directly controllable input. ReCache trains via policy gradients, sidestepping backpropagation through full diffusion inference, and uses no labelled data. Generations from uncached inference serve as matching targets, paired with a reward for generation quality. ReCache is compatible with any caching mechanism, including feature reuse and feature forecasting; for each mechanism, a single trained policy adapts across computational budgets at inference time. ReCache consistently outperforms scheduling baselines: under a $\times5.04$ FLOPs reduction on FLUX, it reduces LPIPS by 31% (from 0.456 to 0.316) compared to DiCache; on Wan 2.1 at a $\sim \times2.6$ speedup, it drops LPIPS by 65% (from 0.480 to 0.169) and boosts the VBench score by 7% (5.6 points, from 70.4 to 76.0) over uniform HiCache. Code is available at https://github.com/thecrazymage/ReCache.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。