EM Distillation for One-step Diffusion Models

Authors: sirui  xie
zhisheng  xiao
diederik p kingma
tingbo  hou
ying nian wu
kevin  murphy
tim  salimans
ben  poole
ruiqi  gao
google  deepmind
google  research
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 06:52:47

Abstract: While diffusion models  can learn complex distributions, sampling requires a computationally expensive iterative process.  Existing distillation methods enable efficient sampling, but have notable limitations, such as performance degradation with very few sampling steps, reliance on training data access, or mode-seeking optimization that may fail to capture the full distribution. We propose EM Distillation (EMD), a maximum likelihood-based approach that distills a diffusion model to a one-step generator model with minimal loss of perceptual quality. Our approach is derived through the lens of Expectation-Maximization (EM), where the generator parameters are updated using samples from the joint distribution of the diffusion teacher prior and inferred generator latents. We develop a reparametrized sampling scheme and a noise cancellation technique that together stabilizes the distillation process. We further reveal an interesting connection of our method with existing methods that minimize mode-seeking KL. EMD outperforms existing one-step generative methods in terms of FID scores on ImageNet-64 and ImageNet-128, and compares favorably with prior work on distilling text-to-image diffusion models.

Main paper: 
 introduction: 
 Diffusion models [1] [2] [3] have enabled high-quality generation of images [4] [5] [6] , videos [7, 8] , and other modalities [9] [10] [11] . Diffusion models use a forward process to create a sequence of distributions that transform the complex data distribution into a Gaussian distribution, and learn the score function for each of these intermediate distributions. Sampling from a diffusion model reverses this forward process to create data from random noise by solving an SDE, or an equivalent probability flow ODE [12] . Typically, solving this differential equation requires a significant number of evaluations of the score function, resulting in a high computational cost. Reducing this cost to single function evaluation would enable applications in real-time generation.To enable efficient sampling from diffusion models, two distinct approaches have emerged: (1) trajectory distillation methods [13] [14] [15] [16] [17] [18] that accelerate solving the differential equation, and (2) distribution matching approaches [19] [20] [21] [22] [23] that learn implicit generators to match the marginals learned by the diffusion model. Trajectory distillation-based approaches have greatly reduced the number of steps required to produce samples, but continue to face challenges in the 1-step generation regime. Distribution matching approaches can enable the use of arbitrary generators and produce more compelling results in the 1-step regime, but often fail to capture the full distribution due to the mode-seeking nature of the divergences they minimize.In this paper, we propose EM Distillation (EMD), a diffusion distillation method that minimizes an approximation of the mode-covering divergence between a pre-trained diffusion teacher model and a latent-variable student model. The student enables efficient generation by mapping from noise to data in just one step. To achieve Maximum Likelihood Estimation (MLE) of the marginal teacher distribution for the student, we propose a method similar to the Expectation-Maximization (EM) framework [24] , which alternates between an Expectation-step (E-step) that estimates the learning gradients with Monte Carlo samples, and a Maximization-step (M-step) that updates the student through gradient ascent. As the target distribution is represented by the pre-trained score function, the E-step in the original EM that first samples a datapoint and then infers its implied latent variable would be expensive. We introduce an alternative MCMC sampling scheme that jointly updates the data and latent pairs initialized from student samples, and develop a reparameterized approach that simplifies hyperparameter tuning and improves performance for short-run MCMC [25] . For the optimization in the M-step given these joint samples, we discover a tractable linear noise term in the learning gradient, whose removal significantly reduces variances. Additionally, we identify a connection to Variational Score Distillation [9, 26] and Diff-Instruct [22] , and show how the strength of the MCMC sampling scheme can interpolate between mode-seeking and mode-covering divergences. Empirically, we first demonstrate that a special case of EMD, which is equivalent to the Diff-Instruct [22] baseline, can be readily scaled and improved to achieve strong performance. We further show that the general formulation of EMD that leverages multi-step MCMC can achieve even more competitive results. For ImageNet-64 and ImageNet-128 conditional generation, EMD outperforms existing one-step generation approaches with FID scores of 2.20 and 6.0. EMD also performs favorably on one-step text-to-image generation by distilling from Stable Diffusion models. 
