PaperlayerView exact revision on arXiv

Source-linked AI summary

Restart Sampling for Improving Generative Processes

Yilun Xu, Mingyang Deng, Xiang Cheng, Yonglong Tian, Ziming Liu, Tommi Jaakkola

arXiv:2306.14878v2cs.LGcs.CVstat.COstat.ML

TL;DR

ODE samplers are fast but plateau, while SDE samplers can attain higher quality at greater sampling time. Restart alternates substantial forward noise with backward ODE steps to combine error contraction and ODE-like discretization. The paper reports improved speed and quality across benchmarks and a better alignment/visual-quality versus diversity trade-off for Stable Diffusion.

  • Problem

    ODE and SDE samplers occupy different speed-quality regimes, motivating a method that balances ODE’s smaller discretization errors with SDE’s contraction of accumulated errors.

    Source: §1 · p. 1, §3 · p. 3

  • Method

    Restart repeats substantial-noise forward steps and backward ODE steps within a selected interval, separating stochasticity from drift.

    Source: §1 · p. 2

  • Results

    Restart surpasses previous ODE and SDE solvers in speed and quality across settings, using 10× fewer steps at the same FID on CIFAR-10 using VP.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Restart improves the balance between text-image alignment/visual quality and diversity in Stable Diffusion models.

    Source: §6 · p. 10

    The method lacks a principled way to select hyperparameters, including the iteration count K and time interval [tmin, tmax].

    Source: §6 · p. 10

Abstract

from arXiv · show

Generative processes that involve solving differential equations, such as diffusion models, frequently necessitate balancing speed and quality. ODE-based samplers are fast but plateau in performance while SDE-based samplers deliver higher sample quality at the cost of increased sampling time. We attribute this difference to sampling errors: ODE-samplers involve smaller discretization errors while stochasticity in SDE contracts accumulated errors. Based on these findings, we propose a novel sampling algorithm called Restart in order to better balance discretization errors and contraction. The sampling method alternates between adding substantial noise in additional forward steps and strictly following a backward ODE. Empirically, Restart sampler surpasses previous SDE and ODE samplers in both speed and accuracy. Restart not only outperforms the previous best SDE results, but also accelerates the sampling speed by 10-fold / 2-fold on CIFAR-10 / ImageNet $64 \times 64$. In addition, it attains significantly better sample quality than ODE samplers within comparable sampling times. Moreover, Restart better balances text-image alignment/visual quality versus diversity than previous samplers in the large-scale text-to-image Stable Diffusion model pre-trained on LAION $512 \times 512$. Code is available at https://github.com/Newbeeer/diffusion_restart_sampling

1 Introduction

ODE and SDE samplers trade speed and quality: ODE discretization errors are smaller, while SDE stochasticity contracts accumulated errors. Restart combines substantial forward noise with backward ODE steps, and experiments report better speed and quality across benchmarks and models.

1 Introduction

  • ODE samplers have smaller discretization errors, while SDE stochasticity contracts accumulated errors, explaining their different performance regimes.SDE contraction applies to both discretization and learned-neural-network approximation errors.Source: §1 · pp. 1–2
  • Restart alternates substantial forward noise with backward ODE steps, separating stochasticity from drift to strengthen error contraction while retaining ODE-like discretization.Repeated cycles strengthen contraction; intervals are typically placed late in simulation, with multiple intervals used for more challenging tasks.Source: §1 · p. 2
  • Restart surpasses previous ODE and SDE solvers in speed and quality across NFEs, datasets, and pre-trained models.It uses 10× fewer steps than the previous best-performing SDEs at the same FID on CIFAR-10 using VP, and 2× fewer steps on ImageNet 64 × 64 with EDM.Source: §1 · p. 2
  • Restart improves the quality-versus-time trade-off in Stable Diffusion, balancing text-image alignment and visual quality against diversity as guidance strength varies.The reported measures are CLIP/Aesthetic scores and FID, respectively.Source: §1 · p. 2
  • The paper reports FID scores of 1.88 on unconditional CIFAR-10 with PFGM++ and 1.36 on class-conditional ImageNet 64 × 64 with EDM.The authors describe these as the best FID scores on commonly used UNet architectures for diffusion models without additional training.Source: §1 · p. 2
  • The paper theoretically analyzes stochasticity’s contraction effect using a Wasserstein-distance upper bound and reports that Restart improves this bound and benchmark performance.The contributions also report improved trade-offs between key metrics on Stable Diffusion.Source: §1 · p. 2

2 Background on Generative Models with Differential Equations

The section introduces differential-equation generative models and describes how diffusion models use forward noise processes and learned reverse dynamics. It also situates Poisson flow generative models as a related approach based on electric-field simulations.

