Improved Distribution Matching Distillation for Fast Image Synthesis

Summary

The authors propose an improved way of distilling image-generating diffusion models into fast models capable of generating high-quality images with as few as 1-4 steps. Compared to prior work using distribution-matching (DMD), they do away with the regression loss term that tied the teacher path to the student path. They also introduce a GAN-style loss in the pipeline, and introduce a few other tricks to squeeze out extra performance. They demonstrate SOTA performance among efficient models, and even beat the teacher model (made possible by training using real images and a GAN loss).

Strengths

Overall, I found this to be an excellent paper. The work is well-motivated, addressing an important problem, and will likely be interesting to a large audience. The writing is excellent, with all concepts being well explained. The experimental coverage is excellent, with evaluation on two datasets and with a user study, leaving nothing more to be desired. Their performance numbers are also convincing. Finally, their attention to detail on experimental parameters also looks very thorough, giving confidence that their results can be reproduced. On a more detailed level, I personally read the original DMD paper not too long ago, and found their regression loss to be slightly unsatisfactory, since it ties the student generation paths to the teacher paths in a way that seems contrary to the idea of the distribution matching loss. Therefore, I was happy to see in this paper that one can do away with this term.

Weaknesses

There's not so much to say here, since I found most aspects of the paper to be excellent. However, I would have liked to see some more details in how their approach relates to competing approaches. Most notably, the line of work by Sauer et al [23, 24] use SDS and a GAN-style loss. Now that a GAN-style loss is introduced in the DMD framework, the gap between these two approaches gets smaller, and it would be nice to see some more explanation about what the core difference is. The progress is mostly empirical in nature. For example, the authors note that the training becomes more stable using the two-scale update rule, but they don't present any theoretical convergence guarantees (which is perfectly fine, the empirical progress is definitely good enough in my opinion). There are a few minor typos (e.g. line 122, should be "gradient of the data log likelihood"), but few enough not to impair the overall understanding (and I trust the authors to do a final proof reading for the camera ready version).

Questions

What is the most significant difference between this work and the line of work by Sauer et al ([23, 24])? Could one sentence about that be added to the "related work" section? The original DMD paper [22] demonstrated examples of mode collapse when omitting the regression loss. GANs (at least some) are also known to be prone to mode collapse. Could you mention anything about the mode collapse situation in this method? (Since the FID numbers are good, I assume that this is also good, but since this was a main point of analysis in [22]?)

Rating

8

Confidence

4

Soundness

4

Presentation

3

Contribution

4

Limitations

Yes, limitation and potential social impact are well-described in section 6 and A.

Summary

This work proposed an improved training method for distribution matching distillation, named DMD2. Notably, compared to DMD, it does need the regression loss which relies on constructing the synthetic noise-data pairs. Instead, DMD2 introduces three new features: 1) a two time-scale update rule for fake score and generator training, 2) a GAN loss which extracts the bottleneck features of fake score network for a training prediction head as the discriminator, and 3) backward simulation for few-step distillation that produces inference-time generator inputs during training.

Strengths

- The paper is well written and easy to read. - The new distribution matching distillation technique (DMD2) significantly improves over DMD with several well-justified innovations, such as TTUR, GAN loss and backward simulation. - DMD2 achieves SOTA performance on ImageNet-64 and COCO 2014 with one-step generation, and can achieve SOTA performance in distilling SDXL to a 4-step generator, measured by sufficient automatic metrics and human studies. - Ablation studies have been well executed to highlight the importance of each introduced feature.

Weaknesses

