SoftCache: Adaptive Step Size Caching for Diffusion Transformers

Shimon Malnick; Nadav Orenstein; Ohad Fried; Shai Avidan

cs.CVOct 15, 2026

SoftCache: Adaptive Step Size Caching for Diffusion Transformers

Authors: Shimon Malnick, Nadav Orenstein, Ohad Fried, Shai Avidan

Organizations: Tel Aviv University · Reichman University

Abstract

Diffusion and flow models generate a sample over iterative steps, each a full pass through a large network, and that pass dominates generation cost. Training-free caching accelerates them by reusing activations from recent evaluations instead of recomputing them. Dynamic caching methods decide when to refresh by accumulating a cheap probe of the network input into a staleness signal and spending a fresh evaluation once it crosses a threshold. They perform hard caching: the signal is a binary switch between refresh and reuse, and every step, fresh or cached, advances by the prescribed step size. We observe that the signal carries much more than this switch. How far it sits below the threshold, its headroom, predicts the error induced by reuse. We propose SoftCache, which performs soft caching: instead of a binary switch, the headroom sets how far each cached step goes. Where the headroom is large, the sampler jumps past several scheduled nodes, and as a refresh nears, its steps shrink back toward the original step. These adaptive jumps may land between the scheduled nodes, taking the sampler off its fixed grid. SoftCache requires no training and adds no per-backbone settings: on FLUX for images and on HunyuanVideo and Wan2.1 for video, at speedups from $3.5\times$ to $6.9\times$ , it is ahead of every baseline on all four fidelity metrics.

Figures & tables

Panel of Figure 1 — Figure 1: SoftCache: how far, not just whether. Hard caching walks the fixed grid of scheduled nodes one node at a time, deciding only whether to refresh or reuse . SoftCache also sets the step size: a reused step the cache still supports skips several nodes and lands between them (inset).

Panel of Figure 2 — Figure 2: The refresh signal carries more than a binary decision. (a) On uncached FLUX trajectories, the error $E_{i}$ of reusing the cached residual (binned median, interquartile band) grows steadily with the refresh signal $A_{i}$ . (b) Hard caching only compares $A_{i}$ with $\tau$ and takes the scheduled step $\Delta_{i}$ at every cached node. SoftCache also reads how far below $\tau$ it sits: the lower $A_{i}$ , the longer the step, up to $J\Delta_{i}$ ( Eq. 9 ). Stars mark the same $A_{i}$ in both panels: a low error buys a long jump , a high one only a short jump , and past $\tau$ the node refreshes .

Figure 3: Step placement on FLUX. Fresh and cached steps for 25 prompts, one per row.

Panel of Figure 4 — Figure 4: FLUX on DrawBench at $\approx\!5.2\times$ speedup. The speedup under each method is its band mean. Red ellipses mark structural errors, and each inset enlarges one region of its tile where shown.

Method	Efficiency	Quality
Method	Speedup	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	DrmSim $\downarrow$
TeaCache	$4.02\times$	16.02	0.660	0.396	0.140
TaylorSeer	$3.70\times$	15.50	0.637	0.397	0.128
SeaCache	$3.25\times$	22.41	0.830	0.173	0.037
DPCache	$3.49\times$	21.27	0.803	0.182	0.041
SoftCache	$3.79\times$	23.20	0.847	0.146	0.028

Table 1: FLUX DrawBench-200. One block per speed band. DrmSim is DreamSim.

Method	Efficiency	Quality
Method	Speedup	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$	DrmSim $\downarrow$
SeaCache	$3.53\times$	23.63	0.778	0.203	0.102
DPCache	$3.28\times$	23.70	0.786	0.182	0.092
SoftCache	$3.52\times$	26.48	0.843	0.141	0.067
SeaCache	$5.33\times$	21.82	0.713	0.292	0.162
DPCache	$5.42\times$	20.34	0.682	0.293	0.162

Table 2: HunyuanVideo VBench 480p. One block per matched speed band. DrmSim is DreamSim.

Panel of Figure 5 — Figure 5: HunyuanVideo on VBench at $\approx\!5.4\times$ speedup. Each clip shows its first, middle and last frame, and the speedup beside each method is its band mean. Insets enlarge one region, and red ellipses mark a person the prompt does not ask for and a malformed duck in mid-air.

