Optimize Your Sampling: Tuned Diffusion Sampling with Bayesian Optimization

Travis Zhang; Christian Belardi; Justin Lovelace; Jin Peng Zhou; Saebyeol Shin; Carla P. Gomes; Kilian Q. Weinberger

cs.LGAug 18, 2026

Optimize Your Sampling: Tuned Diffusion Sampling with Bayesian Optimization

Authors: Travis Zhang, Christian Belardi, Justin Lovelace, Jin Peng Zhou, Saebyeol Shin, Carla P. Gomes, Kilian Q. Weinberger

Organizations: Cornell University

Abstract

Sampling from a diffusion model typically requires many forward passes through a large neural network, making generation computationally expensive. While much work has focused on efficient solvers and samplers, comparatively little attention has been paid to selecting the sampling timesteps themselves. A recent line of work optimizes theoretically derived surrogates for sample quality rather than the quality metric itself. We propose Optimizing Your Sampling (OYS), which instead treats timestep selection as a black-box optimization problem, optimizing the target metric directly with Bayesian optimization. OYS outperforms both the default schedules and those of Align Your Steps on text-to-image generation, and improves over the default schedules on inpainting and other image tasks, in both quantitative and human evaluations. OYS requires no additional training, is applicable even to distilled models, and improves both simple and sophisticated samplers such as Euler and DPM-Solver++. A 5-step OYS schedule retains 89%-94% of the quality of a 50-step schedule while reducing inference cost by 10x.

Figures & tables

Figure 1: Qualitative comparisons between Default, AYS ( Sabour et al., 2024 ) , and ours (OYS) using a 5 step sampling schedule. All prompts are held-out test examples, unseen during schedule optimization for either method.

Figure 2: Overview of our Optimize Your Sampling (OYS) framework using Bayesian Optimization. The process begins with Acquisition Maximization (top left), where promising sampling configurations are identified. The selected configuration undergoes Black Box Metric Evaluation (right). This consists of generating images with the diffusion model and scoring them using metrics like FID and HPS. Results update the surrogate model’s posterior distribution (bottom left), refining our understanding of the configuration space. This cycle repeats until convergence or budget exhaustion, efficiently discovering effective sampling parameters.

Table 3

Figure 3: Text-to-Image User Study. Percent Win Rate on SDXL measuring image fidelity and image alignment. Significance was assessed with two-sided exact binomial tests on the win and lose counts, excluding ties, with all differences being highly significant (

p<0.001

). OYS significantly outperforms both AYS and default sampling in quality (68.8% win rate vs. AYS, 70.6% vs. Default; ties split equally) and in alignment (62.9% win rate vs. AYS, 66.8% vs. Default).

Figure 4: Inpainting Results for SDv1.5 and SDXL against default and OYS schedule. OYS reconstructions remain faithful to the unmasked regions, whereas the 5-step default schedule visibly corrupts the full image.

Figure 5: 5-step text-to-image sampling schedules in log-SNR space. AYS closely tracks the default schedule’s near-uniform log-SNR spacing, while OYS consistently places more steps at low log-SNR, i.e., high noise levels.

Task	Method	PSNR $\uparrow$	HPS Average $\uparrow$	HPS Win Rate $\uparrow$
Inverse HED	Default	59.79	0.197	-
Inverse HED	OYS	60.46	0.206	65.43%
Inverse Depth	Default	58.51	0.198	-
Inverse Depth	OYS	59.05	0.202	57.62%
Inverse Segment	Default	57.97	0.196	-
Inverse Segment	OYS	57.60	0.201	58.80%

Table 4: Prompt Diffusion. PSNR, HPS, and Win Rate on Inverse HED, Inverse Depth, and Inverse Segmentation comparing Default and OYS schedules. HPS Win Rates are with respect to the Default Schedule.

Figure 8: Prompt Diffusion Results for Inverse Depth and Inverse HED against default and OYS schedule. Best viewed in color.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Sampler	Steps	HPS Average $\uparrow$	OYS vs. Row HPS Win Rate $\uparrow$
Default	EA/ED	1	0.277	83.94%
Default	ED	3	0.287	83.35%
Default	EA	3	0.287	78.55%
OYS	ED	3	0.298	-

