Multistep Distillation of Diffusion Models via Moment Matching
Authors: Tim Salimans
Thomas Mensink
Jonathan Heek
Emiel Hoogeboom
Venue: NeurIPS 2024 poster (2024)
Keywords: generative modeling
diffusion
distillation
Primary area: generative_models

Abstract: We present a new method for making diffusion models faster to sample. The method distills many-step diffusion models into few-step models by matching conditional expectations of the clean data given noisy data along the sampling trajectory. Our approach extends recently proposed one-step methods to the multi-step case, and provides a new perspective by interpreting these approaches in terms of moment matching. By using up to 8 sampling steps, we obtain distilled models that outperform not only their one-step versions but also their original many-step teacher models, obtaining new state-of-the-art results on the Imagenet dataset. We also show promising results on a large text-to-image model where we achieve fast generation of high resolution images directly in image space, without needing autoencoders or upsamplers.

Multistep Distillation of Diffusion Models via Moment Matching
Tim Salimans
Thomas Mensink Jonathan Heek Emiel Hoogeboom
{sal imans, iaeiiimmnmmmmiieeee mww.smmnnnnnnnneeeee jheek, emi
Google DeepMind, Amsterdam
Abstract
We present new method for making diffusion models faster to sample. The method distills many-step diffusion models into few-step models by matching conditional expectations of the clean data given noisy data along the sampling trajectory. Our approach extends recently proposed one-step methods to the multi- step case, and provides moment matching. By using up to 8 sampling steps, we obtain distilled models that outperform not only their one-step versions but also their original many-step teacher models, obtaining new state-of-the-ar results on the Imagenet dataset. We also show promising results large text-to-image model where achieve fast generation of high resolution images directly in image space, without needing autoencoders or upsamplers.
perspective by interpreting these approaches in terms
A teddy bear playing guitar next to a cactus
A frog wearing sweater
Figure 1: Selected 8-step samples from our distilled text-to-image model.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Introduction
Diffusion models (Ho et al., 2020; Song & Ermon 2019 Sohl-Dickstein et al. 2015) have recently become the state-of-the-arr model class for generating images, video, audio, and other modalities. By casting the generation of high dimensional outputs as an iterative denoising process, these models have made the problem of learning synthesize complex outputs tractable. Although this decomposition simplifies the training objective compared to alternatives like GANs, shifts the computational burden to inference: Sampling from diffusion models usually requires hundreds of neural network evaluations, making these models expensive to use in applications.
To reduce the cost of inference, recent work has moved towards distilling diffusion models into generators that are faster to sample. The methods proposed so far can be subdivided into 2 classes: deterministic methods that aim to directly approximate the output of the iterative denoising process in fewer steps, and distributional methods that try to generate output with the same approximate distribution as learned by the diffusion model. Here we propose a new method for distilling diffusion models of the second type: We cast the problem conditional expectations of the clean data given the noisy data along the sampling trajectory of the diffusion process. The proposed method is closely related to previous approaches applying score matching with an auxiliary model to distilled one-step generators, but the moment matching perspec- tive allows us to generalize these methods to the few-step setting where we obtain large improvements in output quality, even outperforming the many-step base models our distilled generators are learned from. Finally, the moment matching perspective allows us to also propose a second variant of our algorithm that eliminates the need for the auxiliary model in exchange for processing two independent minibatches per parameter update.
distribution matching terms of matching
2 Background
2.1 Diffusion Models
Diffusion models are trained by learning to invert a noise process that gradually destroys data from a clean data sample æ according to 2t = a, + OiE; with 2, ~ (0, I), where (tz, ar are monotonic functions of diffusion time / E [0, 1]. The coefficients a,, C; may be specified in multiple equivalent ways. Here, we use the variance preserving specification (Ho et al. 2020 that has of ag 01 1, and 01 do 0, such that we have that Zo x and Z1 (N(0,I), I). When using a different specification of the noise process we can always convert to the variance preserving specification by rescaling the data. The quantity of importance is thus the signal-to-noise ratio: (tR(t)) = of /o?, rather than the coefficients individually (Kingma et al. 2021 To invert the specified diffusion process, we can sample from the posterior distribution:
((2, 22: x) = =N(z, Me x)), (c+t))
with 22. ( a and Me-s o7
2zx)
(2)
To sample from a learned diffusion model, we replace x by a prediction from a neural network à go (2.)) with weighting function (t) and where q g(z,, x) denotes sampling æ from the data and then producing =; by forward diffusion. The sampling process starts with pure noise 2 ~ (O, I) and iteratively denoises the data according to (2z,12,, 2) for discrete number of timesteps k., following Algorithm [
go(2,, i) that is fit the data by minimizing Ex-p(t) *..ppt) ax) (o(())(
If we attain the optimal solution o E[æ|2,] and let k - oc the sampling process becomes exact, then the learned diffusion model can be shown to be universal distribution approximator (Song al. (2021b). To get close to this ideal, typically needs to be quite large, making di diffusion models very computationally expensive class of models (Luccioni et al. 2023).
2.2 Generalized method of moments
An alternative to the well-known maximum likelihood estimation method is the method of moments, also known moment matching. Traditionally for univariate distributions, one matches moments my = Exxpxx [a*] of a random variable X. The canonical example is Gaussian distribution, which defined by the first two moments (i.c. the mean and variance) and all (centered) higher order moments

---

Algorithm Ancestral sampling algorithm used for both standard denoising diffusion models as well as our distilled models. For standard models typically 256 s k S 1000, for distilled 1 1 16.
Require:
Denoising model go 2t. go t), number of sampling steps k Initialize noisy data 21 ~ NN 1)
for t E (k - 1)/k, ..... 2/k: 1/k} do
Predict clean data using x go (Z.,t)
Set next timestep L 1/k
Sample next noisy data point 23 ~ g(2,) [t, x)
end for
Return approximate sample x
are zero. Fitting a distribution by setting its moments equal to the moments of the data is then consistent parameter estimation method, and can be readily extended to multivariate distributions, e.g. by matching the mean and covariance matrix for multivariate Gaussian.
One can generalize the method of moments to arbitrary high dimensional functions f F R and match the moment vector Method of Moments (GMM, Haasen ( (1922). Matching such moments can be done by minimizing distance between the moments such ||E (a:) Ex J (a=) where pe is the generative model and px the data distribution. The distillation method we propose in the next section can be interpreted special case of this class of estimation methods.
defined by: m Em--px [f(=))] which is called the Generalized
3 Moment Matching Distillation
Many-step sampling from diffusion models starts by initializing noisy data z ~ (O I), which is then iteratively refined by predicting the clean data using & = go (2,,)) t). and sampling a slightly less noisy data point n, ~ q(z,, |,, x) for new timestep < f, until the final sample is obtained as described is described in Algorithm 1 from the data distribution (xx) if the number of sampling steps grows infinitely large. Here aim to achieve a similar result while taking many fewer sampling steps than would normally be required. To achieve this we finetune our denoising model go into new model g,,2,, () which sample from using the same algorithm, but with a strongly reduced number of sampling steps k for say 1 1 k 8.
0, If & E,(x z this procedure is guaranteed to sample
To make our model produce accurate samples for a small number of sampling steps k, the goal is now no longer for x 9n, (Ze, f) to approximate the expectation E. [[]]]] but rather to produce an approximate sample from this distribution. In particular, if x ~ ((x then Algorithm produces exact samples from the data distribution q for any choice of the number of sampling steps. If 9., perfectly approximates q(x ) as intended, we have that
(xx= (x 2. 12, x) [x [2.] E. [x [2.]
Exc-q(x)) (0 (x|2.)
E.y [x/)x
In words: The conditional expectation of clean data should be identical between the data distribution and the sampling distribution g of the distilled model.
Equation ggives us a set of moment conditions that uniquely identifies the target distribution, similar to how the regular diffusion training loss identifies the data distribution (Song et al., (021b). These moment conditions can be used the basis of distillation method to finetune gn, (ze, t) from the denoising model go. In particular, we can fit 9n to g by minimizing the L2-distance between these moments:
i(n) = マ B
Bytz (z. B, x|,]] - E. [x 2.11
(4)
In practice, we evaluate the moments using a sample z, from our generator distribution, but do not incorporate its dependence on the parameters 1) when calculating gradients of the loss. This decision is purely empirical, as we find it results in more stable training compared to using the full gradient. The approximate gradient of [(n) is then given by
1,,2,((812)) T (E, [l2.] E5, [x|2.)) v, (1 [x|z=] T Eg
2.1) (V,))" (E,g(//2.) Ex[x/2=)), (5)
3

---

where we approximate the first expectation using a single Monte-Carlo sample x and where the second term zero does not depend on 9y. Following this approximate gradient then equivalent to minimizing the loss
L(n) Exer-gga (ze) ( (. ,x., [x" sg (E, X [x|]]] - [x(x)/z))]
(6)
where sg denotes stop-gradient. This loss is minimized if E.g [x|2,] F., [x] lz,] [x|2,] as required. Unfortu- nately, the expectation Fg,g [ 2.3] is not analytically available, which makes the direct application of
Equation 6 impossible. We therefore explore two variations on this moment matching procedure: In
Section 3.1 we approximate E, [x|z.] by a second denoising model, and in Section 3.2 we instead
apply moment matching directly in parameter space rather than x-space.
3.1 Alternating optimization of the moment matching objective
Our first approach to calculating the moment matching objective in equation to approximate Ey [x],]] with an auxiliary denoising model Is trained using standard diffusion loss on samples from our generator model 9g- We then update Is and 9r, in alternating steps, resulting in Algorithm
Algorithm 2 Moment matching algorithm with alternating optimization of generator Oy and auxiliary denoising model 9s-
Require: Pretrained denoising model go()) generator 9n to distill, auxiliary denoising model ge. number of sampling steps k. time sampling distribution p(s), loss weight w(s), and dataset D.
for = O:N do
Sample target time p(s), sample time delta 6, U (0 1/k).
Set sampling time minimum(s + 5z, 1).
Sample clean data from D and do forward diffusion to produce Zr.
Sample 2, from the distilled generator using x = (Z:), 2s ~ a [z, x).
if is even then
Minimize L(Ø) = (ll) - (2.)) + (lgo(z.) - go (z.))1 w.r.t. Φ
else
Minimize L(n) = w(s)x sg.) - go()))] w.r.t.
end if end for
Here we have chosen to train our generator In on all continuous times 1 E (0, 1] even though at inference time (Algorithm I we only evaluate k discrete timesteps. Similarly train with randomly sampled time delta ón rather than fixing this to single value. These choices were found to increase the stability and performance of the proposed algorithm. Further, we optimize Io not just to predict the sampled data x but also regularize it to stay close to the teacher model ge: On convergence this would cause I+ predict the average of x and go- which has the effect of multiplying the generator loss L(π) by 1/2 compared to the loss we introduced in Equation 6
The resulting algorithm resembles the alternating optimization of a GAN (Goodfellow et al. 2020). and like a GAN is generally not guaranteed to converge. In practice, we find that Algorithm 2 is stable for the right choice of hyperparameters, especially when taking k: 2 8 sampling steps. The algorithm also closely resembles Variational Score Distillation as previously used for distilling 1-step generators In in Difs-Instruct. We discuss this relationship in Section
3.2 Parameter-space moment matching
Alternatigg optimization of the moment matching objective (Algorithm{2) is difficult to analyze theo- retically, and the requirement to keep track of two different models adds engineering complexity. We therefore also experiment with an instantaneous version of the auxiliary denoising model gø where 6² is determined using a single infinitesimal gradient descent step on L(Ø) (defined in Algorithm 2 evaluated on single minibatch. Starting from teacher parameters 0, and preconditioning the loss gradient with a pre-determined scaling matrix A, we can define:
((A) = 8 AAV L()))a=s so that 0" lim (()).
(7)
X

---

Now we use o(A) in calculating [n) from Algorithm
take the first-order Taylor expansion for (L(()) @= and scale the loss with the
(O(A) (z+) - go (z,))) ss X
(gg(a.))
(a.) (o(A) 0) A
dgo (z.)
09
09
inverse of A to get:
(z.) AV d (o)oooo)
Linstant (1) = lim X (1) - w()) 2-0
(8)
00
where @20(2=) is the Jacobian of go. and where VoL(e) is evaluated on an independent minibatch 89
from X and zg. In modern frameworks for automatic differentiation, like JAX (Bradbury et al. 2018),
the quantity within the curly braces can be most easily expressed using specialized functions for calculating Jacobian-vecte vector products.
The loss can now equivalently be expressed as performing moment matching in teacher-parameter space rather than x-space. Denoting Lo xx, z,) = wss) (o(2))l and letting Limaam (1) Limstana (17) + constant, we have (as derived fully in Appendix
Linstane (7) =
[x -0.8 (x) 20 x) L ( sg(2.))ll
(9)
|E 1p.00p.8pp00 (() V L (x, - VoLo (x, z (1)))
where the gradient of the teacher training loss is zero when sampling from the training distribution, Exxat V T.g xx x²-) = 0, if the teacher attained minimum of its training loss.
The instantaneous version of our moment matching loss can thus be interpreted as trying to match teacher gradients between the training data and generated data. This makes it a special case of the Efficient Method of Moments (Gallant & Tauchen 1996), classic method in statistics where a teacher model pe is first estimated using maximum likelihood, after which its gradient is used to define a moment matching loss for learning a second model On- Under certain conditions, the second model then attains the statistical efficiency of the maximum likelihood teacher model. The difference between our version of this method and that proposed by Gallant & Tauchen (1996) is that in our case the loss of the teacher model is a weighted denoising loss, rather than the log-likelihood of the data.
The moment matching loss (T)) is minimized if the teacher model has zero loss gradient when evaluated on data generated by the distilled student model In- In other words, optimization successful if the teacher model cannot see the difference between real and generated data and would not change its parameters when trained on the generated data. We summarize the practical implementation of moment matching in parameter-space in Algorithm and Figure
Algorithm 3 Parameter-space moment matching algorithm with instant denoising model Jø-.
Require: Pretrained denoising model go (Z()),, generator 9, to distill, gradient scaling matrix A, number of sampling steps k. time sampling distribution p(s), loss weight w(s). and dataset D)
for n = O:N do
Sample target time p(s), sample time delta 6, U 1/k).
Set sampling time = minimum(s + 8g 1).
Sample two independent batches of data from D and do forward diffusion to produce Z,, z For both batches sample 2, z,, from the distilled generator using x = (z,()), Ze ~ q(z, Ze, x).
Evaluate teacher gradient on one batch:
AV, Lo zz,)
On the other batch, minimize Linstant (1) = w/ xx)))* sg
(go 88
v) w.r.t.
end for
9,(22,)
min [E, V,4( 2F
2, ~ g(2./2..)) x)
0
8 Uie - -1/K,0)
gz(z,)
2, 0 0,E
Figure 2: Visualization of Algorithm 3 Moment matching in parameter space starts with applying forward diffusion to data from our dataset, mapping this to clean samples using the distilled generator model, and then minimizes the gradient of the teacher loss on this generated data.
5