Panel of Figure 6 — Figure 6: Component leave-one-out (FLUX, $50$ DPG-Bench prompts). Top: one prompt at $\approx\!3\times$ with its per-pixel $L_{1}$ error to the uncached image. Bottom: means at matched wall-clock speedup. Every removal trails SoftCache at $3\times$ , and removing the jump costs the most at $7\times$ . Tab. C.4.1 adds $5\times$ and residual motion.

Appendix figures & tables23 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Parameter	Swept values
FLUX
SoftCache, $\beta{=}1$	$\tau$	$0.3,\mathbf{0.5},0.8,1.1$
SoftCache, $\beta{=}0.5$	$\tau$	$\mathbf{0.8},\mathbf{1.1}$
SeaCache	$\tau$	$0.3,\mathbf{0.5},0.8,\mathbf{1.1},\mathbf{1.5},2$
TeaCache	rel. L1 threshold	$0.5,\mathbf{0.8},\mathbf{1.2},\mathbf{1.6}$
TaylorSeer	interval $N$	$\mathbf{6},7,10,\mathbf{16},\mathbf{25}$

Appendix

Table A.1.1: Swept speed parameter per method. Bold values appear in Sec. B.1 .

VBench dimension	Population ( $n{=}946$ )	$n{=}100$ subset
human_action	100 (10.6%)	10 (10%)
temporal_style	100 (10.6%)	12 (12%)
overall_consistency	93 (9.8%)	9 (9%)
aesthetic_quality	93 (9.8%)	9 (9%)
imaging_quality	93 (9.8%)	9 (9%)
appearance_style	90 (9.5%)	9 (9%)

Appendix

Table A.2.1: VBench dimension coverage. Number of prompts carrying each dimension label in the

946

-prompt set and in our

100

-prompt subset. A prompt may carry several labels, so columns do not sum to the prompt count.

Method	Efficiency				Quality
Method	Speedup	Latency (s) ( $\downarrow$ )	TFLOPs ( $\downarrow$ )	Mem (GB) ( $\downarrow$ )	PSNR ( $\uparrow$ )	SSIM ( $\uparrow$ )	LPIPS ( $\downarrow$ )	DreamSim ( $\downarrow$ )
FLUX	$1.00\times$	15.49	3735	33.8	$-$	$-$	$-$	$-$
TeaCache ( Liu et al., 2025a )	$4.02\times$	3.86	931	33.8	16.02	0.660	0.396	0.140
TaylorSeer ( Liu et al., 2025b )	$3.70\times$	4.18	760	37.9	15.50	0.637	0.397	0.128
SeaCache ( Chung et al., 2026 )	$3.25\times$	4.77	1132	33.8	22.41	0.830	0.173	0.037
DPCache ( Cui et al., 2026 )	$3.49\times$	4.44	983	34.0	21.27	0.803	0.182	0.041

Appendix

Table B.1.1: FLUX DrawBench-200. Each block groups the runs that land nearest one speed band, and the speedup each method actually reached is given per row. All four metrics measure fidelity to the uncached FLUX sample for the same prompt and seed. The efficiency columns add the latency, TFLOPs and peak memory behind each speedup.

Method	Efficiency				Quality
Method	Speedup	Latency (s) ( $\downarrow$ )	TFLOPs ( $\downarrow$ )	Mem (GB) ( $\downarrow$ )	PSNR ( $\uparrow$ )	SSIM ( $\uparrow$ )	LPIPS ( $\downarrow$ )	DreamSim ( $\downarrow$ )
HunyuanVideo	$1.00\times$	212.99	30579	55.5	$-$	$-$	$-$	$-$
SeaCache ( Chung et al., 2026 )	$3.53\times$	60.26	8906	55.5	23.63	0.778	0.203	0.102
DPCache ( Cui et al., 2026 )	$3.28\times$	64.95	9744	56.2	23.70	0.786	0.182	0.092
SoftCache	$3.52\times$	60.52	9114	55.5	26.48	0.843	0.141	0.067
SeaCache ( Chung et al., 2026 )	$5.33\times$	39.98	6172	55.5	21.82	0.713	0.292	0.162

Appendix