preliminary 2.1 diffusion models and score matching: 
 Diffusion models [1, 2] , also known as score-based generative models [27, 3] , consist of a forward process that gradually injects noise to the data distribution and a reverse process that progressively denoises the observations to recover the original data distribution p data (x 0 ). This results in a sequence of noise levels t ∈ (0, 1] with conditional distributions q t (x t |x 0 ) = N (α t x 0 , σ 2 t I), whose marginals are q t (x t ). We use a variance-preserving forward process [3, 28, 29] such that σ 2 t = 1 -α 2 t . Song et al. [3] showed that the reverse process can be simulated with a reverse-time Stochastic Differential Equation (SDE) that depends only on the time-dependent score function ∇ xt log p t (x t ) of the marginal distribution of the noisy observations. This score function can be estimated by a neural network s ϕ (x t , t) through (weighted) denoising score matching [30, 31] : J (ϕ) = E p data (x0),p(t),qt(xt|x0) w(t)∥s ϕ (x t , t) -∇ xt log q t (x t |x 0 )∥ 2 2 ,where w(t) is the weighting function and p(t) is the noise schedule.While solving the reverse-time SDE results in a sampling process that traverses noise levels, simulating Langevin dynamics [32] results in a sampler that converges to and remains at the data manifold of a target distribution. As a particularly useful Markov Chain Monte Carlo (MCMC) sampling method for continuous random variables, Langevin dynamics generate samples from a target distribution ρ(x) by iterating throughEQUATIONwhere γ is the stepsize, n ∼ N (0, I), and i indexes the sampling timestep. Langevin dynamics has been widely adopted for sampling from diffusion models [27, 3] and energy-based models [33] [34] [35] [36] . Convergence of Langevin dynamics requires a large number of sampling steps, especially for highdimensional data. In practice, short-run variants with early termination have been succesfully used for learning of EBMs [25, 37, 38] .Expectation-Maximization (EM) [24] is a maximum likelihood estimation framework to learn latent variable models: p θ (x, z) = p θ (x|z)p(z), such that the marginal distribution p θ (x) = p θ (x, z)dz approximates the target distribution q(x). It originates from the generic training objective of maximizing the log-likelihood function over parameters: L(θ) = E q(x) [log p θ (x)], which is equivalent to minimizing the forward KL divergence D KL (q(x)||p θ (x)) [39] . Since the marginal distribution p θ (x) is usually analytically intractable, EM involves an E-step that expresses the gradients over the model parameters θ with an expectation formula where p θ (z|x) = p θ (x|z)p(z)EQUATIONp θ (x)is the posterior distribution of z given x. See Appendix A for a detailed derivation. The expectation can be approximated by Monte Carlo samples drawn from the posterior using e.g. MCMC sampling techniques. The estimated gradients are then used in an M-step to optimize the parameters. Han et al. [40] learned generator networks with an instantiation of this EM framework where E-steps leverage Langevin dynamics for drawing samples.Our method is also closely related to Score Distillation Sampling (SDS) [9] , Variational Score Distillation (VSD) [26] and Diff-Instruct [22] , which have been used for distilling diffusion models into a single-step generator [23, 41] . The generator produces clean images x 0 = g θ (z) with p(z) = N (0, I), and can be diffused to noise level t to form a latent variable modelp θ,t (x t , z) = p θ,t (x t |z)p(z), p θ,t (x t |z) = N (α t g θ (z), σ 2 t I). This model is trained to match the marginal distributions p θ,t (x t ) and q t (x t ) by minimizing their reverse KL divergence. Integrating over all noise levels, the objective is to minimize J (θ) whereEQUATIONWhen parametrizing x t = α t g θ (z) + σ t ϵ, the gradient for this objective in Eq. ( 4) can be written asEQUATIONwhere p(ϵ) = N (0, I), the teacher score is provided by the pre-trained diffusion model. In SDS, ∇ xt log p θ,t (x t ) is the known analytic score function of the Gaussian generator. In VSD and Diff-Instruct, an auxiliary score network s ϕ (x t , t) is learned to estimate it. The training alternates between learning the generator network g θ with the gradient update in Eq. ( 5) and learning the score network s ϕ with the denoising score matching loss in Eq. (1). 
