Search papers, labs, and topics across Lattice.
This paper introduces BAG (Budget-Aware Gating), a novel caching policy for Diffusion Transformers that integrates global budget awareness with dynamic, instance-adaptive feature reuse. By employing a lightweight gating network, BAG intelligently decides between executing full computations or reusing cached features based on the current budget state and local trajectory feedback. Experimental results on FLUX.1-dev and Wan2.1 show that BAG outperforms existing caching methods across multiple speedup tiers, demonstrating robustness across various conditions.
BAG achieves superior performance by dynamically balancing computational efficiency and budget constraints, outperforming traditional caching methods in diverse scenarios.
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 and Wan2.1 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.
