BAG: Budget-Aware Gating for Diffusion Caching

Chi Zhang; Mingkun Lei; Tong Zhao; Yucheng Han

BAG: Budget-Aware Gating for Diffusion Caching

T0 review · 4 major / 5 minor · reviewed 2026-08-11 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗ Review JSON

Pith's one-line read A learned gate with under 1,000 parameters decides per denoising step whether to recompute or reuse cached features, and this paper shows it improves all three reconstruction metrics over static and heuristic schedulers at matched compute…

desk verdict A solid empirical method paper that delivers a practical 2-5x caching speedup with a tiny learned gate, but the distillation story would be stronger with a direct gate-vs-oracle comparison. read the letter →

arxiv 2608.09231 v2 pith:5NVAWUFJ submitted 2026-08-10 cs.CV

Tong Zhao , Mingkun Lei , Yucheng Han , Chi Zhang This is my paper · ORCID

classification cs.CV

keywords diffusioncachingbudget-awareschedulingtransformeraccelerationlearnedgatingscheduledistillationfeaturereuseimageandvideogenerationcache

verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

Diffusion caching speeds up image and video generation by reusing intermediate features across denoising steps instead of running the full network every time. Existing caching policies are either online heuristics, which react to the current trajectory but ignore how much budget remains, or static schedules, which meet the budget exactly but apply the same plan to every prompt. This paper claims that both lose quality compared with per-prompt searched schedules, and that the loss can be recovered by a learned gate of under 1,000 parameters. The gate takes six scalar signals—three budget indicators and three trajectory indicators—and decides at each step whether to compute or reuse, spending exactly the requested number of function evaluations by construction. On held-out prompts from an image and a video diffusion transformer, the paper reports that this single gate per backbone improves PSNR, SSIM, and LPIPS over the strongest baselines at every tested budget tier.

What carries the argument

The load-bearing object is the gate $g_\theta$, an MLP with under 1K parameters that outputs a compute-or-reuse decision from six scalar features: budget state $[B/T,\ (B-c_t)/B,\ (B-c_t)/(T-t)]$ and trajectory state $[t-t_{\mathrm{last}},\ \lVert x_t-x_{t_{\mathrm{last}}} \rVert / \lVert x_{t_{\mathrm{last}}} \rVert,\ \lVert x_t-x_{t-1} \rVert / \lVert x_{t-1} \rVert]$, where $B$ is the budget, $T$ the total steps, $c_t$ evaluations spent, and $x_t$ the post-patch-embedding token tensor. These features are exactly the two contexts prior paradigms each lack: global budget pacing and local rollout feedback. The gate is trained by offline-to-online schedule distillation: a matched-NFE search minimizes LPIPS over masks with exactly $B$ ones, its rollouts supply labeled per-step states, and a class-balanced BCE loss fits the gate in about a minute. At deployment, boundary rules—compute the remainder when budget equals remaining steps, reuse when budget is exhausted—make the realized NFE exactly $B$ without prescribing where the evaluations go.

What would settle it

Take a held-out prompt, run the paper's search to obtain the best schedule for a budget, then run the gate at that budget; if the gate's reconstruction is farther from the full-compute output than the searched schedule is, the distillation has lost information.

Watch

Extended reading notes

Core claim

The central claim is that one budget-aware, trajectory-conditioned gating network can reproduce most of the benefit of expensive per-prompt schedule search at runtime. For training prompts and budgets, the authors run a matched-NFE offline search that minimizes LPIPS distance to the full-compute output over binary cache masks with exactly $B$ evaluations, then decompose the searched rollouts into per-step examples. A small MLP is trained with binary cross-entropy to map the six-feature state to the search's compute-or-reuse decision. At inference, the same frozen gate makes closed-loop decisions on unseen prompts, with boundary rules that force the realized NFE to equal $B$. The paper reports that at every tested budget on FLUX.1-dev and Wan2.1, this gate improves all three reconstruction metrics over the strongest baseline using the same or fewer evaluations, with PSNR gains up to about 2.6 dB, and stays ahead under changed seeds, resolutions, guidance scales, and 25- or 28-step samplers.