em distillation: 
 We consider formulating the problem of distilling a pre-trained diffusion model to a deep latentvariable model p θ,t (x t , z) defined in Section 2.4 using the EM framework introduced in Section 2.3.For simplicity, we begin with discussing the framework at a single noise level and drop the subscript t. We will revisit the integration over all noise levels in Section 3.3. Assume the target distribution q(x) is represented by the diffusion model where we can access the score function ∇ x log q(x).Theoretically speaking, the generator network g θ (z) can employ any architecture including ones where the dimensionality of the latents differs from the data dimensionality. In this work, we reuse the diffusion denoiser parameterization as in other work on one-step distillation: g θ (z) = xθ (z, t * ), where xθ is the x-prediction function inherited from the teacher diffusion model, and t * remains a hyper-parameter.A naive implementation of the E-step involves two steps: (1) draw samples from the target diffusion model q(x) and (2) sample the latent variable z from p θ (z|x) with e.g. MCMC techniques. Both steps can be highly non-trivial and computationally expensive, so here we present an alternative approach to sampling the same target distribution that avoids directly sampling from the pretrained diffusion model, by instead running MCMC from the joint distribution of (x, z). We initialize this sampling process using a joint sample from the student: drawing z ∼ p(z) and x ∼ p θ (x|z). This sampled x is no longer drawn from q(x), but z is guaranteed to be a valid sample from the posterior p θ (z|x). We then run MCMC to correct the sampled pair towards the desired distribution: ρ θ (x, z) := q(x)p θ (z|x) = p θ (x, z) q(x) p θ (x) (see Fig. 1 for a visualization of this process). If q(x) and p θ (x) are close to each other, ρ θ (x, z) is close to p θ (x, z). In that case, initializing the joint sampling of ρ θ (x, z) with pairs of (x, z) from p θ (x, z) could significantly accelerate both sampling of x and inference of z. Assuming MCMC converges, we can use the resulting samples to estimate the learning gradients for EM:EQUATION)We abbreviate our method as EMD hereafter. To successfully learn the student network with EMD, we need to identify efficient approaches to sample from ρ θ (x, z).As an initial strategy, we consider Langevin dynamics which only requires the score functions:EQUATIONWhile we do not have access to the score of the student, ∇ x log p θ (x), we can approximate it with a learned score network s ϕ estimated with denoising score matching as in VSD [26] and Diff-Instruct [22] . As will be covered in Section 3.3, this score network is estimated at all noise levels. The Langevin dynamics defined in Eq. ( 7) can therefore be simulated at any noise level.Running Langevin MCMC is expensive and requires careful tuning, and we found this challenging in the context of diffusion model distillation where different noise levels have different optimal step sizes. We leverage a reparametrization of x and z to accelerate the joint MCMC sampling and simplify step size tuning, similar to Nijkamp et al. [36] , Xiao et al. [42] . Specifically, the parametrization x = αg θ (z) + σϵ defines a deterministic transformation from the pair of (ϵ, z) to the pair of (x, z), which enables us to push back the joint distribution ρ θ (x, z) to the (ϵ, z)-space. The reparameterized distribution isEQUATIONThe score functions becomeEQUATIONAlgorithm 1: EM Distillation Input: Teacher score functions ∇ xt log q t (x t ), generator network g θ , prior p(z), score network s ϕ , noise scheduler p(t), weighting functions w(t) and w(t), # of MCMC steps K, MCMC step size γ.Output: Generator network g θ , score network s ϕ .while not converged do Sampling a batch of t, z, ϵ from p(t), p(z), N (0, I) to obtain x t Updating s ϕ via Stochastic Gradient Descent with the batch estimate of Eq. ( 12) Sampling x K t and z K with (ϵ, z)-corrector(x 0 , ϵ, z, t, ∇ xt log q t (x t ), g θ , s ϕ , K, γ) Updating g θ via Stochastic Gradient Ascent with the batch estimate of Eq. ( 11) end while Algorithm 2: (ϵ, z)-corrector Input: x 0 , ϵ, z, t, teacher score function ∇ xt log q t (x t ), generator network g θ , prior p 0 (z), score network s ϕ , # of MCMC steps K, MCMC step size γ.Output:x K t , z K . Sampling Langevin noise n 1 , n 2 , ..., n K from N (0, I), letting ϵ 0 = ϵ, z 0 = z for i in [1, K] doUpdating (ϵ i , z i ) with 1-step Langevin update over scores Eq. ( 9), with ϵ i updated using n i end for Pushing (ϵ K , z K ) forward to (x K t , z K ) and then canceling the noises in x K t See Appendix B for a detailed derivation. We found that this parameterization admits the same step sizes across noise levels and results in better performance empirically (Table 1 ). Still, learning the student with these samples continued to present challenges. When visualizing samples x produced by MCMC (see Fig. 2a ), we found that samples contained substantial noise. While this makes sense given the level of noise in the marginal distributions, we found that this inhibited learning of the student. We identify that, due to the structure of Langevin dynamics, there is noise added to x at each step that can be linearly accumulated across iterations. By removing this accumulated noise along with the temporally decayed initial ϵ, we recover cleaner x samples (Fig. 2b ). Sincex is effectively a regression target in Eq. ( 6), and the expected value of the noises is 0, canceling these noises reduces variance of the gradient without introducing bias. Empirically, we find bookkeeping the sampled noises in the MCMC chain and canceling these noises after the loop significantly stabilize the training of the generator network. This noise cancellation was critical to the success of EMD, and is detailed in Appendix B and ablated in experiments (Fig. 3ab ).The derivation above assumes smoothing the data distribution with a single noise level. In practice, the diffusion teachers always employ multiple noise levels t, coordinated by a noise schedule p(t). Therefore, we optimize a weighted loss over all noise levels of the diffusion model, to encourage that the marginals of the student network match the marginals of the diffusion process at all noise levels:EQUATIONwhere p θ,t (x t , z) are a series of latent-variable models as defined in Section 2.4, with a shared generator g θ (z) across all noise levels. Empirically, we findw(t) = σ 2 t /α t or w(t) = σ 2 t /α 2 t perform well.Denote the resulted distribution after K steps of MCMC sampling with noise cancellation as ρ K t (x K t , z K ), the final gradient for the generator network g θ isEQUATIONThe final gradient for the score network s ϕ (x t , t) isEQUATION)Similar to VSD [26, 22] , we employ alternating update for the generator network g θ and the score network s ϕ (x t , t). See summarization in Algorithm 1.In this subsection, we reveal an interesting connection between EMD and Variational Score Distillation (VSD) [26, 22] , i.e., although motivated by optimizing different types of divergences, VSD [26, 22] is equivalent to EMD with a special sampling scheme.To see this, consider the 1-step EMD with noise cancellation, stepsize γ = 1 in x, and no update on zEQUATIONSubstitute it into Eq. ( 11), we haveEQUATIONwhich is exactly the gradient for VSD (Eq. ( 5)), up to a sign difference. This insight demonstrates that, EMD framework can flexibly interpolate between mode-seeking and mode-covering divergences, by leveraging different sampling schemes from 1-step sampling in only x (a likely biased sampler) to many-step joint sampling in (x, z) (closer to a mixing sampler). Notably, for image generation, some believe that forward KL divergence may fail to achieve better fidelity compared to reverse KL divergence. The interpolation enabled by EMD can thus be very useful in practice.If we further assume the marginal p θ (x) is a Gaussian, then EMD update in Eq. 14 would resemble Score Distillation Sampling (SDS) [9] .Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?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 have not open sourced the model or code, but our approach is data-free so no training data is required. We also provide implementation details in the appendix that we hope are sufficient for reproducing our results.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. 
