GENIE: Higher-Order Denoising Diffusion Solvers

Introduction

Denoising diffusion models (DDMs) offer both state-of-the-art synthesis quality and sample diversity in combination with a robust and scalable learning objective. DDMs have been used for image and video synthesis, super-resolution , deblurring , image editing and inpainting , text-to-image synthesis , conditional and semantic image generation , image-to-image translation and for inverse problems in medical imaging . They also enable high-quality speech synthesis , 3D shape generation , molecular modeling , maximum likelihood training , and more . In DDMs, a diffusion process gradually perturbs the data towards random noise, while a deep neural network learns to denoise. Formally, the problem reduces to learning the score function, i.e., the gradient of the log-density of the perturbed data. The (approximate) inverse of the forward diffusion can be described by an ordinary or a stochastic differential equation (ODE or SDE, respectively), defined by the learned score function, and can therefore be used for generation when starting from random noise .

A crucial drawback of DDMs is that the generative ODE or SDE is typically difficult to solve, due to the complex score function. Therefore, efficient and tailored samplers are required for fast synthesis. In this work, building on the generative ODE , we rigorously derive a novel second-order ODE solver using truncated Taylor methods . These higher-order methods require higher-order gradients of the ODE—in our case this includes higher-order gradients of the log-density of the perturbed data, i.e., higher-order score functions. Because such higher-order scores are usually not available, existing works typically use simple first-order solvers or samplers with low accuracy , higher-order methods that rely on suboptimal finite difference or other approximations , or alternative approaches for accelerated sampling. Here, we fundamentally avoid such approximations and directly model the higher-order gradient terms: Importantly, our novel Higher-Order Denoising Diffusion Solver (GENIE) relies on Jacobian-vector products (JVPs) involving second-order scores. We propose to calculate these JVPs by automatic differentiation of the regular learnt first-order scores. For computational efficiency, we then distill the entire higher-order gradient of the ODE, including the JVPs, into a separate neural network. In practice, we only need to add a small head to the first-order score network to predict the components of the higher-order ODE gradient. By directly modeling the JVPs we avoid explicitly forming high-dimensional higher-order scores. Intuitively, the higher-order terms in GENIE capture the local curvature of the ODE and enable larger steps when iteratively solving the generative ODE (Fig. 1).

Experimentally, we validate GENIE on multiple image modeling benchmarks and achieve state-of-the-art performance in solving the generative ODE of DDMs with few synthesis steps. In contrast to recent methods that fundamentally modify the generation process of DDMs by training conditional GANs or by distilling the full sampling trajectory , GENIE solves the true generative ODE. Therefore, we also show that we can still encode images in the DDM’s latent space, as required for instance for image interpolation, and use techniques such as guided sampling .

We make the following contributions: (i) We introduce GENIE, a novel second-order ODE solver for fast DDM sampling. (ii) We propose to extract the required higher-order terms from the first-order score model by automatic differentiation. In contrast to existing works, we explicitly work with higher-order scores without finite difference approximations. To the best of our knowledge, GENIE is the first method that explicitly uses higher-order scores for generative modeling with DDMs. (iii) We propose to directly model the necessary JVPs and distill them into a small neural network. (iv) We outperform all previous solvers and samplers for the generative differential equations of DDMs.

Background

We consider continuous-time DDMs whose forward process can be described by

where ${\mathbf{x}}_{0}\sim p_{0}({\mathbf{x}}_{0})$ is drawn from the empirical data distribution and ${\mathbf{x}}_{t}$ refers to diffused data samples at time $t\in$ along the diffusion process. The functions $\alpha_{t}$ and $\sigma_{t}$ are generally chosen such that the logarithmic signal-to-noise ratio $\log\frac{\alpha_{t}^{2}}{\sigma_{t}^{2}}$ decreases monotonically with $t$ and the data diffuses towards random noise, i.e., $p_{1}({\mathbf{x}}_{1})\!\approx\!{\mathcal{N}}({\mathbf{x}}_{1};\bm{0},{\bm{I}})$. We use variance-preserving diffusion processes for which $\sigma_{t}^{2}=1-\alpha_{t}^{2}$ (however, all methods introduced in this work are applicable to more general DDMs). The diffusion process can then be expressed by the (variance-preserving) SDE

where $\beta_{t}=-\frac{d}{dt}\log\alpha_{t}^{2}$, ${\mathbf{x}}_{0}\sim p_{0}({\mathbf{x}}_{0})$ and ${\mathbf{w}}_{t}$ is a standard Wiener process. A corresponding reverse diffusion process that effectively inverts the forward diffusion is given by

and this reverse-time generative SDE is marginally equivalent to the generative ODE

where $\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})$ is the score function. Eq. 4 is referred to as the Probability Flow ODE , an instance of continuous Normalizing flows . To generate samples from the DDM, one can sample ${\mathbf{x}}_{1}\sim{\mathcal{N}}({\mathbf{x}}_{1};\bm{0},{\bm{I}})$ and numerically simulate either the Probability Flow ODE or the generative SDE, replacing the unknown score function by a learned score model ${\bm{s}}_{\bm{\theta}}({\mathbf{x}}_{t},t)\approx\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})$.

The DDIM solver has been particularly popular to simulate DDMs due to its speed and simplicity. It has been shown that DDIM is Euler’s method applied to an ODE based on a re-parameterization of the Probability Flow ODE : Defining $\gamma_{t}=\sqrt{\frac{1-\alpha_{t}^{2}}{\alpha_{t}^{2}}}$ and $\bar{\mathbf{x}}_{t}={\mathbf{x}}_{t}\sqrt{1+\gamma_{t}^{2}}$, we have

where we inserted Eq. 4 for $\frac{d{\mathbf{x}}_{t}}{dt}$ and used $\beta(t)\frac{dt}{d\gamma_{t}}=\frac{2\gamma_{t}}{\gamma_{t}^{2}+1}$. Letting ${\bm{s}}_{\bm{\theta}}({\mathbf{x}}_{t},t)\coloneqq-\frac{{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)}{\sigma_{t}}$ denote a parameterization of the score model, the approximate generative DDIM ODE is then given by

where we used $\sigma_{t}=\sqrt{1-\alpha_{t}^{2}}=\frac{\gamma_{t}}{\sqrt{\gamma_{t}^{2}+1}}$ (see App. A for a more detailed derivation of Eq. 6). The model ${\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ can be learned by minimizing the score matching objective

Higher-Order Denoising Diffusion Solver

As discussed in Sec. 2, the so-known DDIM solver is simply Euler’s method applied to the DDIM ODE (cf. Eq. 6). In this work, we apply a higher-order method to the DDIM ODE, building on the truncated Taylor method (TTM) . The $p$-th TTM is simply the $p$-th order Taylor polynomial applied to an ODE. For example, for the general $\frac{d{\mathbf{y}}}{dt}={\bm{f}}({\mathbf{y}},t)$, the $p$-th TTM reads as

where $h_{n}=t_{n+1}-t_{n}$ (see Sec. B.1 for a truncation error analysis with respect to the exact ODE solution). Note that the first TTM is simply Euler’s method. Applying the second TTM to the DDIM ODE results in the following scheme:

where $h_{n}=\gamma_{t_{n+1}}-\gamma_{t_{n}}$. Recall that $\gamma_{t}=\sqrt{\frac{1-\alpha_{t}^{2}}{\alpha_{t}^{2}}}$, where the function $\alpha_{t}$ is a time-dependent hyperparameter of the DDM. The total derivative $d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}\coloneqq\frac{d{\bm{\epsilon}}_{\bm{\theta}}}{d\gamma_{t}}$ can be decomposed as follows

