Few-Step Diffusion Sampling Through Instance-Aware Discretizations

Liangyu Yuan; Ruoyu Wang; Tong Zhao; Dingwen Fu; Mingkun Lei; Beier Zhu; Chi Zhang

cs.CVMar 18, 2026

Few-Step Diffusion Sampling Through Instance-Aware Discretizations

Authors: Liangyu Yuan, Ruoyu Wang, Tong Zhao, Dingwen Fu, Mingkun Lei, Beier Zhu, Chi Zhang

Organizations: Westlake University · Tongji University · Zhejiang University · Nanyang Technological University

Abstract

Diffusion and flow matching models generate high-fidelity data by simulating paths defined by Ordinary or Stochastic Differential Equations (ODEs/SDEs), starting from a tractable prior distribution. The probability flow ODE formulation enables the use of advanced numerical solvers to accelerate sampling. Orthogonal yet vital to solver design is the discretization strategy. While early approaches employed handcrafted heuristics and recent methods adopt optimization-based techniques, most existing strategies enforce a globally shared timestep schedule across all samples. This uniform treatment fails to account for instance-specific complexity in the generative process, potentially limiting performance. Motivated by controlled experiments on synthetic data, which reveals the suboptimality of global schedules under instance-specific dynamics, we propose an instance-aware discretization framework. Our method learns to adapt timestep allocations based on input-dependent priors, extending gradient-based discretization search to the conditional generative setting. Empirical results across diverse settings, including synthetic data, pixel-space diffusion, latent-space images and video flow matching models, demonstrate that our method consistently improves generation quality with marginal tuning cost compared to training and negligible inference overhead.

Figures & tables

Figure 1: Our effective instance-aware discretization improves sampling quality, by generating a tailored discretization $\xi^{\phi}$ for each initial noise $\mathbf{x}_{T}$ and condition $\mathbf{c}$ , outperforming heuristic and globally optimized schedules. Orange contour represents the ground truth data distribution, blue dots represent the generated samples across different discretizations. ( $\Psi(\cdot,\cdot,\cdot)$ represents the ODE path.)

Figure 2: Comparison of endpoint(accumulated) errors for different timestep strategies (NFE=3). Each point is an initial noise sample $\mathbf{x}_{T}\sim\mathcal{N}(0,\sigma_{T}^{2}\mathbf{I})$ , colored by L2 error relative to 100-step euler as the ground truth. Methods: (a) Uniform timesteps. (b) Globally optimized timesteps. (c) Instance-specific timesteps (overfitted). (d) Instance-specific timesteps (learned through network $\phi$ ).

Figure 3: Quantitative comparison on synthetic experiments, evaluating MSE to teacher samples, KL divergence, and Wasserstein distance across various NFEs(log-scale). Methods include: uniform heuristics, globally optimized timesteps, and our proposed instance-level optimized timesteps conditioned on prior sample.

Figure 4: Architectural design of the proposed lightweight prior conditioning network. When conditional information is available, class indices are first scaled by a factor of $\frac{1}{\sqrt{\text{label\_dim}}}$ and then processed through a linear layer. For prompt embeddings (FLUX.1-dev), T5 embeddings undergo mean pooling to reduce dimensionality before being concatenated with CLIP embeddings.

	CIFAR10 32 $\times$ 32			FFHQ 64 $\times$ 64			AFHQv2 64 $\times$ 64			ImageNet 64 $\times$ 64
Method	NFE=3	NFE=5	NFE=7	NFE=3	NFE=5	NFE=7	NFE=3	NFE=5	NFE=7	NFE=3	NFE=5	NFE=7
Best Heu.	57.39	17.35	7.61	72.29	17.52	8.76	40.24	9.01	4.73	44.93	15.53	8.64
DMN	77.69	12.93	5.42	178.09	20.93	10.17	178.76	26.11	13.03	33.72	10.47	5.39
AMED	18.49	7.59	4.36	26.87	12.49	6.64	31.82	7.34	4.39	28.06	10.74	6.66
GITS	25.98	6.77	3.43	26.41	8.85	5.36	24.17	8.72	5.48	26.41	9.85	6.44
LD3	16.52	5.32	3.37	23.86	8.56	4.69	17.94	6.09	2.77	27.82	11.55	5.63

Table 1: FID Comparison for pixel-space DPMs on CIFAR10, FFHQ, AFHQv2 and class conditional ImageNet64, reporting for NFE=3, 5, and 7. Best Heu. represent the best heuristics schedule among uniform, logSNR and polynomial schedules. The complete results are provided in the appendix.