related work: 
 Diffusion acceleration. Diffusion models have the notable issue of slowness in inference, which motivates many research efforts to accelerate the sampling process. One line of work focuses on developing numerical solvers [43, 12, [44] [45] [46] [47] for the PF-ODE. Another line of work leverages the concept of knowledge distillation [48] to condense the sampling trajectory of PF-ODE into fewer steps [13, 49, 15, 14, 50, 51, 18, [52] [53] [54] [55] . However, both approaches have significant limitations and have difficulty in substantially reducing the sampling steps to the single-step regime without significant loss in perceptual quality.Single-step diffusion models. Recently, several methods for one-step diffusion sampling have been proposed, sharing the same goal as our approach. Some methods fine-tune the pre-trained diffusion model into a single-step generator via adversarial training [20, 21, 56] , where the adversarial loss enhances the sharpness of the diffusion model's single-step output. Adversarial training can also be combined with trajectory distillation techniques to improve performance in few or single-step regimes [52, 57, 58] . Score distillation techniques [9, 26] have been adopted to match the distribution of the one-step generator's output with that of the teacher diffusion model, enabling single-step generation [22, 41] . Additionally, Yin et al. [23] introduces a regression loss to further enhance performance. These methods achieve more impressive 1-step generation, some of which enjoy additional merits of being data-free or flexible in the selection of generator architecture. However, they often minimizes over mode-seeking divergences that can fail to capture the full distribution and therefore causes mode collapse issues. We discuss the connection between our method and this line of work in Section 3.4. Concurrent with our work, Zhou et al. [59] adopt Fisher divergence as the distillation objective and propose a novel decomposition that alleviates the dependency on the approximation accuracy of the auxiliary score network. Although the adopted Fisher divergence is similar to reverse KL in terms of the reparametrization and hence the risk of mode collapse, Zhou et al. [59] demonstrate impressive performance gain. 
experiments: 
 We employ EMD to learn one-step image generators on ImageNet 64×64, ImageNet 128×128 [60] and text-to-image generation. The student generators are initialized with the teacher model weights.Results are compared according to Frechet Inception Distance (FID) [61] , Inception Score (IS) [62] , Recall (Rec.) [63] and CLIP Score [64] . Throughout this section, we will refer to the proposed EMD with K steps of Langevin updates on (x, z) as EMD-K, and we use EMD-1 to describe the DiffInstruct/VSD-equivalent formulation with only one update in x as presented in Section 3.4.We start from showcasing the effect of the key components of EMD, namely noise cancellation, multi-step joint sampling, and reparametrized sampling. We then summarize results on ImageNet 64×64 with Karras et al. [47] as teacher, and ImageNet 128×128 with Kingma and Gao [29] as teacher.Noise cancellation During our development, we observed the vital importance of canceling the noise after the Langevin update. Even though theoretically speaking our noise cancellation technique does not guarantee reducing the variance of the gradients for learning, we find removing the accumulated noise term from the samples (including the initial diffusion noise ϵ) does give us seemingly clean images empirically. See Fig. 2 for a comparison. These updated x K can be seen as regression targets in Eq. (11) . Intuitively speaking, regressing a generator towards clean images should result in more stable training than towards noisy images. Reflected in the training process, canceling the noise significantly decreases the variance in the gradient (Fig. 3a ) and boosts the speed of convergence (Fig. 3b ). We also compare with another setting where only the noise in the last step gets canceled, which is only marginally helpful.We scrutinize the effect of multi-step joint update on (ϵ, z). Empirically, we find a constant step size of Langevin dynamics across all noise levels in the (ϵ, z)-space works well: γ = (γ ϵ , γ z ) = (0.4 2 , 0.004 2 ), which simplifies the process of step size tuning. Fig. 1 shows results of running this (ϵ, z)-corrector for 300 steps. We can see that the (ϵ, z)-corrector removes visual artifacts and improves structure. Fig. 3cd illustrates the relation between the distilled generator's performance and the number of Langevin steps per distillation iteration, measured by FID and Recall respectively. Both metrics show clear improvement monotonically as the number of Langevin steps increases. Recall is designed for measuring mode coverage [63] , and has been widely adopted in the GAN literature. A larger number of Langevin steps encourages better mode coverage, likely because it approximates the mode-covering forward KL better. Sampling z is more expensive than sampling ϵ, requiring back-propagation through the generator g θ . An alternative is to only sample ϵ while keeping z fixed, with the hope that if x does not change dramatically with a finite number of MCMC updates, the initial z remains a good approximation of samples from ρ θ (z|x). As shown in Fig. 3cd , sampling ϵ performs similarly to the joint sampling of (ϵ, z) when the number of sampling steps is small, but starts to fall behind with more sampling steps. Reparametrized sampling As shown in Appendix B, the noise cancellation technique does not depend on the reparametrization. One can start from either the score functions of (x, z) in Eq. ( 7) or the score functions of (ϵ, z) in Eq. ( 9) to derive something similar. We conduct a comparison between the two parameterizations for joint sampling, (x, z)-corrector and (ϵ, z)-corrector.For the (x, z)-corrector, we set the step size of x as σ 2 t γ ϵ to align the magnitude of update with the one of the (ϵ, z)-corrector, and keep the step size of z the same (see Appendix B for details). This also promotes numerical stability in (x, z)-corrector by canceling the σ 2 t in the denominator of the term ∇ x log p θ (x|z) = -x-αg θ (z)σ 2in the score function (Eq. ( 7)). A similar design choice was proposed in Song and Ermon [27] . Also note that adjusting the step sizes in this way results in an equivalence between (ϵ, )-corrector and (x, )-corrector without sampling in z, which serves as the baseline for the joint sampling.Table 1 reports the quantitative comparisons with EMD-8 on ImageNet 64×64 after 100k steps of training. While joint sampling with (ϵ, z)-corrector improves over (ϵ, )-corrector, (x, z)-corrector struggles to even match the baseline. Possible explanations include that the space of (ϵ, z) is more MCMC friendly, or it requires more effort on searching for the optimal step size of z for the (x, z)-corrector. We leave further investigation to future work.We report the results from our full-fledged method, EMD-16, which utilizes a (ϵ, z)-corrector with 16 steps of Langevin updates, and compare with existing approaches. We train for 300k steps on ImageNet 64×64, and 200k steps on ImageNet 128×128. Other hyperparameters can be found in Appendix C. Samples from the distilled generator can be found in Fig. 4 . We also include additional samples in Appendix D.1. We summarize the comparison with existing methods for few-step diffusion generation in Table 2 and Table 3 for ImageNet 64×64 and ImageNet 128×128 respectively. Note that we also tune the baseline EMD-1, which in formulation is equivalent to Diff-Instruct [22] , to perform better than their reported numbers. The improvement mainly comes from a fine-grained tuning of learning rates and enabling dropout for both the teacher and student score functions. Our final models outperform existing approaches for one-step distillation of diffusion models in terms of FID scores on both tasks. On ImageNet 64 × 64, EMD achieves a competitive recall among distribution matching approaches but falls behind trajectory distillation approaches which maintain individual trajectory mappings from the teacher. We further test the potential of EMD on text-to-image models at scale by distilling the Stable Diffusion v1.5 [6] model. Note that the training is image-free and we only use text prompts from the LAION-Aesthetics-6.25+ dataset [75] . On this task, DMD [23] is a strong baseline, which introduced an additional regression loss to VSD or Diff-Instruct to avoid mode collapse. However, we find the baseline without regression loss, or equivalently EMD-1, can be improved by simply tuning the hyperparameter t * . Empirically, we find it is better to set t * to intermediate noise levels, consistent with the observation from Luo et al. [22] . In Appendix C.4 we discuss the selection of t * . The intuition is that by choosing the value of t * , we choose a specific denoiser at that noise level for initialization. Other hyperparameters can be found in Appendix C.3.We evaluate the distilled one-step generator for text-to-image generation with zero-shot generalization on MSCOCO [76] and report the FID-30k in Table 4 and CLIP Score in Table 5 . Yin et al. [23] uses the guidance scale of 3.0 to compose the classifer-free guided teacher score (we refer to this guidance scale of teacher as tCFG) in the learning gradient of DMD, for it achieves the best FID for DDIM sampler. However, we find EMD achieves a lower FID at the tCFG of 2.0. Our method, EMD-8, trained on 256 TPU-v5e for 5 hours (5000 steps), achieves the FID=9.66 for one-step text-to-image generation. Using a higher tCFG, similar to DMD, produces a model with competitive CLIP Score. In Fig. 5 , we include some samples for qualitative evaluation. Additional qualitative results (Tables 14 and 15 ), as well as side-by-side comparisons (Tables 10 to 13 ) with trajectory-based distillation baselines [55, 74] and adversarial distillation baselines [21] can be found in Appendix D.2. Despite EMD being more expensive per training iteration compared to the baseline approach Diff-Instruct, we find the performance gain of EMD cannot be realized by simply running Diff-Instruct for the same amount of time or even longer than EMD. In fact, the additional computational cost that EMD introduced is moderate even with the most expensive EMD-16 setting. In Table 6 we report some quantitative measurement of the computation overhead. Since it is challenging to time each python method's wall-clock time in our infrastructure, we instead logged the sec/step for experiments with various algorithmic ablations on ImageNet 64×64. EMD-16 only doubles the wall-clock time of Diff-Instruct when taking all other overheads into account. 
discussion and limitation: 
 We present EMD, a maximum likelihood-based method that leverages EM framework with novel sampling and optimization techniques to learn a one-step student model whose marginal distributions match the marginals of a pretrained diffusion model. EMD demonstrates strong performance in classconditional generation on ImageNet and text-to-image generation. Despite exhibiting compelling results, EMD has a few limitations that call for future work. Empirically, we find that EMD still requires the student model to be initialized from the teacher model to perform competitively, and is sensitive to the choice of t * (fixed timestep conditioning that repurposes the diffusion denoiser to become a one-step genertor) at initialization. While training a student model entirely from scratch is supported theoretically by our framework, empirically we were unable to achieve competitive results. Improving methods to enable generation from randomly initialized generator networks with distinct architectures and lower-dimensional latent variables is an exciting direction of future work. Although being efficient in inference, EMD introduces additional computational cost in training by running multiple sampling steps per iteration, and the step size of MCMC sampling can require careful tuning. There remains a fundamental trade-off between training cost and model performance. Analysis and further improving on the Pareto frontier of this trade-off would be interesting for future work.To learn a latent variable model p θ (x, z) = p θ (x|z)p θ (z), p θ (x) = p θ (x, z)dz from a target distribution q(x), the EM-like transformation on the gradient of the log-likelihood function is:EQUATION)Line 3 is due to the simple equality that E p θ (z|x) [∇ θ log p θ (z|x)] = 0.Reparametrization. The EM-like algorithm we propose requires joint sampling of (x, z) from ρ θ (x, z). Similar to [36, 42] , we utilize a reparameterization of x and z to overcome challenges in joint MCMC sampling, such as slow convergence and complex step size tuning. Notice that x = αg θ (z) + σϵ defines a deterministc mapping from (ϵ, z) to (x, z). Then by change of variable we have:EQUATIONwhere p(ϵ) and p(z) are standard Normal distributions.The score functions becomeEQUATION)Noise cancellation. The single-step Langevin update on ϵ is then:EQUATION)Interestingly, we find the particular form of ∇ ϵ log ρ(ϵ, z) results in a closed-form accumulation of multi-step updatesEQUATION)which, after the push-forward, gives usEQUATIONAs x i+1 is effectively a regression target in Eq. ( 6), and the expected value of the noise is 0, we can remove it without biasing the gradient. Empirically, we find book-keeping the sampled noises in the MCMC chain and canceling these noises after the loop significantly stabilize the training of the generator network.The same applies to the (x, z) sampling (with step size γσ 2 ):EQUATIONWe train the teacher model using the best setting of EDM [47] with the ADM UNet architecture [77] .We inherit the noise schedule and the score matching weighting function from the teacher. We run the distillation training for 300k steps (roughly 8 days) on 64 TPU-v4. We use (ϵ, z)-corrector, in which both the teacher and the student score networks have a dropout probability of 0.1. We list other hyperparameters in Table 7 . Instead of listing t * , we list the corresponding log signal-to-noise ratio λ * . σ 2 t α 2 t C.2 ImageNet 128×128We train the teacher model following the best setting of VDM++ [29] with the 'U-ViT, L' architecture [78] . We use the 'cosine-adjusted' noise schedule [78] and 'EDM monotonic' weighting for student score matching. We run the distillation training for 200k steps (roughly 10 days) on 128 TPU-v5p. We use (ϵ, z)-corrector, in which both the teacher and the student score networks have a dropout probability of 0.1. We list other hyperparameters in Table 8 . We adopt the public checkpoint of Stable Diffusion v1.5 [6] as the teacher. We inherit the noise schedule from the teacher model. The student score matching uses the same weighting function as the teacher. We list other hyperparameters in Table 9 . The intuition is that by choosing the value of t * , we choose a specific denoiser at that noise level. When parametrizing t, the log-signal-to-noise ratio λ is more useful when designing noise schedules, a strictly monotonically decreasing function [28] . Due to the monotonicity, λ * is an alternative representation for t * that actually reflects the noise levels more directly.Fig. 6 shows the denoiser generation at the 0th training iteration for different λ * in ImageNet 128×128. When λ * = 0, the generated images are no different from Gaussian noises. When λ * = -6, the generated images have more details than λ * = -10. In the context of EMD, these samples help us understand the initialization of MCMC. According to our experiments, setting λ * ∈ [-6, -3] results in similar performance. For the numbers reported in the manuscript, we used the same λ * as the baseline Diff-Instruct on ImageNet 64×64 and only did a very rough grid search on ImageNet 128×128 and Text-to-image.In this section, we present additional qualitative samples for our one-step generator on ImageNet 64×64 and ImageNet 128×128 in Fig. 7 to help further evaluate the generation quality and diversity.In this section, we present additional qualitative samples from our one-step generator distilled from Stable Diffusion 1.5. In Table 10 , 11, 12, and 13, we visually compare the sample quality of our method with open-source competing methods for few-or single-step generation. We also include the teacher model in our comparison. We use the public checkpoints of LCM1 and InstaFlow2 , where From the comparison, we observe that our model significantly outperforms distillation-based methods including LCM and InstaFlow, and it demonstrates better diversity and quality than GAN-based SD-turbo. The visual quality is on-par with 50-step generation from the teacher model.We show additional samples from our model on a more diverse set of prompts in Table 14 and 15 . A close-up photo of a intricate beautiful natural landscape of mountains and waterfalls.A hyperrealistic photo of a fox astronaut; perfect face, artstation.Large plate of delicious fried chicken, with a side of dipping sauce, realistic advertising photo, 4k.A DSLR photo of a golden retriever in heavy snow.Masterpiece color pencil drawing of a horse, bright vivid color.Oil painting of a wise old man with a white beard in the enchanted and magical 3D render baby parrot, Chibi, adorable big eyes. In a garden with butterflies, greenery, lush whimsical and soft, magical, octane render, fairy dust.Dreamy puppy surrounded by floating bubbles.A painting of an adorable rabbit sitting on a colorful splash.Macro photo of a miniature toy sloth drinking a soda, shot on a light pastel cyclorama.A traditional tea house in a tranquil garden with blooming cherry blossom trees.Three cats having dinner at a table at new years eve, cinematic shot, 8k. Justification: In the last section, we discuss the limitation, including the assumptions and their related constraints of the proposed method, the efficiency of the method, and the most sensitive hyperparameters. 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.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: All implementation details are enclosed in the appendix.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: [Yes]Justification: We include the standard deviation for the Inception Score, following the convention in the literature.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: We currently discuss mostly the positive impacts of this work, but can add more acknowledgement of the potential for harm caused by text-to-image generative models if the reviewers believe it is useful. Guidelines:• The answer NA means that there is no societal impact of the work performed.• 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. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
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 introduced.Guidelines:• The answer NA means that the paper does not release new • 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 was 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. 
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: The paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
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: the paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on 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 NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.https://huggingface.co/latent-consistency/lcm-lora-sdv1-5https://huggingface.co/XCLiu/instaflow_0_9B_from_sd_1_5https://huggingface.co/stabilityai/sd-turbo 
 
 Appendix:  


