Multistep Distillation of Diffusion Models via Moment Matching

Authors: tim  salimans
thomas  mensink
jonathan  heek
emiel  hoogeboom
google  deepmind
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 20:38:18

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.

Main paper: 
 introduction: 
 Diffusion models (Ho et al., 2020; Song & Ermon, 2019; Sohl-Dickstein et al., 2015) have recently become the state-of-the-art 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 to synthesize complex outputs tractable. Although this decomposition simplifies the training objective compared to alternatives like GANs, it 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 of distribution matching in terms of matching 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 perspective 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. 
diffusion models: 
 Diffusion models are trained by learning to invert a noise process that gradually destroys data from a clean data sample x according to z t = α t x + σ t t with z t ∼ N (0, I), where α t , σ t are monotonic functions of diffusion time t ∈ [0, 1]. The coefficients α t , σ t may be specified in multiple equivalent ways. Here, we use the variance preserving specification (Ho et al., 2020) that has σ 2 t = 1 -α 2 t , α 0 = σ 1 = 1, and α 1 = σ 0 = 0, such that we have that z 0 = x and z 1 ∼ N (0, 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: SNR(t) = α 2 t /σ 2 t , rather than the coefficients individually (Kingma et al., 2021) . To invert the specified diffusion process, we can sample from the posterior distribution: q(z s |z t , x) = N (z s |µ t→s (z t , x), σ t→s ), (2)To sample from a learned diffusion model, we replace x by a prediction from a neural network x = g θ (z t , t) that is fit to the data by minimizing E t∼p(t),zt,x∼q(zt,x) w(t) x -g θ (z t ) 2 , with weighting function w(t) and where q(z t , x) denotes sampling x from the data and then producing z t by forward diffusion. The sampling process starts with pure noise z 1 ∼ N (0, I) and iteratively denoises the data according to q(z s |z t , x) for a discrete number of timesteps k, following Algorithm 1.If we attain the optimal solution x = E[x|z t ] and let k → ∞ the sampling process becomes exact, then the learned diffusion model can be shown to be a universal distribution approximator (Song et al., 2021b) . To get close to this ideal, k typically needs to be quite large, making diffusion models a very computationally expensive class of models (Luccioni et al., 2023) .An alternative to the well-known maximum likelihood estimation method is the method of moments, also known as moment matching. Traditionally for univariate distributions, one matches moments m k = E x∼p X [x k ] of a random variable X. The canonical example is a Gaussian distribution, which is defined by the first two moments (i.e. the mean and variance) and all (centered) higher order moments Algorithm 1 Ancestral sampling algorithm used for both standard denoising diffusion models as well as our distilled models. For standard models typically 256 ≤ k ≤ 1000, for distilled 1 ≤ k ≤ 16.Require: Denoising model g θ (z t , t), number of sampling steps k Initialize noisy data z 1 ∼ N (0, I) for t ∈ {1, (k -1)/k, . . . , 2/k, 1/k} do Predict clean data using x = g θ (z t , t) Set next timestep s = t -1/k Sample next noisy data point z s ∼ q(z s |z 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 a consistent parameter estimation method, and can be readily extended to multivariate distributions, e.g. by matching the mean and covariance matrix for a multivariate Gaussian.One can generalize the method of moments to arbitrary high dimensional functions f : R d → R k and match the moment vector m as defined by: m= E x∼p X [f (x)],which is called the Generalized Method of Moments (GMM, Hansen (1982) ). Matching such moments can be done by minimizing a distance between the moments such as||E x∼p θ f (x) -E x∼p X f (x)|| 2where p θ is the generative model and p X the data distribution. The distillation method we propose in the next section can be interpreted as a special case of this class of estimation methods.Question: Does the paper discuss the limitations of 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 1-step and 2-step regime.Guidelines:• The answer NA means that the paper has no 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 a 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 (e.g., independence assumptions, noiseless settings, model well-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 a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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 fairness. • 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 acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
moment matching distillation: 
 Many-step sampling from diffusion models starts by initializing noisy data z 1 ∼ N (0, I), which is then iteratively refined by predicting the clean data using x = g θ (z t , t), and sampling a slightly less noisy data point z s ∼ q(z s |z t , x) for new timestep s < t, until the final sample is obtained at s = 0, as described is described in Algorithm 1. If x = E q [x|z t ] this procedure is guaranteed to sample from the data distribution q(x) if the number of sampling steps grows infinitely large. Here we 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 g θ into a new model g η (z t , t) which we sample from using the same algorithm, but with a strongly reduced number of sampling steps k, for say 1 ≤ k ≤ 8.To make our model produce accurate samples for a small number of sampling steps k, the goal is now no longer for x = g η (z t , t) to approximate the expectation E q [x|z t ] but rather to produce an approximate sample from this distribution. In particular, if x ∼ q(x|z t ) then Algorithm 1 produces exact samples from the data distribution q for any choice of the number of sampling steps. If g η perfectly approximates q(x|z t ) as intended, we have thatEQUATIONIn words: The conditional expectation of clean data should be identical between the data distribution q and the sampling distribution g of the distilled model.Equation 3 gives 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., 2021b) . These moment conditions can be used as the basis of a distillation method to finetune g η (z t , t) from the denoising model g θ . In particular, we can fit g η to q by minimizing the L2-distance between these moments:EQUATION)In practice, we evaluate the moments using a sample z s from our generator distribution, but do not incorporate its dependence on the parameters η 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 L(η) is then given byEQUATIONwhere we approximate the first expectation using a single Monte-Carlo sample x and where the second term is zero as it does not depend on g η . Following this approximate gradient is then equivalent to minimizing the lossEQUATIONwhere sg denotes stop-gradient. This loss is minimized if E g [x|z s ] = E q [x|z s ] as required. Unfortunately, the expectation E g [x|z s ] 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 g [x|z s ] by a second denoising model, and in Section 3.2 we instead apply moment matching directly in parameter space rather than x-space.Our first approach to calculating the moment matching objective in equation 6 is to approximateE g [x|z s] with an auxiliary denoising model g φ trained using a standard diffusion loss on samples from our generator model g η . We then update g φ and g η in alternating steps, resulting in Algorithm 2. x = g η (z t ), z s ∼ q(z s |z t , x). if n is even then Minimize L(φ) = w(s){ x -g φ (z s ) 2 + g θ (z s ) -g φ (z s ) 2 } w.r.t. φ else Minimize L(η) = w(s)x T sg[g φ (z s ) -g θ (z s )] w.r.t. η end if end forHere we have chosen to train our generator g η on all continuous times t ∈ (0, 1] even though at inference time (Algorithm 1) we only evaluate on k discrete timesteps. Similarly we train with randomly sampled time delta δ t rather than fixing this to a single value. These choices were found to increase the stability and performance of the proposed algorithm. Further, we optimize g φ not just to predict the sampled data x but also regularize it to stay close to the teacher model g θ : On convergence this would cause g φ to predict the average of x and g θ , 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 ≥ 8 sampling steps. The algorithm also closely resembles Variational Score Distillation as previously used for distilling 1-step generators g η in Diff-Instruct. We discuss this relationship in Section 4.Alternating optimization of the moment matching objective (Algorithm 2) is difficult to analyze theoretically, 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 φ * is determined using a single infinitesimal gradient descent step on L(φ) (defined in Algorithm 2), evaluated on a single minibatch. Starting from teacher parameters θ, and preconditioning the loss gradient with a pre-determined scaling matrix Λ, we can define:EQUATIONNow we use φ(λ) in calculating L(η) from Algorithm 2, take the first-order Taylor expansion for g φ(λ) (z s ) -g θ (z s ) ≈ λ ∂g θ (zs) ∂θ (φ(λ) -θ) = λ ∂g θ (zs) ∂θ Λ∇ φ L(φ)| φ=θ , and scale the loss with the inverse of λ to get:EQUATIONwhere ∂g θ (zs) ∂θ is the Jacobian of g θ , and where ∇ φ L(φ) is evaluated on an independent minibatch from x and z s . 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-vector products.The loss can now equivalently be expressed as performing moment matching in teacher-parameter space rather than x-space. Denoting L θ (x, z s ) ≡ w(s) x -g θ (z s ) 2 , and letting Linstant (η) = L instant (η) + constant, we have (as derived fully in Appendix A): 10) where the gradient of the teacher training loss is zero when sampling from the training distribution, E x,z s ∼q ∇ θ L θ (x, z s ) = 0, if the teacher attained a 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) , a classic method in statistics where a teacher model p θ is first estimated using maximum likelihood, after which its gradient is used to define a moment matching loss for learning a second model g η . 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.Linstant (η) ≡ 1 2 E zt∼q,x=gη(zt),zs∼q(zs|zt,x) ∇ θ L θ (x, sg(z s )) 2 Λ (9) = 1 2 E zt∼q,x=gη(zt),zs∼q(zs|zt,x) ∇ θ L θ (x, sg(z s ))-E x,z s ∼q ∇ θ L θ (x, z s ) 2 Λ ,(The moment matching loss Linstant (η) is minimized if the teacher model has zero loss gradient when evaluated on data generated by the distilled student model g η . In other words, optimization is 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 3 and Figure 2 .Algorithm 3 Parameter-space moment matching algorithm with instant denoising model g φ * .Require: Pretrained denoising model g θ (z t ), generator g η to distill, gradient scaling matrix Λ, number of sampling steps k, time sampling distribution p(s), loss weight w(s), and dataset D.for n = 0:N do Sample target time s ∼ p(s), sample time delta δ t ∼ U [0, 1/k].Set sampling time t = minimum(s + δ t , 1).Sample two independent batches of data from D and do forward diffusion to produce z t , z t .For both batches sample z s , z s from the distilled generator using x = g η (z t ), z s ∼ q(z s |z t , x).Evaluate teacher gradient on one batch:ν = Λ∇ θ L θ (x , z s )On the other batch, minimize L instant (η) = w(s)x T sg ∂g θ (zs) ∂θ ν w.r.t. η end for 0 1Figure 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.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 Diff-Instruct (Luo et al., 2024) and DMD (Yin et al., 2023) . We use the EDM test time noise schedule for p(s), as well as their training loss weighting for w(s), but we shift all log-signal-to-noise ratios with the resolution of the data following Hoogeboom et al. (2023) . For our gradient preconditioner Λ, 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, Λ is not updated.To get stable results for small numbers of sampling steps (k = 1, 2) we find that we need to use a weighting function w(s) with less emphasis on high-signal (low s) data than in the EDM weighting.Using a flat weight w(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 g θ . 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.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
related work: 
 In the case of one-step sampling, our method in Algorithm 2 is a special case of Variational Score Distillation, Diff-Instruct, and related methods (Wang et al., 2024; Luo et al., 2024; Yin et al., 2023; Nguyen & Tran, 2023) which distill a diffusion model by approximately minimizing the KL divergence between the distilled generator and the teacher model:EQUATION≈ E p(s),pη(zs) [(w(s)/2σ 2 s )( z s -αsg[x θ (z s )] 2 -z s -αsg[x φ (z s )] 2 )] (12) = E p(t),pη(zs) [(w(s)α 2 s /σ 2 s )x T sg[x φ (z s ) -xθ (z s )]] + constant (13)Here sg again denotes stop gradient, p η (z s ) is defined by sampling x = g η (z 1 ) with z 1 ∼ N (0, I), and z s ∼ q(z s |x) is sampled using forward diffusion starting from x. The auxiliary denoising model xφ is fit by minimizing E gη w(z s ) xxφ (z s ) 2 , which can be interpreted as score matching because z s is sampled using forward diffusion started from x. In our proposed algorithm, we sample z s from the conditional distribution q(z s |x, z t ): If z t = z 1 ∼ N (0, I) is assumed to be fully independent of x, i.e. that α 2 1 /σ 2 1 = 0, we have that q(z s |x, z 1 ) = q(z s |x) so the two methods are indeed the same. However, this correspondence does not extend to the multi-step case: When we sample z s from q(z s |z t , x) for α 2 t /σ 2 t > 0, fitting xφ through minimizing E gη w(z s ) xxφ (z s ) 2 no longer corresponds to score matching. 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 (see Appendix D for more detail). Instead, our moment matching perspective offers a justification for extending this class of distillation methods to the multistep case without changing the way we fit xφ . 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.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 q(z s |z t , 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 (g θ ). This discriminator is then used to define an adversarial divergence which is minimized w.r.t. g η :EQUATION)The methods differ in their exact formulation of the adversarial divergence D adv , in the sampling of time steps, and in the use of additional losses. For example Xu et al. (2023a) train unconditional discriminators D φ (•, t) and decompose the adversarial objective in a marginal (used in the discriminator) and a conditional distribution approximated with an additional regression model. Xiao et al. (2021) instead use a conditional discriminator of the form D φ (•, z t , t). 
experiments: 
 We evaluate our proposed methods in the class-conditional generation setting on the ImageNet dataset (Deng et al., 2009) , which is the most well-established benchmark for comparing image quality. On this dataset we also run 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 to show our approach can also be scaled to this setting. (Jabri et al., 2023) 281M 1000 1.23 67 our base model 400M 1024 1.42 84 DDIM (Song et al., 2021a) 10 18.7 TRACT (Berthelot et al., 2023) 7.43 4.97 2.93 2.41 CD (LPIPS) (Song et al., 2023) 6.20 4.70 4.32 iCT-deep (Song & Dhariwal, 2023) 3.25 2.77 PD (Salimans & Ho, 2022 (Heek et al., 2024) 1.2B 3.2 1.9 1.6 1.4 73 CTM (Kim et al., 2024) 1.73 64 DMD (Yin et al., 2023) 2.62 Diff-Instruct (Luo et al., 2023) 5.57Alternating (c.f. Sect. We begin by evaluating on class-conditional Im-ageNet generation, at the 64 × 64 and 128 × 128 resolutions (Tables 1 and 2 ). Our results here are for a relatively small model with 400 million parameters based on Simple Diffusion (Hoogeboom et al., 2023) . We distill our models for a 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.For our base models we report results with slight classifier-free guidance of w = 0.1, which gives the optimal FID. We also use an optimized amount of sampling noise, following Salimans & Ho (2022), which is slightly higher compared to equation 2. For our distilled models we obtained better results without classifier-free guidance, and we use standard ancestral sampling without tuning the sampling noise. We use identical 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-Instruct, adversarial methods).Ranking the different methods by FID, we find that our moment matching distillation method is especially competitive when using 8+ sampling steps, where it sets new state-of-the-art results, beating out even the best undistilled models using more than 1000 sampling steps, as well as its teacher model. For 1 sampling step some of the other methods show better results: Further improvement is likely possible by separately optimizing hyperparameters for this setting. For 8+ sampling steps we get similar results for our alternating optimization version (Section 3.1) and the instant 2-batch version (Section 3.2) of our method. For fewer sampling steps, the alternating version performs better.We find that our distilled models also perform very well in terms of Inception Score (Salimans et al., 2016 ) even though we did not optimize for this. By using classifier-free guidance the Inception Score can be improved further, as we show in Section 5.3.How can a distilled model improve upon its teacher? On Imagenet our distilled diffusion model with 8 sampling steps and no classifier-free guidance outperforms its 512-step teacher with optimized guidance level, for both the 64 × 64 and 128 × 128 resolution. This result might be surprising since the many-step teacher model is often seen as the gold standard for sampling quality. However, even the teacher model has prediction error 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 E, which offers scope for improvement.Similarly, prediction error will not be constant over the model inputs z t , 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 to study the improvement of moment matching over our base sampler in further detail, and test our hypotheses about its causes.The distilled generator in our proposed method samples from the conditional q(z s |x, z t ), whereas existing distillation methods based on score matching typically don't condition on z t . Instead they apply noise independently, mirroring the forward diffusion process used during training the original model. When using a 1-step sampling setup, the two approaches are equivalent since any intermediate z s will be independent from the starting point z 1 if that point has zero signal-to-noise. In the multistep setup the two approaches are meaningfully different however, and sampling from the conditional q(z s |x, z t ) or the marginal q(z s |x) are both valid choices. We ablate our choice of conditioning on z t versus applying noise independently, and find that conditioning leads to much better sample diversity in the distilled model, as shown in Figure 3 . Our choice of sampling from the conditional yields much better sample diversity.Our distillation method can be used with or without 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 2batch version we apply guidance and clipping to the teacher model and then calculate its gradient with a straight through approximation. Exper-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. A unique advantage of the instant 2-batch version of our moment matching approach is that, unlike most other distillation methods, it has a simple loss function (equation 9) that is minimized without adversarial techniques, bootstrapping, 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 × 128 in the adjacent figure : The typical behavior we see is that the loss tends to go up slightly for the first few optimization steps, after which it exponentially falls to zero with increasing number of parameter updates. To investigate our proposed method's potential to scale to large text-to-image models we train a pixel-space model (no encoder/decoder) on a licensed dataset of text-image pairs at a resolution of 512 × 512, using the UViT model and shifted noise schedule from Simple Diffusion (Hoogeboom et al., 2023) and using a T5 XXL text encoder following Imagen (Saharia et al., 2022) . We compare the performance of our base model against an 8-step distilled model obtained with our moment matching method. In Table 3 we report zero-shot FID (Heusel et al., 2017) and CLIP Score (Radford et al., 2021) on MS-COCO (Lin et al., 2014) : Also in this setting we find that our distilled model with alternating optimization exceeds the metrics 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 1 and in Figure 7 in the appendix.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [No] Justification: We are currently unable to share code but hope to be able to do so in the future. ImageNet data is publicly available, but we're not able to share the used text-to-image data.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
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 of the clean data given noisy data along the sampling trajectory. The moment matching framework provides a 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 quality are highly imperfect, and in future work we plan to run a full set of human evaluations on the 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 a distilled generator and an auxiliary denoising model, and another using two minibatches to allow only updating the generator. In future work we intend to further explore the space of algorithms spanned by these choices, and gain additional insight into the costs and benefits of both approaches.A Instant moment matching = matching expected teacher gradientsIn section 3.2 we propose an instantaneous version of our moment matching loss that does not require alternating optimization of an auxiliary denoising model g φ . This alternative version of our algorithm uses the loss in equation 8, which we reproduce here for easy readibility:EQUATIONwhere ∂g θ (zs) ∂θ is the Jacobian of g θ , and where ∇ φ L(φ) is evaluated on an independent minibatch from x and z s .It 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 L θ (x, z s ) ≡ w(s) x -g θ (z s ) 2 , we can rewrite the term ∇ φ L(φ)| φ=θ = ∇ θ L θ (x, g θ (z s )) because the first term of L(φ) can be seen as a standard diffusion loss at θ for a generated x, and the second term of L(φ) is zero whenφ = θ. Futher observe that ∇ θ L θ (x, z s ) = 2w(s)(g θ (z s ) -x) T ∂g θ (zs) ∂θ which means that ∇ η ∇ θ L θ (x, z s ) = ∇ η 2w(s)x T ∂g θ (zs) ∂θ . Now letting Linstant (η) = L instant (η)+w(s)g θ (z s ) T ∂g θ (zs) ∂θ ∇ θ L θ (x, z s )where the latter term is constant w.r.t. η, we can write instant moment-matching (Equation 15) as momentmatching of the teacher gradients where again stop gradients are again placed on z s : 17) where we assume that the gradient of the teacher training loss is zero when sampling from the training distribution, E x,z s ∼q ∇ θ L θ (x, z s ) = 0, which is true if the teacher attained a minimum of its training loss. Our instant moment matching variant is thus indeed equivalent to matching expected teacher gradients in parameter space.Linstant (η) ≡ 1 2 E zt∼q,x=gη(zt),zs∼q(zs|zt,x) ∇ θ L θ (x, sg(z s )) 2 Λ (16) = 1 2 E zt∼q,x=gη(zt),zs∼q(zs|zt,x) ∇ θ L θ (x, sg(z s ))-E x,z s ∼q ∇ θ L θ (x, z s ) 2 Λ ,(Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: The paper contains a separate section on hyperparameter choices and a 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 in the core of 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, or as supplemental material. 
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.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at 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 or the standard error of the mean. 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, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as 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 it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?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 circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: This work presents a generic algorithm that contributes to improving the quality and cost of generative models in 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 do 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 no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to 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 to 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 as 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: No models or data are released.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 a 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 we are not able to disclose full details on the latter.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 a 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 are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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.  
 
 Appendix: All experiments were run on TPUv5e, using 256 chips per experiment. For ImageNet we used a global batch size of 2048, while for text-to-image we used a global batch size of 512. The base models were trained for 1M steps, requiring between 2 days (Imagenet 64) to 2 weeks (text-to-image). We use the UViT architecture from Hoogeboom et al. (2023) . Configurations largely correspond to those in the appendix of Hoogeboom et al. (2023) , 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 & Ba, 2014) with β 1 = 0, β 2 = 0.99, = 1e -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 a maximum norm of 1. We don't use an EMA, weight decay, or dropout. When distilling a diffusion model using moment matching with alternating parameter updates (Section 3.1) the auxiliary denoising model xφ is fit by minimizing E pη w(z s ) xxφ (z s ) 2 , with z s and x sampled from our distilled model. In the one-step case this is equivalent to performing score matching, as x = g η (z 1 ) with z 1 ∼ N (0, I), and z s ∼ q(z s |x) is sampled using forward diffusion starting from x. Recall here that q(z s |x, z 1 ) = q(z s |x) as z 1 is pure noise, uncorrelated with z s . Upon convergence, we'll have that(α s xφ (z s ) -z s )/σ 2 s = E pη(x|zs) [(α s x -z s )/σ 2 s ] = ∇ zs log p η (z s ) ∀z s , i.e. our auxiliary model will match the score of the marginal sampling distribution of the distilled model p η (z s ) (because the optimal solution is xφ (z s ) = E x∼pη [x]). In the multi-step case this equivalence does not hold, since the sampling distribution of z s depends on z t . Instead the proper multistep score is given by:EQUATIONwithσ 2 t→s = 1 σ 2 s + α 2 t|s σ 2 t|s -1andEQUATION)This expression suggests we could perform score matching by denoising towards µ t→s (x, z t ), a linear combination of x and z t , rather than just towards x. We tried this in early experiments, but did not get good results.However, moment matching can still be seen to match the proper score expression (equation 19) approximately, if we assume that the forward processes match, meaning p η (z t |z s ) ≈ q(z t |z s ). This then gives:EQUATION)When this is used to fit the auxiliary score s φ (z s ) = (α s xφ (z s ) -z s )/σ 2 s , it is equivalent to fitting xφ (z s ) against just x, so under this approximation moment matching and score matching once again become equivalent.If our distilled model has p η (x|z t ) = q(x|z t ), and if z t ∼ q(z t ) (which is true during training), the approximation p η (z t |z s ) ≈ q(z t |z s ) 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 p η (z t |z s ) = q(z t |z s ) until convergence, so during optimization moment matching and score matching indeed optimize different objectives.In Section 5 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 an 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 a more optimal weighting can be found.Sampling from a diffusion models consists of repeatedly predicting the clean data given an intermediate 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 a simple weighted average of these predictions. In Figure E 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 v-parameterization of Salimans & Ho (2022) . Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in 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 to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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?