where $\tfrac{\partial{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)}{\partial{\mathbf{x}}_{t}}$ denotes the Jacobian of ${\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ and

If not explicitly stated otherwise, we refer to the second TTM applied to the DDIM ODE, i.e., the scheme in Eq. 9, as Higher-Order Denoising Diffusion Solver (GENIE). Intuitively, the higher-order gradient terms used in the second TMM model the local curvature of the ODE. This translates into a Taylor formula-based extrapolation that is quadratic in time (cf. Eqs. 8 and 9) and more accurate than linear extrapolation, as in Euler’s method, thereby enabling larger time steps (see Fig. 1 for a visualization). In App. B, we also discuss the application of the third TTM to the DDIM ODE. We emphasize that TTMs are not restricted to the DDIM ODE and could just as well be applied to the Probability Flow ODE (also see App. B) or neural ODEs more generally.

The Benefit of Higher-Order Methods: We showcase the benefit of higher-order methods on a 2D toy distribution (Fig. 2(a)) for which we know the score function as well as all higher-order derivatives necessary for GENIE analytically. We generate 1k different accurate “ground truth” trajectories ${\mathbf{x}}_{t}$ using DDIM with 10k steps. We compare these “ground truth” trajectories to single steps of DDIM and GENIE for varying step sizes $\Delta t$. We then measure the mean $L_{2}$-distance of the single steps $\hat{\mathbf{x}}_{t}(\Delta t)$ to the “ground truth” trajectories ${\mathbf{x}}_{t}$, and we repeat this experiment for three starting points $t\in\{0.1,0.2,0.5\}$. We see (Fig. 3 (top)) that GENIE can use larger step sizes to stay within a certain error tolerance for all starting points $t$. We further show samples for DDIM and GENIE, using 25 solver steps, in Fig. 2. DDIM has the undesired behavior of sampling low-density regions between modes, whereas GENIE looks like a slightly noisy version of the ground truth distribution (Fig. 2(a)).

Comparison to Multistep Methods: Linear multistep methods are an alternative higher-order method to solve ODEs. Liu et al. applied the well-established Adams–Bashforth [AB, 77] method to the DDIM ODE. AB methods can be derived from TTMs by approximating higher-order derivatives $\frac{d^{p}{\mathbf{y}}}{dt^{p}}$ using the finite difference method . For example, the second AB method is obtained from the second TTM by replacing $\frac{d^{2}{\mathbf{y}}}{dt^{2}}$ with the first-order forward difference approximation ${(f({\mathbf{y}}_{t_{n}},t_{n})-f({\mathbf{y}}_{t_{n-1}},t_{n-1})})/{h_{n-1}}$. In Fig. 3 (bottom), we visualize the mean $L_{2}$-norm of the difference $\xi_{t}(\Delta t)$ between the analytical derivative $d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}$ and its first-order forward difference approximation for varying step sizes $\Delta t$ for the 2D toy distribution. The approximation is especially poor at small $t$ for which the score function becomes complex (App. E for details on all toy experiments).

The above observations inspire to apply GENIE to DDMs of more complex and high-dimensional data such as images. Regular DDMs learn a model ${\bm{\epsilon}}_{\bm{\theta}}$ for the first-order score; however, the higher-order gradient terms required for GENIE (cf. Eq. 10) are not immediately available to us, unlike in the toy example above. Let us insert Eq. 11 into Eq. 10 and analyze the required terms more closely:

We see that the full derivative decomposes into two JVP terms and one simpler time derivative term. The term $\frac{\partial{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)}{\partial{\mathbf{x}}_{t}}$ plays a crucial role in Eq. 12. It can be expressed as

which means that GENIE relies on second-order score functions $\nabla_{{\mathbf{x}}_{t}}^{\top}\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})$ under the hood.

Given a DDM, that is, given ${\bm{\epsilon}}_{\bm{\theta}}$, we could compute the derivative $d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}$ for the GENIE scheme in Eq. 9 using automatic differentiation (AD). This would, however, make a single step of GENIE at least twice as costly as DDIM, because we would need a forward pass through the ${\bm{\epsilon}}_{\bm{\theta}}$ network to compute ${\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ itself, and another pass to compute the JVPs and the time derivative in Eq. 12. These forward passes cannot be parallelized, since the vector-part of $\textrm{JVP}_{1}$ in Eq. 12 involves ${\bm{\epsilon}}_{\bm{\theta}}$ itself, and needs to be known before computing the JVP. To accelerate sampling, this overhead is too expensive.

Gradient Distillation: To avoid this overhead, we propose to first distill $d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}$ into a separate neural network. During distillation training, we can use the slow AD-based calculation of $d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}$, but during synthesis we call the trained neural network. We build on the observation that the internal representations of the neural network modeling ${\bm{\epsilon}}_{\bm{\theta}}$ (in our case a U-Net architecture) can be used for downstream tasks : specifically, we provide the last feature layer from the ${\bm{\epsilon}}_{\bm{\theta}}$ network together with its time embedding as well as ${\mathbf{x}}_{t}$ and the output ${\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ to a small prediction head ${\bm{k}}_{\bm{\psi}}({\mathbf{x}}_{t},t)$ that models the different terms in Eq. 12 (see Fig. 4). The overhead generated by ${\bm{k}}_{\bm{\psi}}$ is small, for instance less than 2% for our CIFAR-10 model (also see Sec. 5), and we found this approach to provide excellent performance. Note that in principle we could also train an independent deep neural network, which does not make use of the internal representations of ${\bm{\epsilon}}_{\bm{\theta}}$ and could therefore theoretically be run in parallel to the ${\bm{\epsilon}}_{\bm{\theta}}$ model. We justify using small prediction heads over independent neural networks because AD-based distillation training is slow: in each training iteration we first need to call the ${\bm{\epsilon}}_{\bm{\theta}}$ network, then calculate the JVP terms, and only then can we call the distillation model. By modeling $d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}$ via small prediction heads, while reusing the internal representation of the score model, we can make training relatively fast: we only need to train ${\bm{k}}_{\bm{\psi}}$ for up to 50k iterations. In contrast, training score models from scratch takes roughly an order of magnitude more iterations. We leave training of independent networks to predict $d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}$ to future work.

Mixed Network Parameterization: We found that learning $d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}$ directly as single output of a neural network can be challenging. Assuming a single data point distribution $p_{0}({\mathbf{x}}_{0})=\delta({\mathbf{x}}_{0}=\mathbf{0})$, for which we know the diffused score function and all higher-order derivatives analytically, we found that the terms in Eq. 12 all behave very differently within the $t\in$ interval (for instance, the prefactor of $\textrm{JVP}_{1}$ in Eq. 12 approaches $1$ as $t\rightarrow 0$, while $\textrm{JVP}_{2}$’s prefactor vanishes). As outlined in detail in Sec. C.2.3, this simple single data point assumption implies an effective mixed network parameterization, an approach inspired by the “mixed score parametrizations” in Vahdat et al. and Dockhorn et al. . In particular, we model

where ${\bm{k}}_{\bm{\psi}}^{(i)}({\mathbf{x}}_{t},t)$, $i\in\{1,2,3\}$, are different output channels of the neural network (i.e. the additional head on top of the ${\bm{\epsilon}}_{\bm{\theta}}$ network). The three terms in Eq. 14 exactly correspond to the three terms of Eq. 12, in the same order. We show the superior performance of this parametrization in Sec. 5.3.

Learning Objective: Ideally, we would like our model ${\bm{k}}_{\bm{\psi}}$ to match $d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}$ exactly, for all $t\in[0,T]$ and ${\mathbf{x}}_{t}$ in the diffused data distribution, which the generative ODE trajectories traverse. This suggests a simple (weighted) $L_{2}$-loss, similar to regular score matching losses for DDMs :

