Title: Restart Sampling for Improving Generative Processes
Venue: NeurIPS 2023
Accepted: yes
Keywords: Generative models

diffusion models

PFGM

sampling
Abstract: 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 \textit{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

Restart Sampling for Improving Generative Processes

Abstract

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 × 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 × 512 . Code is available at

https://github.com/Newbeeer/diffusion_restart_sampling

1 Introduction

Deep generative models based on differential equations, such as diffusion models and Poission flow generative models, have emerged as powerful tools for modeling high-dimensional data, from image synthesis [23, 9, 13, 27, 28] to biological data [10, 26]. These models use iterative backward processes that gradually transform a simple distribution ( e.g . , Gaussian in diffusion models) into a complex data distribution by solving a differential equations. The associated vector fields (or drifts) driving the evolution of the differential equations are predicted by neural networks. The resulting sample quality can be often improved by enhanced simulation techniques but at the cost of longer sampling times.

Prior samplers for simulating these backward processes can be categorized into two groups: ODEsamplers whose evolution beyond the initial randomization is deterministic, and SDE-samplers where the generation trajectories are stochastic. Several works [23, 12, 13] show that these samplers demonstrate their advantages in different regimes, as depicted in Fig. 1(b). ODE solvers [22, 16, 13] result in smaller discretization errors, allowing for decent sample quality even with larger step sizes ( i.e . , fewer number of function evaluations (NFE)). However, their generation quality plateaus rapidly. In contrast, SDE achieves better quality in the large NFE regime, albeit at the expense of increased sampling time. To better understand these differences, we theoretically analyze SDE performance: the

(a)

(b)

0=1

D(Pmodel Paata)

ODE

ODE

SDE

SmallNFERegime

LargeNFERegime

SDE

t=tmax

t=tmin

Restart ..

xK

p!

noise

NFE

stochasticity in SDE contracts accumulated error, which consists of both the discretization error along the trajectories as well as the approximation error of the learned neural network relative to the ground truth drift ( e.g . , score function in diffusion model [23]). The approximation error dominates when NFE is large (small discretization steps), explaining the SDE advantage in this regime. Intuitively, the stochastic nature of SDE helps "forget" accumulated errors from previous time steps.

Inspired by these findings, we propose a novel sampling algorithm called Restart , which combines the advantages of ODE and SDE. As illustrated in Fig. 1(a), the Restart sampling algorithm involves K repetitions of two subroutines in a pre-defined time interval: a Restart forward process that adds a substantial amount of noise, akin to "restarting" the original backward process, and a Restart backward process that runs the backward ODE. The Restart algorithm separates the stochasticity from the drifts, and the amount of added noise in the Restart forward process is significantly larger than the small single-step noise interleaving with drifts in previous SDEs such as [23, 13], thus amplifying the contraction effect on accumulated errors. By repeating the forward-backward cycle K times, the contraction effect introduced in each Restart iteration is further strengthened. The deterministic backward processes allow Restart to reduce discretization errors, thereby enabling step sizes comparable to ODE. To maximize the contraction effects in practice, we typically position the Restart interval towards the end of the simulation, where the accumulated error is larger. Additionally, we apply multiple Restart intervals to further reduce the initial errors in more challenging tasks.

Experimentally, Restart consistently surpasses previous ODE and SDE solvers in both quality and speed over a range of NFEs, datasets, and pre-trained models. Specifically, Restart accelerates the previous best-performing SDEs by 10 × fewer steps for the same FID score on CIFAR-10 using VP [23] ( 2 × fewer steps on ImageNet 64 × 64 with EDM [13]), and outperforms ODE solvers even in the small NFE regime. When integrated into previous state-of-the-art pre-trained models, Restart further improves performance, achieving FID scores of 1.88 on unconditional CIFAR-10 with PFGM++ [28], and 1.36 on class-conditional ImageNet 64 × 64 with EDM. To the best of our knowledge, these are the best FID scores obtained on commonly used UNet architectures for diffusion models without additional training. We also apply Restart to the practical application of text-to-image Stable Diffusion model [19] pre-trained on LAION 512 × 512 . Restart more effectively balances text-image alignment/visual quality (measured by CLIP/Aesthetic scores) and diversity (measured by FID score) with a varying classifier-free guidance strength, compared to previous samplers.

Our contributions can be summarized as follows: (1) We investigate ODE and SDE solvers and theoretically demonstrate the contraction effect of stochasticity via an upper bound on the Wasserstein distance between generated and data distributions (Sec 3); (2) We introduce the Restart sampling, which better harnesses the contraction effect of stochasticity while allowing for fast sampling. The sampler results in a smaller Wasserstein upper bound (Sec 4); (3) Our experiments are consistent with the theoretical bounds and highlight Restart's superior performance compared to previous samplers on standard benchmarks in terms of both quality and speed. Additionally, Restart improves the trade-off between key metrics on the Stable Diffusion model (Sec 5).

2 Background on Generative Models with Differential Equations

Many recent successful generative models have their origin in physical processes, including diffusion models [9, 23, 13] and Poisson flow generative models [27, 28]. These models involve a forward process that transforms the data distribution into a chosen smooth distribution, and a backward process that iteratively reverses the forward process. For instance, in diffusion models, the forward process is the diffusion process with no learned parameters:

where σ ( t ) is a predefined noise schedule increasing with t , and W t ∈ R d is the standard Wiener process. For simplicity, we omit an additional scaling function for other variants of diffusion models as in EDM [13]. Under this notation, the marginal distribution at time t is the convolution of data distribution p 0 = p data and a Gaussian kernel, i.e . , p t = p 0 ∗ N ( 0 , σ 2 ( t ) I d × d ) . The prior distribution is set to N ( 0 , σ 2 ( T ) I d × d ) since p T is approximately Gaussian with a sufficiently large T . Sampling of diffusion models is done via a reverse-time SDE [1] or a marginally-equivalent ODE [23]:

where ∇ x log p t ( x ) in the drift term is the score of intermediate distribution at time t . W.l.o.g we set σ ( t ) = t in the remaining text, as in [13]. Both processes progressively recover p 0 from the prior distribution p T while sharing the same time-dependent distribution p t . In practice, we train a neural network s θ ( x, t ) to estimate the score field ∇ x log p t ( x ) by minimizing the denoising score-matching loss [25]. We then substitute the score ∇ x log p t ( x ) with s θ ( x, t ) in the drift term of above backward SDE (Eq. (1))/ODE (Eq. (2)) for sampling.

Recent work inspired by electrostatics has not only challenged but also integrated diffusion models, notably PFGM/PFGM++, enhances performance in both image and antibody generation [27, 28, 10]. They interpret data as electric charges in an augmented space, and the generative processes involve the simulations of differential equations defined by electric field lines. Similar to diffusion models, PFGMs train a neural network to approximate the electric field in the augmented space.

3 Explaining SDE and ODE performance regimes

To sample from the aforementioned generative models, a prevalent approach employs general-purpose numerical solvers to simulate the corresponding differential equations. This includes Euler and Heun's 2nd method [2] for ODEs (e.g., Eq. (2)), and Euler-Maruyama for SDEs (e.g., Eq. (1)). Sampling algorithms typically balance two critical metrics: (1) the quality and diversity of generated samples, often assessed via the Fréchet Inception Distance (FID) between generated distribution and data distribution [7] (lower is better), and (2) the sampling time, measured by the number of function evaluations (NFE). Generally, as the NFE decreases, the FID score tends to deteriorate across all samplers. This is attributed to the increased discretization error caused by using a larger step size in numerical solvers.

However, as illustrated in Fig. 1(b) and observed in previous works on diffusion models [23, 22, 13], the typical pattern of the quality vs time curves behaves differently between the two groups of samplers, ODE and SDE. When employing standard numerical solvers, ODE samplers attain a decent quality with limited NFEs, whereas SDE samplers struggle in the same small NFE regime. However, the performance of ODE samplers quickly reaches a plateau and fails to improve with an increase in NFE, whereas SDE samplers can achieve noticeably better sample quality in the high NFE regime. This dilemma raises an intriguing question: Why do ODE samplers outperform SDE samplers in the small NFE regime, yet fall short in the large NFE regime?

The first part of the question is relatively straightforward to address: given the same order of numerical solvers, simulation of ODE has significantly smaller discretization error compared to the SDE. For example, the first-order Euler method for ODE results in a local error of O ( δ 2 ) , whereas the first-order Euler-Maruyama method for SDEs yeilds a local error of O ( δ 3 2 ) (see e.g . , Theorem 1 of [4]), where δ denotes the step size. As O ( δ 3 2 ) ≫ O ( δ 2 ) , ODE simulations exhibit lower sampling errors than SDEs, likely causing the better sample quality with larger step sizes in the small NFE regime.

In the large NFE regime the step size δ shrinks and discretization errors become less significant for both ODEs and SDEs. In this regime it is the approximation error -error arising from an

inaccurate estimation of the ground-truth vector field by the neural network s θ -starts to dominate the sampling error. We denote the discretized ODE and SDE using the learned field s θ as ODE θ and SDE θ , respectively. In the following theorem, we evaluate the total errors from simulating ODE θ and SDE θ within the time interval [ t min , t max ] ⊂ [0 , T ] . This is done via an upper bound on the Wasserstein-1 distance between the generated and data distributions at time t min. We characterize the accumulated initial sampling errors up until t max by total variation distances. Below we show that the inherent stochasticity of SDEs aids in contracting these initial errors at the cost of larger additional sampling error in [ t min , t max ] . Consequently, SDE results in a smaller upper bound as the step size δ nears 0 (pertaining to the high NFE regime).

Theorem 1 (Informal) . Let t max be the initial noise level and p t denote the true distribution at noise level t . Let p ODE θ t , p SDE θ t denote the distributions of simulating ODE θ , SDE θ respectively. Assume that ∀ t ∈ [ t min , t max ] , ∥ x t ∥ < B/ 2 for any x t in the support of p t , p ODE θ t or p SDE θ t . Then

In the above, U = BL 1 /t min + L 2 1 t 2 max /t 2 min , λ < 1 is a contraction factor, L 1 and ϵ approx are uniform bounds on ∥ ts θ ( x t , t ) ∥ and the approximation error ∥ t ∇ x log p t ( x ) -ts θ ( x, t ) ∥ for all x t , t , respectively. O () hides polynomial dependency on various Lipschitz constants and dimension.

We defer the formal version and proof of Theorem 1 to Appendix A.1. As shown in the theorem, the upper bound on the total error can be decomposed into upper bounds on the contracted error and additional sampling error . TV ( p ODE θ t max , p t max ) and TV ( p SDE θ t max , p t max ) correspond to the initial errors accumulated from both approximation and discretization errors during the simulation of the backward process, up until time t max. In the context of SDE, this accumulated error undergoes contraction by a factor of 1 -λe -BL 1 /t min -L 2 1 t 2 max /t 2 min within [ t min , t max ] , due to the effect of adding noise. Essentially, the minor additive Gaussian noise in each step can drive the generated distribution and the true distribution towards each other, thereby neutralizing a portion of the initial accumulated error.

The other term related to additional sampling error includes the accumulation of discretization and approximation errors in [ t min , t max ] . Despite the fact that SDE incurs a higher discretization error than ODE ( O ( √ δ ) versus O ( δ ) ), the contraction effect on the initial error is the dominant factor impacting the upper bound in the large NFE regime where δ is small. Consequently, the upper bound for SDE is significantly lower. This provides insight into why SDE outperforms ODE in the large NFE regime, where the influence of discretization errors diminishes and the contraction effect dominates. In light of the distinct advantages of SDE and ODE, it is natural to ask whether we can combine their strengths. Specifically, can we devise a sampling algorithm that maintains a comparable level of discretization error as ODE, while also benefiting from, or even amplifying, the contraction effects induced by the stochasticity of SDE? In the next section, we introduce a novel algorithm, termed Restart , designed to achieve these two goals simultaneously.

4 Harnessing stochasticity with Restart

In this section, we present the Restart sampling algorithm, which incorporates stochasticity during sampling while enabling fast generation. We introduce the algorithm in Sec 4.1, followed by a theoretical analysis in Sec 4.2. Our analysis shows that Restart achieves a better Wasserstein upper bound compared to those of SDE and ODE in Theorem 1 due to greater contraction effects.

4.1 Method

In the Restart algorithm, simulation performs a few repeated back-and-forth steps within a pre-defined time interval [ t min , t max ] ⊂ [0 , T ] , as depicted in Figure 1(a). This interval is embedded into the simulation of the original backward ODE referred to as the main backward process , which runs from T to 0 . In addition, we refer to the backward process within the Restart interval [ t min , t max ] as the Restart backward process , to distinguish it from the main backward process.

Starting with samples at time t min, which are generated by following the main backward process, the Restart algorithm adds a large noise to transit the samples from t min to t max with the help of

the forward process. The forward process does not require any evaluation of the neural network s θ ( x, t ) , as it is generally defined by an analytical perturbation kernel capable of transporting distributions from t min to t max. For instance, in the case of diffusion models, the perturbation kernel is N ( 0 , ( σ ( t max ) 2 -σ ( t min ) 2 ) I d × d ) . The added noise in this step induces a more significant contraction compared to the small, interleaved noise in SDE. The step acts as if partially restarting the main backward process by increasing the time. Following this step, Restart simulates the backward ODE from t max back to t min using the neural network predictions as in regular ODE. We repeat these forward-backward steps within [ t min , t max ] interval K times in order to further derive the benefit from contraction. Specifically, the forward and backward processes in the i th iteration ( i ∈ { 0 , . . . , K -1 } ) proceed as follows:

where the initial x 0 t min is obtained by simulating the ODE until t min: x 0 t min = ODE θ ( x T , T → t min ) , and the noise ε t min → t max is sampled from the corresponding perturbation kernel from t min to t max. The Restart algorithm not only adds substantial noise in the Restart forward process (Eq. (3)), but also separates the stochasticity from the ODE, leading to a greater contraction effect, which we will demonstrate theoretically in the next subsection. For example, we set [ t min , t max ] = [0 . 05 , 0 . 3] for the VP model [13] on CIFAR-10. Repetitive use of the forward noise effectively mitigates errors accumulated from the preceding simulation up until t max. Furthermore, the Restart algorithm does not suffer from large discretization errors as it is mainly built from following the ODE in the Restart backward process (Eq. (4)). The effect is that the Restart algorithm is able to reduce the total sampling errors even in the small NFE regime. Detailed pseudocode for the Restart sampling process can be found in Algorithm 2, Appendix B.2.

4.2 Analysis

We provide a theoretical analysis of the Restart algorithm under the same setting as Theorem 1. In particular, we prove the following theorem, which shows that Restart achieves a much smaller contracted error in the Wasserstein upper bound than SDE (Theorem 1), thanks to the separation of the noise from the drift, as well as the large added noise in the Restart forward process (Eq. (3)). The repetition of the Restart cycle K times further leads to a enhanced reduction in the initial accumulated error. We denote the intermediate distribution in the i th Restart iteration, following the discretized trajectories and the learned field s θ , as p Restart θ ( i ) t ∈ [ t min ,t max ] .

Theorem 2 (Informal) . Under the same setting of Theorem 1, assume K ≤ C L 2 ( t max -t min ) for some universal constant C . Then

where λ < 1 is the same contraction factor as Theorem 1. O () hides polynomial dependency on various Lipschitz constants, dimension.

Proof sketch. To bound the total error, we introduce an auxiliary process q Restart θ ( i ) t ∈ [ t min ,t max ] , which initiates from true distribution p t max and performs the Restart iterations. This process differs from p Restart θ ( i ) t ∈ [ t min ,t max ] only in its initial distribution at t max ( p t max versus p Restart θ (0) t max ). We bound the total error by the following triangular inequality:

To bound the contracted error, we construct a careful coupling process between two individual trajectories sampled from p Restart θ ( i ) t min and q Restart θ ( i ) t min , i = 0 , . . . , K -1 . Before these two trajectories converge, the Gaussian noise added in each Restart iteration is chosen to maximize the probability of the two trajectories mapping to an identical point, thereby maximizing the mixing rate in TV. After converging, the two processes evolve under the same Gaussian noise, and will stay converged as their drifts are the same. Lastly, we convert the TV bound to W 1 bound by multiplying B . The bound on the additional sampling error echoes the ODE analysis in Theorem 1: since the noise-injection and ODE-simulation stages are separate, we do not incur the higher discretization error of SDE.

We defer the formal version and proof of Theorem 2 to Appendix A.1. The first term in RHS bounds the contraction on the initial error at time t max and the second term reflects the additional sampling error of ODE accumulated across repeated Restart iterations. Comparing the Wasserstein upper bound of SDE and ODE in Theorem 1, we make the following three observations: (1) Each Restart iteration has a smaller contraction factor 1 -λ compared to the one in SDE, since Restart separates the large additive noise (Eq. (3)) from the ODE (Eq. (4)). (2) Restart backward process (Eq. (4)) has the same order of discretization error O ( δ ) as the ODE, compared to O ( √ δ ) in SDE. Hence, the Restart allows for small NFE due to ODE-level discretization error. (3) The contracted error further diminishes exponentially with the number of repetitions K though the additional error increases linearly with K . It suggests that there is a sweet spot of K that strikes a balance between reducing the initial error and increasing additional sampling error. Ideally, one should pick a larger K when the initial error at time t max greatly outweigh the incurred error in the repetitive backward process from t max to t min. We provide empirical evidences in Sec 5.2.

While Theorem 1 and Theorem 2 compare the upper bounds on errors of different methods, we provide empirical validation in Section 5.1 by directly calculating these errors, showing that the Restart algorithm indeed yields a smaller total error due to its superior contraction effects. The main goal of Theorem 1 and Theorem 2 is to study how the already accumulated error changes using different samplers, and to understand their ability to self-correct the error by stochasticity. In essence, these theorems differentiate samplers based on their performance post-error accumulation. For example, by tracking the change of accumulated error, Theorem 1 shed light on the distinct "winning regions" of ODE and SDE: ODE samplers have smaller discretization error and hence excel at the small NFE regime. In contrast, SDE performs better in large NFE regime where the discretization error is negligible and its capacity to contract accumulated errors comes to the fore.

4.3 Practical considerations

The Restart algorithm offers several degrees of freedom, including the time interval [ t min , t max ] and the number of restart iterations K . Here we provide a general recipe of parameter selection for practitioners, taking into account factors such as the complexity of the generative modeling tasks and the capacity of the network. Additionally, we discuss a stratified, multi-level Restart approach that further aids in reducing simulation errors along the whole trajectories for more challenging tasks.

Where to Restart? Theorem 2 shows that the Restart algorithm effectively reduces the accumulated error at time t max by a contraction factor in the Wasserstein upper bound. These theoretical findings inspire us to position the Restart interval [ t min , t max ] towards the end of the main backward process, where the accumulated error is more substantial. In addition, our empirical observations suggest that a larger time interval t max -t min is more beneficial for weaker/smaller architectures or more challenging datasets. Even though a larger time interval increases the additional sampling error, the benefits of the contraction significantly outweighs the downside, consistent with our theoretical predictions. We leave the development of principled approaches for optimal time interval selection for future works.

Multi-level Restart For challenging tasks that yield significant approximation errors, the backward trajectories may diverge substantially from the ground truth even at early stage. To prevent the ODE simulation from quickly deviating from the true trajectory, we propose implementing multiple Restart intervals in the backward process, alongside the interval placed towards the end. Empirically, we observe that a 1 -level Restart is sufficient for CIFAR-10, while for more challenging datasets such as ImageNet [5], a multi-level Restart results in enhanced performance [5].

5 Experiments

In Sec 5.1, we first empirically verify the theoretical analysis relating to the Wasserstein upper bounds. We then evaluate the performance of different sampling algorithms on standard image generation benchmarks, including CIFAR-10 [14] and ImageNet 64 × 64 [5] in Sec 5.2. Lastly, we employ Restart on text-to-image generation, using Stable Diffusion model [19] pre-trained on LAION-5B [21] with resolution 512 × 512 , in Sec 5.3.

5.1 Additional sampling error versus contracted error

Our proposed Restart sampling algorithm demonstrates a higher contraction effect and smaller addition sampling error compared to SDE, according to Theorem 1 and Theorem 2. Although our

0.65

0.70

0.75

0.80

0.85

0.90

Additional Sampling error

0.725

0.750

0.775

0.800

0.825

0.850

0.875

Contracted error

ODE

SDE

(a)

0.65

0.70

0.75

0.80

0.85

0.90

0.95

1.00

Additional Sampling error

0.70

0.75

0.80

0.85

0.90

0.95

1.00

Total Error

ODE

SDE

(b)

20

40

80

160

320

NFE

0.750

0.775

0.800

0.825

0.850

0.875

Total Error

ODE

SDE

(c)

theoretical analysis compares the upper bounds of the total, contracted and additional sampling errors, we further verify their relative values through a synthetic experiment.

Setup We construct a 20 -dimensional dataset with 2000 points sampled from a Gaussian mixture, and train a four-layer MLP to approximate the score field ∇ x log p t . We implement the ODE, SDE, and Restart methods within a predefined time range of [ t min , t max ] = [1 . 0 , 1 . 5] , where the process outside this range is conducted via the first-order ODE. To compute various error types, we define the distributions generated by three methods as outlined in the proof of Theorem 2 and directly gauge the errors at end of simulation t = 0 instead of t = t min: (1) the generated distribution as p Sampler 0 , where Sampler ∈ { ODE θ , SDE θ , Restart θ ( K ) } ; (2) an auxiliary distribution q Sampler 0 initiating from true distribution p t max at time t max. The only difference between p Sampler 0 and q Sampler 0 is their initial distribution at t max ( p ODE θ t max versus p t max ); and (3) the true data distribution p 0 . In line with Theorem 2, we use Wasserstein-1 distance W 1 ( p Sampler 0 , q Sampler 0 ) / W 1 ( q Sampler 0 , p 0 ) to measure the contracted error / additional sampling error, respectively. Ultimately, the total error corresponds to W 1 ( p Sampler 0 , p 0 ) . Detailed information about dataset, metric and model can be found in the Appendix C.5.

Results In our experiment, we adjust the parameters for all three processes and calculate the total, contracted, and additional sampling errors across all parameter settings. Figure 2(a) depicts the Pareto frontier of additional sampling error versus contracted error. We can see that Restart consistently achieves lower contracted error for a given level of additional sampling error, compared to both the ODE and SDE methods, as predicted by theory. In Figure 2(b), we observe that the Restart method obtains a smaller total error within the additional sampling error range of [0 . 8 , 0 . 85] . During this range, Restart also displays a strictly reduced contracted error, as illustrated in Figure 2(a). This aligns with our theoretical analysis, suggesting that the Restart method offers a smaller total error due to its enhanced contraction effects. From Figure 2(c), Restart also strikes an better balance between efficiency and quality, as it achieves a lower total error at a given NFE.

5.2 Experiments on standard benchmarks

32

64

128

256

512

1024

NFE

2.0

2.2

2.4

2.6

2.8

3.0

3.2

3.4

FID score (50K)

ODE

Gonna Go Fast

Improved SDE

(a) CIFAR-10, VP

32

64

128

256

512

1024

NFE

1.50

1.75

2.00

2.25

2.50

2.75

3.00

FID score (50K)

128

256

512

1024

1.4

1.6

1.8

ODE

Improved SDE

(b) ImageNet

64

×

64

, EDM

To evaluate the sample quality and inference speed, we report the FID score [7] (lower is better) on 50K samplers and the number of function evaluations (NFE). We borrow the pretrained VP/EDM/PFGM++

15

20

25

30

35

NFE

2.0

2.2

2.4

2.6

2.8

3.0

3.2

3.4

FID score (50K)

DPM Solver

models on CIFAR-10 or ImageNet 64 × 64 from [13, 28]. We also use the EDM discretization scheme [13] (see Appendix B.1 for details) during sampling.

For the proposed Restart sampler, the hyperparameters include the number of steps in the main/Restart backward processes, the number of Restart iteration K , as well as the time interval [ t min , t max ] . We pick the t min and t max from the list of time steps in EDM discretization scheme with a number of steps 18 . For example, for CIFAR-10 (VP) with NFE =75 , we choose t min =0 . 06 , t max =0 . 30 , K =10 , where 0 . 30 / 0 . 06 is the 12 th / 14 th time step in the EDM scheme. We also adopt EDM scheme for the Restart backward process in [ t min , t max ] . In addition, we apply the multi-level Restart strategy (Sec 4.3) to mitigate the error at early time steps for the more challenging ImageNet 64 × 64 . We provide the detailed Restart configurations in Appendix C.2.

For SDE, we compare with the previously best-performing stochastic samplers proposed by [13] ( Improved SDE ). We use their optimal hyperparameters for each dataset. We also report the FID scores of the adaptive SDE [12] ( Gonna Go Fast ) on CIFAR-10 (VP). Since the vanilla reversediffusion SDE [23] has a significantly higher FID score, we omit its results from the main charts and defer them to Appendix D. For ODE samplers, we compare with the Heun's 2 nd order method [2] ( Heun ), which arguably provides an excellent trade-off between discretization errors and NFE [13]. To ensure a fair comparison, we use Heun's method as the sampler in the main/Restart backward processes in Restart.

We report the FID score versus NFE in Figure 3(a) and Table 1 on CIFAR-10, and Figure 3(b) on ImageNet 64 × 64 with EDM. Our main findings are: (1) Restart outperforms other SDE or ODE samplers in balancing quality and speed, across datasets and models. As demonstrated in the figures, Restart achieves a 10-fold / 2-fold acceleration compared to previous best SDE results on CIFAR-10 (VP) / ImageNet 64 × 64 (EDM) at the same FID score. In comparison to ODE sampler (Heun), Restart obtains a better FID score, with the gap increasing significantly with NFE. (2) For stronger models such as EDM and PFGM++, Restart further improve over the ODE baseline on CIFAR-10. In contrast, the Improved SDE negatively impacts performance of EDM, as also observed in [13]. It suggests that Restart incorporates stochasticity more effectively. (3) Restart establishes new state-of-the-art FID scores for UNet architectures without additional training. In particular, Restart achieves FID scores of 1.36 on class-cond. ImageNet 64 × 64 with EDM, and 1.88 on uncond. CIFAR-10 with PFGM++.

To further validate that Restart can be applied in low NFE regime , we show that one can employ faster ODE solvers such as the DPM-solver [16] to further accelerate Restart. Fig. 4 shows that the Restart consistently outperforms the DPM-Solver with an NFE ranging from 16 to 36. This demonstrates Restart's capability to excel over ODE samplers, even in the small NFE regime. It also suggests that Restart can consistently improve other ODE samplers, not limited to the DDIM, Heun. Surprisingly, when paired with the DPM-Solver, Restart achieves an FID score of 2 . 11 on VP setting when NFE is 30, which is significantly

NFE FID

EDM-VP

[13]

ODE (Heun)

63

1.97

35

1.97

Improved SDE

63

2.27

35

2.45

43

1.90

PFGM++

[28]

ODE (Heun)

63

1.91

35

1.91

43

1.88

0

20

40

60

Number of Restart iterations

K

1.8

2.0

2.2

2.4

2.6

2.8

3.0

3.2

FID score (50K)

EDM (ODE)

VP (ODE)

lower than any previous numbers (even lower than the SDE sampler with an NFE greater than 1000 in [23]), and make VP model on par with the performance with more advanced models (such as EDM).

Theorem 4 shows that each Restart iteration reduces the contracted errors while increasing the additional sampling errors in the backward process. In Fig. 5, we explore the choice of the number of Restart iterations K on CIFAR-10. We find that FID score initially improves and later worsens with increasing iterations K , with a smaller turning point for stronger EDM model. This supports the theoretical analysis that sampling errors will eventually outweigh the contraction benefits as K increases, and EDM only permits fewer Restart iterations due to smaller accumulated errors. It also suggests that, as a rule of thumb, we should apply greater Restart strength ( e.g . , larger K ) for weaker or smaller architectures and vice versa.

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

0.290

0.295

0.300

0.305

0.310

0.315

0.320

CLIP score (ViT-g/14)

14

16

18

20

22

FID score (5K)

DDIM (Steps=50)

DDIM (Steps=100)

DDPM (Steps=100)

DDPM (Steps=200)

Restart (Steps=66)

(a) FID versus CLIP score

5.15

5.20

5.25

5.30

5.35

5.40

Aesthetic score

14

16

18

20

22

FID score (5K)

DDIM (Steps=50)

DDIM (Steps=100)

DDPM (Steps=100)

DDPM (Steps=200)

Restart (Steps=66)

(b) FID versus Aesthetic score

We further apply Restart to the text-to-image Stable Diffusion v1.5 2 pre-trained on LAION-5B [21] at a resolution of 512 × 512 . Weemploy the commonly used classifier-free guidance [8, 20] for sampling, wherein each sampling step entails two function evaluations - the conditional and unconditional predictions. Following [18, 20], we use the COCO [15] validation set for evaluation. We assess text-image alignment using the CLIP score [6] with the open-sourced ViT-g/14 [11], and measure diversity via the FID score. We also evaluate visual quality through the Aesthetic score, as rated by the LAION-Aesthetics Predictor V2 [24]. Following [17], we compute all evaluation metrics using 5K captions randomly sampled from the validation set and plot the trade-off curves between CLIP/Aesthetic scores and FID score, with the classifier-free guidance weight w in { 2 , 3 , 5 , 8 } .

We compare with commonly used ODE sampler DDIM [22] and the stochastic sampler DDPM [9]. For Restart, we adopt the DDIM solver with 30 steps in the main backward process, and Heun in the Restart backward process, as we empirically find that Heun performs better than DDIM in the Restart. In addition, we select different sets of the hyperparameters for each guidance weight. For instance, when w = 8 , we use [ t min , t max ]=[0 . 1 , 2] , K =2 and 10 steps in Restart backward process. We defer the detailed Restart configuration to Appendix C.2, and the results of Heun to Appendix D.1.

As illustrated in Fig. 6(a) and Fig. 6(b), Restart achieves better FID scores in most cases, given the same CLIP/Aesthetic scores, using only 132 function evaluations ( i.e . , 66 sampling steps). Remarkably, Restart achieves substantially lower FID scores than other samplers when CLIP/Aesthetic scores are high ( i.e . , with larger w values). Conversely, Restart generally obtains a better text-image alignment/visual quality given the same FID. We also observe that DDPM generally obtains comparable performance with Restart in FID score when CLIP/Aesthetic scores are low, with Restart being more time-efficient. These findings suggest that Restart balances diversity (FID score) against text-image alignment (CLIP score) or visual quality (Aesthetic score) more effectively than previous samplers.

In Fig. 7, we visualize the images generated by Restart, DDIM and DDPM with w = 8 . Compared to DDIM, the Restart generates images with superior details ( e.g . , the rendition of duck legs by DDIM is less accurate) and visual quality. Compared to DDPM, Restart yields more photo-realistic images ( e.g . , the astronaut). We provide extended of text-to-image generated samples in Appendix E.

(b) DDIM (Steps=100)

(c) DDPM (Steps=100)

6 Conclusion and future direction

In this paper, we introduce the Restart sampling for generative processes involving differential equations, such as diffusion models and PFGMs. By interweaving a forward process that adds a significant amount of noise with a corresponding backward ODE, Restart harnesses and even enhances the individual advantages of both ODE and SDE. Theoretically, Restart provides greater contraction effects of stochasticity while maintaining ODE-level discretization error. Empirically, Restart achieves a superior balance between quality and time, and improves the text-image alignment/visual quality and diversity trade-off in the text-to-image Stable Diffusion models.

A current limitation of the Restart algorithm is the absence of a principled way for hyperparameters selection, including the number of iterations K and the time interval [ t min , t max ] . At present, we adjust these parameters based on the heuristic that weaker/smaller models, or more challenging tasks, necessitate a stronger Restart strength. In the future direction, we anticipate developing a more principled approach to automating the selection of optimal hyperparameters for Restart based on the error analysis of models, in order to fully unleash the potential of the Restart framework.

References

Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications , 12(3):313-326, 1982.

Uri M. Ascher and Linda R. Petzold. Computer methods for ordinary differential equations and differential-algebraic equations. In SIAM , 1998.

Andrei N Borodin and Paavo Salminen. Handbook of Brownian motion-facts and formulae . Springer Science & Business Media, 2015.

Arnak S Dalalyan and Avetik Karagulyan. User-friendly guarantees for the langevin monte carlo with inaccurate gradient. Stochastic Processes and their Applications , 129(12):5278-5311, 2019.

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, K. Li, and Li Fei-Fei. Imagenet: A largescale hierarchical image database. 2009 IEEE Conference on Computer Vision and Pattern Recognition , pages 248-255, 2009.

Jack Hessel, Ari Holtzman, Maxwell Forbes, Ronan Joseph Le Bras, and Yejin Choi. Clipscore: A reference-free evaluation metric for image captioning. In Conference on Empirical Methods in Natural Language Processing , 2021.

Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. In NIPS , 2017.

Jonathan Ho. Classifier-free diffusion guidance. ArXiv , abs/2207.12598, 2022.

Jonathan Ho, Ajay Jain, and P. Abbeel. Denoising diffusion probabilistic models. ArXiv , abs/2006.11239, 2020.

Chutian Huang, Zijing Liu, Shengyuan Bai, Linwei Zhang, Chencheng Xu, Zhe Wang, Yang Xiang, and Yuanpeng Xiong. Pf-abgen: A reliable and efficient antibody generator via poisson flow. Machine Learning for Drug Discovery Workshop, International Conference on Learning Representations , 2023.

Gabriel Ilharco, Mitchell Wortsman, Ross Wightman, Cade Gordon, Nicholas Carlini, Rohan Taori, Achal Dave, Vaishaal Shankar, Hongseok Namkoong, John Miller, Hannaneh Hajishirzi, Ali Farhadi, and Ludwig Schmidt. Openclip. Zenodo , 2021.

Alexia Jolicoeur-Martineau, Ke Li, Remi Piche-Taillefer, Tal Kachman, and Ioannis Mitliagkas. Gotta go fast when generating data with score-based models. ArXiv , abs/2105.14080, 2021.

Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. ArXiv , abs/2206.00364, 2022.

Alex Krizhevsky. Learning multiple layers of features from tiny images. Citeseer , 2009.

Tsung-Yi Lin, Michael Maire, Serge J. Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C. Lawrence Zitnick. Microsoft coco: Common objects in context. In European Conference on Computer Vision , 2014.

Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. arXiv preprint arXiv:2206.00927 , 2022.

Chenlin Meng, Ruiqi Gao, Diederik P. Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models. ArXiv , abs/2210.03142, 2022.

Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen. Glide: Towards photorealistic image generation and editing with text-guided diffusion models. In International Conference on Machine Learning , 2021.

Robin Rombach, A. Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 10674-10685, 2021.

Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L. Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, Seyedeh Sara Mahdavi, Raphael Gontijo Lopes, Tim Salimans, Jonathan Ho, David J. Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffusion models with deep language understanding. ArXiv , abs/2205.11487, 2022.

Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, Patrick Schramowski, Srivatsa Kundurthy, Katherine Crowson, Ludwig Schmidt, Robert Kaczmarczyk, and Jenia Jitsev. Laion-5b: An open large-scale dataset for training next generation image-text models. ArXiv , abs/2210.08402, 2022.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. ArXiv , abs/2010.02502, 2020.

Yang Song, Jascha Narain Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. ArXiv , abs/2011.13456, 2020.

LAION-AI Team. Laion-aesthetics predictor v2. https://github.com/ christophschuhmann/improved-aesthetic-predictor , 2022.

Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Computation , 23:1661-1674, 2011.

Joseph L. Watson, David Juergens, Nathaniel R. Bennett, Brian L. Trippe, Jason Yim, Helen E. Eisenach, Woody Ahern, Andrew J. Borst, Robert J. Ragotte, Lukas F. Milles, Basile I. M. Wicky, Nikita Hanikel, Samuel J. Pellock, Alexis Courbet, William Sheffler, Jue Wang, Preetham Venkatesh, Isaac Sappington, Susana Vázquez Torres, Anna Lauko, Valentin De Bortoli, Emile Mathieu, Regina Barzilay, T. Jaakkola, Frank DiMaio, Minkyung Baek, and David Baker. Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models. bioRxiv , 2022.

Yilun Xu, Ziming Liu, Max Tegmark, and T. Jaakkola. Poisson flow generative models. ArXiv , abs/2209.11178, 2022.

Yilun Xu, Ziming Liu, Yonglong Tian, Shangyuan Tong, Max Tegmark, and T. Jaakkola. Pfgm++: Unlocking the potential of physics-inspired generative models. ArXiv , abs/2302.04265, 2023.

Appendix

A Proofs of Main Theoretical Results

In this section, we provide proofs of our main results. We define below some crucial notations which we will use throughout. We use ODE ( . . . ) to denote the backwards ODE under exact score ∇ log p t ( x ) . More specifically, given any x ∈ R d and s > r > 0 , let x t denote the solution to the following ODE:

ODE ( x, s → r ) is defined as "the value of x r when initialized at x s = x ". It will also be useful to consider a "time-discretized ODE with drift ts θ ( x, t ) ": let δ denote the discretization step size and let k denote any integer. Let δ denote a step size, let x t denote the solution to

where for any t , k is the unique integer such that t ∈ (( k -1) δ, kδ ] . We verify that the dynamics of Eq. (6) is equivalent to the following discrete-time dynamics for t = kδ, k ∈ Z :

We similarly denote the value of x r when initialized at x s = x as ODE θ ( x, s → r ) . Analogously, we let SDE ( x, s → r ) and SDE θ ( x, s → r ) denote solutions to

respectively. Finally, we will define the Restart θ process as follows:

where ε i t min → t max ∼ N ( 0 , ( t 2 max -t 2 min ) I ) . We use Restart θ ( x, K ) to denote x K t min in the above processes, initialized at x 0 t min = x . In various theorems, we will refer to a function Q ( r ) : R + → [0 , 1 / 2) , defined as the Gaussian tail probability Q ( r ) = Pr ( a ≥ r ) for a ∼ N (0 , 1) .

A.1 Main Result

Theorem 3. [Formal version of Theorem 1] Let t max be the initial noise level. Let the initial random variables x t max = y t max , and

Let p t denote the true population distribution at noise level t . Let p ODE θ t , p SDE θ t denote the distributions for x t , y t respectively. Assume that for all x, y, s, t , s θ ( x, t ) satisfies ∥ ts θ ( x, t ) -ts θ ( x, s ) ∥ ≤ L 0 | s -t | , ∥ ts θ ( x, t ) ∥ ≤ L 1 , ∥ ts θ ( x, t ) -ts θ ( y, t ) ∥ ≤ L 2 ∥ x -y ∥ , and the approximation error ∥ ts θ ( x, t ) -t ∇ log p t ( x ) ∥ ≤ ϵ approx . Assume in addition that ∀ t ∈ [ t min , t max ] , ∥ x t ∥ < B/ 2 for any x t in the support of p t , p ODE θ t or p SDE θ t , and K ≤ C L 2 ( t max -t min ) for some universal constant C . Then

Proof. Let us define x t max ∼ p t max , and let x t min = ODE ( x t max , t max → t min ) . We verify that x t min has density p t min . Let us also define ˆ x t min = ODE θ ( x t max , t max → t min ) . We would like to bound the Wasserstein distance between ¯ x t min and x t min ( i.e . , p ODE θ t min and p t min ), by the following triangular inequality:

By Lemma 2, we know that

where we use the fact that ∥ ˆ x t max -x t max ∥ = 0 . Thus we immediately have

On the other hand,

where the last equality is due to the data-processing inequality. Combining Eq. (11) , Eq. (12) and the triangular inequality Eq. (10), we arrive at the upper bound for ODE (Eq. (8)). The upper bound for SDE (Eq. (9)) shares a similar proof approach. First, let y t max ∼ p t max . Let ˆ y t min = SDE θ ( y t max , t max → t min ) . By Lemma 5,

On the other hand, by Lemma 4,

The SDE triangular upper bound on W 1 (¯ y t min , y t min ) follows by multiplying the first inequality by B (to bound W 1 (¯ y t min , ˆ y t min ) ) and then adding the second inequality (to bound W 1 ( y t min , ˆ y t min ) ). Notice that by definition, TV ( ˆ y t max , y t max ) = TV ( y t max , y t max ) . Finally, because of the assumption that K ≤ C L 2 ( t max -t min ) for some universal constant, we summarize the second term in the Eq. (8) and Eq. (9) into the big O in the informal version Theorem 1.

Theorem 4. [Formal version of Theorem 2] Consider the same setting as Theorem 3. Let p Restart θ ,i t min denote the distributions after i th Restart iteration, i.e . , the distribution of x i t min = Restart θ ( x 0 t min , i ) . Given initial x 0 t max ∼ p Restart , 0 t max , let x 0 t min = ODE θ ( x 0 t max , t max → t min ) . Then

where λ = 2 Q ( B 2 √ t 2 max -t 2 min ) .

Proof. Let x 0 t max ∼ p t max . Let x K t min = Restart ( x 0 t min , K ) . We verify that x K t min has density p t min . Let us also define ˆ x 0 t min = ODE θ ( x 0 t max , t max → t min ) and ˆ x K t min = Restart θ (ˆ x 0 t min , K ) . By Lemma 1,

The second inequality holds by data processing inequality. The above can be used to bound the 1-Wasserstein distance as follows:

On the other hand, using Lemma 3,

We arrive at the result by combining the two bounds above (Eq. (14), Eq. (15)) with the following triangular inequality,

A.2 Mixing under Restart with exact ODE

Lemma 1. Consider the same setup as Theorem 4. Consider the Restart θ process defined in equation 7. Let

Let p Restart θ ( i ) t and q Restart θ ( i ) t denote the densities of x i t and y i t respectively. Then

Proof. Conditioned on x i t min , y i t min , let x i +1 t max = x i t min + √ t 2 max -t 2 min ξ x i and y i +1 t max = y i t min + √ t 2 max -t 2 min ξ y i . We now define a coupling between x i +1 t min and y i +1 t min by specifying the joint distribution over ξ x i and ξ y i .

̸

If x i t min = y i t min , let ξ x i = ξ y i , so that x i +1 t min = y i +1 t min . On the other hand, if x i t min = y i t min , let x i +1 t max and y i +1 t max be coupled as described in the proof of Lemma 7, with x ′ = x i +1 t max , y ′ = y i +1 t max , σ = √ t 2 max -t 2 min . Under this coupling, we verify that,

̸

Applying the above recursively,

̸

̸

The conclusion follows by noticing that TV ( p Restart θ ( K ) t min , q Restart θ ( K ) t min ) ≤ Pr ( x K t min = y t K min ) = E [ ✶ { x K t min = y t K min }] , and by selecting the initial coupling so that Pr ( x 0 t min = y 0 t min ) = TV ( p Restart θ (0) t min , q Restart θ (0) t min ) .

̸

̸

̸

̸

̸

̸

A.3 W 1 discretization bound

Lemma 2 (Discretization bound for ODE) . Let x t min = ODE ( x t max , t max → t min ) and let x t min = ODE θ ( x t max , t max → t min ) . Assume that for all x, y, s, t , s θ ( x, t ) satisfies ∥ ts θ ( x, t ) -ts θ ( x, s ) ∥ ≤ L 0 | s -t | , ∥ ts θ ( x, t ) ∥ ≤ L 1 and ∥ ts θ ( x, t ) -ts θ ( y, t ) ∥ ≤ L 2 ∥ x -y ∥ . Then

Proof. Consider some fixed arbitrary k , and recall that δ is the step size. Recall that by definition of ODE and ODE θ , for t ∈ (( k -1) δ, kδ ] ,

For t ∈ [ t min , t max ] , let us define a time-reversed process x ← t := x -t . Let v ( x, t ) := ∇ log p -t ( x ) . Then for t ∈ [ -t max , -t min ]

Similarly, define x ← t := x -t and v ( x, t ) := s θ ( x, -t ) . It follows that

where k is the unique (negative) integer satisfying t ∈ [ kδ, ( k +1) δ ) . Following these definitions,

Applying Gronwall's Lemma over the interval t ∈ [ -t max , -t min ] ,

Lemma 3. Given initial x 0 t max , let x 0 t min = ODE ( x 0 t max , t max → t min ) , and let ˆ x 0 t min = ODE θ ( x 0 t max , t max → t min ) . We further denote the variables after K Restart iterations as x K t min = Restart ( x 0 t min , K ) and ˆ x K t min = Restart θ (ˆ x 0 t min , K ) , with true field and learned field respectively. Then there exists a coupling between x K t min and ˆ x K t min such that

Proof. We will couple x i t min and ˆ x i t min by using the same noise ε i t min → t max in the Restart forward process for i = 0 . . . K -1 (see Eq. (7)). For any i , let us also define y i,j t min := Restart θ ( x i t min , j -i ) , and this process uses the same noise ε i t min → t max as previous ones. From this definition, y K,K t min = x K t min . We can thus bound

Using the assumption that ts θ ( · , t ) is L 2 Lipschitz,

where the last equality is because we add the same additive Gaussian noise ε i t min → t max to y 0 ,i t min and ˆ x i t min in the Restart forward process. Applying the above recursively, we get

where the last line follows by Lemma 2 when setting x t max = ¯ x t max . We will now bound ∥ ∥ ∥ y i,K t min -y i +1 ,K t min ∥ ∥ ∥ for some i ≤ K . It follows from definition that

For the remaining steps from i +2 . . . K , both y i, · and y i +1 , · evolve with ODE θ in each step. Again using the assumption that ts θ ( · , t ) is L 2 Lipschitz,

Summing the above for i = 0 ...K -1 , and combining with Eq. (16) and Eq. (17) gives

Lemma 4. Consider the same setup as Theorem 3. Let x t min = SDE ( x t max , t max → t min ) and let x t min = SDE ( x t max , t max → t min ) . Then there exists a coupling between x t and x t such that

Proof. Consider some fixed arbitrary k , and recall that δ is the stepsize. By definition of SDE and SDE θ , for t ∈ (( k -1) δ, kδ ] ,

Let us define a coupling between x t and x t by identifying their respective Brownian motions. It will be convenient to define the time-reversed processes x ← t := x -t , and x ← t := x -t , along with v ( x, t ) := ∇ log p -t ( x ) and v ( x, t ) := s θ ( x, -t ) . Then there exists a Brownian motion B ← t , such that for t ∈ [ -t max , -t min ] ,

where k is the unique negative integer such that t ∈ [ kδ, ( k +1) δ ) . Thus

By Gronwall's Lemma,

A.4 Mixing Bounds

Lemma 5. Consider the same setup as Theorem 3. Assume that δ ≤ t min. Let

Then there exists a coupling between x s and y s such that

̸

Proof. We will construct a coupling between x t and y t . First, let ( x t max , y t max ) be sampled from the optimal TV coupling, i.e . , Pr ( x t max = y t max ) = TV ( x t max , y t max ) . Recall that by definition of SDE θ , for t ∈ (( k -1) δ, kδ ] ,

Let us define a time-rescaled version of x t : x t := x t 2 . We verify that

where k is the unique integer satisfying t ∈ [(( k -1) δ ) 2 , k 2 δ 2 ) . Next, we define the time-reversed process x ← t := x -t , and let v ( x, t ) := s θ ( x, -t ) . We verify that there exists a Brownian motion B x t such that, for t ∈ [ -t 2 max , -t 2 min ] ,

where v x t = s θ ( x ← -( kδ ) 2 , -kδ ) , where k is the unique positive integer satisfying -t ∈ ((( k -1) δ ) 2 , ( kδ ) 2 ] . Let dy ← t = v y t dt + dB y t , be defined analogously. For any positive integer k and for any t ∈ [ -( kδ ) 2 , -(( k -1) δ ) 2 ) , let us define

Let γ t := z t ∥ z t ∥ . We will now define a coupling between dB x t and dB y t as

where ✶ {} denotes the indicator function, i.e. ✶ { t ≤ τ } = 1 if t ≤ τ , and τ is a stopping time given by the first hitting time of z t = 0 . Let r t := ∥ z t ∥ . Consider some t ∈ ( -i 2 δ 2 , -( i -1) 2 δ 2 ) , and Let j := t max δ (assume w.l.o.g that this is an integer), then

where dB 1 s = ⟨ γ t , dB x s -dB y s ⟩ is a 1-dimensional Brownian motion. We also verify that

where the third relation is by adding and subtracting B x -( j -1) 2 δ 2 -B x t and using triangle inequality. The fourth relation is by noticing that x ← -t 2 max + (2 j -1) δ 2 v x -t 2 max + ( B x -( j -1) 2 δ 2 -B x -t 2 max ) = x ← -( j -1) 2 δ 2 and that y ← -t 2 max +(2 j -1) δ 2 v y -t 2 max + ( B x -( j -1) 2 δ 2 -B x t + B y t -B y -t 2 max ) d = y ← -( j -1) 2 δ 2 , and then using our assumption in the theorem statement that all processes are supported on a ball of radius B/ 2 .

We now define a process s t defined by ds t = 2 L 1 /t min dt +2 dB 1 t , initialized at s -t 2 max = B ≥ r -t 2 max . We can verify that, up to time τ , r t ≤ s t with probability 1. Let τ ′ denote the first-hitting time of s t to 0 , then τ ≤ τ ′ with probability 1. Thus

where we apply Lemma 6. The proof follows by noticing that, if τ ≤ -t 2 min , then x t min = y t min . This is because if τ ∈ [ -k 2 δ 2 , -( k -1) 2 δ 2 ] , then x ← -( k -1) 2 δ 2 = y ← -( k -1) 2 δ 2 , and thus x ← t = y ← t for all t ≥ -( k -1) 2 δ 2 , in particular, at t = -t 2 min .

Lemma 6. Consider the stochastic process

Assume that r 0 ≤ B/ 2 . Let τ denote the hitting time for r t = 0 . Then for any T ∈ R + ,

where Q is the tail probability of a standard Gaussian defined in Definition 1.

Proof. We will use he following facts in our proof:

Let dr t = dB 1 t + cdt , with r 0 = a . The density of the hitting time τ is given by

(see e.g. [3]). From item 2 above,

̸

0

Therefore,

Pr

(

τ

T

A.5 TV Overlap

Definition 1. Let x be sampled from standard normal distribution N (0 , 1) . We define the Gaussian tail probability Q ( a ) := Pr ( x ≥ a ) .

Lemma 7. We verify that for any two random vectors ξ x ∼ N ( 0 , σ 2 I ) and ξ y ∼ N ( 0 , σ 2 I ) , each belonging to R d , the total variation distance between x ′ = x + ξ x and y ′ = y + ξ y is given by

where r = ∥ x -y ∥ 2 σ , and Q ( r ) = Pr ( ξ ≥ r ) , when ξ ∼ N (0 , 1) .

Proof. Let γ := x -y ∥ x -y ∥ . We decompose x ′ , y ′ into the subspace/orthogonal space defined by γ :

We verify the independence ξ ⊥ x ⊥ ⊥ ξ ∥ x and ξ ⊥ y ⊥ ⊥ ξ ∥ y as they are orthogonal decompositions of the standard Gaussian. We will define a coupling between x ′ and y ′ by setting ξ ⊥ x = ξ ⊥ y . Under this coupling, we verify that

Therefore, x ′ = y ′ if and only if x ∥ + ξ ∥ x = y ∥ + ξ ∥ y . Next, we draw ( a, b ) from the optimal coupling between N (0 , 1) and N ( ∥ x -y ∥ σ , 1) . We verify that x ∥ + ξ ∥ x and y ∥ + ξ ∥ y both lie in the span of γ . Thus it suffices to compare 〈 γ, x ∥ + ξ ∥ x 〉 and 〈 γ, y ∥ + ξ ∥ y 〉 . We verify that 〈 γ, x ∥ + ξ ∥ x 〉 =

where we define

) =

∫

T

f

(

In the case of a general c = 0 , we can bound ( a + ct ) 2 2 t = a 2 2 t + ac + c 2 t 2 . Consequently,

∫

T

a, c, t

)

dt

f

(

a,

0

, t

)

dte

-

c

= 2

0

Q

(

2

B

T

√

)

·

e

-

ac

-

c

2

T

2

.

Thus TV ( x ′ , y ′ ) = TV ( σa, σb ) = 1 -2 Q ( ∥ x -y ∥ 2 σ ) . The last inequality follows from

B More on Restart Algorithm

B.1 EDMDiscretization Scheme

[13] proposes a discretization scheme for ODE given the starting t max and end time t min. Denote the number of steps as N , then the EDM discretization scheme is:

with t 0 = t max and t N -1 = t min. ρ is a hyperparameter that determines the extent to which steps near t min are shortened. We adopt the value ρ = 7 suggested by [13] in all of our experiments. We apply the EDM scheme to creates a time discretization in each Restart interval [ t max , t min ] in the Restart backward process, as well as the main backward process between [0 , T ] (by additionally setting t min = 0 . 002 and t N = 0 as in [13]). It is important to note that t min should be included within the list of time steps in the main backward process to seamlessly incorporate the Restart interval into the main backward process. We summarize the scheme as a function in Algorithm 1.

Algorithm 1 EDM_Scheme( t min , t max , N, ρ = 7)

B.2 Restart Algorithm

We present the pseudocode for the Restart algorithm in Algorithm 2. In this pseudocode, we describe a more general case that applies l -level Restarting strategy. For each Restart segment, the include the number of steps in the Restart backward process N Restart, the Restart interval [ t min , t max ] and the number of Restart iteration K . We further denote the number of steps in the main backward process as N main. We use the EDM discretization scheme (Algorithm 1) to construct time steps for the main backward process ( t 0 = T, t N main = 0 ) as well as the Restart backward process, when given the starting/end time and the number of steps.

Although Heun's 2 nd order method [2] (Algorithm 3) is the default ODE solver in the pseudocode, it can be substituted with other ODE solvers, such as Euler's method or the DPM solver [16].

The provided pseudocode in Algorithm 2 is tailored specifically for diffusion models [13]. To adapt Restart for other generative models like PFGM++ [28], we only need to modify the Gaussian perturbation kernel in the Restart forward process (line 10 in Algorithm 2) to the one used in PFGM++.

C Experimental Details

In this section, we discuss the configurations for different samplers in details. All the experiments are conducted on eight NVIDIA A100 GPUs.

Algorithm 2 Restart sampling

Algorithm 3 OneStep_Heun( s , x , t , t

̸

C.1 Configurations for Baselines

We select Vanilla SDE [23], Improved SDE [13], Gonna Go Fast [12] as SDE baselines and the Heun 's 2 nd order method [2] (Alg 3) as ODE baseline on standard benchmarks CIFAR-10 and ImageNet 64 × 64 . We choose DDIM [22], Heun 's 2 nd order method, and DDPM [9] for comparison on Stable Diffusion model.

Vanilla SDE denotes the reverse-time SDE sampler in [23]. For Improved SDE, we use the recommended dataset-specific hyperparameters ( e.g . , S max , S min , S churn) in Table 5 of the EDM paper [13]. They obtained these hyperparameters by grid search. Gonna Go Fast [12] applied an adaptive step size technique based on Vanilla SDE and we directly report the FID scores listed in [12] for Gonna Go Fast on CIFAR-10 (VP). For fair comparison, we use the EDM discretization scheme [13] for Vanilla SDE, Improved SDE, Heun as well as Restart.

Weborrow the hyperparameters such as discretization scheme or initial noise scale on Stable Diffusion models in the diffuser 3 code repository. We directly use the DDIM and DDPM samplers implemented in the repo. We apply the same set of hyperparameters to Heun and Restart.

C.2 Configurations for Restart

We report the configurations for Restart for different models and NFE on standard benchmarks CIFAR-10 and ImageNet 64 × 64 . The hyperparameters of Restart include the number of steps in the main backward process N main, the number of steps in the Restart backward process N Restart , the Restart interval [ t min , t max ] and the number of Restart iteration K . In Table 3 (CIFAR-10, VP)

we provide the quintuplet ( N main , N Restart , t min , t max , K ) for each experiment. Since we apply the multi-level Restart strategy for ImageNet 64 × 64 , we provide N main as well as a list of quadruple { ( N Restart ,i , K i , t min ,i , t max ,i ) } l i =1 ( l is the number of Restart interval depending on experiments) in Table 5. In order to integrate the Restart time interval to the main backward process, we round t min ,i to its nearest neighbor in the time steps of main backward process, as shown in line 2 of Algorithm 2. We apply Heun method for both main/backward process. The formula for NFE calculation is

[13], we inflate the additive noise in the Restart forward process by multiplying S noise = 1 . 003 on ImageNet 64 × 64 , to counteract the over-denoising tendency of neural networks. We also observe that setting γ = 0 . 05 in Algorithm 2 of EDM [13] would sligtly boost the Restart performance on ImageNet 64 × 64 when t ∈ [0 . 01 , 1] .

We further include the configurations for Restart on Stable Diffusion models in Table 10, with a varying guidance weight w . Similar to ImageNet 64 × 64 , we use multi-level Restart with a fixed number of steps N main = 30 in the main backward process. We utilize the Euler method for the main backward process and the Heun method for the Restart backward process, as our empirical observations indicate that the Heun method doesn't yield significant improvements over the Euler method, yet necessitates double the steps. The number of steps equals to N main + ∑ l i =1 K i · (2 · ( N Restart ,i -1)) in this case. We set the total number of steps to 66 , including main backward process and Restart backward process.

Given the prohibitively large search space for each Restart quadruple, a comprehensive enumeration of all possibilities is impractical due to computational limitations. Instead, we adjust the configuration manually, guided by the heuristic that weaker/smaller models or more challenging tasks necessitate a stronger Restart strength (e.g., larger K , wider Restart interval, etc). On average, we select the best configuration from 5 sets for each experiment; these few trials have empirically outperformed previous SDE/ODE samplers. We believe that developing a systematic approach for determining Restart configurations could be of significant value in the future.

C.3 Pre-trained Models

For CIFAR-10 dataset, we use the pre-trained VP and EDM models from the EDM repository 4 , and PFGM++ ( D = 2048 ) model from the PFGM++ repository 5 . For ImageNet 64 × 64 , we borrow the pre-trained EDM model from EDM repository as well.

C.4 Classifier-free Guidance

We follow the convention in [20], where each step in classifier-free guidance is as follows:

where c is the conditions, and s θ ( x, c, t ) /s θ ( x, t ) is the conditional/unconditional models, sharing parameters. Increasing w would strength the effect of guidance, usually leading to a better text-image alignment [20].

C.5 More on the Synthetic Experiment

C.5.1 Discrete Dataset

We generate the underlying discrete dataset S with | S | = 2000 as follows. Firstly, we sample 2000 points, denoted as S 1 , from a mixture of two Gaussians in R 4 . Next, we project these points onto R 20 . To ensure a variance of 1 on each dimension, we scale the coordinates accordingly. This setup aims to simulate data points that primarily reside on a lower-dimensional manifold with multiple modes.

The specific details are as follows: S 1 ∼ 0 . 3 N ( a, s 2 I ) + 0 . 7( -a, s 2 I ) , where a = (3 , 3 , 3 , 3) ⊂ R 4 and s = 1 . Then, we randomly select a projection matrix P ∈ R 20 × 4 , where each entry is drawn from N (0 , 1) , and compute S 2 = PS 1 . Finally, we scale each coordinate by a constant factor to ensure a variance of 1.

0.65

0.70

0.75

0.80

0.85

0.90

Additional Sampling error

0.65

0.70

0.75

0.80

0.85

0.90

Contracted error

ODE

Improved SDE

(a)

0.65

0.70

0.75

0.80

0.85

0.90

0.95

1.00

Additional Sampling error

0.70

0.75

0.80

0.85

0.90

0.95

1.00

Total Error

ODE

Improved SDE

(b)

20

40

80

160

320

NFE

0.750

0.775

0.800

0.825

0.850

0.875

Total Error

ODE

Improved SDE

(c)

C.5.2 Model Architecture

We employ a common MLP architecture with a latent size of 64 to learn the score function. The training method is adapted from [13], which includes the preconditioning technique and denoising score-matching objective [25].

C.5.3 Varying Hyperparameters

To achieve the best trade-off between contracted error and additional sampling error, and optimize the NFE versus FID (Fréchet Inception Distance) performance, we explore various hyperparameters. [13] shows that the Vanilla SDE can be endowed with additional flexibility by varying the coefficient β ( t ) (Eq.(6) in [13]). Hence, regarding SDE, we consider NFE values from { 20 , 40 , 80 , 160 , 320 } , and multiply the original β ( t ) = ˙ σ ( t ) /σ ( t ) [13] with values from { 0 , 0 . 25 , 0 . 5 , 1 , 1 . 5 , 2 , 4 , 8 } . It is important to note that larger NFE values do not lead to further performance improvements. For restarts, we tried the following two settings: first we set the number of steps in Restart backward process to 40 and vary the number of Restart iterations K in the range { 0 , 5 , 10 , 15 , 20 , 25 , 30 , 35 } . We also conduct a grid search with the number of Restart iterations K ranging from 5 to 25 and the number of steps in Restart backward process varying from 2 to 7. For ODE, we experiment with the number of steps set to { 20 , 40 , 80 , 160 , 320 , 640 } .

Additionally, we conduct an experiment for Improved SDE in EDM. We try different values of S churn in the range of { 0 , 1 , 2 , 4 , 8 , 16 , 32 , 48 , 64 } . We also perform a grid search where the number of steps ranged from 20 to 320 and S churn takes values of [0 . 2 × steps , 0 . 5 × steps , 20 , 60] . The plot combines the results from SDE and is displayed in Figure 8.

To mitigate the impact of randomness, we collect the data by averaging the results from five runs with the same hyperparameters. To compute the Wasserstein distance between two discrete distributions, we use minimum weight matching.

C.5.4 Plotting the Pareto frontier

We generate the Pareto frontier plots as follows. For the additional sampling error versus contracted error plot, we first sort all the data points based on their additional sampling error and then connect the data points that represent prefix minimums of the contracted error. Similarly, for the NFE versus FID plot, we sort the data points based on their NFE values and connect the points where the FID is a prefix minimum.

D Extra Experimental Results

D.1 Numerical Results

In this section, we provide the corresponding numerical reuslts of Fig. 3(a) and Fig. 3(b), in Table 2, 3 (CIFAR-10 VP, EDM, PFGM++) and Table 4, 5 (ImageNet 64 × 64 EDM), respectively. We also include the performance of Vanilla SDE in those tables. For the evaluation, we compute the Fréchet distance between 50000 generated samples and the pre-computed statistics of CIFAR-10 and ImageNet 64 × 64 . We follow the evaluation protocol in EDM [13] that calculates each FID scores three times with different seeds and report the minimum.

We also provide the numerical results on the Stable Diffusion model [19], with a classifier guidance weight w = 2 , 3 , 5 , 8 in Table 6, 7, 8, 9. As in [17], we report the zero-shot FID score on 5K random prompts sampled from the COCO validation set. We evaluate CLIP score [6] with the open-sourced ViT-g/14 [11], Aesthetic score by the more recent LAION-Aesthetics Predictor V2 6 . We average the CLIP and Aesthetic scores over 5K generated samples. The number of function evaluations is two times the sampling steps in Stable Diffusion model, since each sampling step involves the evaluation of the conditional and unconditional model.

D.2 Sensitivity Analysis of Hyper-parameters

We also investigate the impact of varying t min when t max = t min + 0 . 3 , and the length the restart interval when t min = 0 . 06 . Fig. 10(a) reveals that FID scores achieve a minimum at a t min close to 0

0.285

0.290

0.295

0.300

0.305

0.310

0.315

0.320

CLIP score (ViT-g/14)

14

16

18

20

22

FID score (5K)

Restart (Steps=66)

DDIM (Steps=50)

DDIM (Steps=100)

Heun (Steps=51)

Heun (Steps=101)

DDPM (Steps=100)

DDPM (Steps=200)

(a) FID versus CLIP score

5.15

5.20

5.25

5.30

5.35

5.40

Aesthetic score

14

16

18

20

22

FID score (5K)

Restart (Steps=66)

DDIM (Steps=50)

DDIM (Steps=100)

Heun (Steps=51)

Heun (Steps=101)

DDPM (Steps=100)

DDPM (Steps=200)

(b) FID versus Aesthetic score

on VP, indicating higher accumulated errors at the end of sampling and poor neural estimations at small t . Note that the Restart interval 0 . 3 is about twice the length of the one in Table 1 and Restart does not outperform the ODE baseline on EDM. This suggests that, as a rule of thumb, we should apply greater Restart strength ( e.g . , larger K , t max -t min) for weaker or smaller architectures and vice versa.

In theory, a longer interval enhances contraction but may add more additional sampling errors. Again, the balance between these factors results in a V-shaped trend in our plots (Fig. 10(b)). In practice, selecting t max close to the dataset's radius usually ensures effective mixing when t min is small.

4

2

0

2

log(

t

min

)

2.0

2.2

2.4

2.6

2.8

3.0

3.2

FID score (50K)

EDM (ODE)

VP (ODE)

(a)

(b)

intervallengthvsFID

2.8

2.6

vprestart

vpODE

edmrestart

edmODE

2.2

2.0

0.0

0.1

0.2

0.3

0.4

0.5

0.6

intervallength

E Extended Generated Images

In this section, we provide extended generated images by Restart, DDIM, Heun and DDPM on text-to-image Stable Diffusion v1.5 model [19]. We showcase the samples of four sets of text prompts in Fig. 11, Fig. 12, Fig. 13, Fig. 14, with a classifier-guidance weight w = 8 .

F Heun's method is DPM-Solver-2 (with r 2 = 1 )

The first order ODE in DPM-Solver [16] (DPM-Solver-1) is in the form of:

The first order ODE in EDM is in the form of

(b) DDIM (Steps=100)

(c) Heun (Steps=101)

(d) DDPM (Steps=100)

When x t = ˆ x t α t , ˆ σ t = σ t α t , we can rewrite the DPM-Solver-1 (Eq. (19)) as:

where the expression is exact the same as the ODE in EDM [13]. It indicates that the sampling trajectory in DPM-Solver-1 is equivalent to the one in EDM, up to a time-dependent scaling ( α t ). As lim t → 0 α t = 1 , the two solvers will leads to the same final points when using the same time discretization. Note that the DPM-Solver-1 is also equivalent to DDIM ( c.f. Section 4.1 in [16]), as also used in this paper.

With that, we can further verify that the Heun's method used in this paper corresponds to the DPM-Solver-2 when setting r 1 = 1 .

(b) DDIM (Steps=100)

(c) Heun (Steps=101)

(d) DDPM (Steps=100)

(c) Heun (Steps=101)

(d) DDPM (Steps=100)

G Broader Impact

The field of deep generative models incorporating differential equations is rapidly evolving and holds significant potential to shape our society. Nowadays, a multitude of photo-realistic images generated by text-to-image Stable Diffusion models populate the internet. Our work introduces Restart, a novel sampling algorithm that outperforms previous samplers for diffusion models and PFGM++. With applications extending across diverse areas, the Restart sampling algorithm is especially suitable for generation tasks demanding high quality and rapid speed. Yet, it is crucial to recognize that the utilization of such algorithms can yield both positive and negative repercussions, contingent on their specific applications. On the one hand, Restart sampling can facilitate the generation of highly realistic images and audio samples, potentially advancing sectors such as entertainment, advertising, and education. On the other hand, it could also be misused in deepfake technology, potentially leading to social scams and misinformation. In light of these potential risks, further research is required to develop robustness guarantees for generative models, ensuring their use aligns with ethical guidelines and societal interests.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes analyzes SDE and ODE-based samplers for diffusion models. Based on the analysis, this paper introduces a new solver, Restart, for sampling from diffusion models. The effectiveness of Restart is demonstrated on various unconditional and conditional generation tasks.\n\nSTRENGTHS:\n- The paper is well-written.\n- The proposed method shows better performance than EDM at moderate NFE regions.\n\nWEAKNESSES:\nI am willing to raise the score by 1 or 2 points if the authors address my concerns satisfactorily.\n\n**Weakness 1 : Ambiguity regarding Theorems 1 and 2.**\n- For Theorem 1, if we set $[t_{\\min},t_{\\max}] = [0,T]$, the terms for contracted errors $TV(p_T^{ODE_\\theta},p_T)$ and $TV(p_T^{SDE_\\theta},p_T)$ vanish because $p_T^{ODE_\\theta}$, $p_T^{SDE_\\theta}$, and $p_T$ are all identically Gaussian distributions. Then, we end up with terms depending on $\\delta$, $\\epsilon_{approx}$, and $t_{\\max} - t_{\\min}$ only, so Theorem 1 does not provide any insight into how ODE and SDE have distinct \"winning regions\", as illustrated in Figure 1 (b). The proof and claim for Theorem 1 should be reformulated such that even with $[t_{\\min},t_{\\max}] = [0,T]$, Theorem 1 explains how ODE and SDE have winning regions.\n- Likewise, I think Theorem 2 also should be proven for the entire interval $[0,T]$, so we can directly compare the errors for SDEs, ODEs, and Restart.\n\n**Weakness 2 : (Possibly) weak performance on the small NFE regime.**\n- How does Restart perform in the small NFE regime (NFE $\\leq 30$)? In Figure 3, the figure cuts off just before Restart and ODE intersect in the small NFE regime. This seems to contradict the claim that Restart combines the best of both ODE and SDE. Moreover, given the large size of SOTA diffusion models, it is crucial that diffusion samplers work well in the small NFE regime as well.\n-  How does Restart compare to recent fast samplers such as [1], [2], [3] in the small NFE regime?\n\n[1] DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps, NeurIPS, 2022.\n\n[2] Fast Sampling of Diffusion Models with Exponential Integrator, ICLR, 2023.\n\n[3] Denoising MCMC for Accelerating Diffusion-Based Generative Models, ICML, 2023.\n\nQUESTIONS:\n**Question 1** : How does Restart perform on FFHQ-1024 image generation (score function in https://github.com/yang-song/score_sde) compared to Improved SDE or ODE (Heun in EDM)?\n\nLIMITATIONS:\nDiscussed in Section 6.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

comments: SUMMARY:
ODE-based samplers plateau in performance while SDE-based samplers deliver higher sample quality. The paper attributes this difference to discretization errors and accumulated errors. Based on these, the authors propose a sampling algorithm called Restart which alternates between the forward diffusion process and backward ODE.

STRENGTHS:
1. The authors provide a theoretical explanation of the phenomenon that ODE samplers outperform SDE samplers in the small NFE regime but fall short in the large NFE regime.

2. The experimental results on image generation tasks validate the effectiveness of the method.

WEAKNESSES:
The proposed method relies on several hyperparameters (e.g. $S_{noise}, N_{restart}, i, K_i, t_{min, i}, t_{max, i}$), and the hyperparameters differ in different tasks. It would be hard to effectively tune these parameters in real application.

QUESTIONS:
Since the derivation of the Restart algorithm is motivated by the theoretical analysis of sampling error, is there a way to choose those hyperparameters based on the error analysis?

LIMITATIONS:
The author addressed the limitations of their work.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe papers propose a new sampling method for diffusion models, termed Restart Sampling. The authors first theoretically analyze the error propagation in diffusion models for stochastic and deterministic samplers under Wasserstein-1 distance and show that ODE samplers have a lower-discretization error but SDE samplers contract the initial distribution error as we run more steps. This agrees with the intuition and the experimental findings that support that ODE samplers are better for low NFEs but their performance flattens for more NFEs. Based on this analysis, the authors propose a method that tries to achieve the best of both worlds: it contracts the initial error with more steps and achieves the same discretization error as the ODE samplers. The implementation of the new method is very straightforward: one runs the ODE sampler and every K steps reverts back to some prior diffusion time using the forward model.\n\nSTRENGTHS:\nThe authors study a relevant problem in diffusion models. The proposed solution is simple, effective and novel. The theoretical results motivate the method and show clearly the differences between the SDE and the ODE samplers. The presentation of the paper is excellent. The authors show many experimental results, starting from toy models and going all the way to state-of-the-art text-to-image diffusion models. I think the paper and the method is of interest to the community and the audience of NeurIPS.\n\nWEAKNESSES:\nThe error propagation of diffusion models has been studied before. The results that I am aware of are from the papers \"Sampling is as easy as learning the score\", \"Restoration-degradation beyond linear diffusions: A non-asymptotic analysis for DDIM-type samplers\" and \"The probability flow ODE is provably fast\". The first studies the error propagation of the SDE sampling method and the latter two the propagation of errors for deterministic samplers. It would be beneficial to compare with these works, highlight potential differences in the approach and the final results, etc.\n\nAlso, apart from the Stochastic and the ODE samplers, there is a whole family of samplers that satisfy the same Fokker-Planck equations and hence give the same marginals, e.g. see the work \"Fast Sampling of Diffusion Models with Exponential Integrator\" and also some of the samplers used in the \"Elucidating the Design Space of Diffusion-Based Generative Models\" (EDM) paper. It would be interesting to compare theoretically and experimentally to these samplers.\n\nAnother concern I have is that the evaluation is only as thorough as it should have been and it is only done for relative high NFEs. Since evaluating the performance of a trained model is relatively easy, I would expect a more thorough benchmarking. If the performance of Restart sampler breaks for low NFEs, it is useful to know it and acknowledge it in the paper.\n\nQUESTIONS:\nI think it would be useful to know:\n\n* how the performance is for lower NFEs, e.g. 20 sampling steps?\n* the behavior of the sampler across more datasets and NFEs. For example, it would be useful to include more comparisons with the EDM paper (and the references therein) for different datasets.\n* what are the differences in the theoretical analysis compared to prior results known for error-propagation in diffusion models.\n\nLIMITATIONS:\nThe authors adequately addressed the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nBy analyzing of the trade-off between good sample quality and sampling time of both ODE and SDE-based generative models, a restart sampling strategy is proposed by this paper to combine the advantages of ODE and SDE sampling methods. The author proves two theorems that estimate the upper bound on the total error measured by the Wasserstein distance between generated and data distributions of ODE, SDE, and restart sampling methods respectively. It is illustrated that the total error can be decomposed into two parts: additional sampling error generated by discretization error and contracted error generated by the accumulated total error from previous sampling steps. Moreover, it is proved that ODE-based samplers have smaller additional sampling errors and SDE-based samplers have smaller contracted errors. So Comparing the three upper bounds on the total error, it can be proved theoretically that restart sampling yields a smaller total error because its additional sampling error and contracted error are both small. Finally, a range of experiments are done by authors which shows empirically that: 1. The total error of the restart sampler is indeed smaller than that of others. 2. Restart sampler surpasses previous SDE and ODE samplers in both speed and accuracy. 3. Restart sampler better balances text-image alignment/visual quality versus diversity than previous samplers.\n\nSTRENGTHS:\n1. This paper is written with meaningful motivation and a clear structure. The restart sampling method proposed by this paper is innovative, simple, and effective.\n2. The reasonableness and effectiveness of the resampling method are proved both theoretically and experimentally.\n3. The upper bounds on the total error of the three sampling methods give us an intuitive understanding of the advantages and disadvantages of the three sampling methods.\n\nWEAKNESSES:\n1. The effectiveness of restart sampling method on high-resolution image synthesis is not confirmed. For example, a comparison of sampling speed and accuracy on ImageNet 128*128, 512*512 should be added.\n2. The paper experiments on the sensitivity analysis of the number of restart iterations K, but there is no experiment on the sensitivity of another hyperparameter: the position and length of the restart interval.\n3. As pointed out in the paper, the contracted error further diminishes exponentially with the number of repetitions K though the additional error increases linearly with K. Figure 4 illustrates this trade-off phenomenon, too. So it may be hard to find a suitable K to make the restart algorithm work for different datasets or tasks globally, making it difficult to apply.\n\nQUESTIONS:\n1. Although the sampling step decreases by applying the restart strategy, the repetiness of ODE solver in the restart interval seems time-consuming. Moreover, does NFE computed in the restart interval added to the total NFE in Figure 3? \n2. Why does the total error remain the same using the ODE sampler when NFE>20? What are the results of three samplers with NFE<20? Is it consistent with the trend in Figure 1(b)？\n3. Why not compare to DPM-solver? Moreover, why not apply restart to DPM-solver to see if there are consistent benefits?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-11"}

comments: SUMMARY:
The paper proposed a sampler that balances sampling speed and quality by adding noises and restarting the process. They provide theoretical analysis to show a better upper bound of this method compared to original ODE and SDE samplers. Experiments are done to verify their claims.

STRENGTHS:
1. Authors identify the main cause of different performances of SDE and ODE in different regimes. And by taking advantage of the contraction by adding noises, they balance the speed and quality of the sampler.

2. The theoretical analysis is clear and well-written. 

3. The experiments are thorough with a good explanation of the choices of hyperparameters. 

4. The experiments show good results for the proposed method.

WEAKNESSES:
See questions.

QUESTIONS:
I wonder what the wall clock time looks like between restart, sde, and ode.

LIMITATIONS:
yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-13