Load-bearing premise

The load-bearing premise is that the offline search over compute-or-reuse patterns on the training prompts finds schedules good enough to supervise the gate, and that six hand-picked scalar signals can express what makes those schedules good; if either fails on unseen prompts and budgets, the gate inherits the search's blind spots or cannot represent the searched policy.

Editorial extensions

If this is right

One checkpoint per backbone serves every tested budget and step count; the same gate produced all reported BAG numbers without retraining or re-search.
The budget is a runtime input, so an operator can set the speedup directly and the gate will spend exactly that many evaluations, eliminating threshold sweeps.
Because BAG changes only the schedule, it can be combined with orthogonal accelerators such as fast ODE solvers or mechanism-level feature reuse.
The reported robustness to seed, resolution, guidance, and 25/28-step samplers means the trained gate transfers beyond its training configuration.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

A testable extension the paper does not run: train one gate on a small backbone family and check whether it transfers to a larger DiT without retraining.
Because the search objective is a free choice, the same distillation pipeline could be aimed at human-preference scores; the paper's own HPSv2 variant already points this way.
The paper acknowledges an off-policy state-distribution shift; on-policy relabeling in the style it cites is the obvious next step and would clarify how much headroom remains.
The budget features are all ratio-scaled, so the gate may generalize to arbitrary step counts and budgets; only a few out-of-range points are tested, so this is an inference, not a claim.

Share X Bluesky LinkedIn Reddit HN

Signed reviews

No signed human review yet.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Desk Editor's Note

read the letter

Bottom line: BAG is worth a serious referee. The deployed object is a <1K-parameter gate that reads six budget and trajectory scalars and decides compute-or-reuse at each step, trained by imitating offline-searched schedules. On FLUX and Wan, one gate per backbone beats strong baselines at matched or lower NFE across every tested budget, with the claimed budget-exact property by construction. That is a real result. The paper also does several things right: baselines' thresholds are swept on the test set (favoring baselines), prompts are held out, the ablations show both budget and trajectory inputs matter, and the off-policy shift is explicitly acknowledged. I believe the end-to-end claim.

The soft spots are real but not load-bearing. The biggest is exactly what your stress-test note says: the paper never compares the deployed gate against the offline-searched reference on the same held-out prompts. Without that, the 'we learned the scheduler, not the schedule' claim is partially inferred. The ablations soften the concern—removing either state half collapses the gate into a static or online paradigm and loses several dB—but they don't tell us how close the gate's decisions are to the search's. I'd add that check in revision: a column with oracle-policy results on the same test prompts, plus a measure of decision overlap between gate and oracle under deployment states. The paper already admits the off-policy shift; measuring it would close the loop.

Other weaknesses are smaller. The main tables have no error bars; significance is only tested against SeaCache in the appendix, and at the loosest tier the perceptual gains are not significant. LPIPS also serves as both the search objective and the headline metric, so part of the LPIPS gain is encoded in the labels. None of this changes my read that the method works, but it does mean the headline numbers are optimistic in a direction that matters. Code is promised but not out; the offline search is expensive (days on 8 H100s per backbone), so reproducibility without code will be painful.

Who is this for? People building diffusion serving stacks and researchers working on caching/acceleration. It deserves a full peer review, not a desk reject. I'd recommend conditional accept after code release, an oracle-vs-gate comparison, and proper variance reporting.

Referee Report

4 major / 5 minor

