EM Distillation for One-step Diffusion Models

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.

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. 2. 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. 3. 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. 1. The additional computational overhead needs additional clarity. For instance, from Algorithm 1 and 2, it seems that this method differentiates through the Langevin update steps. If K=16 for instance, this would mean that the computational graph will be $16\times$ larger compared to the baseline $K=1$ like Diff-Instruct, as it would require backprop through the generator network $K$ times. 2. The paper unfortunately has some significant typos. Some equations in Section 3.2 are missing some multiplicative factors (See more in the questions below). As a result, I’m not sure if the experimental observations made from the experiments in Section 3.2 are correct. I hope the authors can clarify my questions below so that I can adjust my score accordingly. 3. This method introduces additional hyper-parameters but the corresponding ablation studies for hyper-parameter sensitivity are missing. For instance, how does the performance of EDM vary with K, the steps of MCMC, during training? Currently, the paper mostly considers two cases K=1 and K=16. It is also unclear how $t^\star$ is selected and how it is used. It appears suddenly in the main text in Section 5.2. (See more questions below)

Questions

1. The expectation in Eq 5 seems incorrect. $\epsilon$ is undefined. I think the expectation should be w.r.t $p(t), p(z), x_0 \sim g_\theta(z)$, and $p(x_t|x_0)$ ie. $E_{p(t), p(z), x_0 \sim g_\theta(z), x_t \sim p(x_t|x_0) }[\cdot]$. 2. There seems to be a parameter $\alpha_t$ missing in equation 7. $\nabla_z \log p_\theta(x|z) = \dfrac{\alpha_t}{\sigma_t^2}(x_t - \alpha_t g_\theta(z))^\top \nabla_\theta g_\theta(z)$. Also, why is score $\nabla_z \log p_\theta(z)$ simplified to $z$ in this equation? 3. I find line 142 confusing. $\epsilon$ is again undefined here but from Algorithm 1, it seems to be i.i.d. sample from $\mathcal{N}(0, I)$. If so, how can $\alpha g_\theta(z) + \epsilon$ be a deterministic transformation? Is $\epsilon$ always fixed for a given $z$? This seems like the reparametrization trick used in VAEs but that still doesn’t make this transformation deterministic. Also, if $p_\theta(x_t|z) = \mathcal{N}(\alpha_t g_\theta(z), \sigma_t^2I)$, then shouldn’t this transformation be $x_t = \alpha g_\theta(z) + \sigma_t \epsilon$? 4. By noise cancellation, does it mean that we collect all terms $\sqrt{2 \gamma} n^i$ from equation 16 for K steps, and then subtract it after K steps? 5. How sensititive is the training and final performance of the model to the specific choices of $K, \gamma_e$ and $\gamma_x$? 6. What is $\lambda^\star$ and how is its value determined? [Table 6, hyperparameters] How sensitive is the final performance to the choice of this hyperparameter? 7. What does $\epsilon$-correction without $z$-correction mean? Is $g(z)$ fixed for all MCMC steps and only $\epsilon$ updated? (This is used in Table 1).

Rating

6

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

The paper discusses its limitations.

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.

Strengths

- The resulting methodology becomes an extension of VSD and Diff-instruct, which is interesting and novel. - The EMD method holds the flexibility to trade off training efficiency and final performance by adjusting the MCMC steps. - The empirical results of the paper are strong.

Weaknesses

- The noise cancellation trick is not well justified in theory, although with empirical evidence. - It seems that there are more hyper-parameters (for the MCMC steps) to tune. I wonder about the cost of doing so. Do you have any guidance on tuning them? - Is the trick of tuning t* used by previous works on text-to-image diffusion distillation? What is the result of doing so for prior works like DMD?

Questions

See above

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

None

Summary

This paper introduces a novel distillation method for converting a diffusion process into a one-step generator. The theoretical foundation is closely tied to the Expectation Maximization (EM) algorithm. The authors aim to minimize the forward Kullback-Leibler (KL) divergence between the target generator distribution and the one-step sampler's output distribution. This minimization is approached using an EM-like algorithm. In the expectation step, Markov Chain Monte Carlo (MCMC) sampling is used to sample from the joint distribution of noise and the generated image, followed by an MCMC correction to adjust towards the distribution of noise and the target image. The estimated gradients from this process are then used to update the generator. Simultaneously, an auxiliary diffusion model is trained to approximate the score of the output distribution. The authors draw an interesting connection between their proposed EM distillation algorithm and previous score distillation-based methods. The final model is evaluated on both class-conditional and text-conditional image generation tasks.

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. S2. The writing is clear, with intuitive mathematics and excellent presentation. S3. 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. W2. While the forward KL divergence is theoretically mode-covering, it still has significantly worse recall compared to the teacher model or trajectory-preserving approaches like the consistency model. Could the authors comment on why this is the case? W3. It seems that the code will not be released. While this is not a critique, providing additional details, such as pseudocode for the MCMC correction, would be helpful for reimplementation.

Questions

All my concerns are detailed in the weakness section.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors adequately discussed the limitations.

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.

Strengths

1. The paper is well-written and easy to understand. 2. It demonstrates the effectiveness of the proposed method on a large scale through text-to-image experiments. 3. Although 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: 1. The requirement for at least K steps of MCMC sampling significantly increases training costs, in contrast to other methods such as SDS, VSD, SiD and consistency distialltion (CD), which generally require only one-step sampling for distillation. 2. While the performance with a large number of MCMC steps is robust, it still does not achieve state-of-the-art results when compared to SiD, particularly with a smaller number of steps. Missing Comparisons: SiD[1] derived a new distillation method originated from fisher divergence, and reach FID 1.52 on ImageNet-64x64. [1] Zhou, Mingyuan, et al. "Score identity distillation: Exponentially fast distillation of pretrained diffusion models for one-step generation." Forty-first International Conference on Machine Learning. 2024.

Questions

Questions: 1. Why modeling the joint distribution of (x, z)? Will this bring any benefits? 2. Without MCMC sampling on $z$, the loss function will be reformed to VSD loss. VSD doesn't require more than one-step MCMC. Does this mean that involving $z$ in the sampling makes the algorithm slower? Then what is the meaning of it?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors discussed the limitations.

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

Strengths

1. Implementing forward KL divergence is important because it can leverage the good statistical property of MLE. 2. 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? 2. Training cost seems extremely expensive. This requires 1) separate approximation of student score, 2) 16 MCMC steps per iterations. Expensive costs itself is bad thing, but the paper does not analysis on the costs rigorously. You must add the portion of training costs (approximated student score training, MCMC, student loss computation, student model back-propagation) in one iteration. You must add performance per iteration (e.g. x-axis:iteration / y-axis: FID). 3. Adding all the metric of (NFE, FID, Pec, Rec, IS) for table2 and table3. 4. Forward KL divergence often fails to achieve better fidelity compared to reverse KL divergence. Did you see any similar situations? If not, you'd better discuss the reason.

Questions

1. What if the MCMC steps becomes different? (e.g. 8, 4, 2) 2. Can you expand the method to f-divergence? 3. Will you release the code?

Rating

3

Confidence

5

Soundness

2

Presentation

2

Contribution

2

Limitations

None.

Thanks for the rebuttal

Thanks for the rebuttal. I recommend the authors clarify these things in the revision: 1. Regarding noise cancellation, though the mean is not affected, the variance can also significantly affect the optimization. So, more analyses are required. 2. What is the performance of the baselines combined with some used tricks in the paper?