Decision: This paper proposes a new approach for one step diffusion model sampling. The idea is to derive an EM style algorithm, which involves running a MCMC sampling process of the joint distribution of noise and data. This algorithm is shown to achieve competitive results in various settings. Most reviewers find the algorithm novel and interesting, and results strong, and the authors did a great job in the rebuttal period answering a lot of the questions. The remaining issues are mostly around MCMC, both its efficiency and whether it's the best way of approximating the forward KL. The AC leans towards accepting this work, primarily based on the considerations that this is a novel algorithm that works pretty well. It will be certainly nice to improve its training efficiency over other methods but that should not outweigh its contributions.

Reviews: review: paper_summary: paper_summary: This paper distill the diffusion model into single step generator through forward KL (mode-coverage) divergence. Apart from previous reverse KL divergence, it requires the joint samples z, x from student distribution. To do this, this paper utilize MCMC method to achieve the samples (z, x).
strength_weakness: strength_weakness: 1. There are two erroneous parts. First one is approximation of student score function through another neural network. Second one is discretization error of MCMC sampling. Can you justify the effects of this two errors in training both theoretical way and empirical way?


2. Training cost seems extremely expensive. This requires 1) separate approximation of student score, 2) 16 MCMC steps per iterations. Expensive costs itself is bad thing, but the paper does not analysis on the costs rigorously. You must add the portion of training costs (approximated student score training, MCMC, student loss computation, student model back-propagation) in one iteration. You must add performance per iteration (e.g. x-axis:iteration / y-axis: FID).


