Restart Sampling for Improving Generative Processes
Authors: Yilun Xu
Mingyang Deng
Xiang Cheng
Yonglong Tian
Ziming Liu
Tommi S. Jaakkola
Venue: NeurIPS 2023 poster (2023)
Keywords: Generative models
diffusion models
PFGM
sampling
TL;DR: We propose a new stochastic sampling algorithm that better balances speed and quality than previous ones

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
Yilun Xu" MIT ylxu@mit.. edu
Mingyang Deng" MIT dengm@mit.edi edu
Xiang Cheng" MIT chengx@mit edu
Yonglong Tian
Google Research
yonglong@ggoglle com
Ziming Liu MIT zml amliu@mit.edt edu
Tommi Jaakkola MIT tommi@csail wetissaal.iii.e...
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 sam- pling errors: ODE 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 ad- ditional forward steps and strictly following a backward ODE. Empirically, Restart sampler surpasses previous SDE and ODE samplers 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-IO/ ImageNet 64x64, 64. 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 the large-scale text-to-image Stable Diffusion model pre-trained LAION 512x512. Code is available com www.///weee///////eee/eeeaae htffffffffuss/ wwwffusioooeetttarrrrreeeaaeeeee restart, sampling
Introduction
Deep generative models based on differential equations, such 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 iterative backward processes that gradually transform a simple distribution (e.g-, Gaussian in diffusion models) into a complex data distribution by solving 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: ODE- samplers whose evolution beyond the initial randomization is deterministie, and SDE-samplery 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. ((b). ODE solvers [22, 16, 13] result 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 the expense of increased sampling time. To better understand these differences, we theoretically analyze SDE performance: the
'Equal Contribution.
37th Conference on Neural Information Processing Systems (NeurIPS 2023).

---

t =
PiPmote Páata)
006
ODE
Smalt NER Regien
torge NFF. Reginer
SDE
Feer
Fast
Restart
XK
drift
noise
NFE
(a)
(b)
Figure 1: (a) Illustration of the implementation of drift and noise terms in ODE, SDE, and Restart. (b) Sample quality versus number of function evaluations (NFE) for different approaches. ODE (Green) provides fast speeds but attains only mediocre quality, even with a large NFE. SDE (Yellow) obtains good sample quality but necessitates substantial sampling time. In contrast to ODE and SDE, which have their own winning regions, Restart (Red) achieves the best quality across all NFEs.
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 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, propose a novel sampling algorithm called Restart, which combines the advantages of ODE and SDE. As illustrated in Fig. (a), the Restart sampling algorithm involves K repetitions of two subroutines in pre-defined time interval: Restart forward process that adds a substantial amount 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 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 larger. Additionally, we apply multiple Restart intervals to further reduce the initial errors in more challenging tasks.
cycle
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 10x fewer steps for the same FID score on CIFAR- 100 using VP [23] (2x fewer steps on ImageNet 64 x 64 with EDM [13]), and outperforms ODE solvers even in the small NFE regime. When integrated into previous state-of-the-ar pre-trained models, Restart further improves performance, achieving FID scores of 1.88 on unconditional CIFAR-11 10 with PFGM++ [28], and 1.36 on class-conditional ImageNet 64 x 64 with EDM. To the best of our knowledge, these are the best FID scores obtained 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 x 512. Restart more effectively balances text-image alignment/visual quality (measured by CLIP/Aesthetic scores) and diversity (measured by FID score) with varying dassifier-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 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 chosen smooth distribution, and backward process that iteratively reverses the forward process. For instance, in diffusion models, the forward process is the diffusion process with no learned parameters:
dr = (f)dW.
where (6) is predefined noise schedule increasing with t, and W. E R 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 time t is the convolution of data distribution Pu = Pdata and a Gaussian kemel,, i..,, Pt = Po * (O, o² (t) The prior distribution is set to (0, o² (T)) since pr is approximately Gaussian with sufficiently large T. Sampling of diffusion models done via a reverse-time SDE [1] or a marginally-equivall
ODE [23]:
(SDE)
dr 26 ( log (c + + (t))()W)
(I)
(ODE)
dr = (()))(())))))))))))))))))))))))) (o)
(2)
where v log (x( in the drift term is the score of intermediate distribution at time t. W.l.o.g we set (1) 1 in the remaining text, as in [13]. Both processes progressively recover Po from the prior distribution pr while sharing the same time-dependent distribution pt- In practice, we train a neural network s0(x, i.) estimate the score field V, log (ar) by minimizing the denoising score-matching loss [25]. We then substitute the score V. log pr (z) with so (a, 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 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 's 2nd method [2] for ODEs (e.g., Eq. (2)), and
for SDEs (e.g.. Eq. (l)). 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 the increased discretization error caused by using a larger step size in numerical solvers.
However, illustrated in Fig. 1(b) and observed in previous works 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 decent quality with limited NFEs, whereas SDE samplers struggle in the same small NFE regime. However, the performance of ODE samplers quickly reaches 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 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 local error of O68)) whereas the first-order
Maruyaaa method for SDEs yeilds a local error of 0(83) (see e.g., Theorem 1 of [4]), where ó denotes the step size. As 0(83) >> O(82), 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 8 shrinks and discretization errors become less significant for both ODEs and SDEs. In this regime it is the approximation error error arising from an
3

---

inaccurate estimation of the ground-trum -tunh vector field by the neural network Sg starts to dominate the sampling error. We denote the discretized ODE and SDE using the learned field so as ODEg and SDEg. respectively. In the following theorem, we evaluate the total enrors from simulating ODE, and SDE, within the time interval [minn mmax] C [0, T). This is done via upper bound the Wasserstein-1 I distance between the generated and data distributions at time mmin- We characterize the accumulated initial sampling errors up until mmax by total variation distances. Below we show that the inherent stochasticity of SDEs aids in contracting these imitial errors at the cost of larger additional sampling error in [fmin / Consequently, SDE results in smaller upper bound as the step size ò nears O (pertaining to the high NFE regime).
Theorem 1 (Informal). Let tmax be the initial noise level and Pt denote the true distributiin at noise level t. Let Pt .Pr
ODE SDEe denote the distributions of simulating ODE,, SDEa respectively. Assume
that vt E [tmin tmax), || < f/2 any 2; in the support of Pt. Pt.
ODEs or Pt
SDE,, Then
Wa ODEE )leses) ) S B BV ODEE :Ptnss)
+ O + Capproz)) (tmex {mi))
W1 We SDE
S (1-te xe B BV SDE, Pl ) + O Ötmax + Cupproxe) (twex (mii))
Intal sreor
apper bound om contracted error
upper bound additional sampling crear
1. P X < 1 is a contraction factor, Ia and Capproz are uniform bounds on ||tso (24, t)]l and the approximation error ||tV, log pe(x) tss tso t) for all wt, t, respectively. O() hides polynomial dependency on various Lipschitz constants and dimension.
In the above, U = BL1/fmin +
We defer the formal version and proof of Theorem 1 to Appendix A.I. As shown in the theorem, the upper bound on the total error can be decomposed into upper bounds on the contructed error and additional sampling error. TV TV(ptun (DDEE Ptsms)) and TV TV(p (pE PP..)) correspond the initial errors accumulated from both approximation and discretization errors during the simulation of the backward process, up until time mma.. In the context of SDE, this accumulated error undergoes contraction by factor of 1 -e- /La L2 2he within [' m 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 [fminn m Despite the fact that SDE incurs a higher discretization error than ODE (O(v)) versus (d)), the contraction effect on the initial error is the dominant factor impacting the upper bound in the large NFE regime where d 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 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, introduce novel algorithm, termed Restart, designed to achieve these two goals simultaneously.
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 due to greater contraction effects.
4.1 Method
In the Restart algorithm, simulation performs few repeated back-and- steps within a pre-defined time interval [/min / C [0, T). as depicted in Figure (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 O. In addition, we refer to the backward process within the Restart interval [tminn mmax] as the
Restart backward process, to distinguish it from the main backward process.
Starting with samples at time fmin- which are generated by following the main backward process, the Restart algorithm adds large noise to transit the samples from tmin tmax with the help of

---

the forward process. The forward process does not require any evaluation of the neural network sa(x,t), t), as is generally defined by an analytical perturbation kemnel capable of transporting distributions from Eii to Emax. For instance, in the case of diffusion models, the perturbation kernel is (0, (mmmms))2 o ) The added noise in this step induces 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 fax back to tmin using the neural network predictions as in regular ODE. We repeat these borward-backward steps within mim mmx] interval K times in order to further derive the benefit from contraction. Specifically, the forward and backward processes in the 1th iteration (i E {0..... K 1}) proceed as follows:
(Restart forward process)
Feme X Lmnn + Etain
(Restart backward process)
I Pann = ODEg(a)) 1+1 tmax - mminn)
(4)
where the initial =Lnm. is obtained by simulating the ODE until /min: Tyo.. ODBg(DD - bmin). and the noise et tass is sampled from the corresponding perturbation kernel from miin mmax- 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 greater contraction effect, which we will demonstrate theoretically in the next subsection. For example, we set [imin mmax] = [0.05, 0.3] for the VP model [13] CIFAR-10. Repetitive use of the forward noise effectively mitigates errors accumulated from the preceding simulation up until fmax. 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 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 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 7th Restart iteration, following the discretized (estiaa())
trajectories and the learned field 89, as PiE(fmen 'amd''
Theorem 2 (Informal). Under the same setting of Theorem 1, assume K S L2z(tnae ban)) for some universal constant C. Then Kestung(K) Wi (P .Plais) (B- TV (p
Resuarn())) (1)
(K + 1). O (6 + Capprosr) ) (kmox Umin)
orror
upper conrencted
upper bound additional sampling
where X < is the same contraction factor as Theorem 1. 00) hides polynomial dependency on various Lipschitz constants, dimension.
Proof sketch. To bound the total error, we introduce an auxiliary process fmes which initiates
Restaria (i)
from true distribution Pt... and performs the Restart iterations. This process differs from Pte fnve]
Restarig
only in its initial distribution at tmax (Pt... versus P Prass triangular inequality:
Restarta (0)), We bound the total error by the following
Wi (p
Resarta (K)
-t )
Wi (p
Restart (K)
Itmm
Restart, (K)
+ W1 Wi(g Restaria (K) .Ptain)
total erroe
contracttdd aire
additional sampling error
To bound the contracted error, we construct a careful coupling process between two individual trajectories sampled from Ptain (esirtaf)) and Restmy(() i O...., 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, bound by multiplying B. The bound on the additional sampling error echoes the ODE analysis in Theorem I: since the noise-injection and ODE-simulation stages are separate, we do not incur the higher discretization emro of SDE.
5

---

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 tmax 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 I, 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 (Fq. (4) has the same order of discretization error (6)) as the ODE, compared to O(v5) 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 K that strikes a balance between reducing the initial error and increasing additional sampling error. Ideally, one should pick larger K when the initial error at time mmax greatly outweigh the incurred error in the repetitive backward process from max to min- We provide empirical evidences in Sec 5.2.
While Theorem 1 and Theorem 2 compare the upper bounds on errors of different methods, 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 and Theorem 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 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 contract accumulated errors comes to the fore.
to study how the already accumulated error changes
4.3 Practical considerations
The Restart algorithm offers several degrees of freedom, including the time interval [tmin; mmax]] and the number of restart iterations K. Here we provide 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 Restart? Theorem 2 shows that the Restart algorithm effectively reduces the accumulated етror at time mmax by a contraction factor in the Wasserstein upper bound. These theoretical findings inspire us to position the Restart interval [tmin fmaxx] 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 mmin 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.
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, 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].
Multi-level Restart
5 Experiments
In See 5.1, we first empirically verify the theoretical analysis relating to the Wasserstein upper bounds. We then evaluate the performance of different sampling algorithms generation benchmarks, including CIFAR-10 [14] and ImageNet 64 x 64 [5] in Sec 5.2. Lastly, we employ Restart on text-to-imasl toaeeeeimeeaaam generation, using Stable Diffusion model [19] pre-trained on LAION-5B [21] with resolution 512 x 512, in Sec 5.3.
standard image
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 and Theorem 2. Although our

