O PTIMIZING F EW -S TEP D IFFUSION S AMPLERS BY
G RADIENT
D
ESCENT
Anonymous authors
Paper under double-blind review
A BSTRACT
Denoising Diffusion Probabilistic Models (DDPMs) have emerged as a ﬂexible family of generative models rivaling GANs and autoregressive models in sample quality and likelihoods.
DDPMs however typically require hundreds of inference steps to generate a high-ﬁdelity image, despite recent progress on speeding up diffusion model sampling.
We introduce Differentiable Diffusion Sampler
Search (DDSS): a method that learns few-step samplers for any pre-trained DDPM by using gradient descent.
We propose Generalized Gaussian Diffusion Models (GGDM), a family of non-Markovian samplers for diffusion models, and we show how to improve the generated samples of pre-trained DDPMs by optimizing the degrees of freedom of the GGDM sampler family with respect to a perceptual loss.
Our optimization procedure backpropagates through the sampling process using the reparametrization trick and gradient rematerialization.
Searching our novel GGDM family with DDSS, we achieve strong results on unconditional image generation on both CIFAR and ImageNet 64x64 ( e.g., FID scores on CIFAR of 7.59 with only 10 inference steps, and 4.67 with 25 steps, compared to 13.
and 6.56 with the strongest respective DDIM( η = 0 ) baselines).
Our method is compatible with any pre-trained DDPM without re-training, only needs to be applied once, and does not ﬁnetune the parameters of the pre-trained DDPM.
I NTRODUCTION
Denoising Diffusion Probabilistic Models (DDPM)
(Sohl-Dickstein et al., 2015; Song & Ermon, 2019; Ho et al., 2020) have emerged as a powerful family of generative models, capable of synthesizing high-quality samples across various modalities, e.g., images, audio, and 3D shapes (Ho et al., 2020; 2021;
Cai et al., 2020; Chen et al., 2021a;b; Luo & Hu, 2021).
Recent work (Dhariwal &
Nichol, 2021; Ho et al., 2021) shows that DDPMs can outperform Generative Adversarial Networks (GAN) (Goodfellow et al., 2014;
Brock et al., 2018) in generation quality, but unlike GANs, DDPMs admit likelihood computation and much more stable training dynamics (Arjovsky et al., 2017; Gulrajani et al., 2017).
That being said, GANs are much more efﬁcient at generation time, typically
[Figure]
requiring a single forward pass through the generator network, whereas DDPMs require hundreds of inference steps.
Instead of learning a generator directly, DDPMs learn to model the score function of the marginals of a diffusion process that interpolates between the data distribution and a simple prior, leading to a wide variety of feasible strategies for sampling (Song et al., 2021b).
In particular, at inference time, DDPMs allow controlling the number of forward passes to the modeled score function (a.k.a. inference steps ) (Song et al., 2020; Nichol & Dhariwal, 2021).
It has been shown both empirically and mathematically that, for any sufﬁciently good DDPM, more inference steps leads to better log-likelihood and sample quality (Nichol & Dhariwal, 2021; Kingma et al., 2021).
In practice, the minimum number of inference steps to achieve competitive sample quality is highly problem-dependent, e.g., depends on the complexity of the dataset, and the strength of the conditioning signal if the task is conditional.
Given the importance of generation speed, recent work (Song et al., 2020; Chen et al., 2021a; Watson et al., 2021) has explored reducing the number of steps required for high quality sampling with pretrained diffusion models.
This paper treats the selection of efﬁcient samplers for diffusion models as a differentiable optimization problem, and proposes Differentiable Diffusion Sampler Search (DDSS).
Our key observation is that one can unroll the sampling chain of a diffusion model and use reparametrization trick (Kingma & Welling, 2013) and gradient rematerialization (Kumar et al., 2019a) to optimize over a class of parametric few-step samplers with respect to a global objective function.
Our class of parameteric samplers, which we call Generalized Gaussian Diffusion Model (GGDM), includes Denoising Diffusion Implicit Models (DDIM) (Song et al., 2020) as a special case and is motivated by the success of DDIM on few-step diffusion sampling.
An important challenge for efﬁcient DDPM sampling is the mismatch between the training objective ( e.g., ELBO or weighted ELBO) and sample quality.
Prior work (Watson et al., 2021; Song et al., 2021a) ﬁnds that samplers that are optimal with respect to ELBO often lead to worse sample quality and Fr ´ echet Inception Distance (FID) scores (Heusel et al., 2017), especially with few inference steps.
We propose the use of a perceptual loss within the DDSS framework to ﬁnd high-ﬁdelity diffusion samplers, motivated by prior work showing that their optimization leads to solutions that correlate better with human perception of quality.
We empirically ﬁnd that using DDSS with the
Kernel Inception Distance (KID) (Bi´nkowski et al., 2018) as the perceptual loss indeed leads to fewstep samplers with signiﬁcantly better image quality than prior work (see Figure 1).
Moreover, our method is robust to different choices of kernels.
Our main contributions are as follows:
We propose Differentiable Diffusion Sampler Search (DDSS), using the reparametrization trick and gradient rematerialization to optimize any parametric space of few-step diffusion samplers.
We identify a parametric family of Generalized Gaussian Diffusion Model (GGDM) that admits high-ﬁdelity few-step samplers for diffusion models.
We show that using DDSS to optimize samplers by minimizing the Kernel Inception Distance leads to few-step diffusion samplers with state-of-the-art sample quality scores.
B ACKGROUND ON D ENOISING D IFFUSION I MPLICIT M ODELS
We start with a brief review on DDPM (Ho et al., 2020) and DDIM (Song et al., 2020).
DDPMs ﬁx a Markovian forward process
[Equation]
(1) where q ( x 0 ) is the data distribution, and, following the notation of Watson et al.
[Equation]
DDPMs are trained to model ∇ x t log q ( x t ) by maximizing a (possibly reweighted) evidence lower bound (ELBO), which is equivalent to minimizing a weighted sum of denoising score matching objectives (DSM) (Vincent, 2011):
[Equation] [Equation]
and the model is parametrized as
[Equation]
for each contiguous s < t .
The seminal work of Song et al.
(2020) presents Denoising Diffusion
Implicit Models (DDIM): a family of evidence lower bounds (ELBOs) with corresponding forward diffusion processes and samplers.
All of these ELBOs share the same marginals as DDPM, but allow arbitrary choices of posterior variances.
Speciﬁcally, Song et al.
(2020) note that is it possible to construct alternative
ELBOs with only a subsequence of the original timesteps S ⊂ { 1 , ..., T } that shares the same marginals as the construction above (i.e., q S ( x t | x 0 ) =
q ( x t | x 0 ) for every t ∈ S , so q S deﬁnes a faster sampler compatible with the pre-trained model) by simply using the new contiguous timesteps in the equations above.
They also show it is also possible to construct an inﬁnite family of nonMarkovian processes { q σ : σ ∈ [0 , 1] T − 1 } where each q σ also shares the same marginals as the original forward process with:
[Equation]
and where the posteriors are deﬁned as
[Equation]
Song et al.
(2020) empirically ﬁnd that the extreme case of using all-zero variances ( a.k.a.
DDIM( η = 0 )) consistently helps with sample quality in the few-step regime.
Combined with a good selection of timesteps to evaluate the modeled score function ( a.k.a. strides ), DDIM( η = 0 ) establishes the current state-of-the-art for few-step diffusion model sampling with the smallest inference step budgets.
G ENERALIZED G AUSSIAN D IFFUSION M ODELS
We now present Generalized Gaussian Diffusion Model (), a family of Gaussian diffusion processes that includes DDIM as a special case.
We deﬁne a joint distribution with no independence assumptions
[Equation]
where the new factors are deﬁned as
[Equation]
(9) (letting S t = { 0 , ..., T } \ { 1 , ..., t } for notation compactness), with σ t and µ tu free parameters ∀ t ∈ { 1 , ..., T } , u ∈ S t .
In other words, when predicting the next, less noisy image, the sampler can take into account all the previous, noisier images in the sampling chain, and similarly to DDIM, we can also control the sampler’s variances.
As we prove in the appendix (A.2), this construction admits
Gaussian marginals, and we can differentiably compute the marginal coefﬁcients from arbitrary choices of µ and σ :
Theorem 1.
Given some t ∈ { 1 , ..., T } , let a (1) tu = µ tu ∀ u ∈ S t and v (1) t = σ 2 t .
For each i ∈ { 1 , ..., T − t } , recursively deﬁne
[Equation]
Then, it follows that
[Equation]
In other words, instead of letting f t 0 and g t 0 (or equivalently, each of the f t,t − 1 and g t,t − 1 , i.e., the β t coefﬁcients in Ho et al.
(2020)) deﬁne the forward process as done by a DDPM, the GGDM family lets the µ tu and σ t deﬁne the process, and consequently, f t 0 = (cid:80) u ∈ S t
+ i a ( i +1) tu x u , and g 2 t
= v ( i +1) t .
The reverse process is thus deﬁned as p ( x T ) (cid:81)
Tt =1 p ( x t − 1 | x t ) with p ( x T ) ∼ N ( 0 , I d )) and p θ ( x t
[Equation]
x >t )
[Equation]
The corresponding ELBO is
[Equation]
where we recover the usual connection to DSM:
[Equation]
3. I GNORING THE MATCHING MARGINALS CONDITION
Unlike DDIM, the GGDM family does not guarantee that the marginals of the new forward process match that of the original DDPM.
We empirically ﬁnd, however, that this condition can often be too restrictive and better samplers exist where the marginals q µ,σ ( x t | x 0 ) =
N (cid:16) x t (cid:12)(cid:12)(cid:12) a ( T − t +1) tx 0 , v ( T − t +1) t
I d (cid:17) of the new forward process differ from the original DDPM’s marginals.
We verify this empirically by applying DDSS to both the family of DDIM sigmas and
DDPM variances: both sampler families have the same number of parameters (the reverse process variances), but the latter does not adjust the mean coefﬁcients like DDIM to ensure matching marginals and still achieves similar or better scores than the former across sample quality metrics (and even outperforms the DDIM( η = 0 ) baseline); see Section 5.2.
D IFFERENTIABLE D IFFUSION S AMPLER S EARCH
We now describe DDSS, our approach to search for high-ﬁdelity samplers with a limited budget of
K < T steps.
We observe we can backpropagate through the diffusion sampling process via the reparamterization trick.
Equipped with this, we can now use gradient-descent to optimize any given differentiable loss function over model samples.
We begin with a pre-trained DDPM and a family of K -step samplers that we wish to search for this DDPM.
We parametrize this family’s degrees of freedom as simple transformations of trainable variables.
We experiment with the following families in this paper, but emphasize that DDSS is applicable to any other family where model samples are differentiable with respect to the trainable variables:
• DDIM : we parametrize the posterior variances with σ t = g s 0 sigmoid( v t ) , where v 1 , ..., v K are trainable variables (the g s 0 constant ensures real-valued mean coefﬁcients; see the square root in Equation 7).
• DDPM-G : we parametrize the marginal variances with g t 0 =
cumsum(softmax([ v ; 1])) t , ensuring they are monotonically increasing with respect to t (concatenating a one to ensure K degrees of freedom).
• GGDM : we parametrize all the µ tu and σ t as sigmoids of trainable variables.
• GGDM
+PRED : we parametrize all the µ tu and σ t identically to GGDM, but also learn the marginal coefﬁcients with a cumsum ◦ softmax parametrization (identical to DDPM-G) instead of computing them via Theorem 1, as well as the coefﬁcients that predict x 0 from a t x t − b t (cid:15) with 1 + softplus and softplus parametrizations, instead of using 1 f t 0 and g t 0 f t 0 .
• [family]+TIME : for any sampler family, we additionally parametrize the timesteps used to call the pre-trained score model with a cumsum ◦ softmax parametrization (identical to DDPM-G).
As we will show in our experiments, despite the fact that out pre-trained DDPMs were trained with discrete timesteps, we ﬁnd that learning the timesteps is still helpful.
In principle, this should only be possible for DDPMs trained with continuous time sampling (Chen et al., 2021a; Song et al., 2021b;
Kingma et al., 2021), but in practice we ﬁnd that DDPMs trained with continuously embedded discrete timesteps are still well-behaved when applied at a timestep not present during training.
We hypothesize that this is due to the regularity of the sinusoidal positional encodings (following Vaswani et al.
(2017)) used by these models when trained with a sufﬁciently large number of timesteps T .
4. P ERCEPTUAL L OSS
We can ﬁnd our sampler in this family by optimizing the desired search family’s trainable variables; however, the question of which objective to use remains.
Prior work has shown that optimizing loglikelihood can actually worsen sample quality and FID scores in the few-step regime (Watson et al., 2021; Song et al., 2021a).
Thus, we instead design a perceptual loss which simply compares mean statistics between the DDPM’s samples against real samples in neural network feature space.
These types of objectives have been shown in prior work to prefer minima that better correlate with human perception of sample quality (Johnson et al., 2016; Heusel et al., 2017), which we also observe in our experiments.
Speciﬁcally, we rely on the representations of the penultimate layer of a pre-trained
InceptionV3 classiﬁer (Szegedy et al., 2016), and we optimize the Kernel Inception Distance (KID) (Bi´nkowski et al., 2018).
Letting the inception features of an image be φ ( x ) and ψ the collection of all the introduced trainable variables, our exact objective is the unbiased estimator of the squared maximum mean discrepancy (MMD) used by Bi ´ nkowski et al.
[Equation]
i = n
[Equation]
j
[Equation]
(15) where p ψ ( x p ) is the marginal distribution of the ﬁnal generated sample of the sampler family parametrized by ψ , q ( x 0 ) is the distribution of real data, k is any differentiable, positive deﬁnite kernel, k φ ( x , y )
=
k ( φ ( x ) , φ ( y )) and f ∗ ( x ) =
E x (cid:48)
p ∼ p ψ ( x p ) k
φ ( x , x (cid:48) p ) − E x (cid:48)
q ∼ q ( x 0 ) k φ ( x , x (cid:48) q ) is the witness function attaining the supremum in the MMD.
Since the sampling chain of any Gaussian diffusion process admits using the reparametrization trick (Kingma & Welling, 2013), our loss function is fully differentiable with respect to the trainable variables ψ .
We empirically ﬁnd that using the perceptual features is crucial; i.e., by trying φ ( x ) = x to compare images directly on pixel space rather than neural network feature space (as above), we observe that our method makes samples consistently worsen in apparent quality as training progresses.
[Table] [Table]
4. G RADIENT R EMATERIALIZATION
In order for backpropagation to be feasible under reasonable memory constraints, one ﬁnal problem must be addressed: since we are taking gradients with respect to model samples, the cost in memory to maintain the state of the forward pass scales linearly with the number of inference steps, which can quickly become unfeasible considering the large size of typical DDPM architectures.
To address this issue, we use gradient rematerialization (Kumar et al., 2019b).
Instead of storing a particular computation’s output from the forward pass required by the backward pass computation, we recompute it on the ﬂy.
To trade O ( K ) memory cost for O ( K ) computation time, we simply rematerialize calls to the pre-trained DDPM (i.e., the estimated score function), but keep in memory all the progressively denoised images from the sampling chain.
In JAX (Bradbury et al., 2018), this is trivial to implement by simply wrapping the score function calls with jax.remat .
E XPERIMENTS
In order to emphasize that our method is compatible with any pre-trained DDPM, we apply our method on pre-trained DDPM checkpoints from prior work.
Speciﬁcally, we experiment with the
DDPM trained by Ho et al.
(2020) with L simple on CIFAR10, as well as a DDPM following the exact conﬁguration of Nichol & Dhariwal (2021) trained on ImageNet 64x64 (Deng et al., 2009) with their
L hybrid objective (with the only difference being that we trained the latter ourselves for 3M rather than 1.5M steps).
Both of these models utilize adaptations of the UNet architecture (Ronneberger et al., 2015) that incorporate self-attention layers (Vaswani et al., 2017).
We evaluate all of our models on both FID and Inception Score (IS) (Salimans et al., 2016), comparing the samplers discovered by DDSS against DDPM and DDIM baselines with linear and quadratic strides.
As previously mentioned, more recent methods for fast sampling are outperformed by DDIM when the budget of inference steps is as small as those we utilize in this work (5, 10, 15, 20, 25).
All reported results on both of these approximate sample quality metrics were computed by comparing 50K model and training data samples, as is standard in the literature.
Also as is standard, IS scores are computed 10 times, each time on 5K samples, and then averaged.
In all of our experiments, we optimize the DDSS objective presented in Section 4.1 with the following hyperparameters:
For every family of models we search over, we initialize the degrees of freedom such that training begins with a sampler matching DDPM with K substeps following Song et al.
; Nichol & Dhariwal (2021).
[Table]
apply gradient updates using the Adam optimizer
(Kingma & Ba, 2015).
We sweeped over the learning rate and used λ = 0 . 0005 .
We did not sweep over other Adam hyperparameters and kept β 1 = 0 .
9 , β 2 = 0 .
999 , and (cid:15) = 1 × 10 − 8 .
We tried batch sizes of 128 and 512 and opted for the latter, ﬁnding that it leads to better sample quality upon inspection.
Since the loss depends on averages over examples as our experiments are on unconditional generation, this choice was expected.
We run all of our experiments for 50K training steps and evaluate the discovered samplers at this exact number of training steps.
We did not sweep over this value.
We include our main results in Table 1 for CIFAR10 and Table 2 for ImageNet 64x64, comparing DDSS applied to GGDM +PRED+TIME against DDPM and DDIM baselines with linear and quadratic strides.
All models use a linear kernel, i.e., k φ ( x , y )
=
φ ( x ) (cid:62) φ ( y ) , which we found to perform slightly better than the cubic kernel used by Bi ´ nkowski et al.
(we ablate this in section 5.1).
We omit the use of the learned variances of the ImageNet 64x64 model (i.e., following
Nichol & Dhariwal (2021)), as we search for the variances ourselves via DDSS.
We include samples for 5, 10 and 25 steps comparing the strongest DDIM baselines to DDSS + GGDM with a learned stride; see Figures 1 and 3.
We include additional ImageNet 64x64 samples in the appendix (A.1).
5. A BLATIONS FOR KID KERNEL AND GGDM VARIANTS
As our approach is compatible with any choice of KID kernel, we experiment with different choices of kernels.
Namely, we try the simplest possible linear kernel, k φ ( x , y ) =
φ ( x ) (cid:62) φ ( y ) , as well as the cubic kernel k φ ( x , y ) =
(cid:0) 1 d φ ( x ) (cid:62) φ ( y ) + 1 (cid:1) 3 used by Bi´nkowski et al.
We compare the performance of these two kernels, as well as different variations of GGDM (i.e., with and without
TIME and PRED as deﬁned in Section 4).
Results are included for CIFAR10 across all budgets in Table 3.
We also include a smaller version of this ablation for ImageNet 64x64 in the appendix (A.3).
The results in this ablation show that the contributions of the linear kernel, timestep learning, and the empirical choice of learning the coefﬁcients that predict x 0 all slightly contribute to better FID and IS scores.
Importantly, however, removing any of these additions still allows us to comfortably outperform the strongest baselines.
5. S EARCH SPACE ABLATION
Now, in order to further demonstrate the key importance of optimizing our GGDM family to ﬁnd high-ﬁdelity samplers, we also apply DDSS to the less general DDIM and DDPM-G families.
We show that, while we still attain better scores than a regular DDPM, searching these less ﬂexible families of samplers does not yield improvements as signiﬁcant as with out novel GGDM family.
In particular, optimizing the DDIM sigma coefﬁcients does not outperform the corresponding
DDIM( η = 0 ) baseline on CIFAR10, which is not a surprising result as Song et al.
(2020) show empirically that most choices of the σ t degrees of freedom lead to worse FID scores than setting them all to 0.
These results also show that optimizing the DDPM-G can outperform DDSS applied
[Table]
to the DDIM family
[Table]
and even the strongest DDIM( η = 0 ) baselines for certain budgets, justifying our choice of not enforcing the marginals to match (as discussed in Section 3.1)
[Table]
D ISCUSSION
When applied to a sufﬁciently ﬂexible family (such as the GGDM family proposed in this work)
[Table]
DDSS consistently ﬁnds samplers that achieve better image generation quality than the strongest baselines in the literature for very few steps.
This is qualitatively apparent in non-cherrypicked samples ( e.g., DDIM( η = 0 ) tends to generate blurrier images and with less background details as the budget decreases), and multiple quantitative sample quality metrics (FID and IS) also reﬂect these results.
Still, we observe limitations to our method.
Finding samplers with inference step budgets as small as K < 10 that have little apparent loss in quality remains challenging with our proposed search family.
And, while on CIFAR10 the metrics indicate signiﬁcant relative improvement over sample quality metrics, the relative improvement on ImageNet 64x64 is less pronounced.
We hypothesize that this is an inherent difﬁculty of ImageNet due to its high diversity of samples, and that in order to retain sample quality and diversity, it might be impossible to escape some minimum number of inference steps with score-based models as they might be crucial to mode-breaking.
Beyond the empirical gains of applying our procedure, our ﬁndings shed further light into properties of pre-trained score-based generative models.
First, we show that without ﬁne-tuning a DDPM’s parameters, these models are already capable of producing high-quality samples with very few inference steps, though the default DDPM sampler in this regime is usually suboptimal when using a few-step sampler.
We further show that better sampling paths exist, and interestingly, these are determined by alternative variational lower bounds to the data distribution that make use of the score-based model but do not necessarily share the same marginals as the original DDPM forward process.
Our ﬁndings thus suggest that enforcing this marginal-sharing constraint is unnecessary and can be too restrictive in practice.
O THER R ELATED W ORK
Besides DDIM (Song et al., 2020), there have been more recent attempts at reducing the number of inference steps for DDPMs.
Jolicoeur-Martineau et al.
(2021) proposed a dynamic step size SDE solver that can reduce the number of calls to the modeled score function to ∼ 150 on CIFAR (Krizhevsky et al., 2009) with minimal cost in FID scores, but quickly falls behind DDIM( η = 0 ) with as many as 50 steps.
Watson et al.
(2021) proposed a dynamic programming algorithm that chooses log-likelihood optimal strides, but ﬁnd that log-likelihood reduction has a mismatch with
FID scores, particularly with in the very few step regime, also falling behind DDIM( η = 0 ) in this front.
Other methods that have been shown to help sample quality in the few-step regime include non-Gaussian variants of diffusion models (Nachmani et al., 2021) and adaptively adjusting the sampling path by introducing a noise level estimating network (San-Roman et al., 2021), but more thorough evaluation of sample quality achieved by these approaches is needed with budgets as small as those considered in this work.
Other approaches to sampling DDPMs have also been recently proposed, though not for the explicit purpose of efﬁcient sampling.
Song et al.
(2021b) derive a reverse SDE that, when discretized, uses different coefﬁcients than the ancestral samplers considered in this work.
The same authors also
[Figure]
derive “corrector” steps, which introduce additional calls to the pre-trained DDPM as a form of gradient ascent (Langevin dynamics) that help with quality but introduce computation cost, as well as an alternative sampling procedure using a probability ﬂow ODE that shares the same marginals as the DDPM’s original forward process.
Huang et al.
(2021) generalize this family of samplers to a “plug-in reverse SDE” that interpolates between a probability ﬂow ODE and the reverse SDE, similarly to how the DDIM η interpolates between an implicit probabilistic model and a stochastic reverse process.
Our proposed search family includes discretizations of most of these cases for
Gaussian processes, notably missing corrector steps, where reusing a single timestep is considered.
C ONCLUSION AND F UTURE W ORK
We propose Differentiable Diffusion Sampler Search (DDSS), a method for ﬁnding few-step samplers for Denoising Diffusion Probabilistic Models.
We show how to optimize a perceptual loss over a space of diffusion processes that makes use of a pre-trained DDPM’s samples by leveraging the reparametrization trick and gradient rematerialization.
Our results qualitatively and quantitatively show that DDSS is able to signiﬁcantly improve sample quality for unconditional image generation over prior methods on efﬁcient DDPM sampling.
The success of our method hinges on searching a novel, wider family of Generalized Gaussian Diffusion Model (GGDM) than those identiﬁed in prior work (Song et al., 2020).
DDSS does not ﬁne-tune the pre-trained DDPM, only needs to be applied once, has few hyperparameters, and does not require re-training the DDPM.
Our ﬁndings suggest future directions to further reduce the number of inference steps while retaining high ﬁdelity in generated samples.
For instance, it is plausible to use different representations for the perceptual loss instead of those of a classiﬁer, e.g., use representations from an unsupervised model such as SimCLR (Chen et al., 2020), to using internal representations learned by the pre-trained
DDPM itself, which would eliminate the burden of additional computation.
Moreover, considering the demonstrated beneﬁts of applying DDSS to our proposed GGDM family of samplers (as opposed to narrower families like DDIM), we motivate future work on identifying more general families of samplers and investigating whether they help uncover even better samplers or lead to overﬁtting.
Finally, identifying other variants of perceptual losses ( e.g., that do not sample from the model), or alternative optimization strategies ( e.g., gradient-free methods) that lead to similar results is important future work.
This would make DDSS itself a more efﬁcient procedure, as gradient-based optimization of our proposed loss requires extensive memory or computation requirements to backpropagate through the whole sampling chain.
R EPRODUCIBILITY S TATEMENT
We plan to open source our work for the camera-ready version of the paper.
Still, we provide sufﬁcient details in the paper to completely reproduce our work, and the datasets and model checkpoints we experiment with are already publicly available.
We include a theorem that shows how to compute the coefﬁcients of the marginals of GGDM (Equation 10), which is a necessary and the most non-trivial step to reproduce our results.
We also include a proof in the paper, immediately after the theorem statement.
Additionally, we indicate the reader how to implement key ideas in the paper, e.g., in Section 4.1 we provide pseudocode showing exactly how we parametrize learnable, monotonically increasing timesteps, and in Section 4.2 we explain how to implement gradient rematerialization with JAX.
Finally, our experiments section is thorough on all the hyperparameters we use, and any exceptions such as early stopping are always documented where appropriate.
E THICS S TATEMENT
We acknowledge that generative models have potentially dangerous applications, and that making them more efﬁcient (as we do in our work) allows more people to utilize them.
Still, methods outside the scope of this work like GANs are already much more efﬁcient and can still achieve high-ﬁdelity samples.
Bias is also an important concern when working on such image generative models, more work needs to be done to investigate and mitigate bias of our samplers.
We do not believe that our research will be directly used for harmful purposes, nor that its release or open sourcing will pose any risk to the community.
A A PPENDIX A. A DDITIONAL I MAGE N ET 64 X 64 SAMPLES
We provide additional samples for our results on ImageNet 64x64.
The DDIM( η = 0 ) samples (left) use a linear stride, while our DDSS + GGDM samples (right) use a learned stride.
For reference, we include DDPM samples with all 4,000 steps (bottom left) real samples bottom middle
DDPM
DDIM
Ours
[Figure]
A. P ROOF FOR T HEOREM
Theorem 1.
Given some t ∈ { 1 , ..., T } , let a (1) tu = µ tu ∀ u ∈ S t and v (1) t = σ 2 t .
For each i ∈ { 1 , ..., T − t } , recursively deﬁne
• a ( i +1) tu = a ( i ) t,t
+ i µ
[Equation]
+ i,
[Equation]
+
[Equation]
• v ( i +1)
t
[Equation]
Then, it follows that
[Equation]
Proof.
Let us prove this result with mathematical induction.
Note that, for each such t , we have by deﬁnition that
[Equation]
and
[Equation]
Therefore, following Svens ´ en & Bishop (2007) (2.115), by prior conjugacy it follows that q
[Equation]
This proves the base case for our induction argument.
Now, let us prove the inductive step.
Suppose there exists some integer j ∈ { 1 , ..., T − t + 1 } such that
[Equation]
By deﬁnition
[Equation]
we already know q ( x t + j +1 | x >t + j +1 , x 0 ) , so we have
[Equation]
and (rewriting the inductive hypothesis)
[Equation]
Therefore, by prior conjugacy again
[Equation]
it follows that
[Equation]
x
[Equation]
| x >t + j +
[Equation]
This concludes the proof of the inductive step.
Hence, we have proven the result for any i ∈ { 1 , ..., T − t } .
In particular,
[Equation]
) = N (cid:16) x t (cid:12)(cid:12)(cid:12) a ( T − t +1) tx 0 , v ( T − t
[Equation]
1) t
I d (cid:17) .
(cid:3)
A. A DDITIONAL ABLATION OF KID KERNEL AND GGDM VARIANTS FOR I MAGE N ET 64 X
We also ran a smaller version of the ablation results presented in Section 5.
[Table]
but for ImageNet 64x64 instead of CIFAR10, as these are more computationally intensive to do a full grid search.
Results for a step budget K = 15 are included below.
When not learning the timesteps, we ﬁx them to a linear stride, as Table 2 shows this performs best on ImageNet 64x64.
[Table]