SlowGuess/NeurIPS_2023_OCR huggingface dataset:
To bound the contracted error, we construct a careful coupling process between two individual trajectories sampled from  $p_{t_{\min}}^{\text{Restart}_{\theta}(i)}$  and  $q_{t_{\min}}^{\text{Restart}_{\theta}(i)}$ ,  $i = 0, \dots, K - 1$ . Before these two trajectories converge, the Gaussian noise added in each Restart iteration is chosen to maximize the probability of the two trajectories mapping to an identical point, thereby maximizing the mixing rate in TV. After converging, the two processes evolve under the same Gaussian noise, and will stay converged as their drifts are the same. Lastly, we convert the TV bound to  $W_1$  bound by multiplying  $B$ . The bound on the additional sampling error echoes the ODE analysis in Theorem 1: since the noise-injection and ODE-simulation stages are separate, we do not incur the higher discretization error of SDE.
We defer the formal version and proof of Theorem 2 to Appendix A.1. The first term in RHS bounds the contraction on the initial error at time  $t_{\mathrm{max}}$  and the second term reflects the additional sampling error of ODE accumulated across repeated Restart iterations. Comparing the Wasserstein upper bound of SDE and ODE in Theorem 1, we make the following three observations: (1) Each Restart iteration has a smaller contraction factor  $1 - \lambda$  compared to the one in SDE, since Restart separates the large additive noise (Eq. (3)) from the ODE (Eq. (4)). (2) Restart backward process (Eq. (4)) has the same order of discretization error  $O(\delta)$  as the ODE, compared to  $O(\sqrt{\delta})$  in SDE. Hence, the Restart allows for small NFE due to ODE-level discretization error. (3) The contracted error further diminishes exponentially with the number of repetitions  $K$  though the additional error increases linearly with  $K$ . It suggests that there is a sweet spot of  $K$  that strikes a balance between reducing the initial error and increasing additional sampling error. Ideally, one should pick a larger  $K$  when the initial error at time  $t_{\mathrm{max}}$  greatly outweighs the incurred error in the repetitive backward process from  $t_{\mathrm{max}}$  to  $t_{\mathrm{min}}$ . We provide empirical evidences in Sec 5.2.
While Theorem 1 and Theorem 2 compare the upper bounds on errors of different methods, we provide empirical validation in Section 5.1 by directly calculating these errors, showing that the Restart algorithm indeed yields a smaller total error due to its superior contraction effects. The main goal of Theorem 1 and Theorem 2 is to study how the already accumulated error changes using different samplers, and to understand their ability to self-correct the error by stochasticity. In essence, these theorems differentiate samplers based on their performance post-error accumulation. For example, by tracking the change of accumulated error, Theorem 1 shed light on the distinct "winning regions" of ODE and SDE: ODE samplers have smaller discretization error and hence excel at the small NFE regime. In contrast, SDE performs better in large NFE regime where the discretization error is negligible and its capacity to contract accumulated errors comes to the fore.
4.3 Practical considerations
The Restart algorithm offers several degrees of freedom, including the time interval  $[t_{\mathrm{min}}, t_{\mathrm{max}}]$  and the number of restart iterations  $K$ . Here we provide a general recipe of parameter selection for practitioners, taking into account factors such as the complexity of the generative modeling tasks and the capacity of the network. Additionally, we discuss a stratified, multi-level Restart approach that further aids in reducing simulation errors along the whole trajectories for more challenging tasks.
Where to Restart? Theorem 2 shows that the Restart algorithm effectively reduces the accumulated error at time  $t_{\mathrm{max}}$  by a contraction factor in the Wasserstein upper bound. These theoretical findings inspire us to position the Restart interval  $[t_{\mathrm{min}}, t_{\mathrm{max}}]$  towards the end of the main backward process, where the accumulated error is more substantial. In addition, our empirical observations suggest that a larger time interval  $t_{\mathrm{max}} - t_{\mathrm{min}}$  is more beneficial for weaker/smaller architectures or more challenging datasets. Even though a larger time interval increases the additional sampling error, the benefits of the contraction significantly outweighs the downside, consistent with our theoretical predictions. We leave the development of principled approaches for optimal time interval selection for future works.
Multi-level Restart For challenging tasks that yield significant approximation errors, the backward trajectories may diverge substantially from the ground truth even at early stage. To prevent the ODE simulation from quickly deviating from the true trajectory, we propose implementing multiple Restart intervals in the backward process, alongside the interval placed towards the end. Empirically, we observe that a 1-level Restart is sufficient for CIFAR-10, while for more challenging datasets such as ImageNet [5], a multi-level Restart results in enhanced performance [5].
In Sec 5.1, we first empirically verify the theoretical analysis relating to the Wasserstein upper bounds. We then evaluate the performance of different sampling algorithms on standard image generation benchmarks, including CIFAR-10 [14] and ImageNet  $64 \times 64$  [5] in Sec 5.2. Lastly, we employ Restart on text-to-image generation, using Stable Diffusion model [19] pre-trained on LAION-5B [21] with resolution  $512 \times 512$ , in Sec 5.3.
5.1 Additional sampling error versus contracted error
Our proposed Restart sampling algorithm demonstrates a higher contraction effect and smaller addition sampling error compared to SDE, according to Theorem 1 and Theorem 2. Although our
(b)
Figure 2: Additional sampling error versus (a) contracted error, where the Pareto frontier is plotted and (b) total error, where the scatter plot is 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 a synthetic experiment.
Setup We construct a 20-dimensional dataset with 2000 points sampled from a Gaussian mixture, and train a four-layer MLP to approximate the score field  $\nabla_{x}\log p_{t}$ . We implement the ODE, SDE, and Restart methods within a predefined time range of  $[t_{\mathrm{min}}, t_{\mathrm{max}}] = [1.0, 1.5]$ , where the process outside this range is conducted via the first-order ODE. To compute various error types, we define the distributions generated by three methods as outlined in the proof of Theorem 2 and directly gauge the errors at end of simulation  $t = 0$  instead of  $t = t_{\mathrm{min}}$ : (1) the generated distribution as  $p_{0}^{\mathrm{Sampler}}$ , where Sampler  $\in \{\mathrm{ODE}_{\theta}, \mathrm{SDE}_{\theta}, \mathrm{Restart}_{\theta}(K)\}$ ; (2) an auxiliary distribution  $q_{0}^{\mathrm{Sampler}}$  initiating from true distribution  $p_{t_{\mathrm{max}}}$  at time  $t_{\mathrm{max}}$ . The only difference between  $p_{0}^{\mathrm{Sampler}}$  and  $q_{0}^{\mathrm{Sampler}}$  is their initial distribution at  $t_{\mathrm{max}}$  ( $p_{t_{\mathrm{max}}}^{\mathrm{ODE}_{\theta}}$  versus  $p_{t_{\mathrm{max}}}$ ); and (3) the true data distribution  $p_{0}$ . In line with Theorem 2, we use Wasserstein-1 distance  $W_{1}(p_{0}^{\mathrm{Sampler}}, q_{0}^{\mathrm{Sampler}}) / W_{1}(q_{0}^{\mathrm{Sampler}}, p_{0})$  to measure the contracted error / additional sampling error, respectively. Ultimately, the total error corresponds to  $W_{1}(p_{0}^{\mathrm{Sampler}}, p_{0})$ . Detailed information about dataset, metric and model can be found in the Appendix C.5.
Results In our experiment, we adjust the parameters for all three processes and calculate the total, contracted, and additional sampling errors across all parameter settings. Figure 2(a) depicts the Pareto frontier of additional sampling error versus contracted error. We can see that Restart consistently achieves lower contracted error for a given level of additional sampling error, compared to both the ODE and SDE methods, as predicted by theory. In Figure 2(b), we observe that the Restart method obtains a smaller total error within the additional sampling error range of [0.8, 0.85]. During this range, Restart also displays a strictly reduced contracted error, as illustrated in Figure 2(a). This aligns with our theoretical analysis, suggesting that the Restart method offers a smaller total error due to its enhanced contraction effects. From Figure 2(c), Restart also strikes an better balance between efficiency and quality, as it achieves a lower total error at a given NFE.
5.2 Experiments on standard benchmarks
(a) CIFAR-10, VP
Figure 3: FID versus NFE on (a) unconditional generation on CIFAR-10 with VP; (b) class-conditional generation on ImageNet with EDM.
(b) ImageNet  $64 \times 64$ , EDM
To evaluate the sample quality and inference speed, we report the FID score [7] (lower is better) on 50K samplers and the number of function evaluations (NFE). We borrow the pretrained VP/EDM/PFGM++
Figure 4: CIFAR-10, VP, in the low NFE regime. Restart consistently outperforms the DPM-Solver with an NFE ranging from 16 to 36.
models on CIFAR-10 or ImageNet  $64 \times 64$  from [13, 28]. We also use the EDM discretization scheme [13] (see Appendix B.1 for details) during sampling.
For the proposed Restart sampler, the hyperparameters include the number of steps in the main/Restart backward processes, the number of Restart iteration  $K$ , as well as the time interval  $[t_{\mathrm{min}}, t_{\mathrm{max}}]$ . We pick the  $t_{\mathrm{min}}$  and  $t_{\mathrm{max}}$  from the list of time steps in EDM discretization scheme with a number of steps 18. For example, for CIFAR-10 (VP) with NFE=75, we choose  $t_{\mathrm{min}} = 0.06$ ,  $t_{\mathrm{max}} = 0.30$ ,  $K = 10$ , where  $0.30 / 0.06$  is the  $12^{\mathrm{th}} / 14^{\mathrm{th}}$  time step in the EDM scheme. We also adopt EDM scheme for the Restart backward process in  $[t_{\mathrm{min}}, t_{\mathrm{max}}]$ . In addition, we apply the multi-level Restart strategy (Sec 4.3) to mitigate the error at early time steps for the more challenging ImageNet  $64 \times 64$ . We provide the detailed Restart configurations in Appendix C.2.
For SDE, we compare with the previously best-performing stochastic samplers proposed by [13] (Improved SDE). We use their optimal hyperparameters for each dataset. We also report the FID scores of the adaptive SDE [12] (Gonna Go Fast) on CIFAR-10 (VP). Since the vanilla reverse-diffusion SDE [23] has a significantly higher FID score, we omit its results from the main charts and defer them to Appendix D. For ODE samplers, we compare with the Heun's  $2^{\text{nd}}$  order method [2] (Heun), which arguably provides an excellent trade-off between discretization errors and NFE [13]. To ensure a fair comparison, we use Heun's method as the sampler in the main/Restart backward processes in Restart.
We report the FID score versus NFE in Figure 3(a) and Table 1 on CIFAR-10, and Figure 3(b) on ImageNet  $64 \times 64$  with EDM. Our main findings are: (1) Restart outperforms other SDE or ODE samplers in balancing quality and speed, across datasets and models. As demonstrated in the figures, Restart achieves a 10-fold / 2-fold acceleration compared to previous best SDE results on CIFAR-10 (VP) / ImageNet  $64 \times 64$  (EDM) at the same FID score. In comparison to ODE sampler (Heun), Restart obtains a better FID score, with the gap increasing significantly with NFE. (2) For stronger models such as EDM and PFGM++, Restart further improve over the ODE baseline on CIFAR-10. In contrast, the Improved SDE negatively impacts performance of EDM, as also observed in [13]. It suggests that Restart incorporates stochasticity more effectively. (3) Restart establishes new state-of-the-art FID scores for UNet architectures without additional training. In particular, Restart achieves FID scores of 1.36 on class-cond. ImageNet  $64 \times 64$  with EDM, and 1.88 on uncond. CIFAR-10 with PFGM++.
To further validate that Restart can be applied in low NFE regime, we show that one can employ faster ODE solvers such as the DPM-solver [16] to further accelerate Restart. Fig. 4 shows that the Restart consistently outperforms the DPM-Solver with an NFE ranging from 16 to 36. This demonstrates Restart's capability to excel over ODE samplers, even in the small NFE regime. It also suggests that Restart can consistently improve other ODE samplers, not limited to the DDIM, Heun. Surprisingly, when paired with the DPM-Solver, Restart achieves an FID score of 2.11 on VP setting when NFE is 30, which is significantly
Table 1: Uncond. CIFAR-10 with EDM and PFGM++
NFE	FID
EDM-VP [13]
ODE (Heun)	63	1.97
35	1.97
Improved SDE	63	2.27
35	2.45
Restart	43	1.90
PFGM++ [28]
ODE (Heun)	63	1.91
35	1.91
Restart	43	1.88
Figure 5: FID score with a varying number of Restart iterations  $K$
lower than any previous numbers (even lower than the SDE sampler with an NFE greater than 1000 in [23]), and make VP model on par with the performance with more advanced models (such as EDM).
Theorem 4 shows that each Restart iteration reduces the contracted errors while increasing the additional sampling errors in the backward process. In Fig. 5, we explore the choice of the number of Restart iterations  $K$  on CIFAR-10. We find that FID score initially improves and later worsens with increasing iterations  $K$ , with a smaller turning point for stronger EDM model. This supports the theoretical analysis that sampling errors will eventually outweigh the contraction benefits as  $K$  increases, and EDM only permits fewer Restart iterations due to smaller accumulated errors. It also suggests that, as a rule of thumb, we should apply greater Restart strength (e.g., larger  $K$ ) for weaker or smaller architectures and vice versa.
5.3 Experiments on large-scale text-to-image model
(a) FID versus CLIP score
(b) FID versus Aesthetic score
Figure 6: FID score versus (a) CLIP ViT-g/14 score and (b) Aesthetic score for text-to-image generation at  $512 \times 512$  resolution, using Stable Diffusion v1.5 with a varying classifier-free guidance weight  $w = 2,3,5,8$ .
We further apply Restart to the text-to-image Stable Diffusion v1.5  $^{2}$  pre-trained on LAION-5B [21] at a resolution of  $512 \times 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, as rated by the LAION-Aesthetics Predictor V2 [24]. Following [17], we compute all evaluation metrics using 5K captions randomly sampled from the validation set and plot the trade-off curves between CLIP/Aesthetic scores and FID score, with the classifier-free guidance weight  $w$  in  $\{2,3,5,8\}$ .
We compare with commonly used ODE sampler DDIM [22] and the stochastic sampler DDPM [9]. For Restart, we adopt the DDIM solver with 30 steps in the main backward process, and Heun in the Restart backward process, as we empirically find that Heun performs better than DDIM in the Restart. In addition, we select different sets of the hyperparameters for each guidance weight. For instance, when  $w = 8$ , we use  $[t_{\mathrm{min}}, t_{\mathrm{max}}] = [0.1, 2]$ ,  $K = 2$  and 10 steps in Restart backward process. We defer the detailed Restart configuration to Appendix C.2, and the results of Heun to Appendix D.1.
As illustrated in Fig. 6(a) and Fig. 6(b), Restart achieves better FID scores in most cases, given the same CLIP/Aesthetic scores, using only 132 function evaluations (i.e., 66 sampling steps). Remarkably, Restart achieves substantially lower FID scores than other samplers when CLIP/Aesthetic scores are high (i.e., with larger  $w$  values). Conversely, Restart generally obtains a better text-image alignment/visual quality given the same FID. We also observe that DDPM generally obtains comparable performance with Restart in FID score when CLIP/Aesthetic scores are low, with Restart being more time-efficient. These findings suggest that Restart balances diversity (FID score) against text-image alignment (CLIP score) or visual quality (Aesthetic score) more effectively than previous samplers.
In Fig. 7, we visualize the images generated by Restart, DDIM and DDPM with  $w = 8$ . Compared to DDIM, the Restart generates images with superior details (e.g., the rendition of duck legs by DDIM is less accurate) and visual quality. Compared to DDPM, Restart yields more photo-realistic images (e.g., the astronaut). We provide extended of text-to-image generated samples in Appendix E.
$^{2}$ https://huggingface.co/runwayml/stable-diffusion-v1-5
(a) Restart (Steps=66)
(b) DDIM (Steps=100)
Figure 7: Visualization of generated images with classifier-free guidance weight  $w = 8$ , using four text prompts ("A photo of an astronaut riding a horse on mars.", "A raccoon playing table tennis", "Intricate origami of a fox in a snowy forest" and "A transparent sculpture of a duck made out of glass") and the same random seeds.
(c) DDPM (Steps=100)
6 Conclusion and future direction
In this paper, we introduce the Restart sampling for generative processes involving differential equations, such as diffusion models and PFGMs. By interweaving a forward process that adds a significant amount of noise with a corresponding backward ODE, Restart harnesses and even enhances the individual advantages of both ODE and SDE. Theoretically, Restart provides greater contraction effects of stochasticity while maintaining ODE-level discretization error. Empirically, Restart achieves a superior balance between quality and time, and improves the text-image alignment/visual quality and diversity trade-off in the text-to-image Stable Diffusion models.
A current limitation of the Restart algorithm is the absence of a principled way for hyperparameters selection, including the number of iterations  $K$  and the time interval  $[t_{\mathrm{min}}, t_{\mathrm{max}}]$ . At present, we adjust these parameters based on the heuristic that weaker/smaller models, or more challenging tasks, necessitate a stronger Restart strength. In the future direction, we anticipate developing a more principled approach to automating the selection of optimal hyperparameters for Restart based on the error analysis of models, in order to fully unleash the potential of the Restart framework.
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-2112665 (TILOS AI Research Institute).
[19] Robin Rombach, A. Blattmann, Dominik Lorenz, Patrick Esser, and Björn 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 Sahara, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L. Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, Seyedeh Sara Mahdavi, Raphael Gontijo Lopes, Tim Salimans, Jonathan Ho, David J. Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffusion models with deep language understanding. ArXiv, abs/2205.11487, 2022.
[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 Kaczmarczyk, and Jenia Jitsev. Laion-5b: An open large-scale dataset for training next generation image-text models. ArXiv, abs/2210.08402, 2022.
[22] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. ArXiv, abs/2010.02502, 2020.
[23] Yang Song, Jascha Narain Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. ArXiv, abs/2011.13456, 2020.
[24] LAION-AI Team. Laion-aesthetics predictor v2. https://github.com/ christophschuhmann/improved-aesthetic-predictor, 2022.
[25] Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Computation, 23:1661-1674, 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 F. Milles, Basile I. M. Wicky, Nikita Hanikel, Samuel J. Pellock, Alexis Courbet, William Sheffler, Jue Wang, Preetham Venkatesh, Isaac Sappington, Susana Vázquez Torres, Anna Lauko, Valentin De Bortoli, Emile Mathieu, Regina Barzilay, T. Jaakkola, Frank DiMaio, Minkyung Baek, and David Baker. Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models. bioRxiv, 2022.
[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.
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(\ldots) to denote the backwards ODE under exact score  $\nabla \log p_t(x)$ . More specifically, given any  $x \in \mathbb{R}^d$  and  $s > r > 0$ , let  $x_t$  denote the solution to the following ODE:
$$
d x _ {t} = - t \nabla \log p _ {t} \left(x _ {t}\right) d t. \tag {5}
$$
$\mathrm{ODE}(x,s\to r)$  is defined as "the value of  $x_{r}$  when initialized at  $x_{s} = x$ ". It will also be useful to consider a "time-discretized ODE with drift  $ts_{\theta}(x,t)$ ": let  $\delta$  denote the discretization step size and let  $k$  denote any integer. Let  $\delta$  denote a step size, let  $\overline{x}_t$  denote the solution to
$$
d \bar {x} _ {t} = - t s _ {\theta} \left(x _ {k \delta}, k \delta\right) d t, \tag {6}
$$
where for any  $t$ ,  $k$  is the unique integer such that  $t \in ((k - 1)\delta, k\delta]$ . We verify that the dynamics of Eq. (6) is equivalent to the following discrete-time dynamics for  $t = k\delta$ ,  $k \in \mathbb{Z}$ :
$$
\overline {{x}} _ {(k - 1) \delta} = \overline {{x}} _ {k \delta} - \frac {1}{2} \left(\left((k - 1) \delta\right) ^ {2} - (k \delta) ^ {2}\right) s _ {\theta} \left(x _ {k \delta}, k \delta\right).
$$
We similarly denote the value of  $\overline{x}_r$  when initialized at  $\overline{x}_s = x$  as  $\mathrm{ODE}_{\theta}(x,s\to r)$ . Analogously, we let  $\mathrm{SDE}(x,s\rightarrow r)$  and  $\mathrm{SDE}_{\theta}(x,s\rightarrow r)$  denote solutions to
$$
d y _ {t} = - 2 t \nabla \log p _ {t} (y _ {t}) d t + \sqrt {2 t} d B _ {t}
$$
$$
d \bar {y} _ {t} = - 2 t s _ {\theta} (\bar {y} _ {t}, t) d t + \sqrt {2 t} d B _ {t}
$$
respectively. Finally, we will define the  $\mathrm{Restart}_{\theta}$  process as follows:
(Restart  $\theta$  forward process)  $x_{t_{\mathrm{max}}}^{i + 1} = x_{t_{\mathrm{min}}}^{i} + \varepsilon_{t_{\mathrm{min}}\to t_{\mathrm{max}}}^{i}$
(  $\mathrm{Restart}_{\theta}$  backward process)  $x_{t_{\min}}^{i + 1} = \mathrm{ODE}_{\theta}(x_{t_{\max}}^{i + 1},t_{\max}\to t_{\min}),$  (7)
where  $\varepsilon_{t_{\min}\to t_{\max}}^i\sim \mathcal{N}\left(\mathbf{0},\left(t_{\max}^2 -t_{\min}^2\right)\mathbf{I}\right)$ . We use  $\operatorname{Restart}_{\theta}(x,K)$  to denote  $x_{t_{\min}}^{K}$  in the above processes, initialized at  $x_{t_{\min}}^{0} = x$ . In various theorems, we will refer to a function  $Q(r):\mathbb{R}^{+}\rightarrow [0,1 / 2)$ , defined as the Gaussian tail probability  $Q(r) = Pr(a\geq r)$  for  $a\sim \mathcal{N}(0,1)$ .
A.1 Main Result