Alternative Learning Approaches: As shown in Eq. 13, GENIE relies on second-order score functions. Recently, Meng et al. directly learnt such higher-order scores with higher-order score matching objectives. Directly applying these techniques has the downside that we would need to explicitly form the higher-order score terms $\nabla^{\top}_{{\mathbf{x}}_{t}}{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$, which are very high-dimensional for data such as images. Low-rank approximations are possible, but potentially insufficient for high performance. In our approach, we are avoiding this complication by directly modeling the lower-dimensional JVPs. We found that the methods from Meng et al. can be modified to provide higher-order score matching objectives for the JVP terms required for GENIE and we briefly explored this (see App. D). However, our distillation approach with AD-based higher-order gradients worked much better. Nevertheless, this is an interesting direction for future research. To the best of our knowledge, GENIE is the first solver for the generative differential equations of DDMs that directly uses higher-order scores (in the form of the distilled JVPs) for generative modeling without finite difference or other approximations.

Related Work

Accelerated Sampling from DDMs. Several previous works address the slow sampling of DDMs: One line of work reduces and readjusts the timesteps used in time-discretized DDMs . This can be done systematically by grid search or dynamic programming . Bao et al. speed up sampling by defining a new DDM with optimal reverse variances. DDIM , discussed in Sec. 2, was also introduced as a method to accelerate DDM synthesis. Further works leverage modern ODE and SDE solvers for fast synthesis from (continuous-time) DDMs: For instance, higher-order Runge–Kutta methods and adaptive step size SDE solvers have been used. These methods are not optimally suited for the few-step synthesis regime, in which GENIE shines; see also Sec. 5. Most closely related to our work is Liu et al. , which simulates the DDIM ODE using a higher-order linear multistep method . As shown in Sec. 3, linear multistep methods can be considered an approximation of the TTMs used in GENIE. Furthermore, Tachibana et al. solve the generative SDE via a higher-order Itô–Taylor method and in contrast to our work, they propose to use an “ideal derivative trick” to approximate higher-order score functions. In Sec. B.2, we show that applying this ideal derivative approximation to the DDIM ODE does not have any effect: the “ideal derivatives” are zero by construction. Note that in GENIE, we in fact use the DDIM ODE, rather than, for example, the regular Probability Flow ODE , as the base ODE for GENIE.

Alternatively, sampling from DDMs can also be accelerated via learning: For instance, Watson et al. learn parameters of a generalized family of DDMs by optimizing for perceptual output quality. Luhman and Luhman and Salimans and Ho distill a DDIM sampler into a student model, which enables sampling in as few as a single step. Xiao et al. replace DDMs’ Gaussian samplers with expressive generative adversarial networks, similarly allowing for few-step synthesis. GENIE can also be considered a learning-based approach, as we distill a derivative of the generative ODE into a separate neural network. However, in contrast to the mentioned methods, GENIE still solves the true underlying generative ODE, which has major advantages: for instance, it can still be used easily for classifier-guided sampling and to efficiently encode data into latent space—a prerequisite for likelihood calculation and editing applications . Note that the learnt sampler defines a proper probabilistic generalized DDM; however, it isn’t clear how it relates to the generative SDE or ODE and therefore how compatible the method is with applications such as classifier guidance.

Other approaches to accelerate DDM sampling change the diffusion itself or train DDMs in the latent space of a Variational Autoencoder . GENIE is complementary to these methods.

Higher-Order ODE Gradients beyond DDMs. TTMs and other methods that leverage higher-order gradients are also applied outside the scope of DDMs. For instance, higher-order derivatives can play a crucial role when developing solvers and regularization techniques for neural ODEs . Outside the field of machine learning, higher-order TTMs have been widely studied, for example, to develop solvers for stiff and non-stiff systems.

Concurrent Works. Zhang and Chen motivate the DDIM ODE from an exponential integrator perspective applied to the Probability Flow ODE and propose to apply existing solvers from the numerical ODE literature, namely, Runge–Kutta and linear multistepping, to the DDIM ODE directly. Lu et al. similarly recognize the semi-linear structure of the Probability Flow ODE, derive dedicated solvers, and introduce new step size schedulers to accelerate DDM sampling. Karras et al. propose new fast solvers, both deterministic and stochastic, specifically designed for the differential equations arising in DDMs. Both Zhang et al. and Karras et al. realize that the DDIM ODE has “straight line solution trajectories” for spherical normal data and single data points—this exactly corresponds to our derivation that the higher-order terms in the DDIM ODE are zero in such a setting (see Sec. B.2). Bao et al. learn covariance matrices for DDM sampling using prediction heads somewhat similar to the ones in GENIE; in Sec. G.1, we thoroughly discuss the differences between GENIE and the method proposed in Bao et al. .

Experiments

Datasets: We run experiments on five datasets: CIFAR-10 (resolution 32), LSUN Bedrooms (128), LSUN Church-Outdoor (128), (conditional) ImageNet (64), and AFHQv2 (512). On AFHQv2 we only consider the subset of cats; referred to as “Cats” in the remainder of this work.

Architectures: Except for CIFAR-10 (we use a checkpoint by Song et al. ), we train our own score models using architectures introduced by previous works . The architecture of our prediction heads is based on (modified) BigGAN residual blocks . To minimize computational overhead, we only use a single residual block. See App. C for training and architecture details.

Evaluation: We measure sample quality via Fréchet Inception Distance [FID, 102] (see Sec. F.1).

Synthesis Strategy: We simulate the DDIM ODE from $t{=}1$ up to $t{=}10^{-3}$ using evaluation times following a quadratic function (quadratic striding ). For variance-preserving DDMs, it can be beneficial to denoise the ODE solver output at the cutoff $t{=}10^{-3}$, i.e., ${\mathbf{x}}_{0}=\frac{{\mathbf{x}}_{t}-\sigma_{t}{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)}{\alpha_{t}}$ . Note that the denoising step involves a score model evaluation, and therefore “loses” a function evaluation that could otherwise be used as an additional step in the ODE solver. To this end, denoising the output of the ODE solver is left as a hyperparameter of our synthesis strategy.

Analytical First Step (AFS): Every additional neural network call becomes crucial in the low number of function evaluations (NFEs) regime. We found that we can improve the performance of GENIE and all other methods evaluated on our checkpoints by replacing the learned score with the (analytical) score of ${\mathcal{N}}(\bm{0},{\bm{I}})\approx p_{t=1}({\mathbf{x}}_{t})$ in the first step of the ODE solver. The “gained” function evaluation can then be used as an additional step in the ODE solver. Similarly to the denoising step mentioned above, AFS is treated as a hyperparameter of our Synthesis Strategy. AFS details in Sec. F.2.

Accounting for Computational Overhead: GENIE has a slightly increased computational overhead compared to other solvers due to the prediction head ${\bm{k}}_{\bm{\psi}}$. The computational overhead is increased by 1.47%, 2.83%, 14.0%, and 14.4% on CIFAR-10, ImageNet, LSUN Bedrooms, and LSUN Church-Outdoor, respectively (see also Sec. C.2.5). This additional overhead is always accounted for implicitly: we divide the NFEs by the computational overhead and round to the nearest integer. For example, on LSUN Bedrooms, we compare baselines with 10/15 NFEs to GENIE with 9/13 NFEs.

In Fig. 5 we compare our method to the most competitive baselines. In particular, on the same score model checkpoints, we compare GENIE with DDIM , S-PNDM , and F-PNDM . For these four methods, we only include the best result over the two hyperparameters discussed above, namely, the denoising step and AFS (see Sec. F.6 for tables with all results). We also include three competitive results from the literature that use different checkpoints and sampling strategies: for each method, we include the best result for their respective set of hyperparameters. We do not compare in this figure with Knowledge Distillation [KD, 68], Progressive Distillation [PG, 69] and Denoising Diffusion GANs [DDGAN, 67] as they do not solve the generative ODE/SDE and use fundamentally different sampling approaches with drawbacks discussed in Sec. 4.

For NFEs $\in\{10,15,20,25\}$, GENIE outperforms all baselines (on the same checkpoint) on all four datasets (see detailed results in Sec. F.6 and GENIE image samples in Sec. F.7). On CIFAR-10 and (conditional) ImageNet, GENIE also outperforms these baselines for NFEs=5, whereas DDIM outperforms GENIE slightly on the LSUN datasets (see tables in Sec. F.6). GENIE also performs better than the three additional baselines from the literature (which use different checkpoints and sampling strategies) with the exception of the Learned Sampler [LS, 66] on LSUN Bedrooms for NFEs=20. Though LS uses a learned striding schedule on LSUN Bedrooms (whereas GENIE simply uses quadratic striding), the LS’s advantage is most likely due to the different checkpoint. In Tab. 1, we investigate the effect of optimizing the striding schedule, via learning (LS) or grid search (DDIM & GENIE), on CIFAR-10 and find that its significance decreases rapidly with increased NFEs (also see Sec. F.6 for details). In Tab. 1, we also show additional baseline results; however, we do not include commonly-used adaptive step size solvers in Fig. 5, as they are arguably not well-suited for this low NFE regime: for example, on the same CIFAR-10 checkpoint we use for GENIE, the adaptive SDE solver introduced in Jolicoeur-Martineau et al. obtains an FID of 82.4 at 48 NFEs. Also on the same checkpoint, the adaptive Runge–Kutta 4(5) method applied to the ProbabilityFlow ODE achieves an FID of 13.1 at 38 NFEs (solver tolerances set to $10^{-2}$).

The results in Fig. 5 suggest that higher-order gradient information, as used in GENIE, can be efficiently leveraged for image synthesis. Despite using small prediction heads our distillation seems to be sufficiently accurate: for reference, replacing the distillation heads with the derivatives computed via AD, we obtain FIDs of 9.22, 4.11, 3.54, 3.46 using 10, 20, 30, and 40 NFEs, respectively (NFEs adjusted assuming an additional computational overhead of 100%). As discussed in Sec. 3, linear multistep methods such as S-PNDM and F-PNDM can be considered (finite difference) approximations to TTMs as used in GENIE. These approximations can be inaccurate for large timesteps, which potentially explains their inferior performance when compared to GENIE. When compared to DDIM, the superior performance of GENIE seems to become less significant for large NFE: this is in line with the theory, as higher-order gradients contribute less for smaller step sizes (see the GENIE scheme in Eq. 9). Approaches such as FastDDIM and AnalyticDDIM , which adapt variances and discretizations of discrete-time DDMs, are useful; however, GENIE suggests that rigorous higher-order ODE solvers leveraging the continuous-time DDM formalism are still more powerful. To the best of our knowledge, the only methods that outperform GENIE abandon this ODE or SDE formulation entirely and train NFE-specific models which are optimized for the single use-case of image synthesis.

2 Guidance and Encoding

As discussed in Sec. 4, one major drawback of approaches such as KD , PG and DDGAN is that they abandon the ODE/SDE formalism, and cannot easily use methods such as classifier(-free) guidance or perform image encoding. However, these techniques can play an important role in synthesizing photorealistic images from DDMs , as well as for image editing tasks .

Classifier-Free Guidance : We replace the unconditional model ${\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ with $\hat{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t,c,w)=(1+w){\bm{\epsilon}}_{{\bm{\theta}}}({\mathbf{x}}_{t},t,c)-w{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ in the DDIM ODE (cf Eq. 6), where ${\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t,c)$ is a conditional model and $w>1.0$ is the “guidance scale”. GENIE then requires the derivative

for guidance. Hence, we need to distill $d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t,c)$ and $d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$, for which we could also share parameters . We compare GENIE with DDIM on ImageNet in Fig. 6. GENIE clearly outperforms DDIM, in particular for few NFEs, and GENIE also synthesizes high-quality images (see Fig. 7).

Image Encoding: We can use GENIE also to solve the generative ODE in reverse to encode given images. Therefore, we compare GENIE to DDIM on the “encode-decode” task, analyzing reconstructions for different NFEs (used twice for encoding and decoding): We find that GENIE reconstructs images much more accurately (see Fig. 8). For more details on this experiment as well as the guidance experiment above, see Sec. F.4 and Sec. F.3, respectively. We also show latent space interpolations for both GENIE and DDIM in Sec. F.5.

3 Ablation Studies

4 Upsampling

Cascaded diffusion model pipelines and DDM-based super-resolution have become crucial ingredients in DDMs for large-scale image generation . Hence, we also explore the applicability of GENIE in this setting. We train a $128\times 128$ base model as well as a $128\times 128\rightarrow 512\times 512$ diffusion upsampler on Cats. In Tab. 3, we compare the generative performance of GENIE to other fast samplers for the upsampler (in isolation). We find that GENIE performs very well on this task: with only five NFEs GENIE outperforms all other methods at NFEs=15. We show upsampled samples for GENIE with NFEs=5 in Fig. 9. For more quantitative and qualitative results, we refer to Sec. F.6 and Sec. F.7, respectively. Training and inference details for the score model and the GENIE prediction head, for both base model and upsampler, can be found in App. C.

Conclusions

We introduced GENIE, a higher-order ODE solver for DDMs. GENIE improves upon the commonly used DDIM solver by capturing the local curvature of its ODE’s gradient field, which allows for larger step sizes when solving the ODE. We further propose to distill the required higher-order derivatives into a small prediction head—which we can efficiently call during inference—on top of the first-order score network. A limitation of GENIE is that it is still slightly slower than approaches that abandon the differential equation framework of DDMs altogether, which, however, comes at the considerable cost of preventing applications such as guided sampling. To overcome this limitation, future work could leverage even higher-order gradients to accelerate sampling from DDMs even further (also see Sec. G.2).

Broader Impact. Fast synthesis from DDMs, the goal of GENIE, can potentially make DDMs an attractive method for promising interactive generative modeling applications, such as digital content creation or real-time audio synthesis, and also reduce DDMs’ environmental footprint by decreasing the computational load during inference. Although we validate GENIE on image synthesis, it could also be utilized for other tasks, which makes its broader societal impact application-dependent. In that context, it is important that practitioners apply an abundance of caution to mitigate impacts given generative modeling can also be used for malicious purposes, discussed for instance in Vaccari and Chadwick , Nguyen et al. , Mirsky and Lee .

Acknowledgements

We thank Yaoliang Yu for early discussions. Tim Dockhorn acknowledges additional funding from the Vector Institute Research Grant, which is not in direct support of this work.

References

Appendix A DDIM ODE

The DDIM ODE has previously been shown to be a re-parameterization of the Probability Flow ODE . In this section, we show an alternative presentation to the ones given in Song et al. and Salimans and Ho . We start from the Probability Flow ODE for variance-preserving continuous-time DDMs , i.e.,

where $\beta_{t}=-\frac{d}{dt}\log\alpha_{t}^{2}$ and $\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})$ is the score function. Replacing the unknown score function with a learned score model ${\bm{s}}_{\bm{\theta}}({\mathbf{x}}_{t},t)\approx\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})$, we obtain the approximate Probability Flow ODE