Table B.1.2: HunyuanVideo VBench 480p. Blocks are three increasing speed bands. All four metrics measure fidelity to the uncached HunyuanVideo clip for the same prompt and seed, averaged over all

65

frames. The efficiency columns add the latency, TFLOPs and peak memory behind each speedup.

Method	Efficiency				Quality
Method	Speedup	Latency (s) ( $\downarrow$ )	TFLOPs ( $\downarrow$ )	Mem (GB) ( $\downarrow$ )	PSNR ( $\uparrow$ )	SSIM ( $\uparrow$ )	LPIPS ( $\downarrow$ )	DreamSim ( $\downarrow$ )
Wan2.1	$1.00\times$	440.02	28578	6.7	$-$	$-$	$-$	$-$
SeaCache ( Chung et al., 2026 )	$3.73\times$	117.93	7771	7.2	24.89	0.838	0.139	0.058
DPCache ( Cui et al., 2026 ) †	$3.60\times$	122.35	7209	5.9	23.20	0.791	0.155	0.071
SoftCache	$3.83\times$	115.03	7770	7.1	25.95	0.869	0.102	0.038
SeaCache ( Chung et al., 2026 )	$5.30\times$	83.03	5522	7.2	21.17	0.736	0.262	0.144

Appendix

Table B.1.3: Wan2.1-T2V-1.3B, VBench 100-prompt subset. Blocks are two increasing speed bands. All four metrics measure fidelity to the uncached Wan2.1 clip for the same prompt and seed.

Figure B.2.1: FLUX DrawBench-200. The four fidelity metrics against achieved speedup for TeaCache, TaylorSeer, SeaCache, DPCache and SoftCache. Each faint dot is one prompt. Solid markers are per-method means, connected across the three speed bands of Sec. B.1 .

Figure B.2.2: HunyuanVideo VBench 480p. The four fidelity metrics against achieved speedup. Each faint dot is one prompt. Solid markers are per-method means, connected across the three speed bands of Sec. B.1 . SoftCache holds the best mean PSNR, LPIPS, SSIM and DreamSim at all three.

Panel of Figure 16 — Figure B.3.1: ImageReward on DrawBench-200 with FLUX (one seed, $200$ prompts), paired by prompt. (a) Per-prompt change relative to the uncached image, for each cache at the $\approx\!3.7\times$ and $\approx\!5.2\times$ bands of Tab. 1 . (b) Mean per-prompt difference between SoftCache and each baseline in the same band, with bars at one standard deviation of the per-prompt differences. ImageReward separates the methods only where the fidelity metrics already do: TaylorSeer at the $\approx\!5.2\times$ band collapses, and every other difference is small against its spread.

Method	Speedup	Subj. ( $\uparrow$ )	Bkg. ( $\uparrow$ )	Flicker ( $\uparrow$ )	Smooth ( $\uparrow$ )	Aesth. ( $\uparrow$ )	Imag. ( $\uparrow$ )	Dynamic ( $\uparrow$ )	Quality ( $\uparrow$ )
HunyuanVideo	$1.00\times$	0.963	0.965	0.991	0.992	0.601	0.619	0.347	0.810
SeaCache ( Chung et al., 2026 )	$3.53\times$	0.962	0.968	0.991	0.992	0.589	0.591	0.306	0.801
DPCache † ( Cui et al., 2026 )	$3.28\times$	0.964	0.964	0.991	0.993	0.595	0.607	0.319	0.805
SoftCache	$3.52\times$	0.964	0.968	0.992	0.993	0.592	0.597	0.333	0.805
SeaCache ( Chung et al., 2026 )	$5.33\times$	0.960	0.970	0.992	0.992	0.575	0.544	0.306	0.791
DPCache † ( Cui et al., 2026 )	$5.42\times$	0.961	0.959	0.989	0.991	0.589	0.597	0.292	0.797

Appendix

Table B.3.1: HunyuanVideo VBench quality dimensions.

