GENIE: Higher-Order Denoising Diffusion Solvers
Abstract
Denoising diffusion models (DDMs) have emerged as a powerful class of generative models.
A forward diffusion process slowly perturbs the data, while a deep model learns to gradually denoise.
Synthesis amounts to solving a differential equation (DE) defined by the learnt model.
Solving the DE requires slow iterative solvers for high-quality generation.
In this work, we propose Higher-Order Denoising Diffusion Solvers (GENIE):
Based on truncated Taylor methods, we derive a novel higher-order solver that significantly accelerates synthesis.
Our solver relies on higher-order gradients of the perturbed data distribution, that is, higher-order score functions.
In practice, only Jacobian-vector products (JVPs) are required and we propose to extract them from the first-order score network via automatic differentiation.
We then distill the JVPs into a separate neural network that allows us to efficiently compute the necessary higher-order terms for our novel sampler during synthesis.
We only need to train a small additional head on top of the first-order score network.
We validate GENIE on multiple image generation benchmarks and demonstrate that GENIE outperforms all previous solvers.
Unlike recent methods that fundamentally alter the generation process in DDMs, our GENIE solves the true generative DE and still enables applications such as encoding and guided sampling.
Project page and code: https://nv-tlabs.github.io/GENIE .
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 [1–5] and video [6, 7] synthesis, super-resolution [8, 9], deblurring [10, 11], image editing and inpainting [5, 12–14], text-to-image synthesis [15–17], conditional and semantic image generation [18–22], imageto-image translation [14, 23, 24] and for inverse problems in medical imaging [25–31].
They also enable high-quality speech synthesis [32–37], 3D shape generation [38–42], molecular modeling [43– 46], maximum likelihood training [47–50], and more [51–56].
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 [47, 57].
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 [47, 57, 58], we rigorously derive a novel second-order ODE solver using truncated Taylor methods [59].
These higher-order methods require higher-order gradients of the ODE—in our case this includes higher-order gradients of the
NVIDIA.
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 [1, 57, 58, 60], higher-order methods that rely on suboptimal finite difference or other approximations [61–63], or alternative approaches [64–66] 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-ofthe-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 [67] or by distilling the full sampling trajectory [68, 69], 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 [4, 57, 70].
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 [1, 57, 71] whose forward process can be described by
[Equation]
where x 0 ∼ p 0 ( x 0 ) is drawn from the empirical data distribution and x t refers to diffused data samples at time t ∈ [0 , 1] along the diffusion process.
The functions α t and σ t are generally chosen such that the logarithmic signal-to-noise ratio [48] log α 2 t σ 2 t decreases monotonically with t and the data diffuses towards random noise, i.e., p 1 ( x 1 ) ≈N ( x 1 ; 0 , I ) .
We use variance-preserving [57] diffusion processes for which σ 2 t = 1 − α 2 t (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
[Equation]
where β t = − ddt log α 2 t , x 0 ∼ p 0 ( x 0 ) and w t is a standard Wiener process.
A corresponding reverse diffusion process that effectively inverts the forward diffusion is given by [57, 72, 73]
[Equation]
and this reverse-time generative SDE is marginally equivalent to the generative ODE [47, 57]
[Equation]
where ∇ x t log p t ( x t ) is the score function .
Eq.
(4) is referred to as the Probability Flow ODE [57], an instance of continuous Normalizing flows [74, 75].
To generate samples from the DDM, one can sample x 1 ∼ N ( x 1 ; 0 , I ) and numerically simulate either the Probability Flow ODE or the generative
SDE, replacing the unknown score function by a learned score model
s θ ( x t , t )
≈ ∇ x t log p t ( x t ) .
The DDIM solver [58] 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 [58, 69]:
Defining γ t = (cid:113) 1 − α 2 t α 2 t and ¯ x t = x t (cid:112) 1 + γ 2 t , we have
[Equation]
[Equation]
where we inserted Eq.
(4) for d x t dt and used β ( t ) dtdγ t = 2 γ t γ 2 t +1 .
Letting s θ ( x t , t ) :
= − ϵ θ ( x t ,t ) σ t denote a parameterization of the score model, the approximate generative DDIM ODE is then given by
[Equation]
where we used σ t = (cid:112) 1 − α
[Equation]
= γ t √ γ 2 t +1 (see App.
A for a more detailed derivation of Eq.
The model ϵ θ ( x t , t ) can be learned by minimizing the score matching objective [1, 76]
[Equation]
for small 0 < t cutoff ≪ 1 .
As is standard practice, we set g ( t ) = 1 .
Other weighting functions g ( t ) are possible; for example, setting g ( t ) =
β t 2 σ 2 t recovers maximum likelihood learning [47–50].
Higher-Order Denoising Diffusion Solver
As discussed in Sec. 2, the so-known DDIM solver [58] is simply Euler’s method applied to the DDIM
ODE (cf.
Eq.
In this work, we apply a higher-order method to the DDIM ODE, building on the truncated Taylor method (TTM) [59].
The p -th TTM is simply the p -th order Taylor polynomial applied to an ODE.
For example, for the general d y dt = f
( y , t ) , the p -th TTM reads as
[Equation]
where h n = t n +1 − t n (see App.
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:
[Equation]
(9) where h n = γ t n +1 − γ t n .
Recall that γ t = (cid:113) 1 − α 2 t α 2 t , where the function α t is a time-dependent hyperparameter of the DDM.
The total derivative d γ t ϵ θ := d ϵ θ dγ t can be decomposed as follows
[Equation]
where ∂
ϵ θ ( x t ,t ) ∂ x t denotes the Jacobian of ϵ θ ( x t , t ) and
[Equation]
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 [57] (also see App. B) or neural ODEs [74, 75] more generally.
[Figure]
[Figure]
The Benefit of Higher-Order Methods: We showcase the benefit of higher-order methods on a 2D toy distribution (Fig. 2a) 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 x t using DDIM with 10k steps.
We compare these “ground truth” trajectories to single steps of DDIM and GENIE for varying step sizes ∆ t .
We then measure the mean
L 2 -distance of the single steps ˆ x t (∆ t ) to the “ground truth” trajectories x t , and we repeat this experiment for three starting points t ∈ { 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, usingsolver 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. 2a).
Comparison to Multistep Methods: Linear multistep methods are an alternative higher-order method to solve ODEs.
Liu et al.
[63] 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 d p y dt p using the finite difference method [78].
For example, the second AB method is obtained from the second TTM by replacing d 2 y dt 2 with the first-order forward difference approximation ( f ( y t n , t n ) − f ( y t n − 1 , t n − 1 )) /h
n − 1 .
In Fig.
3 ( bottom ), we visualize the mean L 2 -norm of the difference ξ t (∆ t ) between the analytical derivative d γ t ϵ θ and its first-order forward difference approximation for varying step sizes ∆ 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).
3.1 Learning Higher-Order Derivatives
The above observations inspire to apply GENIE to DDMs of more complex and high-dimensional data such as images.
Regular DDMs learn a model ϵ θ 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:
[Equation] [Equation]
∂t dt
[Equation]
We see that the full derivative decomposes into two JVP terms and one simpler time derivative term.
The term ∂ ϵ θ ( x t ,t ) ∂ x t plays a crucial role in Eq.
It can be expressed as
[Equation]
which means that GENIE relies on second-order score functions ∇ ⊤ x t ∇ x t log p t ( x t ) under the hood.
Given a DDM, that is, given ϵ θ , we could compute the derivative d γ t ϵ θ 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 ϵ θ network to compute ϵ θ ( x t , t ) itself, and another pass to compute the JVPs and the time derivative in Eq.
These forward passes cannot be parallelized, since the vector-part of JVP 1 in Eq.
(12) involves ϵ θ 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 γ t ϵ θ into a separate neural network.
During distillation training, we can use the slow AD-based calculation of d γ t ϵ θ , but during synthesis we call the trained neural network.
We build on the observation that the internal representations of the neural network modeling ϵ θ (in our case a U-Net [79] architecture) can be used for downstream tasks [80, 81]: specifically, we provide the last feature layer from the ϵ θ network together with its time embedding as well as x t and the output ϵ θ ( x t , t ) to a small prediction head k ψ ( x t , t ) that models the different terms in Eq.
(12) (see Fig. 4).
The overhead generated by k ψ 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 ϵ θ and could therefore theoretically be run in parallel to the ϵ θ 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 ϵ θ network, then calculate the JVP terms, and only then can we call the distillation model.
By modeling d γ t ϵ θ via small prediction heads, while reusing the internal representation of the score model, we can make training relatively fast: we only need to train k ψ 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 γ t ϵ θ to future work.
Mixed Network Parameterization:
We found that learning d γ t ϵ θ directly as single output of a neural network can be challenging.
Assuming a single data point distribution p 0 ( x 0 )
=
δ ( x 0 = 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 ∈ [0 , 1] interval (for instance, the prefactor of JVP 1 in Eq.
(12) approaches 1 as t → 0 , while JVP 2 ’s prefactor vanishes).
As outlined in detail in App.
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.
[49] and Dockhorn et al.
In particular, we model
[Equation]
(14) where k ( i ) ψ ( x t , t ) , i ∈ { 1 , 2 , 3 } , are different output channels of the neural network (i.e. the additional head on top of the ϵ θ 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.
Learning Objective: Ideally, we would like our model k ψ to match d γ t ϵ
θ exactly, for all t ∈ [0 , T ] and 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 [1, 57]: min
[Equation]
ϵ , t ) − d γ t ϵ θ ( α t x 0 + σ t ϵ , t ) ∥ 22 (cid:3) (15) for diffused data points α t x 0 + σ t ϵ and g d ( t )
=
γ 2 t to counteract the 1 /γ t in the first and third terms of Eq.
This leads to a roughly constant loss over different time values t .
During training we compute d γ t ϵ θ via AD; however, at inference time we use the learned prediction head k ψ to approximate d γ t
ϵ θ .
In App.
C.2.4, we provide pseudo code for training and sampling with heads k ψ .
Note that our distillation objective is consistent and principled: if k ψ matches d γ t ϵ θ exactly, the resulting GENIE algorithm recovers the second TTM exactly (extended discussion in App. B.4).
Alternative Learning Approaches: As shown in Eq.
(13), GENIE relies on second-order score functions.
Recently, Meng et al.
[82] 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 ∇
⊤ x t ϵ θ ( 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.
[82] 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 [3, 64] used in time-discretized DDMs [1, 71].
This can be done systematically by grid search [32] or dynamic programming [83].
Bao et al.
[65] speed up sampling by defining a new DDM with optimal reverse variances.
DDIM [58], 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 [57, 84] and adaptive step size SDE solvers [62] 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.
[63], which simulates the DDIM ODE [58] using a higher-order linear multistep method [77].
As shown in Sec. 3, linear multistep methods can be considered an approximation of the TTMs used in GENIE.
Furthermore, Tachibana et al.
[61] solve the generative SDE via a higher-order Itô–Taylor method [59] and in contrast to our work, they propose to use an “ideal derivative trick” to approximate higher-order score functions.
In App.
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 [57], as the base ODE for GENIE.
Alternatively, sampling from DDMs can also be accelerated via learning: For instance, Watson et al.
[66] learn parameters of a generalized family of DDMs by optimizing for perceptual output quality.
Luhman and Luhman [68] and Salimans and Ho [69] distill a DDIM sampler into a student model, which enables sampling in as few as a single step.
Xiao et al.
[67] 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 [4, 57, 70] and to efficiently encode data into latent space—a prerequisite for likelihood calculation [47, 57] and editing applications [17].
Note that the learnt sampler [66] 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 [60, 85, 86] or train DDMs in the latent space of a Variational Autoencoder [49].
GENIE is complementary to these methods.
Higher-Order ODE Gradients beyond DDMs.
TTMs [78] and other methods that leverage higherorder gradients are also applied outside the scope of DDMs.
For instance, higher-order derivatives can play a crucial role when developing solvers [87] and regularization techniques [88, 89] for neural
ODEs [74, 75].
Outside the field of machine learning, higher-order TTMs have been widely studied, for example, to develop solvers for stiff [90] and non-stiff [90, 91] systems.
Concurrent Works.
Zhang and Chen [92] 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.
[93] 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.
[94] propose new fast solvers, both deterministic and stochastic, specifically designed for the differential equations arising in DDMs.
Both Zhang et al.
[95] and Karras et al.
[94] 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 App.
B.2).
Bao et al.
[96] learn covariance matrices for DDM sampling using prediction heads somewhat similar to the ones in GENIE; in App.
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 [97] (resolution 32), LSUN Bedrooms [98] (128), LSUN Church-Outdoor [98] (128), (conditional) ImageNet [99] (64), and AFHQv2 [100] (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.
[57]), we train our own score models using architectures introduced by previous works [1, 4].
The architecture of our prediction heads is based on (modified)
BigGAN
residual blocks [57, 101].
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 App. 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 [58]).
For variance-preserving DDMs, it can be beneficial to denoise the ODE solver output at the cutoff t =10 − 3 , i.e., x 0 = x
t − σ t ϵ θ ( x t ,t ) α t [57, 103].
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 N ( 0 , I )
≈ p t =1 ( 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 App.
F.2.
Accounting for Computational Overhead: GENIE has a slightly increased computational overhead compared to other solvers due to the prediction head k ψ .
The computational overhead is increased by 1.47%, 2.83%, 14.0%, and 14.4% on CIFAR-10, ImageNet, LSUN Bedrooms, and LSUN ChurchOutdoor, respectively (see also App. 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.
5.1 Image Generation
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 [58], S-PNDM [63], and F-PNDM [63].
For these four methods, we only include the best result over the two hyperparameters discussed above, namely, the denoising step and AFS (see App.
F.6 for tables with all results).
We also include three competitive results from the literature [64–66] 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 ∈ { 10 , 15 , 20 , 25 } , GENIE outperforms all baselines (on the same checkpoint) on all four datasets (see detailed results in App.
F.6 and GENIE image samples in App. F.7).
On CIFARand (conditional) ImageNet, GENIE also outperforms these baselines for NFEs=5, whereas DDIM outperforms GENIE slightly on the LSUN datasets (see tables in App. 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 App.
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.
[62] obtains an FID of 82.4 at 48 NFEs.
Also on the same checkpoint, the adaptive Runge–Kutta 4(5)
[84] method applied to the ProbabilityFlow
ODE achieves an FID of 13.1 at 38 NFEs (solver tolerances set to 10 − 2 ).
[Figure]
CIFARImageNet
[Figure]
GENIE (ours)
DDIM [58]
S-PNDM [63] F-PNDM [63]
Learned Sampler [66] (†)
FastDDIM [64] (†)
Analytic DDIM [65] (†)
[Figure] [Figure]
Euler–Maruyama
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
Analytic DDIM [65] (†)
[Figure]
[Table]
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 [63] and F-PNDM [63] 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.
Approaches such as FastDDIM [64] and AnalyticDDIM [65], 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 [67, 69] which are optimized for the single use-case of image synthesis.
5.2 Guidance and Encoding
[Figure]
As discussed in Sec. 4, one major drawback of approaches such as KD [68], PG [69] and DDGAN [67] is that they abandon the ODE/SDE formalism, and cannot easily use methods such as classifier(-free) guidance [57, 70] or perform image encoding.
However, these techniques can play an important role in synthesizing photorealistic images from DDMs [3, 4, 15, 17], as well as for image editing tasks [12, 17].
Classifier-Free Guidance [70]: We replace the unconditional model ϵ θ ( x t , t ) with ˆ ϵ θ ( x t , t, c, w ) =
(1 + w ) ϵ θ ( x t , t, c ) − w ϵ θ ( x t , t ) in the DDIM ODE (cf Eq.
(6)), where ϵ θ ( x t , t, c ) is a conditional model and w > 1 . 0 is the “guidance scale”.
GENIE then requires the derivative
[Figure]
[Figure] [Figure]
[Equation]
for guidance.
Hence, we need to distill d γ t ϵ θ ( x t , t, c ) and d γ t ϵ θ ( x t , t ) , for which we could also share parameters [70].
We compare GENIE with DDIM on ImageNet in Fig.
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.
For more details on this experiment as well as the guidance experiment above, see App.
F.4 and App.
F.3, respectively.
We also show latent space interpolations for both GENIE and DDIM in App. F.5.
5.3 Ablation Studies
[Table]
We perform ablation studies over architecture and training objective for the prediction heads used in GENIE:
In Tab. 2, “No mixed” refers to learning d γ t
ϵ θ directly as single network output without mixed network parameterization
; “No weighting” refers to setting g d ( t ) = 1 in Eq.
(15); “Standard” uses both the mixed network parameterization and the weighting function g d ( t )
= γ 2 t .
We can see that having both the mixed network parametrization and the weighting function is clearly beneficial.
We also tested deeper networks in the prediction heads: for “Bigger model” we increased the number of residual blocks from one to two.
The performance is roughly on par with “Standard”, and we therefore opted for the smaller head due to the lower computational overhead.
5.4 Upsampling
[Table]
Cascaded diffusion model pipelines [2] and DDM-based super-resolution [8] have become crucial ingredients in
DDMs for large-scale image generation [105].
Hence, we also explore the applicability of GENIE in this setting.
We train a 128 × 128 base model as well as a 128 × 128 → 512 × 512 diffusion upsampler [2, 8] 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.
For more quantitative and qualitative results, we refer to App.
F.6 and App.
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 App.
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 [106], Nguyen et al.
[107], Mirsky and Lee [108].
Acknowledgements
This work was funded by NVIDIA.
Tim Dockhorn acknowledges additional funding from the Vector Institute Research Grant, not in direct support of this work. We thank Yaoliang Yu for discussions.
References
A DDIM ODE
B Synthesis from Denoising Diffusion Models via Truncated Taylor Methods
B.1 Theoretical Bounds for the Truncated Taylor Method . .
B.2 Approximate Higher-Order Derivatives via the “Ideal Derivative Trick” . . . . . .
B.3 3rd TTM Applied to the DDIM ODE . . . . . . . . . . . . . . . . . . . . . . . . .
B.4 GENIE is Consistent and Principled . . . . . . . . . . . . . . . . . . . . . . . . .
C Model and Implementation Details
C.1 Score Models . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C
Prediction
Heads . .
C.2.
Model Architecture . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.2.
Training Details . .
C.2.
Mixed Network Parameterization . .
C.2.
Pseudocode . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.2.
Measuring Computational Overhead . . . . . . . . . . . . . . . . . . . . .
D Learning Higher-Order Gradients without Automatic Differentiation and DistillationE Toy Experiments
F Image Experiments
F.
Evaluation Metrics, Baselines, and Datasets . . . . . . . . . . . . . . . . . . . . .
F.
Analytical First Step (AFS) . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
F.
Classifier-Free Guidance . .
F.
Encoding . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
F.
Latent Space Interpolation . .
F.
Extended Quantitative Results . . . . . . . . . . . . . . . . . . . . . . . . . . . .
F.
Extended Qualitative Results . .
F.
Computational Resources . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
G Miscellaneous
G.1 Connection to Bao et al.
G.2 Combining GENIE with Progressive Distillation . . . . . . . . . . . . . . . . . . .
A DDIM ODE
The DDIM ODE has previously been shown [58, 69] to be a re-parameterization of the Probability
Flow ODE [57].
In this section, we show an alternative presentation to the ones given in Song et al.
[58] and Salimans and Ho [69].
We start from the Probability Flow ODE for variance-preserving continuous-time DDMs [57], i.e.,
[Equation]
17) where β t = − ddt log α 2 t and ∇ x t log p t ( x t ) is the score function .
Replacing the unknown score func- tion with a learned score model s θ ( x t , t ) ≈ ∇ x t log p t ( x t ) , we obtain the approximate Probability Flow ODE
[Equation]
[Equation]
Let us now define γ t = (cid:113) 1 − α 2 t α 2 t and ¯ x t = x t (cid:112) 1 + γ 2 t , and take the (total) derivative of ¯ x t with respect to γ t :
[Equation]
x t
[Equation]
The derivative d x t dγ t can be computed as follows
[Equation] [Equation]
( inserting Eq.
[Equation]
( by definition of β t ) (23)
[Equation]
We can write α 2 t as a function of γ t , i.e., α 2 t = (cid:0) γ 2 t + 1 (cid:1) − 1 , and therefore
[Equation]
Inserting Eq.
(27) into Eq.
(26), we obtain
[Equation] [Equation]
Lastly, inserting Eq.
(28) into Eq.
(20), we have
[Equation]
Letting s θ ( x t , t ) :
= − ϵ θ ( x t ,t ) σ t , where σ t = (cid:112) 1 − α 2 t = γ t √ γ 2 t +1 , denote a particular parameterization of the score model, we obtain the approximate generative DDIM ODE as
[Equation]
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) [78].
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 d y dt = f
( y , t ) , the p -th
TTM reads as
[Equation]
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:
[Equation]
33) where f ( x t , t ) =
− 12 β ( t ) (cid:104) x
t − ϵ θ ( x t ,t ) σ t (cid:105) , and
[Equation]
In this work, we generate samples from DDMs using the scheme in Eq.
We distill the derivative d γ t ϵ θ :=
d ϵ θ dγ t into a small neural network k ψ .
For training, d γ t ϵ θ is computed via automatic differentiation, however, during inference, we can efficiently query the trained network k ψ .
B.1 Theoretical Bounds for the Truncated Taylor Method
Consider the p -TTM for a general ODE d y dt =
f
( y , t ) :
[Equation] [Equation]
We represent, the exact solution y ( t n +1 ) using the ( p + 2) -th
Taylor expansion
[Equation]
The local truncation error (LTE) introduced by the p -th TTM is given by the difference between the two equations above
[Equation]
For small h n , the LTE is proportional to h p +1 n .
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.
[61] sample from DDMs using (an approximation to) a higher-order Itô-Taylor method [59].
In their scheme, they approximate higher-order score functions with the “ideal derivative trick”, essentially assuming simple single-point ( 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 x t , there is a single unique x 0 from the input data to be predicted with the score model).
In that case, further assuming the score model ϵ θ ( x t , t ) is learnt perfectly (i.e., it perfectly predicts the noise that was used to generate x t from x 0 ), one has
[Equation]
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.
[Equation]
and
[Equation]
Rearranging Eq.
(38), we have
[Equation]
Inserting this expression, Eq.
(40) becomes
[Equation]
We will now proceed to show that the “ideal derivative trick”,
i
[Equation]
e. using the approximations in Eqs.
(39) and (42), results in d γ t ϵ θ = 0 .
As in Sec.
[Equation]
the total derivative d γ t ϵ θ is composed as
[Equation]
Inserting the “ideal derivative trick”, the above becomes
[Equation]
where we have inserted Eq.
(26) for d x t dγ t and used the usual parameterization s θ ( x t , t ) :
= − ϵ θ ( x t ,t ) σ t .
Using d log α 2 t dt = 1 α 2 t dα 2 t dt and dα 2 t dt dtdγ t = dα 2 t dγ t , we can see that the right-hand side of Eq.
(44) is 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 [58], 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 f ( x t , t ) = −
[Equation]
β
[Equation]
t ,
[Equation]
(cid:105) (right-hand side of Probability Flow ODE), then
[Equation] [Equation]
46) where β ′ ( t ) :
= dβ ( t ) dt .
Using the “ideal derivative trick”, we have d ϵ θ dt =
d γ t ϵ θ d t γ t ≈ 0 , and therefore the above becomes
[Equation]
The derivative dσ t dt can be computed as follows
[Equation]
Putting everything back together, we have
[Equation]
which is clearly not 0 for all 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.
[94] 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 γ t ϵ θ = 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 σ 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:
[Equation]
52) where h n = ( γ t n +1 − γ t n ) .
In the remainder of this section, we derive a computable formula for d 2 ϵ θ dγ 2 t , only containing partial derivatives.
Using the chain rule, we have
[Equation]
where, using Eq.
[Equation]
γ ϵ
[Equation]
[Equation] [Equation]
and
[Equation] [Equation]
The remaining terms in Eq.
(55) can be computed as
[Equation]
and
[Equation]
Step size ∆ t 0 .0 .0 .0 .
.k
ξ t ( ∆ t ) k
GENIE (2nd TTM), t = 0 . GENIE (2nd TTM), t = 0 .
GENIE (2nd TTM), t = 0 .3rd
TTM, t=0.
3rd TTM, t=0. 3rd TTM, t=0.
where, inserting Eq.
(28) for d x t dγ t as well as using the usual parameterization s θ ( x t , t ) :
= − ϵ θ ( x t ,t ) σ t , ∂
[Equation]
d x t dγ
[Equation] [Equation]
[Equation]
We now have a formula for d 2 ϵ θ dγ 2 t containing only partial derivatives, and therefore we can compute d 2 ϵ θ dγ 2 t using automatic differentiation.
Note that we could follow the same procedure to compute even higher derivatives of ϵ θ .
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 CIFAR10.
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 θ ( 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 θ ( 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 [109] could potentially alleviate this issue.
Also the higher-order score matching techniques derived by Meng et al.
[82] 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 [78]) with provable local and global truncation errors (see also App. 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 k ψ ( x t , t ) .
Consequently, if k ψ ( x t , t ) modeled the ground truth second-order derivatives exactly, i.e. k ψ ( x t , t ) =
d γ t ϵ θ ( x t , t ) for all 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 k ψ ( x t , t ) towards d γ t ϵ θ ( x t , t ) is well motivated and theoretically grounded.
In particular, during training we are calculating exact ODE gradients using automatic differentiation on the firstorder 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.
perfectly and obtain k ψ ( x t , t ) =
d γ t ϵ θ ( 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 s θ ( x t , t ) and the spatial derivative of the log-perturbation kernel, i.e., ∇ x t log p t ( x t | 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 [63], 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.
C Model and Implementation Details C.1 Score Models
We train variance-preserving DDMs [57] for which σ 2 t = 1 − α 2 t .
We follow Song et al.
[57] and set β ( t ) = 0 .
1 + 19 .
9 t ; note that α t = e − 12 (cid:82) t 0 β ( t ′ )
dt ′ .
All score models are parameterized as either s θ ( x t , t ) :
= − ϵ θ ( x t ,t ) σ t ( ϵ -prediction) or s θ ( x t , t ) := − α t v θ ( x t ,t )+ σ t x t σ t ( v -prediction), where ϵ θ ( x t , t ) and v θ ( x t , t ) are U-Nets [79].
The ϵ -prediction model is trained using the following score matching objective [1]
[Equation]
The v -prediction model is trained using the following score matching objective [69]
[Equation]
which is referred to as “SNR+1” weighting [69].
The neural network v θ is now effectively tasked with predicting v :=
α t ϵ − σ t x 0 .
CIFAR-10:
On this dataset, we do not train our own score model, but rather use a checkpointprovided by Song et al.
The model is based on the DDPM++ architecture introduced in Song et al.
[57] and predicts ϵ θ .
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.
[1] and predicts ϵ θ .
ImageNet: This model is based on the architecture introduced in Dhariwal and Nichol [4].
We make a small change to the architecture and replace its sinusoidal time embedding by a Gaussian Fourier projection time embedding [57].
The model is class-conditional and we follow Dhariwal and Nichol [4] and simply add the class embedding to the (Gaussian Fourier projection) time embedding.
The model predicts ϵ θ .
Cats (Base): This model is based on the architecture introduced in Dhariwal and Nichol [4].
We make a small change to the architecture and replace its sinusoidal time embedding by a Gaussian
Fourier projection time embedding [57].
The model predicts v θ .
Cats (Upsampler): This model is based on the architecture introduced in Dhariwal and Nichol [4].
We make a small change to the architecture and replace its sinusoidal time embedding by a
Gaussian Fourier projection time embedding [57].
The upsampler is conditioned on noisy upscaled lower-resolution images, which are concatenated to the regular channels that form the synthesized outputs of the diffusion model.
Therefore, we expand the number of input channels from three to six.
We use augmentation conditioning [105] to noise the lower-resolution image.
In particular, we upscale α t ′ x low + σ t ′ z , where x low is the clean lower-resolution image.
During training t ′ is sampled from U [ t cutoff , 1] .
During inference, t ′ is a hyper-parameter which we set to 0 .
1 for all experiments.
We use two-independent Gaussian Fourier projection embeddings for t and t ′ 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 γ t ϵ θ using a small prediction head k ψ on top of the first-order score model ϵ θ .
In particular, we provide the last feature layer from the ϵ θ network together with its time embedding as well as x t and the output of ϵ ( x t , t ) to the prediction head (see Fig.
4 for a visualization).
We found modeling d γ t ϵ θ to be effective even for our Cats models that learn to predict v = α t ϵ − σ t x 0 rather than ϵ .
Directly learning d γ t v θ and adapting the mixed network parameterization (see App.
C.2.3) could potentially improve results further.
We leave this exploration to future work.
We provide additional details on our architecture next.
[Table]
[Figure]
[Table]
C.2.1 Model Architecture
The architecture of our prediction heads is based on (modified)
BigGAN
residual blocks [57, 101].
To minimize computational overhead, we only use a single residual block.
In particular, we concatenate the last feature layer with x t as well as ϵ θ ( 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 k
( i ) ψ , i ∈ { 1 , 2 , 3 } , in Eq.
All model hyperparameters can be found in Tab.
We also include the additional computational overhead induced by the prediction heads in Tab. 5; see App.
C.2.5 for details on how we measured the overhead.
C.2.2 Training Details
We train for 50k iterations using Adam [110].
We experimented with two base learning rates: 10 −and 5 · 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 [69], 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 ∂
ϵ θ ( x t ,t ) ∂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 ∂
ϵ θ ( x t ,t ) ∂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 ( x t ) =
N ( x t ; 0 , σ 2 t I ) .
This assumption leads to ϵ θ ( x t , t )
≈ x t σ t which we can plug into the three
[Table]
terms of Eq.
[Equation]
and −
[Equation]
(64) and finally ∂ ϵ
[Equation] [Equation]
[Equation] [Equation]
where we have used σ t = γ t √ γ 2 t +1 .
This derivation therefore implies the following mixed network parameterization
[Equation]
where k ( i ) ψ ( x t , t ) , i ∈ { 1 , 2 , 3 } , are different output channels of the neural network (i.e. the additional head on top of the ϵ θ network).
To provide additional intuition, we basically replaced the − x t σ t terms in Eqs.
(63) to (65) by neural networks.
However, we know that for approximately Normal data x t σ t
≈ ϵ θ ( x t , t ) , where ϵ θ ( x t , t ) predicts “noise” values ϵ 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 k ( i ) ψ
( 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.
[49] and Dockhorn et al.
C.2.4 Pseudocode
In this section, we provide pseudocode for training our prediction heads k ψ and using them for sampling with GENIE.
In Alg. 1, the analytical dtdγ t is an implicit hyperparameter of the DDM as it depends on α t .
For our choice of α t = e − 12 (cid:82) t 0 0 .
1+19 .
9 t ′ dt ′
(see App. C.1), we have
[Equation]
(67) where γ t = (cid:113) 1 − α 2 t α 2 t .
In Alg. 2, we are free to use any time discretization t 0 = 1 . 0 > t 1 > · · · > t N = t cutoff .
When referring to “linear striding” in this work, we mean the time discretization t n = 1 .− (1 .
− t cutoff ) nN .
When referring to “quadratic striding” in this work, we mean the time discretization t n = (cid:0) 1 . 0 − (1 . 0 − √ t cutoff ) nN (cid:1) 2 .
Input: Score model s θ := − ϵ θ
x t ,t ) σ t , number of training iterations N .
Output: Trained prediction head k ψ .
for n = 1 to N do
Sample x 0 ∼ p 0 ( x 0 ) , t ∼ U
[ t cutoff , 1] , ϵ ∼ N ( 0 , I ) Set x t
=
α t x 0 + σ t
ϵ Compute
ϵ θ
( x t , t )
Compute the exact spatial Jacobian-vector product JVP s = ∂
ϵ θ ( x t ,t ) ∂ x t (cid:18) 1 √ γ 2 t +1 ϵ
θ
( x t , t ) − γ t 1+ γ 2 t x t (cid:19) via automatic differentiation Compute the exact temporal Jacobian-vector product JVP t =
∂
ϵ θ ( x t ,t ) ∂t dtdγ t via automatic differentiation ( dtdγ t can be computed analytically)
Compute k ψ ( x t , t ) using the mixed parameterization in Eq.
(66) Update weights ψ to minimize γ 2 t ∥ k ψ ( x t , t ) − d γ t ϵ θ ( x t , t ) ∥ 22 , where d γ t ϵ θ ( x t , t )
=
JVP s − JVP t end for
Input: Score model s θ := − ϵ θ
x t ,t ) σ t , prediction head k ψ , number of sampler steps N , time discretization { t n } Nn =0 .
Output: Generated GENIE output sample y .
Sample x t 0 ∼ N ( 0 , I )
Set ¯ x t 0 = (cid:113) 1 + γ 2 t 0 x t▷ Note that ¯ x t n = (cid:113) 1 + γ 2 t n x t n for all t n for n = 0 to N − 1 do if AFS and n =
then ¯ x t n
[Equation]
= ¯ x t n + ( γ t n +1 − γ t n ) x
t n else ¯ x
[Equation]
=
[Equation]
x t
[Equation] [Equation]
γ
[Equation]
− γ t
[Equation]
k ψ ( x t
[Equation]
, t
[Equation]
) end if x t n +1 =
¯ x tn +1 (cid:113) 1+ γ 2 tn + end for if Denoising then
[Equation]
else y = x
t N end if
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.
D Learning Higher-Order Gradients without Automatic Differentiation and Distillation
In this work
[Equation]
we learn the derivative d γ t ϵ θ , 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
[Equation]
where s 1 ( x t , t ) :
= ∇ x t log p t ( x t ) and S 2 ( x t , t ) :
= ∇ ⊤ x t ∇ x t log p t ( x t ) .
The above formula is derived in Meng et al.
[Theorem 1, 82].
Adding x t x
⊤ t to Eq.
(68) and subsequently dividing by σ 2 t , we have
[Equation]
where we could pull the 1 σ 2 t x t x ⊤ t term into the expectation because it is conditioned on t and x t .
Using x t = α t x 0 + σ t ϵ
[Equation]
we can rewrite the above as
[Equation]
For an arbitrary v := v ( x t , t ) , we then have
[Equation]
Therefore, we can develop a score matching-like learning objective for the (general) spatial JVP o θ ( x t , t )
≈ S 2 ( x , t ) v as
[Equation]
for some weighting function g no − ad ( t ) .
Setting v ( x t , t ) = − σ t
[Equation]
would recover the spatial JVP needed for the computation of d γ t ϵ .
In the initial phase of this project, we briefly experimented with learning the spatial JVP using this approach; however, we found that our distillation approach worked significantly better.
E Toy Experiments
For all toy experiments in Sec. 3, we consider the following ground truth distribution:
[Equation]
where
[Equation]
We set σ = 10 − 2 , s 1 = 0 .
9 , s 2 = 0 .
2 , and µ 1 = (cid:18)
[Equation]
The ground truth distribution is visualized in Fig.
2a.
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.
[Figure]
F Image Experiments F.1 Evaluation Metrics, Baselines, and Datasets
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 [4] 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 ( x 1 ) ≈ N ( x 1 ; 0 , 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
N ( 0 , I ) ≈ p 1 ( 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.
In particular, we have
[Equation] [Equation]
and d ϵ θ ( x 1 , 1) dγ≈ 0 as shown below:
[Equation]
Given this, the AFS step becomes identical to the Euler update that uses the Normal score function for 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 ϵ θ ( x t , t ) with
[Equation]
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 ϵ θ ( x t , t, c ) trained before (see App.
C.1 for details), and train an additional unconditional ImageNet score model ϵ θ ( 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 [70].
The modified GENIE scheme for classifier-free guidance is then given as
[Equation]
where
[Equation]
F.4 Encoding
To encode a data point x 0 into latent space, we first “diffuse” the data point to t = 10 − 3 , i.e., x t =
α t x 0 + σ t ϵ , ϵ ∼ N ( 0 , I ) .
We subsequently simulate the generative ODE (backwards) from t = 10 − 3 to t = 1 , obtaining the latent point x 1 .
To decode a latent point x 1 , we simulate the generative ODE (forwards) from t = 1 . 0 to t = 10 − 3 .
We then denoise the data point, i.e., x 0 = x t − σ t ϵ θ ( x t ,t ) α 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 x (0)0 and x (1)0 .
We first encode both data points, following the encoding setup from App.
F.4, and obtain x (0)1 and x (1)1 , respectively.
We then perform spherical interpolation of the latent codes:
[Equation]
Subsequently, we decode the latent code x ( b ) 1 following the decoding setup from App.
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 [63] and F-PNDM [63].
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 to 13.
As expected, AFS can considerably improve results for almost all methods, in particular for NFEs ≤ 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 [111].
In particular, we follow DDGAN [67] and use the improved recall score [112]; 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 ≥ 15 .
However, F-PNDM cannot be run for NFE ≤ 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.
Striding Schedule Grid Search.
As discussed in Sec.
5 the fixed quadratic striding schedule (for choosing the times t for evaluating the model during synthesis under fixed NFE budgets) used in
GENIE may be sub-optimal, in particular for small NFEs.
To explore this, we did a small grid search over three different striding schedules.
As described in App.
C.2.4, the quadratic striding schedule can be written as t n = (cid:0) 1 . 0 − (1 . 0 − √ t cutoff ) nN (cid:1) 2 , and easily be generalized to
[Equation]
In particular, besides the quadratic schedule ρ = 2 , we also tested the two additional values ρ = 1 .and ρ = 2 . 5 .
We tested these schedules on GENIE as well as DDIM [58]; note that the other two comptetive baselines, S-PNDM [63] and F-PNDM [63], 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 ( ρ = 2 . 0 ) also to ρ = 1 . 5 and ρ = 2 . 5 .
As can be seen in the table, ρ = 1 . 5 improves for both DDIM and GENIE for NFE = 5 (over the quadratic schedule
Since this encode-decode loop is itself not perfect, the references are not perfectly reproduced at b = 0 and b = 1 .
[Table]
Euler–Maruyama
[Table] [Table]
ρ = 2 ), whereas larger ρ are preferred for larger NFE.
The improvement of GENIE from 13.9 to 11.
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 x T ∼ N ( 0 , 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 [58], S-PNDM [63], F-PNDM [63], 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 trainingfree 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 [104]) 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
[Figure]
NFEs
M ea n L- d i s t a n ce
GENIE
DDIM [58]
S-PNDM [63] F-PNDM [63]
GENIE to the LTE of DDIM.
In particular, we compute LTEs at three starting times t ∈ { 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 ∆ 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 ChurchOutdoor (Fig. 16), ImageNet (Fig. 17), and Cats (upsampler conditioned on test set images) (Fig.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.
[Figure]
t − ˆ x t ( ∆ t ) ∥
GENIE, t = 0 .
GENIE, t = 0 . GENIE, t = 0 .
DDIM, t = 0 . DDIM, t = 0 .
DDIM, t = 0 .
[Table] [Table]
S
[Table]
[63] ✗ ✗ 60.
[Table]
62 ✗ ✓ 17.
[Table]
62 ✓ ✗ 35.
[Table] [Table]
51 ✓ ✓ 56.
[Table] [Table]
[Table]
63] ✗ ✗
N/A
N/A 12.
4.89 ✗ ✓
N/A
N/A 19.
[Table]
43 ✓ ✗
N/A
N/A 10.
[Table] [Table]
73 ✓ ✓
N/A
N/A 15.
Euler–Maruyama ✗ ✗    85.
[Table]
✓    89.9 ✓ ✗    80.3 ✓ ✓    83.
[Table] [Table]
✗ ✓ -
[Table] [Table]
✗ ✓ 12.
[Table] [Table]
✗ ✓ - 14.
[Table]
60] ✗ ✗    106 VESDE-PC [57] ✗ ✓
[Table] [Table]
[Table] [Table]
[Table]
G Miscellaneous G.1 Connection to Bao et al. [96]
The concurrent Bao et al.
[96] learn covariance matrices for diffusion model sampling using prediction heads somewhat similar to the ones in GENIE.
Specifically, both Bao et al.
[96] 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.
[96] learn the DDM’s sampling covariance matrices, while we learn higher-order ODE gradients.
More generally, Bao et al.
[96] 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.
[96]’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 x t and t as inputs), as in GENIE,
Bao et al.
[96] would require to model an entire distribution for each 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.
[96] 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.
[96] 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.
[96], 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 [69]:
In every distillation stage of [69], 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 stagewise distillation and maybe also lead to an accuracy and performance improvement.
In particular, we could replace the DDIM predictions in Algorithm 2 of [69] 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.