PaperlayerView exact revision on arXiv

Source-linked AI summary

OMS-DPM: Optimizing the Model Schedule for Diffusion Probabilistic Models

Enshu Liu, Xuefei Ning, Zinan Lin, Huazhong Yang, Yu Wang

arXiv:2306.08860v1cs.LG

TL;DR

DPM generation can require hundreds or thousands of neural-network evaluations, making sampling slow and impractical for real-time applications. OMS-DPM uses a predictor and evolutionary search to choose a model schedule for pre-trained DPMs under a generation-time budget. Using different models at different denoising steps can significantly improve DPM generation quality and sampling speed.

  • Problem

    DPM generation can require hundreds or thousands of neural-network evaluations, making sampling slow and impractical for real-time applications.

    Source: §1 · pp. 1–2

  • Method

    OMS-DPM uses a predictor and evolutionary search to choose a model schedule for pre-trained DPMs under a generation-time budget.

    Source: Given a set of pre-trained DPM models… · p. 2, §4.1 · p. 4, §4.2.1 · p. 6

  • Results

    Using different models at different denoising steps can significantly improve DPM generation quality and sampling speed.

    Source: Given a set of pre-trained DPM models… · p. 2

  • Takeaways & Limitations

    With Stable Diffusion, OMS-DPM achieves better FID with 12 steps than the best single model with 24 steps, potentially saving computation resources.

    Source: §5.2 · p. 7

    For a new dataset or task, OMS-DPM requires preparing predictor training data and training a new predictor, which incurs substantial overhead.

    Source: §6 · p. 9

Abstract

from arXiv · show

Diffusion probabilistic models (DPMs) are a new class of generative models that have achieved state-of-the-art generation quality in various domains. Despite the promise, one major drawback of DPMs is the slow generation speed due to the large number of neural network evaluations required in the generation process. In this paper, we reveal an overlooked dimension -- model schedule -- for optimizing the trade-off between generation quality and speed. More specifically, we observe that small models, though having worse generation quality when used alone, could outperform large models in certain generation steps. Therefore, unlike the traditional way of using a single model, using different models in different generation steps in a carefully designed \emph{model schedule} could potentially improve generation quality and speed \emph{simultaneously}. We design OMS-DPM, a predictor-based search algorithm, to optimize the model schedule given an arbitrary generation time budget and a set of pre-trained models. We demonstrate that OMS-DPM can find model schedules that improve generation quality and speed than prior state-of-the-art methods across CIFAR-10, CelebA, ImageNet, and LSUN datasets. When applied to the public checkpoints of the Stable Diffusion model, we are able to accelerate the sampling by 2$\times$ while maintaining the generation quality.

1. Introduction

DPMs offer strong generation quality but require many neural-network evaluations, making sampling slow. The paper introduces model scheduling and OMS-DPM to choose among pre-trained models at different denoising steps under a time budget, reporting improved quality-speed trade-offs.

1. Introduction

  • DPM sampling can require hundreds or thousands of neural-network evaluations, contributing to slow generation and limiting real-time applications.Prior work has mainly addressed speed through denoising formation, including noise schedules, discretization schemes, and solver formulas.Source: §1 · pp. 1–2
  • Smaller models can outperform larger ones at particular denoising steps, so assigning different models across steps can improve generation quality and speed.The paper names these per-step assignments a model schedule.Source: §1 · p. 2

Given a set of pre-trained DPM models and a generation

  • Because DPM training and sampling can be decoupled, the paper studies schedule optimization with pre-trained models rather than training from scratch.The introduction notes the prevalence of public pre-trained DPMs, including Stable Diffusion.Source: §1 · p. 2
  • OMS-DPM uses a predictor trained on a small amount of data to estimate schedule quality, then an evolutionary algorithm searches schedules under different generation-time budgets.The schedule search addresses a space that grows exponentially with the number of sampling steps.Source: Given a set of pre-trained DPM models… · p. 2
  • Across CIFAR-10, CelebA, ImageNet-64, and LSUN, OMS-DPM achieves better quality-speed trade-offs than baselines; on CIFAR-10, it reaches FID 3.19 vs. 3.56 and is 2.8× faster than a single model.The introduction also reports over 2× faster sampling while maintaining quality on Stable Diffusion checkpoints.Source: Given a set of pre-trained DPM models… · p. 2

2. Background and Related Work

The background describes DPM generation as reverse diffusion, then reviews training-free samplers and AutoML approaches. It motivates predictor-based search as a way to make expensive schedule evaluation more efficient.