Method	Object ( $\uparrow$ )	Multi-obj. ( $\uparrow$ )	Action ( $\uparrow$ )	Color ( $\uparrow$ )	Spatial ( $\uparrow$ )	Scene ( $\uparrow$ )	Appear. ( $\uparrow$ )	Temporal ( $\uparrow$ )	Overall ( $\uparrow$ )	Semantic ( $\uparrow$ )	Total ( $\uparrow$ )
HunyuanVideo	0.820	0.658	0.900	0.856	0.621	0.416	0.198	0.237	0.267	0.715	0.791
SeaCache ( Chung et al., 2026 )	0.860	0.683	0.910	0.888	0.603	0.403	0.200	0.237	0.264	0.723	0.785
DPCache † ( Cui et al., 2026 )	0.843	0.704	0.910	0.863	0.586	0.388	0.197	0.239	0.264	0.716	0.787
SoftCache	0.832	0.707	0.920	0.900	0.600	0.389	0.199	0.234	0.263	0.722	0.788
SeaCache ( Chung et al., 2026 )	0.779	0.606	0.890	0.846	0.574	0.355	0.206	0.227	0.261	0.687	0.771
DPCache † ( Cui et al., 2026 )	0.831	0.668	0.910	0.874	0.619	0.392	0.199	0.238	0.264	0.717	0.781

Appendix

Table B.3.2: HunyuanVideo VBench semantic dimensions and total.

	SoftCache, SeaCache probe			SoftCache, TeaCache probe
$\tau$	Speedup	$\Delta$ PSNR	$\Delta$ PSNR	Speedup	$\Delta$ PSNR	$\Delta$ PSNR	$\Delta$ LPIPS
		vs. SeaCache	vs. TeaCache		vs. SeaCache	vs. TeaCache	vs. TeaCache
0.3	2.95 $\times$	$+2.93$	$+10.00$	3.13 $\times$	$-3.61$	$\mathbf{+2.92}$	$+0.138$
0.5	3.80 $\times$	$+2.09$	$+7.75$	4.44 $\times$	$-3.77$	$\mathbf{+1.32}$	$+0.094$
0.8	5.43 $\times$	$+0.41$	$+5.12$	6.09 $\times$	$-4.16$	$+0.29$	$+0.079$
1.1	6.89 $\times$	$+0.04$	$+3.94$	7.18 $\times$	–	$+0.25$	$+0.037$

Appendix

Table C.2.1: SoftCache on either probe. SoftCache driven by the SeaCache probe (used throughout the paper) or the TeaCache probe, scored at matched wall-clock budget against both caches. Positive values favour SoftCache. Means over 50 DPG prompts, FLUX.1-dev, 50 steps. A dash marks a budget beyond the SeaCache frontier, where the comparison would be extrapolated.

	LPIPS ( $\downarrow$ )			DreamSim ( $\downarrow$ )			PSNR ( $\uparrow$ )
	$3\times$	$5\times$	$7\times$	$3\times$	$5\times$	$7\times$	$3\times$	$5\times$	$7\times$
w/o jumps ( $J{=}1$ )	0.103	0.245	0.548	0.018	0.048	0.199	24.30	18.88	13.67
w/o headroom ( $H_{i}{=}1$ )	0.122	0.215	0.330	0.016	0.041	0.082	25.70	21.36	18.17
w/o re-spacing ( $m_{i}$ fixed)	0.109	0.231	0.371	0.018	0.045	0.101	24.64	19.65	16.64
w/o residual motion ( $\beta{=}0$ )	0.117	0.246	0.363	0.018	0.051	0.087	25.50	20.24	18.28
SoftCache	0.093	0.207	0.329	0.013	0.037	0.078	26.50	20.61	17.85

Appendix

Table C.4.1: Component leave-one-out (FLUX, DPG-50). Each row removes one component of SoftCache at matched wall-clock speedup, with

\beta{=}1

3\times

and

\beta{=}0.5

above it.

Variant	$-\Delta$ LPIPS			$\Delta$ PSNR
Variant	$3\times$	$5\times$	$7\times$	$3\times$	$5\times$	$7\times$
$\beta{=}0.5$	$-0.009$	$+0.004$	$+0.080$	$-0.31$	$+0.75$	$+1.65$
cap $J{=}3$	$+0.005$	$-0.020$	$-0.085$	$-0.47$	$-0.63$	$-1.28$
cap $J{=}5$	$+0.004$	$-0.014$	$-0.067$	$-0.50$	$-0.37$	$-1.02$
cap $J{=}12$	$-0.006$	$-0.014$	$+0.021$	$+0.83$	$+0.18$	$+0.47$
cap $J{=}20$	$-0.025$	$-0.005$	$+0.050$	$+0.68$	$+1.19$	$+1.12$