Figure 5: Ablation study on FFHQ, LSUN-Bedroom and FLUX.1-dev, instance condition is observed to be the most contributing factor, while the effect of shifted factors varies across pretrained models.

Method	NFE=3	NFE=5	NFE=7
best heu.	41.99	6.38	4.39
DMN	28.11	6.15	5.16
AMED	58.21	13.20	7.00
GITS	44.78	17.29	9.59
LD3	14.62	5.93	4.31
INDIS	12.44	4.99	3.81

Table 2: FID on latent space LSUN 256x256.

Figure 6: Qualitative Results on FlUX.1-dev (NFE=7) of our instance-level INDIS method, compared with Global Heuristics (RDS) and globally optimized discretization (GOD). Corresponding prompts could be referred in the appendix.

	CIFAR10 (pixel-space)			LSUN 256 $\times$ 256 (latent-space)
Method	NFE=3	NFE=5	NFE=7	NFE=3	NFE=5	NFE=7
EPD	10.40	4.33	2.82	13.21	7.52	5.97
AdaSDE	12.62	4.18	2.88	18.03	6.96	5.16
INDIS	9.26	3.31	2.60	12.44	4.99	3.81

Table 3: Comparison with solver distillation on FID.

Metrics	Method	NFE=3	NFE=5	NFE=7
FID ( $\downarrow$ )	RDS	64.50	30.12	22.58
	GOD	56.82	28.52	22.77
	INDIS	44.35	24.89	22.70
CLIP ( $\uparrow$ )	RDS	23.29	29.66	30.76
	GOD	24.41	29.70	30.80
	INDIS	26.33	30.01	30.86

Table 4: Performance comparison on FLUX.1-dev (NFE=3, 5, 7).

Method	Aes. Quality	Imaging Quality	Subject Consis.
RDS	0.579	0.597	0.963
GOD	0.583	0.603	0.963
INDIS	0.593	0.613	0.964

Table 5: Performance comparison on LTX-Video (NFE=5).

Figure 7: Visualization of instance-level discretizations on FLUX.

Appendix figures & tables24 assets

Supplementary material from the paper’s appendix.

Appendix

Steps	MSE ↓			KL Divergence ↓			Wasserstein ↓
Steps	Uniform	Global	Instance	Uniform	Global	Instance	Uniform	Global	Instance
3	0.0588	0.0245	0.0158	5.1574	1.4452	0.6091	0.2284	0.1466	0.0925
4	0.0213	0.0223	0.0115	1.1362	1.3032	0.6065	0.1361	0.1403	0.0872
5	0.0238	0.0122	0.0093	1.2130	0.5819	0.3463	0.1496	0.1039	0.0771
6	0.0147	0.0094	0.0076	0.6775	0.3521	0.2434	0.1139	0.0869	0.0716
7	0.0121	0.0072	0.0066	0.4180	0.2246	0.2194	0.1038	0.0787	0.0722

Appendix

Table 6: Comparison of different timestep scheduling strategies across MSE, KL divergence, and Wasserstein distance under varying step numbers.

Figure 8: Qualitative comparison of sampling trajectories: Ground Truth ( gray , 100 NFE), Globally Optimized Timesteps ( green ), and Instance-Specific Timesteps ( purple ). Orange coutour represents the data manifold.

Figure 9: Architectural design of the proposed lightweight prior conditioning network. When conditional information is available, class indices are first scaled by a factor of $\frac{1}{\sqrt{\text{label\_dim}}}$ and then processed through a linear layer. For prompt embeddings (FLUX.1-dev), T5 embeddings undergo mean pooling to reduce dimensionality before being concatenated with CLIP embeddings.

Hyperparameter	CIFAR10	FFHQ	AFHQv2	ImageNet64	LSUN	FLUX.1-dev	LTX-Video
Learning Rate	0.05	0.05	0.05	0.05	0.05	0.001	0.001
Batch Size	16	16	16	8	4	4	4
GPUs (A100)	1	1	1	1	4	4	4
Nimgs/videos	10k	10k	10k	10k	10k	10k	5k
$b_{\gamma}$	0.05	0.05	0.05	0.05	0.05	0.01	0.01
$b_{\Delta\tau}$	0.05	0.05	0.05	0.05	0.05	0.01	0.01

Appendix

Table 7: Hyperparameter settings.

Models	Training Time	Sampling Overhead (%)	Parameter Overhead (%)
CIFAR10 ( $32\times 32$ )	8 min	2.5%	0.43%
ImageNet64 ( $64\times 64$ )	40 min	1.9%	0.32%
FLUX.1-dev ( $512\times 512$ )	8 h †	2.3%	0.21%