2.1. Diffusion Probabilistic Models

  • DPMs define a forward diffusion process and generate data by solving a reverse SDE or its equivalent probability-flow ODE using a neural network to estimate the score.The forward process adds noise; the reverse process runs from T to 0.Source: §2.1 · p. 2

2.2. Training-Free Samplers

  • Training-free samplers use pre-trained models to discretize timesteps and apply a solver; DDIM and DPM-Solver are two samplers compatible with OMS-DPM.DPM-Solver uses a Taylor expansion and uniform log(SNR) discretization; its order k requires k NFEs to compute derivatives.Source: §2.2 · p. 3

2.3. AutoML

  • AutoML selects machine-learning systems for particular conditions, with model selection, hyperparameter tuning, and architecture design among its research problems.This work focuses on automatic optimization of DPM model schedules, including model selection and sampling-schedule design.Source: §2.3 · p. 3

2.3.1. PREDICTOR-BASED NEURAL ARCHITECTURE SEARCH

  • Predictor-based NAS accelerates architecture evaluation by training a predictor on a small set of architecture-performance pairs; OMS-DPM adapts this idea because DPM sampling evaluations are expensive.The predictor can efficiently evaluate new architectures, paralleling the paper’s goal of accelerating schedule search.Source: §2.3.1 · p. 3

3. Model Schedule: A New Dimension in DPM Design

Denoising ability varies across timesteps, so a model that performs poorly overall may excel late in sampling. The paper uses this observation to motivate mixing models across steps as a potential way to improve generation quality and latency together.

3. Model Schedule: A New Dimension in DPM Design

  • Architecture 1 has the worst overall generation quality but outperforms most other models late in denoising, where model rankings are inconsistent across time.A smaller model used at late steps can reduce denoising loss and latency versus architecture 6 alone.Source: §3 · pp. 3–4
  • Figure 2a relates architecture latency to FID, while Figure 2c compares random model combinations with the best single-model FID.The figure’s red line marks the best single-model FID.Source: §3 · p. 4, Figure 2 · p. 4
  • FID ≈5.4: a randomly mixed schedule outperforms the best-quality single model, architecture 6, in both generation quality and speed.The authors caution that lower denoising loss does not always mean better FID.Source: §3 · p. 4
  • The motivating research question is how to choose a sequence of pre-trained models to achieve the best generation-quality and latency trade-off.Prior training-free work focused on sampler design to reduce the number of model evaluations; model schedule is presented as another optimization dimension.Source: §3 · pp. 3–4

4. OMS-DPM: Optimizing the Model Schedule

OMS-DPM frames model scheduling as a quality-maximization problem under a generation-time budget, then uses a learned predictor and evolutionary search to find schedules.

4.1. Problem Definition: Model Schedule Optimization

  • The optimization chooses the number of sampling steps and the model assigned to each timestep to maximize sample quality within a time budget.Sample quality is measured by a score such as FID, and the models are applied in reverse timestep order to transform noise into data.Source: §4.1 · pp. 4–5
  • Empirical timestep discretization and a null model let the search vary timestep spacing and leave selected steps unused.The paper uses linear discretization for DDIM experiments and uniform logSNR for DPM-Solver experiments.Source: §4.1 · p. 5

4.2.1. OVERALL WORKFLOW

  • The workflow separates schedule-data preparation and predictor training from budget-conditioned search, so the budget affects only the search phase.The workflow takes a model zoo and ODE solver as inputs; schedules are evaluated using FID during training-data preparation.Source: Figure 4 · p. 5, §4.2.1 · p. 5
  • A predictor scores candidate schedules, and evolutionary search selects the best predicted schedule that satisfies the generation-time budget.Training data are schedule–FID pairs; the predictor is trained with pairwise ranking loss, and a simple evolutionary algorithm iteratively modifies promising schedules.Source: §4.2.1 · pp. 5–6

4.2.2. PREDICTOR DESIGN

  • The predictor combines model and timestep embeddings into a sequence, processes it with an LSTM, then averages its outputs and uses an MLP to score the schedule.The model embedding is learned, while the timestep embedder maps a sinusoidal timestep representation to an embedding.Source: §4.2.2 · p. 6

5. Experiments

Across four image datasets, OMS-DPM is evaluated with DPM-Solver and DDIM against single-model, training-set-best, and random schedules; the experiments also test Stable Diffusion, ablations, and searched-schedule patterns.