Summary. The paper proposes BAG, a budget-aware gating policy for diffusion transformer caching. An offline matched-NFE local search over cache masks (Eq. 2) generates reference schedules for each training prompt and budget; the per-step decisions of those schedules are distilled into a sub-1K-parameter MLP that reads six scalar budget and trajectory features (Eq. 3) and outputs a binary compute-or-reuse decision (Eq. 4). At inference the frozen gate runs closed-loop inside the sampler, with two deterministic boundary rules in Algorithm 2 guaranteeing that the realized NFE equals the requested budget. Experiments on FLUX.1-dev and Wan2.1 report consistent PSNR/SSIM/LPIPS improvements over TeaCache, MagCache, SeaCache, TaylorSeer, BudCache, and naive step reduction at matched or lower NFE, using one gate per backbone across multiple budgets and across shifts in step count, seed, resolution, and guidance.

Significance. If the results hold, BAG is a practically useful and well-engineered advance: the deployed policy is budget-exact by construction, instance-adaptive through trajectory feedback, and tiny (449–785 parameters), and one checkpoint per backbone serves many budgets. The protocol is unusually careful in several respects: evaluation is on held-out prompts (GenEval→DrawBench, VBench splits with overlap filtering), baseline thresholds are swept on the test set (which favors the baselines), BudCache is run under both its official protocol and a variant using the authors' search, and ablations in Tab. 3 and Tab. A5 show that both halves of the state matter. Tab. A6 is a strong control, showing that the budget-conditioned gate rather than the specific feature set is the key ingredient. The main weakness is that the mechanistic claim of 'learning the scheduler, not the schedule' is never tested against its own offline oracle: the paper validates the search (Fig. 2) and the final gate (Tabs. 1–2) but not the transfer between them. Main-table uncertainty is also unreported, with significance tests limited to SeaCache.

major comments (4)

[§3.2, §4.3; Eqs. (2)–(4)] The paper's central claim is that offline-searched schedules are distilled into a compact online gate, but no experiment compares the deployed gate with the searched reference on the same held-out prompts at the same budget. The off-policy state-distribution shift is explicitly acknowledged in App. A.2 ('distillation incurs an off-policy state-distribution shift') and in App. E, yet its effect is never quantified; the gains in Tabs. 1–2 could in principle be driven largely by the deterministic budget-exact boundary rules of Algorithm 2 rather than by successful transfer of the searched policy. Please report, for the held-out DrawBench/VBench prompts at each budget, the gate's per-prompt LPIPS (and ideally the agreement rate between gate decisions and reference mask) against the reference schedule, e.g., a scatter plot analogous to Fig. 2, together with the realized-NFE distribution. This is the direct evidence needed for the 'learning the scheduler, not the schedule' claim.
[§4.1, Tabs. 1–2, Tab. A2] The main tables report a single run per cell with no error bars or per-prompt variance, so the claim that BAG 'consistently outperforms' all baselines on all three metrics is supported only by point estimates. The paired Wilcoxon tests in Tab. A2 are against SeaCache only, and several margins are small (e.g., FLUX B=20 LPIPS 0.0835→0.0773, Wan B=24 SSIM 0.943→0.947). Per-prompt difference distributions or bootstrap intervals for the runner-up at each tier are needed to establish that the ranking is not noise; this matters especially for the looser tiers where the paper itself notes the gains are not significant on the perceptual metrics.
[Eq. (2), Tabs. 1–2] The offline search in Eq. (2) minimizes LPIPS, and LPIPS is one of the three headline metrics in the main tables, so part of the LPIPS gain is inherited from the search objective and the comparison on that metric is partly circular. PSNR and SSIM gains provide independent corroboration, and the paper should state this explicitly in the main text, or report the LPIPS improvement of the searched reference itself on held-out prompts (Fig. 2 currently shows only aggregate comparisons against one baseline per backbone) so the reader can separate objective-match effects from policy-transfer effects.
[§3.2, Eq. (3), Tab. A6] The six hand-picked features are shown to be necessary (Tab. 3, Tab. A5) but not sufficient to express the reference policy; Tab. A6 shows that SeaCache's own signals fed through the BAG gate recover most of the gain, which is reassuring, but the expressiveness loss of the specific feature set relative to the searched oracle is still unmeasured. In the requested oracle-versus-gate comparison, please also report the per-step agreement and the per-prompt quality gap, which would quantify how much of the reference policy is lost by the six-scalar state representation.

