Source-linked AI summary
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
TL;DR
Diffusion sampling is expensive, and timestep selection remains underexplored despite its impact at low step budgets. OYS tunes timestep schedules directly with Bayesian optimization, outperforming standard alternatives across image tasks while retaining 89%–94% of 50-step quality with five steps.
Problem
Timestep selection remains underexplored despite substantially affecting quality, especially when diffusion generation uses few denoising steps.
Method
OYS casts sampling-parameter selection as black-box optimization, using Bayesian optimization to tune the evaluation metric directly without retraining.
Results
OYS outperforms AYS and default schedules on text-to-image generation and improves over defaults on inpainting and inverse image tasks; five steps retain 89%–94% of 50-step quality.
Takeaways & Limitations
OYS makes domain-specific schedule tuning practical across pretrained models and samplers, including distilled models, while reducing five-step inference cost to one tenth.
Takeaways & Limitations
OYS does not seek universal schedules because optimal noise schedules are inherently domain-dependent.
Abstract
from arXiv · showhide
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.
1 Introduction
OYS reframes diffusion timestep selection as direct black-box optimization of sampling quality, using Bayesian optimization without additional training. It improves over established schedules across text-to-image, inpainting, and inverse image tasks, while retaining much of 50-step quality with only 5 steps.
- Method: OYS treats timestep selection as a black-box optimization problem and solves it with Bayesian optimization.The method searches the full configuration space without requiring gradient information.
- Method: OYS requires no additional training and applies to any diffusion model.Its gradient-free search can reallocate steps across the entire timestep range rather than refining locally.
- Results: OYS consistently outperforms AYS and the log-linearly downsampled default schedule on text-to-image generation, and the default schedule on inpainting and inverse image tasks.These gains are verified by quantitative metrics and human evaluation.
- Results: OYS allocates more steps to high-noise timesteps, unlike the approximately uniform log-SNR spacing of AYS and the default schedule.This pattern is observed by inspecting the optimized schedules.
- Results: OYS retains much of the quality of a 50-step schedule using only 5 steps.The framework therefore targets substantial quality retention under reduced sampling budgets.
2 Related Work
Prior work optimizes sampling schedules using parameter searches, dynamic programming, differentiable objectives, or theoretically derived surrogates. OYS instead optimizes perceptual quality directly through black-box forward evaluations without gradients.
- Customized schedules: Karras et al. define sampling parameters including σmin, σmax, and ρ, but optimizing their configuration requires an expensive grid search.σmin and σmax bound the noise range, while ρ controls step-size distribution.
- Customized schedules: Lugmayr et al. improve inpainting coherence with a zigzag denoising schedule that repeatedly steps backward and forward, at the cost of additional function evaluations.
- Direct schedule optimization: Watson et al. select K-step subsequences by dynamic programming to maximize an evidence lower bound, improving likelihood but not perceptual quality.
- Direct schedule optimization: DDSS optimizes perceptual scores by backpropagating through parametric samplers, but its memory grows linearly with sampling steps and it rematerializes score-function calls.OYS shares DDSS’s goal of directly optimizing perceptual quality but treats the objective as a black box, using ordinary forward sampling passes without gradients.
- Surrogate-based optimization: AYS, ART, and HSO optimize theoretically derived surrogates—KLUB, Euler discretization error, and midpoint error proxy—rather than the quantity of interest itself.Their surrogate-based searches avoid generating images during optimization but depend on how well each surrogate tracks perceptual quality.
3 Optimize Your Sampling (OYS)
OYS treats sampling-timestep selection as black-box optimization of an evaluation metric using Bayesian optimization. It supports direct timestep optimization and joint optimization of parametric schedule parameters while enforcing valid monotonic schedules when needed.
- Bayesian Optimization: OYS optimizes sampling parameters p for a fixed pretrained diffusion model and tuning inputs X using Bayesian optimization until convergence or budget exhaustion.The evaluation metric e quantifies task performance, and optimization is phrased as minimization, including negated higher-is-better metrics such as HPS.
- Direct schedule optimization: For nonparametric schedules, OYS directly optimizes timesteps and uses cumulative products of parameters pk ∈[0, 1] to guarantee monotonically decreasing schedules.Each pk for k > 1 represents the ratio between consecutive timesteps, tk/tk−1.
- Direct schedule optimization: Discrete-time schedules are rounded after scaling so their timesteps belong to the model’s fixed integer timestep set.The parameterization does not produce a uniform distribution over valid monotonic schedules.
- Direct schedule optimization: Despite a bias toward smaller timesteps, OYS consistently discovers schedules allocating more steps to high-noise regions.Later cumulative-product timesteps concentrate near smaller values, while the Bayesian optimization surrogate compensates for this non-uniformity.
- Parametric schedules: For EDM, OYS jointly optimizes real-valued schedule parameters such as σmin, σmax, and ρ, tuning timesteps indirectly.The EDM formulation can also include EMA length and classifier-free guidance strength among the optimized parameters.
- Bayesian Optimization: A Gaussian process models observed evaluations, and its posterior mean and uncertainty guide qLogNEI-based candidate selection after Sobol initialization.The number of initial Sobol samples is twice the number of parameters.
4 Experiments
Experiments show that OYS is most valuable in aggressive low-step regimes, where it improves text-to-image quality across models and samplers, including distilled models. Human evaluations and results on other metrics and tasks support that these gains are not limited to the tuning objective.
- Text-to-Image: OYS focuses on 5-step text-to-image sampling, where schedule choices meaningfully separate after substantial degradation from 10-step performance.At 10 steps, downsampled defaults remain competitive with AYS; the study therefore emphasizes the low-step regime.
- Text-to-Image: OYS improves 3-step SDXL-Turbo and 5-step FLUX.1-dev and QwenImage schedules over their corresponding defaults.These results cover standard and distilled models, discrete-time and continuous-time formulations, and multiple samplers.
- Human Evaluation: 68.8% and 70.6% quality win rates show OYS outperforming AYS and Default, while alignment win rates reach 62.9% and 66.8%.The SDXL user study involved 58 participants, and all reported differences were highly significant with p < 0.001.
- Metric Generalization: Tuning on HPS improves FID substantially on DeepFloyd and SDXL without degrading distributional quality, while SDv1.5 schedules have comparable 5-step FID.The paper cautions that small zero-shot COCO FID differences on SDv1.5 do not strongly track perceived quality.
- Schedule Analysis: OYS allocates more log-SNR timesteps to higher noise levels, unlike the locally refined AYS schedule that remains close to Default.The schedule comparison uses log-SNR as a unified representation of noise levels across diffusion models.
- Optimization Efficiency: Performance saturates before 170K images for all models and before 17K generations for SDXL, making its 5-step tuning cost about 1.7K 50-step generations.The reported AYS tuning upper bound is approximately 2.4 million generations.
5 Conclusion
OYS frames diffusion sampling schedule selection as direct black-box optimization of the evaluation metric, without retraining. It supports pretrained models, samplers, and objectives while enabling domain-specific tuning that is cheaper to apply repeatedly.
- OYS casts diffusion sampling parameter selection as black-box optimization of the evaluation metric itself rather than a surrogate.The framework is presented as the paper’s central contribution.
- OYS requires no retraining and applies to any pretrained model, sampler, or objective.
- OYS improves over AYS and the default schedule for text-to-image generation, and over the default schedule for inpainting and inverse image tasks.The gains also extend to SDXL-Turbo, a model distilled for few-step sampling.
- In log-SNR space, short schedules benefit from allocating more steps to high-noise regions, unlike the approximately uniform spacing of default and AYS schedules.For EDM2, OYS instead tunes a parametric schedule and drives σmax far below its default across model sizes.
- OYS does not seek universal schedules because optimal noise schedules are domain-dependent; its one-time model-and-task cost is amortized across subsequent sampling.The authors position this as a practical route to more efficient sampling across architectures and tasks.
A Additional Results
Additional results show that OYS remains effective for distilled SDXL-Turbo and improves FID on several tasks and models, while FID does not always track perceived quality. On SDv1.5, OYS is slightly behind baselines in FID despite winning HPS comparisons.
- SDXL-Turbo: 0.298 mean HPS: OYS’s 3-step Euler Discrete schedule achieves the highest score on SDXL-Turbo.SDXL-Turbo is a distilled few-step model with already highly compressed default schedules.
- SDXL-Turbo: 78.55%-83.94% win rates: OYS consistently outperforms all SDXL-Turbo default baselines.The comparison uses OYS at 3 steps against default schedules at 1 and 3 steps with Euler Discrete and Euler Ancestral samplers.
- FID evaluation: FID is reported for text-to-image, inpainting, and Prompt Diffusion, although zero-shot COCO FID may not correspond to perceived performance.Text-to-image and inpainting use the COCO Captions validation split, while Prompt Diffusion uses the Wang et al. (2023) dataset test split.
- FID evaluation: Within 0.7 FID: all three 5-step SDv1.5 schedules are closely matched, with OYS slightly behind the baselines.This contrasts with OYS winning the HPS comparisons in Table 1, consistent with loose coupling between COCO FID and perceived quality.
B User Study
The user study used a two-phase, anonymized pairwise evaluation to separately assess image quality and prompt alignment. Participants strongly preferred OYS images over Default and AYS for both criteria.
- Methodology: The study evaluated image quality and prompt alignment across different sampling schedules using a two-phase process for each image pair.Figure 9 shows the study template.
- Methodology: Participants judged image quality before seeing the prompt, then evaluated text-image alignment to prevent prompt content from influencing quality judgments.Quality was assessed in Question 1 and alignment in Question 2.
- Methodology: Participants compared randomly shuffled and anonymized image pairs generated with Default, AYS, and OYS schedules.Anonymization concealed which sampling technique generated each image.
- Results: Participants strongly preferred OYS images over both Default and AYS schedules for image quality and prompt alignment.These user-study results were presented in the main body of the paper.
C EDM2 Sampling Configurations
OYS discovers markedly different EDM2 sampling configurations from the defaults, especially by reducing σmax and shifting σmin to search-range boundaries. On text-to-image generation, OYS also substantially lowers FID for SDXL and SDv1.5.
- Noise-level parameters: σmax reaches 20 for every model except EDM2-S, which settles at 22.1, versus the default of 80.Both OYS values are at or near the lower search bound, suggesting the preferred σmax may be lower still.
- Noise-level parameters: σmin reaches a search-range boundary for every model: EDM2-XS reaches 0.0001, while all others reach 0.01 versus the default of 0.002.The upper-bound value 0.01 is five times the default.
D Notation
This section summarizes the notation used throughout the paper in Table 10.
- Table 10 provides a summary of the notation used throughout the paper.