{"id":"990a589c-efbc-4209-bfd3-74661c982068","arxiv_id":"2608.09231","paper_version":1,"verdict":"CONDITIONAL","confidence":"HIGH","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":3,"one_line_summary":"A single learned gate, trained by copying offline-searched cache schedules, decides per step how to spend a diffusion sampling budget and beats static and heuristic caching baselines at matched compute.","lead":"This paper introduces BAG, a small learned gate that decides at every denoising step whether a diffusion model should recompute features or reuse cached ones, using both the remaining compute budget and how much the image has changed. It is trained by copying schedules found by expensive offline search, so the same gate can run cheaply at test time on new prompts and budgets.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The paper never quantifies the gap between the deployed gate and the offline-searched oracle on held-out prompts; the core 'distillation' claim rests on an unmeasured assumption.","rationale":"The reader's weakest assumption combines search quality and feature sufficiency. My concern is the second half: the paper never directly measures whether the gate reproduces the searched references on held-out prompts. This is the load-bearing link between the two stages; if the gate is not faithful, the method may still beat baselines, but the 'offline-to-online schedule distillation' explanation is unsupported, and the novelty claim would rest on the budget-exact inference rule rather than on learning. I do not think this warrants rejection: the end-to-end results are consistent across two backbones, multiple budgets, ablations, and held-out prompts, and the authors disclose the off-policy shift. However, the missing oracle-versus-gate comparison is exactly the experiment that would convert the mechanism claim from plausible to demonstrated. The reader's conditional verdict already requires stronger evidence; this check would supply it.","tokens_in":20842,"tokens_out":10003,"duration_ms":113008,"concrete_test":"Take 20 held-out FLUX DrawBench prompts and 10 held-out Wan prompts. For each, run the same matched-NFE offline search (Eq. 2) at B=9/13/20 (FLUX) and B=15/19/24 (Wan) to obtain oracle schedules and oracle PSNR/SSIM/LPIPS. Then deploy BAG on the same prompts and seeds at the same budgets. Report paired differences (BAG minus oracle) for all three metrics, and compare realized schedules by Hamming distance or placement similarity. If BAG is within about 0.5 dB PSNR or 0.02 LPIPS of the oracle, distillation is faithful; if the gap is large (e.g., more than 2 dB), the gate is not reproducing the search and the paper's mechanism claim is unsupported, though end-to-end superiority may still hold.","verdict_should_be":"UNCHANGED","load_bearing_attack":"BAG's central novelty is distilling offline-searched schedules (Eq. 2) into a compact online gate (Eq. 4). The paper validates the search (Fig. 2) and the final gate (Tabs. 1-2), but never compares the gate against the searched reference on the same held-out prompts at the same budget. All training examples come from reference rollouts; deployment states come from the gate's own rollouts, an off-policy shift the authors acknowledge but do not measure. If the realized schedules diverge from the references, the reported gains could come from the deterministic budget-exact boundary rules in Algorithm 2 or from the budget features alone, rather than from successful transfer of the searched policy. The six scalar features in Eq. (3) are also hand-picked; the paper shows they are necessary (Tab. 3) but not sufficient to express the search's decisions. Without an oracle-versus-gate comparison on held-out prompts, the core 'learning the scheduler, not the schedule' claim is not directly supported, even though the end-to-end numbers are positive.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","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.","tokens_in":21070,"tokens_out":7588,"duration_ms":79708,"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":[{"comment":"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.","section":"§3.2, §4.3; Eqs. (2)–(4)"},{"comment":"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.","section":"§4.1, Tabs. 1–2, Tab. A2"},{"comment":"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.","section":"Eq. (2), Tabs. 1–2"},{"comment":"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.","section":"§3.2, Eq. (3), Tab. A6"}],"minor_comments":[{"comment":"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.","section":"§3.2, Eq. (3)"},{"comment":"The resolution formatting '1024 2' should read '1024×1024'; the same rendering artifact may affect other dimensions in the supplementary material.","section":"§4.1"},{"comment":"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.","section":"App. B"},{"comment":"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.","section":"References"},{"comment":"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.","section":"Algorithm 2"}],"recommendation":"major_revision","confidential_remarks":"The stress-test concern about the unmeasured gate-versus-oracle gap lands: the manuscript's own App. A.2 and limitations section admit the off-policy shift but never measure it, and that gap is load-bearing for the distillation claim. The protocol is otherwise above the field median, and the PSNR/SSIM corroboration makes the empirical result credible. I would also ask the editor to watch the positioning against the concurrent ReCache work: the authors are appropriately candid about concurrency, but the 'what is deployed' distinction in Tab. A1 should be sharpened in the final version to avoid novelty disputes."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","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.\n\nThe 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.\n\nOther 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.\n\nWho 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.","headline":"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.","tokens_in":21583,"tokens_out":2637,"would_cite":true,"duration_ms":28348,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"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…","keywords":["diffusion caching","budget-aware scheduling","diffusion transformer acceleration","learned gating","schedule distillation","feature reuse","image and video generation","cache schedule"],"falsifier":"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.","tokens_in":20596,"feed_emoji":"⚡","tokens_out":8306,"duration_ms":80858,"temperature":0.7,"pith_summary":"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.","feed_headline":"One tiny gate beats static and heuristic caching at every matched budget","feed_subtitle":"A learned gate with under 1,000 parameters spends exactly its budget and wins every reconstruction metric at matched compute.","key_machinery":"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.","core_discovery":"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.","pith_inferences":["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."],"forward_implications":["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."],"supporting_citations":[{"why":"Defines the residual-reuse mechanism that the schedule operates on; BAG keeps this fixed and changes only when compute happens.","marker":"[6]"},{"why":"The strongest online-threshold baseline in the paper's runs; supplies the spectral-evolution signal that also seeds the offline search.","marker":"[7]"},{"why":"The representative online threshold heuristic for video caching and a matched-NFE baseline in the main tables.","marker":"[23]"},{"why":"The static-in-effect baseline on FLUX whose calibration curve produces the same schedule for every prompt; a key comparison and search start.","marker":"[30]"},{"why":"The budget-constrained schedule-search baseline, re-run under the paper's search to isolate its one-prompt calibration choice.","marker":"[20]"},{"why":"The forecasting-based reuse mechanism baseline, measured at matched NFE as a reference on the mechanism axis.","marker":"[24]"},{"why":"Simulated annealing, the optimizer used in the offline matched-NFE reference search.","marker":"[17]"},{"why":"LPIPS, the distance minimized by the offline search and one of the three reconstruction metrics.","marker":"[52]"}],"fun_headline_variants":["Budget-aware gate beats static schedules on every reconstruction metric","Learned gate with under 1K params wins at matched compute","Tiny gate adapts caching to each prompt and budget","BAG: one gate, any budget, best reconstruction","Instance-adaptive caching without the search cost"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"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.","fun_headline_variants_meta":{"raw":{"variants":["Budget-aware gate beats static schedules on every reconstruction metric","Learned gate with under 1K params wins at matched compute","Tiny gate adapts caching to each prompt and budget","BAG: one gate, any budget, best reconstruction","Instance-adaptive caching without the search cost"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000558,"raw_usage":{"total_tokens":2653,"prompt_tokens":942,"completion_tokens":1711,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":558,"completion_tokens_details":{"reasoning_tokens":1632}},"tokens_in":558,"tokens_out":1711,"duration_ms":11799,"temperature":1.0,"reasoning_tokens":1632,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-11T21:06:13.683288+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"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.","supporting_citations":[{"cited_title":"SeaCache: Spectral-evolution-aware cache for acceler- ating diffusion models","cited_arxiv_id":null,"evidence_quote":"The strongest online-threshold baseline in the paper's runs; supplies the spectral-evolution signal that also seeds the offline search."},{"cited_title":"Timestep embedding tells: It’s time to cache for video diffu- sion model","cited_arxiv_id":null,"evidence_quote":"The representative online threshold heuristic for video caching and a matched-NFE baseline in the main tables."},{"cited_title":"MagCache: Fast video generation with magnitude-aware cache.Advances in Neural Information Processing Systems, 38:34348–34380, 2025","cited_arxiv_id":null,"evidence_quote":"The static-in-effect baseline on FLUX whose calibration curve produces the same schedule for every prompt; a key comparison and search start."},{"cited_title":"Budget-constrained step-level diffusion caching","cited_arxiv_id":null,"evidence_quote":"The budget-constrained schedule-search baseline, re-run under the paper's search to isolate its one-prompt calibration choice."},{"cited_title":"From reusing to forecasting: Accelerat- ing diffusion models with TaylorSeers","cited_arxiv_id":null,"evidence_quote":"The forecasting-based reuse mechanism baseline, measured at matched NFE as a reference on the mechanism axis."},{"cited_title":"Opti- mization by simulated annealing.science, 220(4598):671–680,","cited_arxiv_id":null,"evidence_quote":"Simulated annealing, the optimizer used in the offline matched-NFE reference search."},{"cited_title":"The unreasonable effectiveness of deep fea- tures as a perceptual metric","cited_arxiv_id":null,"evidence_quote":"LPIPS, the distance minimized by the offline search and one of the three reconstruction metrics."}],"review_version":1}