BAG: Budget-Aware Gating for Diffusion Caching

• URL: http://arxiv.org/abs/2608.09231v2
• Date: Fri, 14 Aug 2026 07:23:26 GMT
• Title: BAG: Budget-Aware Gating for Diffusion Caching
• Abstract summary: BAG (Budget-Aware Gating) is a novel caching policy that unifies global budget pacing with dynamic, instance-adaptive feature reuse.<n>We train this policy via offline-to-online schedule distillation, transferring the decision-making of offline-searched schedules into a compact online gate.
• Score: 13.854727125627107
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Evolving Cache Schedules for Fast Diffusion Policy Inference [6.879031540914794]
  Diffusion policies achieve strong visuomotor control by iteratively denoising action chunks, but repeated denoising makes real-time deployment computationally demanding.<n>We introduce Evolving Cache Schedules (EVO), a training-free acceleration framework that globally schedules cache refreshes via evolutionary search.<n>EVO preserves near-full performance while substantially reducing computation, achieving up to 8.05x action-generation speedup and reducing FLOPs from 15.77G to as low as 1.96G.
  arXiv  Detail & Related papers  (2026-07-22T15:40:11Z)
• AdaptiveSD A Stability-Aware, Runtime-Adaptive Speculative Decoding Framework with Multi-Policy Orchestration for CPU-Constrained LLM Inference [0.0]
  Fixed-depth speculative decoding has emerged as one promising technique, but in practice, it often leads to performance degradation due to bandwidth saturation, instability, or even catastrophic resource exhaustion resulting in system failure.<n>Our solution consists of four tightly-coupled components working together in a continuous feedback loop: an Adaptive Draft Controller, a Dynamic Policy Engine, a KV Cache Coordination Layer, and a Monitoring Engine.<n>While conventional approaches focus solely on maximizing throughput, we instead assess the effectiveness of our approach based on several key metrics including wasted drafted compute and inter-token latency dispersion alongside standard measures of speculative efficiency.
  arXiv  Detail & Related papers  (2026-07-04T13:39:40Z)
• Budget-Constrained Step-Level Diffusion Caching [9.784100090561045]
  Existing methods make per-step cache decisions using threshold-baseds, without directly optimizing for final output quality.<n>We propose BudCache, which fixes the compute budget in advance and search for the cache policy that best preserves the final output.<n>BudCache achieves better generation quality than deterministic caching baselines under the same inference budgets.
  arXiv  Detail & Related papers  (2026-06-11T15:45:05Z)
• BUDDY: BUdget-Driven DYnamic Depth Routing for Adaptive Large Language Model Inference [8.041559188948003]
  Buddy is a budget-driven dynamic depth routing framework for large language models.<n>It supports strict budget control, decode-time rerouting, and multiple budgets within a single trained model.
  arXiv  Detail & Related papers  (2026-06-08T14:06:35Z)
• Aligning Flow Map Policies with Optimal Q-Guidance [50.514994916864275]
  Flow map policies are designed for fast action generation by learning to take arbitrary-size jumps.<n>We derive FLOW MAP Q-GUIDANCE (FMQ), a principled closed-form learning target that is optimal for adapting offline flow map policies under a critic-guided trust-region constraint.<n>FMQ achieves state-of-the-art performance in offline-to-online RL, outperforming the previous one-step policy MVP by a relative improvement of 21.3% on the average success rate.
  arXiv  Detail & Related papers  (2026-05-12T17:12:29Z)
• WorldCache: Content-Aware Caching for Accelerated Video World Models [50.7543797435026]
  We present textbfWorldCache, a Perception-Constrained Dynamic CaCache framework.<n>WorldCache improves both when and how to reuse features.<n>On Cosmos-2.5-2B evaluated on PAI-Bench, WorldCache achieves textbf$2.3times$ inference speedup while preserving textbf99.4% of baseline quality.
  arXiv  Detail & Related papers  (2026-03-23T17:59:54Z)
• Adaptive Linear Path Model-Based Diffusion [52.84663832658799]
  We introduce Linear Path Model-Based Diffusion (LP-MBD), which replaces the variance-preserving schedule with a flow-matching linear probability path.<n>We also propose Adaptive LP-MBD (ALP-MBD), which leverages reinforcement learning to adjust diffusion steps and noise levels according to task complexity and environmental conditions.
  arXiv  Detail & Related papers  (2026-02-02T21:33:03Z)
• SPA-Cache: Singular Proxies for Adaptive Caching in Diffusion Language Models [56.45983529954998]
  We present SPA-Cache that jointly optimize update identification and budget allocation in DLM cache.<n>First, we derive a low-dimensional singular proxy that enables the identification of update-critical tokens in a low-dimensional subspace.<n>Second, we introduce an adaptive strategy that allocates fewer updates to stable layers without degrading generation quality.
  arXiv  Detail & Related papers  (2026-01-30T05:22:44Z)
• Sparse ActionGen: Accelerating Diffusion Policy with Real-time Pruning [9.92274274871221]
  Diffusion Policy has dominated action generation due to its strong capabilities for modeling multi-modal action distributions.<n>We propose $underlinetextbfS$parse $underlinetextbfA$ction$underlinetextbfG$en for extremely sparse action generation.
  arXiv  Detail & Related papers  (2026-01-19T09:50:36Z)
• TS-DP: Reinforcement Speculative Decoding For Temporal Adaptive Diffusion Policy Acceleration [64.32072516882947]
  Diffusion Policy excels in embodied control but suffers from high inference latency and computational cost.<n>We propose Temporal-aware Reinforcement-based Speculative Diffusion Policy (TS-DP)<n>TS-DP achieves up to 4.17 times faster inference with over 94% accepted drafts, reaching an inference frequency of 25 Hz.
  arXiv  Detail & Related papers  (2025-12-13T07:53:14Z)
• Iterative Refinement of Flow Policies in Probability Space for Online Reinforcement Learning [56.47948583452555]
  We introduce the Stepwise Flow Policy (SWFP) framework, founded on the key insight that discretizing the flow matching inference process via a fixed-step Euler scheme aligns it with the variational Jordan-Kinderlehrer-Otto principle from optimal transport.<n>SWFP decomposes the global flow into a sequence of small, incremental transformations between proximate distributions.<n>This decomposition yields an efficient algorithm that fine-tunes pre-trained flows via a cascade of small flow blocks, offering significant advantages.
  arXiv  Detail & Related papers  (2025-10-17T07:43:51Z)
• Semantic Caching for Low-Cost LLM Serving: From Offline Learning to Online Adaptation [54.61034867177997]
  Caching inference responses allows them to be retrieved without another forward pass through the Large Language Models.<n>Traditional exact-match caching overlooks the semantic similarity between queries, leading to unnecessary recomputation.<n>We present a principled, learning-based framework for semantic cache eviction under unknown query and cost distributions.
  arXiv  Detail & Related papers  (2025-08-11T06:53:27Z)
• Block-wise Adaptive Caching for Accelerating Diffusion Policy [10.641633189595302]
  Block-wise Adaptive Caching(BAC) is a method to accelerate Diffusion Policy by caching intermediate action features.<n>BAC achieves up to 3x inference speedup for free on robotic benchmarks.
  arXiv  Detail & Related papers  (2025-06-16T13:14:58Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.