minor comments (5)

[§3.2, Eq. (3)] The 'post-patch-embedding token tensor' used for the trajectory features should be specified precisely (which block or layer, whether the conditional or unconditional branch is used under CFG, and how the norm is computed) so that the six features are exactly reproducible.
[§4.1] The resolution formatting '1024 2' should read '1024×1024'; the same rendering artifact may affect other dimensions in the supplementary material.
[App. B] The statement that 'SeaCache's δ=0.6 spends 26% of the steps at T=50 but 40% at T=25' should clarify whether this refers to a particular prompt set or to the average over the evaluation prompts, since Tab. 5 keeps δ=0.6 and lists different realized NFEs for the 25-step sampler.
[References] There are minor formatting issues in the bibliography, such as 'PmLR' in [33], lowercase 'science' in [17], and inconsistent arXiv/venue formatting; these should be normalized.
[Algorithm 2] The two boundary rules in lines 6–7 are correct because T−t>0 at t<T, but adding a one-line comment explaining why they cannot conflict would help readers verify the budget-exactness guarantee.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity: the gate is a learned policy trained from offline-searched labels and evaluated on held-out prompts; the budget-exactness is a counting constraint, and no headline metric reduces to a fitted parameter.

full rationale

I walked the derivation chain: Eq. (2) searches reference schedules that minimize LPIPS per (prompt, budget) cell; Eq. (4) trains a sub-1K-parameter gate by per-step BCE against those searched decisions; Algorithm 2 deploys the frozen gate with count-exact boundary rules. The headline tables measure PSNR/SSIM/LPIPS against same-seed full-compute outputs on held-out prompt sets (GenEval to DrawBench on FLUX; a disjoint VBench split on Wan), and the evaluation includes budgets not seen in training on FLUX (B=9, B=20) and on Wan (B=19). No equation in the paper defines a reported result in terms of the trained gate's own parameters or in terms of the evaluation metrics: the gate's decisions are not derived from the metric values, and the exact-NFE property is a forced counting rule, not a predicted quantity. The self-citations (BudCache as a baseline, DAgger as future work) are not load-bearing; BudCache is re-run under both its official protocol and the authors' search, with official numbers reported in App. B. The one legitimate closeness concern is that the offline search objective is LPIPS, the same perceptual metric used in the main evaluation, so part of the LPIPS gain is inherited from the label generation procedure rather than independently predicted. The paper discloses this trade-off, and on held-out prompts the gate must generalize searched decisions, which is an objective-alignment choice rather than a circular reduction. The absence of a direct oracle-versus-gate comparison on held-out prompts is a validation gap, not circularity. The paper is self-contained against external benchmarks, so the fair finding is no significant circularity.

Assumptions & free parameters 3 free parameters · 4 assumptions · 0 invented entities

The central claim rests on no invented physical or mathematical entities. The free parameters are the learned gate weights, the decision cutoff, and the offline search effort, all disclosed. The main assumptions are about search quality, feature sufficiency, evaluation objective, and cross-prompt generalization. None of these is circular by construction, but the LPIPS label/evaluation overlap is the closest thing to a fitted objective.

free parameters (3)

Gate parameters (FLUX 449, Wan 785) = 449 (FLUX), 785 (Wan)
The decision rule itself is learned by supervised classification on reference decisions; there is no closed-form justification for these weights.
Inference cutoff tau = 0.5
Chosen in advance as the default decision threshold. A sensitivity sweep shows LPIPS is lowest at the default among tested values, but the cutoff still influences which steps are computed.
Offline search effort per (prompt, budget) cell = about 450 rollouts on FLUX, about 100 on Wan
Chosen as a compute budget for reference search; the paper shows cutting search time degrades quality gracefully, so the number is a hand-set cost-quality knob.