Appendix

Table 8: Efficiency Analysis.

{{\dagger}}

: for FLUX.1-dev, we train

\phi

for 2 hours on 4 A100 GPUs with a batch size of 4.

Figure 10: Qualitative comparison against teacher and global baseline.

	NFE
Ablation (on CIFAR-10)	3	4	5	6	7	8
w/o instance level	12.05	14.44	5.36	3.83	3.84	3.20
w/o shifted factors	11.12	8.35	4.64	3.30	2.92	3.24
w/o SVD	9.60	4.31	3.27	2.65	2.77	2.62
full	9.26	4.14	3.31	2.81	2.60	2.34

Appendix

Table 9: Ablation study on the components of our method. Metrics are FID

\downarrow

	NFE
Solver choices	3	4	5	6	7	8
Uni_PC	68.25	43.92	24.01	13.12	6.63	4.41
Uni_PC (+INDIS)	10.88	5.88	3.84	3.57	2.59	2.48
DPM-solver++	68.43	46.59	24.99	12.16	6.88	4.62
DPM-solver++ (+INDIS)	9.92	6.52	4.62	3.55	3.17	2.92
iPNDM	57.39	29.78	17.35	9.95	7.61	5.41

Appendix

Table 10: Ablation study on different solver choices, with and without our method (INDIS), on CIFAR-10.

Teacher NFEs.	3	4	5	6	7	8
10	11.57	8.34	5.52	4.01	4.45	3.38
20	9.73	6.59	3.48	3.11	2.81	2.55
30	9.26	4.14	3.31	2.81	2.60	2.34

Appendix

Table 11: Ablations on teacher solver steps.

Table 22

	CIFAR10 32 $\times$ 32			FFHQ 64 $\times$ 64			ImageNet64 (class-cond.)			LSUN 256 $\times$ 256 (latent)
Method	NFE=3	NFE=5	NFE=7	NFE=3	NFE=5	NFE=7	NFE=3	NFE=5	NFE=7	NFE=3	NFE=5	NFE=7
EPD	10.40	4.33	2.82	21.74	7.84	4.81	18.28	6.35	5.26	13.21	7.52	5.97
AdaSDE	12.62	4.18	2.88	23.80	8.05	5.11	18.51	6.90	5.26	18.03	6.96	5.16
INDIS	9.26	3.31	2.60	17.72	6.91	3.90	18.96	7.28	4.94	12.44	4.99	3.81

Appendix

Table 16: Comparison with solver distillation. FID at NFE=3,5,7 on CIFAR10, FFHQ, class-conditional ImageNet64, and LSUN 256

\times

256 (latent-space).

Method	NFE=3	NFE=4	NFE=5	NFE=6	NFE=7	NFE=8
best heu.	41.99	11.93	6.38	5.08	4.39	4.88
DMN	28.11	11.82	6.15	4.71	5.16	4.55
AMED	58.21	15.67	13.20	8.92	7.00	4.19
GITS	44.78	21.67	17.29	11.52	9.59	8.82
LD3	14.62	8.48	5.93	4.52	4.31	4.22
INDIS	12.44	6.55	4.99	3.84	3.81	3.66

Appendix

Table 17: FID results on LSUN-Bedroom.

Metrics	Method	NFE=3	NFE=4	NFE=5	NFE=6	NFE=7
FID( $\downarrow$ )	RDS	64.50	30.36	30.12	23.16	22.58
	GOD	56.82	29.33	28.52	23.32	22.77
	INDIS	44.35	26.07	24.89	22.93	22.70
CLIP( $\uparrow$ )	RDS	23.29	27.77	29.66	30.42	30.76
	GOD	24.41	28.24	29.70	30.55	30.80
	INDIS	26.33	28.70	30.01	30.67	30.86

Appendix

Table 18: Full comparison on FLUX.1-dev(MS-COCO). For each column, we bold the best performing method. We found AFS to be negative on this large scale pretrained model, thus reporting the result w/o AFS.

Prompt set	Method	Aes.	Img.	Subj.	PAS	TCS	CAS
VBench	RDS	0.579	0.597	0.963	2.162	95.643	49.615
	GOD	0.583	0.603	0.963	3.155	96.672	49.776
	INDIS	0.593	0.613	0.964	3.368	98.337	49.852
VMBench	RDS	0.552	0.601	0.955	2.956	96.053	48.981
	GOD	0.561	0.607	0.956	3.220	96.892	49.693
	INDIS	0.583	0.611	0.961	3.347	97.187	50.138

Appendix

Table 19: Performance on VBench and VMBench prompt subsets.

