Optimize Your Sampling: Tuned Diffusion Sampling with Bayesian Optimization
Proposes Bayesian Optimization-based sampling schedule (OYS) that reduces steps to 5 while retaining 89%-94% of quality, greatly lowering inference costs.
Key Findings
Methodology
This paper introduces OYS, which formulates the selection of diffusion timesteps as a black-box optimization problem. Using Gaussian Process-based Bayesian optimization, the method iteratively searches for the optimal sampling configuration by modeling the performance metrics (e.g., FID, HPS) as a function of scheduling parameters. The process involves initial Sobol sampling to generate diverse configurations, followed by Gaussian process fitting to predict performance, and acquisition function maximization (e.g., qLogNEI) to select the next promising configuration. Each evaluation involves generating images with the current schedule and scoring them, which feeds back into the model for subsequent iterations. This approach supports both continuous and discrete parameterizations, enabling flexible adaptation across different diffusion models, including distilled and complex samplers like Euler and DPM-Solver++. Experimental results demonstrate that optimizing the schedule directly on target metrics yields significant improvements in low-step regimes, with 5-step schedules achieving near-50-step quality levels while reducing inference costs by over tenfold.
Key Results
- On COCO captions, the 5-step optimized schedule improves HPS from 0.215 (default) to 0.252, with a win rate of 83.99% over the default schedule. Similar gains are observed across models like Stable Diffusion XL and v1.5, with win rates exceeding 70%.
- In the SDXL model, a 3-step schedule tuned via OYS achieves an HPS of 0.298, surpassing the default 0.287 by 4%.
- In image inpainting tasks, PSNR increases by 6.8dB and HPS improves by over 30%, with human evaluations favoring OYS-generated images in quality and content alignment.
- Across multiple tasks, OYS consistently outperforms baseline schedules, demonstrating its broad applicability and robustness.
Significance
This work fundamentally shifts the paradigm of diffusion schedule design by directly optimizing for perceptual and task-specific metrics, bypassing the limitations of surrogate-based methods. The ability to achieve high-quality generation with minimal steps dramatically enhances the practicality of diffusion models, especially for real-time applications and resource-constrained environments. The approach’s flexibility allows it to be extended to various tasks and models, fostering broader adoption in industry and research. By enabling efficient low-step sampling, this method paves the way for deploying diffusion models in edge devices, mobile applications, and interactive systems, where speed and efficiency are critical. Moreover, the framework sets a new standard for automated, data-driven schedule tuning, inspiring future work on multi-objective and adaptive scheduling strategies.
Technical Contribution
The core technical innovation lies in casting the timestep selection as a black-box optimization problem and solving it with Bayesian optimization, which is a departure from previous surrogate-based or heuristic methods. The paper introduces a flexible parameterization of the schedule, supporting both continuous and discrete models, and employs Gaussian processes to model the performance landscape. The acquisition function (qLogNEI) guides the search efficiently, enabling the discovery of schedules that allocate more steps to high-noise regions, which are crucial for maintaining quality at low steps. The approach is model-agnostic, requiring no retraining or gradient information, and supports multi-task optimization. Extensive experiments across diverse datasets and models validate its effectiveness, demonstrating that it can reliably produce low-step schedules that nearly match the quality of full-step schedules, significantly reducing inference costs.
Novelty
This research is the first to treat diffusion timestep selection as a black-box optimization problem, leveraging Bayesian optimization to directly maximize target metrics rather than relying on surrogate bounds like KLUB. Unlike prior methods such as AYS, HSO, or ART, which optimize theoretical bounds or local heuristics, OYS performs a global search over the entire schedule space, enabling more effective reallocation of steps. Its flexibility supports both continuous and discrete parameterizations, and it can optimize any user-defined metric, including perceptual scores and human preferences. The method’s ability to produce high-quality results with as few as 5 steps, while maintaining near-full-step quality, marks a significant advancement in low-step diffusion sampling.
Limitations
- While OYS excels in low-step scenarios, its benefits diminish at higher step counts where traditional schedules already perform well, making the optimization less impactful. The computational cost of repeated image generation for evaluation remains high, especially for large models or multi-objective settings, potentially limiting scalability in resource-constrained environments.
- The approach relies on the availability of a suitable evaluation metric and may require task-specific tuning to achieve optimal results. Its performance can vary across different models and datasets, necessitating further validation for broader applications.
- Current experiments focus on static tasks; extending the framework to dynamic or multi-modal scenarios, such as video synthesis or 3D content, poses additional challenges. Incorporating adaptive or online scheduling mechanisms could further enhance its practicality.
Future Work
Future directions include integrating multi-objective optimization to balance quality, speed, and resource consumption, developing adaptive online scheduling algorithms, and exploring reinforcement learning-based approaches for real-time adjustment. Additionally, extending the framework to multi-modal and temporal content generation, such as video or 3D scene synthesis, is a promising avenue. Combining schedule optimization with model fine-tuning or guidance techniques could further boost performance. Lastly, efforts to reduce evaluation costs, perhaps through surrogate models or learned proxies, will be critical for scaling this approach to large-scale industrial deployment.
AI Executive Summary
Diffusion models have revolutionized generative AI, enabling the creation of highly realistic images, videos, and audio through iterative denoising processes. Despite their success, a persistent challenge remains: the high computational cost associated with the numerous forward passes required for sampling. Traditional solutions have focused on designing efficient samplers or training models with fewer steps, but these approaches often rely on heuristics or theoretical bounds that do not directly optimize the ultimate quality of generated samples.
This paper introduces a novel approach called 'Optimize Your Sampling' (OYS), which directly addresses the core issue of timestep scheduling in diffusion models. Instead of relying on predefined schedules or surrogate bounds like KL divergence upper bounds (KLUB), OYS formulates the selection of sampling timesteps as a black-box optimization problem. Using Bayesian optimization with Gaussian processes, the method iteratively searches for the optimal sequence of timesteps that maximize a target metric—such as Fréchet Inception Distance (FID) or Human Preference Score (HPS)—by generating images, evaluating their quality, and updating the surrogate model. This process enables the discovery of highly effective schedules that allocate more steps to high-noise regions, which are crucial for maintaining quality at low step counts.
Experimental results demonstrate that OYS can reduce the number of sampling steps from 50 to just 5, while still retaining 89%-94% of the quality achieved by full-step schedules. In text-to-image tasks, this translates to over tenfold reduction in inference cost, making real-time generation feasible on resource-limited devices. The method also shows significant improvements in image inpainting and inverse tasks, with PSNR gains of up to 6.8dB and HPS increases exceeding 30%. Human evaluations further confirm that images generated with OYS are perceived as higher quality and more content-aligned than those produced by traditional schedules.
The broader impact of this work lies in its potential to democratize high-quality content generation. By drastically reducing computational requirements, OYS paves the way for deploying diffusion models in edge devices, mobile apps, and interactive systems. Its flexible, model-agnostic framework can be extended to multi-objective optimization, adaptive scheduling, and even multi-modal content creation, opening new avenues for research and industrial applications. Despite these advances, challenges remain, including the high cost of repeated evaluations during optimization and the need for further validation across diverse tasks. Future work will focus on integrating multi-objective and online adaptive strategies, aiming to develop smarter, faster, and more resource-efficient generative models.
Deep Analysis
Background
近年来，扩散模型在生成任务中取得了突破性进展，成为图像、视频、语音等多模态内容生成的主流技术。早期代表作如Ho et al. (2020)提出的Denoising Diffusion Probabilistic Models（DDPM），通过逐步去噪实现高质量图像生成。随后，Song et al. (2021)引入连续时间扩散过程，提升了模型的灵活性和性能。近年来，调度参数化（Karras et al., 2022）、自定义调度（Lugmayr et al., 2022）以及调度优化（Sabour et al., 2024）等技术不断推动采样效率的提升。然而，尽管在采样器设计和模型训练方面取得了显著进展，采样时间点的选择仍是瓶颈。传统方法多依赖理论推导或局部优化，难以实现全局最优，且在实际应用中存在调度不适应不同模型和任务的缺陷。因此，如何在保证生成质量的同时，减少采样步骤，成为当前研究的热点与难点。
Core Problem
核心问题在于扩散模型的采样调度设计。现有调度策略多为预定义或局部优化，难以适应不同模型架构和任务需求，尤其在低步数场景下，样本质量迅速下降。传统调度依赖理论推导（如KLUB）或启发式方法，存在局限性：一方面无法实现全局最优，另一方面在实际应用中调度效果不稳定。随着模型规模和复杂度的增加，调度策略的优化变得更加困难，尤其在没有明确梯度信息的情况下，如何找到最优的时间点序列，成为制约低成本高质量采样的关键瓶颈。此外，调度优化的计算成本较高，尤其在需要多次生成与评估的情况下，限制了其在工业应用中的推广。
Innovation
本研究的创新点主要体现在：1）将采样时间点的选择问题转化为黑箱优化任务，利用贝叶斯优化（Gaussian Process）实现全局搜索，避免了梯度依赖和理论推导的局限；2）设计了支持连续与离散模型的参数化调度空间，采用乘积参数化和log-SNR空间映射，确保调度的单调性与可调性；3）实现无需模型再训练的调度优化流程，支持多任务、多模型的泛用性；4）在多个任务（文本到图像、修复、反向）中验证调度优化的有效性，显著提升低步数采样的质量，降低推理成本。这种全局优化策略，突破了传统调度的局限，为扩散模型的高效应用提供了新思路。
Methodology
- �� 定义调度参数空间：采用乘积参数化（pk ∈ [0, 1]）确保调度单调递减，将时间点表示为t(p) = t_max * ∏_{i=1}^k p_i。
- �� 采样空间支持连续与离散模型：对离散模型中的时间点进行四舍五入处理。
- �� 构建目标指标：根据任务定义不同的性能指标（如FID、HPS、LPIPS、PSNR），作为优化目标。
- �� 贝叶斯优化流程：
- 初始化：采用Sobol采样生成初始配置，数量为参数个数的两倍。
- 建模：利用高斯过程（Gaussian Process）对采样配置的性能进行建模，形成后验分布。
- 采集：使用qLogNEI等采集函数，选择下一组最有潜力的采样配置。
- 评估：生成图像并计算目标指标，将结果反馈到高斯过程模型中。
- 迭代：重复采集与建模，直到满足收敛条件或预算耗尽。
- �� 支持多模型、多任务：调度参数化支持连续和离散模型，优化过程可扩展到多目标、多指标场景。
Experiments
实验在多个公开数据集（如COCO Captions、DiffusionDB）上进行，包括文本到图像、图像修复和反向任务。采用多种模型（如DeepFloyd、Stable Diffusion XL、v1.5、QwenImage）进行调度优化，比较默认调度、AYS调度和OYS调度的性能。指标包括HPS、FID、PSNR等，评估生成质量与内容一致性。调优过程采用贝叶斯优化，调度参数在不同步数（如5步、10步、50步）下进行。每次实验都进行多轮随机初始化，确保结果的稳健性。通过人类评测验证生成内容的感知质量，结合自动指标，全面评估调度方案的优劣。
Results
调度优化显著提升了低步数采样的性能。例如，在COCO字幕数据集上，5步调度的HPS由默认的0.215提升至0.252，win率达到83.99%。在SDXL模型中，3步调度经过调优后，HPS提升至0.298，优于默认的0.287，提升幅度达4%。在图像修复任务中，调度优化带来PSNR提升6.8dB，HPS提升超过30%。人类评测显示，OYS调度生成的图像在质量和内容一致性方面，明显优于传统调度方案，用户偏好率超过70%。这些结果表明，调度参数的全局优化在低步数采样中具有巨大潜力，能在保持高质量的同时大幅度降低计算成本。
Applications
该调度优化方法适用于多种场景，包括快速文本到图像生成、图像修复、逆向图像任务等。工业界可以利用OYS实现低延迟、高质量的内容生成，尤其适合边缘设备和实时交互场景。未来，结合多目标优化和自适应调度，有望在视频生成、3D内容等更复杂任务中发挥作用。此外，该方法还可用于模型微调和指导，提升模型的泛化能力和鲁棒性。
Limitations & Outlook
尽管OYS在低步数场景表现优异，但在高步数（如50步）下，调度优化的边际收益有限，可能不如传统调度方案直观有效。优化过程依赖大量的图像生成与评估，计算成本较高，限制了其在资源有限环境中的应用。调度效果在不同任务和模型间存在一定差异，需针对具体应用进行调参和验证，泛化能力仍需提升。此外，当前方法主要在静态任务中验证，动态场景和多模态任务的适应性有待进一步研究。
Plain Language Accessible to non-experts
想象你在一家工厂里，生产线上的每个工序都需要一定的时间和资源。传统上，工厂设计了一个固定的生产流程，比如从原料到成品，按照一定的顺序和时间点进行操作。这就像在扩散模型中，生成一张图片需要经过多次“去噪”步骤，每个步骤的时间点和顺序都预先设定好。可是，工厂的效率其实可以更高，如果我们能根据实际情况动态调整每个步骤的时间和顺序，让生产线更合理、更节省资源。这个想法就是本文提出的“调度优化”。作者用一种叫贝叶斯优化的方法，像在工厂里不断试验不同的流程，找到最节省时间又能保证产品质量的方案。通过这种方法，即使只用很少的步骤，也能生产出高质量的图片，就像用少量的时间和资源，做出和传统方法一样好的产品。这种创新让我们可以用更少的努力，得到更好的结果，特别是在需要快速生成内容的场景中，非常实用。
ELI14 Explained like you're 14
你知道吗？有一种叫扩散模型的技术，就像在画画一样，它会一步步把一张模糊的图片变得越来越清晰。可是，要画一张漂亮的图片，传统方法需要很多次“去噪声”的步骤，就像画家用很多细笔一笔一笔慢慢画出来。这就很慢，也很费力。科学家们一直在想办法，让这个过程变得快一些，但又不影响画质。这个研究就像是给画画的步骤安排一个聪明的“时间表”，让每一步都用得恰到好处。作者用一种叫贝叶斯优化的方法，就像是让一个聪明的助手帮你试验不同的画画顺序，找到最节省时间又能画出漂亮图片的方案。结果发现，只用很少的步骤，比如5步，就能画出和用50步差不多好的图片，节省了很多时间和计算资源。这就像用少量的笔，就能画出一幅细腻的画，非常厉害！这项技术让未来的图片生成变得更快、更便宜，也更适合在手机或边缘设备上使用。
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.
References (20)
Human Preference Score v2: A Solid Benchmark for Evaluating Human Preferences of Text-to-Image Synthesis
Xiaoshi Wu, Yiming Hao, Keqiang Sun et al.
Align Your Steps: Optimizing Sampling Schedules in Diffusion Models
Amirmojtaba Sabour, Sanja Fidler, Karsten Kreis
Adversarial Diffusion Distillation
Axel Sauer, Dominik Lorenz, A. Blattmann et al.
GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium
M. Heusel, Hubert Ramsauer, Thomas Unterthiner et al.
Microsoft COCO Captions: Data Collection and Evaluation Server
Xinlei Chen, Hao Fang, Tsung-Yi Lin et al.
Elucidating the Design Space of Diffusion-Based Generative Models
Tero Karras, M. Aittala, Timo Aila et al.
In-Context Learning Unlocked for Diffusion Models
Zhendong Wang, Yifan Jiang, Yadong Lu et al.
Analyzing and Improving the Training Dynamics of Diffusion Models
T. Karras, M. Aittala, J. Lehtinen et al.
SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis
Dustin Podell, Zion English, Kyle Lacey et al.
DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models
Cheng Lu, Yuhao Zhou, Fan Bao et al.
Hierarchical Text-Conditional Image Generation with CLIP Latents
A. Ramesh, Prafulla Dhariwal, Alex Nichol et al.
Efficient Global Optimization of Expensive Black-Box Functions
Donald R. Jones, Matthias Schonlau, W. Welch
Constrained Bayesian Optimization with Noisy Experiments
Benjamin Letham, B. Karrer, G. Ottoni et al.
Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding
Chitwan Saharia, William Chan, Saurabh Saxena et al.
Scale-wise Distillation of Diffusion Models
Nikita Starodubcev, Denis Kuznedelev, A. Babenko et al.
Learning to Efficiently Sample from Diffusion Probabilistic Models
Daniel Watson, Jonathan Ho, Mohammad Norouzi et al.
DiffusionDB: A Large-scale Prompt Gallery Dataset for Text-to-Image Generative Models
Zijie J. Wang, Evan Montoya, David Munechika et al.
Latent Consistency Models: Synthesizing High-Resolution Images with Few-Step Inference
Simian Luo, Yiqin Tan, Longbo Huang et al.
Pseudo Numerical Methods for Diffusion Models on Manifolds
Luping Liu, Yi Ren, Zhijie Lin et al.
Learning Fast Samplers for Diffusion Models by Differentiating Through Sample Quality
Daniel Watson, William Chan, Jonathan Ho et al.