3. Adding all the metric of (NFE, FID, Pec, Rec, IS) for table2 and table3.


4. Forward KL divergence often fails to achieve better fidelity compared to reverse KL divergence. Did you see any similar situations? If not, you'd better discuss the reason.
questions: questions: 1. What if the MCMC steps becomes different? (e.g. 8, 4, 2)

2. Can you expand the method to f-divergence?

3. Will you release the code?
limitations: limitations: None.
paper_summary: paper_summary: This paper introduces the EM Distillation (EMD) method, which efficiently distills diffusion models into a one-step generator model. It utilizes a maximum likelihood approach grounded in Expectation-Maximization (EM) and maintains good image generation quality. The method incorporates a reparametrized sampling scheme and a noise cancellation technique, enhancing the stability of the distillation process. EMD demonstrates good FID scores relative to existing one-step generative models on ImageNet-64 /128 and exhibits capabilities in distilling text-to-image diffusion models.
strength_weakness: strength_weakness: Weaknesses: 

1.	The requirement for at least K steps of MCMC sampling significantly increases training costs, in contrast to other methods such as SDS, VSD, SiD and consistency distialltion (CD), which generally require only one-step sampling for distillation. 
2.	While the performance with a large number of MCMC steps is robust, it still does not achieve state-of-the-art results when compared to SiD, particularly with a smaller number of steps.