5.1. The Effectiveness of OMS-DPM

  • OMS-DPM achieves the best quality-speed trade-offs against state-of-the-art methods and outperforms single-model schedules across datasets, budgets, and both samplers.The reported comparisons cover DPM-Solver and DDIM.Source: Table 1 · p. 7, §5.1 · p. 7, Table 2 · p. 8
  • OMS-DPM outperforms training-set-best and random schedules, while training-set-best schedules often outperform single-model baselines.The authors attribute the gains to predictor generalization from limited training schedules and search for better schedules.Source: §5 · p. 6, §5.1 · p. 7

5.2. Results with Stable Diffusion

  • On Stable Diffusion, OMS-DPM achieves better FID with 12 steps than the best single model with 24 steps.The experiment uses four official checkpoints and DPM-Solver; the checkpoints come from different iterations of the same model.Source: §5.2 · pp. 7–8, Table 3 · p. 8

5.3. Ablation Study

  • Using a model zoo of six models gives the best ablation results, while OMS-DPM improves over baseline at every tested zoo size.The N = 2 result is slightly better than N = 4; the authors suggest its smaller search space may permit more thorough exploration.Source: §5.3 · p. 8, Table 4 · p. 8
  • Less predictor training data degrades performance, but OMS-DPM still beats the baseline in all three tested cases.The authors conclude that predictor training can be made more efficient while retaining promising results.Source: §5.3 · p. 8

5.4. Empirical Observations

  • Under low budgets, smaller models with more steps are more likely to improve quality; with adequate budgets, larger models are suggested.The authors attribute the low-budget pattern to rapidly increasing solver-formula and time-discretization error as NFE decreases.Source: §5.4 · p. 9
  • Where large models work best varies by dataset and sampler: near generated data on ImageNet-64 and some CIFAR-10 runs, but near noise on LSUN-Church.The reported comparisons use DPM-Solver and DDIM on ImageNet-64, and DDIM on CIFAR-10.Source: §5.4 · p. 9
  • Discovered DDIM timestep spacing varies across datasets, and DPM-Solver schedules favor first- or second-order solvers at tight budgets but use third-order solvers with adequate budgets.For Stable Diffusion, schedules mix solver orders at all budgets and prefer first-order when t is close to 0.Source: §5.4 · p. 9

6. Limitations and Future Work

OMS-DPM requires new predictor training data and a newly trained predictor for each new dataset or task, creating substantial overhead. The authors identify efficient model-zoo construction and broader use of combined off-the-shelf models and APIs as future directions.

6. Limitations and Future Work

  • For a new dataset or task, OMS-DPM requires preparing predictor training data and training a new predictor, which incurs substantial overhead.The authors also identify efficient construction of a good model zoo as an open topic.Source: §6 · p. 9
  • The authors suggest that combining off-the-shelf models and APIs with different expertise and complexity could improve performance-efficiency trade-offs across applications.This is presented as a broader perspective beyond schedules of pretrained DPMs.Source: §6 · p. 9

A.1. Results of ablation study

The appendix points to ablation results and searched-schedule visualizations, then describes model-zoo construction and inference-latency measurement.

A.1. Results of ablation study

  • The ablation results using less predictor training data are reported in Table 5.The passage identifies the table but does not provide its numerical results.Source: §A.1 · p. 12

A.2. Demonstration of searched model schedules

  • Searched schedules are shown for all four datasets with both samplers, with ImageNet-64 and LSUN-Church DDIM patterns placed in Figure 6.Other searched patterns are shown in Figure 7.Source: §A.2 · p. 12

B.1. Model Zoo Construction

  • The model zoo can use models with differing complexity or functionality, built by training models or using public ones.For the four unconditional datasets, the authors vary U-Net architecture; for Stable Diffusion, they use four official pretrained models.Source: §B.1 · p. 12
  • The Stable Diffusion models share a VAE and CLIP text encoder, allowing the authors to compose their U-Nets for latent diffusion.The four models share an architecture but differ in training datasets and parameters.Source: §B.1 · p. 13

B.2. Inference Latency

  • Inference latency is averaged over 500 batch inferences on a single A100 GPU, with batch sizes varying by dataset.The batch sizes are 64 for LSUN-Church, 128 for CelebA and ImageNet-64, and 512 for CIFAR-10.Source: §B.2 · p. 13

C.1. Evaluation of DPMs

