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

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. Code is available at https://github.com/thecrazymage/ReCache.

• The paper introduces a reinforcement learning policy for budget-aware caching in diffusion models, optimizing inference fidelity within explicit compute budgets.
• It employs a budget-conditioned Plackett-Luce distribution with Gumbel-Top-k sampling to dynamically select critical timesteps for full recomputation.
• Empirical results show significant reductions in LPIPS and FLOPs while preserving visual and semantic integrity in both image and video generation.

ReCache: Budget-Aware Scheduling for Diffusion Model Caching

Motivation and Problem Formulation

Diffusion models, including DiT-based architectures, provide state-of-the-art results for image and video generation but suffer from intensive inference costs due to repeated evaluation of large denoising networks. Feature caching has emerged as a practical method for amortizing redundant computations by reusing or approximating intermediate activations across timesteps. However, the question of when to recompute features and when to reuse cached activations (the caching schedule) is typically resolved via static rules (uniform spacing) or per-step heuristics, both lacking explicit budget control and global optimization with respect to final output quality.

ReCache introduces a policy-driven, budget-conditioned approach: given an explicit compute budget $k$, a learned policy selects the most critical timesteps for full recomputation, optimizing directly for perceptual/semantic output quality and fidelity to unaccelerated inference. This repositions scheduling as a reinforcement learning (RL) problem, sidestepping hand-crafted schedule design and enabling adaptivity across a range of budgets.

Approach: Reinforcement Learning for Schedule Selection

ReCache parameterizes the schedule policy as a budget-conditioned Plackett-Luce distribution over $k$-subsets of inference steps. Given a prompt, noise latent, caching mechanism, and explicit budget $k$, a lightweight MLP predicts step importance logits and samples $k$ steps (via Gumbel-Top-$k$) for full computation; the remainder are handled by caching, either direct feature reuse or forecasting (e.g., TaylorSeer/HiCache).

Figure 1: Overview of the ReCache method. The policy network predicts importance scores, selecting top-$k$ steps for full computation; it is trained via REINFORCE to optimize a fidelity-quality reward.

The reward comprises (a) a fidelity term (patchwise LPIPS) to match full inference and (b) a quality term (human preference score, HPSv2/ImageReward), weighted by a tunable coefficient $\alpha_{iq}$. Entropy regularization encourages exploration during RL training and avoids collapse to a deterministic schedule. Notably, ReCache does not require labeled data or backpropagation through the generative model: only reward evaluation on generated outputs is needed.

Figure 2: Impact of the quality reward coefficient $\alpha_{iq}$. Higher values improve perceptual scores at some cost to LPIPS, reflecting a controlled trade-off.

Empirical Results: Images and Videos

Evaluations span FLUX.1-dev (image), Wan2.1 and HunyuanVideo (video), using multiple modern caching mechanisms (DiCache, TaylorSeer, HiCache, DPCache, FORA, etc.) and budgets (e.g., 7, 9, 13 steps). Across all regimes, ReCache achieves strict computational budget compliance and outperforms static and adaptive baselines. Notable quantitative results include:

• On FLUX.1-dev (image) at $k=9$ (a $\times 5.04$ FLOPs reduction): LPIPS reduced by 31% (from 0.456 to 0.316) vs. DiCache; ImageReward and HPSv2 consistently improved.
• On Wan2.1 (video) at $k$0 speedup: LPIPS reduced by 65% (0.480 to 0.169), VBench score increased by 7% (70.4 to 76.0) over uniform HiCache.

Qualitative results demonstrate that, under tight budgets, heuristic schedules produce visible degradation (text corruption, object layout distortion, color artifacts), while ReCache preserves fidelity and semantic structure.

Figure 3: ReCache maintains text and visual fidelity under aggressive acceleration on FLUX.1-dev, succeeding where heuristic baselines fail.

Figure 4: Qualitative comparison on FLUX.1-dev; ReCache preserves essential structures, while uniform/heuristic schedules lead to text corruption and compositional failures.

Figure 5: Video comparison (Wan2.1, HunyuanVideo prompt: "zebra drinking water"); ReCache eliminates artifacts and restores quality.

Performance scales favorably with budget, with greatest gains in the most aggressive regimes. ReCache is mechanism-agnostic: training one policy per mechanism suffices to support multiple FLOPs budgets, and policies generalize well—even when transferring schedules across backbone architectures.

Ablation and Analysis

Ablation studies reveal:

• The nested structure of learned schedules: expanding budget never reorders prior selections, confirming an essentially budget-invariant step importance ranking.
• Entropy regularization is crucial under looser budgets for exploration and schedule discovery.
• Gumbel-Top-$k$1 sampling (vs. Bernoulli) delivers strict budget adherence and superior results.

Budget-adaptive policies outperform per-budget specialists, offering statistical robustness and practical deployment simplicity.

Figure 6: Cache step selection visualized across budgets with TaylorSeer and Taylor-DP; additive/nested schedule structure emerges.

Figure 7: Entropy regularization encourages optimal schedule exploration; however, excessive values can slow convergence.

Implications and Future Directions

ReCache achieves explicit, reliable control over computational budgets for diffusion inference, optimizing output quality in conjunction with training-free, mechanism-independent acceleration. This allows for:

• More predictable operational cost in deployed generative systems.
• Higher output quality under strict latency/constrained hardware settings.
• Composable integration with architectural/solver-based model acceleration (e.g., pruning, quantization, advanced solvers, distillation).

The structural stability and cross-backbone transferability of learned schedules suggest underlying denoising step importance is intrinsic to diffusion trajectories, not idiosyncratic to architecture or data. As diffusion-based generative modeling expands into longer sequences and higher resolutions, learned budget-aware scheduling becomes increasingly relevant.

Potential directions include schedule distillation for ultra-low budget regimes (1–4 steps), real-time schedule adaptation for interactive tasks, and joint optimization with downstream content controls or safety mechanisms.

Conclusion

ReCache reframes scheduling of diffusion model caching as a learned, budget-aware policy optimization—maximizing generation quality directly under explicit computational constraints. Empirically, ReCache outperforms static and heuristic baselines across image and video domains, multiple architectures, and caching strategies, while strictly adhering to resource limits. Its reinforcement learning-based schedule selection delivers practical acceleration for high-fidelity generative modeling, with implications for scalable, efficient deployment.