---

0.875
0.35
0.875
0,850
Total Emaar
Contracted
ODE
0.825
0.800
Total
0.85
Restart
Restart
6.775
DGE
0.80
0.775
6.750
0.75
6.725
0.70
0.75
0.85
0.85
0.70
0.80 0.85 0.90 0.95 1.00 0.75
40
160
320
Additional Sampling error
Additional Sampling entor
NFE
(a)
(b)
(c)
Figure 2: Additional sampling error versus (a) contracted error, where the Pareto frontier is plotted
and (b) total error, where the scatter plot provided. (c) Pareto frontier of NFE versus total error.
theoretical analysis compares the upper bounds of the total, contracted and additional sampling errors, we further verify their relative values through synthetic experiment.
Setup We construct a 20-dimensional dataset with 2000 points sampled from Gaussian mixture, and train four-layer MLP to approximate the field V log Pt. We implement the ODE, SDE, and Restart methods within predefined time range of [fminn m = [1.0, 1.5], where the process outside this range is conducted via the first-order ODE. To compute various error types, define the distributions generated by three methods as outlined in the proof of Theorem 2 and directly gauge the errors at end of simulation & 0 instead of tmin (1) the generated distribution as Sampler where Sampler E (ODEs, SDE,, Restart, (K) (2) an auxiliary distribution To
Sumpler initiating Sampler their initial Sampler and go
Po
from true distribution Pt.sa. at time max. The only difference between Po
distribution at mmx (PDD versus ptmms): and (3) the data distribution pn- In line with Theorem we use Wasserstein-1 I distance Wi WP(Po Sumpler Sunnlee .40 Sample))
Wi W Sampler Po) measure the contracted error Wa
/ additional sampling error, respectively. Ultimately, the total error corresponds to Wi (i(ro (po .P)). Detailed information about dataset, metric and model can be found the Appendix C.5.
Results In our experiment, we adjust the parameters for all three processes and calculate the total, contracted, and additional sampling
across all parameter settings. Figure 2(a) depicts the Pareto frontier additional sampling error versus contracted error. We can see that Restart consistently achieves lower contracted error for given level of additional sampling error, compared to both the ODE and SDE methods, as predicted by theory. Figure 2(b), we observe that the Restart method obtains smaller total error within the additional sampling error range of [0.8, 0.85). During this range, Restart also displays 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 its enhanced contraction effects. From Figure 2(c), Restart also strikes an better balance between efficiency and quality, as it achieves a lower total eтror at a given NFE.
5.2 Experiments on standard benchmarks
34
ODE
3.00
ODE
FID score (50K)
3.2
Improved SDE
Gonna Fast
2.75
Improved SDE
1.8
3.0
Restarl
Restart
2/8
score
1.8
28
2.00
1.4
1.75
1.50
266
-612
1024
1024
32
64
128
258
512
1024
NFE
NFE
(a) CIFAR-10, VP
(b) ImageNet 61x64, EDM
Figure 3: FID versus NFE on (a) unconditional generation on CIFAR-10 with VP: (b) class- conditional generation ImageNet with 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++

---

3.4
Restart
3.2
Score (5D()
2.0
2.4
15
Figure 4: CIFAR-10, VP, in the low NFE regime. Restart consistently outperforms the DPM-Solver with NFE ranging from 16 36.
64 from [13, 28]. We also use the EDM discretization scheme [13] (see Appendix B.1 for details) during sampling.
models on CIFAR-10 or ImageNet
For the proposed Restart sampler, the hyperparameters include the number of steps in the main/ backward processes, the number of Restart iteration K, as well as the time interval [tmin fmax]. We pick the mmin and max from the list of time steps in EDM discretization scheme with number of steps 18. For example, for CIFAR- 1-10 (VP) with NFE=75, we choose miin =0.06, m =0.30, K=10, where 0.30/0.06 00 is the /1 time step in the EDM scheme. We also adopt EDM scheme for the Restart backward process in [fmin. m 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 x 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 10 (VP). Since the vanilla reverse- diffusion SDE [23] has significantly higher FID score, we omit its results from the main charts and
defer them Appendix D. For ODE samplers, we compare with the Heun's 2nd order method [2] (Heun). which arguably provides an excellent trade-off between discretization errors and NFE [13]. To ensure fair comparison, we use Heun's method the sampler in the main/Restart backward processes in Restart.
We report the FID score versus NFE in Figure 3(a) and Table 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 10-fold / 2-fold acceleration compared previous best SDE results on CIFAR-10 (VP) ImageNet 64 x 64 (EDM) at the same FID score. In comparison to ODE sampler (Heun), Restart obtains 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, also observed [13]. suggests that Restart incorporates stochasticity more effectively. (3) Restart establishes new state-of-the-an FID scores for UNet architectures without additional training. In particular, Restart achieves FID scores of 1.36 on class-cond. ImageNet 64 x on uncond. CIFAR-10 with PFGM++
64 with EDM, and 1.88
To further validate that Restart can be ap- plied 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 out- performs the JPM-Solvee with an NFE rang- ing from 16 to 36. This demonstrates Restart's capability 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. Sur- prisingly, when paired with the DPM-Soll Solver, Restart achieves an FID score of 2.11 on VP setting when NFE is 30, which is significantly
Table 1: Uncond. CIFAR-10 with EDM
3.2
EDM (Restort)
and PFGM++
3.0
EDM (DDil)
VP (Restatt)
2.8
(ODE)
NFE FID
2.6
FDM-VP [13]
ODE (Ieen)
1.97
2.
1.97
Improved SDE
2.27
2.45
Restart
43
1.90
20
40 60
Number of Restart flerations
PFGM++ [28]
Figure 5: FID score with a
ODE (Teun)
1.91
varying number of Restart itera-
1.91
tions K.
Restart
1.88
8

---

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 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, 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
DDIM (Steps=50)
22
DDIM (Steps=60)
DDIM (Steps=100)
DDIM (Steps=100)
FID score (SK).
20
DDPM (0t)ps=100)
FID score (5K)
20
DDPM (Steps=100)
DDPM (0t)
DDPM (Stepe=200)
18
Restart (Steps=66)
18
Restart (Stsps-66)
16
16
14
14
0.290
0,295
0.300 0.305 0.310 0,315 0.320
5:15
5.20
5.25
5.30
5.35
5.40
CLIP score (ViT-g/14)
Aesthetic score
(a) FID versus CLIIP score
(b) FID versus Aesthetic score
Figure 6: FID versus CLIP ViT-g/14 and (b) Aesthetic score for text-to-image generation at 512 x 512 resolution, using Stable Diffusion v1.5 with a varying classifier-free guidance weight w 2,,
We further apply Restart to the text-to- rext-to-image Stable Diffusion v1.5 pre- pre-trained LAION-5B [21] a resolution of 512. We employ 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, rated by the LAION-Aesto Aesthetics Predictor V2 [24]. Following [17], we compute all evaluation metrics using 5K captions randomly sampled from the validation and plot the trade-off curves between CLIP/Aesthetic scores and FID score, with the classifier- guidance weight w in (2,3,5,4 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 the Restart backward process, we empirically find that Heun performs better than DDIM in the Restart. In addition, we select different of the hyperparameters for each guidance weight. For instance, when w = 8, we use ['min, m [0.1, 2), , K-2 and 10 steps in Restart backward process. We defer the detailed Restart configuration Appendix C.2, and the results of Heun to Appendix D.I.
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). Remark- ably, Restart achieves substantially lower FID scores than other samplers when CLIP/Aesthetic scores are high (i.e., with larger values). Conversely, Restart generally obtains a better text-image align- ment/visual quality given the same FID. We also observe that DDPM generally obtains comparable performance with Restart in FID score when CLIPP/esthetic 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 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
hattpe://uugiiggen http://gg....aacoooooo///////tta wtt.:/taabbbelfffooooooo-o-oo

---

(a) Restart (Steps=66)
(b) DDIM (Steps=100)
(c) DDPM (Steps=100)
Figure 7: Visualization of generated images with classifier-free guidance weight 8, using four text prompts (*A photo of an astronaut riding horse on mars.", "A raccoon playing table tennis" "Intricate origami of a fox in a snowy forest" and "A transparent sculpture of duck made out of glass") and the same random seeds.
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 forward process that adds 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 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 [fmin {max] At present, adjust these parameters based on the heuristic that weaker/smaller models, or more challenging tasks, necessitate stronger Restart strength. In the future direction, anticipate developing more principled approach to automating the selection optimal hyperparameters for Restart based on the error analysis of models, in order to fully unleash the potential of the Restart framework.
Acknowledgements
YX and TJ acknowledge support from MIT-DSTA Singapore collaboration, from NSF Expeditions grant (award 1918839) "Understanding the World Through Code", and from MIT-IBM Grand Challenge project. Xiang Cheng acknowledges support from NSF CCF-21 12665 (TILOS Research Institute).
10

---

References
[1] Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313-33 32,,6,,, 1982.
[2] Uri M. Ascher and Linda R. Petzold. Computer methods for ordinary differential equations and differentiallalgebrai equations. In SIAM, 1998.
[3] Andrei Borodin and Paavo Salminen. Handbook of Brownian motion-facts and formulae. Springer Science & Business Media, 2015.
[4] Arnak Dalalyan and Avetik Karagulyan. User-friendly guarantees for the langevin monte carlo with inaccurate gradient. Stochastic Processes and their Applications, 229(12):577-53111 2019.
[5] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, K. Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. 2009 IEEE Conference on Computer Vision and Pattern Recognition, pages 248-255, 2009.
[6] 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.
[7] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge local nash equilibrium. In NIPS, 2017.
[8] Jonathan Ho. Classifier-free diffusion guidance. ArXiv, abs/2207.12592 2022.
(9] Jonathan Ho, Ajay Jain, and P. Abbeel. Denoising diffusion probabilistic models. ArXiv, abs/2006.1 1239, 2020.
[10] 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.
[11] 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.
[12] Alexia
Ke Li, Remi Taillef,,,,,,er, Tal Kachman, and loannis Mitliagkas. Gotta go fast when generating data with score-based models. ArXiv, 0iss21055.14080 2021.
[13] Tero Karras, Miia Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. ArXiv, abs/2206. 000422006000364 2022.
[14] Alex Krizhevsky. Learning multiple layers of features from tiny images. Citeseer, 2009.
[15] Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pictro Perona, Deva Ramanan, Piotr Dollár, and C. Lawrence Zitnick. Microsoft coco: Common objects in context. In European Conference Computer Vision, 2014.
[16] 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 2.Xiv:206.000927,, 2022.
[17] 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.
[18] 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 Machine Learning, 1021.

---

[19] Robin Rombach, A. Blattmann, Dominik Lorenz, Patrick Esser, and Björ Ommer. High- resolution image synthesis with latent diffusion models. 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 10674-10685, 2021.
[20] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L. Denton, Seyed Kamyar Seyed Ghasemipour, Bureu 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.
[21] 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 Kaezmarezyk, and Jenia Jitsev. Laion-5b: An open large-scale dataset for training next generation image-text models. ArXiv, 0tts/2210004400 2022.
[22] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. ArXiv, 0as/2000.02502, 2020.
[23] Yang Song, Jascha Narain Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Er- mon, and Ben Poole. Score-based generative modeling through stochastic differential equations. ArXiv, abs/201 1.14456, 2020.
[24] LAION-AI Team.
Laion- aesthetssthetccs predictor v2.
https://github. com/
christophscchuhmam httpss/shhhummmmmmmmmmeeeeeee imprmoedd www.aaaeeteteeeeeeeeeeeeeeeeeeee prociccorrrrtoroorr 2022.
[25] Pascal Vincent. connection between score matching and denoising autoencoders. Neural
Computation, 23:1661-1677 2011.
[26] 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 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.
[27] Yilun Xu, Ziming Liu, Max Tegmark, and T. Jaakkola. Poisson flow generative models. ArXiv, abs/2209. 11178, 2022.
[28] 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.
12

---

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 v log (æ)) More specifically, given any E Fd and following ODE:
0, let denote the solution to the
dx. = VV log (1)
(5)
ODE r) is defined "the value of 2, when initialized I x". It will also be useful to consider "time-disceetized ODE with drift tsg t)": I denote any integer. Let ö denote a step size, let T, denote the solution to
: let denote the discretization step size and let
di, tsg (#&6, kö)dt,
where for any 4, is the unique integer such that E ((k - 1)6, A:8]. We verify that the dynamics of Eq. (6) is equivalent the following discrete-time dynamics for = kō, k Z:
((k 1)5)² (1:6)2 se (Zk5, ko).
F(k- 1}6
We similarly denote the value of h.r when initialized J., let SDE(2, 8 - r) and SDE (z, - r) denote solutions to
ODEg (z, + r). Analogously,
dy: = 22 log Pr()) + 2tdB
dy: 21:0 t)dt 2tddB
respectively. Finally, we will define the Restart, process as follows:
(Restarts forward process) (Restarto backward process)
1/111 aux =t.... + E tm
1-1 1 3 2mm
ODE ( Imax + bmin).
(7)
where of Him ~ N (0, (Laxx ( I). We use Restarta (x, K) to denote of in the above processes, initialized at 2b = x. In various theorems, we will refer to function Q(r) : R [0, 1/2), defined as the Gaussian tail probability Q(r) = Pr(a 2 r) for a ~ NN(00)) 1).
A.1 Main Result
Theorem 3. [Formal version of Theorem 1] Let tmax be the initial noise level. Let the initial random variables Time Vtaes and
Ft
ODE,
(Ii max -> tmin)
D SDEg (T. mma > tmin);
Let pt denote the true population distribution at noise level t. Let Pt ODFs SDF, denote the distributions for Ft. Vt respectively. Assume that for all x, y, S. s0(x, satisfies [tss(2, (2.t) ts s)ll < Lols t). [ t)ll < L1, [ -tsgyy, t) < L2 || vll. and the approximation error (x,
log Pe (x)(( Capproz- Assume in addition that Vt E [min, max)) || < B/2 for any
D, the support of Pt- Pt
ODEs pt
SDE,, and K S L2(Amo for some universal constant Then
Wi Wi(p .Pinn)
ODEs
- B TV (ODED .Pinns
+ ba(tuu tmaa)) tam))
(6(TeTT1 T.o) + (appraz) (Amax (min)
(8)
Wi Wi SDEE
SDE .Pima) S (1 11- -BLi/hhaa -312 1la0
BLi/tao
TV SDEo .Prmn)
+ c 2. tma)) (Capprox + 8Lo + L2 (BL1 / mddtt (tmax tmina) (9)
where À :0
2
13

---

Proof. Let us define ti Plnne and let Beii ODE(E mmax + mmin). We verify that Fim has density Pimin Let us also define 2tmi -DEE (m tmax Emin). We would like to bound the Wasserstein distance between Itan and Eitni (I.e., Piam inequality:
ODEE and Pt by the following triangular
(I Dimin S W1 (I Btmm) W1 Wi (#t Ftrin)
(10)
By Lemma 2, we know that
|| Ftuaill S (frer ellera-tau tnn) (ő(L2h) + Lo) + capprox) (4max (min))
where use the fact that |itmm Ptam 0. Thus immediately have
Wi Dtmmn S ( tain) (8(L241 + Lo) + Eapproz) (tmax min)
(11)
On the other hand,
Wi Witti t
<B TV (frais Train)
SB TV (itam: :
(12)
where the last equality is due to the patceesonsssing inequality. Combining Eq. (11), Eq. (12) and the triangular inequality Eq. (10), we arrive the upper bound for ODE (Eq. (8)). The upper bound for
SDE (Eq. (9)) shares a similar proof approach. First, let ylme ~ Pinas Let Dt...
= SDE₁ yttme: tmax 7
min)). By Lemma
TV (irasi Mrrain
<1-20
B
BL1/tm L (Pan
TV (itran BeeeassBher Itana
2
2
P
$2 min
On the other hand, by Lemma 4,
E Illie.n Vemm II] 2ealialtun-aaii tmin)
(Cppprox + 8Lv L2 (ŌLs 25dttmux
(fmaxx
tminn)
The SDE triangular upper bound on W1 (Ttem Vtmn. follows by multiplying the first inequality by B (to bound Wi (Nt.. Htmm) )) and then adding the second inequality (to bound W1 Dtmn)) )). Notice that by definition, TV (itmer Dram. TV (Vim. Tr...). Finally, because of the assumption that K Ly(tan- tmm) 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 Ptan
Restarto &
denote the distributions after 7th Restart iteration, i.e., the distribution of T Restarta (T) i).
Given initial F.l t
Ressart.0 Ptace
let
xo fass
ODEg(T7 tanss" Ima + bmin). Then
Wi Wi(Pe (Pemm
Restarto.K K
PPacce) B x)x TV (p Restart O
upper bound on contracted error
+ e (1))a((ma lasie) (K 1) (8(L2L1 Lo) + Capproz) (tmax
t
apper bound on additional sampling error
(13)
where x = 2Q
B
2
Proof. Let x
~ Ptiaa Let x = Restartt(22)
(2) K).
K has density Ptm. Let us We verify that t
also define 20
ODEa (x) Imax bmin) and i.K Restarts (59 (#om.. K).
By Lemma 1.
K
B
TV (IK 2K 1
1 1 1-2Q
TV
(Tl)
P 2
max
K
B
1 2Q
TV
2
max
mnn
B
1 2Q
TV
(x.)
2
max
min
14

---

The second inequality holds by data processing inequality. The above can be used to bound the W-Wassesstein distance follows:
K TV
W1 (Ik Fiin"
..K Fani
BV
K
1 1 1 2Q
B
(I) F
Aain
2 max
1
(14)
On the other hand, using Lemma 3,
W1
a.K tain
- $1
Le (amm)))) (K 1) (ō(L2Iz + La) + Capprox) (famx mmin)
(15)
We arrive at the result by combining the two bounds above (Eq. (14), Eq. (15)) with the following triangular inequality,
Wi ( F Cam
WI
S tmm Wi Ftann
A.2 Mixing under Restart with exact ODE
Lemma 1. Consider the same setup as Theorem 4. Consider the Restarty process defined in equation Let
E Restarto(s
Vi Restarto (Y.... i).
Let Pt.
Restartull (6) and 4t
(estarta(i)
denote the densities of xt and vi respectively. Then
TV
Ptan
Restarto (K)
Itain
Restarto (K) (1 A) TV
Prain
Restarto (0) Restarto (0)
It
where A = 2Q
2v Casr
Proof. Conditioned on zi... Vim. let 21t1 Paax 12 S. We now define coupling between Tf and yiun by specifying the joint dis- tribution over E, and C."
2inm + Vthan 12 57 and Vieu vi...
If af vi... let 57' C that (+1 be coupled as described in the proof of Lemma 7, with 2-1 of y 5 / Pm.... min Under this coupling, verify that.
yf+1 On the other hand, if art F Vina let s and yiur
E
x fris
i+1
#
[+11] }]
<
1
i+1
* i+1l] }]
frass
| 2/ max
y
<E
1
2Q
I a
# werm
12
1 2Q
E 1
F Vimn
2
12
12
max
min
Applying the above recursively,
E [ fatk +
y H <
K
1 - 2Q
E
[1 (2 # Vlnnall
2 2/
12
min
The conclusion follows by noticing that TV
Restarto (K)
Mtam
Restarto (K)
< Pr (atte + vk
[1 { 2xKK # [KK]] and by selecting the initial coupling so that Pr (2.) # P Restarta(1) Ros
TV Pinie 4mme
15

---

A.3 Wi discretization bound
Lemma 2 (Discretization bound for ODE). Let Ita
ODE (I tmax - bmin) and let Ilan ODE, (Tteer Emax - tminn). Assume that for all 8, so s0(x, t) satisfies ltte (a 8) - Lols t. [z, t)ll L1 and [[tsg(x,t) ty t)ll I Le || y1l. Then
|| It |l S e ellarobbaa t|
Itm l (0(L211 + Lo) + capprox) (bmxx
(min))
Proof. Consider some fixed arbitrary k. and recall that ó is the step size. Recall that by definition of ODE and ODEg, for & E ((A: - 1)ó, kō),
dx; -t log (e(e
dx, -sso (Itds kd) kd)dt.
For E [fmin: tmax], let define time-reversed process ot x- Let u(r,) V log (ogp-((x). Then for t [-tmax: - tminl
dxt = tv(xt t)ds.
Similarly, define o == F- and J(, t) == 59 (z. -t). It follows that
dे = to पह kō)ds,
where & the unique (negative) integer satisfying t E [kō; (It + 1)6). Following these definitions,
d dt
| - x|
||tv(x] ,t) - tt(0x;; ,t]l + |[w(( ,t) tu( 1)ll
| ,t) -tū(xz k6)1l
+ | ,kŌ) [ kō)
Caprroz T-2 | xx{ || + 5To + T.2 - 2k8
Capproz L2 |, x; || + ôLo + 8L2L1.
Applying Gronwall's Lemma over the interval E [ tmaax min))
frii Immbl
x tais
La(hua -tuis) (11== tom x tmm + (Capprox + 8Lo + (LLL1) (tmax tmin))
= L2ma
(1128 Itmaal || + (Capprox + 8L0 أاطعع) (tmax tmmin))
Lemma 3. Given initial 2imm. let 20
ODE (Ior tmax - ttmin.. and let 20 ODEa (a)) tmax - tmmn). We further denote the variables after K Restart iterations xi Restart(x) K) and at Reestartt (49 K), with true field and learned field respectively. Then there exists coupling between K and 2K such that
| 2 S e((++))) ((K+1))ra(tmm (mm) (an) 1))( + Lo) + Capprox) (tmax tmin))
Proof. We will couple ain.. and f by using the same noise ei thn in the Restart forward process for 0... ...... K (see Eq. (7)). For any i, let us also define y Restarto
(x) . i), and this
process uses the same noise et_ +tt as previous ones. From this definition, yt Vtam Fain" thus bound
K.K We can
K-1
tuin Prmil <
Viuin
D.K
+
Σ
i.K ++1, Viai
V
(16)
16

---

Using the assumption that tso(:,t) is L2 Lipschitz,
0,1+1
y min
ODEa (Blma tmax - tmin)
ODE (2) (aim fmax + Emin)
6 La(tm bmio
Fain
where the last equality because we add the same additive Gaussian noise Et -ia to Vtam and i in the Restart forward process. Applying the above recursively, we get
0,0
2K Emin
L2(ttm - fma))
Se. K L:2(taa tria) I
<el (K+ +1)aatta (min (8(I.2I11 + Lo) Capproz) (Imax bmiin)
(17)
where the last line follows by Lemma when setting Ptrm Fl We will now bound
141 for some S K. It follows from definition that 2m
Mlam
i,++1 ODEg (1 Amax (min
Vini
++1,1+1 2
1+1 ODE (Ii Imux _ frmin)
Vtui
Enie
By Lemma 2,
(+1,+11 bann
S e (2ttam (8(L241 + Lu) + capproz) (frmx - (minn)
For the remaining steps from + ..... K, both y and yi+ evolve with ODE, in each step. Again using the assumption that tsg (-) is L2 Lipschitz,
1,1KK
(K)-)))a(tm bmn) (5(L2 + Lo) + Eapprox) (Imax Lmin)
Vinin
Summing the above for i = O...K - 1, and combining with Eq. (16) and Eq. (17) gives
K
s (k(K)))aaltss-t Immi) (K + 1) + Lo) + Capprox) (fmax trnin)
Aais
SDE (Itmm tmas tmii) and let Tt.... SDE (It Imaxx > bmin). Then there exists a coupling between Ft and T.p such that
Lemma 4. Consider the same setup Theorem 3. Let Zt,is
Dhan - Ztam 111 S e 2La(therttaa brai) (l Ieaee II]
+ 2222(trwe tuni) tmu) (Capprox + ōLv + L2 (L1 + v 222dtmme (fmax
(min)
Proof. Consider some fixed arbitrary k. and recall that is the stepsize. By definition of SDE and SDE,, for & E ((k 1)6, k6),
dx. - 21V lo + 2tdd
dr, (Tso kō) 2tdB.
Let us define coupling between 2, and F; by identifying their respective Brownian motions. It will be convenient define the time-reversed processes z x-ts and कद := T-1. along with var 1) = v log (x)) and T(x, 1) == (0( - -). Then there exists a Brownian motion B such that for E {- -mmax Imin).
dxf 20v(2 t)dt 22d
dz; = -2tT(Ikk, k6) + -
d(x; bt) -21 (v(zī, 1) T( kō)) dl,
17

---

where k is the unique negative integer such that & [kō, (k + 1)6). Thus
d E | Tt511 111
di
2 (E |||ev(at ; to B [1Iu tu(a ,t)l] ,
(E [||t t) ,,[ ko)ll + E [|] kd) to( kō)
22 (Eapproz + [ -[] 5L + L2 ب| - T IlI)
2 Capprox + L2 |||e - के I| + 8L0 L2 مان + v 22ddt
By Gronwall's Lemma,
E
Dlgii Ftai I]
E
x
x -fnii
2 22mos (nii))
E 1|
xt tea |] (apprrox + öLg + L2 (8L2 v mx2 (fmux tmmin)
21.2(mm tris)
E [1 Fe |lI Cappror 8hn + La (8hs + 2ddlmaax ) (Xmax lmiin)
A.4 Mixing Bounds
Lemma 5. Consider the same setup Theorem 3. Assume that 8 S Imin- Let
Ttam SDEa (Stawe bmax mmii))
Vt.m SDEg (ytmme tmax + min))
Then there exists coupling between x and ys such that
TV (Itar Yluce ) S 1 1 2Q
BL1 Li Paa
TV (I Blanes
2 Prax
min
Proof. We will construct coupling between 24 and yt. First, let (Itmm It..e) ) be sampled from the optimal TV coupling, i.e., Pr(It.. # yt = TV(2 (Xtmm Vires ). Recall that by definition of SDE,, fort E ((k- 1)8,k)) kö).
dæ, -2ts (xKS, ko)dt 2tdB.
Let us define time-rescaled version of x/: 34 3y22 We verify that
dr. so ( k:))d dB:,
where k is the unique integer satisfying t [[(k: 2) Next, we define the time-reversed process w := H _-,, and let v(x, t) = s0 t). We verify that there exists Brownian motion B such that, for & E [-12 12ax:
dæ = of + dB,,
where or = so [1k8}= -ō)) where & is the unique positive integer satisfying -t E (((k 1)8)2 (k8) Let dy. v.d dt + dB,, be defined analogously. For any positive integer and for any / E E (k6)) -(k - 1)8)2), let us define
2, oz kх
(2k 1)5² (v= - 2 ((a)))
B B (x(y)))) -
- (By- B" (=(()))))
Let 7/ =
We will now define coupling between dBy and dB as |
dB = (I - 21 (t < Thin? dB;
18

---

where 1 { denotes the indicator function, i.c. 1 r)= T, and T is a stopping time given by the first hitting time of 2; = O. Let r; = | |- Consider some E (-188, -(i 16-11 ( and Let j == tys (assume w.l.o.g that this is an integer), then
(v (u= - v <<<p)pp) 2k 1)0²
1 1tt $ -
Tt
-
k=i
(k² (K: 1) 2222 (tmin) + f
1lt S 72ddB
k=i
(i- 1)6²
2L1 ds min
'"
1 1 SttSB)2dBt
where dB. = (,t, d
dBg) a
a 1-dimensional Brownian motion. We also verify that
T-
3-12,
T 12m 1_1 + (2j 1)0²
-1- -1aa
+Br
Br
B²
-(B3
B
x + (2j 1)6²
+ vー 1)6²2²==
(B" -B 12-2m
g (2j 1) amn +
B BB- Bi By
B 1am
SB
where the third relation is by adding and subtracting 111--1))28-
B and using triangle inequality.
The fourth relation is by noticing that xI e2
+ (2j - 1)0²²² t + B
B 1mn 1
11-1)
(= 82 and that TT + (2j - 1)62
G B + BY
B" -
-
and then using our assumption in the theorem statement that all processes are supported on ball of radius B/2.
We now define process s, defined by ds; = 2L1 di 2dB initialized at s-12 We can verify that, up to time т. Tt St with probability Let denote the first-hitting time of st to O, then r I r' with probability 1. Thus
B r- -
Pr(t & (in) 2Pr(T' < m 2 2Q
B
Bl.1//am 1.3 /
2
12
12
where we apply Lemma 6. The proof follows by noticing that, i S m then Tiai = Yttma This because r 1- ((k - 1))2 then x- (kk-1)282 T 1-(N-1)055 and thus , = D. for all
t (k 1)28², in particular, t -t2 min
Lemma 6. Consider the stochastic process
dr: = dB, + edt.
Assume that To S B/2. Let denote the hitting time for = 0. Then for any T e IR+,
B
2x
Pr(s TT)
e -ac
2VT
where Q is the tail probability of a standard Gaussian defined in Definition 1.
Proof. We will use he following facts in our proof:
1. For x ~ N(O, Pr(x > r) = (1 er (-)) = ferfee (
dexp 2 dl er erfe
2. [T 12
2Pr (N(0, T) > a) 2Q a
by definition Q.
VIT
VT
19

---

Let dr, = dB, + cdt, with = The density of the hitting time is given by
(a+et)
exp
pT= f(,,,,))
(18)
V2rt²
(see e.g. [3]. From item above,
l f(a,0,t)dt = 2Q
In the case of a general c + 0, we can bound (a+tof)
a² 21
+ ac + S/4 Consequently,
2 a
ac
Therefore,
Pr(r ST) [ f(a,c,t)dt [[ fa.00,)dtte-=
e ac- Sr
22T
A.5 TV Overlap
Definition 1. Let be sampled from standard normal distribution (0, 1). We define the Gaussian tail probability Q(a):= = Pr(x 2 a).
Lemma 7. We verify that for any two random vectors & (N r2 and Ey NN0, o2),, each belonging Rdd the total variation distance between x' x E. and y = y Ey given by
=1 11 (r) S
2r
--1
72 + 2T
where 20
||=-v|| and Q(r) = Pr(E 2 r), when & ~ 1..
Proof. Let ;= =
2 M
We decompose x', , into the subspace/orthogonnal space defined by
x' xt + c + x!l + gll
y' = y + E + vy" + c!l
where we define
xl = R/2"" x
x = x x
n = 77 T
yy = y
!! r7 T &
E = Ex- El
< = 7> T
c = Ey c!]
We verify the independence + IL e!! and e IL su as they are orthogonal decompositions of the standard Gaussian. We will define coupling between x' and y' by setting & = Ey. Under this coupling, verify that
(x'+ (.) - (y+ + Ey = - V- m>_ (x y)
Therefore, x' "=yf y if and only if x!ll + c!! nlll EV- Next, we draw (a, b) from the optimal coupling
between (0) and N |I=-v 1). We verify that æ!l + En and ylll ell both lie in the span of 7. Thus suffices to compare (7, ll + e!! and (7.1)) + <ll). We verify that (7,
xr l <!!
20

---

(n,31 + (,,,z) / 71 e! ~N((7, ) || Ull ,ll _ (,,81)) -bb. We similarly verify that (7: { - (r.y)) + (7. ~ (7((7 إلر .²2 (7.1!)) + ou.
22 ( (12-~11) The last inequality follows from
Thus (V(x', ./ (oa, ob) =1 122
1)) 7 2m
-2/2
B More on Restart Algorithm
B.1 EDM Discretization Scheme
[13] proposes discretization scheme for ODE given the starting mmax and end time émin- Denote the number of steps as N, then the EDM discretization scheme is:
tNN
Emax
min
tmnxx)
with to max and tN-1 Emin- P is hyperparameter that determines the extent to which steps near min are shortened. We adopt the value the EDM scheme to creates time discretization in each Restart interval [frmxs fmin] in the Restart backward process, as well as the main backward process between [0, T (by additionally setting fmin 0.002 and fn = 0 as in [13]). It is important to note that min should be included within the list 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.
suggested by [13] in all our experiments. We apply
Algorithm EDM Scheme(tmin fmax, N, p = 7)
N-1
1: return (trixx + N-1 (trin tmmx)) ))
i=0
B.2 Restart Algorithm
We present the pseudocode for the Restart algorithm in Algorithm 2. this pseudocode, we describe more general case that applies 1-level Restarting strategy. For each Restart segment, the include the number of steps in the Restart backward process NRestat the Restart interval [fmin [max]] and the number of Restart iteration K. We further denote the number of steps in the main backward process Nmain We the EDM discretization scheme (Algorithm 1) construct time steps for the main backward process (4g T, Ne.. 0) as well as the Restart backward process, when given the starting/end time and the number of steps.
Although Heun's 's 2nd order method [2] (Algorithm 3) is the default ODE solver in the pseudocode, 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 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 A.00 GPUs.
21

---

Algorithm 2 Restart sampling
1: Input: Score network sp, time steps in main backward process 4e (0. Restart parameters
(NRresart... Kj. Iminn.j- f)
2: Round f11.1 to its nearest neighbor in 4e (o,
3: Sample 20 ~ NO0, II)
4: for O... Nmain do
P Main backward process
5:
Tt++1 = OneSttp__eunn(M5 L2, li+1)
Di Running single step ODE
6:
if 3j E {l..... l++1 = f then
7:
mmin Emin.j> mmax I ij j K
Kj, Resttrr R
8:
Fimin Iit+1
9:
for k O....
do ~ NO, (L
> Restart for K iterations
10:
Etuia
2 I)
11:
x+1
F + fnit Inas
Restart forward process
{fm] J m=
Nasam
EDM Scheme (Emin mmax N NRestart)
13:
for m = 0.. Resstrt - 1 do
D Restart backward process
14:
2 Em+
OneStep_Hor Heui(p__eee(ssp.. Em Em+1)
15:
end for
end for
17:
end if
end for
19:
return Nunn
Algorithm OneStep HetrStep_Heun(SSS
Heun(sp. s0, Fty l; (4+1)
1: d; = 1,80 ( 44)
2: Ft,t1 It. (1++1 d:
3: if f1+1 # then
4: de d== =t
5: Tta-1 Ita (ti-1 t.)( (d))
6: end if
7: return T1,41
C.1
Configurations for Baselines
We select Vanilla SDE [23]. Improved SDE [13], Gonna Go Fast [12] as SDE baselines and the Heun's s 2nd order method [2] (Alg 3) as ODE baseline on standard benchmarks CIFAR-10 and ImageNet 64 64. We choose DDIM [22]. Heun's 2nd 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 recom- mended dataset-specific hyperparameters (e.g., Smax, Sminy Setur) 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 Vanilla SDE and we directly report the FID scores listed in [12] for Gonna Go Fast on CIFAR-10 10 (VP), For fair comparison, we use the EDM discretization scheme [13] for Vanilla SDE, Improved SDE, Heun as well as Restart.
We borrow the hyperparameters such as discretization scheme or initial noise scale on Stable Diffusion models in the diffuser" code repository. We directly use the DDIM and DDPM samplers implemented in the repo. We apply the same set of hyperparameters Heun and Restart.
C.2 Configurations for Restart
We report the configurations for Restart for different models and NFE on standard benchmarks CIFAR-16 10 and ImageNet 64 x 64. The hyperparameters of Restart include the number of steps in the main backward process Nmain the number of steps in the Restart backward process NRestart the Restart interval [fmin, mmax] and the number of Restart iteration K. In Table 3 10, VP)
22

---

we provide the quintuplet (NV NRestart miinn mmax: K) for each experiment. Since we apply the multi-level Restart strategy for ImageNet 64 64, we provide Nmain as well as list of quadruple {(NRestatt K2, Emin, a [)= 1 (1 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 m to its nearest neighbor the time steps of main backward process, as shown in line of Algorithm 2. We apply Heun method for both main/backward process. The formula for NFE calculation is NFE
2. N
K2
(2 (Neestant, 1)) in this case. Inspired by
number repetitions per iteration in 1h Resiart interval [13], we inflate the additive noise in the Restart forward process by multiplying Spoise 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 x 64 when E [0.01, 1].
main backward process
We further include the configurations for Restart on Stable Diffusion models Table 10, with a varying guidance weight w. Similar to ImageNet 64 x 64, we use multi-level Restart with fixed number of steps Nmain 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 't yield significant improvements over the Euler method, yet necessitates double the steps. The number steps equals to Nmain + E K; (2 RNeestart. 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, adjust the configuration manually, guided by the heuristic that weaker/smaller models or more challenging tasks necessitate stronger Restart strength (e.g., larger K, wider Restart interval, etc). On average, we select the best configuration from sets for each experiment; these few trials have empirically outperformed previous SDE/ODE samplers. We believe that developing 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 *, and PFGM++ (D) 2048) model from the PFGM++ repository S. 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:
ăg (z, c, t) = wse (z, c, t) + (1 - ())
where e is the conditions, and so c, t)/ (z, t) is the conditional/uncondtiionn models, sharing parameters. Increasing would strength the effect of guidance, usually leading to 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, from mixture of two Gaussians in R.4 Next, we project these points onto 20 To ensure a variance of I on each dimension, we scale the coordinates accordingly. This setup aims simulate data points that primarily reside on a lower-dimensional manifold with multiple modes.
The specific details are as follows: S ~ 0.3.N (a, s² + 0.7(-at 20 where a = (3, 3,33 R² and s = 1. Then, we randomly select projection matrix P E R where each entry is drawn from (0,, 1), and compute S2 PS,. Finally, we scale each coordinate by constant factor to ensure variance of 1.
http::::///gituuu.oo/////// 'https://githuu wttptttwww..o//eebbee////ttgge
23

---

ODE SDF
ODE
0.95
Improved SDE
0,875
Impr
Vanila SDE
Contracted
0.80
Restart
Restart.
Total Ergar
0.025
ODE
0.75
0.85
Improved SDE
0.80
Restart
0.775
0.75
0.65
0.70
0.85
0.70
0.75
0.50
0.85
0.50
0.85
0.70
0.80 0.85 0.90 0.95 Additional Sampling enter
0.75
1.00
40
180
320
Additional Sampling arror
NFE
(a)
(b)
(c)
Figure 8: Comparison of additional sampling error versus (a) contracted error (plotring the Pareto
frontier) and (b) total error (using scatter plot). (c) Pareto frontier of NFE versus total error.
C.5.2 Model Architecture
We employ a common MLP architecture with 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 8(t) (6q.(6) in [13]). Hence, regarding SDE, we consider NFE values from {20, 40, 80, 160, 320], and multiply the original (t) = o(t)/a(( (t [13] with values from {0, 0.25, 0.5, 1.5, 2, 4, 8). is important 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 {20, 40, 80, 160, 320, 640).
Additionally, we conduct an experiment for Improved SDE in EDM. We try different values of Sahum 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 Sahurn takes values of [0.2 x steps, 0.5 x steps, 20, 60). The plot combines the results from SDE and is displayed in Figure 8.
To mitigate the impact of randomness, 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 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 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 Ta- ble 2,3 (CIFAR-10 VP, EDM, PFGM++) and Table 4, 5 (ImageNet 64 x 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 statisties of CIFAR-10 and ImageNet 64 x 64. We follow the evaluation protocol in EDM [13] that calculates each FID scores three times with different seeds and report the minimum.
24

---

We also provide the numerical results on the Stable Diffusion model [19], with a classifier guidance weight w 2, 5, in Table 6, prompts sampled from the COCO validation set. We evaluate CLIP score [6] with the open-sourced ViT-g/14 [1 Aesthetic score by the
As [17], we report the zero-shot FID score on 5K random
recent LAION-Aesthetics Predictor V2 o We average the CLIP and Aesthetic scores over 5K generated samples. The number of function evaluations two times the sampling steps in Stable Diffusion model, since each sampling step involves the evaluation of the conditional and unconditional model.
Table 2: CIFAR-10 sample quality (FID score) and number of function evaluations (NFE) on VP [23] for baselines
NFE
FID
ODE (Heun) [13]
1023
2.90
511
2.90
255
2.90
127
2.90
63
2.89
2.97
Vanilla SDE [23]
1024
2.79
512
4.01
256
4.79
128
12.57
Gonna Go Fast [12]
1000
2.55
329
2.70
274
2.74
179
2.59
147
2.95
72.29
Improved SDE [13]
1023
2.35
511
2.37
255
2.40
127
2.58
63
2.88
35
3.45
Table 3: CIFAR-10 sample quality (FID score), number of function evaluations (NFE) and Restart configurations VP [23], EDM [13] and PFGM+ [28]
Configuration
Method
NFE
FID
(N (Mminn- NRestant K₁, fmin, tmax,s)
VP
519
2.11
(20, 9, 30, 0.06. 0.20)
115
2.21
(18. 3, 20, 0.06, 0.30)
75
2.27
(18,3, 3, 10, 0.06, 0.30)
55
2.45
(18,3, 5, 0.06, 0.30)
43
2.70
(18,3, 2, 0.06,.0..... 0.300
EDM
43
1.90
(18, 2.114,0..333
PFGM++
43
1.88
(18, 2, 0.14 0.300
25

---

Table 4: ImageNet x sample quality (FID score) and number of function evaluations (NFE) on
EDM [13] for baselines
NFE
FID (50k)
ODE (Heun) [13]
1023
2.24
511
2.24
255
2.24
127
2.25
63
2.30
35
2.46
Vanilla SDE [23]
1024
1.89
512
3.38
256
11.91
128
59.71
Improved SDE [13]
1023
1.40
511
1,45
255
1.50
127
1.75
63
2.24
35
2.97
Table ImageNet 64 x 64 sample quality (FID score), number of function evaluations (NFE) and
Restart configurations on EDM [13]
NFE FID (50k)
Configuration
Nmain- {(NRestan,, Ki, min,, 4 .)
36, ((10, 3, 19.35, 4007 3, 1.09, 1.92). (7, 6, 0.59, 1.09), (7, 6, 0.30, 0.59), (7, 25, 0.06, 0.30))
623
1.36
36, (66, 1, 19.35, 40. (6,79).66, 1, 1.09, 1.92),
535
1.39
(7, 6, 0.59, 1.09), (7, 6, 0.30, 0.59), (7, 25, 0.06. 0.30)) 1.09. 1.92),
36. ((3, 1, 19.35, 40.79)
385
1.41
(6, 5, 0.59, 1.09), (6, 5, 0.30, 0.59), (6, 20, 0.06, 0.30))
36, ((4, 1, 19.35, 40.79) 1, 1.09, 1.92), (4, 0.59, 1.09), (4, 5, 0.30, 0.59), 6, 0.06, 0.30)
203
1.46
18. ((3, 1, 19.35, 40.79) 1, 1.09, 1.92),
165
1.51
(4, 5, 0.59, 1.09), (4, 5, 0.30, 0.59), (4, 10, 0.06, 0.30))
18. ((3, 1, 19.35, (0.79)
1.09. 1.92),
99
1.71
(4, 4, 0.59, 1.09), (4, 1, 0.30, 0.59), (4. 0.06, 0.30))
18, ((5, 1, 19.35, 40.79), 1, 1.09, 1.92),
67
1.95
(5, 1, 0.59, 15, 0.06, 0.30))
14, ((3, 19,35,
39
2.38
(3. 1, 1.09, (3, 0.30)
26

---

Table 6: Numerical results on Stable Diffusion v1.5 with classifier-free guidance weight 2
Steps
FID (5k) +
CLIP score f
Aesthetic score f
DDIM [22]
50
16.08
0,2905
5.13
100
15.35
0.2920
5.15
Heun
51
18.80
0,2865
5.14
101
18.21
0.2871
5.15
DDPM [9]
100
13.53
0,3012
5.20
200
13.22
0.2999
5.19
Restart
66
13.16
0.2987
5.19
Table 7: Numerical results Stable Diffusion v1.5 with classifier-free guidance weight 3
Steps
FID (5k) 4
CLIP score f
Aesthetic score f
DDIM [22]
50
14.28
0.3056
5.22
100
14.30
0.3056
5.22
Heun
51
15.63
0,3022
5,20
101
15.40
0.3026
5.21
DDPM [9]
100
15.72
0.3129
5.28
200
15.13
0,3131
5,28
Restart
66
14.48
0.3079
5.25
Table 8: Numerical results Stable Diffusion v1.5 with classifier-free guidance weight 5
Steps
FID (5k) +
CLIP score f
Aesthetic score f
DDIM [22]
50
16.60
0.3154
5.31
100
16.80
0.3157
5.31
Heun
51
16.26
0.3135
5.28
101
16.38
0.3136
5.29
DDPM [9]
100
19.62
0,3197
5,36
200
18.88
0,3200
5.35
Restart
66
16.21
0.3179
5.33
Table 9: Numerical results on Stable Diffusion v1.5 with classifier-free guidance weight
8
Steps
FID (5k)
CLIP score f
Aesthetic score f
DDIM [22]
50
19.83
0,3206
5.37
100
19.82
0.3200
5.37
Heun
51
18.44
0.3186
5,35
101
18.72
0,3185
5.36
DDPM [9]
100
22.58
0.3223
5.39
200
21.67
0.3212
5.38
Restart
47
18.40
0.3228
5.41
D.2 Sensitivity Analysis of Hyper-parameters
We also investigate the impact of varying Emin when tmax = Emin + 0.3, and the length the restart interval when tmin 0.06. Fig. 10(a) reveals that FID scores achieve minimum a fmin close to 0
27

---

Table 10: Restart (Steps=66) configurations on Stable Diffusion v1.5
Configuration
N Mainn { Ki, (min, /max,
30, [(5, 2, 1, 9), (5, 2, 10)1
30, [(10, 2, 0.1. 3)]
30, [(10, 0.1, 2)
30, ((10, 2, 0.1, 2))
Restart (8tsps-86)
(Step=ddd)
(Stxpe-50)
DDIM (Stepen50)
DDIM (Stpp=-00) (0)
DDM (3tpps-100)
20
(Stepa=611
20
(Slepse61)
Heur (Stepe=101) 001)
Heun (Stspe-101)
DDPM (Stspa=1000 100)
DDPM (Stppa-100 100)
DDPM (Stups=200)
DOPM (Steps=200)
14
14
0,300 0.310 CLIP score (ViT-gi11)
5.15
5.20
5.25
5.30
5.35
5.40
Aesthelic score
(a) FID versus CLIP score
(b) FID versus Aesthetic score
Figure 9: FID score versus (a) CLIP ViT-g/14 score and (b) Aesthetic score for text-to- image generation at 512 512 resolution, using Stable Diffusion v1.5 with varying classifier-free guidance weight w 5,
VP, indicating higher accumulated errors at the end of sampling and poor neural estimations small t. Note that the Restart interval 0.3 is about twice the length of the one in Table and Restart does not outperform the ODE baseline EDM. This suggests that, as a rule of thumb, we should apply greater Restart strength (e.g., larger K, mmxx (min) for weaker or smaller architectures and Vice versa.
In theory, 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 tmax close to the dataset's radius usually ensures effective mixing when fii small.
3.2
interval length FID
3.0
2.8
FID score (50K)
2.8
EDM (Restart)
EDM (ODE)
2.6
restart vp
VP (Restart)
2.4
VP (ODE)
edm restart
edm
2.2
2.2
2.0
-4
-2
2
(tg(tmii
0.0
0.1
0.2
interval length
0.3
0.4
0.5
0.5
(a)
(b)
Figure 10: (a): Adjusting Emin in Restart on VP/EDM; (b): Adjusting the Restart interval length when Imin 0.06.
28

---

E Extended Generated Images
In this section, provide extended generated images by Restart, DDIM, Heun and DDPM 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 classifier-guidance weight 8.
(a) Restart (Steps=66)
(b) DDIM (Steps=100)
(c) Heun (Steps=101)
(d) DDPM (Steps=100)
Figure 11: Generated images with text prompt="A A photo of an astronaut riding horse on mars" and
8.
Heun's method is DPM-Solver-2 (with 7'2 = 1)
The first order ODE in DPM-Solver [16] (DPM-Solver-I) 1) is in the form of:
Olt:
at. - )..) -6. )o, Vz log (ft.)
iow (Ô).-) a:
(19)
The first order ODE in EDM is in the form of
Tr 1 It, - (ot.-n o1,)ot. Vz log Pa (st.)
(20)
29

---

(a) Restart (Steps=66)
(b) DDIM (Steps=100)
(c) Heun (Steps=101)
(d) DDPM (Steps=100)
Figure 12: Generated images with text prompt=" A raccoon playing table tennis" and 8.
21, न. JeC,, we can rewrite the DPM-Solver-1 (Eq. (19)) as: When x, aq"
Ft..1 -T;; - (att.1 - -1.) 0r. V log pa. (0t,)
1
Ft, (Ott.- obe 01, V log Par, (Zt.) an.
(change-of-variabee
-(tt.. 0r. v log Po., (xt.)
where the expression exact same as the ODE in EDM [13]. indicates that the sampling trajectory in DPM-Solver-1 is equivalent to the one in EDM, up to time-dependent scaling (ox). As lim,-oo a, 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 this paper corresponds to the 2 when setting 1.
30

---

(a) Restart (Steps=66)
(b) DDIM (Steps=100)
(c) Heun (Steps=101)
(d) DDPM (Steps=100)
Figure 13: Generated images with text prompt="Intriee IIntical origami of a fox in a snowy forest" and
w 8
31

---

(a) Restart (Steps=66)
(b) DDIM (Steps=100)
(c) Heun (Steps=101)
(d) DDPM (Steps=100)
Figure 14: Generated images with text prompt="A A transparent sculpture of duck made of glass" and to 8.
32

---

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, 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 entertainment, advertising, and education. On the other hand, 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.
33


Reviews:
[{"review_id": "PWdc2fke48", "reviewer": "Reviewer_HtQK", "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. ", "questions": "I wonder what the wall clock time looks like between restart, sde, and ode.", "limitations": "yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "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.\n\n2. The theoretical analysis is clear and well-written. \n\n3. The experiments are thorough with a good explanation of the choices of hyperparameters. \n\n4. The experiments show good results for the proposed method.", "weaknesses": "See questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review and thoughtful feedback. Below we address specific questions.\n\n**Q1: I wonder what the wall clock time looks like between Restart, SDE, and ODE.**\n\nA: Thank you for the question. The wall clock time during sampling is approximately proportional to the NFE (number of function evaluations), which is reported in the paper. This is because the primary computational bottleneck during sampling lies in evaluating the neural networks. Any other overhead, in comparison, is negligible. Therefore, when comparing Restart, SDE, and ODE, their wall clock time will closely align with their respective NFEs.\n"}, {"review_id": "EESeFobKzY", "reviewer": "Reviewer_YGP7", "summary": "By 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.", "questions": "1. 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?", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. 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.", "weaknesses": "1. 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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review and thoughtful feedback. Below we address specific questions.\n\n**Q1: The effectiveness of Restart sampling method on high-resolution image synthesis is not confirmed. A comparison of sampling speed and accuracy on ImageNet-128/-512 should be added.**\n\nA: Thanks for the suggestion. We agree with the reviewer that including the suggested datasets can strengthen our experimental results. However, we would like to note that we have already conducted experiments at a resolution of 512x512, on the Stable Diffusion model. In particular, Restart demonstrates a superior FID and CLIP/aesthetic score trade-off, underscoring its scalability to higher-resolution images. As suggested by reviewer kog8, we've also presented a qualitative demonstration of Restart's effectiveness on FFHQ-1024 (please refer to Fig.2 in the rebuttal PDF).\n\n**Q2: There is no experiment on the sensitivity of another hyperparameter: the position and length of the restart interval.**\n\nA: Thank you for the suggestion. In Fig.9 of our paper, we illustrated the sensitivity of $t_{min}$ while keeping $t_{max}$ constant. We additionally add more sensitivity analysis of both the position and length of the restart interval on CIFAR-10. We only implement one Restart interval in all the experiments. We include the results in Fig.3 in the rebuttal PDF.\n\nFor sensitivity to Restart length $t_{max}-t_{min}$, we fix $t_{min}$ at 0.06 for VP and 0.14 for EDM. 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.3(a)). In practice, selecting $t_{max}$ close to the dataset's radius usually ensures effective mixing when $t_{min}$ is small.\n\nFor sensitivity to $t_{min}$, Fig.3(b) shows that a moderately small $t_{min}$ minimizes the approximation error post-restart on CIFAR-10. However, the contraction effect diminishes as $t_{max} - t_{min}$ diminishes. \n\n**Q3: 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\nA: Thanks for pointing this out. We agree that fully optimizing the hyper-parameter $K$ would be challenging. However, in general, the quality of generated images initially improves and later worsens with an increasing $K$. This trend makes it feasible to identify an appropriate $K$ value through a straightforward binary search.\n\nIn our experiments, we've found that choosing a suitable value for $K$ that leads to improved performance is relatively easy.  For example, introducing a Restart interval with $K=2$ at small time consistently outperforms the baselines for all the datasets given the same NFE. In addition, as a heuristic, one could set $K$ to reasonable values by following the recipe that at a smaller time $t$, a larger $K$ is necessary to contract more accumulated errors. Nevertheless, we acknowledge that determining the optimal $K$ for different Restart intervals could be intricate. We will delve deeper into this in future studies.\n\n**Q4: 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?**\n\nA: Yes, the total NFE reported in the paper includes both the NFE in Restart intervals as well as the NFE in the main backward process. Even though each Restart interval involves several function evaluations, the overall FID-NFE trade-off in Restart remains superior to previous methods. This is primarily because Restart allows for a reduced NFE during the main backward process.\n\n**Q5: Why does the total error remain the same using the ODE sampler when NFE>20 (Fig.2)? What are the results of three samplers with NFE<20? Is it consistent with the trend in Figure 1(b)**\n\nA: Fig.2 plots the Pareto frontier of total Error versus NFE. For ODE, we have conducted experiments with NFE 20,40,80,160,320,640, with total error 0.89, 0.90, 0.90, 0.90, 0.90, 0.90 respectively. This reveals that the Pareto front stabilizes at an NFE of 20 with an error of 0.89. A larger NFE didn’t reduce the error since the discretization error is already small. \n\nWe conducted additional experiments (please see Fig.1(a) in the rebuttal PDF) and verified that the trend is consistent when NFE is less than 20. We will include the result in the updated draft.\n\n**Q6: Why not compare to DPM-solver? Why not apply Restart to DPM-solver to see if there are consistent benefits?**\n\nA: Thanks for the suggestions. As suggested by the reviewer, we compare Restart with DPM-Solver. In order to further accelerate Restart, we also use DPM-Solver in the main/Restart backward processes of Restart. We’ve included the FID versus NFE curves in Fig.1(a)  in the rebuttal PDF in the “Summary of Updates” comment above. The results show 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 in the paper. Surprisingly, when paired with the DPM-Solver, Restart achieves an FID score of 2.11 on VP [1] when NFE=30, which is significantly lower than any previous numbers (even lower than the SDE sampler with an NFE $\\ge 1000$ in [1]), and make VP model on par with the performance with more advanced models (such as EDM in [2]). We will include these results in our updated draft.\n\n*[1]  Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. ICLR 2021.*\n\n*​​[2] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the Design Space of Diffusion-Based Generative Models. NeurIPS 2022.*\n\n\n"}, {"review_id": "TGgZqjoiWg", "reviewer": "Reviewer_kn2f", "summary": "The 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.", "questions": "I 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.", "limitations": "The authors adequately addressed the limitations of their work.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The 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.", "weaknesses": "The 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\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review and thoughtful feedback. Below we address specific questions.\n\n\n**Q1: The error propagation of diffusion models has been studied before ... What are the differences in the theoretical analysis compared to prior results known for error-propagation in diffusion models.**\n\n\nA: Thank you for pointing out these related works. We will add a more detailed comparison to the draft. Briefly, our result differs from the above-mentioned papers in the following key way: [1,2] aim to control the discretization error going from $T \\to 0$ during the backward process.  [1] does this using Girsanov’s theorem, and [2] does this using a more involved analysis as one cannot apply Girsanov in the absence of diffusion noise. In contrast, our analysis focuses on how Restart iterations over the sub-interval $[t_{min} , t_{max}]$ reduces the accumulated discretization error from $[t_{max},T]$. There are some similarities with [3], notably they also have a corrector stage where they run overdamped LMC. One important difference from [3] is that our Restart iterations are not interleaved with ODE integration.\n\nA remarkable flexibility of our analysis is that **we need to make very little assumption about the accumulated error up to $t_{max}$**. Consequently, one can adopt a wide range of integration techniques from [t_{max}, T], and apply Restart at the end (over $[t_{max}, t_{min}]$) to reduce the accumulated integration error between $[t_{max}, T]$. In contrast, the analysis in [1,2,3] is focused on specific ODE/SDE integration algorithms.\n\nWe also note that [3] is a concurrent work and appeared on Arxiv after our submission to NeurIPS. We will include these discussions in our revised version.\n\n\n**Q2: 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\"(DEIS) and also some of the samplers used in the EDM paper. It would be interesting to compare theoretically and experimentally to these samplers.**\n\nA: We would like to note that the samplers highlighted by the reviewer are encompassed within the either stochastic or ODE samplers we've discussed, and we've already made some comparisons in our paper. Specifically, we've compared the ODE sampler (Heun) in EDM and the recommended stochastic sampler in EDM (Improved SDE). Please refer to Fig.3, Table.1, Fig.8, and Fig.10 in the draft for details.\nWe appreciate the suggestion to compare faster ODEs with exponential integrators. We've incorporated DDIM in the draft, which serves as the first-order counterpart to DEIS. Our new empirical results demonstrate that Restart can improve over DPM-Solver, which is also a higher-order ODE solver that utilizes exponential integrators. The subsequent question's response provides more detailed results.\n\n**Q3: Another concern I have is that the evaluation is only as thorough as it should have been and it is only done for relatively high NFEs. How is the performance for lower NFEs, e.g. 20 sampling steps?**\n\nA: Thanks for the suggestion. We employ the recommended Heun’s 2nd order ODE sampler [4] in the main backward process of Restart in our draft. Since Heun’s 2nd order ODE is not targeting the low NFE regime, we did not emphasize the performance of such regime. To validate the effectiveness of the Restart sampler in the low NFE regime, we use the faster ODE sampler in [5] (DPM-Solver) in the backward processes of Restart. We’ve included the FID versus NFE curves in Fig.1(a)  in the rebuttal PDF in the “Summary of Updates” comment above. The results show 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. We will include these results in our updated draft.\n\n\n**Q4: 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\nA: Thanks for the suggestions. Following the suggestion, we evaluated the Restart sampler on smaller NFEs, particularly when paired with the DPM-Solver, as previously discussed. We've also highlighted the efficiency of the Restart sampler on the FFHQ-1024 dataset, as depicted in Fig.2 of the rebuttal PDF. We will include more datasets in the updated draft.\n\n*[1] Sampling is as easy as learning the score, Chen et al., ICLR 2023*\n\n*[2] Restoration-degradation beyond linear diffusions: A non-asymptotic analysis for DDIM-type samplers, Chen et al., ICML 2023*\n\n*[3] The probability flow ODE is provably fast, Chen et al., arXiv 2305.11798*\n\n*​​[4] Elucidating the Design Space of Diffusion-Based Generative Models, Karras et al., NeurIPS 2022.*\n\n*[5]  Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps, Lu et al. NeurIPS 2022.*\n"}, {"review_id": "G0yr7jkDn8", "reviewer": "Reviewer_csQs", "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.", "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.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "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.\n\n2. 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.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review and thoughtful feedback. Below we address specific questions.\n\n**Q1:  The proposed method relies on several hyperparameters (e.g. $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 applications. 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?**\n\nA: Thanks for the question. Our choice of hyperparameter is partially motivated by theory. For instance, for a small $t_{min}$, we usually pick the parameter $t_{max} \\approx B$, where $B$ is the radius of the dataset; this ensures that the contraction factor $\\lambda$ in Theorem 4 is sufficiently large. \n\nNote that our theoretical results contain a number of Lipschitz constants, which are difficult to estimate in practice; thus our theoretical upper bounds help us decide high-level scaling but do not accurately prescribe a precise parameter choice. As another example, a larger accumulated error requires bigger Restart intervals and more Restart iterations, thus we use a larger $K$ and Restart interval for the weaker VP model [1] on CIFAR-10, compared to the EDM model [2].\n\n*[1] Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. ICLR 2021.*\n\n*[2] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the Design Space of Diffusion-Based Generative Models. NeurIPS 2022*\n\n​​\n\n"}, {"review_id": "LGU86eoXGT", "reviewer": "Reviewer_kog8", "summary": "This 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.", "questions": "**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)?", "limitations": "Discussed in Section 6.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written.\n- The proposed method shows better performance than EDM at moderate NFE regions.", "weaknesses": "I 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review and thoughtful feedback. Below we address specific questions.\n\n**Q1: For Theorem 1, if we set $[t_{min},t_{max}]=[0,T]$, the terms for contracted errors TV(p_{T}^{ODE_θ},pT) and TV(p_{T}^{SDE_θ},pT) vanish because p_{T}^{ODE_θ}, p_{T}^{SDE_θ}, and p_T are all identically Gaussian distributions. Then, we end up with terms depending on δ, ϵ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. 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\nA:  Thank you for raising this question. First, we agree that if we set $[t_{min},t_{max}]=[0,T]$, Restart is the same as ODE and has no theoretical advantage. **However, we emphasize that the utility of restart is to reduce the accumulated error between $[t_{max}, T]$.** The interesting point is how this accumulated error is corrected during \"Restart iterations\". Therefore, setting $t_{min}=0$ and $t_{max}=T$ is a mis-application of Restart or SDE, as the accumulated error = 0 by definition. \n\nThe main goal of Theorem 1 and 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. We will clarify this point in our updated draft.\n\n**Q2: How does Restart perform in the small NFE regime (NFE ≤30)? How does Restart compare to recent fast samplers in the small NFE regime?**\n\nA: Thanks for the suggestions. We would like to highlight that the Restart sampler is compatible with fast ODE samplers, as they can be integrated in the deterministic backward process of Restart. As suggested by the reviewer, we compare Restart with DPM-Solver [3], a commonly-used fast ODE solver. In order to further accelerate Restart, we also use DPM-Solver in the main/Restart backward processes of Restart. We’ve included the FID versus NFE curves in Fig.1(a)  in the rebuttal PDF in the “Summary of Updates” comment above. The results show that the Restart consistently outperforms the DPM-Solver with a NFE ranging from 16 to 36. This demonstrates Restart's capability to excel over ODE samplers, even in the small NFE regime. Surprisingly, when paired with the DPM-Solver, Restart achieves a FID score of 2.11 on VP [1] when NFE=30, which is significantly lower than any previous numbers (even lower than the SDE sampler with a NFE $\\ge 1000$ in [1]), and make VP model on par with the performance with more advanced models (such as EDM in [2]). We will include these results in our updated draft.\n\n**Q3: 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\nA: Thanks for the suggestion. We use the pre-trained FFHQ-1024 checkpoints in the code base pointed out by the reviewer (https://github.com/yang-song/score_sde), which is based on the Score-SDE-VE model with NCSN++ architecture. For stochastic baseline, we compare with the default Predictor-Corrector (PC) sampler [1] in the code base instead of Improved SDE, due to its need for time-consuming tuning of the hyper-parameters ($S_{tmax}, S_{tmin}, S_{noise}, S_{churn}$) on FFHQ. For ODE baseline, we compare with Heun as suggested by the reviewer. We set the NFE to 300 for all samplers.\n\n\nSince it’s prohibitively expensive to compute the FID score on this dataset, we qualitatively assess the visual quality from different samplers. We’ve included these images in Fig.2 in the rebuttal PDF in the “Summary of Updates” comment above. We observe that the Restart sampler produces notably superior image quality compared to other baselines. Heun sampler fails to generate clean images, and there is noticeable noise and artifacts of images generated by PC sampler compared to Restart. It indicates that the Restart sampler can successfully scale to 1024 resolution, and better balance the speed and quality in comparison to both ODE and SDE samplers. We will include the experiment in our revised version.\n\n\n*[1] Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. ICLR 2021.*\n\n*​​[2] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the Design Space of Diffusion-Based Generative Models. NeurIPS 2022.*\n\n*[3] Lu, Cheng, et al. \"Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps.\" Advances in Neural Information Processing Systems 35 (2022): 5775-5787*\n\n"}]

Author rebuttal:
# Summary of Updates

We would like to thank all reviewers for their constructive feedback. We have revised our draft according to all the valuable comments. Below we summarize updates in the revised version. We also include all the new figures in the PDF files attached.

## 1. More experiments

In response to Reviewer YGP7, kn2f and kog8, we have compared Restart with DPM-Solver and incorporated experiments in the low NFE regime. We’ve included the FID versus NFE curves in Fig.1(a)  in the rebuttal PDF. The results show that the Restart consistently outperforms the DPM-Solver with an NFE ranging from 16 to 36. As recommended by Reviewer kog8, we qualitatively validate the effectiveness of Restart on FFHQ-1024 dataset (Fig.2 in the rebuttal PDF). Additionally, we have included smaller NFEs in the study of total error versus NFE (Fig.1(b) in the rebuttal PDF), and additional sensitivity analysis of hyperparameters (Fig.3 in the rebuttal PDF), as suggested by Reviewer YGP7.

## 2. Discussion on related works / Clarification on theorems

As suggested by Reviewer k2nf, we’ve added the discussion of the difference between Restart and prior or concurrent works. In response to Reviewer kog8, we have also provided clarifications on our theorems, emphasizing their objectives in examining the behavior of various samplers with accumulated errors.


Meta-review: Based on my assessment of the reviews, I recommend accepting this paper. The work proposes a novel sampling algorithm for diffusion models that aims to balance speed and quality. The core idea of adding noise and restarting sampling shows promise in improving over prior methods.

The reviewers rate the paper mostly as Weak Accept or Accept, with multiple increasing scores after the rebuttal. They find the theory and experiments to be technically solid. The qualitative gains on high-resolution image synthesis and comparisons to recent fast samplers address initial concerns. The discussion helped clarify ambiguities in the theorems.

While further experiments on scalability would be beneficial, the paper makes a meaningful contribution in improving sampling for generative diffusion models. The combination of new algorithm design, analysis, and empirical gains over strong baselines highlights the potential of this approach. I recommend acceptance to share this work with the broader community through publication. Additional experiments can be incorporated into the camera ready version.
Decision: Accept (poster)