questions: 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?
1.  Why does the instantaneous version of the algorithm have no reported results for 1 and 2 steps?
limitations: limitations: Although there is no limitations section in the paper, the discussion of the limitations is well addressed throughout the paper.
paper_summary: paper_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
strength_weakness: strength_weakness: 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.

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?

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.

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？
questions: questions: Please see Weakness.
limitations: limitations: yes
paper_summary: paper_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.
strength_weakness: strength_weakness: 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)? 

2. 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.
questions: questions: 1. In tables 1 and 2, by increasing NFE to the level of the base model, what the FIDs and IS will be?

2. 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?

3. How is the generation quality and FID by using instantaneous version of the moment matching (algorithm 3) on ImageNet $128\times128$?

4. Can the proposed method be compatible with Consistency Model for further timestep distillation?

5. Also, can the method support training from scratch?

6. 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: limitations: The authors provided a discussion on potential limitations on evaluations and 1 or 2 steps generation.
paper_summary: paper_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.
strength_weakness: strength_weakness: Minor Typo: There is a typo on line 72: "as described is described."

Lack 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:

- Exclusion of the z_s dependence on the parameters when calculating gradients of the moment matching loss (line 90).
- Training the auxiliary diffusion model using a mix of teacher predictions and generated samples (line 110).