- The proposed method introduces many hyperparameters and seems to be sensitive to these hyperparameters, such as batch size, guidance scale for teacher score, GAN loss weighting and fake score updating frequency. For instance, in different distillation tasks (EDM on ImageNet, SDv1.5 and SDXL), these hyperparameters are different (as shown in Appendix G). I’m not sure if the hyperparameters need to be specifically tuned for good performance in each distillation task. - The authors claimed that “SDXL remains challenging to distill into a one-step generator because of limited model capacity and a complex optimization landscape to the direct mapping from noise to highly diverse and detailed images”. For the point of “limited model capacity”, does it mean that if the student has the same network with SDXL, by any means, we are not able to achieve a one-step generation that matches SDXL’s performance? For the point of “a complex optimization landscape”, it seems that both SDv1.5 and SDXL are trained on LAION dataset, which means they are both trying to learn the same mapping from noise to data. Does it mean the higher-quality generation of SDXL (rather the training data of teacher models) hinders the one-step distillation? On the other hand, I wonder if it is possible to tune the hyperparameters of distilling SDXL for a better one-step generator. For example, in Appendix G, the batch size for distilling SDXL is only 128 while the batch size for distilling SDv1.5 is only 2048. If DMD2 is sensitive to batch size in the large-scale text-to-image case, can we increase the batch size for distilling SDXL to 2048 during training for improved performance? - There are some inconsistencies: 1) In Figure 4, the caption says the teacher uses 50 sampling steps while the main text says “while requiring 25x fewer forward passes (4 vs 100)”. 2) EDM (Teacher ODE) originally reported their FID as 2.22, while this work reports 2.32 in Table 1. 3) Both $\mu_{\text{real}}$ and $\mu_{\text{fake}}$ are introduced without definition.

Questions

- I’m curious about the extra cost of backward simulation, i.e., producing synthetic images with the current student generator running several steps. Is it possible to compare the training time per iteration with and without backward simulation? - From the numbers in Table 2, it looks like the 4-step distillation improves Patch FID but gets worse FID and CLIP, compared to 1-step distillation. Any justification?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors adequately addressed the limitations and potential negative societal impact of their work.

Summary

This paper introduces DMD2, a few-step distilled generator to achieve fast sampling while maintaining the decent generation quality of the multi-step diffusion models. DMD2 proposes several new improvements to the training procedure of the original DMD, including (1) replacing the regression loss with the Two-Time scale Update Rule (TTUR) to stabilize the training process, (2) incorporating the standard GAN loss to achieve better quality, (3) and utilizing the backward simulation to alleviate the potential mismatch of training and inference. Built upon these modifications, DMD2 achieves excellent results on few-step image generation.

Strengths

1. The source-intensive process of generating noise-image pairs is replaced by a simple TTUR strategy. 2. DMD2 achieves SOTA results on one-step image generation on ImageNet64 and shows its effectiveness on distilling SDXL into few steps.

Weaknesses

1. It would be better to include a detailed training algorithm to clearly showcase the modifications over the original DMD training process. 2. In practice, the real data used to train the (teacher) diffusion models may not be accessible to the users who hope to distill a small (student) generation model (due to privacy, storage, …). In that case, one limitation of DMD2 is the calculation of GAN loss may become inapplicable. Could you share your opinions about this matter?

Questions

1. In the original DMD paper, the regression loss is capable of mitigating the issue of mode collapse. Would DMD2 also suffer from this issue as the regression loss is removed and an extra GAN loss is introduced? 2. Why the intermediate outputs of DMD2 shown in the right subfigure of Figure 3 are so similar and seem to follow a certain trajectory? As far as I know, distribution matching-based methods do not guarantee the specific paths of the teacher diffusion model and the student generation model are aligned (as mentioned in Lines 35-39). Could authors provide more explanations about this phenomenon (Fig 3)? Note that the samples generated by few-step consistency models may also switch between different paths and generate very different images.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Please refer to the weakness and question sections above.

Summary

This work addresses identifies reasons for training instability of one of competitive diffusion distillation approaches based on distribution matching, using bi-level optimization and also adopt a GAN based feature space feedback for improved quality. Overall demonstrate very good performance on SDXL, SD checkpoints demonstrating effectiveness on large models too.

Strengths