---

3.3 Hyperparameter choices
In our choice of hyperparameters we choose to stick as closely as possible to the values recommended in EDM (Karras et al. 2022), some of which were also used in lnff-Instruct (Luo et al., 2024) and DMD (Yin al. 2023). We use the EDM test time noise schedule for p(s). as well as their training loss weighting for wi s), but we shift all log-signal-to-niis ratios with the resolution of the data following Hoogeboom et al. (2023). For our gradient preconditioner A, as used in Section 3.2] we use the preconditioner defined in Adam (Kingma & Ba 2014), which can be loaded from the teacher checkpoint or calculated fresh by running a few training steps before starting distillation. During distillation, A is not updated.
To get stable results for small numbers of sampling steps (k 1, 2) we find that we need use weighting function w(s) with less emphasis on high-signal (low s) data than in the EDM weighting. Using a flat weight w(s) (s) 1 or the adaptive weight from DMD (Yin et al., 2023) works well.
As with previous methods, it's possible to enable classifier-free guidance (Ho & Salimans 2022) when evaluating the teacher model go- We find that guidance is typically not necessary if output quality is measured by FID, though it does increase Inception Score and CLIP score. To enable classifier-free guidance and prediction clipping for the teacher model in Algorithm [3] we need to define how to take gradients through these modifications: Here we find that a simple straight-through approximation works well, using the backward pass of the unmodified teacher model.
4 Related Work
In the case of one-step sampling, our method in Algorithm 2 is a special case of Variational Score Distillation, hiff-Instuuc,, and related methods (Wang et al., 2024; Luo et al.] 20224] Yin et al.]
2023) Nguyynn & Tran] (2023) [ra3) which distill a diffusion model by approximately minimizing the KL divergence between the distilled generator and the teacher model:
[(1) =
(x)) (2o))o(22))
(11)
Pept) (())()()))))((((() (1z, (2x))11 - ||z -asgassleee asg[&x(2, [13]
(12)
= Ept),,, ( [w] (s) ())2/)/(//))) sg.) - X0 (2.)]l + constant
(13)
Here sg again denotes stop gradient, (z.) is defined by sampling x = (21) with 2j ~ N I). and n, ~ (x) is sampled using forward diffusion starting from x. The auxiliary denoising model & is fit by minimizing Eg, u(zz) |xx (2.) )|2, which can be interpreted as score matching because z, is sampled using forward diffusion started from x. In our proposed algorithm, we sample z, from the conditional distribution (2,, [x, 2(): If 2; Z1 ~ NO, I) is assumed to be fully independent of &, i.e. that oof = O, have that q(z X Z1 = g(z, (X) the two methods are indeed the same. However, this correspondence does not extend to the multi-step case: When we sample z. from qzz,, i., x) for of /ot > O, fitting &a through minimizing E,, (2.)) corresponds to score matching. One could imagine fitting X through score matching against the conditional distribution (z, 2c. x) but this did not work well when we tried it (see Appendix D for more detail). Instead, our moment matching perspective offers justification for extending this class of distillation methods to the multistep case without changing the way we fit Xz. Indeed, we find that moment matching distillation also works when using deterministic samplers like DDIM (Song et al. 2021a) which also do not fit with the score matching perspective.
8z(2.)) no longer
In addition to the one-step distillation methods based on score matching, our method is also closely related to adversarial multistep distillation methods, such as Xiao et al. (2021) and Xu et al. (2023a) which use the same conditional g(z., 2., x) we use. These methods train a discriminator model to tell apart data generated from the distilled model (g.,) from data generated from the base model (go). This discriminator is then used to define an adversarial divergence which is minimized w.r.t. 9n²
LTn) Ee=p(t) Dunv Dntv (po())))
(14)
The methods differ in their exact formulation of the adversarial divergence Dav. in the sampling of
time steps, and the use of additional losses. For example Xu et al. (2023a) train unconditional discriminators Da (,(-, L) and decompose the adversarial objective in a marginal (used in the discrimi-
nator) and a conditional distribution approximated with an additional regression model. Xiao et al.
(2021) instead use conditional discriminator of the form Do (,, Zz, t).

---

Experiments
We evaluate our proposed methods in the class-conditional generation setting the ImageNet dataset (Deng et al. 2009), which is the most well-established benchmark for comparing image quality. On this dataset also several ablations to show the effect of classifier-free guidance and other hyperparameter choices on our method. Finally, we present an experiment with a large text-to-image model show our approach can also be scaled to this setting.
5.1 Class-conditioaal generation on ImageNet
We begin by evaluating on class-conditional Im- ageNet generation, at the 64 64 and 128 128
Table 1: Results ImageNet 64x64.
resolutions (Tables 1 and 2) Our results here for relatively small model with 400 million parameters based on Simple Diffussin (Hooge- boom et al., 2023). We distill our models for maximum of 200,000 steps at batch size 2048, calculating FID every 5,000 steps. We report the optimal FID seen during the distillation process, keeping evaluation data and random seeds fixed across evaluations to minimize bias.
Method
param
NFE
FID,
VDM++ Kingma & Gao. 2025
2B
1024
1.43
64
RIN
Jabn ot 40461
281M
1000
1.23
67
our buse model
400M
1024
1.42
DDIM TRACT Berthelot
Song 202
10
18.7
2023
7.43
4.97
2.93
2.41
(LPIPS) Song et al. 2023
6.20
4.70
4.32
For our base models report results with slight classifier-free guidance of w 0.1, which gives optimized amount sampling noise, following (Salimans & Ho (2022), which is slightly higher compared to equation 2 For our distilled models we ob- tained better results without classifier-fre: guid- ance, and we use standard ancestral sampling without tuning the sampling noise. We use iden- tical hyperparameters across all our experiments. We compare against various distillation methods from the literature, including both distillation methods that produce deterministic samplers (progressive distillation, consistency distillation) and stochastic samplers (Diff--Instrul adversar- ial methods).
iCT-deep Song & Dhariwal 2023
3.25
2.77
the optimal FID. We also
PD (Salimans & Ho. 2022 (reimpl. From
400M
10.7
FromHeck Heck
2.4
1.7
63
MultiStep- CD
2024
1.2B
1.9
1.6
73
CTM Kim al.
2024
1.73
64
DMD Yin et
2.62
Diff Instruct
2023
5.57
Moment Matching
400M
Altemating (c.f. Sect
3.1
3.0
89
3.86
1.50
75
1.24
78
Instant (c.f.
98
1.35
81
Ranking the different methods by FID, find that our moment matching distillation method is especially competitive when using 8+ sampling steps, where it sets new state-of-the-as results, beating out even the best undistilled models us- ing more than 1000 sampling steps, as well as its teacher model. For sampling step some of the other methods show better results: Further improvement is likely possible by separately op- timizing hyperparameters for this setting. For 8+ sampling steps we get similar results for our
Table 2: Results on ImageNet 128x 128
Method
# param
NFE
FID.
ISa
VDM++ Kingma & Gao 2023
2B
1024
1.75
171
base moud
400M
1024
1.76
194
PD (alimanns Ho 2023
400M
(reimpi. tronklleck leck 2024
3.8
162
MultiStep-CC CD Hlekk 2024
1.2B
7.0
3.1
alternating optimization version (Section 3.1 and the instant 2-batch version (Section (.2) of method. For fewer sampling steps, the alter- nating version performs better.
2.1
160
Moment Matching
400M
Alternoting (c.f. Scet 3.1
170
3.14
163
1.72
184
1.49
184
We find that our distilled models also perform very well in terms Inception Score (Salimans et al. 2016) even though we did not optimize for this. By using classifier-free guidance the Inception Score can improved further, as we show in Section 5.3
Instant (c.f. Sect. 3.2
3.48
232
1.54
183
7

---

How can a distilled model improve upon its teacher? On Imagenet our distilled diffusion model with sampling steps and no classifier-free guidance outperforms 512-step teacher with optimized guidance level, for both the 64 x 64 and 128 x 128 resolution. This result might be surprising since the many-step teacher model often seen as the gold standard for sampling quality. However, even the teacher model has prediction emorr that makes it possible to improve upon it. In theory, predictions of the clean data at different diffusion times are all linked and should be mutually consistent, but since the diffusion model is implemented with an unconstrained neural network this generally will not be the case in practice. Prediction errors will thus be different across timesteps which opens up the possibility of improving the results by averaging over these predictions in the right way. Different sampling algorithms average over these predictions differently, as shown in Appendix which offers scope for improvement.
Similarly, prediction error will not be constant over the model inputs z,, and biasing generation away from areas of large error could also yield sampling improvements. Although many-step ancestral sampling typically gives good results, and is often better than deterministic samplers like DDIM, it's not necessarily optimal. In future work we hope study the improvement of moment matching over our base sampler in further detail, and test our hypotheses about its causes.
5.2 Ablating conditional sampling
The distilled generator in our proposed method samples from the conditional q(z, |x, zz). whereas existing distillation methods based on score matching typically don't condition Ze. Instead they apply noise independently. mirroring the forward diffusion process used during training the original model. When using 1-step sampling setup, the two approaches are equivalent since any intermediate ze will be independent from the starting point 21 if that point has zero signal-to-noise. In the multistep meaningfully different however, and sampling from the conditional (Zz,) X, 2,) or the marginal gzz,|8i are both valid choices. We ablate our choice conditioning 2, versus applying noise independently, and find that conditioning leads much better sample diversity in the distilled model, as shown in Figure 3
setup the
approaches
Conditional sampling
Unconditional sampling
Figure 3: Multistep distillation results for a single Imagenet class obtained with two different methods of sampling from the generator during distillation; Conditional (zz x z,), and unconditional q(z, X). Our choice of sampling from the conditional yields much better sample diversity.
FID/IS for different guidance levels on Imagenet 128x128 Our distillation method be used with or with- out guidance. For the alternating optimization version of our method we only apply guidance in the teacher model, but not in the generator or auxiliary denoising model. For the instant 2- batch version we apply guidance and clipping to the teacher model and then calculate its gradient with straight through approximation. Exper-
5.3 Effect of classifier-free guidance
2.4
2.2
2.0
1.8
190
200
210
220
230
240
8
15

---

imenting with different levels of guidance, we find that increasing guidance typically increases Inception Score and CLIP Score, while reducing FID, as shown in the adjacent figure.
5.4 Distillation loss is informative for moment matching
A unique advantage of the instant 2-batch ver- sion of moment matching approach is that, unlike most other distillation methods. it has simple loss function (equation (9) that mini- mized without adversarial techniques, bootstrap- ping, or other tricks. This means that the value of the loss is useful for monitoring the progress of the distillation algorithm. We show this for Imagenet 128 x 128 in the adjacent figure: The typical behavior we see that the loss tends to go up slightly for the first few optimization steps, after which exponentially falls to zero with increasing number of parameter updates.
moment matching loss
22-1
22-1
2500 5000
7500 10000
12500 15000 17500 20000
distillation steps
5.5 Text to image
To investigate our proposed method's potential to scale to large text-to-ing to-ttoo-mage models we train a pixel-space model (no encoder/decoder) on licensed dataset of text-image pairs at a res- olution of 512 x 512, using the UVIT model and shifted noise schedule from Simple Diffu- sion (Hoogeboom et al., 2023) and using a T5 XXL text encoder following Imagen (Saharia ct al. 2022). We compare the performance of base model against 8-step distilled model obtained with our moment matching method. In
Table 3: Results on text-to-image, 512 x 512.
COCO FID
CLIP
Method
NFE
guidance
Score f
our model
9.6
0.290
512
0.5
7.9
0.305
512
12.7
0.315
$12
13.4
0,316
StableDiffusion v1.5*
512
low
8.78
(Rombach al. 2022
512
high
13.5
0.322
DMD
low
11.5
2023
high
14.9
0.32
LFO
12.8
0.311
Table 3 we report zero-shot FID (Heusel al.
Xu et al. 2023b
2017) and CLIP Score (Radford et al.) [0221) MS-COCO (Lin 2014): Also in this setting we find that our distilled model with al- ternating optimization exceeds the metries for our base model. The instant 2-batch version of our algorithm performs somewhat less well at 8 sampling steps. Samples from our distilled text-to-image model are shown in Figure and in Figure 7 in the appendix.
Swiltbrush
16.67
0.29
Nguyen & Tran 2023
InstaFlow- 1.7 1.7B
11.8
0.309
al. 2023
PeRFlow
11.3
Yan et al. 2024
Moment Matching Altemating (Sec
3.1
0
7.25
0.297
14.15
0.319
Instant (Sec.
3.2
9.5
0.300
19.0
0.306
Reported results for StubleDiftusion v1.5 are trom Yin et 20023
6 Conclusion
We presented Moment Matching Distillation, a method for making diffusion models faster to sample. The method distills many-step diffusion models into few-step models by matching conditional expectations the clean data given noisy data along the sampling trajectory. The moment matching framework provides new perspective on related recently proposed distillation methods and allows us to extend these methods to the multi-step setting. Using multiple sampling steps, our distilled models consistently outperform their one-step versions, and often even exceed their many-step teachers, setting new state-of-the-art results on the Imagenet dataset. However, automated metrics of image full set of human evaluations on the
quality are highly imperfect, and in future work plan outputs of our distilled models to complement the metrics reported here.
We presented two different versions of our algorithm: One based on alternating updates of distilled generator and an auxiliary denoising model, and another using two minibatches to allow only updating the generator. In future work intend to further explore the space of algorithms spanned by these choices, and gain additional insight into the costs and benefits of both approaches.

---

References
Berthelot, D., Autef, A., Lin, J., Yap. D. A., Zhai, S., Hu, s., Zheng, D., Talbott, W., and Gu, E. TRACT: denoising diffusion models with transitive closure time-distillation. CoRR, abs/2303.04248, 2023.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke,
A., VanderPlas, J. Wanderman- S., and Zhang, Q. JAX: composable transformations of Python+NumPP programs. 2018. URL http: //i thub. com com/ commgoogle/jar
Deng, J., Dong, W., Socher, R., Li. L.-J,, Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. leee, 2009.
Gallant, A. and Tauchen, G. Which moments to match? Econometric theory, 12(4):657-681, 1996.
Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial networks. Communications of the ACM, 63(11):139-144 2020.
Hansen, L. P. Large sample properties of generalized method of moments estimators. Econometrica: Journal of the econometric society, pp. 1029-1054, 1982.
Heek, J., Hoogeboom, E., and Salimans, T. Multistep consistency models. arXiv:2403.06887, 2024.
arXiv preprint
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by two time-scale update rule converge to local nash equilibrium. Advances in neural information processing systems, 30, 2017.
Ho, and Salimans, T. Classifier-fres diffusion guidance. CoRR, abs/2207. 2022. doi: 00448500arXi..2227.1335 /ttp://///////////././////////// 200.
htt.hhttss///oi.orgg 10555////.....3//......
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS, 2020.
Hoogeboom, E., Heek. J., and Salimans, T. simple diffuusio:: End-to-end diffusion for high resolution images. In International Conference on Machine Learning, pp. 13213-13232. PMLR, 2023.
Jabri, A., Fleet, D. J., and Chen, T. Scalable adaptive computation for iterative generation. In International Conference on Machine Learning, pp. 14569-14589 PMLR, 2023.
Karras, T., Aittala, M., Aila, T., and Laine, S. Elucidating the design space of diffusion-based generative models. In Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS, 2022.
Kim, D., Lai, C.-H., Liao, W.-H., Murata, N. Takida, Y., Uesaka, T., He. Y., Mitsufuji, Y., and Ermon, S. Consistency trajectory models: Learning probability flow ODE trajectory diffusion. In The Twelfth International Conference Learning Representations, 2024. URL https:// openneview net. forum? iat.uumiiidyjIIIIfIT..
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. 1rXiv:1412.6988 69800, 2014.
arXiv preprint
Kingma, D. P. and Gao, R. Understanding the diffusion objective as a weighted integral of elbos. CoRR, abs/2303.00848, 2023.
Kingma, D. P., Salimans, T., Poole, B., and Ho, J. Variational diffusion models. abs/2107.00630, 2021.
CoRR,
Lin, T., Maire, M., Belongie, S. J., Bourdev, L. D., Girshick, R. B., Hays, J., Perona, P., Ramanan, D., Doll'a r, P. and Zitnick, C. I.. Microsoft COCO: common objects context. CoRR, abs/1 405.0312, 2014. URL http:: /tappv////xii....a.m
10

---

Liu. X., Zhang, X., Ma, J.. Peng, J., and Liu, Q. Instaflow: One step is enough for high-quality diffusion-based text-to-image generation. CoRR, abs/2309.06380, 2023.
Luccioni, A. S., Jernite, Y., and Strubell, E. Power hungry processing: Watts driving the cost of ai deployment? arXiv preprini urXiv:2311. 2023.
Luo, W., Hu, T., Zhang, S., Sun, J., Li, Z., and Zhang, Z. Diff-instruct: A universal approach for transferring knowledge from pre-trained diffusion models. CoRR, abs/2305.18455. 2023.
Luo, W., Hu, T., Zhang, S., Sun, J., Li, Z., and Zhang, Diff-instruct: A universal approach for transferring knowledge from pre-trained diffusion models. Advances in Neural Information Processing Systems, 36, 2024.
Nguyen, T. H. and Tran, A. Swiftbrush: One-step text-to-imaae diffusion model with variational score distillation. arXiv preprint arXiv:2312.052399 2023.
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Míshkin, P., Clark, J., et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 8748-8763. PMLR, 2021.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image syn- thesis with latent diffusion models. IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pp. 10674-10685. IEEE, 2022.
Saharia, C., Chan, W., Saxena, S., L., Whang, J., Denton, E., Ghasemipour, S. K. S., Ayan, B. K., Mahdavi, S. S., Lopes, R. G., Salimans, T., Ho. J., Fleet, D. J., and Norouzi, M. Photorealistic text-to-image diffusion models with deep language understanding. CoRR, abs/2205. 1487,, 2022.
Salimans, T. and Ho, Progressive distillation for fast sampling of diffusion models. In The Tenth International Conference on Learning Representations, ICLR. OpenReview.net, 2022.
Salimans, T., Goodfellow, I., Zaremba, W., Cheung, V., Radford, A., and Chen, X. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
Sohl-Dickstein, J., Weiss, E. A., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequilibrium thermodynamics. In Bach, R. and Blei, M. (eds.), Proceedings of the 32nd International Conference on Machine Learning, ICML, 2015.
Song. J., Meng, C., and Ermon, Denoising diffusion implicit models. In 9th International Conference on Learning Representations, ICLR, 2021a.
Song, Y. and Dhariwal, P. Improved techniques for training consistency models. CoRR, abs/2310.14182 2023.
Song. Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. In Advances Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS, 2019.
Song. Y.. Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021b.
Song, Y., Dhariwal, P., Chen, M., and Sutskever, Consistency models. In International Conference on Machine Learning, ICML, 2023.
TPUv5e. Google cloud tpu training. Accessed: 2024-05-21.
Wang, Z., Lu, C., Wang, Y., Bao, F. Li, C., Su, H., and Zhu, J. Prolificdreamer: High-fidelity and diverse text-to-3d generation with variational score distillation. Advances in Neural Information Processing Systems, 36, 2024.

---

Xiao, Z., Kreis, K., and Vahdat, A. Tackling the generative learning trilemma with denoising diffusion gans. In International Conference on Learning Representations, 2021.
Xu, Y. Gong, M., Xie, S., Wei, W., Grundmann, M., Hou, T., et al. Semi-implicit denoising diffusion models (siddms). arXiv preprini arXiv:2306. 12511, 2023a.
Xu, Y., Zhao, Y., Xiao, Z.., and Hou, T. Ufogen: You forward once large scale text-to-image generation via diffusion gans. arXiv preprint arXiv:2311.09257, 2023b.
Yan, H., Liu, X., Pan, J., Liew, J. H., Liu, Q. and Feng. J. Perflow: Piecewise rectified flow as universal plug-and-play accelerator. arXiv preprint arXiv:2405.077 07510, 2024.
Yin. T., Gharbi, M., Zhang, R., Shechtman, E., Durand, Freeman, W. T., and Park, T. One-step diffusion with distribution matching distillation, arXiv preprint arXiv:2311.18822 2023.
12

---

A Instant moment matching matching expected teacher gradients
In section 3.2 we propose an instantaneous version of our moment matching loss that does not require alternating optimization of an auxiliary denoising model Ja- This alternative version of our algorithm uses the loss equation which we reproduce here for easy readibility:
1 (n)) w(s)x² sg
(ge(Z.) 00
Linstant (1)
lim x
AV AVL(())0=
(15)
where (ge(x.) the Jacobian of go. and where VoL(e) is evaluated on an independent minibatch 0y from X and Ze.
turns out that this loss can be expressed as performing moment matching in teacher-parameter space rather than x-space. Denoting the standard diffusion loss as (x, z,) = w(s) x go (z.))l² we can rewrite the term (Ø)( -=0 Vo (x. go (zx))) because the first term ((Ø) can be seen as standard diffusion loss at & for a generated &, and the second term of L(ő) is zero when Φ = A. Futher
observe that VxLg(X, ).)
2wo(s) (z.) x)- @gu(x.) which means that V, VoLo z.)
v,,2w(()) Agg (x=), Now letting Trinstam (1) Tinsstmm (1) +w()) (8) lo (2,)² 0gg (x,) VoTa (x, 2n) where the latter term constant w.r.t. 1). we can write instant moment- matthhng (Equation [5 as moment- matching of the teacher gradients where again stop gradients are again placed on Z,:
1 Tinstannt (1)
T.g &X, ((.)
(16)
[2, (t. V
PL (X, sg(2.)) -xxx)))Exxxx,0 VoLo (x, z (17
where assume that the gradient of the teacher training loss zero when sampling from the training distribution, ~x,;;~g Lo z = 0, which is true if the teacher attained a minimum of its training
loss. Our instant moment matching variant is thus indeed equivalent matching expected teacher gradients in parameter space.
B Experimental details
experiments were run on TPUv5e using 256 chips per experiment. For ImageNet we used global batch size of 2048, while for text-to-image used global batch size of 512. The base models were trained for 1M steps, requiring between days (Imagenet 64) to 2 weeks (text-to-image). We use the UViT architecture from (Hoogeboom al. 2023). Configurations largely correspond to those in the appendix of Hoogeboom et al. (2033) where we used their small model variant for our Imagenet experiments.
For Imagenet we distill the trained base models for a maximum of 200,000 steps, and for text-to-image we use a maximum of 50,000 steps. We report the best FID obtained during distillation, evaluating every 5,000 steps. We fix the random seed and data used in each evaluation to minimize biasing our results. We use the Adam optimizer (Kingma & Bal 2014) with 33 0, 3 0.99, c le- 12, We use learning rate warmup for the first 1,000 steps and then linearly anneal the learning rate to zero over the remainder of the optimization steps. We use gradient clipping with maximum norm of 1. We don't use an EMA, weight decay, or dropout.
13

---

C More model samples
Figure 4: Random samples for random ImageNet classes at the 64 x 64 resolution, from our 8-step distilled model, using the alternating optimization version of our algorithm.
Figure 5: Random samples for a single ImageNet class at the 64 x 64 resolution, from our 8-step distilled model. Visualizing samples from single class helps to assess sample diversity.
14

---

Figure 6: Random samples for random ImageNet classes at the 128 128 resolution from our 8-step distilled model, using the alternating optimization version of our algorithm.
15

---

Figure Selected 8-step samples from our distilled text-to-image model.
16

---

D Relationship to score matching
When distilling a diffiusion model using moment matching with alternating parameter updates [Section3.11 the auxiliary denoising model & is fit by minimizing Ep, w(zz) x & (2, )l2², with Z, and sampled from our distilled model. In the one-step case this is equivalent to performing score matching, as x (21 with ~ N(0, I), and 2, ~ q(z, (x) is sampled using forward diffusion starting from x. Recall here that q(z, |X, 2 q(z xx) as Z1 pure noise, uncorrelated with Z,. Upon convergence, we'll have that
(a. 2.)/o²- E (8x, | 2 Vz, log Py, Vz.:
i.e. our auxiliary model will match the score of the marginal sampling distribution of the distilled
model Pr, (2,) (because the optimal solution is (2.) = Ex-p,, {x). In the multi-step case this
equivalence does not hold, since the sampling distribution of n depends on zz. Instead the proper multistep score is given by:
V 22, log Pry (Za) (x. x.). V., log (og (11)
(18)
= p2, (8.2.122) (Mi-to (x,, 21) - -2a)/ (
(19)
with o2- =
(=1² 03
1
+
-1
and Pe-s = 02-05
attss
zt +
ag
(20)
z/e
This expression suggests we could perform score matching by denoising towards p4-- (x, 2,), linear combination and Zt, rather than just towards x. We tried this in early experiments, but did not get good results.
However, moment matching can still be seen match the proper score expression (equation 19 approximately, if we assume that the forward processes match, meaning Pe, (2.l2a) R (2,.22.). This then gives:
22, (ze)) 2
(x n. 2.),
( (x, 21)
(21)
Orls
ag x
1
1
= Ep (x (1(2)))))))
21
+
+
2
(22)
0²²
0/
1/8
0
1
a
Ex, 2.)
2x
x
+
2z
(23)
1/1
0²
0²
1/
= Ep, x.)
x
2
(24)
-0²
0
= (o, Ea, [x [,] - z,)) 20)
(25)
When this used to fit the auxiliary score 8 (2.) = (a.2x (2.) 2/// is equivalent to fitting (2.)) against just x, so under this approximation moment matching and score matching once again become equivalent.
If our distilled model has Pr, (x]2,) = q(x|z,). and if ~ (Zz)) (which is true during training), the approximation (z)())).) x q (z.) would become exact. Both multistep moment matching and multistep score matching thus have a fixed point that corresponds to the correct target distribution q. We currently do not have any results on guaranteeing when this fixed point is indeed attained for both methods, and exploring this further would make for useful future research. Note that in general Pr, (2: (z,) # q(z, (z.) until convergence, so during optimization moment matching and score matching indeed optimize different objectives.
E How different sampling algorithms average over model predictions
In Section we note that our distilled models often improve over their teacher, and we speculate that this may be due to the way the student model averages over teacher predictions. In theory, predictions of the clean data at different diffusion times should be mutually consistent, but since the diffusion model is implemented with unconstrained neural network this generally will not be the case in practice. Different sampling algorithms average over these predictions differently, which offers scope for improvement if more optimal weighting can be found.
17

---

Sampling from a diffusion models consists of repeatedly predicting the clean data given an inter- mediate noisy data point. Ignoring the dependence between these predictions (earlier predictions determine the input of later predictions), we can say that the final sample is then simple weighted average these predictions. In Figure [ we show how the standard ancestral (DDPM) sampler (Ho) et al., 2020) and the deterministic DDIM sampler (Song et al. 2021a) perform this weighting, where our neural network uses the y-parameterization of Salimans Ho (2022).
ddim
0.025
wowggttt
0.015
0.010
0.000
0.0
0.2
0.4
0.6
0.8
1.0
time
Figure 8: Implied weighting of v-predictions made during sampling in determining the final sample, for the DDPM and DDIM sampling algorithms.
18

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
All claims are supported by experimental evidence.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We're explicit about the limitations of automated metrics of image quality, and the need to do human evaluations in future work. We're also explicitly mentioning that other methods are currently better in the I-step and 2-step regime.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model cell-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
19

---

Answer: [Yes]
Justification: No formal theorems are included in the paper, but results like e.g. the equiv- alence between our method and Diff-Instruct in the 1-step case is fully derived in the paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
The description in the paper and technical appendix are enough to reproduce our main claim: Applying our distillation method on conditonnii ImageNet models yields 8-step samplers that match or outperform their teacher in FID. We are currently not able share data, code, and some technical details needed fully reproduce the text-to-image experiment.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky some cases, in which case authors welcome describe the particular way they provide for reproducibility.
20

---

In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [No]
Justification: We are currently unable share code but hope to be able do in the future. ImageNet data is publicly available, but we're 're not able to share the used text-to-image data.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines /tttps://nips.c/
public//uuddes/CoeeSuutttt wttpeeuuubiisiiooooooooooo for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed run reproduce the results. See the NeurIPS code and data submission guidelines (https
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The paper contains separate section hyperparameter choices and technical appendix with additional training/testing details.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Since running these experiments is expensive, we were not able to accurately characterize error bars and statistical significance.
Guidelines:
The answer NA means that the paper does not include experiments.
21

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The computational resources used are described in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental rums well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
We complied with all guidelines.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
22

---

Justification: This work presents a generic algorithm that contributes to improving the quality and cost of generative models general, and we do not see any impact of this particular work that differs from the societal impact of the broader field of research. We not release any model, code, or data that would benefit from implementing a mitigation strategy specific to this particular work.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
No models or data are released.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We cite sources for public data (ImageNet) and for our base models (Simple Diffusion) and explicitly mention using licensed data for the text-to-image experiment, but
not able to disclose full details on the latter.
23

---

Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e-g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of in the package should be provided. For popular datasets, paperssi thcode. com/ com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
No new assets are introduced.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
No crowdsoureing and research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]

---

Justification:
No human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
25


Reviews:
[{"review_id": "MpS9p22Rs5", "reviewer": "Reviewer_5zBL", "summary": "The paper presents a novel method for distilling diffusion models to require fewer function evaluations.  The method is based on moment-matching, and two practical algorithms are presented and evaluated.  Empirical results are state-of-the-art in few-step regimes.", "questions": "1.  In Eqn. (6), how is $\\mathbb{E}_q[x\\mid z_s]$ approximated?  Is it through samples from the dataset?  If so, are there additional assumptions that need to be made?  Naively, I'd expect that, especially as $s \\to 1$, the true data distribution would not necessarily be unimodal around the empirical samples in the training set.  Am I overthinking this?\n1.  Why does the instantaneous version of the algorithm have no reported results for 1 and 2 steps?", "limitations": "Although there is no limitations section in the paper, the discussion of the limitations is well addressed throughout the paper.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents a well-motivated novel method for distillation fo diffusion models.  This method achieves impressive results on standard benchmarks, and represents a solid contribution to the state of diffusion models for image generation.  The presentation is clear and complete, and the method is well situated in the literature both through empirical comparison and discussion.", "weaknesses": "This is a strong paper overall, without many weaknesses I noticed.  I would have liked to see more discussion of why the 1- and 2- step regimes do not perform as well as other approaches.  Is this a characteristic of the moment-matching approach?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your kind words.\nPlease find our answers to your questions below:\n\n> I would have liked to see more discussion of why the 1- and 2- step regimes do not perform as well as other approaches. Is this a characteristic of the moment-matching approach?\n\nAs the number of sampling steps is decreased, this class of methods becomes increasingly sensitive to its hyperparameters and parameter initialization, and the variance of results over different runs increases. The 1-step case can be made to work well if everything is tuned carefully, as shown by the DMD and DiffInstruct papers which can be interpreted as special cases of our method in this setting. Using their hyperparameters is not optimal for the multistep case however, so we chose not to focus on this. We will add some discussion on this in the revision of our paper.\n\n\n> In Eqn. (6), how is $\\mathbb{E}_q[x|z_s]$ is approximated? [...] are there any additional assumptions that need to be made?\n\nThis expectation is approximated with the pre-trained diffusion model, which in turn is trained by denoising samples from the dataset. No additional assumptions are required, though of course our distilled model will be affected by any errors in the teacher diffusion model that is used to approximate $\\mathbb{E}_q[x|z_s]$, as distillation methods typically are.\n\n> Why does the instantaneous version of the algorithm have no reported results for 1 and 2 steps?\n\nUsing the standard hyperparameters, the instantaneous algorithm does not provide competitive results when using fewer than 4 sampling steps, and can even diverge. We will add some discussion on this in the captions of the relevant results tables.", "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": "", "extra_scores": {}}, {"review_id": "DZcl3hseMO", "reviewer": "Reviewer_14eo", "summary": "This paper proposes a diffusion distillation algorithm based on moment matching. The starting point of the paper is to achieve distribution consistency by ensuring that the denoising process $\\widetilde{x} = g_\\eta(z_t,t)$ with fewer steps conforms to the true distribution, incorporating moment matching. To estimate the first moment of the student models' $x_0$ predictions, the paper suggests using a small surrogate denoising network to learn moment information. Simultaneously, the moment information from the surrogate denoising network is used to optimize the student models. The distillation model proposed in this paper demonstrates good performance with 8-step sampling", "questions": "Please see Weakness.", "limitations": "yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Research based on moment matching requires moment information. This paper proposes a surrogate denoising network for the first moment, based on the principle that one-step sampling in diffusion is unbiased for the first moment. The distillation model needed is trained using this moment network. The approach of using this surrogate is both reasonable and solid\n\n2. The experimental section of the paper demonstrates that, compared to different distillation models, the MM-based distillation model achieves the best results with 8-step sampling", "weaknesses": "1. Theoretically, according to the current setting of the paper, the discretization error in the distillation model sampling originates from the gap between the one-step prediction of the distillation model and the first moment of the true teacher models' denoising. It is worth exploring the values of L(ϕ)  at different timesteps after training.\n\n2. The experiments revealed that the performance of the distillation model worsens as the number of steps increases after training (FID of 3.0 for one step and 3.86 for two steps). Could you provide a theoretical explanation for this phenomenon?\n\n3. For the distillation model, the number of steps in the solver of pretrained diffusion models significantly affects the results. This paper lacks an analysis of the solvers' step counts, such as the impact on results when reducing from DDIM 50 steps to DDIM 25 steps.\n\n4. Theoretically, the method proposed in this paper can be applied to any distillation model. If the distillation model uses a pre-trained distillation model for initialization, the accuracy of the x0x0​ predictions will be higher in the early stages of training. Consequently, the optimization effect of the moment estimation surrogate network will be better. Further fine-tuning on this basis could additionally demonstrate the extent to which moment matching improves the distillation model. Could this experiment be supplemented？", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your review. Please find our response to your remarks below:\n\n\n> 1. [...] It is worth exploring the values of L(ϕ) at different timesteps after training.\n\nThe instantaneous version of L(ϕ) is indeed informative, and we investigate its value over different training steps in section 5.4 (averaged over all sampling timesteps). Can we interpret your suggestion as adding a figure like this for each of the 4-8 sampling timesteps individually? Alternatively, we could make an interpretable version of L(ϕ) in our alternating algorithm by freezing the student model after distillation and then optimizing the auxiliary diffusion model (ϕ parameters) until convergence, which we could then again report for each individual sampling timestep. We’d be happy to include either of these with a discussion in the appendix.\n\n\n> 2. The experiments revealed that the performance of the distillation model worsens as the number of steps increases after training (FID of 3.0 for one step and 3.86 for two steps). Could you provide a theoretical explanation for this phenomenon?\n\nIn general we see improved performance when increasing the number of sampling steps. The only exception to this are our results on Imagenet 64x64 when taking 1-2 steps, which yield the FIDs mentioned in your comment. The explanation here is that we find that this class of methods is sensitive to hyperparameters and random seeds when taking very few sampling steps. We did not separately optimize hyperparameters for this case. The relationship between sampling steps and performance is discussed already, but we will add some discussion on the specific Imagenet results that you pointed out.\n\n> 3. For the distillation model, the number of steps in the solver of pretrained diffusion models significantly affects the results. This paper lacks an analysis of the solvers' step counts, such as the impact on results when reducing from DDIM 50 steps to DDIM 25 steps.\n\nHere there might be a slight misunderstanding: Unlike some other distillation methods, we never run a full multi-step solver (e.g. DDIM with 50 or 25 steps) on our pretrained diffusion model during distillation. Instead, we sample from the student model and use a single-step evaluation of the pretrained model to provide a training signal to the student. Hence, the distilled model has no dependence on a choice of solver for the base model. Of course the reported baseline (undistilled) numbers in our results tables do depend on the choice of solver: Here we use a tuned 1024 step stochastic sampler to provide the strongest possible baseline for comparison. If useful we could add some results here using a faster sampler (e.g. DDIM with 25 steps). Please let us know if this is what you have in mind.\n\n> 4. Theoretically, the method proposed in this paper can be applied to any distillation model. If the distillation model uses a pre-trained distillation model for initialization, the accuracy of the x0x0​ predictions will be higher in the early stages of training. Consequently, the optimization effect of the moment estimation surrogate network will be better. Further fine-tuning on this basis could additionally demonstrate the extent to which moment matching improves the distillation model. Could this experiment be supplemented?\n\nPlease clarify your suggestion for the additional experiment and we’d be happy to add it if feasible. Both the distilled (student) model and the auxiliary denoising model are currently initialized from the pretrained diffusion model. Is your suggestion to run moment matching distillation warm-started from a model distilled using another algorithm such as progressive distillation or consistency distillation? If so, that sounds like an interesting experiment we’d be happy to run.", "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": "", "extra_scores": {}}, {"review_id": "WCxEmlNsVr", "reviewer": "Reviewer_Xksr", "summary": "The authors proposed a fast sampling method by distilling a diffusion model to model $q(x|z_t)$. This is achieved by matching moments, with two novel approaches proposed to implement it in practice.", "questions": "1. In tables 1 and 2, by increasing NFE to the level of the base model, what the FIDs and IS will be?\n\n2. I understand that the sampling in the base diffusion model may not be optimal. However, which mechanisms in the proposed methods can achieve better FID than the base model with distillation?\n\n3. How is the generation quality and FID by using instantaneous version of the moment matching (algorithm 3) on ImageNet $128\\times128$?\n\n4. Can the proposed method be compatible with Consistency Model for further timestep distillation?\n\n5. Also, can the method support training from scratch?\n\n6. In Algorithm 3, storing and computing both teacher gradients and $\\eta$-gradients appears costly, particularly as the model's scale (in terms of trainable parameters) increases. I'm curious about the scalability of this distillation method to large latent diffusion models.", "limitations": "The authors provided a discussion on potential limitations on evaluations and 1 or 2 steps generation.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "- The idea of distilling to model the conditional distribution $q(x|z_t)$ by using matching the moments is interesting and novel. \n\n- The equivalence between the proposed instantaneous version of the moment matching loss and matching expected teacher gradients is interesting and insightful.", "weaknesses": "1. Some of the notations should be clarified or made coherent. For instance, do both $g_{\\theta}(z_t, t)$ and $g_{\\theta}(z_t)$ indicate the parametrized denoiser? What's the definition of $q(z_t)$ in Eq. (6)? In Eq. (4) and Eq. (6), what depends on $\\eta$ (please be consistent about dropping this parameter)? \n\n2. Algorithm 2 is somehow heuristics. It may be hard to see if $\\mathbb{E}_{g}[\\tilde{x}|z_s]$ can be approximated well with the auxiliary network model.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your review. Please find our response to your comments and questions here:\n\n> Notation\n\nThanks! We’ll clarify the notation issues you identified in the paper.\n$g_{\\theta}(z_t, t)$ and $g_{\\theta}(z_t)$ indeed refer to the same denoising model. Dropping the dependence on $t$ is customary, but we’ll make this explicit in the text.\nq() refers to the forward diffusion process, i.e. $q(z_t) = \\int q(z_t | x) q(x) dx$ as defined in equation 3. We’ll add a comment on this after equation 6.\nDependency on $\\eta$ in equations 4 and 6 is through $g$. We’ll make this explicit.\n\n\n> It may be hard to see if $\\mathbb{E}_{g}[\\tilde{x}|z_s]$ can be approximated well with the auxiliary network model.\n\nThis is indeed an approximation. However since this auxiliary denoising task is almost identical to the original denoising problem, just using different samples, we do not expect this to be a practical limitation.\n\n> In tables 1 and 2, by increasing NFE to the level of the base model, what the FIDs and IS will be?\n\nAs the NFE is increased to the level of the base model, the FID and IS converge to those of the base model. For sufficiently large NFE, the distillation loss becomes zero at initialization (starting from the parameters of the base model) and hence distillation does not change the model.\n\n> Which mechanisms in the proposed methods can achieve better FID than the base model with distillation?\n\nWe hypothesize the improvement with respect to ancestral sampling on the base model to derive from having a different implied weighting on the denoising predictions made at different timesteps during sampling. The final sample $x$ produced by a diffusion model sampler can be written as a linear combination of predictions made along the sampling path. For different types of samplers the weighting of these predictions look different: e.g. deterministic DDIM sampling puts more weight on the first few predictions compared to stochastic ancestral sampling (also see the discussion in the Karras et al. EDM paper on this topic). Our distillation method trains a student model that similarly tries to match an implicit weighted average of teacher model predictions, with a weighting that apparently works better than the standard samplers. Of course this is all dependent on the chosen hyperparameters and sampling schedule, and for different choices our method could potentially perform less well than standard samplers applied to the base model. Our explanation here is speculative, and we do not know of a simple experiment that would clearly confirm this, but we’d be happy to include some of this discussion in the paper.\n\n> How is the generation quality and FID by using instantaneous version of the moment matching (algorithm 3) on ImageNet 128x128?\n\nThese results are shown at the bottom of Table 2 in the paper\n\n> Can the proposed method be compatible with Consistency Model for further timestep distillation?\n\nOne possibility would be to first distill a diffusion model using consistency distillation, and then further improve it by finetuning with our moment matching objective. Another reviewer may have asked for an experiment like this (if we understand them correctly), so we’re likely to include this in the updated paper. Would this be of interest?\n\n> Can the method support training from scratch?\n\nThe proposed distillation methods (as well as other related methods from the literature) are sensitive to the initialization of the student, and do not give good results when training from scratch. We hope that future work can improve this class of methods further to enable training from scratch.\n\n> In Algorithm 3, storing and computing both teacher gradients and $\\eta$ gradients appears costly, particularly as the model's scale (in terms of trainable parameters) increases. I'm curious about the scalability of this distillation method to large latent diffusion models.\n\nRegarding storage (memory): Since gradients and parameters can be fully sharded over devices (see e.g. the ZeRO optimizer work) this is not a bottleneck in practice. In our experiments we went up to 5B parameters, using small-memory accelerators (TPU v5e), and we could have scaled even further. Regarding compute: A full iteration of algorithm 3 costs approximately the equivalent of 8 forward passes through the diffusion model (2x forward on student, fwd + bwd on the teacher for batch 1, fwd on the teacher for batch 2, backward on student for batch 2, where full backward passes are generally about twice as expensive as forward passes). This scaling does not change with the number of parameters. In absolute terms this makes one iteration of our method about equally expensive as the approaches in DMD and DiffInstruct, and about 60% more expensive than one iteration of Progressive Distillation. We could add some discussion on this to the paper if useful.", "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": "", "extra_scores": {}}, {"review_id": "uSveHyMtc2", "reviewer": "Reviewer_cU8D", "summary": "This paper proposes an approach to distill a diffusion model into a multi-step generator. Building on previous works that use distribution matching to train a few-step student generator, the paper introduces a novel method of matching the conditional expectation of clean data given noisy data along the sampling trajectory. By focusing on the matching of the first-order moment, the authors derive a gradient for training the generator, which can be approximated using two diffusion models (as seen in Diff-Instruct, DMD, and SwiftBrush) or estimated using batch statistics. The primary contribution is conditioning on noisy data during expectation matching, distinguishing this approach from previous score distillation methods that match unconditioned expectations. Extensive results demonstrate the effectiveness of this method.", "questions": "Please find a set of suggestions in the weakness section above.", "limitations": "Yes.", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "S1. This paper provides a new perspective on recent distribution matching-based diffusion distillation approaches such as Diff-Instruct, DMD, and SwiftBrush. The moment matching formulation is novel and theoretically sound, leading to practical improvements including conditioning on noisy input to the generator, as well as a different way of training the model that doesn't rely on an auxiliary diffusion model\n\nS2. The experimental validation is comprehensive, with the 8-step model achieving exceptionally strong image generation results.", "weaknesses": "Minor Typo: There is a typo on line 72: \"as described is described.\"\n\nLack of Detailed Ablation Studies: There are several instances where the authors claim one formulation is better based on their trials but do not provide exact ablation studies. Including these studies would help readers understand the effectiveness of alternative formulations. Specific examples include:\n\n- Exclusion of the z_s dependence on the parameters when calculating gradients of the moment matching loss (line 90).\n- Training the auxiliary diffusion model using a mix of teacher predictions and generated samples (line 110).\n\nComparison Between Optimization Setups: Further comparison between the alternating and instant optimization setups would be beneficial. Currently, alternating optimization appears better, but instant optimization feels more principled. The authors should discuss this discrepancy in more detail and highlight potential challenges to improve performance.\n\nEnhancement of Conditional Sampling Ablation: The ablation study on the importance of conditional sampling could be clearer. The paper shows that without conditioning on noisy input, the generator produces less diverse images with drifted tones. However, other methods like Diff-Instruct, DMD, and SwiftBrush, which don't condition on noisy input, work without these diversity/tone issues. This seems contradictory to existing literature. Does this issue arise as the number of sampling steps increases? What is the failure point for other unconditional methods? These experiments and discussions are crucial, as this is the paper's most significant contribution and claim.\n\nOn L181, the authors mention that \"One could imagine fitting x through score matching against the conditional distribution q(z_s | z_t, x_~) but this did not work well when we tried it.\" Could the authors elaborate on \"score matching against the conditional distribution\"?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your review and kind words. Please find our response to your comments below:\n\n\n* Requests for additional ablations:\n\n> Exclusion of the z_s dependence on the parameters when calculating gradients of the moment matching loss (line 90).\n\nWe’d be happy to include this ablation in the camera ready update of the paper.\n\n> Training the auxiliary diffusion model using a mix of teacher predictions and generated samples (line 110).\n\nWe are experimenting with this for video and other modalities, but consider this to be out of scope for the current paper.\n\n* Additional comparison between alternating and instant optimization setups.\n> Further comparison between the alternating and instant optimization setups would be beneficial. Currently, alternating optimization appears better, but instant optimization feels more principled. The authors should discuss this discrepancy in more detail and highlight potential challenges to improve performance.\n\nWhen taking very few sampling steps, the main discrepancy is that the instant approach drops modes more quickly due to the conditioning of the optimization problem. For 8+ steps both approaches are very similar. We can add some additional discussion on this to the paper.\n\n> The ablation study on the importance of conditional sampling could be clearer. The paper shows that without conditioning on noisy input, the generator produces less diverse images with drifted tones. However, other methods like Diff-Instruct, DMD, and SwiftBrush, which don't condition on noisy input, work without these diversity/tone issues. This seems contradictory to existing literature.\n\nIn our experience, methods from the literature (e.g. Diff-Instruct, DMD) are actually extremely sensitive to hyperparameters, weighting functions, parameter initialization etc. They work for a much smaller range of hyperparameters than our multistep distillation method, and even runs that produce the FIDs reported in the literature often need to be stopped early because they eventually still end in mode collapse. A full investigation of the stability of these other methods is out of scope for this paper, but we could add some more discussion on this if useful.\n\n> On L181, the authors mention that \"One could imagine fitting $\\hat{x}_{\\phi}$ through score matching against the conditional distribution $q(z_s | z_t, \\tilde{x})$ but this did not work well when we tried it.\" Could the authors elaborate on \"score matching against the conditional distribution\"?\n\nIn standard score matching, like used in DDPM and Diff-Instuct/DMD, we minimize the average L2 distance between an estimate score $s_{\\phi}(z_s)$ and the score of the forward process $\\nabla_{z_s} \\log q(z_s | \\tilde{x})$. One way of extending this to the conditional multistep setting would be to make both expressions conditional on $z_t$, i.e. minimize the average L2 between $s_{\\phi}(z_s; z_t)$ and $\\nabla_{z_s} \\log q(z_s | z_t, \\tilde{x})$. In early experiments we did not find this to work well for distillation, so we chose to go with the current moment matching objective instead, which is fundamentally different. In the updated version of the paper we have an appendix with a more detailed discussion of the relationship to score matching. (This is already finished, but NeurIPS does not allow us to add new math/text content at this stage of the reviewing process)", "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": "", "extra_scores": {}}, {"review_id": "GAZnnCHDdz", "reviewer": "Reviewer_V8xX", "summary": "This paper presents a new method for making diffusion models faster to sample. The method distills many-step diffusion models into few-step models and extends recently proposed one-step methods to the multistep case by moment matching. By using up to 8 sampling steps, the obtained distilled models outperform not only their one-step versions but also their original many-step teacher models, obtaining new SOTA results on the Imagenet dataset.", "questions": "1. the analysis of Ablating conditional sampling is confusing. It’s not unusual that conditioning leads to much better sample diversity in the distilled model because you use extra label information. Could the authors explain more about the ablating conditional sampling?\n2. Why you select 8 sampling steps? How more or less eight steps will affect experimental results, run speeds, etc?", "limitations": "A limitation of this study is that the theoretical part is difficult and too mathematical to understand.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper presents a new method for making diffusion models faster to sample. The method distills many-step diffusion models into few-step models and extends recently proposed one-step methods to the multistep case by moment matching. \n2. The paper is well writing. The proposed moment matching proposed in this paper is a technique worthy of reference and further in-depth study by researchers. \n3. Experimental results show the effectiveness of the proposed method and the results of experiments achieve SOTA.", "weaknesses": "1. Moment Matching Distillation Part is difficult and too mathematical to understand. The authors could improve the writing to make it easier for the reader to understand.\n2. In line 446, \"because the ﬁrst term of $L( \\phi)$ can be seen as a standard diffusion loss at $\\theta$.\" I don't understand, the author needs to give more interpretations.\n2. The experiment of Class-conditional generation is only on ImageNet. The authors could provide more experiment results on other datasets, such as celeba, lsun and so on in later versions.\n3. The experiments lack ablating sampling steps. The authors could provide more experiment results adapting other sampling steps in later versions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your review. Please find our response to your comments below:\n\n> Moment Matching Distillation Part is difficult and too mathematical to understand.\n\nAppendix A of the paper contains the mathematical details we felt we could reasonably omit from the main text. In our updated version of the paper we’ll have another mathematical appendix with some of the more tangential mathematical analysis. The remaining math (e.g. the definition of the sampling distributions) we think is essential to describing the core contributions of the paper. Apart from the method section we usually describe things in words rather than math. If you have concrete suggestions for making things easier to understand we’d be very happy to consider them.\n\n> In line 446, \"because the ﬁrst term of $L(\\phi)$ can be seen as a standard diffusion loss at $\\theta$.\" I don't understand, the author needs to give more interpretations.\n\nThis refers to the loss $L(\\phi)$ as defined in Algorithm 2. The first term is a standard diffusion loss, which is evaluated at $\\phi = \\theta$, i.e. using the parameters of the pretrained teacher model. We’ll make this easier to read by repeating the definition for $L(\\phi)$ and by referring to Algorithm 2.\n\n> The experiment of Class-conditional generation is only on ImageNet. The authors could provide more experiment results on other datasets, such as celeba, lsun and so on in later versions.\n\nWe could easily add one of these additional datasets to the appendix of the updated paper, but in our experience they are not very informative as these datasets are of much lower diversity than Imagenet and are more prone to overfitting. Please note that the paper does include experiments on a large high-resolution text-to-image dataset (section 5.5) which we think is more informative.\n\n> The experiments lack ablating sampling steps. The authors could provide more experiment results adapting other sampling steps in later versions.\n\nThe paper currently presents results using 1, 2, 4, and 8 sampling steps. Are you suggesting we add results for 16 steps and higher? Qualitatively, results for between 16 and 1024 sampling steps lie somewhere in between the presented results for our distilled model at 8 sampling steps and the undistilled baseline model at 1024 steps. We could add some discussion on this if useful.\n\n> The analysis of Ablating conditional sampling is confusing. It’s not unusual that conditioning leads to much better sample diversity in the distilled model because you use extra label information. Could the authors explain more about the ablating conditional sampling?\n\nThere seems to be a misunderstanding here: The experiment in Section 5.2 ablates conditioning on intermediate steps z_t, which is something that distinguishes our approach from other methods in the literature. We do not ablate conditioning on label information.\n\n> Why you select 8 sampling steps? How more or less eight steps will affect experimental results, run speeds, etc?\n\nWe obtain the best results for our method using 8 sampling steps, but the results tables also report using 1, 2, and 4 steps. Results for 4 steps are still very competitive. Speed / running time scales linearly with the number of steps.", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes a new distillation method for accelerating diffusion model inference. The key contribution is the notion of distribution matching between the teacher and student models, via an approximation of moment matching. Reviewers are all positive about the idea and empirical results, reflected by unanimous votings of acceptance. The AC believes that this is a great paper and recommends acceptance.
Decision: Accept (poster)