Comparison 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.

Enhancement 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.

On 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"?
questions: questions: Please find a set of suggestions in the weakness section above.
limitations: limitations: Yes.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2. 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.
2. 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.
3. The experiments lack ablating sampling steps. The authors could provide more experiment results adapting other sampling steps in later versions.
questions: 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?
2. Why you select 8 sampling steps? How more or less eight steps will affect experimental results, run speeds, etc?
limitations: limitations: A limitation of this study is that the theoretical part is difficult and too mathematical to understand.

Comments: title: Thanks. We'll include both points in the revised paper.
we'll add this experiment
Thanks for clarification.
comment: Thanks. We'll include both points in the revised paper.
We agree that this experiment (moment matching warm-started from another distilled model) would be interesting and have started work on it. Thanks!
I appreciate the reviewers' clarification and support the acceptance of this paper. In the revised edition, it would be beneficial to discuss the "scalability of this distillation method to large latent diffusion models" and to further investigate "why the student outperforms the teacher".
*"Is your suggestion to run moment matching distillation warm-started from a model distilled using another algorithm such as progressive distillation or consistency distillation?"*

Yes, that’s precisely my idea. This isn’t a mandatory experiment, but it could further demonstrate the advantages of Moment Matching. Thank you for your response and I will raise my score to 6.
I thank the authors for their response. Consequently, I have raised my score.
I'd like to thank the authors for responding to the review.  After reading other reviews and the author's responses, I continue to recommend this paper for acceptance.
I thank the authors for their response. I will keep my original rating.