Let us now define $\gamma_{t}=\sqrt{\frac{1-\alpha_{t}^{2}}{\alpha_{t}^{2}}}$ and $\bar{\mathbf{x}}_{t}={\mathbf{x}}_{t}\sqrt{1+\gamma_{t}^{2}}$, and take the (total) derivative of $\bar{\mathbf{x}}_{t}$ with respect to $\gamma_{t}$:

The derivative $\frac{d{\mathbf{x}}_{t}}{d\gamma_{t}}$ can be computed as follows

We can write $\alpha_{t}^{2}$ as a function of $\gamma_{t}$, i.e., $\alpha_{t}^{2}=\left(\gamma_{t}^{2}+1\right)^{-1}$, and therefore

Lastly, inserting Eq. 28 into Eq. 20, we have

Letting ${\bm{s}}_{\bm{\theta}}({\mathbf{x}}_{t},t)\coloneqq-\frac{{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)}{\sigma_{t}}$, where $\sigma_{t}=\sqrt{1-\alpha_{t}^{2}}=\frac{\gamma_{t}}{\sqrt{\gamma_{t}^{2}+1}}$, denote a particular parameterization of the score model, we obtain the approximate generative DDIM ODE as

Appendix B Synthesis from Denoising Diffusion Models via Truncated Taylor Methods

In this work, we propose Higher-Order Denoising Diffusion Solvers (GENIE). GENIE is based on the truncated Taylor method (TTM) . As outlined in Sec. 3, the $p$-th TTM is simply the p-th order Taylor polynomial applied to an ODE. For example, for the general $\frac{d{\mathbf{y}}}{dt}={\bm{f}}({\mathbf{y}},t)$, the $p$-th TTM reads as

where $h_{n}=t_{n+1}-t_{n}$. To generate samples from denoising diffusion models, we can, for example, apply the second TTM to the (approximate) Probability Flow ODE or the (approximate) DDIM ODE, resulting in the following respective schemes:

where ${\bm{f}}({\mathbf{x}}_{t},t)=-\tfrac{1}{2}\beta(t)\left[{\mathbf{x}}_{t}-\tfrac{{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)}{\sigma_{t}}\right]$, and

In this work, we generate samples from DDMs using the scheme in Eq. 34. We distill the derivative $d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}\coloneqq\frac{d{\bm{\epsilon}}_{\bm{\theta}}}{d\gamma_{t}}$ into a small neural network ${\bm{k}}_{\bm{\psi}}$. For training, $d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}$ is computed via automatic differentiation, however, during inference, we can efficiently query the trained network ${\bm{k}}_{\bm{\psi}}$.

Consider the $p$-TTM for a general ODE $\frac{d{\mathbf{y}}}{dt}={\bm{f}}({\mathbf{y}},t)$:

We represent, the exact solution ${\mathbf{y}}(t_{n+1})$ using the $(p+2)$-th Taylor expansion

The local truncation error (LTE) introduced by the $p$-th TTM is given by the difference between the two equations above

For small $h_{n}$, the LTE is proportional to $h_{n}^{p+1}$. Consequently, using higher orders $p$ implies lower errors, as $h_{n}$ usually is a small time step.

In conclusion, this demonstrates that it is preferable to use higher-order methods with lower errors when aiming to accurately solve ODEs like the Probability Flow ODE or the DDIM ODE of diffusion models.

B.2 Approximate Higher-Order Derivatives via the “Ideal Derivative Trick”

