Technical Summary: Optimizing Your Sampling (OYS) with Bayesian Optimization

Problem Statement

Diffusion models have achieved state-of-the-art performance in generative tasks such as text-to-image synthesis, super-resolution, and inpainting. However, the iterative denoising process inherent to these models requires numerous forward passes through large neural networks, rendering generation computationally expensive. While significant research has focused on developing efficient solvers (e.g., DPM-Solver++) and few-step distilled models (e.g., consistency models), comparatively little attention has been paid to the selection of sampling timesteps themselves.

The choice of timesteps significantly impacts sample quality, particularly under low step budgets where default schedules degrade rapidly. Prior work, such as Align Your Steps (AYS), attempts to optimize timestep selection by minimizing a theoretically derived surrogate: the Kullback-Leibler divergence Upper Bound (KLUB) between the true generative process and its discretization. However, minimizing this bound does not necessarily minimize the mismatch between output distributions. Consequently, AYS relies on early stopping from a heuristic initialization, and its resulting schedules remain close to the default log-linear spacing in log-SNR space, failing to fully exploit the potential gains of global reallocation of steps.

Methodology: Optimizing Your Sampling (OYS)

The authors propose Optimizing Your Sampling (OYS), a framework that treats timestep selection as a black-box optimization problem. Instead of optimizing a theoretical surrogate, OYS directly optimizes the target metric (e.g., Human Preference Score, FID, PSNR) using Bayesian Optimization.

Core Framework

1. Black-Box Formulation: OYS assumes access to a fixed pretrained diffusion model $m$, a set of tuning inputs $X$, and a black-box evaluation metric $e$. The goal is to find a sampling configuration $p$ that minimizes $e(m(X; p))$.
2. Parameterization:
   • Direct Schedule Optimization: For models without a parametric schedule, OYS optimizes the timesteps directly. To ensure monotonicity (high noise to low noise), the authors use a cumulative product reparameterization where each step $k$ is defined by a continuous parameter $p_k \in [0, 1]$. For discrete-time models, timesteps are rounded to the nearest integer.
   • Parametric Schedules: For models like the EDM family, OYS jointly optimizes the functional parameters (e.g., $\sigma_{min}$, $\sigma_{max}$, $\rho$, EMA length, guidance strength) rather than just the timesteps.
3. Bayesian Optimization Loop:
   • Initialization: Initial configurations are selected via Sobol sampling.
   • Surrogate Modeling: A Gaussian Process (GP) is fitted to observed evaluations to approximate the objective function, providing mean and uncertainty estimates.
   • Acquisition: The search for promising configurations is guided by the $q$LogNEI acquisition function.
   • Evaluation: Each iteration involves generating images with the diffusion model and scoring them. This process repeats until convergence or budget exhaustion.

Key Distinctions from Prior Work

Unlike methods that optimize surrogates (e.g., KLUB for AYS, Euler discretization error for ART, or midpoint error for HSO), OYS optimizes the target metric itself. While this incurs the cost of generating images during tuning, it avoids the mismatch between the surrogate and perceptual quality. Furthermore, OYS does not require gradient information through the sampler, making it applicable to non-differentiable operations (like rounding) and objectives without tractable surrogates (like human preference).

Key Contributions

• Direct Metric Optimization: OYS is the first framework to treat diffusion sampling parameter selection as a black-box optimization problem, directly optimizing perceptual quality metrics rather than theoretical bounds.
• No Retraining Required: The method requires no additional training of the diffusion model and is applicable to any pretrained model, sampler, or objective.
• Broad Applicability: OYS improves performance across diverse tasks (text-to-image, inpainting, inverse tasks) and model types (standard, distilled, discrete-time, continuous-time).
• Efficiency: The tuning process is highly efficient. For SDXL, tuning plateaus after approximately 17K five-step generations (85K forward passes), which is an order of magnitude fewer than the upper bound reported for AYS tuning costs.

Experimental Results

The authors evaluated OYS across multiple benchmarks, comparing it against default schedules and AYS schedules.

Text-to-Image Generation

• Metrics: Evaluated using Human Preference Score (HPS) and FID on COCO Captions and DiffusionDB.
• Performance: OYS consistently outperformed both default and AYS schedules, particularly in the aggressive low-step regime (5 steps).
  • On SDXL (5 steps), OYS achieved a 70.70% win rate against the default schedule and a 68.8% win rate against AYS in human evaluations.
  • On DeepFloyd and Stable Diffusion v1.5, OYS showed significant improvements in HPS win rates (83.99% and 59.29% respectively) over the default at 5 steps.
  • OYS also improved performance on SDXL-Turbo (a distilled model), demonstrating that schedule tuning complements distillation.

Inpainting and Inverse Tasks

• Inpainting: Tuning on LPIPS for 5-step schedules improved PSNR by 6.9 dB (SDXL) and 2.9 dB (SDv1.5) compared to defaults. Human evaluation showed OYS winning 93.4% and 85.1% of comparisons, respectively. Crucially, OYS preserved unmasked regions, whereas default schedules often corrupted them.
• Inverse Tasks: For Prompt Diffusion tasks (Inverse HED, Depth, Segmentation), OYS improved HPS win rates by 57.6%–65.4% over defaults.

EDM2 Image Generation

• OYS optimized the parametric schedule ($\sigma_{min}, \sigma_{max}, \rho$) and guidance for the EDM2 family on ImageNet-512.
• Results showed FID improvements of 3.57% to 7.00% across all model sizes (XS to XXL) compared to the grid-searched configurations of Karras et al. (2024).

Qualitative and Structural Insights

• Schedule Analysis: Inspection of optimized schedules in log-SNR space revealed a consistent pattern: OYS allocates more steps to high-noise (low log-SNR) regions. This contrasts with the approximately uniform log-SNR spacing of both default and AYS schedules.
• Human Evaluation: A user study with 58 participants confirmed that OYS generations were preferred over both AYS and Default in terms of image fidelity and prompt alignment (all comparisons significant at $p < 0.001$).
• Generalization: Improvements held across metrics not used during tuning (e.g., tuning on HPS improved FID), suggesting genuine quality gains rather than metric overfitting.

Significance and Claims

The paper claims that OYS offers a practical path to more efficient sampling across architectures and tasks. By casting sampling parameter selection as a black-box optimization problem, OYS bypasses the limitations of theoretical surrogates that fail to capture perceptual quality.

The authors emphasize that while optimal noise schedules are inherently domain-dependent, OYS makes the necessary process of domain-specific tuning simpler and cheaper. The cost is incurred once per model and task but is amortized over all subsequent sampling. The results demonstrate that conventional spacing and noise-range heuristics leave significant quality on the table at low step budgets, and that the parameters worth tuning extend beyond timesteps to include guidance strength and noise range bounds. A 5-step OYS schedule retains 89%–94% of the quality of a 50-step schedule while reducing inference costs by 10x.