EM Distillation for One-step Diffusion Models
Authors: Sirui Xie
Zhisheng Xiao
Diederik P Kingma
Tingbo Hou
Ying Nian Wu
Kevin Patrick Murphy
Tim Salimans
Ben Poole
Ruiqi Gao
Venue: NeurIPS 2024 poster (2024)
Keywords: Generative models
Primary area: diffusion_based_models
TL;DR: 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.

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.

EM Distillation for One-step Diffusion Models
Sirui Xie 12,3
Zhisheng Xiao²
Diederik P. Kingma'
Tingbo Hou
Ying Nian Wu
Kevin Murphy'
Tim Salimans'
Ben Poole' Ruiqi Gao
Google DeepMind
Google Research
UCLA
Abstract
While diffusion models can learn complex distributions, sampling requires compu- tationally 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
approach that distills a diffusion model a one-step generator model with minimal loss of perceptual quality. Our approach is derived (EM), where the generator parame- ters updated using samples from the joint distribution of the diffusion teacher prior and inferred generator latents. We develop a reparametrized sampling scheme and noise cancellation technique that together stabilize 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 distilling text-to-image diffusion models.
through the lens of Dnpectation-Maximizallee
Introduction
Diffusion models [1-3] have enabled high-quality generation of images [4-6], videos [7, 8]. and other modalities [9-11]. Diffusion models use forward process to create sequence of distributions that transform the complex data distribution into Gaussian distribution, and learn the score function for each of these intermediate distributions. Sampling from 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 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: (I) trajectory distillation methods [13-18] that accelerate solving the differential equation, and (2) distribution matching approaches [19-23] that learn implicit generators to match the marginals learned by the diffusion model. Trajectory distillation-basee 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 pre-train pre-trained diffusion teacher model and 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 method similar to the Expectation-Me Docerattion-Miiizzatioon (EM) framework [24], which alternates between an Expectation-steg (E-step) that estimates the learning gradients with Monte Carlo samples, and a Maximization-stey (M-step) that updates the student
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