Tachibana et al. sample from DDMs using (an approximation to) a higher-order Itô-Taylor method . In their scheme, they approximate higher-order score functions with the “ideal derivative trick”, essentially assuming simple single-point (${\mathbf{x}}_{0}$) data distributions, for which higher-order score functions can be computed analytically (more formally, their approximation corresponds to ignoring the expectation over the full data distribution when learning the score function. They assume that for any ${\mathbf{x}}_{t}$, there is a single unique ${\mathbf{x}}_{0}$ from the input data to be predicted with the score model). In that case, further assuming the score model ${\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ is learnt perfectly (i.e., it perfectly predicts the noise that was used to generate ${\mathbf{x}}_{t}$ from ${\mathbf{x}}_{0}$), one has

This expression can now be used to analytically calculate approximate spatial and time derivatives (also see App. F.1 and App. F.2 in Tachibana et al. ):

Inserting this expression, Eq. 40 becomes

We will now proceed to show that the “ideal derivative trick”, i.e. using the approximations in Eqs. 39 and 42, results in $d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}=\bm{0}$.

As in Sec. 3, the total derivative $d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}$ is composed as

Inserting the “ideal derivative trick”, the above becomes

where we have inserted Eq. 26 for $\tfrac{d{\mathbf{x}}_{t}}{d\gamma_{t}}$ and used the usual parameterization ${\bm{s}}_{\bm{\theta}}({\mathbf{x}}_{t},t)\coloneqq-\frac{{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)}{\sigma_{t}}$. Using $\frac{d\log\alpha_{t}^{2}}{dt}=\frac{1}{\alpha_{t}^{2}}\frac{d\alpha_{t}^{2}}{dt}$ and $\frac{d\alpha_{t}^{2}}{dt}\frac{dt}{d\gamma_{t}}=\frac{d\alpha_{t}^{2}}{d\gamma_{t}}$, we can see that the right-hand side of Eq. 44 is $\bm{0}$. Hence, applying the second TTM to the DDIM ODE and using the “ideal derivative trick” is equivalent to the first TTM (Euler’s method) applied to the DDIM ODE. We believe that this is potentially a reason why the DDIM solver , Euler’s method applied to the DDIM ODE, shows such great empirical performance: it can be interpreted as an approximate (“ideal derivative trick”) second order ODE solver. On the other hand, our derivation also implies that the “ideal derivative trick” used in the second TTM for the DDIM ODE does not actually provide any benefit over the standard DDIM solver, because all additional second-order terms vanish. Hence, to improve upon regular DDIM, the “ideal derivative trick” is insufficient and we need to learn the higher-order score terms more accurately without such coarse approximations, as we do in our work.

Furthermore, it is interesting to show that we do not obtain the same cancellation effect when applying the “ideal derivative trick” to the Probability Flow ODE in Eq. 18: Let ${\bm{f}}({\mathbf{x}}_{t},t)=-\frac{1}{2}\beta(t)\left[{\mathbf{x}}_{t}-\tfrac{{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)}{\sigma_{t}}\right]$ (right-hand side of Probability Flow ODE), then

where $\beta^{\prime}(t)\coloneqq\tfrac{d\beta(t)}{dt}$. Using the “ideal derivative trick”, we have $\tfrac{d{\bm{\epsilon}}_{\bm{\theta}}}{dt}=d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}\,d_{t}\gamma_{t}\approx\bm{0}$, and therefore the above becomes

The derivative $\frac{d\sigma_{t}}{dt}$ can be computed as follows

Putting everything back together, we have

which is clearly not $\bm{0}$ for all ${\mathbf{x}}_{t}$ and $t$. Hence, in contrast to the DDIM ODE, applying Euler’s method to the Probability Flow ODE does not lead to an approximate (in the sense of the “ideal derivative trick”) second order ODE solver.

Note that very related observations have been made in the concurrent works Karras et al. and Zhang et al. . These works notice that when the data distribution consist only of a single data point or a spherical Gaussian distribution, then the solution trajectories of the generative DDIM ODE are straight lines. In fact, this exactly corresponds to our observation that in such a setting we have $d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}=\bm{0}$, as shown above in the analysis of the “ideal derivatives approximation”. Note in that context that our above derivation considers the “single data point” distribution assumption, but also applies to the setting where the data is a spherical normal distribution (only $\sigma_{t}$ would be different, which would not affect the derivation).

B.3 3rd TTM Applied to the DDIM ODE

As promised in Sec. 3, we show here how to apply the third TTM to the DDIM ODE, resulting in the following scheme:

where $h_{n}=(\gamma_{t_{n+1}}-\gamma_{t_{n}})$. In the remainder of this section, we derive a computable formula for $\frac{d^{2}{\bm{\epsilon}}_{\bm{\theta}}}{d\gamma_{t}^{2}}$, only containing partial derivatives.

The remaining terms in Eq. 55 can be computed as

where, inserting Eq. 28 for $\tfrac{d{\mathbf{x}}_{t}}{d\gamma_{t}}$ as well as using the usual parameterization ${\bm{s}}_{\bm{\theta}}({\mathbf{x}}_{t},t)\coloneqq-\frac{{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)}{\sigma_{t}}$,

We now have a formula for $\frac{d^{2}{\bm{\epsilon}}_{\bm{\theta}}}{d\gamma_{t}^{2}}$ containing only partial derivatives, and therefore we can compute $\frac{d^{2}{\bm{\epsilon}}_{\bm{\theta}}}{d\gamma_{t}^{2}}$ using automatic differentiation. Note that we could follow the same procedure to compute even higher derivatives of ${\bm{\epsilon}}_{\bm{\theta}}$.

We repeat the 2D toy distribution single step error experiment from Sec. 3 (see also Fig. 3 (top) and App. E for details). As expected, in Fig. 10 we can clearly see that the third TTM improves upon the second TTM.

In Fig. 11, we compare the second TTM to the third TTM applied to the DDIM ODE on CIFAR-10. Both for the second and the third TTM, we compute all partial derivatives using automatic differentiation (without distillation). It appears that for using 15 or less steps in the ODE solver, the second TTM performs better than the third TTM. We believe that this could potentially be due to our score model $s_{\bm{\theta}}({\mathbf{x}}_{t},t)$ not being accurate enough, in contrast to the above 2D toy distribution experiment, where we have access to the analytical score function. Furthermore, note that when we train $s_{\bm{\theta}}({\mathbf{x}}_{t},t)$ via score matching, we never regularize (higher-order) derivatives of the neural network, and therefore there is no incentive for them to be well-behaved. It would be interesting to see if, besides having more accurate score models, regularization techniques such as spectral regularization could potentially alleviate this issue. Also the higher-order score matching techniques derived by Meng et al. could help to learn higher-order derivates of the score functions more accurately. We leave this exploration to future work.

B.4 GENIE is Consistent and Principled

GENIE is a consistent and principled approach to developing a higher-order ODE solver for sampling from diffusion models: GENIE’s design consists of two parts: (1) We are building on the second Truncated Taylor Method (TTM), which is a well-studied ODE solver (see Kloeden and Platen ) with provable local and global truncation errors (see also Sec. B.1). Therefore, if during inference we had access to the ground truth second-order ODE derivatives, which are required for the second TTM, GENIE would simply correspond to the exact second TTM.

(2) In principle, we could calculate the exact second-order derivatives during inference using automatic differentiation. However, this is too slow for competitive sampling speeds, as it requires additional backward passes through the first-order score network. Therefore, in practice, we use the learned prediction heads $\mathbf{k}_{\psi}(\mathbf{x}_{t},t)$.

Consequently, if $\mathbf{k}_{\psi}(\mathbf{x}_{t},t)$ modeled the ground truth second-order derivatives exactly, i.e. $\mathbf{k}_{\psi}(\mathbf{x}_{t},t)=d_{\gamma_{t}}\mathbf{\epsilon}_{\theta}(\mathbf{x}_{t},t)$ for all $\mathbf{x}_{t}$ and $t$, we would obtain a rigorous second-order solver based on the TTM, following (1) above.

In practice, distillation will not be perfect. However, given the above analysis, optimizing a neural network $\mathbf{k}_{\psi}(\mathbf{x}_{t},t)$ towards $d_{\gamma_{t}}\mathbf{\epsilon}_{\theta}(\mathbf{x}_{t},t)$ is well motivated and theoretically grounded. In particular, during training we are calculating exact ODE gradients using automatic differentiation on the first-order score model as distillation targets. Therefore, in the limit of infinite neural network capacity and perfect optimization, we could in theory minimize our distillation objective function (Eq. 15) perfectly and obtain $\mathbf{k}_{\psi}(\mathbf{x}_{t},t)=d_{\gamma_{t}}\mathbf{\epsilon}_{\theta}(\mathbf{x}_{t},t)$.

