SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yilun Xu*
ylxu@mit.edu
Mingyang Deng*
dengm@mit.edu
chengx@mit.edu
Yonglong Tian
yonglong@google.com
Ziming Liu
zmliu@mit.edu
Tommi Jaakkola
tommi@csail.mit.edu
Generative processes that involve solving differential equations, such as diffusion models, frequently necessitate balancing speed and quality. ODE-based samplers are fast but plateau in performance while SDE-based samplers deliver higher sample quality at the cost of increased sampling time. We attribute this difference to sampling errors: ODE-samplers involve smaller discretization errors while stochasticity in SDE contracts accumulated errors. Based on these findings, we propose a novel sampling algorithm called Restart in order to better balance discretization errors and contraction. The sampling method alternates between adding substantial noise in additional forward steps and strictly following a backward ODE. Empirically, Restart sampler surpasses previous SDE and ODE samplers in both speed and accuracy. Restart not only outperforms the previous best SDE results, but also accelerates the sampling speed by 10-fold / 2-fold on CIFAR-10 / ImageNet  $64 \times 64$ . In addition, it attains significantly better sample quality than ODE samplers within comparable sampling times. Moreover, Restart better balances text-image alignment/visual quality versus diversity than previous samplers in the large-scale text-to-image Stable Diffusion model pre-trained on LAION  $512 \times 512$ . Code is available at https://github.com/Newbeeer/diffusion_restart_sampling
Deep generative models based on differential equations, such as diffusion models and Poisson flow generative models, have emerged as powerful tools for modeling high-dimensional data, from image synthesis [23, 9, 13, 27, 28] to biological data [10, 26]. These models use iterative backward processes that gradually transform a simple distribution (e.g., Gaussian in diffusion models) into a complex data distribution by solving a differential equations. The associated vector fields (or drifts) driving the evolution of the differential equations are predicted by neural networks. The resulting sample quality can be often improved by enhanced simulation techniques but at the cost of longer sampling times.
Prior samplers for simulating these backward processes can be categorized into two groups: ODE-samplers whose evolution beyond the initial randomization is deterministic, and SDE-samplers where the generation trajectories are stochastic. Several works [23, 12, 13] show that these samplers demonstrate their advantages in different regimes, as depicted in Fig. 1(b). ODE solvers [22, 16, 13] result in smaller discretization errors, allowing for decent sample quality even with larger step sizes (i.e., fewer number of function evaluations (NFE)). However, their generation quality plateaus rapidly. In contrast, SDE achieves better quality in the large NFE regime, albeit at the expense of increased sampling time. To better understand these differences, we theoretically analyze SDE performance: the
(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 is large (small discretization steps), explaining the SDE advantage in this regime. Intuitively, the stochastic nature of SDE helps "forget" accumulated errors from previous time steps.
Inspired by these findings, we propose a novel sampling algorithm called Restart, which combines the advantages of ODE and SDE. As illustrated in Fig. 1(a), the Restart sampling algorithm involves  $K$  repetitions of two subroutines in a pre-defined time interval: a Restart forward process that adds a substantial amount of noise, akin to "starting" the original backward process, and a Restart backward process that runs the backward ODE. The Restart algorithm separates the stochasticity from the drifts, and the amount of added noise in the Restart forward process is significantly larger than the small single-step noise interleaving with drifts in previous SDEs such as [23, 13], thus amplifying the contraction effect on accumulated errors. By repeating the forward-backward cycle  $K$  times, the contraction effect introduced in each Restart iteration is further strengthened. The deterministic backward processes allow Restart to reduce discretization errors, thereby enabling step sizes comparable to ODE. To maximize the contraction effects in practice, we typically position the Restart interval towards the end of the simulation, where the accumulated error is larger. Additionally, we apply multiple Restart intervals to further reduce the initial errors in more challenging tasks.
Experimentally, Restart consistently surpasses previous ODE and SDE solvers in both quality and speed over a range of NFEs, datasets, and pre-trained models. Specifically, Restart accelerates the previous best-performing SDEs by  $10 \times$  fewer steps for the same FID score on CIFAR-10 using VP [23] ( $2 \times$  fewer steps on ImageNet  $64 \times 64$  with EDM [13]), and outperforms ODE solvers even in the small NFE regime. When integrated into previous state-of-the-art pre-trained models, Restart further improves performance, achieving FID scores of 1.88 on unconditional CIFAR-10 with PFGM++ [28], and 1.36 on class-conditional ImageNet  $64 \times 64$  with EDM. To the best of our knowledge, these are the best FID scores obtained on commonly used UNet architectures for diffusion models without additional training. We also apply Restart to the practical application of text-to-image Stable Diffusion model [19] pre-trained on LAION  $512 \times 512$ . Restart more effectively balances text-image alignment/visual quality (measured by CLIP/Aesthetic scores) and diversity (measured by FID score) with a varying classifier-free guidance strength, compared to previous samplers.
Our contributions can be summarized as follows: (1) We investigate ODE and SDE solvers and theoretically demonstrate the contraction effect of stochasticity via an upper bound on the Wasserstein distance between generated and data distributions (Sec 3); (2) We introduce the Restart sampling, which better harnesses the contraction effect of stochasticity while allowing for fast sampling. The sampler results in a smaller Wasserstein upper bound (Sec 4); (3) Our experiments are consistent with the theoretical bounds and highlight Restart's superior performance compared to previous samplers on standard benchmarks in terms of both quality and speed. Additionally, Restart improves the trade-off between key metrics on the Stable Diffusion model (Sec 5).
2 Background on Generative Models with Differential Equations
Many recent successful generative models have their origin in physical processes, including diffusion models [9, 23, 13] and Poisson flow generative models [27, 28]. These models involve a forward process that transforms the data distribution into a chosen smooth distribution, and a backward process that iteratively reverses the forward process. For instance, in diffusion models, the forward process is the diffusion process with no learned parameters:
$$
\mathrm {d} x = \sqrt {2 \dot {\sigma} (t) \sigma (t)} \mathrm {d} W _ {t},
$$
where  $\sigma(t)$  is a predefined noise schedule increasing with  $t$ , and  $W_t \in \mathbb{R}^d$  is the standard Wiener process. For simplicity, we omit an additional scaling function for other variants of diffusion models as in EDM [13]. Under this notation, the marginal distribution at time  $t$  is the convolution of data distribution  $p_0 = p_{\mathrm{data}}$  and a Gaussian kernel, i.e.,  $p_t = p_0 * \mathcal{N}(\mathbf{0}, \sigma^2(t) \mathbf{I}_{d \times d})$ . The prior distribution is set to  $\mathcal{N}(\mathbf{0}, \sigma^2(T) \mathbf{I}_{d \times d})$  since  $p_T$  is approximately Gaussian with a sufficiently large  $T$ . Sampling of diffusion models is done via a reverse-time SDE [1] or a marginally-equivalent ODE [23]:
(SDE)  $\mathrm{dx} = -2\dot{\sigma} (t)\sigma (t)\nabla_{x}\log p_{t}(x)dt + \sqrt{2\dot{\sigma} (t)\sigma (t)}\mathrm{d}W_{t}$  (1)
(ODE)  $\mathrm{d}x = -\dot{\sigma} (t)\sigma (t)\nabla_{x}\log p_{t}(x)dt$  (2)
where  $\nabla_{x}\log p_{t}(x)$  in the drift term is the score of intermediate distribution at time  $t$ . W.l.o.g we set  $\sigma (t) = t$  in the remaining text, as in [13]. Both processes progressively recover  $p_0$  from the prior distribution  $p_T$  while sharing the same time-dependent distribution  $p_t$ . In practice, we train a neural network  $s_\theta (x,t)$  to estimate the score field  $\nabla_{x}\log p_{t}(x)$  by minimizing the denoising score-matching loss [25]. We then substitute the score  $\nabla_{x}\log p_{t}(x)$  with  $s_\theta (x,t)$  in the drift term of above backward SDE (Eq. (1))/ODE (Eq. (2)) for sampling.
Recent work inspired by electrostatics has not only challenged but also integrated diffusion models, notably PFGM/PFGM++, enhances performance in both image and antibody generation [27, 28, 10]. They interpret data as electric charges in an augmented space, and the generative processes involve the simulations of differential equations defined by electric field lines. Similar to diffusion models, PFGMs train a neural network to approximate the electric field in the augmented space.
3 Explaining SDE and ODE performance regimes
To sample from the aforementioned generative models, a prevalent approach employs general-purpose numerical solvers to simulate the corresponding differential equations. This includes Euler and Heun's 2nd method [2] for ODEs (e.g., Eq. (2)), and Euler-Maruyama for SDEs (e.g., Eq. (1)). Sampling algorithms typically balance two critical metrics: (1) the quality and diversity of generated samples, often assessed via the Fréchet Inception Distance (FID) between generated distribution and data distribution [7] (lower is better), and (2) the sampling time, measured by the number of function evaluations (NFE). Generally, as the NFE decreases, the FID score tends to deteriorate across all samplers. This is attributed to the increased discretization error caused by using a larger step size in numerical solvers.
However, as illustrated in Fig. 1(b) and observed in previous works on diffusion models [23, 22, 13], the typical pattern of the quality vs time curves behaves differently between the two groups of samplers, ODE and SDE. When employing standard numerical solvers, ODE samplers attain a decent quality with limited NFEs, whereas SDE samplers struggle in the same small NFE regime. However, the performance of ODE samplers quickly reaches a plateau and fails to improve with an increase in NFE, whereas SDE samplers can achieve noticeably better sample quality in the high NFE regime. This dilemma raises an intriguing question: Why do ODE samplers outperform SDE samplers in the small NFE regime, yet fall short in the large NFE regime?
The first part of the question is relatively straightforward to address: given the same order of numerical solvers, simulation of ODE has significantly smaller discretization error compared to the SDE. For example, the first-order Euler method for ODE results in a local error of  $O(\delta^2)$ , whereas the first-order Euler-Maruyama method for SDEs yields a local error of  $O\left(\delta^{\frac{3}{2}}\right)$  (see e.g., Theorem 1 of [4]), where  $\delta$  denotes the step size. As  $O\left(\delta^{\frac{3}{2}}\right) \gg O(\delta^2)$ , ODE simulations exhibit lower sampling errors than SDEs, likely causing the better sample quality with larger step sizes in the small NFE regime.
In the large NFE regime the step size  $\delta$  shrinks and discretization errors become less significant for both ODEs and SDEs. In this regime it is the approximation error — error arising from an
inaccurate estimation of the ground-truth vector field by the neural network  $s_{\theta}$  — starts to dominate the sampling error. We denote the discretized ODE and SDE using the learned field  $s_{\theta}$  as  $\mathrm{ODE}_{\theta}$  and  $\mathrm{SDE}_{\theta}$ , respectively. In the following theorem, we evaluate the total errors from simulating  $\mathrm{ODE}_{\theta}$  and  $\mathrm{SDE}_{\theta}$  within the time interval  $[t_{\min}, t_{\max}] \subset [0, T]$ . This is done via an upper bound on the Wasserstein-1 distance between the generated and data distributions at time  $t_{\min}$ . We characterize the accumulated initial sampling errors up until  $t_{\max}$  by total variation distances. Below we show that the inherent stochasticity of SDEs aids in contracting these initial errors at the cost of larger additional sampling error in  $[t_{\min}, t_{\max}]$ . Consequently, SDE results in a smaller upper bound as the step size  $\delta$  nears 0 (pertaining to the high NFE regime).
Theorem 1 (Informal). Let  $t_{\text{max}}$  be the initial noise level and  $p_t$  denote the true distribution at noise level  $t$ . Let  $p_t^{ODE_\theta}, p_t^{SDE_\theta}$  denote the distributions of simulating  $ODE_\theta$ ,  $SDE_\theta$  respectively. Assume that  $\forall t \in [t_{\text{min}}, t_{\text{max}}}$ ,  $\| x_t \| < B/2$  for any  $x_t$  in the support of  $p_t$ ,  $p_t^{ODE_\theta}$  or  $p_t^{SDE_\theta}$ . Then
$$
\begin{array}{l} W _ {1} \left(p _ {t _ {\min }} ^ {O D E _ {\theta}}, p _ {t _ {\min }}\right) \leq B \cdot T V \left(p _ {t _ {\max }} ^ {O D E _ {\theta}}, p _ {t _ {\max }}\right) + O \left(\delta + \epsilon_ {\text {a p p r o x}}\right) \cdot \left(t _ {\max } - t _ {\min }\right) \\ \underbrace {W _ {1} \left(p _ {t _ {\text {m i n}}} ^ {S D E _ {\theta}} , p _ {t _ {\text {m i n}}}\right)} _ {\text {t o t a l e r r o r}} \leq \underbrace {\left(1 - \lambda e ^ {- U}\right) B \cdot T V \left(p _ {t _ {\text {m a x}}} ^ {S D E _ {\theta}} , p _ {t _ {\text {m a x}}}\right)} _ {\text {u p p e r b o u n d o n c o n t r a c t e d e r r o r}} + \underbrace {O \left(\sqrt {\delta t _ {\text {m a x}}} + \epsilon_ {\text {a p p r o x}}\right) \left(t _ {\text {m a x}} - t _ {\text {m i n}}\right)} _ {\text {u p p e r b o u n d o n a d d i t i o n a l s a m p l i n g e r r o r}} \\ \end{array}
$$
In the above,  $U = BL_{1} / t_{\min} + L_{1}^{2}t_{\max}^{2} / t_{\min}^{2}, \lambda < 1$  is a contraction factor,  $L_{1}$  and  $\epsilon_{\text{approx}}$  are uniform bounds on  $\| ts_{\theta}(x_t,t)\|$  and the approximation error  $\| t\nabla_x\log p_t(x) - ts_{\theta}(x,t)\|$  for all  $x_{t}, t$ , respectively.  $O(\cdot)$  hides polynomial dependency on various Lipschitz constants and dimension.
We defer the formal version and proof of Theorem 1 to Appendix A.1. As shown in the theorem, the upper bound on the total error can be decomposed into upper bounds on the contracted error and additional sampling error.  $TV(p_{t_{\max}}^{\mathrm{ODE}_{\theta}}, p_{t_{\max}})$  and  $TV(p_{t_{\max}}^{\mathrm{SDE}_{\theta}}, p_{t_{\max}})$  correspond to the initial errors accumulated from both approximation and discretization errors during the simulation of the backward process, up until time  $t_{\max}$ . In the context of SDE, this accumulated error undergoes contraction by a factor of  $1 - \lambda e^{-BL_1 / t_{\min} - L_1^2 t_{\max}^2 / t_{\min}^2}$  within  $[t_{\min}, t_{\max}]$ , due to the effect of adding noise. Essentially, the minor additive Gaussian noise in each step can drive the generated distribution and the true distribution towards each other, thereby neutralizing a portion of the initial accumulated error.
The other term related to additional sampling error includes the accumulation of discretization and approximation errors in  $[t_{\min}, t_{\max}]$ . Despite the fact that SDE incurs a higher discretization error than ODE  $(O(\sqrt{\delta})$  versus  $O(\delta))$ , the contraction effect on the initial error is the dominant factor impacting the upper bound in the large NFE regime where  $\delta$  is small. Consequently, the upper bound for SDE is significantly lower. This provides insight into why SDE outperforms ODE in the large NFE regime, where the influence of discretization errors diminishes and the contraction effect dominates. In light of the distinct advantages of SDE and ODE, it is natural to ask whether we can combine their strengths. Specifically, can we devise a sampling algorithm that maintains a comparable level of discretization error as ODE, while also benefiting from, or even amplifying, the contraction effects induced by the stochasticity of SDE? In the next section, we introduce a novel algorithm, termed Restart, designed to achieve these two goals simultaneously.
4 Harnessing stochasticity with Restart
In this section, we present the Restart sampling algorithm, which incorporates stochasticity during sampling while enabling fast generation. We introduce the algorithm in Sec 4.1, followed by a theoretical analysis in Sec 4.2. Our analysis shows that Restart achieves a better Wasserstein upper bound compared to those of SDE and ODE in Theorem 1 due to greater contraction effects.
In the Restart algorithm, simulation performs a few repeated back-and-forth steps within a pre-defined time interval  $[t_{\mathrm{min}}, t_{\mathrm{max}}] \subset [0, T]$ , as depicted in Figure 1(a). This interval is embedded into the simulation of the original backward ODE referred to as the main backward process, which runs from  $T$  to 0. In addition, we refer to the backward process within the Restart interval  $[t_{\mathrm{min}}, t_{\mathrm{max}}]$  as the Restart backward process, to distinguish it from the main backward process.
Starting with samples at time  $t_{\mathrm{min}}$ , which are generated by following the main backward process, the Restart algorithm adds a large noise to transit the samples from  $t_{\mathrm{min}}$  to  $t_{\mathrm{max}}$  with the help of
the forward process. The forward process does not require any evaluation of the neural network  $s_{\theta}(x,t)$ , as it is generally defined by an analytical perturbation kernel capable of transporting distributions from  $t_{\mathrm{min}}$  to  $t_{\mathrm{max}}$ . For instance, in the case of diffusion models, the perturbation kernel is  $\mathcal{N}(\mathbf{0}, (\sigma(t_{\mathrm{max}})^2 - \sigma(t_{\mathrm{min}})^2) I_{d \times d})$ . The added noise in this step induces a more significant contraction compared to the small, interleaved noise in SDE. The step acts as if partially restarting the main backward process by increasing the time. Following this step, Restart simulates the backward ODE from  $t_{\mathrm{max}}$  back to  $t_{\mathrm{min}}$  using the neural network predictions as in regular ODE. We repeat these forward-backward steps within  $[t_{\mathrm{min}}, t_{\mathrm{max}}]$  interval  $K$  times in order to further derive the benefit from contraction. Specifically, the forward and backward processes in the  $i^{\mathrm{th}}$  iteration  $(i \in \{0, \dots, K-1\})$  proceed as follows:
$$
\text {(R e s t a r t f o r w a r d p r o c e s s)} \quad x _ {t _ {\max }} ^ {i + 1} = x _ {t _ {\min }} ^ {i} + \varepsilon_ {t _ {\min } \rightarrow t _ {\max }} \tag {3}
$$
$$
\text {(R e s t a r t b a c k w a r d p r o c e s s)} \quad x _ {t _ {\min }} ^ {i + 1} = \mathrm {O D E} _ {\theta} \left(x _ {t _ {\max }} ^ {i + 1}, t _ {\max } \rightarrow t _ {\min }\right) \tag {4}
$$
where the initial  $x_{t_{\mathrm{min}}}^{0}$  is obtained by simulating the ODE until  $t_{\mathrm{min}}$ :  $x_{t_{\mathrm{min}}}^{0} = \mathrm{ODE}_{\theta}(x_{T}, T \to t_{\mathrm{min}})$ , and the noise  $\varepsilon_{t_{\mathrm{min}} \to t_{\mathrm{max}}}$  is sampled from the corresponding perturbation kernel from  $t_{\mathrm{min}}$  to  $t_{\mathrm{max}}$ . The Restart algorithm not only adds substantial noise in the Restart forward process (Eq. (3)), but also separates the stochasticity from the ODE, leading to a greater contraction effect, which we will demonstrate theoretically in the next subsection. For example, we set  $[t_{\mathrm{min}}, t_{\mathrm{max}}] = [0.05, 0.3]$  for the VP model [13] on CIFAR-10. Repetitive use of the forward noise effectively mitigates errors accumulated from the preceding simulation up until  $t_{\mathrm{max}}$ . Furthermore, the Restart algorithm does not suffer from large discretization errors as it is mainly built from following the ODE in the Restart backward process (Eq. (4)). The effect is that the Restart algorithm is able to reduce the total sampling errors even in the small NFE regime. Detailed pseudocode for the Restart sampling process can be found in Algorithm 2, Appendix B.2.
We provide a theoretical analysis of the Restart algorithm under the same setting as Theorem 1. In particular, we prove the following theorem, which shows that Restart achieves a much smaller contracted error in the Wasserstein upper bound than SDE (Theorem 1), thanks to the separation of the noise from the drift, as well as the large added noise in the Restart forward process (Eq. (3)). The repetition of the Restart cycle  $K$  times further leads to an enhanced reduction in the initial accumulated error. We denote the intermediate distribution in the  $i^{\text{th}}$  Restart iteration, following the discretized trajectories and the learned field  $s_{\theta}$ , as  $p_{t \in [t_{\min}, t_{\max}]}^{\text{Restart}_{\theta}(i)}$ .
Theorem 2 (Informal). Under the same setting of Theorem 1, assume  $K \leq \frac{C}{L_2(t_{\max} - t_{\min})}$  for some universal constant  $C$ . Then
$$
\underbrace {W _ {1} \left(p _ {t _ {\min }} ^ {\text {R e s t a r t} _ {\theta} (K)} , p _ {t _ {\min }}\right)} _ {\text {t o t a l e r r o r}} \leq \underbrace {B \cdot (1 - \lambda) ^ {K} T V \left(p _ {t _ {\max }} ^ {\text {R e s t a r t} _ {\theta} (0)} , p _ {t _ {\max }}\right)} _ {\text {u p p e r b o u n d o n c o n t r a c t e d e r r o r}} + \underbrace {(K + 1) \cdot O \left(\delta + \epsilon_ {\text {a p p r o x}}\right) \left(t _ {\max } - t _ {\min }\right)} _ {\text {u p p e r b o u n d o n a d d i t i o n a l s a m p l i n g e r r o r}}
$$
where  $\lambda < 1$  is the same contraction factor as Theorem 1.  $O()$  hides polynomial dependency on various Lipschitz constants, dimension.
Proof sketch. To bound the total error, we introduce an auxiliary process  $q_{t\in [t_{\min},t_{\max}]}^{\text{Restart}_{\theta}(i)}$ , which initiates from true distribution  $p_{t_{\max}}$  and performs the Restart iterations. This process differs from  $p_{t\in [t_{\min},t_{\max}]}^{\text{Restart}_{\theta}(i)}$  only in its initial distribution at  $t_{\max}$  ( $p_{t_{\max}}$  versus  $p_{t_{\max}}^{\text{Restart}_{\theta}(0)}$ ). We bound the total error by the following triangular inequality:
$$
\underbrace {W _ {1} \left(p _ {t _ {\operatorname* {m i n}}} ^ {\text {R e s t a r t} _ {\theta} (K)} , p _ {t _ {\operatorname* {m i n}}}\right)} _ {\text {t o t a l e r r o r}} \leq \underbrace {W _ {1} \left(p _ {t _ {\operatorname* {m i n}}} ^ {\text {R e s t a r t} _ {\theta} (K)} , q _ {t _ {\operatorname* {m i n}}} ^ {\text {R e s t a r t} _ {\theta} (K)}\right)} _ {\text {c o n t r a c t e d e r r o r}} + \underbrace {W _ {1} \left(q _ {t _ {\operatorname* {m i n}}} ^ {\text {R e s t a r t} _ {\theta} (K)} , p _ {t _ {\operatorname* {m i n}}}\right)} _ {\text {a d d i t i o n a l s a m p l i n g e r r o r}}
$$