through gradient ascent. As the target distribution is represented by the pre- pre-trained score function, the E-step the original EM that first samples 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 Disf-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 Inff-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-122 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.
2 Preliminary
2.1 Diffusion models and score matching
Diffusion models [1, 2], also known score-based generative models [27, 3), consist of forward process that gradually injects noise to the data distribution and reverse process that progressively denoises the observations recover the original data distribution Pdata (xo). This results of noise levels & E (0, 1] with conditional distributions qi (x, xo) (axxo, I), whose marginals of
sequence
are (xt). We use pariance-preserving forward process [3, 28, 29] such that of 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 Vx, log (x of the marginal distribution of the noisy observations. This score function can be estimated by neural network sΦ (x, 1) through (weighted) denoising score matching [30, 31]:
J Epsau(Xo)) (()))(()))))) (xx/x))
[w[tt || (x,, Vx, log ae (x, [x0]
(I)
where wt) is the weighting function and p(f) is the noise schedule.
2.2 MCMC with Langevin dynamies
While solving the reverse- time SDE results in a sampling process that traverses noise levels, simulat- ing Langevin dynamics [32] results in 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 target distribution p(x) by iterating through
x - + log p(x') + /2nn,
(2)
where y is the stepsize, n ~ N(O, 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-36]. Convergence of Langevin dynamics requires a large number of sampling steps, especially for high- dimensional data. In practice, short-run variants with early termination have been succesfully used for learning of EBMs [25, 37, 38).
2.3 Maximum Likelihood and Expectation-Maxiiizzaioo
cetitatiio-Maximizzzannn (EM) [24] is a maximum likelihood estimation framework to learn latent variable models: (x z) = pe 2 p(z). such that the marginal distribution pe (x) = Po z)dz approximates the target distribution (x). It originates from the generic training objective of maxi- mizing the log- likeliiod function over parameters: C(0) = Eg(x) [log (logge(()) (x)), which is equivalent to minimizing the forward KI. divergence (x (oo [39]. Since the marginal distribution Pe (o(x) is usually analytically intractable, HM involves an E-step that expresses the gradients over the model parameters 0 with an expectation formula
Vo (()
Ve Ex(x)) [log po (x)) E (x())) Vo log z)].
2

---

(a) ImageNet
(b) Text-to-image (SD embedding space)
(c) Text-to to-ageee (SD image space)
Figure I: Before and after MCMC correction. In (a)(b), the left columns are x go (z), the right columns are updated x after 300 steps of MCMC sampling jointly on and z. (a) illustrates the effect of correction in ImageNet. Note that the off-manifold images are corrected. (b) illustrates
correction in the embedding space of Stable Diffusion v1.5, which decoded to image space in (c). Note the disentanglement the cats and sharpness of the sofa. Zoom in for better viewing.
where Po (z|x) (x(x)))p(2) is the posterior distribution of z given x. See Appendix A for a detailed (o(x)
derivation. The expectation 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 instantiation of this EM framework where E-steps leverage Langevin dynamics for drawing samples.
2.4 Variational Score Distillation and Diff-Instruct
Our method is also closely related to Score Distillation Sampling (SDS) [9], Variational Score Distillation (VSD) [26] and Disf-Instruct [22], which have been used for distilling diffusion mod- els into single-step generator [23, 41J. The generator produces clean images X Ge (z) with p(z) = (O, I). and can be diffused to noise level t to form a latent variable model Pa.: z) Pe. p (x (z)) (x,((x,(z)
(V(agoozz), of This model is trained to match the marginal distributions pe,((x)) and qe (xx) by minimizing their reverse KL divergence. Integrating over all noise levels, the objective is to minimize I(0) where
Epcc) (i(() / (x.) log (x,((Xe) dx: 4s (xx)
T(0) = Epco (xo.
(4)
When parametrizing x, (z)g(()) + oFE, the gradient for this objective in Eq. (4) can be written as VoJ(0)- Eet0 V log (x()))
Vx, log lo. ag (z)),
(5)
teacher score
learned (x(
where p(e) (O, I). the teacher score is provided by the pre-trained diffusion model. In SDS, Vx, log (x,) is the known analytic score function of the Gaussian generator. In VSD and Diff- Instruct, an auxiliary score network so (X,, t) is learned to estimate The training alternates between learning the generator network go with the gradient update in Eq. (5) and learning the score network ss with the denoising score matching loss Eq. (1).
3 Method
3.1 EM Distillation
We consider formulating the problem of distilling pre-trained diffusion model deep latent- variable model pe exx,, z) defined in Section 2.4 using the EM framework introduced in Section 2.3, For simplicity,
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) represented by the diffusion model where we can access the score function Vx log (xx). Theoretically speaking. the generator network (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: (0(2) = *x(2, (*), where &, the x-prediction function inherited from the teacher diffusion model, and remains hyper-paramn 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 po (z(zx) 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, 2). We initialize this sampling process using a joint sample from the student: drawing z ~ p(z) and x ~ pe (x ()). This sampled x is no longer drawn from g(x), but z is guaranteed to be valid sample from the posterior pe (z]x). We then run MCMC to correct the sampled pair towards the desired distribution:
po(x, 2) == (x)) = (x 2)
px(x) (see Fig. 1 for visualization of this process). If ((x)
g(x)
and (x) are close each other, pe (x, z) is close to px z). In that case, initializing the joint sampling of (e z) with pairs of (x, z) from pe(x. z) could significantly accelerate both sampling of x and inference Assuming MCMC converges, we can use the resulting samples estimate the learning gradients for EM:
Ve x ago (2)
V (0) (x.,x [V log (x, z)] = E Epo(x.2)
(6)
20
We abbreviate our method as EMD hereafter. To successfully learn the student network with EMD, we need to identify efficient approaches to sample from pe (xxx z).
3.2 Reparametrized sampling and noise cancellation
As an initial strategy, we consider Langevin dynamics which only requires the score functions:
Vx log (x) = Vx log (x)
Vx log + Vx log po (x(z))... (x|z).
teacher score
leamned (x(x)
-agoo a (z00 2.
Vx log po = Vv (o(( + Vz (z)
r²
While we do not have access to the score of the student, Vx log) 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 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 agga (z) + ae defines a deterministic transformation from the pair of (6, z) to the pair of (x, z), which enables to push back the joint distribution po po(x, (x, z) to the 2e,aaes The reparameterized distribution is
q(ago (z) + oE) Po Po + σc) p())p(z).
Po 2)
(8)
The score functions become
V. log p(e, 2) (Vx log ((x) - Vx log (x))) -,
V₂ pog a log Vx (ogg z.
Algorithm 1: EM Distillation
Input: Teacher score functions Vx, log (Xt), generator network go. prior p(z), score network Sg, noise scheduler ptt), weighting functions wi) and (t)), of MCMC steps K, MCMC step size 7.
Output: Generator network go, score network 8tp- while not converged do
Sampling a batch of t, z, e from p(t), p(z), (0, I) to obtain x, Updating $6 via Stochastic Gradient Descent with the batch estimate of Eq. (12)
Sampling x and K with (e, z) -coorreecoo e, z, f, Vx, log go (xxx,, ), go, 86, K, x) Updating Go via Stochastic Gradient Ascent with the batch estimate of Eq. (11)
end while

---

Algorithm 2: (c, 2) correetor
Input: Xo: 6, 2, i, teacher score function Xx, log qt (Xt), generator network go. prior po pz(z), score
network Sd, # of MCMC steps K, MCMC step size 7.
Output: x z²
Sampling Langevin noise n², n², n from (0, I), letting e. = x = for i in [1, K] do
Updating (E',z") with 1-step Langevin update over scores Eq. (9), with e¹ updated using n²
end for Pushing (E" x²²) forward to (xf _zx) and then canceling the noises in x
See Appendix 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 visu- alizing samples produced by MCMC (see Fig. 2a), we found that samples contained sub- stantial noise. While this makes sense given the level of noise in the marginal distributions, found that this inhibited learning of the stu- dent. We identify that, due to the structure of Langevin dynamics, there is noise added to at each step that can be linearly accumulated across iterations. By removing this accumulated noise along with the temporally decayed initial e, we recover cleaner x samples (Fig. 2b). Since x is effectively regression target in Eq. (6), and the expected value of the noises is O, 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).
(a) x w/ accumulated noise (b) x w/o accumulated noise Figure Images after 8-step Langevin updates with and without accumulated noise.
3.3 Maximum Likelihood across noise levels
The derivation above assumes smoothing the data distribution with single noise level. In practice, the diffusion teachers always employ multiple noise levels coordinated by noise schedule p(t).
Therefore, we optimize 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:
VoC(O)= Ve (xx)) (t) log po.: (x,)] (p(t).::(x..))) (t)(t) Ve log pe, ( z)]
(10)
where (x,, z) are a series of latent-variable models as defined in Section 2.4, with a shared generator (z) across noise levels. Empirically, we find (()) = of /at or w(() = of/0 perform well.
Denote the resulted distribution after steps of MCMC sampling with noise cancellation as of (x/ =x),
), the final gradient for the generator network is
VoL(O) = E PCO (x (x -w(() (() Vollx
aggo (z )113
(11)
20
The final gradient for the score network s (x( 1) is
V TΦ) = (ott)..o..(xx..)) (w)) V
t) Vx.
(12)
Similar to VSD [26, 22], we employ alternating update for the generator network go and the score network so (x,, t). See summarization in Algorithm 1.
5

---

3.4 Connection with VSD and Diff-Instruct
In this subsection, we reveal an interesting connection between EMD and Variational Score Dis- tillation (VSD) [26, 22], i.e. although motivated by optimizing different types of divergences, VSD [26, 22] equivalent to EMD with a special sampling scheme.
To see this, consider the 1-step EMD with noise cancellation, stepsize 7 = 1 in x, and no update on z
04 (z) + ore, x (2)0(() + oz v log ((x))
20n
(13)
po,)
Substitute it into Eq. (11), we have
-w(t) (t)) V x (01g0(z))
Ve((O) (p(())(())p(()))
20
(14)
(p(().(())ppca) ((())(x, x, log d: (x.))) Xx, log (x))))))) Vogo (z)o
which is exactly the gradient for VSD (Eq. (5)), up to 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 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 Po(x) is a Gaussian, then EMD update in Eq. 14 would resemble Score Distillation Sampling (SDS) [9].
Related Work
Diffussin 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-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-55J. 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 al. [23] introduces a regression loss 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 work, Zhou al. [59] adopt Fisher divergence as the distillation objective and propose 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.
5 Experiments
We employ EMD to learn one-step image generators on ImageNet 64x64, ImageNet 128x 128 [60] and text-to-image generation. The student generators are initialized with the teacher model weights. Results are compared according Frechet Inception Distance (FID) [61], Inception Score (IS) [62],

---

Gredient norm
HD
HD
VGG-Recall
all noise conceltd
150
all nolse concelled
last ness cancelied
cancallation
125
last cancelled
cencellation
0.505
100
0.503
2.5
0.575
2.4-
0.565
+
(4. corresrr
Hocemettrr
0.500
20000 40000 Training Berocions
60000
20000
40000
24
2
20
Thaining Rerations
Langecth stepe
Langeedn sbeps
(a)
(b)
(c)
(d)
Figure 3: (a) Gradient norms and FIISs for complete noise cancellation, last-step noise cancellation
and no noise cancellation. (c)() FIDs and Recalls of EMD with different numbers of Langevin steps.
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
DifflnstruurtVSDDeqqiivvaaall formulation with only one update in x as presented in Section 3.4.
5.1 ImageNet
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 with Karras et al. [47] as teacher, and ImageNet 128x 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 tech- nique 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 e) does give us seemingly clean images empirically. See Fig. 2 for a comparison. These updated x" can be seen as regression targets in Eq. (11). Intuitively speaking, regressing a generator towards clean images should result 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.
Multi-step joint sampling We scrutinize the effect of multi-step joint update on (e, z). Empirically, find a constant step size of Langevin dynamics across all noise levels in the (c, -sscce works well: y = (T, 7/) = (0.4², 0.0042), which simplifies the process of step size tuning. Fig. shows results of running this (c. cororeccor for 300 steps. We see that the (c. corrrcctor removes visual artifacts and improves structure. Fig. 3cd illustrates the relation between the distilled generator's 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 more expensive than sampling e, requiring back proppappaaatiion through the generator ge. An alternative is to only sample while keeping z fixed, with the hope that if x does not change dramatically with a finite number MCMC updates, the initial remains a good approximation of samples from po (ozzxx x). As shown in Fig. 3cd, sampling e performs similarly the joint sampling of (c, when the number of sampling steps 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
Table 1: EMD-8 ImageNet 64x64, 100k steps of training
the score functions of (c, z) in Eq. (9) to derive something similar. We conduct a comparison between the two parameterizations for joint sampling, (x, corrrectt and (e, cooorreccoo..
FID (L)
IS (1)
(x, (/(e)
2.829
62.31
(x, z)
3.11
61.08
For the (x, c)ooortct,, we set the step size of x as of to align the magnitude of update with the one of the (e, zi) -correcto,,, and keep
(c,z)
2.77
62.98
7

---

Table 2: Class-conditional genreation on
Table 4: FID-30k for text-to-image generation in
ImageNet 64x64.
MSCOCO. Results are evaluated by Yin et al. [23].
Method
NFF (()
FID (L)
Rec. (1)
Family
Method
Latency
FID (L)
Multiple Steps
DALL-E [66]
27.5
DDIM [12)
50
13.7
DALL-E 2 [4]
10.39
EDM-Heun [47]
10
17:25
Parti-B [67]
6,4s
8.10
DPM Solver [44]
10
7,93
Make-A-Scene [68]
25.0s
11.84
PD [13]
8.95
0.65
Unacoelerated
GLIDE [69]
15.0s
12.24
CD [15]
2
4.70
0.64
Imagen [S]
9.1s
7.27
Multistep CD [18]
2.0
cDiff-I
32.0s
6.95
Single Step
StyleGAN-T[71]
0.10s
13.90
dogG [65]
4.06
0.48
GANs
GigaCAN [72]
0.138
9.09
EDM [47]
154.78
PD
0.62
DPM++ (4 step) [45]
0.26s
22.36
BOOT [16]
16.30
0.36
UniPC (4 step) [73]
0.266
19.57
DENO [17]
7,83
Accelerated
LCM-LoRA (1 step) [74]
0.09s
77.90
TRACT [14]
7.43
LCM-LoRA (4 step) [74]
0.19s
23.62
CD-LPIPS [15]
6.20
0.63
InstaFloo--D999 [55]
0.09%
13.10
Diff-Instruct (22)
5.57
UFOGen [20]
0.09s
12.78
DMD [23]
EMD-I (baseline)
2,62
3.1
0.55
DMD (ICFG=3) | [23]
0.09s
11.49
EMD-16 (ours)
2,20
0.59
EMD-1 (boseline, ICFG=3)
0.09
10.96
EMD-1 (baseline, (CFG=2)
0.09s
9.78
Teacher
256
1.43
EMD-8 (ours, (CFG=2)
Teacher
SDv1.5 [6]
2.59s
8.78
Table 3: Class-conditional generation on
ImageNet 128x 128.
Table 5: CLIP Score in high CFG regime.
Method
(F)
FID (
(t)
Family
Method
Latency (1)
CLJP (t)
Multiple Steps
Multistep CD [18]
2.1
164
DPM++ (4 step) [45]*
0.26s
0.309
Multistep CD [18]
2.3
157
UniPC step) [73]
0.26s
0.308
Multistep CD [18]
3.1
147
Accelerated
LCM-LoRA step) [74]
0.09s
0.238
Single Step
LCM-LoRA step) [74]
0.19
0.297
CD |151
7.0
DMD [23]
0.09s
0.320
EMD-I (boseline)
134 2.75
EMD-8 (ours)
0.09s
0.316
EMD-16 (ours)
6.0
140
2.8.3
Teacher
512 1.75
171.1
2.7
Teacher
SDv1.5 [6]
2.59
0,322
the step size of the same (see Appendix B for details). This also promotes numerical stability in (x, z)correect by canceling the of in the denominator of the term Vx log po (x 2)) x- (x) in 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 (e. c-orrectorr and (x, )-corrector without sampling in z, which serves as the baseline for the joint sampling.
Table reports the quantitative comparisons with EMD-8 on ImageNet 64 64 after 100k steps of training. While joint sampling with (c, z)-orrecctor improves over (c, )-corrector, (x, z)-corrector struggles to match the baseline. Possible explanations include that the space of (c. z) is more MCMC friendly, or it requires more effort on searching for the optimal step size of z for the (x, coorrecoo We leave further investigation future work.
Comparsion with existing methods We report the results from our full-fledged method, EMD-16, which utilizes a (e, 2) c)-corrett with 16 steps of Langevin updates, and compare with existing approaches. We train for 300k steps on ImageNet 64 64x64, and 200k steps on ImageNet 128 x 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. We summarize the comparison with existing methods for few-step diffusion generation Table and Table 3 for ImageNet 64x64 and ImageNet 128% x 128 respectively. Note that we also tune the baseline EMD-1, which in formulation equivalent to Diff- 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 x 64, EMD achieves competitive recall among distribution matching approaches but falls behind trajectory distillation approaches which maintain individual trajectory mappings from the teacher.
8

---

(a) ImageNet 64x64 multi-class
(b) ImageNet 128x x 128 multi-class
(c) ImageNet 128x 128 single-clas class
Figure ImageNet samples from the distilled 1step generator. Models are trained conditondit
with all classes. We provide single-class samples in (c) to demonstrate good mode coverage.
Figure 5: Text-to-image samples from the 1-step student model distilled from Stable Diffusion v1.5.
5.2 Text-to-image generation
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-Acsthetics-.. 6..25 dataset [75]. On this task, DMD [23] is strong baseline, which introduced an additional regression loss VSD or Disf-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 t* intermediate noise levels, consistent with the observation from Luo et al. [22]. In Appendix C.4 we discuss the selection of l". The intuition that by choosing the value of 1*, we choose specific denoiser 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 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 compose the classifer-free guided teacher score (we refer 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 -ottgt generation. Using higher tCFG, similar to DMD, produces 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 13) with trajectory- distillation baselines [55, 74] and adversarial distillation basclines [21] can be found in Appendix D.2.

---

Table 6: Training steps per second in ablations for computation overhead in ImageNet 64
Algorithmic Ablation
sec/step
Student score matching only
0.303
Generator update for EMD-1 based on (c, Z) corrector
0.303
Generator update for EMD-2 based on (c, z corrector
0,417
Generator update for EMD-4 based on (c, z corrretor
0.556
Generator update for EMD-8 based on (e, -corrector
0.714
Generator update for EMD-16 based on (c., -cooreccto
1.111
EMD-16 student score matching + generator update based on (c, corrrectrr)
1.515
Baseline Diff-Instruct (student score matching generator update)
0.703
5.3 Computation overhead in training
Despite EMD being more expensive per training iteration compared to the baseline approach Diff- Instruct, 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 moderate even with the most expensive EM setting. In Table we report some quantitative measurement of the computation overhead. Since it is challenging to time each instead logged the sec/step for experiments
python method's wall-clock time in our infrastructure, with various algorithmic ablations on ImageNet 64.64. EMD-16 only doubles the wall-clock time of Instruct when taking all other overheads into account.
Discussion and limitation
We present EMD, a maximum likelihood beeeihood-bbsed method that leverages EM framework with novel sampling and optimization techniques learn one-step student model whose marginal distributions match the marginals of a pretrained diffusion model. EMD demonstrates strong performance in class- conditional generation 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 perform competitively, and is sensitive to the choice of t* (fixed timestep conditioning that repurposes the diffusion denoiser to become one-step genertor) initialization. While training student model entirely from scratch is supported theoretically by our framework, empirically we were unable to achieve competitive results. Improving methods enable generation from randomly initialized generator networks with distinct architectures and lower- dinee-siinnion 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 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.
ImageNet and text-to-image generation. Despite exhibiting compelling
Acknowledgement
We thank Jonathan Heek and Lucas Theis for their valuable discussion and feedback. We also thank Tianwei Yin for helpful sharing of experimental details in his work. Sirui would like to thank Tao Zhu, Jiahui Yu, and Zhishuai Zhang for their support in prior internship Google DeepMind.
References
[1] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International conference machine learning, pages 2256-2265. PMLR, 2015.
[2] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
10

---

[3] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2020.
[4] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06/25, 1(2);, 2022.
[5] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Bureu Karagol Ayan, Tim Salimans, al. Photorealistic -ext-to-image diffusion models with deep language understanding, Advances in neural information processing systems, (1136479-364999 2022.
[6] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10684-10695, 2022.
[7] Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:22 10.02303, 2022.
[8] Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo. Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, Clarence Ng, Ricky Wang, and Aditya Ramesh. Video generation models world simulators. 2024. URL https://openai. com/research/ video- video-generation
generation coneiiiioommoellssse models as aalassswwoll www.ddsiiulaatnnooo simulators.
[9] Ben Poole, Ajay Jain, Jonathan T Barron, and Ben Mildenhall. Dreamfusion: Text-to-3d using 2d diffusion. Xrviv preprint arXiv:2209. 14988, 2022.
[10] Minkai Xu, Lantao Yu, Yang Song, Chence Shi. Stefano Ermon, and Jian Tang. Geodiff: A geo- metric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.0292. 2022.
[11] Mengjiao Yang, KwangHwan Cho, Amil Merchant, Pieter Abbeel, Dale Schuurmans, Igor Mordatch, and Ekin Dogus Cubuk. Scalable diffusion for materials generation. urXiv preprint arXiv:2311.09235. 2023.
[12] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02550,, 2020.
[13] Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00512 2022.
[14] David Berthelot, Arnaud Autef, Jierui Lin, Dian Ang Yap, Shuangfei Zhai Siyuan Hu, Daniel Zheng, Walter Talbott, and Eric Gu. Tract: Denoising diffusion models with transitive closure time-distillation. arXiv preprint arXiv: 0424444888 2023.
International Conference on Machine Learning, pages 32211-32252 PMLR, 2023. [15] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In
[16] Jiatao Gu, Shuangfei Zhai, Yizhe Zhang, Lingjie Liu, and Joshua M Susskind. Boot: Data-free distillation of denoising diffusion models with bootstrapping. In ICML 2023 Workshop on Structured Probabilistic Inference { Generative Modeling. 2023.
[17] Hongkai Zheng, Weili Nie, Arash Vahdat, Kamyar Azizzadenesheli, and Anima Anandkumar. Fast sampling of diffusion models via operator learning. In International Conference on Machine Learning, pages 42390 42402. PMLR, 2023.
[18] Jonathan Heek, Emiel Hoogeboom, and Tim Salimans. Multistep consistency models. arXiv preprint 200i::24033006877 2024.
[19] Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion gans. In International Conference on Learning Representations, 2021.

---

[20] Yanwu Xu, Yang Zhao, Zhisheng Xiao, and Tingbo Hou. Ufogen: You forward once large scale text-to-image generation via diffusion gans. arXiv preprint arXiv:231 2023.
[21] Axel Sauer, Dominik Lorenz, Andreas Blattmann, and Robin Rombach. Adversarial diffusion distillation. arXiv preprint arXiv:2311. 2023.
[22] Weijian Luo, Tianyang Hu, Shifeng Zhang, Jiacheng Sun, Zhenguo Li, and Zhihua Zhang, Diff- instruct: universal approach for transferring knowledge from pre-trained diffusion models. Advances in Neural Information Processing Systems, 36, 2023.
[23] Tianwei Yin, Michaël Gharbi, Richard Zhang, Eli Shechtman, Frédo Durand, William T Freeman, and Taesung Park. One-step diffusion with distribution matching distillation. In CVPR, 2024.
[24] Arthur P Dempster. Nan M Laird, and Donald B Rubin. Maximum likelihood from incomplete data via the em algorithm. Journal of the royal statistical society: series (methodological), 39 (1):1-22, 1977.
[25] Erik Nijkamp, Mitch Hill, Song-Chun Zhu, and Ying Nian Wu. Learning non-convergent non non-persistent short-run meme toward energy-based model. Advances in Neural Information Processing Systems, 32, 2019.
[26] Zhengyi Wang, Cheng Lu, Yikai Wang, Fan Bao, Chongxuan Li, Hang Su, and Jun Zhu. Pro- lifiedreamer: High-fidelity and diverse text-to-3d generation with variational score distillation.
Advances in Neural Information Processing Systems, 36, 2024.
[27] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances neural information processing systems, 32, 2019.
[28] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems, 34:21696-2117 2021.
[29] Diederik Kingma and Ruiqi Gao. Understanding diffusion objectives as the elbo with simple data augmentation. Advances in Neural Information Processing Systems, 36, 2023.
[30] Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(4), 2005.
[31] Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation, 23(7):1661-1674 2011.
[32] Radford M Neal et al. Meme using hamiltonian dynamics. Handbook of markov chain monte carlo, 2(11):2, 2011.
[33] Jianwen Xie, Yang Lu, Song-Chun Zhu, and Yingnian Wu. A theory of generative convnet. In International Conference on Machine Learning, pages 2635-2644. PMLR, 2016.
[34] Jianwen Xie, Yang Lu, Ruiqi Gao, and Ying Nian Wu. Cooperative learning of energy-based model and latent variable model via meme teaching. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018
[35] Ruiqi Gao, Yang Lu, Junpei Zhou, Song-Cho Chun Zhu, and Ying Nian Wu. Learning generative convnets via multi-grid modeling and sampling. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pages 9155-9164, 2018
[36] Erik Nijkamp, Ruiqi Gao, Pavel Sountsov, Srinivas Vasudevan, Bo Pang, Song- Zhu, and Ying Nian Wu. Meme should mix: Learning energy-based model with neural transport latent space meme. International Conference on Learning Representations, 2021.
[37] Erik Nijkamp, Mitch Hill, Tian Han, Song-Chun Zhu, and Ying Nian Wu. On the anatomy of meme-based maximum likelihood learning of energy-based models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 5272-5280, 2020.
[38] Ruiqi Gao, Yang Song, Ben Poole, Ying Nian Wu, and Diederik Kingma. Learning energy- based models by diffusion recovery likelihood. arXiv preprint arXtv:2012.08125, 2020.
12

---

[39] Christopher M Bishop. Pattern recognition and machine learning. Springer google schola, 2: 1122-1128, 2006.
[40] Tian Han, Yang Lu, Song-Chun Chun Zhu, and Ying Nian Wu. Alternating fack-propagation for generator network. Proceedings of the AAAI Conference on Artificial Intelligence, volume 31, 2017.
[41] Thuan Hoang Nguyen and Anh Tran. Swiftbrush: One-step text-to-image diffusion model with variational score distillation. arXiv preprint arXiv:2312.052239, 2023.
[42] Zhisheng Xiao, Karsten Kreis, Jan Kautz, and Arash Vahdat. Vaebm: symbiosis between variational autoencoders and imergy-based models. In International Conference on Learning Representations, 2020.
[43] Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.02388, 2021.
[44] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 35:5775-5787, 2022.
[45] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm- solver++: Fast solver for guided sampling of diffusion probabilistic models. arXiv preprint wxXv:2211.00109 2022.
[46] Fan Bao, Chongxuan Li. Jun Zhu, and Bo Zhang. Analytic-dpm: analytic estimate of the optimal reverse variance diffusion probabilistic models. arXiv preprint arXiv:2201.066 1.6503, 2022.
[47] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. Advances in Neural Information Processing Systems, 35: 26565 26577, 2022.
[48] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. 2015.
arXiv preprint
[49] Chenlin Meng, Robin Rombach, Ruiqi Gao, Diederik Kingma, Stefano Ermon, Jonathan Ho, and Tim Salimans. On distillation of guided diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 14297-14306, 2023.
[50] Yanyu Li, Huan Wang, Qing Jin, Ju Hu, Pavlo Chemerys, Yun Fu, Yanzhi Wang, Sergey Tulyakov, and Jian Ren. Snapfusion: Text-to-image diffusion model mobile devices within two seconds. arXiv preprint
2023.
[51] Simian Luo, Yigin Tan, Longbo Huang, Jian Li, and Hang Zhao. Latent consistency models: Synthesizing high-resolution images with few-step inference. arXiv preprint arXiv:2310.04377, 2023.
[52] Dongjun Kim, Chieh-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning probability flow ode trajectory of diffusion. arXiv preprint arXiv:2310.02299, 2023.
[53] Jianbin Zheng, Minghui Hu, Zhongyi Fan, Chaoyue Wang, Changxing Ding, Dacheng Tao, and Tat-Jen Cham. Trajectory consistency distillation. arXiv preprint arXiv:2402. 19159, 2024.
[54] Hanshu Yan, Xingchao Liu, Jiachun Pan, Jun Hao Liew, Qiang Liu, and Jiashi Feng. Perflow: Piecewise rectified flow as universal plug-and-play accelerator. arXiv preprint arXiv:2405.075100 2024.
[55] Xingchao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, et al. Instaflow: One step is enough for high-quality diffusion-based text-to-image generation. The Twelfih International Conference on Learning Representations, 2023.
13

---

[56] Axel Sauer, Frederic Boesel, Tim Dockhorn, Andreas Blattmann, Patrick Esser, and Robin Rombach. Fast high-resolution image synthesis with latent adversarial diffusion distillation. arXiv preprini arXiv:2403. /2015, 2024.
[57] Shanchuan Lin, Anran Wang, and Xiao Yang. Sdxl-lighting: Progressive adversarial diffusion distillation. arXiv preprint arXiv:2402. 13929, 2024.
[58] Yuxi Ren, Xin Xia, Yanzuo Lu, Jiacheng Zhang, Jie Wu, Pan Xie, Xing Wang, and Xuefeng Xiao. Hyper-sd: Trajectory segmented consistency model for efficient image synthesis. arXiv preprint arXiv:2404. I 13686, 2024.
[59] Mingyuan Zhou, Huangjie Zheng, Zhendong Wang. Mingzhang Yin, and Hai Huang. Score identity distillation: Exponentially fast distillation of pretrained diffusion models for one-step generation. In Forty-first International Conference on Machine Learning, 2024.
[60] Jia Deng, Wei Dong. Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 245..5.. leee, 2009.
[61] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
[62] Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
[63] Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. Advances neural information processing systems, 32, 2019.
[64] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pages 7768-8773. PMLR, 2021.
[65] Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale gan training for high fidelity natural image synthesis. arXiv preprint urXiv:1809) 10066, 2018.
[66] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and llya Sutskever. Zero-shot text-to-image generation. In International conference on machine learming, pages 8821-8831, Pmlr, 2021.
[67] Jiahui Yu, Yuanzhong Xu, Jing Yu Koh, Thang Luong, Gunjan Baid, Zirui Wang, Vijay Vasudevan, Alexander Ku, Yinfei Yang. Burcu Karagol Ayan, et al. Scaling autoregressive models for content-rich text-to-image generation. arXiv preprint arXiv:2206. 10789, 2(3):5, 2022.
[68] Oran Gafni, Adam Polyak. Oron Ashual, Shelly Sheynin, Devi Parikh, and Yaniv Taigman. Make-a-scene: Scene-based text-to-image generation with human priors. European Confer- ence on Computer Vision, pages 89-106. Springer, 2022.
[69] Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen. Glide: Towards photorealistic image generation and editing with texted diffusion models. arXiv preprint arXiv:2112.10741,
[70] Yogesh Balaji, Seungjun Nah, Xun Huang, Arash Vahdat, Jiaming Song, Qinsheng Zhang, Karsten Kreis, Miika Aittala, Timo Aila, Samuli Laine, al. ediff-i: Text-to-mmage diffusion models with ensemble of expert denoisers. arXiv preprint arXiv:2211.0133 2022.
[71] Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. arxiv e-prints. arXiv preprint 1rXiv:1812.000 2018.
14

---

[72] Minguk Kang, Jun- Yan Zhu, Richard Zhang, Jaesik Park, Eli Shechtman, Sylvain Paris, and Taesung Park. Scaling up gans for text-to-image synthesis. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 10124-10134, 2023.
[73] Wenliang Zhao, Lujia Bai, Yongming Rao, Jie Zhou, and Jiwen Lu. Unipe: A unified predictor- corrector framework for fast sampling of diffusion models. Advances Neural Information Processing Systems, 36, 2024.
[74] Simian Luo, Yiqin Tan, Suraj Patil, Daniel Gu, Patrick von Platen, Apolinário Passos, Longbo Huang, Jian Li, and Hang Zhao. Lem-lora: A universal stable-diffusion acceleration module. arXiv preprini arXiv:231 11.331..0555 1.05556, 2023.
[75] Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, et al. Laion- 5b: An open large-scale dataset for training next generation image-text models. Advances in Neural Information Processing Systems, 01225578-252999 2022.
[76] Tsung-Y Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Purt 13, pages 740-755. Springer, 2014.
[77] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-8788 87944, 2021.
[78] Emiel Hoogehoom, Jonathan Heek, and Tim Salimans. simple diffusion: End-to-end diffusion for high resolution images. In International Conference on Machine Learning, pages 13213- 13232. PMLR, 2023.
A Expectation-Maxiiizaaioo
To learn a latent variable model po(x, z) po (x (z), po(x) px( z)d2 from a target distribution (xx), the EM-like transformation on the gradient of the log- likelihood function is:
VoL(O) Vo (x (lxg
Exo() (x()) V Po (x)ll
(15)
(x(x))n Vo (x) Vo lo
(x(x)pa(x1xx) Vo log z)].
Line 3 due to the simple equality that (p Vo Ve (z Vo
0.
B Reparametrized sampling and noise cancellation
Reparametrization. The EM-like algorithm propose requires joint sampling of (x, from Po(x, z). Similar to [36, 42], we utilize a reparameterizatioo in joint MCMC sampling, such as slow convergence and complex step size tuning. Notice that x age (z) + de defines a deterministc mapping from (e, z) to (x, z). Then by change of variable have:
x and z to overcome challenges
polc, z)dedz = po( z)dxdz
9(x)
Po(x)* po (x, z)dxdz
g(ago (z) + σE)
(16)
po(e, 2)dedz
Po(ago(2 (z) + oc)
9(age (z) + oc)
po(E,z)= Po + oe) p(p) (()))(z)
where p(e) and p(z) standard Normal distributions.
15

---

The score functions become
Ve log 2) dVx log ((x)
Vx logpo(())
(17)
V log p(e, V log g(x)
Vx po V
Noise cancellation.
The single-step Langevin update on e then:
g(x')
= ()--)) 71' + yo Vx log
+ v 2/n'.
(18)
(x(x))
Interestingly, we find the particular form of V log p(e, z) results in closed-form accumulation of multi-step updates
k=0 (1- Vx log g(x')
c+1 = (1- +
+
Σ k=0 (1-) /27n" (19)
px(x')
which, after the push-forward, gives
x x =ag(z + 10 7)t 2²,, log ((x²)
(1
De (x
drift
(20)
+ (1 (1- 0c E 7)th
/ 2
noise
As x 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 yo²):
x
=x² + + log ((x)) - Vx logpo(x))) 7(x²= azg(z))) + 22
=7 Σ k=1 (1 7)² 7)* ag((z))) + E (1 -)= (Vx log g(x') Vx log (og')())))
+ (1 7) (z") + (1 (1 06 + Σ k-0 (1 (1- i /2701*
nosses
(21)
C Implementation details
C.1 ImageNet 64
We 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 days) on 64 TPU-v4. We use (c, -coorrecr, 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 A
Table 7: Hyperparameters for EMD on ImageNet 64x64.
Irg
Ir's
batch size
Adam by
Adam b2
X*
i(1)
2 x 10-6
1 x 10-5
128
0.0
0.99
0.4
0.004
16
332188
o7
16

---

C.2 ImageNet 128 128
We train the teacher model following the best setting of VDM++ [29] with the 'U-ViT, L' architec- ture [78]. We use the "cosine-adie -djine-ddjustd' 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 (e, -ccorrector,, in which both the teacher and the student score networks have dropout probability of 0.1, We list other hyperparameters in Table 8.
Table 8: Hyperparameters for EMD on ImageNet 128 128.
Irg
Ir.
batch size
Adam b1
Adam by
Y
7/
K
X*
(
x 10-6
1 x 10-5
1024
0.0
0.99
0.4
0.004
16
-6
di
C.3 Text-to-image generation
We adopt the public checkpoint of Stable Diffusion v1.5 [6] 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 Table 9.
Table 9: Hyperparameters for EMD on Text-to-image generation.
Irg
Irs
batch size
Adam 6
Adam b2
Ye
7/
t"
(()
2 x 10-6
10-5
1024
0.0
0.99
0.3²
0.005
500
of
C.4
Choosing t" and X
The intuition is that by choosing the value of t*, we choose specific denoiser at that noise level. When parametrizing t, the log-signal-to-nois ratio À is more useful when designing noise schedules, strictly monotonically decreasing function [28]. Due to the monotonicity, X is an alternative representation for t* that actually reflects the noise levels more directly.
Fig. 6 shows the denoiser generation at the Oth training iteration for different X in ImageNet 128 128. When O, the generated images are no different from Gaussian noises. When X -6 the generated images have more details than X =- 110. In the context of EMD, these samples help us understand the initialization of MCMC. According to our experiments, setting results in similar performance. For the numbers reported in the manuscript, we used the same X as the bascline Diff-Instruct on ImageNet 64x64 64 and only did very rough grid search on ImageNet 128 x 128 and Text-to-image.
[-6. -3
D Additional qualitative results
D.1
Additional ImageNet results
In this section, we present additional qualitative samples for our one-step generator on ImageNet 64x64 and ImageNet x28x 128 in Fig. 7 to help further evaluate the generation quality and diversity.
D.2 Additional text-to-image results
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 comparison. We use the public checkpoints of LCM' and InstaFlow where
hnccpe://huggingface http:////haagcaocccccccc//ttttoe htwp:://wwe.ccttaaaaoootoatteete "attpss//huggiingfac.. http://awwa.aco/e//////eeeteeeee
adv1-5
6__
17

---

(a) Initial denoiser generation with X
(b) Initial denoiser generation with X -6
(c) Initial denoiser generation with X = -10
Figure 6: Initial denoiser generation with different X
both checkpoints share the same Stable Diffusion 1.5 as teachers. Note that the SD-turbo results are obtaind from the public checkpoint fine-tuned from Stable Diffusion 2.1, which is different from our teacher model.
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 and 15.
18

---

(a) ImageNet 64x64 Multi-class
ImageNet 128 Multi-class
(c) ImageNet 128x 128 Single-class (Left: Husky, right: Siamese)
Figure 7: Additional qualitative results for ImageNet
19

---

Teacher (50 step)
EMD (1 step)
LCM (1 step)
LCM (2 steps)
SD-Turbo (1 step)
InstaFlow (1 step)
Table 10: Prompt: Dog graduation at university.
20

---

3D
YDD
Teacher (50 step)
EMD (1 step)
LCM (1 step)
LCM (2 steps)
SD-Turbo (1 step)
InstaFlow (1 step)
Table 11: Prompt: 3D animation cinematic style young caveman kid, in its natural environment.
21

---

Teacher (50 step)
EMD (1 step)
LCM (1 step)
LCM (2 steps)
SD-Turbo (1 step)
InstaFlow step)
Table 12: Prompt: An underwater photo portrait of beautiful fluff white cat, hair floating. dynamic swimming pose. The sun rays filters through the water. High-angle shot, Shot on Fujifim X.
22

---

Teacher (50 step)
EMD (1 step)
LCM (1 step)
LCM (2 steps)
SD-Turbo (1 step)
InstaFlow step)
Table 13: Prompt: A minimalist Teddy bear in front of a wall of red roses.
23

---

close-up photo of 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 white beard in the enchanted and mugical forest
Table 14: Additional qualitative results of EMD. Zoom-in for better viewing.
24

---

3D render baby parrot, Chibi, adorable big eyes. 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 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. Table 15: Additional qualitative results of EMD. Zoom-in for better viewing.
25

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the
paper's contributions and scope?
Answer: [Yes]
All the claims the abstract and the introduction are supported by the experiments.
Justification:
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 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 is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
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 limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "I "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 the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-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 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 used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
26

---

Answer: [Yes]
Justification: All derivations are included in the main text and the appendix. All assumptions are explicitly communicated.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We include the appendix.
the algorithm and full implementation details in the main text and
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. 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 have some path reproducing or verifying the results.
5. Open access to data and code
27

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: We have not open sourced the model code, but our approach is data-free so no training data is required. We also provide implementation details in the appendix that we hope sufficient for reproducing results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
public/guides/ http///eeuuiiisisoooooooeee
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/nips. cctpublic/guiii wwwc//gwgeee/eeeeeeeeecee SrndeeuumissionPooo
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/ettails
Question:
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary 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 the core the paper to level detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
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 emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
28

---

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.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The computed resources used in this project are listed in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, 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 the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https https://meett http:r////e//e///eeeeee cc cc/public///ecccccccccccccclllcl
Answer: [Yes]
Justification: confirm the paper conform every respect with the NeurIPS code of ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
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 is useful.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
29

---

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), faimess 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 tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate 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 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 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 system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [No]
Justification:
The model will not be released.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere 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 challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All models that open- open-sourced through HuggingFace are provided with corresponding hyeperlinks.
Guidelines:
The answer NA means that the paper does 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 used and, 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.
30

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: No new assets introduced.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing 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 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 the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then 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.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country 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 research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
31

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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.
32


Reviews:
[{"review_id": "kJrGWPRLWf", "reviewer": "Reviewer_eeRQ", "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).", "questions": "1. What if the MCMC steps becomes different? (e.g. 8, 4, 2)\n\n2. Can you expand the method to f-divergence?\n\n3. Will you release the code?", "limitations": "None.", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. Implementing forward KL divergence is important because it can leverage the good statistical property of MLE. \n\n2. The empirical results that compares (EMD 1 vs EMD 16) shows direct benefit of proposed method. Especially, the improvements of recall in table2 is align with the motivation of forward KL divergence.", "weaknesses": "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?\n\n\n2. 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).\n\n\n3. Adding all the metric of (NFE, FID, Pec, Rec, IS) for table2 and table3.\n\n\n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback. We appreciate the acknowledgment of our motivation in using forward KL and the corresponding validation in the ablation between EMD-16 and EMD-1. However, we would like to point out a misunderstanding in the summary “Apart from previous reverse KL divergence, it requires the joint samples z, x from student distribution”: The joint sampling is towards the target distribution of the teacher, instead of the student. Below is our point-to-point response. \n\n**[Approximation errors in student score and MCMC]**\n\nThis is a valid concern and below we discuss the two errors separately.\n\nThe approximation error of student score is unfortunately a universal issue in the family of distribution matching methods, including VSD, Diff-Instruct, DMD and EMD. Due to the generality of this issue, we believe it is worth dedicated research. Concurrent to our work, [1] finds that more optimization steps leads to better approximation. Another concurrent work mentioned by reviewer 3oLm, SiD, investigates a better way to decompose the distribution divergence. We are happy to mention this issue in the limitation section, provide pointers to these works in the revised version, and explore in future work if these techniques are transferable to our method. \n\nEmpirically, we find that this discretization error in MCMC is not problematic, and short-run MCMC that we use in EMD can instead be motivated as interpolating between Moment-Matching Estimate and Maximum Likelihood Estimate [2]. In Fig 3cd, we empirically show that this error can be reduced by running longer MCMC chains. \n\n\n[1] Yin et al. \"Improved Distribution Matching Distillation for Fast Image Synthesis.\" arXiv 2024.\n\n[2] Nijkamp et al. \"Learning non-convergent non-persistent short-run mcmc toward energy-based model.\" NeurIPS 2019.\n\n**[Computation overhead]**\n\nSee Global Response\n\n**[Performance per iteration (e.g. x-axis:iteration / y-axis: FID)]**\n\nWe did include this curve in Fig. 3b. \n\n**[Adding all the metrics for table2 and table3]**\n\nThanks for the suggestion. We will add the following numbers to Table 2 in the revised draft:\n\nEMD-16 Prec. 0.7559, IS 68.31, EMD-1 Prec. 0.7579, IS 62.43. \n\nFor table 3, unfortunately, we don’t find baseline Prec. and Rec. to compare with.  \n\n**[Forward KL divergence often fails to achieve better fidelity compared to reverse KL divergence]**\n\nThis is a good question. Empirically we observe EMD still gives high per sample fidelity while alleviates the mode coverage problem of Diff-instruct that leverages reverse KL. A possible explanation is that the short-run MCMC sampler only produces an approximation of the posterior mean but is still far from fully mixing. Therefore, EMD is better to be understood as an interpolation between mode-seeking reverse and mode-covering forward KL, which maintains high perceptual quality and meanwhile avoids significant mode seeking.\n\n**[What if the MCMC steps become different? (e.g. 8, 4, 2)]**\n\nSee Global Response.\n\n**[Can you expand the method to f-divergence?]**\n\nUpon the reviewer’s request, we reviewed the literature and found there is an alpha-EM framework [3] that generalizes EM to alpha-divergence, another special case of f-divergence. Happy to iterate with the reviewer on other types of f-divergence you deem to be interesting and better. \n\n[3] Matsuyama, Yasuo. \"The/spl alpha/-EM algorithm: surrogate likelihood maximization using/spl alpha/-logarithmic information measures.\" IEEE Transactions on Information Theory 49.3 (2003)\n\n**[Will you release the code?]**\n\nWe plan to release code for ImageNet-64 by the camera-ready deadline.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MT8tcqotE3", "reviewer": "Reviewer_3oLm", "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.", "questions": "Questions: \n\n1. Why modeling the joint distribution of (x, z)? Will this bring any benefits? \n2. 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": "The authors discussed the limitations.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tThe paper is well-written and easy to understand.\n2.\tIt demonstrates the effectiveness of the proposed method on a large scale through text-to-image experiments.\n3.\tAlthough the methodology builds on existing concepts, its application of latent variable models and the EM algorithm provides a novel perspective on the distillation problem.", "weaknesses": "Weaknesses: \n\n1.\tThe 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. \n2.\tWhile 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.\n\nMissing Comparisons:\n\nSiD[1] derived a new distillation method originated from fisher divergence, and reach FID 1.52 on ImageNet-64x64. \n\n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the feedback! We feel very encouraged to receive the recognition that EMD provides a novel perspective on the distillation problem. Below we respond to your questions and stated weaknesses:\n\n**[Overhead for using MCMC in training]**\n\nSee Global Response\n\n**[Comparison with SiD]**\n\nThanks for bringing this insightful paper to our attention. We weren't aware of it at the time of submission and will definitely include this citation in the revised version. We think it is a very interesting future direction to see if the objective decomposition techniques proposed in SiD can be incorpated in EMD. \n\n**[Why modeling the joint distribution of (x, z)? Why running the “slow” MCMC sampling on (x, z)?]**\n\nWhen viewing the 1-step generator model as a latent-variable model, the generative modeling problem naturally becomes “learning a joint distribution of latent z and data x such that the marginal distribution of x matches”. Then it naturally leads us to the EM framework that matches the marginal distribution with mode-covering forward KL. \n\nWe include the ablation of different MCMC steps K in Fig. 3(c)(d) and discuss it in L231-236. The results read that both FID and VGG-Recall show clear improvement monotonically as the number of Langevin steps increases. We also provide a possible explanation that a larger number of Langevin steps encourages better mode coverage, likely because it approximates the mode-covering forward KL better.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "887Eybmk9k", "reviewer": "Reviewer_p5As", "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.\n\nThis 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.\n\nThe 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.", "questions": "All my concerns are detailed in the weakness section.", "limitations": "The authors adequately discussed the limitations.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "S1. The theoretical framework is both novel and robust, showcasing an innovative adaptation of the EM algorithm for diffusion distillation. The paper clearly delineates the connection with previous methods based on reverse KL minimization.\n\nS2. The writing is clear, with intuitive mathematics and excellent presentation.\n\nS3. A wide range of design choices, such as reparameterized sampling and noise cancellation in gradient estimation, are well-founded and enhance performance.", "weaknesses": "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.\n\nW2. 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?\n\nW3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback. It is our honor to know the reviewer liked the framework, exposition and technical inventions in our paper.  One small correction we want to make in the summary is that in the expectation step, Markov Chain Monte Carlo (MCMC) sampling is initialized with samples from the joint distribution of noise and the generated image, as the initial sampling of z and x does not involve MCMC. We respond below to several questions and concerns:\n\n**[Performance]**\n\nWe want to respectively argue that on ImageNet-64, where the metric is more reliable, EMD improves the FID significantly from 3.1 of the 1-step baseline to 2.2. For text-to-image generation, we only use the zero-shot MSCOCO FID as a proxy for the evaluation following works in the literature However, the metric might be less reliable, as the teacher model’s distribution may be different from the data from MSCOCO. \n\nConcurrent to our work, [1] studies some optimization techniques in the framework of DMD to reduce the bias in the student score model. We would like to try these techniques in future work to see if it improves the MCMC correction in EMD. \n\n[1] Yin et al. \"Improved Distribution Matching Distillation for Fast Image Synthesis.\" arXiv 2024.\n\n**[Worse recall compared to the teacher model or trajectory-preserving approaches?]**\n\nThis is a good question, which we discussed a bit in L235-236: “A larger number of Langevin steps encourages better mode coverage, likely because it approximates the mode-covering forward KL better.” The short-run MCMC sampler only produces an approximation of the posterior mean, so it is possible that there are still mode missing. EMD is better to be understood as an interpolation between mode-seeking and mode-covering KL. It achieves high perceptual quality and avoids significant mode seeking. \n\n**[Code release and pseudocode for the MCMC correction]**\n\nWe plan to release code for ImageNet-64 by the camera-ready deadline. The pseudocode for MCMC correction is provided in Algorithm 2. We will be happy to polish it in the revised version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Q0XWEfUxAD", "reviewer": "Reviewer_fdee", "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.", "questions": "See above", "limitations": "None", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The resulting methodology becomes an extension of VSD and Diff-instruct, which is interesting and novel.\n- The EMD method holds the flexibility to trade off training efficiency and final performance by adjusting the MCMC steps. \n- The empirical results of the paper are strong.", "weaknesses": "- The noise cancellation trick is not well justified in theory, although with empirical evidence. \n- 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?\n- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the feedback! We are very glad to see the reviewer appreciated EMD’s flexibility to trade off training efficiency and final performance by adjusting the MCMC steps. Below we respond to questions and concerns:\n\n**[Theoretical justification for noise cancellation]**\n\nWe agree that there is no rigorous proof for a guaranteed variance reduction at this point. But we also want to highlight that even at worst case, noise cancellation won’t introduce biases, for canceling the noises whose mean is 0 won’t affect the mean of the gradients. \n\n**[More hyper-parameters (for the MCMC steps) to tune]**\n\nSee Global Response for ablation of MCMC steps. \n\nAs for the step size, we found $\\gamma_\\epsilon\\in[0.3^2, 0.4^2]$ and  $\\gamma_z\\in[0.003^2, 0.004^2]$ are generally good for the 3 tasks we experimented with. We reported the best configuration in the manuscript. \n\n**[Tuning t\\* for prior works on text-to-image like DMD?]**\n\nWhile the DMD paper didn’t report the numbers for the Diff-Instruct baseline in their setting of text-to-image, we tuned the t\\* for Diff-Instruct on text-to-image and obtained results better than those reported in DMD (FID 10.96 vs 11.49). We don’t have the resources to tune DMD on text-to-image generation during the limited time of rebuttal.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wFDOf3q7hV", "reviewer": "Reviewer_MtDz", "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.", "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]$.\n2. 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?\n3. 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$?\n4. 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? \n5. How sensititive is the training and final performance of the model to the specific choices of $K, \\gamma_e$ and $\\gamma_x$? \n6. 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?\n7. 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": "The paper discusses its limitations.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The qualitative results for 1-step generation with EMD are impressive (Figure 6-14 in the appendix). The quantitative metrics seems comparable or better on datasets like ImageNet-64X64/128X128, and MS-COCO etc.\n2. The proposed method can generate high quality images in one step. Further, the method also shows improved diversity of generated images for the same prompt. This indicates that the method indeed results in improved mode coverage.\n3. I like the idea of using stochastic Langevin updates to get better mode coverage but later getting rid of the added noise to aid smoother training.", "weaknesses": "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. \n     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.\n2. 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. \n3. 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)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback! It is very encouraging to know the reviewer liked the idea of combining Langevin updates with noise reduction, as well as our experiment results. \n\n**[Computational overhead]**\n\nSee Global Response\n\n**[Typos]**\n\nWe sincerely apologize for the typos in the submitted version. Here are some clarifications:\n1. We should have put the definition of $\\epsilon\\sim\\mathcal{N}(0, I)$ and $x_t = \\alpha_t g_\\theta(z) + \\sigma_t \\epsilon$ earlier such that the expectation in Eq (5) is a result of reparametrization. $x = \\alpha g_\\theta(z) + \\epsilon$ in L142 is a typo (as it misses a $\\sigma$), the correct one is in L177. \n2. In Eq (7), missing the parameter $\\alpha_t$ is a typo. We have double-checked our implementation. Thanks for pointing it out.\n\n**[Noise cancellation]**\n\nYes, you are completely right. We stated this in L521-533: “Empirically, we find book-keeping the sampled noises in the MCMC chain and canceling these noises after the loop significantly stabilizes the training of the generator network.” We will move these lines to the main text in the revised version. \n\n**[Hyperparameters of MCMC (K, stepsize)]**\n\nSee Global Response for the ablation of MCMC steps. \n\nAs for the step size, we found $\\gamma_\\epsilon\\in[0.3^2, 0.4^2]$ and  $\\gamma_z\\in[0.003^2, 0.004^2]$ are generally good for the 3 tasks we experimented with. We reported the best configuration in the manuscript. \n\n**[t\\* and  \\lambda\\*]**\n\nWe respectfully argue that t\\* does not suddenly appear in Section 5.2. In Section 3.1, L113-116, we provide an introduction of t\\* under the context of the diffusion denoiser, i.e. the x-prediction function. 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 $\\lambda$ is more useful when designing noise schedules, a strictly monotonically decreasing function $f_\\lambda$  [1]. Due to the monotonicity, $\\lambda^*$ is an alternative representation for t\\* that actually reflects the noise levels more directly. (During rebuttal, we find there is another typo in Table 6 and Table 7, for they reported -logSNR. We will correct them in the revision. )\n\nThe pdf in the global response provides the denoiser generation at the 0th training iteration for different $\\lambda^*$. When $\\lambda^*=0$, the generated images are no different from Gaussian noises. When $\\lambda^*=-6$, the generated images have more details than $\\lambda^*=-10$. In the context of EMD, these samples help us understand the initialization of MCMC. According to our experiments, setting $\\lambda^*\\in[-3, -6]$ results in similar performance. For the numbers reported in the manuscript, we used the same $\\lambda^*$ as the baseline Diff-Instruct on ImageNet-64 and only did a very rough grid search on ImageNet-128 and Text-to-image. \n\n[1] Kingma et al. \"Variational diffusion models.\" NeurIPS 2021.\n\n**[What does eps-correction without z-correction mean?]**\n\n$\\epsilon$-correction without $z$-correction is to fix $g(z)$ and only update $\\epsilon$. It is not a theoretically rigorous algorithm. The reason we include this is that we need a baseline to show the marginal benefit of the update on $z$. The context here is that when the step sizes in $z$ and $\\epsilon$ are not aligned, the performance of the joint update can be worse than the update on $\\epsilon$ only. In Table 1, we showed that the reparametrized sampling eases the burden of co-adjustment of step sizes in the data and latent spaces.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
# Global Response #

We would like to thank all reviewers for your careful and helpful feedback! Specifically, we want to express our appreciation to reviewer eeRQ for recognizing the motivation of using forward KL, to reviewer 3oLm for identifying the novel perspective that EMD offers on the distillation problem, to reviewer p5As for the encouraging comments of our theoretical framework being novel and robust, writing being clear and intuitive, design choices being well founded and effective, to reviewer fdee for the particular interest in EMD’s flexibility to trade off training efficiency and final performance, and to reviewer MtDz for liking the core designs in our technical contribution. 

We would like to clarify and address several common concerns here:

**[Ablation on MCMC steps]**

In the submitted version, we included the ablation of different MCMC steps K in Fig. 3(c)(d) and discussed it in L231-236. Both FID and VGG-Recall show clear improvement monotonically as the number of Langevin steps increases. We also provide a possible explanation for why a larger number of Langevin steps encourages better mode coverage as it approximates the mode-covering forward KL better. 

**[Computation overhead]**

First, 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. Second, the additional computational cost that EMD introduced is moderate even with the most expensive EMD-16 setting. See below for a detailed time analysis. Finally, for text-to-image generation, it takes EMD-8 and EMD-1 3h50min and  2h14min respectively to converge to the lowest FID.

In response to the reviewer eerQ’s, 3oLm’s and MtDz’s request, here 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. 
| Algorithmic Ablation | sec/step|
|:---------------------------------|:------------------:|
| Student score matching only |  0.303 |
| Generator update for EMD-1 (joint sampling of eps and z) | 0.303 |
| Generator update for EMD-2 (joint sampling of eps and z) | 0.417 |
| Generator update for EMD-4 (joint sampling of eps and z) | 0.556 |
| Generator update for EMD-8 (joint sampling of eps and z) | 0.714 |
| Generator update for EMD-16 (joint sampling of eps and z) | 1.111 |
| EMD-16 ( student score + generator update w/ joint sampling) | 1.515 |
| Baseline Diff-Instruct (student score + generator update) | 0.703 |

So EMD-16 only doubles the wall-clock time of Diff-Instruct when taking all other overheads into account.

Meta-review: 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.
Decision: Accept (poster)