Also recall that regular denoising score matching itself, on which all diffusion models rely, follows the exact same argument. In particular, denoising score matching also minimizes a “simple” (weighted) $L_{2}$-loss between a trainable score model $\mathbf{s}_{\theta}(\mathbf{x}_{t},t)$ and the spatial derivative of the log-perturbation kernel, i.e., $\nabla_{\mathbf{x}_{t}}\log p_{t}(\mathbf{x}_{t}\mid\mathbf{x}_{0})$. From this perspective, denoising score matching itself also simply tries to “distill” (spatial) derivatives into a model. If we perfectly optimized the denoising score matching objective, we would obtain a diffusion model that models the data distribution exactly, but in practice, similar to GENIE, we never achieve that due to imperfect optimization and finite-capacity neural networks. Nevertheless, denoising score matching similarly is a well-defined and principled method, precisely because of that theoretical limit in which the distribution can be reproduced exactly.

We would also like to point out that other, established higher-order methods for diffusion model sampling with the generative ODE, such as linear multistep methods , make approximations, too, which can be worse in fact. In particular, multistep methods always approximate higher-order derivatives in the TTM using finite differences which is crude for large step sizes, as can be seen in Fig. 3 (bottom). From this perspective, if our distillation is sufficiently accurate, GENIE can be expected to be more accurate than such multistep methods.

Appendix C Model and Implementation Details

We train variance-preserving DDMs for which $\sigma_{t}^{2}=1-\alpha_{t}^{2}$. We follow Song et al. and set $\beta(t)=0.1+19.9t$; note that $\alpha_{t}=e^{-\tfrac{1}{2}\int_{0}^{t}\beta(t^{\prime})\,dt^{\prime}}$. All score models are parameterized as either ${\bm{s}}_{\bm{\theta}}({\mathbf{x}}_{t},t)\coloneqq-\frac{{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)}{\sigma_{t}}$ (${\bm{\epsilon}}$-prediction) or ${\bm{s}}_{\bm{\theta}}({\mathbf{x}}_{t},t)\coloneqq-\frac{\alpha_{t}{\mathbf{v}}_{\bm{\theta}}({\mathbf{x}}_{t},t)+\sigma_{t}{\mathbf{x}}_{t}}{\sigma_{t}}$ (${\mathbf{v}}$-prediction), where ${\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ and ${\mathbf{v}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ are U-Nets . The ${\bm{\epsilon}}$-prediction model is trained using the following score matching objective

The ${\mathbf{v}}$-prediction model is trained using the following score matching objective

which is referred to as “SNR+1” weighting . The neural network ${\mathbf{v}}_{\bm{\theta}}$ is now effectively tasked with predicting ${\mathbf{v}}\coloneqq\alpha_{t}{\bm{\epsilon}}-\sigma_{t}{\mathbf{x}}_{0}$.

CIFAR-10: On this dataset, we do not train our own score model, but rather use a checkpointThe checkpoint can be found at https://drive.google.com/file/d/16_-Ahc6ImZV5ClUc0vM5Iivf8OJ1VSif/view?usp=sharing. provided by Song et al. . The model is based on the DDPM++ architecture introduced in Song et al. and predicts ${\bm{\epsilon}}_{\bm{\theta}}$.

LSUN Bedrooms and LSUN Church-Outdoor: Both datasets use exactly the same model structure. The model structure is based on the DDPM architecture introduced in Ho et al. and predicts ${\bm{\epsilon}}_{\bm{\theta}}$.

ImageNet: This model is based on the architecture introduced in Dhariwal and Nichol . We make a small change to the architecture and replace its sinusoidal time embedding by a Gaussian Fourier projection time embedding . The model is class-conditional and we follow Dhariwal and Nichol and simply add the class embedding to the (Gaussian Fourier projection) time embedding. The model predicts ${\bm{\epsilon}}_{\bm{\theta}}$.

Cats (Base): This model is based on the architecture introduced in Dhariwal and Nichol . We make a small change to the architecture and replace its sinusoidal time embedding by a Gaussian Fourier projection time embedding . The model predicts ${\mathbf{v}}_{\bm{\theta}}$.

We use two-independent Gaussian Fourier projection embeddings for $t$ and $t^{\prime}$ and concatenate them before feeding them into the layers of the U-Net.

Model Hyperparameters and Training Details: All model hyperparameters and training details can be found in Tab. 4.

C.2 Prediction Heads

We model the derivative $d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}$ using a small prediction head ${\bm{k}}_{\bm{\psi}}$ on top of the first-order score model ${\bm{\epsilon}}_{\bm{\theta}}$. In particular, we provide the last feature layer from the ${\bm{\epsilon}}_{\bm{\theta}}$ network together with its time embedding as well as ${\mathbf{x}}_{t}$ and the output of ${\bm{\epsilon}}({\mathbf{x}}_{t},t)$ to the prediction head (see Fig. 4 for a visualization). We found modeling $d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}$ to be effective even for our Cats models that learn to predict ${\mathbf{v}}=\alpha_{t}{\bm{\epsilon}}-\sigma_{t}{\mathbf{x}}_{0}$ rather than ${\bm{\epsilon}}$. Directly learning $d_{\gamma_{t}}{\mathbf{v}}_{\bm{\theta}}$ and adapting the mixed network parameterization (see Sec. C.2.3) could potentially improve results further. We leave this exploration to future work.

We provide additional details on our architecture next.

The architecture of our prediction heads is based on (modified) BigGAN residual blocks . To minimize computational overhead, we only use a single residual block.