Appendix

Table C.4.2: One setting at a time (FLUX, DPG-50). Change against SoftCache at

\beta{=}1

at matched wall-clock speedup, signed so that positive is better.

Placement	$\tau{=}0.3$			$\tau{=}0.8$
Placement	Speedup	Jumps	$\Delta$ PSNR	Speedup	Jumps	$\Delta$ PSNR
SoftCache	$2.91\times$	5.0	$+3.08$	$5.34\times$	5.9	$+0.35$
Evenly spaced	$2.56\times$	2.0	$+0.92$	$4.74\times$	3.4	$-0.26$
Near the end	$2.57\times$	3.0	$+1.88$	$5.25\times$	4.0	$+0.10$
Random (best of 5)	$2.43\times$	1.0	$+1.69$	$5.28\times$	5.0	$+0.22$
Random (mean of 5)	$2.49\times$	2.0	$+0.98$	$5.23\times$	4.0	$-0.99$

Appendix

Table C.4.3: Where the jumps fall (FLUX, DPG-50). PSNR gain over SeaCache at each variant’s own wall-clock speedup. Jumps is the mean number of jumps per image.

Panel of Figure 23 — Figure D.1.1: Higher PSNR, blurrier frames. Four frames of one HunyuanVideo 480p clip at a similar speedup, with green marking the better value on each metric. SeaCache leads on PSNR, yet the inset shows its foliage blurred where SoftCache keeps the leaves. SSIM, LPIPS and DreamSim rank the clip the other way.

Panel of Figure 24 — Figure D.2.1: FLUX at the $\approx\!6.6\times$ band, eight DrawBench prompts. Each method’s achieved speedup over DrawBench-200 is under its header ( Tab. B.1.1 ).

Panel of Figure 25 — Figure D.2.2: FLUX on DrawBench at $\approx\!5.2\times$ speedup, two prompts enlarged. The row under each prompt enlarges the region in its red box, the frame of the clock and the middle cat, where SoftCache stays closest to the reference.

Panel of Figure 26 — Figure D.2.3: HunyuanVideo on VBench at $\approx\!5.4\times$ speedup, two further clips. Same band, methods and frames as Fig. 5 . The red ellipse marks a second broom in DPCache, and the inset enlarges the back of the zebra’s head.

Panel of Figure 27 — Figure D.2.4: HunyuanVideo on VBench at $\approx\!3.5\times$ speedup. Frames 1, 33 and 65 of 65, with a zoom inset of the boxed region on the first frame. The red ellipse marks a hallucinated second hand in DPCache.

Panel of Figure 28 — Figure D.2.5: Wan2.1 on VBench at $\approx\!3.7\times$ speedup. Frames 1, 41 and 81 of 81, with each method’s mean speedup at left ( Tab. B.1.3 ) and a zoom inset of the pool on the middle frame. SeaCache and DPCache distort the pool, while SoftCache stays closest to the reference.

Panel of Figure 29 — Figure D.2.6: HunyuanVideo at native 720p, $\approx\!4.5\times$ speedup. Frames 1, 65 and 129 of 129, the last frame carrying a zoom inset of the corgi’s face.

Prompt id	Category	Text
Fig. 4
drawbench_133	Positional	A zebra to the right of a fire hydrant.
drawbench_079	Descriptions	A device consisting of a circular canopy of cloth on a folding metal frame supported by a central rod, used as protection against rain or sometimes sun.
drawbench_027	Conflicting	A bird scaring a scarecrow.
Fig. D.2.1
drawbench_098	Gary Marcus et al.	An elephant is behind a tree. You can see the trunk on one side and the back legs on the other.

Appendix

Table D.2.1: Prompts of the image qualitative figures, verbatim.

Prompt id	VBench dimension	Text
Fig. 5
vbad077	multiple_objects	a dog and a horse
vbad687	overall_consistency, aesthetic_quality, imaging_quality	An astronaut feeding ducks on a sunny afternoon, reflection from the water.
Fig. D.2.3
vbad168	human_action	A person is sweeping floor
vbad351	object_class	a zebra

Appendix