Appendix

Table 5: SDXL-Turbo Results. Mean HPS across the COCO Captions validation split and the HPS win rate of OYS (3-step, ED) against each default baseline. ED and EA denote Euler Discrete and Euler Ancestral, respectively.

Table 11

Panel of Figure 9 — Figure 9: User Study Template. Instructions and question format used in our user study for evaluating image pairs.

Model	$\sigma_{\min}$	$\sigma_{\max}$	$\rho$	EMA length	Guidance
EDM2-XS
Default	0.002	80.0	7.0	0.045	1.40
OYS	0.0001	20.0	4.976	0.025	1.758
EDM2-S
Default	0.002	80.0	7.0	0.025	1.40
OYS	0.010	22.132	9.767	0.028	1.659

Appendix

Table 9: EDM2 Sampling Configurations. Comparison of the default and OYS-optimized sampling configurations for each EDM2 model size.

Problem setup
Notation	Definition
$\mathcal{X}$	Set of task inputs (e.g., text prompts) used for tuning
$m$	Sampling function of the fixed, pretrained diffusion model
$e$	Black-box evaluation metric
$\mathbf{p}$	Sampling configuration, $\mathbf{p}=\{p_{1},\dots,p_{M}\}$ ; the parameters OYS optimizes
$p_{k}$	$k$ -th element of $\mathbf{p}$ ; $p_{k}\in[0,1]$ for the schedule parameters

Appendix

Table 10: Notation. Summary of the notation used throughout the paper.

Explore similar work

Learning Sampling Parameters for Diffusion Models
Jul 26, 2026Arisrei Lim, Yossi GandelsmanDiffusion Model Guidance Diffusion Model Sampling
Reve
Analyzing and Improving Fast Sampling of Text-to-Image Diffusion Models
Feb 28, 2026Zhenyu Zhou, Defang Chen, Siwei Lyu +2Diffusion Model Inference Acceleration T2I Diffusion Models
State Key Laboratory of Blockchain and Data Security, Zhejiang University · Hangzhou High-Tech Zone (Binjiang) Institute of Blockchain and Data Security · University of California, Berkeley +1
Ctrl-Z Sampling: Scaling Diffusion Sampling with Controlled Random Zigzag Explorations
Jun 25, 2025Shunqi Mao, Wei Guo, Chaoyi Zhang +3Diffusion Sampling Conditional Diffusion Models
School of Computer Science, The University of Sydney, Australia

Jul 26, 2026J/K move · Enter open · S save

Learning Sampling Parameters for Diffusion Models

Arisrei Lim, Yossi Gandelsman

Reve

Diffusion Model Guidance Diffusion Model Sampling RL for Diffusion Models T2I Diffusion Models

Open paper AskPDF arXiv

Text-to-image diffusion models expose many inference-time sampling parameters, including prompts, negative prompts, classifier-free guidance scales, and noise schedules. These parameters are typically manually chosen once and then held fixed across prompts and denoising timesteps, even though different prompts and stages of generation can benefit from different parameter values. We introduce LeSAMP, a framework for learning prompt-conditioned, timestep-varying sampling parameters. We formulate parameter selection as a reinforcement learning problem: Given a user prompt, a large language model is trained to emit schedules for the chosen sampling parameters. We optimize our model using rewards from human preference models and VLM-as-a-judge. We evaluate our model on Flux.1 [dev] and Stable Diffusion 3.5, and find that compared to baselines, LeSAMP has a win rate of up to 68.12% using human preference scores and 73.37% using VLM-as-a-judge. These gains are validated in a user study where we achieve win rates of up to 59.46% over previous baselines. Our results suggest that learned sampling-parameter policies provide a complementary approach to existing post-training methods for improving diffusion model outputs.

Abstract

Explore similar work

Learning Sampling Parameters for Diffusion Models

Analyzing and Improving Fast Sampling of Text-to-Image Diffusion Models

Ctrl-Z Sampling: Scaling Diffusion Sampling with Controlled Random Zigzag Explorations