ReCache: Learning Budget-Aware Caching Schedules for Diffusion Models via REINFORCE

Mishan Aliev, Eva Neudachina, Ilya Bykov, Aleksandr Oganov, Kirill Struminsky, Aibek Alanov, Denis Rakitin

Keywords: Diffusion Models, Inference Acceleration, Feature Caching, Reinforcement Learning, Generative AI

Abstract:

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.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 150