In particular, we concatenate the last feature layer with ${\mathbf{x}}_{t}$ as well as ${\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ and feed it into a convolutional layer. For the upsampler, we also condition on the noisy up-scaled lower resolution image. We experimented with normalizing the feature layer before concatenation. The output of the convolutional layer as well as the time embedding are then fed to the residual block. Similar to U-Nets used in score models, we normalize the output of the residual block and apply an activation function. Lastly, the signal is fed to another convolutional layer that brings the number of channels to a desired value (in our case nine, three for each ${\bm{k}}_{\bm{\psi}}^{(i)}$, $i\in\{1,2,3\}$, in Eq. 66).

All model hyperparameters can be found in Tab. 5. We also include the additional computational overhead induced by the prediction heads in Tab. 5; see Sec. C.2.5 for details on how we measured the overhead.

C.2.2 Training Details

We train for 50k iterations using Adam . We experimented with two base learning rates: $10^{-4}$ and $5\cdot 10^{-5}$. We furthermore tried two “optimization setups”: (linearly) warming up the learning rate in the first 10k iterations (score models are often trained by warming up the learning rate in the first 100k iterations) or, following Salimans and Ho , linearly decaying the learning rate to 0 in the entire 50k iterations of training; we respectively refer to these two setups as “warmup” and “decay”. We measure the FID every 5k iterations and use the best checkpoint.

Note that we have to compute the Jacobian-vector products in Eq. 12 via automatic differentiation during training. We repeatedly found that computing the derivative $\tfrac{\partial{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)}{\partial t}$ via automatic differentiation leads to numerical instability (NaN) for small $t$ when using mixed precision training. For simplicity, we turned off mixed precision training altogether. However, training performance could have been optimized by only turning off mixed precision training for the derivative $\tfrac{\partial{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)}{\partial t}$.

All training details can be found in Tab. 5.

C.2.3 Mixed Network Parameterization

Our mixed network parameterization is derived from a simple single data point assumption, i.e., $p_{t}({\mathbf{x}}_{t})={\mathcal{N}}({\mathbf{x}}_{t};\bm{0},\sigma_{t}^{2}{\bm{I}})$. This assumption leads to ${\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)\approx\frac{{\mathbf{x}}_{t}}{\sigma_{t}}$ which we can plug into the three terms of Eq. 12:

where we have used $\sigma_{t}=\frac{\gamma_{t}}{\sqrt{\gamma_{t}^{2}+1}}$. This derivation therefore implies the following mixed network parameterization

where ${\bm{k}}_{\bm{\psi}}^{(i)}({\mathbf{x}}_{t},t)$, $i\in\{1,2,3\}$, are different output channels of the neural network (i.e. the additional head on top of the ${\bm{\epsilon}}_{\bm{\theta}}$ network). To provide additional intuition, we basically replaced the $-\frac{{\mathbf{x}}_{t}}{\sigma_{t}}$ terms in Eqs. 63, 64 and 65 by neural networks. However, we know that for approximately Normal data $\frac{{\mathbf{x}}_{t}}{\sigma_{t}}\approx{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$, where ${\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ predicts “noise” values ${\bm{\epsilon}}$ that were drawn from a standard Normal distribution and are therefore varying on a well-behaved scale. Consequently, up to the Normal data assumption, we can also expect our prediction heads ${\bm{k}}_{\bm{\psi}}^{(i)}({\mathbf{x}}_{t},t)$ in the parameterization in Eq. 66 to predict well-behaved output values, which should make training stable. This mixed network parameterization approach is inspired by the mixed score parameterization from Vahdat et al. and Dockhorn et al. .

C.2.4 Pseudocode

In this section, we provide pseudocode for training our prediction heads ${\bm{k}}_{\bm{\psi}}$ and using them for sampling with GENIE. In Alg. 1, the analytical $\frac{dt}{d\gamma_{t}}$ is an implicit hyperparameter of the DDM as it depends on $\alpha_{t}$. For our choice of $\alpha_{t}=e^{-\tfrac{1}{2}\int_{0}^{t}0.1+19.9t^{\prime}\,dt^{\prime}}$ (see Sec. C.1), we have

where $\gamma_{t}=\sqrt{\frac{1-\alpha_{t}^{2}}{\alpha_{t}^{2}}}$.

C.2.5 Measuring Computational Overhead

Our prediction heads induce a slight computational overhead since their forward pass has to occur after the forward pass of the score model. We measure the overhead as follows: first, we measure the inference time of the score model itself. We do five forward passes to “warm-up” the model and then subsequently synchronize via torch.cuda.synchronize(). We then measure the total wall-clock time of 50 forward passes. We then repeat this process using a combined forward pass: first the score model and subsequently the prediction head. We choose the batch size to (almost) fill the entire GPU memory. In particular we chose batch sizes of 512, 128, 128, 64, 64, and 8, for CIFAR-10, LSUN Bedrooms, LSUN Church-Outdoor, ImageNet, Cats (base), and Cats (upsampler), respectively. The computational overhead for each model is reported in Tab. 5. This measurement was carried out on a single NVIDIA 3080 Ti GPU.

Appendix D Learning Higher-Order Gradients without Automatic Differentiation and Distillation

In this work, we learn the derivative $d_{\gamma_{t}}{\bm{\epsilon}}_{\bm{\theta}}$, which includes a spatial and a temporal Jacobian-vector product, by distillation based on automatic differentiation (AD). We now derive an alternative learning objective for the spatial Jacobian-vector product (JVP) which does not require any AD. We start with the following (conditional) expectation

where ${\bm{s}}_{1}({\mathbf{x}}_{t},t)\coloneqq\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})$ and ${\bm{S}}_{2}({\mathbf{x}}_{t},t)\coloneqq\nabla_{{\mathbf{x}}_{t}}^{\top}\nabla_{{\mathbf{x}}_{t}}\log p_{t}({\mathbf{x}}_{t})$. The above formula is derived in Meng et al. [Theorem 1, 82]. Adding ${\mathbf{x}}_{t}{\mathbf{x}}_{t}^{\top}$ to Eq. 68 and subsequently dividing by $\sigma_{t}^{2}$, we have

where we could pull the $\frac{1}{\sigma_{t}^{2}}{\mathbf{x}}_{t}{\mathbf{x}}_{t}^{\top}$ term into the expectation because it is conditioned on $t$ and ${\mathbf{x}}_{t}$. Using ${\mathbf{x}}_{t}=\alpha_{t}{\mathbf{x}}_{0}+\sigma_{t}{\bm{\epsilon}}$, we can rewrite the above as

For an arbitrary ${\bm{v}}\coloneqq{\bm{v}}({\mathbf{x}}_{t},t)$, we then have

Therefore, we can develop a score matching-like learning objective for the (general) spatial JVP ${\bm{o}}_{\bm{\theta}}({\mathbf{x}}_{t},t)\approx{\bm{S}}_{2}({\mathbf{x}}_{,}t){\bm{v}}$ as

Appendix E Toy Experiments

For all toy experiments in Sec. 3, we consider the following ground truth distribution:

We set $\sigma=10^{-2}$, $s_{1}=0.9$, $s_{2}=0.2$, and

The ground truth distribution is visualized in Fig. 2(a). Note that we can compute the score functions (and all its derivatives) analytically for Gaussian mixture distributions.

In Fig. 2, we compared DDIM to GENIE for sampling using the analytical score function of the ground truth distribution with 25 solver steps. In Fig. 12, we repeated this experiment for 5, 10, 15, and 20 solver steps. We found that in particular for $n=10$ both solvers generate samples in interesting patterns.

Appendix F Image Experiments

Metrics: We quantitatively measure sample quality via Fréchet Inception Distance [FID, 102]. It is common practice to use 50k samples from the training set for reference statistics. We follow this practice for all datasets except for ImageNet and Cats. For ImageNet, we follow Dhariwal and Nichol and use the entire training set for reference statistics. For the small Cats dataset, we use the training as well as the validation set for reference statistics.

Baselines: We run baseline experiments using two publicly available repositories. The score_sde_pytorch repository is licensed according to the Apache License 2.0; see also their license file here. The CLD-SGM repository is licensed according to the NVIDIA Source Code License; see also their license file here.

Datasets: We link here the websites of the datasets used in this experiment: CIFAR-10, LSUN datasets, ImageNet, and AFHQv2.

F.2 Analytical First Step (AFS)

The forward process of DDMs generally converges to an analytical distribution. This analytical distribution is then used to sample from DDMs, defining the initial condition for the generative ODE/SDE. For example, for variance-preserving DDMs, we have $p_{1}({\mathbf{x}}_{1})\approx{\mathcal{N}}({\mathbf{x}}_{1};\bm{0},{\bm{I}})$.

In this work, we try to minimize the computational complexity of sampling from DDMs, and therefore operate in a low NFE regime. In this regime, every additional function evaluation makes a significant difference. We therefore experimented with replacing the learned score with the (analytical score) of ${\mathcal{N}}(\bm{0},{\bm{I}})\approx p_{1}({\mathbf{x}}_{1})$ in the first step of the ODE solver. This “gained” function evaluation can then be used as an additional step in the ODE solver later.

and $\frac{d{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{1},1)}{d\gamma_{1}}\approx\bm{0}$ as shown below:

Given this, the AFS step becomes identical to the Euler update that uses the Normal score function for ${\mathbf{x}}_{1}$. This step is shown in the pseudocode in Alg. 2.

F.3 Classifier-Free Guidance

As discussed in Sec. 5.2, to guide diffusion sampling towards particular classes, we replace ${\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)$ with

where $w>1.0$ is the “guidance scale”, in the DDIM ODE. We experiment with classifier-free guidance on ImageNet. In Eq. 79 we re-use the conditional ImageNet score model ${\bm{\epsilon}}_{{\bm{\theta}}}({\mathbf{x}}_{t},t,c)$ trained before (see Sec. C.1 for details), and train an additional unconditional ImageNet score model ${\bm{\epsilon}}_{{\bm{\theta}}}({\mathbf{x}}_{t},t)$ using the exact same setup (and simply setting the class embedding to zero). We also re-use the conditional prediction head trained on top of the conditional ImageNet score model and train an additional prediction head for the unconditional model. Note that for both the score models as well as the prediction heads, we could share parameters between the models to reduce computational complexity . The modified GENIE scheme for classifier-free guidance is then given as

F.4 Encoding

To encode a data point ${\mathbf{x}}_{0}$ into latent space, we first “diffuse” the data point to $t=10^{-3}$, i.e., ${\mathbf{x}}_{t}=\alpha_{t}{\mathbf{x}}_{0}+\sigma_{t}{\bm{\epsilon}}$, ${\bm{\epsilon}}\sim{\mathcal{N}}(\bm{0},{\bm{I}})$. We subsequently simulate the generative ODE (backwards) from $t=10^{-3}$ to $t=1$, obtaining the latent point ${\mathbf{x}}_{1}$.

To decode a latent point ${\mathbf{x}}_{1}$, we simulate the generative ODE (forwards) from $t=1.0$ to $t=10^{-3}$. We then denoise the data point, i.e., ${\mathbf{x}}_{0}=\frac{{\mathbf{x}}_{t}-\sigma_{t}{\bm{\epsilon}}_{\bm{\theta}}({\mathbf{x}}_{t},t)}{\alpha_{t}}$. Note that denoising is generally optional to sample from DDMs; however, for our encoding-decoding experiment we always used denoising in the decoding part to match the inital “diffusion” in the encoding part.

F.5 Latent Space Interpolation

We can use encoding to perform latent space interpolation of two data points ${\mathbf{x}}_{0}^{(0)}$ and ${\mathbf{x}}_{0}^{(1)}$. We first encode both data points, following the encoding setup from Sec. F.4, and obtain ${\mathbf{x}}_{1}^{(0)}$ and ${\mathbf{x}}_{1}^{(1)}$, respectively. We then perform spherical interpolation of the latent codes:

Subsequently, we decode the latent code ${\mathbf{x}}_{1}^{(b)}$ following the decoding setup from Sec. F.4. In Fig. 13, we show latent space interpolations for LSUN Church-Outdoor and LSUN Bedrooms.

F.6 Extended Quantitative Results

In this section, we show additional quantitative results not presented in the main paper. In particular, we show results for all four hyperparameter combinations (binary choice of AFS and binary choice of denoising) for methods evaluated by ourselves. For these methods (i.e., GENIE, DDIM, S-PNDM, F-PNDM, Euler–Maruyama), we follow the Synthesis Strategy outlined in Sec. 5, with the exception that we use linear striding instead of quadratic striding for S-PNDM and F-PNDM . To apply quadratic striding to these two methods, one would have to derive the Adams–Bashforth methods for non-constant step sizes which is beyond the scope of our work.

Results can be found in Tabs. 8, 9, 10, 11, 12 and 13. As expected, AFS can considerably improve results for almost all methods, in particular for NFEs $\leq 15$. Denoising, on the other hand, is more important for larger NFEs. For our Cats models, we initially found that denoising hurts performance, and therefore did not further test it in all settings.

Recall Scores. We quantify the sample diversity of GENIE and other fast samplers using the recall score . In particular, we follow DDGAN and use the improved recall score ; results on CIFAR-10 can be found in Tab. 6. As expected, we can see that for all methods recall scores suffer as the NFEs decrease. Compared to the baselines, GENIE achieves excellent recall scores, being on par with F-PNDM for NFE$\geq 15$. However, F-PNDM cannot be run for NFE$\leq$10 (due to its additional Runge–Kutta warm-up iterations). Overall, these results confirm that GENIE offers strong sample diversity when compared to other common samplers using the same score model checkpoint.

In particular, besides the quadratic schedule $\rho=2$, we also tested the two additional values $\rho=1.5$ and $\rho=2.5$. We tested these schedules on GENIE as well as DDIM ; note that the other two comptetive baselines, S-PNDM and F-PNDM , rely on linear striding, and therefore a grid search is not applicable. We show results for GENIE and DDIM in Tab. 7; for each combination of solver and NFE we applied the best synthesis strategy (whether or not we use denoising and/or the analytical first step) of quadratic striding ($\rho=2.0$) also to $\rho=1.5$ and $\rho=2.5$. As can be seen in the table, $\rho=1.5$ improves for both DDIM and GENIE for NFE$=$5 (over the quadratic schedule $\rho=2$), whereas larger $\rho$ are preferred for larger NFE. The improvement of GENIE from 13.9 to 11.2 FID for NFE=5 is significant.

Discretization Errors of GENIE compared to other Fast Samplers. We compute discretization errors, in particular local and global truncation errors, of GENIE and compare to existing faster solvers. We are using the CIFAR-10 model. We initially sample 100 latent vectors $\mathbf{x}_{T}\sim\mathcal{N}(\bm{0},\bm{I})$ and then, starting from those latent vectors, synthesize 100 approximate ground truth trajectories (GTTs) using DDIM with 1k NFEs (for that many steps, the discretization error is negligible; hence, we can treat this as a pseudo ground truth).

We then synthesize 100 sample trajectories for DDIM , S-PNDM , F-PNDM , and GENIE (for NFEs=$\{5,10,15,20,25\}$, similar to the main experiments) using the same latent vectors as starting points that were used to generate the GTTs. DDIM, S-PNDM, and F-PNDM are training-free methods that can be run on the exact same score model, which also our GENIE relies on. Thereby, we are able to isolate discretization errors from errors in the learnt score function. We then compute the average $L_{2}$-distance (in Inception feature space ) between the output image of the fast samplers and the “output” of the pseudo GTT. As can be seen in Fig. 14, GENIE outperforms the three other methods on all NFEs.

Comparing the local truncation error (LTE) of different higher-order solvers can unfortunately not be done in a fair manner. Similar to DDIM, GENIE only needs the current value and a single NFE to predict the next step. In contrast, multistep methods rely on a history of predictions and Runge–Kutta methods rely on multiple NFEs to predict the next step. Thus, we can only fairly compare the LTE of GENIE to the LTE of DDIM. In particular, we compute LTEs at three starting times $t\in\{0.1,0.2,.5\}$ (similar to what we did in Fig. 3). For each $t$, we then compare one step predictions for different step sizes $\Delta t$ against the ground truth trajectory ($L_{2}$-distance in data space averaged over 100 predictions; since we are not operating directly in image space at these intermediate $t$, using inception feature would not make sense here). As expected, we can see in Fig. 15 that GENIE has smaller LTE than DDIM for all starting times $t$.

F.7 Extended Qualitative Results

In this section, we show additional qualitative comparisons of DDIM and GENIE on LSUN Church-Outdoor (Fig. 16), ImageNet (Fig. 17), and Cats (upsampler conditioned on test set images) (Fig. 18 and Fig. 19). In all figures, we can see that samples generated with GENIE generally exhibit finer details as well as sharper contrast and are less blurry compared to standard DDIM.

In Fig. 20 and Fig. 21, we show additional high-resolution images generated with the GENIE Cats upsampler using base model samples and test set samples, respectively.

F.8 Computational Resources

The total amount of compute used in this research project is roughly 163k GPU hours. We used an in-house GPU cluster of V100 NVIDIA GPUs.

Appendix G Miscellaneous

The concurrent Bao et al. learn covariance matrices for diffusion model sampling using prediction heads somewhat similar to the ones in GENIE. Specifically, both Bao et al. and GENIE use small prediction heads that operate on top of the large first-order score predictor. However, we would like to stress multiple differences: (i) Bao et al. learn the DDM’s sampling covariance matrices, while we learn higher-order ODE gradients. More generally, Bao et al. rely on stochastic diffusion model sampling, while we use the ODE formulation. (ii) Most importantly, in our case we can resort to directly learning the low-dimensional JVPs without low-rank or diagonal matrix approximations or other assumptions. Similar techniques are not directly applicable in Bao et al. ’s setting. In detail, this is because in their case the relevant matrices (obtained after Cholesky or another applicable decomposition of the covariance) do not act on regular vectors but random noise variables. In other words, instead of using a deterministic JVP predictor (which takes ${\mathbf{x}}_{t}$ and $t$ as inputs), as in GENIE, Bao et al. would require to model an entire distribution for each ${\mathbf{x}}_{t}$ and $t$ without explicitly forming high-dimensional Cholesky decomposition-based matrices, if they wanted to do something somewhat analogous to GENIE’s novel JVP-based approach. As a consequence, Bao et al. take another route to keeping the dimensionality of the additional network outputs manageable in practice. In particular, they resort to assuming a diagonal covariance matrix in their experiments. By directly learning JVPs, we never have to rely on such potentially limiting assumptions. (iii) Experimentally, Bao et al. also consider fast sampling with few neural network calls. However, GENIE generally outperforms them (see, for example, their CIFAR10 results in their Table 2 for 10 and 25 NFE). This might indeed be due to the assumptions made by Bao et al. , which we avoid. Furthermore, their stochastic vs. our deterministic sampling may play a role, too.

G.2 Combining GENIE with Progressive Distillation

We speculate that GENIE could potentially be combined with Progressive Distillation : In every distillation stage of , one could quickly train a small GENIE prediction head to model higher-order ODE gradients. This would then allow for larger and/or more accurate steps, whose results represent the distillation target (teacher) in the progressive distillation protocol. This may also reduce the number of required distillation stages. Overall, this could potentially speed up the cumbersome stage-wise distillation and maybe also lead to an accuracy and performance improvement. In particular, we could replace the DDIM predictions in Algorithm 2 of with improved GENIE predictions.

Note that this approach would not be possible with multistep methods as proposed by Liu et al. . Such techniques could not be used here, because they require the history of previous predictions, which are not available in the progressive distillation training scheme.

We leave exploration of this direction to future work.