2 Background on Generative Models with Differential Equations

  • Diffusion models and Poisson flow generative models use forward processes to transform data into a smooth distribution and backward processes to reverse that transformation.The backward process iteratively recovers the data distribution.Source: §2 · p. 2
  • Diffusion sampling uses a reverse-time SDE or a marginally equivalent ODE, with a Gaussian prior at sufficiently large T.At time t, the marginal distribution is the data distribution convolved with a Gaussian kernel.Source: §2 · p. 2
  • A neural network estimates the score field, and its estimate replaces the score in the backward SDE or ODE used for sampling.The score estimate is trained with denoising score-matching loss.Source: §2 · p. 3
  • Poisson flow models interpret data as electric charges in an augmented space and simulate differential equations defined by electric field lines.Their neural network approximates the electric field in that space.Source: §2 · p. 3

3 Explaining SDE and ODE performance regimes

ODE samplers tend to perform better at small NFE, while SDE samplers can improve at large NFE. The analysis attributes this shift to the balance between discretization error and stochastic contraction of accumulated error.

3 Explaining SDE and ODE performance regimes

  • Sampling balances sample quality and diversity, often assessed by FID, against sampling time measured by NFE; lower NFE generally worsens FID as larger steps increase discretization error.Lower FID indicates better quality and diversity in the stated metric.Source: §3 · p. 3
  • ODE samplers attain decent quality with limited NFE, but plateau as NFE increases; SDE samplers struggle at low NFE yet improve in the high-NFE regime.The section frames this as the difference in the samplers’ quality-versus-time regimes.Source: §3 · p. 3
  • For first-order solvers, Euler ODE local error is O(δ^2), versus O(δ^(3/2)) for Euler-Maruyama SDE, supporting lower ODE sampling error at larger step sizes.The passage attributes ODE’s better quality at small NFE as likely caused by this lower sampling error.Source: §3 · p. 3
  • At large NFE, neural-network approximation error can dominate as the step size shrinks, and SDE stochasticity contracts accumulated initial error despite greater additional sampling error.The analysis compares upper bounds on Wasserstein-1 distance and states that SDE’s bound becomes smaller as δ approaches 0.Source: §3 · p. 4

4 Harnessing stochasticity with Restart

Restart embeds repeated noisy forward and ODE backward steps within a selected interval to combine contraction with ODE-like discretization. Its analysis describes an error trade-off across iterations, and practical guidance varies interval placement and restart levels by task difficulty.

4 Harnessing stochasticity with Restart

  • Restart is presented as a sampler that incorporates stochasticity while enabling fast generation, with a better Wasserstein upper bound attributed to greater contraction.The section covers the algorithm and its theoretical analysis.Source: §4 · p. 4

4.1 Method

  • Restart inserts repeated forward-noise and backward-ODE steps within [tmin, tmax] into the main backward ODE process.The forward step adds substantial noise without evaluating the neural network; the backward step simulates the ODE using network predictions, and K cycles repeat this sequence.Source: §4.1 · pp. 4–5

4.2 Analysis

  • Restart’s separated noise and drift yield a smaller contraction factor than SDE, while its backward process retains ODE-order discretization error.The contracted error decreases exponentially with K, while additional error increases linearly, implying a balance point for choosing K.Source: §4.2 · p. 6

4.3 Practical considerations

  • Practical guidance places the interval late in the main backward process and suggests wider intervals for weaker models or more challenging datasets.One-level Restart is reported sufficient for CIFAR-10, while multi-level Restart improves performance on more challenging ImageNet.Source: §4.3 · p. 6

5 Experiments

Experiments compare Restart with ODE and SDE samplers on synthetic errors, standard image-generation benchmarks, and text-to-image generation. Across these settings, Restart achieves favorable error, speed, and quality trade-offs, with results varying by task and model.

5.1 Additional sampling error versus contracted error

  • Restart has lower contracted error than ODE and SDE at a given additional sampling error, and lower total error in the [0.8, 0.85] additional-error range.Figure 2 also shows lower total error at a given NFE.Source: §5.1 · p. 7, Figure 2 · p. 7

5 Experiments

  • The experiments evaluate quality and inference speed using FID and NFE on pretrained VP/EDM/PFGM++ models and assess text-to-image trade-offs on Stable Diffusion v1.5.Standard benchmarks include CIFAR-10 and ImageNet 64 × 64; the text-to-image model is pretrained on LAION-5B at 512×512.Source: §5 · pp. 6–7, §5.2 · p. 7, §5.3 · p. 9