assumptions (4)

ad hoc to paper The matched-NFE local search in Eq. (2) approximates the true best schedule well enough to serve as training label.
If the search returns poor or biased schedules, distillation inherits the bias. The paper checks reference quality against baselines but provides no proof of search optimality.
ad hoc to paper The six scalar features in Eq. (3), budget ratios and feature-drift norms, capture the state relevant to compute-or-reuse decisions.
Ablations show both halves of the state matter, but there is no guarantee that other signals would not perform better; the feature set is a design choice.
domain assumption LPIPS distance to the same-seed full-compute output is the right objective for schedule quality.
Reconstruction fidelity is one valid goal for caching, but preference metrics can disagree; the authors show BAG-R for preference, confirming the objective is a choice rather than a universal target.
domain assumption A gate trained on 96 GenEval prompts (FLUX) or 50 VBench prompts (Wan) generalizes to disjoint test prompts.
Held-out evaluation supports this empirically, but there is no theoretical guarantee, and the Wan gate is trained without a validation split due to the small label set.

how reviews work

0 comments

Cite this review

Pith. "Pith review of BAG: Budget-Aware Gating for Diffusion Caching." pith.science (2026). https://pith.science/paper/5NVAWUFJ

@misc{pith2026260809231,
  author       = {Pith},
  title        = {Pith review of: BAG: Budget-Aware Gating for Diffusion Caching},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/5NVAWUFJ}},
  note         = {Machine review of arXiv:2608.09231}
}

read the original 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.

Figures

Figures reproduced from arXiv: 2608.09231 by the authors.

**Figure 1.** BAG overview: learn the scheduler, not the schedule. Among the compared scheduling paradigms, BAG deploys the only per-step rule that reads both the remaining budget and the realized trajectory. (1) Search (training only, once per backbone): for each (prompt, budget) cell, a matched-NFE search produces a reference schedule: a full-horizon allocation of exactly B evaluations, optimized to preserve the fullcompute ou… view at source ↗

**Figure 2.** Both paradigms are suboptimal per prompt at matched NFE. (a) FLUX.1-dev, 100 prompts at three budgets: per-prompt reconstruction LPIPS of the searched schedule vs. MagCache at the same NFE; points below the dashed line favor the searched schedule. (b) Wan2.1, 50 prompts: against SeaCache, the mean LPIPS to the full-compute output drops by 22–32% (ticks: 90th percentile). At a given budget the deployed allocation is… view at source ↗

**Figure 3.** Qualitative results on FLUX.1-dev at the ∼5× (B=9) and ∼2.4× (B=20) tiers, with the 50-step original leftmost. At matched compute the accelerated baselines drift in composition and object identity, while BAG stays closer to the full-compute output. More qualitative results in App. D [PITH_FULL_IMAGE:figures/full_fig_p006_3.png] view at source ↗

Figures from the paper (2 more)

**Figure 4.** Figure 4: Qualitative results on Wan2.1 at the ∼3.4× tier (B=15, six frames). At matched NFE BAG follows the full-compute motion (the horse’s gait, the locomotive’s smoke) while SeaCache blurs fast motion and drifts. More qualitative results in App. D. Cache [30], SeaCache [7] (…

**Figure 5.** Figure 5: One BAG checkpoint across the budget range vs. naive step reduction (two of the 50 sweep prompts; the full sweep is in Figs. A2 and A3, metrics over all 50 in Tab. A8). Leftmost: the same-seed 50-step full-compute output. Columns run from NFE 24 down to 10; per case, t…

Discussion (0). Continue with ORCID to comment.

Pith tools

Reviewed August 11, 2026 · model on record in the stance chip above.