Learning Fast Samplers for Diffusion Models by Differentiating Through Sample Quality
Authors: Daniel Watson
William Chan
Jonathan Ho
Mohammad Norouzi
Venue: ICLR 2022 Poster (2022)

Abstract: Diffusion models have emerged as an expressive family of generative models rivaling GANs in sample quality and autoregressive models in likelihood scores. Standard diffusion models typically require hundreds of forward passes through the model to generate a single high-fidelity sample. We introduce Differentiable Diffusion Sampler Search (DDSS): a method that optimizes fast samplers for any pre-trained diffusion model by differentiating through sample quality scores. We also present Generalized Gaussian Diffusion Models (GGDM), a family of flexible non-Markovian samplers for diffusion models. We show that optimizing the degrees of freedom of GGDM samplers by maximizing sample quality scores via gradient descent leads to improved sample quality. Our optimization procedure backpropagates through the sampling process using the reparametrization trick and gradient rematerialization. DDSS achieves strong results on unconditional image generation across various datasets (e.g., FID scores on LSUN church 128x128 of 11.6 with only 10 inference steps, and 4.82 with 20 steps, compared to 51.1 and 14.9 with strongest DDPM/DDIM baselines). Our method is compatible with any pre-trained diffusion model without fine-tuning or re-training required.