Missing Comparisons:

SiD[1] derived a new distillation method originated from fisher divergence, and reach FID 1.52 on ImageNet-64x64. 

[1] Zhou, Mingyuan, et al. "Score identity distillation: Exponentially fast distillation of pretrained diffusion models for one-step generation." Forty-first International Conference on Machine Learning. 2024.
questions: questions: Questions: 

1. Why modeling the joint distribution of (x, z)? Will this bring any benefits? 
2. Without MCMC sampling on $z$, the loss function will be reformed to VSD loss. VSD doesn't require more than one-step MCMC. Does this mean that involving $z$ in the sampling makes the algorithm slower? Then what is the meaning of it?
limitations: limitations: The authors discussed the limitations.
paper_summary: paper_summary: This paper introduces a novel distillation method for converting a diffusion process into a one-step generator. The theoretical foundation is closely tied to the Expectation Maximization (EM) algorithm. The authors aim to minimize the forward Kullback-Leibler (KL) divergence between the target generator distribution and the one-step sampler's output distribution.

This minimization is approached using an EM-like algorithm. In the expectation step, Markov Chain Monte Carlo (MCMC) sampling is used to sample from the joint distribution of noise and the generated image, followed by an MCMC correction to adjust towards the distribution of noise and the target image. The estimated gradients from this process are then used to update the generator. Simultaneously, an auxiliary diffusion model is trained to approximate the score of the output distribution.

The authors draw an interesting connection between their proposed EM distillation algorithm and previous score distillation-based methods. The final model is evaluated on both class-conditional and text-conditional image generation tasks.
strength_weakness: strength_weakness: W1. The performance improvement over baseline approaches (such as 1-step EMD used in VSD, DiffInstruct, DMD, Swiftbrush) is minimal, particularly for text-to-image synthesis. There appears to be a significant gap between the distilled models and the original diffusion teacher.

W2. While the forward KL divergence is theoretically mode-covering, it still has significantly worse recall compared to the teacher model or trajectory-preserving approaches like the consistency model. Could the authors comment on why this is the case?

W3. It seems that the code will not be released. While this is not a critique, providing additional details, such as pseudocode for the MCMC correction, would be helpful for reimplementation.
questions: questions: All my concerns are detailed in the weakness section.
limitations: limitations: The authors adequately discussed the limitations.
paper_summary: paper_summary: The paper proposes EM Distillation (EMD), a maximum likelihood-based approach that distills a diffusion model to a one-step generator model with minimal loss of perceptual quality. Notably, in EMD, the generator parameters are updated using samples from the joint distribution of the diffusion teacher prior and inferred generator latents by MCMC sampling. The method forms an extension of VSD and Diff-instruct. The empirical results are good.
strength_weakness: strength_weakness: - The noise cancellation trick is not well justified in theory, although with empirical evidence. 
- It seems that there are more hyper-parameters (for the MCMC steps) to tune. I wonder about the cost of doing so. Do you have any guidance on tuning them?
- Is the trick of tuning t* used by previous works on text-to-image diffusion distillation? What is the result of doing so for prior works like DMD?
questions: questions: See above
limitations: limitations: None
paper_summary: paper_summary: This work proposes EM Distillation which uses the idea of expectation maximization (EM) to distill a pretrained diffusion model. Naive adaptation of EM algorithm can be computationally expensive as it requires sampling from the teacher diffusion model (which can be slow). This work proposes an alternate approach to avoid sampling from the pretrained diffusion model. The idea is to run MCMC on the joint distribution $p(x,z)$ to generate samples. To further simply Langevin dynamics, the paper uses reparametrization trick and then performs Langevin updates on Gaussian noise $\epsilon$ as well as the latents z (which come from the prior $p(z)$). As each step of Langevin dynamics adds noise, this adversely affects training due to large variance. The paper suggests getting rid of this additional noise to stabilize training (called noise cancellation in paper). The proposed method generalizes previously proposed Variational Score Distillation-based methods such as Diff-Instruct and seems to perform well on conditional image generation tasks.
strength_weakness: strength_weakness: 1. This method needs more calls to teacher model compared to the baseline methods like Diff-Instruct and DMD. The corresponding overhead associated in training, both in terms of additional training time and compute, should be discussed in the paper. 
     1. The additional computational overhead needs additional clarity. For instance, from Algorithm 1 and 2, it seems that this method differentiates through the Langevin update steps. If K=16 for instance, this would mean that the computational graph  will be $16\times$ larger compared to the baseline $K=1$ like Diff-Instruct, as it would require backprop through the generator network $K$ times.