5.2 Experiments on standard benchmarks

  • Restart achieves 10-fold / 2-fold acceleration over previous best SDE results at the same FID on CIFAR-10 (VP) / ImageNet 64 × 64 (EDM), and establishes FID scores of 1.88 on uncond. CIFAR-10 with PFGM++ and 1.36 on class-cond. ImageNet 64 × 64 with EDM.It also obtains better FID than Heun, with the gap increasing with NFE; Figure 3 and Table 1 show benchmark comparisons.Source: §5.2 · p. 8, Figure 3 · p. 8, Table 1 · p. 9
  • With NFE from 16 to 36, Restart consistently outperforms DPM-solver on CIFAR-10 VP, reaching FID 2.11 at NFE 30.The reported FID is lower than previously reported numbers, including an SDE result with NFE greater than 1000.Source: §5.2 · pp. 8–9, Figure 4 · p. 8
  • FID initially improves and then worsens as Restart iterations K increase, with the turning point occurring at a smaller K for the stronger EDM model.The authors suggest greater Restart strength for weaker or smaller architectures.Source: §5.2 · p. 9, Figure 5 · p. 9

5.3 Experiments on large-scale text-to-image model

  • At the same CLIP/Aesthetic scores, Restart achieves better FID in most cases using 132 function evaluations, and generally provides better alignment or visual quality at the same FID.For high CLIP/Aesthetic scores, its FID is substantially lower than other samplers; at low scores, DDPM can have comparable FID but is less time-efficient. Figure 6 plots these trade-offs.Source: §5.3 · pp. 9–10, Figure 6 · p. 9
  • At guidance weight w = 8, Restart images have more accurate details and better visual quality than DDIM, and are more photorealistic than DDPM.Figure 7 compares generated images for four prompts using the same random seeds.Source: §5.3 · p. 10, Figure 7 · p. 10

6 Conclusion and future direction

Restart interweaves substantial forward noise with a backward ODE to balance contraction and discretization error, and the paper reports quality–time and text-to-image trade-off improvements. It also identifies principled hyperparameter selection as an unresolved limitation.

6 Conclusion and future direction

  • Restart combines substantial forward noise with a backward ODE, theoretically strengthening stochastic contraction while maintaining ODE-level discretization error.The paper also reports improved quality–time balance and text-image alignment/visual quality versus diversity in Stable Diffusion.Source: §6 · p. 10
  • The algorithm lacks principled hyperparameter selection for iteration count K and interval [tmin, tmax]; current choices use a heuristic based on model strength and task difficulty.Source: §6 · p. 10

A Proofs of Main Theoretical Results

A.1 Main Result

The appendix bounds sampling error for ODE and SDE solvers and analyzes Restart under stated regularity and support assumptions. It also examines mixing across Restart iterations under an exact ODE.

A.1 Main Result

  • The ODE error bound combines a discretization-and-approximation term with a bound on the distance between the learned and true trajectories.The proof uses a triangular inequality to combine these contributions.Source: §A.1 · p. 14
  • Restart’s stated bound separates an upper bound on contracted error from an upper bound on additional sampling error.Source: §A.1 · p. 14

A.2 Mixing under Restart with exact ODE

  • For exact ODE, the appendix analyzes mixing under Restart by recursively applying couplings across iterations.Source: §A.2 · p. 15

A.3 W1 discretization bound

The appendix develops discretization and mixing bounds for ODE/SDE processes and describes the time discretization and implementation choices used for Restart. Its experimental setup specifies baseline samplers, Restart configurations, and classifier-free guidance.

A.3 W1 discretization bound

  • The ODE discretization bound assumes time regularity, bounded score magnitude, and state-space Lipschitz continuity of the learned score field.The lemma compares the true and learned ODE trajectories over the interval from tmax to tmin.Source: §A.3 · p. 16
  • The appendix defines post-iteration Restart variables and uses recursive comparisons of trajectories to bound accumulated discretization error.The proof excerpts invoke the ODE bound and Lipschitz condition when comparing subsequent backward steps.Source: §A.3 · pp. 16–17

A.4 Mixing Bounds

  • The mixing analysis includes a hitting-time argument under a bounded-support assumption and a Gaussian-tail bound for coupling processes.The excerpts define a stopping time at which coupled processes meet and compare it with a dominating process.Source: §A.4 · p. 19

A.5 TV Overlap

  • For two Gaussian perturbations with common covariance, the stated total variation distance is 1 − 2Q(r), where r is the input separation divided by 2σ.The passage also gives the upper bound 1 − 2r/√(r^2 + 1).Source: §A.5 · p. 20

B.1 EDM Discretization Scheme

  • The EDM schedule uses ρ = 7 and includes tmin among the main-process time steps so the Restart interval joins the backward process.The same discretization scheme is used within Restart backward intervals and the main backward process.Source: §B.1 · p. 21

B.2 Restart Algorithm

  • Restart repeats a forward-noise step and a backward ODE within configured intervals, with the pseudocode also allowing multiple Restart levels.The algorithm specifies steps per Restart backward process, interval endpoints, iteration counts, and main-process steps; Heun can be replaced by other ODE solvers.Source: §B.2 · p. 21