The appendix specifies evaluation, schedule execution for two samplers, predictor-training data generation, the DPM-Solver predictor adjustment, and predictor hyperparameters.

C.1. Evaluation of DPMs

  • Unconditional DPM evaluation generates 50k images and computes FID against the training set, with generation noise fixed across experiments.Stable Diffusion evaluation uses 30k validation captions and guidance scale s = 1.5, then compares generated images with the raw images for those captions.Source: §C.1 · pp. 13–14

C.2. How to Conduct DPM Sampling for A Model Schedule

  • For DPM-Solver, each group of three schedule entries selects an inactive, first-, second-, or third-order solver step; time is split using a uniform logSNR rule.The schedule determines model assignments within solver steps, and inactive groups are omitted from sampling.Source: §C.2 · p. 14
  • For DDIM, nonzero schedule entries select timesteps from a linear discretization of [0,T], while zero entries are skipped.The search space therefore includes timestep discretization.Source: §C.2 · p. 14

C.3. Schedule-FID Data Generation

  • Predictor-training schedules are sampled from multiple manually specified multinomial distributions to diversify the schedule-FID data.Each unconditional schedule uses 5k generated images for FID; Stable Diffusion uses 1.5k sampled MS-COCO captions.Source: §C.3 · p. 15

C.4. Adjustment of the Sequence Predictor for DPM-Solver

  • For DPM-Solver, the predictor combines embeddings for each three-model group into a solver embedding, which is joined with timestep embeddings before the LSTM.This adjustment matches the grouped-step structure of the solver.Source: §C.4 · p. 15

C.5. Hyperparameter of Predictor

  • The predictor uses a 128-unit, one-layer LSTM and a four-layer MLP, with embedding dimensions set by dataset and sampler.The authors report that they conduct no hyperparameter tuning.Source: §C.5 · p. 15

C.6. Training

The predictor is trained with a ranking loss and validated against ground-truth FID rankings. OMS-DPM then uses the predictor to search model schedules through evolutionary mutation.

C.6. Training

  • The predictor is trained with ranking loss on selected pairs whose ground-truth FID difference exceeds 0.15, using compare ratio 2 and margin 1.0.Pairs are sampled up to compare ratio × batch size.Source: §C.6 · p. 15

C.7. Validating the Reliability of Predictor

  • KD 0.9543 on Stable Diffusion validation schedules, using 2070 training examples and 108 validation examples, indicates the predictor’s effectiveness on unseen schedules.Table 11 reports validation Kendall’s Tau between predicted scores and true FID.Source: §C.7 · p. 15, Table 11 · p. 16

C.8. Evolutionary Search

  • Evolutionary search initializes a schedule near the time budget, repeatedly mutates selected parent schedules, and retains candidates by predicted score.Search runs for a configured number of epochs and removes excess schedules when the population exceeds its cap.Source: §C.8 · p. 16

C.9. Implementation Details and Full Results of Baseline (1)

The paper reports that OMS-DPM outperforms single-model baselines across datasets despite dataset-dependent optimal settings. It also provides baseline configurations, search details, and generated-image examples.

C.9. Implementation Details and Full Results of Baseline (1)

  • OMS-DPM outperforms single-model baselines across datasets, although the best model size, solver order, and time discretization scheme vary by dataset.Quadratic discretization performs better on CIFAR-10, CelebA, and ImageNet-64, while linear performs better on LSUN-Church.Source: §C.9 · p. 16, Table 12 · p. 17, Table 13 · p. 20
  • Baseline configurations use fast first-, second-, and third-order DPM-Solver variants, two DDIM time discretizations, and single-step DPM-Solver for Stable Diffusion.The DPM-Solver and Stable Diffusion setups use uniform logSNR discretization; DDIM uses quadratic and uniform time discretization.Source: §C.9 · p. 16

D. Generated Images

  • Generated-image examples compare OMS-DPM with baseline (1) at the lowest budget across CIFAR-10, CelebA, LSUN-Church, and MS-COCO.The figure captions specify budgets of 700 ms, 650 ms, 4000 ms, and 15 NFE, respectively.Source: D. Generated Images · p. 17, Figure 8 · p. 21, Figure 9 · p. 21
  • The evolutionary search selects a parent by predicted score, mutates it to generate candidates within the time budget, and prunes the population by predicted score.The algorithm’s inputs include a time budget and a trained predictor; its cost function sums model inference latencies.Source: D. Generated Images · p. 19
Loading 2306.08860v1…