2. The paper unfortunately has some significant typos. Some equations in Section 3.2 are missing some multiplicative factors (See more in the questions below). As a result, I’m not sure if the experimental observations made from the experiments in Section 3.2 are correct. I hope the authors can clarify my questions below so that I can adjust my score accordingly. 
3. This method introduces additional hyper-parameters but the corresponding ablation studies for hyper-parameter sensitivity are missing. For instance, how does the performance of EDM vary with K, the steps of MCMC, during training? Currently, the paper mostly considers two cases K=1 and K=16. It is also unclear how $t^\star$ is selected and how it is used.  It appears suddenly in the main text in Section 5.2. (See more questions below)
questions: questions: 1. The expectation in Eq 5 seems incorrect. $\epsilon$ is undefined. I think the expectation should be w.r.t $p(t), p(z), x_0 \sim g_\theta(z)$, and  $p(x_t|x_0)$ ie. $E_{p(t), p(z), x_0 \sim g_\theta(z),  x_t \sim p(x_t|x_0) }[\cdot]$.
2. There seems to be a parameter $\alpha_t$ missing in equation 7. $\nabla_z \log p_\theta(x|z) = \dfrac{\alpha_t}{\sigma_t^2}(x_t - \alpha_t g_\theta(z))^\top \nabla_\theta g_\theta(z)$. Also, why is score $\nabla_z \log p_\theta(z)$ simplified to $z$ in this equation?
3. I find line 142 confusing. $\epsilon$ is again undefined here but from Algorithm 1, it seems to be i.i.d. sample from $\mathcal{N}(0, I)$. If so, how can $\alpha g_\theta(z) + \epsilon$ be a deterministic transformation? Is $\epsilon$ always fixed for a given $z$? This seems like the reparametrization trick used in VAEs but that still doesn’t make this transformation deterministic. Also, if $p_\theta(x_t|z) = \mathcal{N}(\alpha_t g_\theta(z), \sigma_t^2I)$, then shouldn’t this transformation be $x_t = \alpha g_\theta(z) + \sigma_t \epsilon$?
4. By noise cancellation, does it mean that we collect all terms $\sqrt{2 \gamma} n^i$ from equation 16 for K steps, and then subtract it after K steps? 
5. How sensititive is the training and final performance of the model to the specific choices of $K, \gamma_e$ and $\gamma_x$? 
6. What is $\lambda^\star$ and how is its value determined? [Table 6, hyperparameters] How sensitive is the final performance to the choice of this hyperparameter?
7. What does $\epsilon$-correction without $z$-correction mean? Is $g(z)$ fixed for all MCMC steps and only $\epsilon$ updated? (This is used in Table 1).
limitations: limitations: The paper discusses its limitations.

Comments: title: Thanks for the rebuttal
comment: Thank you very much for your very supportive comments! We will include a pytorch-style pseudo code in the revised version.
Thanks for raising the score! 

For the MCMC, we hope the demonstration of 300 steps of update in Fig. 1 helps make it more convincing. EMD with fewer steps of updates can be viewed as approximately amortizing these long sampling chains. 

We will include the result of the consistency trajectory model (CTM) in the revised version. We also hope the reviewer would like to notice EMD-16's slightly better VGG-Recall (0.59 vs 0.57) albeit the FID gap.
I personally don't think using MCMC with a few steps for approximating such high dimensional joint distribution is a good choice. Maybe removing the accumulated noise, which makes the denoising process more like an ODE, helps here. Comparison to consistecy trajetory model is also missing, which reaches 1.98 FID on ImageNet (vs 2.2 FID in this paper) with one-step generation. The performance wise is not stunning, and the computational cost is more than doubled. The cost issue is going to be further worse in large scale models, which will limit the scaling up of the method. Given the efforts, I increased but kept as a borderline score.
Thank you very much!
Thanks for answering my questions. I am satisfied with the response and increasing my score.
Thank the reviewer for this comment. We will include the following clarification in the revision:

1. While the noise cancellation technique improved performance on three distinct tasks (ImageNet-64, ImageNet-128, Text-to-image), the theoretical analysis of its variance reduction effect indeed needs more investigation. 
2. As stated in L210-211 and L266-267, the EMD-1 baseline is the Diff-Instruct baseline with all the "tricks" (same $t^*$, same learning rate, etc.). We will reiterate it in the revision.
Thanks for the rebuttal. I recommend the authors clarify these things in the revision:

1. Regarding noise cancellation, though the mean is not affected, the variance can also significantly affect the optimization. So, more analyses are required. 

2. What is the performance of the baselines combined with some used tricks in the paper?
Thanks the reviewer for the comments. Can we ask for a clarification on how to implement forward KL with a discriminator? It would be very helpful if the reviewer would like to provide some pointers to existing works.
Thank you for the response. I read all the reviews and rebuttals. In my opinion, the performance gain of EMD is not enough enduring the twice amount of computational cost. I think MCMC is not good way to implement forward KL in distillation scenario. For example, if you adopt auxiliary density ratio estimator (a.k.a. discriminator) between student and teacher, you can implement forward KL without MCMC. I want to keep my score.
I thank the authors for their response and will maintain my original accept rating. Additionally, I believe the doubling of runtime is not a significant issue, provided that it results in notable improvements in quality. Further enhancements in text-to-image quality and recall would be interesting future directions.

For the pseudo code, I was referring to some pytorch-style pseudo code that can be directly transferred e.g. the one in Moco Algorithm 1 [1]. But of course, it is even better to have the full imagenet code release.

[1] He, Kaiming, et al. "Momentum contrast for unsupervised visual representation learning." CVPR. 2020.