C.1 Configurations for Baselines

  • Baseline comparisons use shared EDM discretization for Vanilla SDE, Improved SDE, Heun, and Restart on standard benchmarks.The configurations also include DDIM, Heun, and DDPM for Stable Diffusion.Source: §C.1 · p. 22

C.2 Configurations for Restart

  • Restart configurations are selected from a limited manual search because the configuration space is too large for comprehensive enumeration.Stable Diffusion experiments use 66 total steps, while the search procedure averages five configuration sets per experiment and calls systematic selection a future need.Source: §C.2 · p. 23

C.5.1 Discrete Dataset

This experimental subsection describes a synthetic dataset, its score-model architecture, hyperparameter sweeps, and the construction of Pareto-frontier plots.

C.5.1 Discrete Dataset

  • The discrete dataset contains 2,000 projected points from a two-component Gaussian mixture, scaled to unit variance per dimension.The mixture is generated in R^4 and projected into R^20 to model points on a lower-dimensional, multimodal manifold.Source: §C.5.1 · p. 23

C.5.2 Model Architecture

  • A latent-size-64 MLP learns the score function using preconditioning and denoising score matching adapted from [13].The training objective is cited to.Source: §C.5.2 · p. 24

C.5.3 Varying Hyperparameters

  • The experiments vary SDE and ODE step counts and Restart iteration counts and backward-step counts across specified grids.For SDEs, the sweep also varies the multiplier on β(t); Improved SDE experiments vary Schurn and step counts.Source: §C.5.3 · p. 24

C.5.4 Plotting the Pareto frontier

  • The plotted Pareto frontiers connect prefix minima after sorting by additional sampling error or NFE, respectively.The corresponding paired quantities are contracted error and FID.Source: §C.5.4 · p. 24, Figure 8 · p. 24

D.1 Numerical Results

The numerical-results section reports benchmark evaluations for image generation and Stable Diffusion, with the supplied passages identifying evaluation protocols and table scopes rather than numerical outcomes.

D.1 Numerical Results

  • FID evaluation uses 50,000 generated samples, with three seeded calculations and the minimum reported for CIFAR-10 and ImageNet 64 × 64.The tables cover baseline and Restart results across CIFAR-10 and ImageNet settings.Source: §D.1 · p. 24, Table 2 · p. 25, Table 3 · p. 26
  • Stable Diffusion results are evaluated at guidance weights 2, 3, 5, and 8 using zero-shot FID, CLIP, and Aesthetic scores.The evaluation uses 5,000 COCO validation prompts and averages CLIP and Aesthetic scores over 5,000 generated samples.Source: §D.1 · p. 25, Table 6 · p. 27, Table 7 · p. 27
  • The supplied Stable Diffusion table captions identify numerical results by guidance weight, including the w = 5 table.The supplied captions do not expose the table cell values.Source: Table 8 · p. 28, Table 10 · p. 28

D.2 Sensitivity Analysis of Hyper-parameters

The sensitivity analysis reports how Restart parameters relate to sample quality and discusses a practical rule for choosing Restart strength. The following image section presents generated samples from four samplers.

D.2 Sensitivity Analysis of Hyper-parameters

  • On VP, FID is minimized near tmin = 0, while Restart does not outperform the EDM ODE baseline in the described setting.The tested interval length is 0.3, about twice that used in Table 1; the authors suggest stronger Restart for weaker or smaller architectures.Source: §D.2 · p. 29
  • Longer Restart intervals may increase contraction while also adding sampling error, producing a V-shaped trend; tmax near the dataset radius usually supports effective mixing when tmin is small.These observations describe the sensitivity to Restart interval length.Source: §D.2 · p. 29

E Extended Generated Images

F Heun’s method is DPM-Solver-2 (with r2 = 1)

The passage relates DPM-Solver-1 to EDM through a time-dependent scaling and connects Heun’s method to DPM-Solver-2.

F Heun’s method is DPM-Solver-2 (with r2 = 1)

G Broader Impact

Restart is presented as a sampler that outperforms previous samplers for diffusion models and PFGM++. The broader-impact discussion notes potential beneficial uses and misuse risks, and calls for robustness guarantees.

G Broader Impact

  • Restart outperforms previous samplers for diffusion models and PFGM++, and is described as suitable for high-quality, rapid generation.Source: G Broader Impact · p. 31
  • Realistic image and audio generation could advance entertainment, advertising, and education.Source: G Broader Impact · p. 34
  • Misuse in deepfakes could enable scams and misinformation, motivating further research on robustness guarantees and ethical use.Source: G Broader Impact · p. 34

Related papers

Loading 2306.14878v2…