BAG: Budget-Aware Gating for Diffusion Caching

One gate · any budget · exact NFE

Tong Zhao^1,2, Mingkun Lei², Yucheng Han³, Chi Zhang^2,📧

¹Zhejiang University ²AGI Lab, Westlake University ³StepFun

📧 Corresponding author

Qwen-Image-2512 accelerated by a single BAG gate — the number on each image is its speedup (2.5× → 5.0×).

📰 News

• 2026-08 🖼️ Added Qwen-Image-2512 results.
• 2026-08 📄 Paper released on arXiv.

🔍 What is BAG?

Diffusion caching skips denoising steps by reusing cached features — but online heuristics ignore the compute budget, and static schedules ignore the trajectory dynamics. BAG distills offline-searched schedules into a gate of under 1K parameters that reads the remaining budget and the realized trajectory at every step, then decides compute-or-reuse on the fly:

• 🎯 Budget-exact — realized NFE equals the requested budget
• 🔁 One gate, any budget — a single checkpoint covers all evaluated budgets, resolutions, and step counts
• 📊 Strong results — outperforms strong static and online caching baselines at matched compute, on FLUX.1-dev, Wan2.1, and Qwen-Image-2512

🎞️ One checkpoint across the budget range

Columns run from NFE 24 down to 10 on FLUX.1-dev; leftmost is the same-seed 50-step output. Per case, top: naive step reduction, bottom: BAG at the same NFE — step reduction settles on a different image as steps shrink, while the single gate stays close to the full-compute output across the whole range.

✅ TODO

• Inference code + trained gates (FLUX.1-dev / Wan2.1 / Qwen-Image-2512)
• Offline search + gate training code (the full pipeline)

📖 Citation

@article{zhao2026bag,
  title={BAG: Budget-Aware Gating for Diffusion Caching},
  author={Zhao, Tong and Lei, Mingkun and Han, Yucheng and Zhang, Chi},
  journal={arXiv preprint arXiv:2608.09231},
  year={2026}
}