Paper is well written and easy to understand, with good benchmarking for large scale models and also comparing to other distillation techniques. Overall, DMD puts lesser constraints on distillation w.r.t underlying map from noise to data space enabling more good formulation for distillation. And improving stability is useful for broader adoption of DMD style formulation for distillation towards practical diffusion based applications.

Weaknesses

As the authors discuss on not using real-data within current formulation and setup, there could be a tendency for model to have model collapse? As proposed distillation objective is not sampling w.r.t teacher's marginal predictive distribution nor data-distribution. It would be useful to get some diversity metric at per-prompt level using e.g., LPIPS Diversity or etc comparing to other Distillation approaches and teacher model. Also it would be useful to understand what stage of training causes this mode collapse, i.e., does small scale training preserve diversity at cost of some quality drop or we observe a consistent drop in diversity?

Questions

What is setup of ablation without backward simulation in multi-step setting? Do forward diffusion and query the student generator and based on that query fake score function estimator? If so is the fake score function also trained on equivalent forward diffuse + generator's predictive distribution? As it is currently unclear is it alignment of fake score function alignment to generator or the backward simulation which is resulting in improved performance. Also given DMD has implicit assumption that the fake score function is capturing predictive distribution of student generator and authors identify its fitting being one of reasons for instability. It might be useful to understand how sensitive is alignment of fake score estimation to student's generator. Also, how good is quality of fake score fun at different stages of training and its implications? As at implementation level we are starting with pretrained model and in the limit if distilled student model matches pre-trained model's weights we are asking fake score function to match pre-trained model again.

Rating

8

Confidence

5

Soundness

4

Presentation

3

Contribution

3

Limitations

This work build on DMD and improves stability, good engineering practice as primary contribution of this work with limited formulation novelty or insights. So more insights on hyperparameter sensitivity, why and how are different design choices effect final performance, diversity etc as discussed above would make it more useful for community and strong contribution!

Summary

The paper introduces an upgraded version of Distribution Matching Distillation (DMD), i.e., DMD2, which addresses the limitation and inefficiency of previous DMD and improves the performance of efficient and high-quality image synthesis using diffusion models. Specifically, the authors identify the limitations of the original Distribution Matching Distillation (DMD), such as the need for a regression loss and extensive dataset construction. DMD2 eliminates the regression loss, integrates a Generative Adversarial Network (GAN) loss, and introduces a two-time-scale update rule to stabilize training. Additionally, a new training procedure is implemented to simulate multi-step sampling, addressing the training-inference mismatch. Experimental results demonstrate that DMD2 achieves state-of-the-art performance, surpassing the original DMD and other competitive models in image quality and efficiency.

Strengths

1. **Elimination of Regression Loss**: By removing the regression loss, DMD2 simplifies the training process and reduces computational costs, making it more scalable and flexible for large-scale applications. 2. **Integration of GAN Loss**: The incorporation of a GAN loss improves the quality of generated images by discriminating between real and generated samples, enhancing the overall distribution matching objective. 3. **Two-Time-Scale Update Rule**: This technique addresses training instability issues, ensuring that the fake score accurately tracks the generator’s output distribution, leading to stable and high-quality image generation. 4. **Multi-Step Sampling**: The introduction of multi-step sampling allows DMD2 to produce high-quality images in fewer steps, addressing the inefficiency of one-step generation while maintaining performance. 5. **Comprehensive Evaluation**: The paper provides extensive experimental results on various benchmarks, demonstrating DMD2's superior performance in both class-conditional and text-to-image synthesis tasks.

Weaknesses

I do not find a specific weakness of this paper.

Questions

- Training with GAN often entails numerical instability. Does DMD2 have such concerns? If it is true, could the author provide some details in overcoming the instability of DMD2? - Besides the evaluation metric such as FID and Inception scores, how about some human-related metrics such as ImageReward or aesthetic scores? Does DMD2 show comparable results to teacher SDXL?

Rating

8

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

The paper includes limitations of the proposed method, e.g., requires multiple steps to generate on par with teacher model, e.g., SDXL.