Published as a conference paper at ICLR 2022
LEARNING FAST SAMPLERS FOR DIFFUSION MODELS BY DIFFERENTIATING THROUGH SAMPLE QUALITY
Daniel Watson, William Chan, Jonathan Ho & Mohammad Norouzi Google Research, Brain Team
(watsondaniel, williamchan hot.aaatttttttoooooooooooogooooa
ABSTRACT
Diffusion models have emerged valing GANs in sample quality and autoregressive models in likelihood scores. Standard diffusion models typically require hundreds of forward passes through the model to generate a single high-fidelity sample. We introduce Differentiable Diffusion Sampler Search (DDSS): a method that optimizes fast samplers for any pre-trained diffusion model by differentiating through sample quality scores. We present Generalized Gaussian Diffusion Models (GGDM), family of flexible Mon-Markovian samplers for diffusion models. We show that optimizing the de- grees of freedom of GGDM samplers by maximizing sample quality scores via gradient descent leads improved sample quality. Our optimization procedure backpropagates through the sampling process using the reparametrization trick and gradient rematerialization. DDSS achieves strong results unconditional im- age generation across various datasets (e.g., FID scores on LSUN church 128x 128 of 1 11.6 with only 10 inference steps, and 4.82 with 20 steps, compared to 51.1 and 14.9 with strongest DDPM/DDIM baselines). Our method is compatible with any pre-trained diffusion model without fine-tuning or re-training required.
expressive family of generative models ri-
INTRODUCTION
Denoising Diffusion Probabilistic Models (DDPM) (Sckl-Dickstein et al. 2015 Song & Ermon
2019 Ho al. 2020) have emerged powerful family of generative models, capable of synthe-
sizing high-quality images, audio, and 3D shapes (Ho et al. 2020 2021 Chen et al. 2222 b Cai
al. 2020 Luo & Hu 2021). Recent work (Dhariwal & Nichol 2021 Ho al. 202 shows that
DDPMs can outperform Generative Adversarial Networks (GAN) (Goodfellow et al, 2014 Brock 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., (017).
S steps
10 steps
20 steps
1000 steps
(smoo) ssaa
Real
DIIN
Figure 1: Non-cherrypicked samples from DDSS (ours) strongest DDIM 0) baseline uncondi- tional DDPMs trained on LSUN churches 128x128. All samples are generated with the same random seed. Original DDPM samples (1000 steps) and training images are shown on the left.
Work done part of the Google AT Residency.

---

Published as a conference paper at ICLR 2022
However, GANs are typically much more efficient than DDPMs generation time, often requiring a single forward pass through the generator network, whereas DDPMs require hundreds of forward passes through U-Net model. Instead of learning generator directly, DDPMs learn to convert noisy data to less noisy data starting from pure noise, which leads to a wide variety of feasible strategies for sampling (Song al., (02Tb). In particular, at inference time, DDPMS allow trolling the number of forward passes (a.k.a. inference steps) through the denoising network (Song al. 2020 Nichol & Dhariwal 2021).
It has been shown both empirically and mathematically that, for any sufficiently good DDPM, more inference steps leads to better likeliool and sample quality (Nichol Dhariwal 2021| Kingma et al.) (022)). 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 a Watson et al. [02T]) has explored reducing the number of steps required for high quality sampling with pretrained diffusion models. See Section 7 for
more complete review of prior work on few-step sampling.
This paper treats the design of fast samplers for diffusion models as a differentiable optimization problem, and proposes Differentiable Diffusion Sampler Search (DDSS). Our key observation is that one 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 samplers, which we call Generalized Gaussian Diffusion Model (GGDM), includes Denoising Dif- fusion Implicit Models (DDIM) (Song et al. 2020) as a special case and is motivated by the success of DDIM on fast sampling of diflusion models.
global objective function. Our class of parameteric
An important challenge for fast 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) finds that samplers that are optimal with respect ELBO often lead to worse sample quality and Fréchet 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 find high-fidelity diffusion samplers, motivated by prior work showing that their optimization leads to solutions that correlate better with human perception of quality. We empirically find that using DDSS with the al.] 2018) as the perceptual loss indeed leads to fast samplers with significantly better image quality than prior work (see Figure D.. Moreover, our method is robust to different choices of kernels for KID.
Kernel Inception Distance (KID) (Binkowski
Our main contributions are as follows:
1. We propose Differentiable Diffusion Sampler Search (DDSS), which uses the reparametrizatioon trick and gradient rematerialization to optimize over a parametric family of fast samplers for diffusion models.
2. We identify a parametric family of Generalized Gaussian Diffusion Model (GGDM) that admits high-fidelity fast samplers for diffusion models.
3. We show that using DDSS to optimize samplers by minimizing the Kernel Inception Distance leads fast diffusion model samplers with state-of-the-ard sample quality scores.
BACKGROUND ON DENOISING DIFFUSION IMPLICIT MODELS
We start with brief review DDPM (Ho et al. 2020 and DDIM (Song al. 2020 DDPMs
pre-specify a Markovian forward diffusion process, which gradually adds noise to data in T steps.
Following the notation of Ho et al. 2020).
( .... xT) = q(xu) ) II (
(1)
q( = Nax / 1eV1/BB-.... 3,I,,
(2)
where q(xn) is the data distribution, and 3, is the variance of Gaussian noise added at step t. To be able to gradually convert noise to data, DDPMs leam invert (1) with model po( [2..) which
2

---

Published as a conference paper at ICLR 2022
is trained by maximizing a (possibly reweighted) evidence lower bound (ELBO):
g DKL (xulg((r\/=)))l(<=>))) Σ
DRL DkL]q((t-1121 2o 20)p: log po (2)))
(3)
DDPMs specifically choose the model to parametrized as
po po(21-1/2z) = xt-1
Zit ty
=
[ - vI ((co(()1))))
vān
(4)
B:
āe
= N
x1-1
VI -
co (2(c,
1 ū;-1 3eId
VI
-
where di I 3,) for each With this parametrization, maximizing the ELBO equivalent to minimizing a weighted sum of denoising score matching objectives (Vincent 2011
The seminal work of [Song al. (2020) presents Denoising Diffusion Implicit Models (DDIM): 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. Specifically, (Sng et al. (2020) note that is it possible to construct alternative ELBOs with only a subsequence of the original timesteps marginals the construction above (i.e., Us (s(x;, aD)) q( (ano) for every faster sampler compatible with the pre-trained model) by simply using the new contiguous timesteps in the equations above. They also show also possible construct infinite family of Markovian processes {go E [0, 1j=-] where each Go also shares the same marginals as the original forward process with:
C [I....,II T} that shares the same
S, so qs defines
T-1
la (~ (Zo,.... vx)) = ((2o))(((rr)) 2o)
do (x, [1+1,, 20)) II t=1
(5)
and where the posteriors are defined as
9 (11-111/2,, xD) = N x Voi-1Va1112 + /1 ā:-1 o/ Dt Vorzo
of
(6)
Song et al.
(2020) empirically find that the extreme case of using all-zero variances (a.ka.
DDIM(1 0)) consistently helps with sample quality in the few-step regime. Combined with
good selection of timesteps to evaluate the modeled score function (a.k.a. strides), DDIM(7 0)
establishes the current state-of-the-ar for few-step diffusion model sampling with the smallest in- ference step budgets. Our key contribution that allows improving sample quality significantly by optimizing sampler families is constructing a family that generalizes DDIM. See Section 4 for more complete treatment of our novel GGDM family.
DIFFERENTIABLE DIFFUSION SAMPLER SEARCH (DDSS)
We now describe DDSS, our approach to search for fast high-fidelity samplers with limited budget of K < T steps. Our key observation is that one can backpropagate through the sampling process of a diffusion model via the reparamterization trick (Kingma & Welling 2013). Equipped with this, we stochastic gradient descent to learn fast samplers by optimizing any given differentiable loss function over a minibatch model samples.
We begin with pre-trained DDPM and family of K-step samplers that we wish to optimize for the given DDPM. We parametrize this family's degrees 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: parametrize the posterior variances with o, vI- d,-1 sig)oid(::), where U..... UK are trainable variables (the /I d constant ensures real-valued mean coef- ficients; see the square root in Equation

---

Published as a conference paper at ICLR 2022
pe (x
X0
Xt
X++
XT
Figure 2: Illustration of GGDM. To improve sample quality, our novel family of samplers combines information from all previous (noisier) images at every denoising step.
we parametrize the marginal variances DDPM as cumsum (softmax([v)) 1))) instead of fixing them to 1 dg.. This ensures they monotonically increasing with respect to t (appending a one to ensure K degrees freedom).
VARS:
GGDM: a new family of non-Markovian samplers for diffusion models with more degrees of freedom illustrated in Figure 2 and defined in Section 4 We parametrize Pta. and og of a GGDM for all as sigmoid functions of trainable variables.
we parametrize all the (eux and identically to GGDM, but also learn the marginal coefficients with cumsum softmax parametrization (identical to VARS) instead of computing them via Theorem I, well the coefficients that predict from ax - b,e with + softplus and softplus parametrizations.
GGDM +PRED:
[family]+TIME: for any sampler family, additionally parametrize the timesteps used to query the score model with a cumsum softmax parametrization (identical to VARS).
As we will show in the experiments, despite the fact that our pre-trained DDPMs are trained with discrete timesteps, learning the timesteps is still helpful. principle, this should only be possible for DDPMs trained with continuous time sampling (Chen et al. 2021a [Song al., 202 Kingma jet al. 2021). but in practice we find that DDPMs trained with continuously embedded discrete timesteps still well-behaved when applied at timesteps not present during training. We think this is due to the regularity of the sinusoidal positional encodings Vaswani et al. (2017) used in these model architectures and training with a sufficiently large number of timesteps T.
3.1 DIFFERENTIABLE SAMPLE QUALITY SCORES
We can differentiate through a stochastic sampler using the reparameterization trick, but the question of which objective to optimize still remains. Prior work has shown that optimizing -ikeeiiiddood 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 perceptual loss which simply compares mean statistics be- tween model samples and real samples in the neural network feature space. These types of objectives have been shown in prior work better correlate with human perception of sample quality (Johnson et al. 2016 Heusel et al. 2017). which we also confirm in our experiments.
We rely on the representations of the penultimate layer a pre-trained Inception V3 classifier (Szegedy et al. 2016) and optimize the Kernel Inception Distance (KID) (Binkowski et al. 2018). Let o( (z) denote the inception features of an image æ and py represent a diftusion sampler with trainable parameters sb. For linear kernel, which works best in experiments, the objective is:
C(
E. ((Tp)) ((ar))
E Bp aara
E d(x,)) d(srg)
More generally, KID can be expressed as:
(kI =
E 1 (xx))
E l (xx
(8)
pprpps
24
where f*(x) Exg kig (x, xb,) WEri ( ( ki (x, x(e) is the witness function for any differen- tiable, positive definite kernel k. and (xx y) = k(o(z), (.))) Note that f attains the supremum of the MMD. To enable stochastic gradient descent, we minibatch of n model samples (
unbiased estimator of KID using (n) xp ~ Pv and n real samples w. ...x." ~ q:
(n)
1
2
x (1)
his (xx( (((i))) C,
(9)
n(n 1)
n
4

---

Published as a conference paper at ICLR 2022
where e is constant in v. Since the sampling chain of any Gaussian diffusion process admits using the reparametrization trick, our loss function is fully differentiable with respect to the trainable vari- ables th. We empirically find that using the perceptual features is crucial; i... by trying ((x) æ to compare images directly on pixel space rather than neural network feature space (as above), we ob- serve that method makes samples consistently worsen in apparent quality as training progresses.
3.2 GRADIENT REMATERIALIZATION
In order for backpropagation to be feasible under reasonable memory constraints, one final problem must be addressed: since we taking gradients with respect 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 s output from the forward pass required by the backward pass computation, we recom- pute it on the fly. To trade (K) memory cost for O ( computation time, we simply rematerialize calls to the pre-trained DDPM (i.e., the estimated score function), but keep in memory the pro- gressively 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.
GENERALIZED GAUSSIAN DIFFUSION MODELS
We now present Generalized Gaussian Diffusion Model (GGDM), our novel family of Gaussian diffusion processes that includes DDIM as a special case mentioned section distribution with no independence assumptions
We define a joint
T-1
qu.a (oo ... ØT) - ((<0))((T vo) II
quro (we piie 2o)
(10)
t=1
where the factors defined as
qu. (xt (x,, N
Σ NES,
Hew o.Id
(11)
(letting S, = (0..... {1. for notation compactness), with and You free parameters Vt {I...., T), St. 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 the appendix (A.2). this construction admits Gaussian marginals, and we can differentiably compute the marginal coefficients from arbitrary choices of μ and o:
Theorem 1. Given some t E [',,,,,,,, let dra pru Vu E S, and vt = of For each i {I, ..., T t), recursively define
((+1) (6) deu =0 oi,tiileti,u + are Yu E Si+s
(1)
and up ((+1) = ui (t) + (1(1) (afhitotts)
Then, it follows that
qu (2: = =
Σ BES
(++1) azu Dus vl
(1+1)
la
(12)
In other words, instead of letting the 3, (or equivalently, the āt) define the forward process as done by a usual DDPM, the GGDM family lets the Pta and a, define the process. In particular, an immediate corollary of Theorem 1 that the marginal coefficients are given by
qu.o ( = N
220
(T-++) (T-1+1) Ia
(13)
The reverse process is thus defined p(xT)
IT pl p( x1) with p(xT) ~ N Ia)) and
xc V Vor (T-1+1) ( ()
Pa (2x xx))) = qu.o
x ( Φ0 = [T-++1)
6o
(14)
a4a
5

---

Published as a conference paper at ICLR 2022
Table 1: FID and IS scores for DDSS against baseline methods for DDPM trained on CIFARIO
with the simple objective proposed by (Ho et al. 2020). FID scores (lower is better) are the numbers the left of each entry, and IS scores (higher Is better) are at the right.
Sampler K
5
10
15
20
25
DDPM (lincar stride)
5,339
43.39/7.088
31.40/ 7.609
25.94 7.879
22.60/ 8.043
DDPM (quadratic stride)
76,25 5,435
42.03 / 6,965
27.78 / 77711
20.225 8.128
16.17 8.350
DDIM (linear stride)
44,41 6.750
19.11 7.966
14.06 / 8.190
11.82/ 8.420
10.52 8.512
DDIM (quadratic stride)
7.090
13.62 8.190
9.318 8.495
7.500 8.641
6,560 8.759
GGDM +PRED+ TIME
13.77/8.522
8.227 8,903
6.115 9,050
4.72279.261
4,250 0225075 99188
Table 2: FID/IS scores for DDSS against baseline methods for a DDPM trained on ImageNet 64x64 with the Lhybrid objective proposed by Nichol & Dhariwal 2021
Sampler K
10
15
20
25
DDPM (linear stride)
5.870/5.878
58.78 10.67
39.30/ / 13.22
31,36 14.72
26,36/ / 15.71
DDPM (quadratic stride)
394.8 / 1,351
129.5 / 5,997
80,10 / 9.595
61.34 / 11.60
49.60/ / 13.01
DDIM (lincar stride)
135.4 5.898
40.70 / 12.225
28.54 / 13.99
24.225 / 14.75
22.13 / 15.16
DDIM (quadratic stride)
409.1 1.380
148.6 5.533
67.65 / /9844
45.60 11.99
36.11 13.225
GGDM +PRED+TIME
55.14/ 12.90
37.32/ / 14.76
24.69 / 17.225
20.69 /17.92
18.40 / 18.12
4.1 IGNORING 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 find, however, that this condition can often be too restrictive and better samplers exist where the marginals qu.o (1g (zu)
N [2 d20 (T-++1)
wo U/ (T-++1) Ia) 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 ("VARS" in Section B): both sampler families have the same number of parame- ters (the reverse process variances), but the latter does not adjust the mean coefficients 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(1 = 0) baseline); see Section 5.2
5 EXPERIMENTS
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. Specifically, we experiment with the DDPM trained by Ho et al. (020)) with Lsimple on CIFAR10, as well as DDPM following the exact configuration Nichol & Dharwal (202I) trained on ImageNet 64x64 (Deng et al.) 2009) with their Liybri objective (with the only difference being that we trained the latter ourselves for 3M rather than 1. 1.5M steps). Both of these models utilize adaptations of the UNet architecture (Ronnebergeri 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), compar- ing 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 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 3.1 with the follow- ing hyperparameters:
1. 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. (020)) Nichol & Dhariwal 2021)

---

Published as a conference paper at ICLR 2022
Table 3: FID IS scores for the KID kernel ablation CIFAR1O. When not learning the timesteps, we fix them to quadratic stride, as Table shows this performs best on CIFARIO. Sampler K
5
10
15
20
25
DDSS (linear kernel)
GGDM +PRED+TIME
13.77 / 8.520
8.227 8,900
6.115 /9.050
4.722 9.261
4.250/ 9.186
GGDM +PRED
14.26 / 8.406
8.617 / 8.842
5.939 / /9.035
4.893 /9.153
4.574 / 99.155
GGDM +TIME
12.85 / 8,383
7.858/ / 8.895
6.265 / 9.075
5,367 / /9.136
4.887 / /9.229
GGDM)
14.45 / 8,281
8.154 / 8,892
7.045 / 8.939
5,477 /9.183
4.815 9.189
DDSS (cubic kernel)
GGDM +PRED+TIME
14.41 / 8.527
8.2257 / 9.0007
5,895 / 9.036
4,932 / 9.092
4,278/ 9.286
GGDM +PRED
14.39 / 8,401
8,977 / 8,870
6.517 / 8.970
4.915 / 19.132
4,471 / 9.244
GGDM +TIME
12.35 / 8.406
7.879 / 8,852
6.682/ 8.999
5.639 / 9.058
4.631 / 9.189
GGDM
14.57 8.297
8.2252 8.836
6.727 / 8.904
5.569 99.177
4.668 9.192
2. We 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 hyperpa- rameters and kept 81 0.9, 92 0.999, and e x 10-5
3. We tried batch sizes of 128 and 512 and opted for the latter, finding 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.
4. 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 I for CIFARIO and Table 2 for ImageNet 64x64, compar- ing DDSS applied GGDM +PRED+TIME against DDPM and DDIM baselines with linear and quadratic strides. All models use linear kerne,, i.e., h y) = o(z) oly). which we found perform slightly better than the cubic kernel used by Binkowski et al. (2018) (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 (202T)), 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 learned stride; see Figures [ and 3. We include additional ImageNet 64x64 samples (A.1 and results for
larger resolution datasets (A.4 in the appendix.
5.1 ABLATIONS FOR KID KERNEL AND GGDM VARIANTS
As our approach is compatible with any choice of KID kernel, experiment with different choices of kernels. Namely, we try the simplest possible linear kernel, Kg (z. y) = <(x) ((y), as well as
the cubic kernel (a, y) (2)()) TTy) + 1)² used by Binkowski al. (2018). We compare the performance of these two kernels, as well as different variatinns of GGDM (I.c., with and without TIME and PRED as defined in Section 3). Results are included for CIFAR 10 across all budgets in Table A.3).
We also include a smaller version of this ablation for ImageNet 64x64 in the appendix
The results in this ablation show that the contributions of the linear kernel, timestep learning, and the empirical choice of learning the coefficients that predict my 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. See also the results on LSUN in the appendix A.44 which also do not include these additional trainable variables.
5.2 SEARCH SPACE ABLATION
Now, in order to further demonstrate the key importance of optimizing our GGDM family to find high-fidelity samplers, we also apply DDSS to the less general DDIM and VARS families. We show that, while we still attain better scores than regular DDPM, searching these less flexible families of samplers does not yield improvements as significant as with out novel GGDM family. In particular, optimizing the DDIM sigma coefficients does not outperform the corresponding DDIM() 0) baseline on 10, which is not a surprising result as [Song et al. (2020) show empirically that most choices of the ae degrees of freedom lead worse FII) scores than setting them all to 0. These

---

Published as a conference paper at ICLR 2022
Table 4: FID / IS scores for the DDSS search space ablation CIFARIO. All runs fix the timesteps to a quadratic stride and use a linear kernel except for the last row (we only include the GGDM results for ease of comparison).
Sampler K
5
10
15
20
25
DDIMO1 = 0)
32.66 / 7.090
13.62 / 8.10
9.318 / 88449
7,500 / 8.641
6,560 / 8.759
DDSS
VARS
33.08 7.096
15.33 8.559
9,693 8.845
7.297 8.924
6.172 99055
DDIM
32.61 / 7.084
16.29 / 7.966
11.31 8.372
9.120 8.563
7.853 8,644
GGDM
14.45 / 8,281
8.154 / 8.892
7.045 8,939
5.477 /9.183
4.815 / 19,189
GGDM +PRED+TIME
13.77 / 8.520
8.227 / 8.903
6.115 99005
4.722 /9.261
4,250 991886
results also show that optimizing the VARS can outperform DDSS applied to the DDIM family, and even the strongest (DIM(1) = 0) baselines for certain budgets,justing justifying our choice of not enforcing the marginals to match (as discussed in Section 4.1,
DISCUSSION
When applied to a sufficiently flexible family (such as the GGDM family proposed in this work), DDSS consistently finds samplers that achieve better image generation quality than the strongest baselines in the literature for very few steps. This is qualitatively apparent in -herryyyree sam- ples (e.g.. DDIM(1 = 0) tends to generate blurrier images and with less background details as the budget decreases), and multiple quantitative sample quality metries (FID and IS) also reflect 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 CIFARIO the metrics indicate significant relative improvement over sample quality metrics, the relative improvement on ImageNet 64x64 is less pronounced. We hy- pothesize that this is an inherent difficulty 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 they might be crucial to bode-breaking.
Beyond the empirical gains of applying our procedure, our findings shed further light into properties of pre-trained score-based generative models. First, we show that without fine-tuning DDPM's parameters, these models are already capable of producing high-quality samples with very few in- ference 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 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 findings thus suggest that enforeing this marginal-sharing constraint unnecessary and can be too restrictive in practice.
OTHER RELATED WORK
Besides DDIM (ong al. 2020), there have been more recent attempts at reducing the number of inference steps for DDPMs. Jolicoeur-Martiean et al. (2021) proposed a dynamic step size SDE solver that can reduce the number of calls the modeled score function to ~ 150 on CIFARIO (Krizhevsky et al., 2009) with minimal cost in FID scores, but quickly falls behind DDIM(1 = 0) with as many as 50 steps. |Watson et al. (20T)) proposed dynamic programming algorithm that chooses log-likelihood optimal strides, but find that log-likelihood reduction has mismatch with FID scores, particularly with in the very few step regime, also falling behind DDIM(1 0) 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 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 efficient sampling. (Song al. (2021b) derive reverse SDE that, when discretized, uses
8

---

Published as a conference paper at ICLR 2022
5 steps
10 steps
steps
Reference
DDIM
Ours
Figure 3: Non-cherrypicked samples for DDPM trained on CIFARIO, comparing the strongest DDIM(1 = 0) baseline and our approach. All samples were generated with the same random seeds. For reference, we include DDPM samples using all 1,000 steps (top right) and real images (bottom right).
different coefficients than the ancestral samplers considered in this work. The same authors also derive "corrector" steps, which introduce additional calls to the pre-trained DDPM as form of gradient ascent (Langevin dynamics) that help with quality but introduce computation cost, as well as an alternative sampling procedure using probability flow ODE that shares the same marginals as the DDPM's original forward process. [Huang al.) (202T) generalize this family of samplers to "plug-in reverse SDE" that interpolates between probability flow ODE and the reverse SDE, similarly to how the DDIM yj 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 single timestep is considered.
CONCLUSION AND FUTURE WORK
We propose Differentiable Diffusion Sampler Search (DDSS), a method for finding few-step sam- plers for Denoising Diffusion Probabilistic Models. We show how to optimize perceptual loss 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 significantly improve sample quality for unconditional image generation over prior methods on efficient DDPM sampling. The success of our method hinges searching a novel, wider family of Generalized Gaussian Diffusion Model (GGDM) than those identified in prior work (Song et al. (2200). DDSS does not fine-tune the pre-trained DDPM, only needs to be applied once, has few hyperparameters, and does not require re-training the DDPM.
Our findings suggest future directions further reduce the number inference steps while retaining high fidelity in generated samples. For instance, is plausible to use different representations for the perceptual loss instead of those of a classifier, 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 benefits 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 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 similar results is im- portant future work. This would make DDSS itself a more efficient procedure, as gradient-based optimization of our proposed loss requires extensive memory or computation requirements back- propagate through the whole sampling chain.
lead to overfitting.

---

Published as a conference paper at ICLR 2022
REFERENCES
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein GAN. In arXiv, 2017.
Mikolaj Binkowski, Danica J Sutherland, Michael Arbel, and Arthur Gretton. Demystifying mmd gans. arXiv preprint arXiv: .11 401, 2018.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal
Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Miine -Milne,, and Qiao
Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
//github. cttl cow//ooo/////a.
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale GAN training for high fidelity natural image synthesis. arXiv preprint arXhv:/809.11096 2018.
Ruojin Cai, Guandao Yang. Hadar Averbuch-Elon Zekun Hao, Serge Belongie, Noah Snavely, and Bharath Hariharan. Learning Gradient Fields for Shape Generation. In ECCV, 2020.
Nanxin Chen, Yu Zhang, Heiga Zen, Ron J. Weiss, Mohammad Norouzi, and William Chan. Wave- Grad: Estimating Gradients for Waveform Generation. In ICLR, 2021a.
Nanxin Chen, Yu Zhang, Heiga Zen, Ron Weiss, Mohammad Norouzi, Najim Dehak, and William Chan. WaveGrad 2: Iterative Refinement for Text-to-Speech Synthesis In INTERSPEECH, 2021b.
Ting Chen, Simon Kornblith, Kevin Swersky, Mohammad Norouzi, and Geoffrey Hinton. Big Self- Supervised Models are Strong Semi-Supervised Learners. In NeurlPS, 2020.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. Icee, 2009.
Prafulla Dhariwal and Alex Nichol. Diffusion models beat gans on image synthesis. arXiv preprint arXiv:2105.05233, 2021.
Ian J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. arXiv preprint arXiv:1406.2660 2014.
Ishaan Guirajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville. Im- proved training of wasserstein gans. arXiv preprint 1XXv:1704.00028. 2017.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge to a local nash equilibrium. arXiv preprint arXiv:1706.08500, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models. NeurIPS, 2020.
Jonathan Ho, Chitwan Saharia, William Chan, David J Fleet, Mohammad Norouzi, and Tim Salimans. Cascaded diffusion models for high fidelity image generation. arXiv preprint 11Xi::2106.15282, 2021.
Chin-We Chin-Wei Huang. Jae Hyun Lim, and Aaron Courville. A variational perspective on diffusion-based generative models and score matching, arXiv preprint 100.0:2000002888,, 2021.
Justin Johnson, Alexandre Alahi, and Li Fei-Fei. Perceptual losses for real-time style transfer and super-resolution. In European conference on computer vision, pp. 694-711. Springer, 2016.
Alexia Jolicoeur-Martineau, Ke Li, Rémi Piché-Taillefer Tal Kachman, and loannis Mitliagkas. Gotta go fast when generating data with score-based models. arXiv preprint arXiv:2 14080, 2021.
Diederik Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. In ICLR, 2015.
10

---

Published as a conference paper at ICLR 2022
Diederik P Kingma and Max Welling. Auto-Encoding Variational Bayes. In ICLR, 2013.
Diederik P Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. arXiv preprint arXiv:2107.00634 2021.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Technical Report, 2009.
Kundan Kumar, Rithesh Kumar, Thibault de Boissiere, Lucas Gestin, Wei Zhen Teoh, Jose Sotelo, Alexandre de Brebisson, Yoshua Bengio, and Aaron Courville, MelGAN: Generative Adversarial Networks for Conditional Waveform Synthesis. In NeurIPS, 2019a.
Ravi Kumar, Manish Purohit, Zoya Svitkina, Erik Vee, and Joshua R Wang. Efficient remateri- alization for deep networks. In Proceedings of the 33rd International Conference on Neural Information Processing Systems, pp. 15172-15181, 2019b.
Shitong Luo and Wei Hu. Diffusion probabilistic models for 3d point cloud generation. In Proceed- ings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 2837-2845, 2021.
Eliya Nachmani, Robin San Roman, and Lior Wolf. Non gaussian denoising diffusion models. arXiv preprint arXiv:2106.077 075500000777..... 2021.
Alex Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. arXiv preprint arXiv:2102.0967. 2021.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedi- cal image segmentation. In International Conference on Medical image computing and computer- assisted intervention, pp. 234 144.241.355 Springer. 2015.
Tim Salimans, Ian Goodfellow, Wojeiech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. Advances in neural information processing systems, 29: 2234 2242, 2016.
Robin San-Roman, Eliya Nachmani, and Lior Wolf. Noise estimation for generative diffusion mod- els. arXiv preprint arXiv: 1et.221440000000 2021.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learn- ing. pp. 2256-2265 PMLR, 2015.
Jiaming Song. Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02500 2020.
Yang Song and Stefano Ermon. Generative Modeling by Estimating Gradients of the Data Distribu- tion. NewrIPS, 2019.
Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score- based diffusion models. arXiv e-prints, pp. arXiv-2101, 2021a.
Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-Based Generative Modeling through Stochastic Differential Equations. In ICLR, 2021b.
Markus Svensén and Christopher M Bishop. Pattern recognition and machine learning, 2007.
Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jonathon Shiens, and Zbigniew Wojna. Re- thinking the Inception Architecture for Computer Vision. In CVPR, 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz. Kaiser, and Illia Polosukhin. Attention Is All You Need. NIPS, 2017.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Compu- fation, 23 2011.
11

---

Published as a conference paper at ICLR 2022
Daniel Watson, Jonathan Ho. Mohammad Norouzi, and William Chan. Learning to efficiently sam- ple from diffusion probabilistic models. arXiv preprint arXiv:2106.03882, 2021.
Fisher Yu, Ari Seff, Yinda Zhang, Shuran Song, Thomas Funkhouser, and Jianxiong Xiao. Lsun: Construction of large-scale image dataset using deep learning with humans in the loop. arXiv preprint arXiv:/5000 506.03365 2015.
12

---

Published as a conference paper at ICLR 2022
A APPENDIX
A.1 ADDITIONAL IMAGENET 64x64 SAMPLES
We provide additional samples for our results on ImageNet 64x64. The DDPM and (DIM 0) samples (left and middle, respectively) use a linear stride, while our DDSS + GGDM samples (right) use learned stride.
DDPM
DDIM
Ours
5 steps
10 steps
20 RaS
Reference
Figure A.1: Additional samples on ImageNet 64x64. For reference, we include DDPM samples with all 4,000 steps (bottom left) and real samples (bottom middle).
13

---

Published as a conference paper at ICLR 2022
A.2 PROOF FOR THEOREM 1
Theorem 1. Given some E TT, let Oia
Peu Vu E S, and vi
o². For each
i E (I,... ..., T 1), recursively define
(++1) 0r.4iiliiin + (0) vu E Si+s aiw
((+1)) (t) + (alet)00t1))
(6)
Then, it follows that
qu.o
N
It
Σ
aiu
(i+1) Lus
U/
(6+1)
I8
Proof. Let us prove this result with mathematical induction. Note that, for each such t, we have by definition that
quo [14+1 [1541, 20)) =N (xr Euesin Le [242
and
(1) qu 1344, 30) N (x,, Eves, Pr of = CwEs, aru Bus vt La
Therefore, following Svensén & Bishop 2007 (2.115), by prior conjugacy it follows that
qu.o (x) [0 zo) = N
2 art+1 Evese Meti,uuu + EuEsi a
(v. + 11 22 Is)
(1)
= a2
Euesco ( (1,411e411 + ata (1) Bns (2)
= 2
EmeScn
aiu tue o/ Ia)
This proves the base case for our induction argument. Now, let us prove the inductive step. Suppose there exists some integer j {l.....? T t + 1} such that
qu.o (x. マز Fo) = N (xs Euesi,,
(i+1) atu
11)
(1+1)
Du.
By definition, we already know g( | 2o), so we have
ya ( vn) = NV (0245-1 Euesuoon Metitla Cn listzIaa
and (rewriting the inductive hypothesis)
qu.o (x #x0+s+1, xo) NV =N(x; Eveseor atu Bu.
(++1) of ((it1)
Therefore, by prior conjugacy again, it follows that
qu.o (x, (e un)
=N
az a (ittit1 Evesey Mt I, Emeseos (1+1), (1.(42)) ala (ve
(1+1)
+ of Ia)
(++1)
(+1)
=N
az Emesc ((itt)+144+4+1, + aiu (1+1) Du. ur (1+2) Ia)
(f+1)
=W az
EueSi (j+2) xu. D1
x
(1+2)
This concludes the proof of the inductive step. Hence, we have proven the result for any i {1, T f). In particular,
qu.o (2c = N (221
a wo, 11 (T-++1) Ia)
(T-++1)
14

---

Published as a conference paper at ICLR 2022
A.3 ADDITIONAL ABLATION KID KERNEL AND GGDM VARIANTS FOR IMAGENET 64x64
We also ran smaller version of the ablation results presented in Section 5.1 but for ImageNet 64x64 instead of CIFARIO, as these are more computationally intensive to do a full grid search. Results for step budget K 15 are included below. When not learning the timesteps, we fix them to linear stride, as Table 2 shows this performs best on ImageNet 64x64.
Sampler \ K
15
DDSS (linear kernel)
GGDM +PRED+TIME
24.69 / 17.225
GGDM +PRED
27.08 16.44
GGDM +TIME
25.73 17.27
GGDM
28.34/ 16.63
DDSS (cubic kernel)
GGDM +PRED+TIME
26.52 16.29
GGDM +PRED
27.82/ 16.3
GGDM +TIME
26.87/ 16.99
GGDM
28.83 / 16.32
A.4 RESULTS ON LARGER RESOLUTION DATASETS
We include results for LSUN Yu et al. 2015) bedrooms and churches at the 128x resolution. We trained the models for 400K and 200K steps (respectively), and all other hyperparameters are identical: we use the Adam optimizer with learning rate 0.0003 (linearly warmed up for the first 1000 training steps), batch size 2048, gradient clipping at norms over 1.0, dropout of 0. and EMA over the weights with decay rate 0.9999. We train the models using a linear stride of 1000 evenly- spaced timesteps. fixing the
schedule to a cosine function monotonically decreasing from 20 to -20. The ELBO is reweighted with Simple following Ho et al. (2020), but additionally reweight each term by max(1, (1, SNR) which we found to be slightly helptul in resulting FII scores (note this is equivalent minimizing the worst mean squared error between either the 2g or e). The UNet employs five down/up-sampling resolutions with 768 x (1, 2, 4, 6, 8) respective channels, 3 ResNet blocks per resolution, and spatial self-attention the 3 smallest resolutions, i.e., 8, 16, and 32.
After training the models, run DDSS using just the GGDM model family for simplicity (i.e., don't t use the +PRED and +TIME we experiment with in the paper) at 5, 10 and 20 evenly-spaced inference steps. DDSS training occurs for 50K steps, using the Adam optimizer with learning rate of 0.0005 and batch size 512, optimizing the linear kernel for the KID loss. We compare against the usual DDPM and DDIM(1 = 0) baselines the same inference budgets and include the FID scores with all 1000 steps for reference. Results and samples are included below.
Sampler
1000
LSUN Bedroom
DDPM
95.38
44.84
16.88
2.547
DDIM(1 = 0)
168.7
56.33
9.527
DDSS (GGDM)
29.15
11.01
4.817
LSUN Church
DDPM
96.67
51.05
16.53
2.718
DDIM() = 0)
133.1
54.39
14.96
DDSS (GGDM)
30.24
11.59
6.736
15

---

Published as a conference paper at ICLR 2022
DDPM
DDIM
Ours
5 steps
10 steps
20 steps
Reference
Figure A.2: Non-cherrypicked samples for a DDPM trained on LSUN bedroom 128x128, comparing DDPM and DDIM(" 0) to our approach. All samples were generated with the same random seeds and linear stride. For reference, we include DDPM samples using all 1,000 steps (bottom left) and real images (bottom middle).
16

---

Published as a conference paper at ICLR 2022
DDPM
DDIM
Ours
5 steps
10 steps
20 steps
Reference
Figure A.3: Non-cherrypicked samples for a DDPM trained on LSUN church 128x128, comparing DDPM and DDIM() 0) to our approach. All samples were generated with the same random seeds and linear stride. For reference, we include DDPM samples using all 1,000 steps (bottom left) and real images (bottom middle).
17


Reviews:
[{"review_id": "qQVi-V6p_dY", "reviewer": "Reviewer_vLgy", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strengths**\n\nThe idea in the paper is relatively simple and straightforward, and has good empirical performance. The paper discusses how to address the engineering challenges in general DDSS search problems (discontinuity of timesteps, memory constraints). Empirically, optimizing perceptual losses benefits sample quality and is better than evaluation likelihoods (especially when most methods need many steps to have marginally better likelihoods).\n\n**Weaknesses**\n\nThe clarity of the paper is a bit on the weaker side, and I think the paper can be significantly improved by having better clarity.\n- It was quite unclear at first glance what the double subscripts mean in Equation (2) and beyond. It should be made explicit that subscript $ts$ means starting from $t$ and ending at $s$. These double subscripts also make it difficult to identify the main idea behind GGDP. Perhaps having a figure illustrating the GGDP idea helps with clarity.\n- It is unclear whether, with the changed GGDP variational distribution, it is still possible to recover the same \"diffusion / score-matching\" solution with a new variational objective (not defined in the paper). I am inclined to believe this is the case, but I think a proof (like the one in DDIM paper) would be warranted.\n- Equation (9) should be written explicitly as the batch version, since the actual algorithm is not an unbiased estimator of (9). Also is there a reason behind only matching the first-order moment, rather than higher-order moments? Is there a possibility that the DDSS sampler collapses to a single point that has the correct mean? This might need a bit of explanation.\n- Minor, it would be helpful to have more details on how the parametrization in DDSS is performed in the appendix. Part of the confusion may be from the fact that we have a lot of notations, such as $\\mu_{tu}$ and $\\sigma_t$, $a$, $v$ etc. It would be great to explicitly explain what these mean (or correspond to in DDPM / DDIM).\n- I have a hard time understanding the code `t = cumsum(softmax([*v, 1.]))[:-1]`. What are the range of v and t? Maybe an example in the appendix can help?\n- I think the proof sketch can be moved into appendix as a more formal proof. At a high-level everything is Gaussian, so the claims themselves are not too hard to understand.\n\n**Comments**\n\n- The GGDP methods seem to be quite related to multistep methods in numerical integration, [here](https://en.wikipedia.org/wiki/Linear_multistep_method). I wonder what would be the connections to established multistep methods, such as Adam-Bashforth? It would be interesting to see what parameters does DDSS learn to further understand what it does. \n- Is there a reason why datasets with larger images such as LSUN (256x256) are not tested?\n\n\n==== Post rebuttal review ===\nI think the rebuttal addressed most of my concerns. I am a bit disappointed that no discussion about larger image datasets are given, but the other points are adequately addressed. So I will keep my score as is.", "summary_of_the_paper": "The paper discusses how to accelerate a diffusion generative model via Differentiable Diffusion Sampler Search (DDSS). First, a family of non-Markovian samplers is introduced, allowing the models to use multiple previous steps for generation. Next, a sampler search method is proposed to optimize the perceptual loss, with certain engineering issues addressed such as bounded and monotonically increasing steps and gradient parametrization. Empirically, this is able to improve sample quality in a few steps compared to other heuristics such as linear and quadratic schedules on CIFAR10 and ImageNet 64x64, especially when there is a very limited amount of steps that can be taken. The proposed DDSS + GGDP method outperforms DDIM and DDSS + DDIM, showing the advantage of DDSS.", "summary_of_the_review": "I think the problem of accelerating diffusion models is important and believe that the paper makes a novel contribution towards it. That being said, I believe that the paper has serious issues regarding clarity and minor issues on experiments over larger datasets. \n\nI am willing to increase the score if the authors address clarity issues well enough in the rebuttal and especially in the revisions that are allowed by the ICLR format. If the revision does not change by much, then my score will be around 5 and 6 (good idea, relatively poorly polished paper).", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "qDoME97JrSq", "reviewer": "Reviewer_owpU", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**The strengths of the paper**  \nIn general, I found that the paper provides interesting values to the ML community:\n\n1. It tackles an important problem in diffusion-based generative models, i.e., preserving the qualities of the generated samples with fewer discretization steps.\n2. The paper proposes a new method using non-Markovian hierarchical generative models with the pre-trained diffusion-based generative models.\n3. The paper discusses practical techniques to support the scalability of the proposed method.\n\n\n**The weaknesses of the paper**  \nWhile the novelty of the proposed method, I found that the following aspects can be improved.\n\nFirst of all, the motivation of the proposed model needs to be clarified, meaning the authors need to discuss the optimality of the non-Markovian but simple Gaussian hierarchical generative models and the proposed posteriors. I also consider this issue connects to the insufficient justification of minimizing the conceptual loss instead of maximizing the ELBO.\n    \n- For example, suppose the proposed model performs significantly worse when it is trained via maximizing ELBOs. In that case, it implies that the proposed posteriors are not well suited for the proposed generative models. This means the ELBO can not be tightened regardless of the number of discretization steps. Consequently, the pair of the proposed posterior and the generative model are inevitably suboptimal, similar to simple variational autoencoders. This raises the question of what family of densities the proposed model can model or not.  \n    \n- On the contrary, DDIM is designed under the premise that the DDIM models will follow the same marginal $p(x_t)$ of the DDPM models. This implies that DDPMs' expressivity and tightness of the ELBO indirectly justify DDIMs'. Thus, DDIMs' approach makes much more sense, as we can expect that DDIM will model arbitrary densities by maximizing their ELBOs (or weighted ones). Thus, the tuning weights for DDIMs and DDPMs are valid engineering trade-offs depending on the application scenarios.  \n    \n    - A side note: Sohl-Dickstein et al., 2015 and DDPMs are inspired by the fact that the first-order discretization of continuous-time It\\^o diffusion is a Markov chain, each of whose transition probability follows a Normal distribution. Moreover, it has been proven that their reverse-time It\\^ diffusion process exists (under mild conditions), and thus their first-order discretization again is Gaussian Markov chains.  \n    \n- To resolve this, I consider that it is necessary to either provide (1) theoretical proofs that the proposed generative models are arbitrary expressive and the ELBO with the posteriors can be tightened, or (2) empirical analyses that the proposed method can achieve comparable performance when it is trained via maximizing ELBOs. Unless either of them cannot be demonstrated, the significance of the proposed method is doubtful.\n    \nSecond, the introduction of the proposed method can be improved. For instance, the authors briefly mention that the proposed method will use pre-trained reconstruction networks in the abstract and Sec 2. However, I found that clarifying how Equation 7 is used to define generative models will be helpful for readers.\n\nThirds, the descriptions of the proposed method in the context of the previous works seem confusing. For example, the usage of the term \"diffusion process\" sounds confusing. \"Diffusion process\" have commonly referred to as It\\^o diffusion process, which is a solution of specific types of stochastic differential equations; thus, continuous-time models. I found that it is a little problematic to use the \"diffusion process\" for the proposed method.", "summary_of_the_paper": "The paper aims at reducing the number of discretization steps for the generation process of diffusion-based generative models. To do that, the paper proposes a new probabilistic generative model, which exploits pre-trained denoising diffusion probabilistic models. In particular, similarly to Denoising Diffusion Implicit Models (DDIM, Song et al. (2020)), the paper proposes a non-markovian hierarchical generative model, whose densities are defined as;\n- $p(x_0, ..., x_T) = \\prod_{t =0}^{T} p(x_{t} | x_{>t})$, where $p(x_T)$ is a prior distribution such as a standard Normal distribution.\n- $p(x_{t} | x_{>t}) = q_{\\mu, \\sigma}(x_{t} |x_{>t}, x_0 = \\hat{x}(x_{t}, t) )$\n- $q_{\\mu, \\sigma}(x_{t} |x_{>t}, x_0) = N(x_t | \\Sigma_{u = t+1}^{T} \\mu_{t, u} x_u, \\sigma_t^2 I_d)$\n- $\\hat{x}(x_t, t)$ is a reconstruction network similar to the one in the DDIM.\n- Note that the proposed method employs $q_{\\mu, \\sigma}(x_{t} |x_{>t}, x_0))$, while DDIM uses conditionally Markov $q_{\\mu, \\sigma}(x_{t} |x_{s}, x_0)$ for any $t < s$.\n\nNoting that the discrete-time diffusion-based generative models are fully differentiable, the paper proposes to train the models while the pre-trained reconstruction network is fixed; here, the authors name this training as Differentiable Diffusion Sampler Search (DDSS). To train the models, the paper first emphasizes that maximizing the ELBO of the proposed method deteriorates the performance wrt FID. Then the authors instead propose to minimize perception scores between models' sample and training data. Here, the perception score is a squared distance between empirical means of two sets of samples on the feature space of a pre-trained classifier.\n\nMoreover, observing that training the entire discrete-time diffusion-based generative models requires huge memory costs, the paper proposes to use gradient rematerialization (Kumar et al., 2019b) methods.\n\nIn the experiments, the paper tests the effectiveness of the proposed method in image generation quality benchmarks, such as CIFAR-10. Specifically, the analysis of the experiments focuses on two perspectives. First, the authors compare the generation qualities of the proposed method and the baseline methods, including denoising diffusion probabilistic models (DDPM), in the few-step regime. Second, the paper compares the generation qualities of the DDIM and the proposed model when both are trained via DDSS. The paper claims that the proposed generative models trained via DDSM outperform the previous methods in the few-step regime.\n\n===== POST-REBUTTAL COMMENTS ========  \nIn the updated version, Theorem 1 shows that under certain choices of model parameterizations, the proposed models' ELBOs will be the same as the DDPMs'. This clarifies that it is valid to plug-in pre-trained DDPM models (or other similar types), potentially without fine-tuning; therefore, it seems valid to treat the proposed method as improving sampling process. ", "summary_of_the_review": "In general, the paper's contributions are unclear, while the paper tackles a very interesting problem in the context of diffusion-based generative models and provides empirical improvements. Unfortunately, the proposed method is not well-motivated, and the discussion about it is insufficient. Thus, I opt to reject the paper. However, I'm also inclined to improve my evaluation if the aforementioned weak points are well addressed.\n\n===== POST-REBUTTAL COMMENTS ========  \nThe rebuttal had addressed some of my concerns. Consequently, I raised my score from 3 to 6, and I also raised the \"Correctness\" & \"Technical Novelty And Significance\" scores from 2 to 3.\n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "AhpKZ_04zok", "reviewer": "Reviewer_u4ae", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "1. It seems that optimizing the perceptual loss might not give you the samples from the true underlying DDPM model. Although this approach might produce reasonable samples empirically, it lacks theoretical insights: it is unclear what distribution these samples are optimized to match with. It seems that the training objective is hand-crafted, and might lack theoretical guarantees. \n2. Although samples from the proposed method have good FID/IS scores, optimizing the model via the perceptual loss could be directly/indirectly overfitting the evaluation metrics. So the improvements in FID/IS scores might be expected. \n3. Are you still able to get likelihoods (lower bound) on the generated samples from GGDP? \n4. As mentioned in the paper, GGDP can use pre-trained DDPM models. It would be more convincing to also show results on more challenging high-resolution datasets (e.g. LSUN, ImageNet, CelebA) using pre-trained models since the inference time for high-resolution datasets are often significantly longer, and improving the inference time would be a more crucial task for larger models. Currently, only small-scale datasets are considered.\n5. For theorem 1, there is only \"proof sketch\". Is there a formal proof?\n6. In section 4.1, it would be better to mention clearly what \"the collection of all the introduced trainable variables\" is.\n7. It seems that the samples of the proposed method in figure 2 look reasonable when using 5 steps on CIFAR-10. However, in the appendix, the samples using 5 steps on ImageNet 64 do not look that realistic. Does it imply that the performance of the proposed approach is dependent on the dataset? If that is the case, more datasets are needed to show the strength of the method.", "summary_of_the_paper": "Denoising Diffusion Probabilistic Models (DDPMs) can generate high-quality samples, but they are not efficient at inference. This paper proposes Differentiable Diffusion Sampler Search (DDSS) to generate high-quality samples while using fewer inference steps than DDPM. The proposed approach uses reparametrization trick and gradient rematerialization to optimize over a class of parametric samplers that use fewer steps than DDPMs. By optimizing perceptual losses, the proposed approach can generate high-quality samples using a smaller number of inference steps compared to existing approaches for sampling from DDPMs.", "summary_of_the_review": "1. Although the proposed approach can generate high-quality samples, the objective used is hand-crafted and might lack theoretical guarantees. As mentioned in the paper \"GGDP family does not guarantee that the marginals of the new forward process match that of the original DDPM.\" Although the sample quality can be reasonable, it is unclear what the model is doing theoretically (e.g., is it optimizing likelihood, or doing score matching).  \n2. The samples, after optimizing perceptual loss, might not correspond to the samples of the underlying DDPM model. Optimizing via perceptual distance is expected to have better performance under FID/IS scores, but it would also change the statistical property of DDPM. In some sense, the samples are not from DDPM, and it might be useful to discuss in more detail the statistical properties of GGDP.\n3. As mentioned in the paper, GGDP can use pre-trained DDPM models. It would be more convincing to also show results on more challenging high-resolution datasets (e.g. LSUN, ImageNet, CelebA) using pre-trained models, since the inference time of high-resolution models are ofter longer, and improving the inference time of larger models is more crucial.\n4. The samples on ImageNet 64 in the appendix are not very good compared to the samples on CIFAR-10 in figure 2, which also raises questions on how general and robust the proposed approach is.\n5. Given the existing work (e.g., DDIM), conditioned on all the previous noisier images in the sampling chain might not be very novel.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "1LjPA5EU0Je", "reviewer": "Reviewer_zfCn", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I find this paper to be very interesting and its results are quite good. I will address the main points of this paper as outlined above. \n\n1. It appears that the introduction of a more flexible family of samplers is introduced mainly to admit the optimization process for finding the variances of a more efficient sampler (DDSS). As optimizing for $v_1 \\dots v_k, \\sigma$ would be difficult if the corresponding marginals were constrained to be equal. Indeed this is shown in 5.1 where optimizing $\\sigma$ for DDIM does not perform well. This is to say that I am not sure there is any benefit or purpose to GGDP beyond admitting the use of DDSS. This is not necessarily a bad thing, but it's worth pointing out that GGPD doesn't look to stand on its own as a better sampler without DDSS. If this is not the case I am happy to be wrong here. \n\n2. I'm not sure about the practicality of DDSS. Given that gradient rematerialization (only feasible in JAX as far as I know) is leveraged to compute the gradient of the objective, I can only assume that the computational cost of a naive implementation is quite high even for CIFAR10. That being said, this work is about efficient (in number of steps) samplers, not scalable ones. I found the results of GGDP+DDSS to be quite impressive. The samples shown in figure 2 are in my opinion strikingly better than the baseline approach. \n\nI found the evaluation to be reasonable if a bit light on the number of datasets. Regardless, the results are compelling, and given that the DDSS objective is quite simple, I am happy with the approach. \n\nAll things considered, I think this optimization approach is a natural direction when considering how to construct more efficient DDPMs, and the results support it. \n\n-------------------------------------------------------------\n### Post Rebuttal Comments. \n\nI appreciate the updates to the submission made by the authors. In particular, the improved motivation of using a perceptual loss instead of maximizing an ELBO. I will retain my score, though I hope that the authors follow through on their claim that they will provide additional results in the camera ready submission. \n", "summary_of_the_paper": "This work examines the problem of the large number of iterations necessary to get high quality samples from denoising diffusion generative models. \nThe paper mostly consists of two new components:\n\n1. A generalized family of probabilistic samplers called Generalized Gaussian Diffusion Processes (GGDP). This construction is similar to the denoising implicit processes found in [1], in the sense that it discards the Markov assumption, and $p(x_{t=i})$ may rely on all $p(x_{t < i})$. The difference is that the Gaussian marginals of the GGPD may not match those of the original process. \n\n2. An optimization procedure for finding the sampler variances. This procedure is general and may be applied to any family of samplers. The idea is to minimize the distance in feature space between a classifier's representation given the real data, and the final generated iterate in the process, optimizing over the parameters of the sampler family. A computationally amenable scheme is proposed to allow for backpropagating through the diffusion process making use of gradient rematerialization (recomputation of intermediate gradients). \n\nThe method is compared against DDPMs [2] and DDIM [1] in CIFAR10 and Imagenet 64x64. \n\n[1] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020\n\n[2] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pp. 2256–2265. PMLR, 2015\n", "summary_of_the_review": "This work proposed a relatively simple, if computationally difficult optimization procedure for finding the parameters of a sampler able to generate compelling images in fewer iterations. I found the approach reasonable given that a more flexible family of samplers is proposed to be amenable to the optimization process. The results are compelling, achieving better sample quality in fewer timesteps. In the regime of fewer than 10 steps, this work achieves better sample quality in half the steps. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper tackles a very interesting problem in the context of diffusion-based generative models and provides empirical improvements. Pre-rebuttal, reviewers' main concerns lie in the motivation and clarification of the method, while after rebuttal, all reviewers satisfied the response and gave positive scores. The authors should include the additional results to well address the reviewers' concerns in the final version.
Decision: Accept (Poster)