Figure 11: Qualitative Comparison (NFE=5) between global heuristics, global optimized and instance level INDIS methods (1/2) on LTX-Video

Figure 12: Qualitative Comparison (NFE=5) between global heuristics, global optimized and instance level INDIS methods (2/2) on LTX-Video

Panel of Figure 13 — Figure 13: Qualitative Comparison (NFE=7) between global heuristics, global optimized and instance level INDIS methods (1/2) on FLUX.1-dev

Panel of Figure 14 — Figure 14: Qualitative Comparison (NFE=7) between global heuristics, global optimized and instance level INDIS methods (2/2) on FLUX.1-dev

Panel of Figure 15 — Figure 15: Qualitative comparison on FFHQ64x64 datasets with NFE=3 settings.

Panel of Figure 16 — Figure 16: Qualitative comparison on AFHQv2 64x64 datasets with NFE=3 settings.

Panel of Figure 17 — Figure 17: Qualitative comparison on ImageNet 64x64 datasets with NFE=3 settings.

Panel of Figure 18 — Figure 18: Qualitative comparison on CIFAR10 32x32 datasets with NFE=3 settings.

Panel of Figure 19 — Figure 19: Qualitative comparison on latent space LSUN-Bedroom 256x256 datasets with NFE=3 settings.

Panel of Figure 20 — Figure 20: Qualitative comparison on latent space LSUN-Bedroom 256x256 datasets with NFE=4 settings.

Few-Step Diffusion Sampling

Few-step diffusion sampling develops ways to generate data with far fewer denoising evaluations while retaining quality and diversity.

See all 174 papers

Formalizing the Sampling Design Space of Diffusion-Based Generative Models via Adaptive Solvers and Wasserstein-Bounded Timesteps
Feb 13, 2026Sangwoo Jo, Sungjoon ChoiDiffusion Model Sampling Diffusion Sampling
Department of Artificial Intelligence, Korea University, Seoul, South Korea.
Few-Step Generation via Data-Space Iteration
Oct 8, 2026Shanchuan Lin, Yansong Peng, Fu-Yun Wang +1Iterative Refinement Flow-Based Generative Modeling
ByteDance Seed
Sharpen Your Flow: Sharpness-Aware Sampling for Flow Matching
May 12, 2026Aditi Gupta, Soon Hoe Lim, Annan Yu +1Flow Matching Flow-Based Generative Modeling
Lawrence Berkeley National Laboratory · International Computer Science Institute · 1Lawrence Berkeley National Laboratory, 2International Computer Science Institute +4

Feb 13, 2026J/K move · Enter open · S save

Formalizing the Sampling Design Space of Diffusion-Based Generative Models via Adaptive Solvers and Wasserstein-Bounded Timesteps

Sangwoo Jo, Sungjoon Choi

Department of Artificial Intelligence, Korea University, Seoul, South Korea.

Diffusion Model Sampling Diffusion Sampling Adaptive Sampling Diffusion Timestep Scheduling Diffusion Model Acceleration

Open paper AskPDF arXiv

Diffusion-based generative models have achieved remarkable performance across various domains, yet their practical deployment is often limited by high sampling costs. While prior work focuses on training objectives or individual solvers, the broader sampling design problem, specifically solver selection and scheduling, remains largely governed by static heuristics. We propose SDM, a principled, training-free sampling framework that adapts both the numerical solver and the timestep schedule to the intrinsic properties of the diffusion trajectory. By analyzing the PF-ODE dynamics, we show that velocity variation is small in high-noise stages and increases near the data manifold, identifying intervals where solver order is most consequential. In parallel, we introduce an offline-calibrated adaptive scheduling method that explicitly controls the local Wasserstein discretization error and projects the calibrated trajectory to a prescribed NFE budget. We further extend the formulation to a mixed-transition Wasserstein error bound, providing a unified error-propagation view of adaptive scheduling and solver selection within the overall SDM framework. Across standard benchmarks, with extensions to modern ODE samplers, high-resolution synthesis, and text-to-image generation, SDM achieves improved sample quality compared to baseline methods, attaining an FID of 1.93 on CIFAR-10, 2.41 on FFHQ, and 1.98 on AFHQv2, with a reduced number of function evaluations compared to existing samplers. Our code is available at https://github.com/aiimaginglab/sdm.

Abstract

Few-Step Diffusion Sampling

Explore similar work

Formalizing the Sampling Design Space of Diffusion-Based Generative Models via Adaptive Solvers and Wasserstein-Bounded Timesteps

Few-Step Generation via Data-Space Iteration

Sharpen Your Flow: Sharpness-Aware Sampling for Flow Matching