# BAG: Budget-Aware Gating for Diffusion Caching

arXiv: https://arxiv.org/abs/2608.09231 · HTML: https://mlsanity.com/paper/2608.09231/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 52 of its 57 references that are in the map.
- First posted on arXiv 10 Aug 2026.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 52)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [The Unreasonable Effectiveness of Deep Features as a Perceptual Metric](https://mlsanity.com/paper/1801.03924/) (2018)
- [Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding](https://mlsanity.com/paper/2205.11487/) (2022)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Wan: Open and Advanced Large-Scale Video Generative Models](https://mlsanity.com/paper/2503.20314/) (2025)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [VBench: Comprehensive Benchmark Suite for Video Generative Models](https://mlsanity.com/paper/2311.17982/) (2023)
- [HunyuanVideo: A Systematic Framework For Large Video Generative Models](https://mlsanity.com/paper/2412.03603/) (2024)
- [ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2304.05977/) (2023)
- [Qwen-Image Technical Report](https://mlsanity.com/paper/2508.02324/) (2025)
- [One-step Diffusion with Distribution Matching Distillation](https://mlsanity.com/paper/2311.18828/) (2023)
- [GenEval: An Object-Focused Framework for Evaluating Text-to-Image Alignment](https://mlsanity.com/paper/2310.11513/) (2023)
- [Human Preference Score v2: A Solid Benchmark for Evaluating Human Preferences of Text-to-Image Synthesis](https://mlsanity.com/paper/2306.09341/) (2023)
- [Improved Distribution Matching Distillation for Fast Image Synthesis](https://mlsanity.com/paper/2405.14867/) (2024)
- [UniPC: A Unified Predictor-Corrector Framework for Fast Sampling of Diffusion Models](https://mlsanity.com/paper/2302.04867/) (2023)
- [DeepCache: Accelerating Diffusion Models for Free](https://mlsanity.com/paper/2312.00858/) (2023)
- [Z-Image: An Efficient Image Generation Foundation Model with Single-Stream Diffusion Transformer](https://mlsanity.com/paper/2511.22699/) (2025)
- [Q-Diffusion: Quantizing Diffusion Models](https://mlsanity.com/paper/2302.04304/) (2023)
- [Timestep Embedding Tells: It's Time to Cache for Video Diffusion Model](https://mlsanity.com/paper/2411.19108/) (2024)
- [Delta-DiT: A Training-Free Acceleration Method Tailored for Diffusion Transformers](https://mlsanity.com/paper/2406.01125/) (2024)
- [Hyper-SD: Trajectory Segmented Consistency Model for Efficient Image Synthesis](https://mlsanity.com/paper/2404.13686/) (2024)
- [Accelerating Diffusion Transformers with Token-wise Feature Caching](https://mlsanity.com/paper/2410.05317/) (2024)
- [Cache Me if You Can: Accelerating Diffusion Models through Block Caching](https://mlsanity.com/paper/2312.03209/) (2023)
- [Real-Time Video Generation with Pyramid Attention Broadcast](https://mlsanity.com/paper/2408.12588/) (2024)
- [Learning-to-Cache: Accelerating Diffusion Transformer via Layer Caching](https://mlsanity.com/paper/2406.01733/) (2024)
- [FORA: Fast-Forward Caching in Diffusion Transformer Acceleration](https://mlsanity.com/paper/2407.01425/) (2024)
- [From Reusing to Forecasting: Accelerating Diffusion Models with TaylorSeers](https://mlsanity.com/paper/2503.06923/) (2025)

## Built on by (most cited first, 1 of 1)

- [MORCA: Offline-to-Online Reinforcement Learning for Adaptive Cache Reuse in Video Diffusion Acceleration](https://mlsanity.com/paper/2610.10457/) (2026)

## Abstract (arXiv)

Diffusion caching is a lightweight strategy that accelerates Diffusion Transformers (DiTs) by reusing intermediate features across denoising steps, but existing paradigms face a fundamental trade-off: online heuristics lack global budget awareness, whereas static schedules lack instance adaptivity and fail to flexibly adapt to varying runtime budget constraints. To bridge this gap, we present BAG (Budget-Aware Gating), a novel caching policy that unifies global budget pacing with dynamic, instance-adaptive feature reuse. Rather than relying on hand-crafted rules, BAG employs a lightweight gating network that dynamically decides whether to execute a full computation or reuse cached features at each step by jointly conditioning on the budget state and local trajectory feedback. We train this policy via offline-to-online schedule distillation, transferring the decision-making of offline-searched schedules into a compact online gate. Extensive experiments on FLUX.1-dev, Wan2.1, and Qwen-Image-2512 demonstrate that BAG consistently outperforms state-of-the-art caching methods across various speedup tiers while remaining robust across different resolutions, seeds, and guidance scales. Code will be released.