Optimize Your Sampling: Tuned Diffusion Sampling with Bayesian Optimization
2026-08-18 • Machine Learning
Machine LearningComputer Vision and Pattern Recognition
AI summaryⓘ
The authors focus on making image generation from diffusion models faster by choosing better sampling steps. Instead of using usual methods, they treat picking these steps like a black-box problem and solve it with Bayesian optimization, which aims directly to improve image quality. Their method, called Optimizing Your Sampling (OYS), works well on different image tasks without extra training and speeds up the process a lot while keeping most of the image quality. It even improves complex samplers and works for models that have been simplified (distilled).
diffusion modelssampling timestepsBayesian optimizationtext-to-image generationinpaintingEuler samplerDPM-Solver++distilled modelsinference costblack-box optimization
Authors
Travis Zhang, Christian Belardi, Justin Lovelace, Jin Peng Zhou, Saebyeol Shin, Carla P. Gomes, Kilian Q. Weinberger
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.