Table D.2.2: Prompts of the video qualitative figures, verbatim.

Diffusion Model Caching

Diffusion model caching reuses features, activations, or attention states across denoising steps to reduce inference cost while preserving output quality.

See all 106 papers

Jul 31, 2026

OnlineCache: Learning Dynamic Caching Policies with Error Correction for Efficient Diffusion Inference

Diffusion models have revolutionized generative tasks but incur high latency due to iterative denoising. While cache-based strategies accelerate inference by reusing intermediate features, they largely rely on static, sample-agnostic schedules. We argue that this rigidity overlooks two facts empirically validated in this paper: (i) generation difficulty varies across prompts, requiring adaptive resource allocation--complex inputs demand more computation while simpler ones require less; (ii) error sensitivity fluctuates across timesteps, where static policies may cache high-error steps or waste computation on low-error ones. We therefore propose OnlineCache, a dynamic caching framework that jointly learns when to cache and how to correct approximation errors. We leverage policy gradient to train a lightweight network for adaptive speed-quality trade-offs, and incorporate a learnable corrector to mitigate caching-induced errors. Both modules are jointly optimized under a bilevel optimization framework, with the policy targeting global generation quality and the corrector minimizing local errors. Our method automatically allocates computational resources across both samples and timesteps, improving overall generation quality. Extensive experiments demonstrate clear superiority. On FLUX.1-dev model, OnlineCache achieves nearly 3 speedup while preserving generation fidelity. On DiT and CogVideoX, it similarly delivers competitive acceleration without compromising quality; across all scenarios, it consistently outperforms existing cache-based acceleration baselines.

Zhikang Xie, Xichen Ye, Yifan Wu +5

College of Computer Science and Artificial Intelligence, Fudan University · School of Data Science, Fudan University · 3Shanghai Key Laboratory of Intelligent Information Processing

Diffusion Model Caching Diffusion Model Inference Acceleration Diffusion Models

Ask PDF

Oct 2, 2026

Rethinking What to Cache in Few-Step Diffusion Transformers: Solver-Aware Target Selection

Diffusion Transformers (DiTs) can generate high-quality images and videos, but generating each sample requires multiple costly DiT forward passes. Two common ways to accelerate DiT sampling are step distillation, which reduces the number of sampling steps, and caching, which skips some DiT evaluations by reusing a tensor computed at an earlier step. Most caching methods decide in advance which tensor to reuse. After distillation, adjacent sampling steps are farther apart. Reusing a tensor across this larger gap introduces more error, so choosing what to cache becomes especially important. We therefore introduce AutoTarget, a method that chooses the cached tensor for a given model, solver, and reuse schedule. AutoTarget uses a small set of runs without cache reuse to measure the error caused by reusing each candidate tensor, then selects the candidate with the lowest error. We also analyze how an error at one reuse step affects the final sample. For Euler sampling, we identify cache targets that produce the same trajectory and show why a stored solver update may not. Experiments on distilled image and video DiTs show that the best cache target changes with the model, image resolution, and solver. AutoTarget reduces DiT evaluations and retained cache storage. Generation quality remains close to the corresponding uncached run. On the tested PixArt-LCM and FLUX.1-schnell settings, its calibration ranking matches the ranking from held-out cached runs. To help others reproduce the method, we provide its core implementation on GitHub at https://github.com/wali1024-offical/AutoTarget.

Shuo Yang, Lihao Fang, Yi Zhang +5

Beijing University of Chemical Technology · ROBOT-XZC · Tsinghua University

Diffusion Model Caching Diffusion Model Inference Acceleration Diffusion Transformer Few-Step Diffusion Sampling

Ask PDF

Aug 10, 2026

BAG: Budget-Aware Gating for Diffusion Caching

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.

Tong Zhao, Mingkun Lei, Yucheng Han +1

Zhejiang University · AGI Lab, Westlake University · StepFun

Diffusion Transformer Diffusion Model Caching Diffusion Model Acceleration

Ask PDF

Abstract

Diffusion Model Caching

Explore similar work

OnlineCache: Learning Dynamic Caching Policies with Error Correction for Efficient Diffusion Inference

Rethinking What to Cache in Few-Step Diffusion Transformers: Solver-Aware Target Selection

BAG: Budget-Aware Gating for Diffusion Caching