Improved Distribution Matching Distillation for Fast Image Synthesis

We sincerely appreciate Reviewer xaYh's constructive feedback. We will fix all typos. Below, we address the remaining concerns.

How is DMD2 related to ADD and LADD? What is the most significant difference?

Thank you for the opportunity to discuss how DMD2 relates to other concurrent GAN based methods such as ADD and LADD. The ADD paper [23] utilizes a pretrained DINO-based classifier in pixel space, which we found to be less efficient in terms of training and leading to reduced diversity, as detailed in Table 6. More recent efforts like LADD, UFOGEN, and SDXL-Lightning employ latent diffusion GAN-based networks akin to our DMD2. A significant difference, which is often overlooked, is that DMD-style training inherently integrates classifier-free guidance directly into the real score of the DMD gradient. This integration simplifies the training process of our model. In contrast, purely GAN-based methods typically struggle to incorporate classifier-free guidance directly. For example, SDXL-Lightning needs to combine GAN with progressive distillation to enable CFG, while LADD relies on diffusion generated images with CFG as real data in its GAN discriminator. These approaches often complicate the training process. We will elaborate further on these distinctions in the related work section of our revised paper.

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]?)

Similar to our response to Reviewer xYYF, we assess the mode collapse situation under two scenarios:

• Class-Conditional Image Generation: We observed no mode collapse, as evidenced by the state-of-the-art FID scores for image generation (see Table 1).
• Text-to-Image Generation: This scenario is more complex. We utilize classifier-free guidance, which typically trades diversity for image quality. At high guidance settings, using the SDXL baseline, we achieved superior image quality with diversity comparable to or better than other distillation methods but slightly worse than the teacher model ((see Table 6 and the new results in our Response to Reviewer aF5v).

Our current framework trains stably and produces excellent image quality. However, there remains a small gap in output diversity compared to the original teacher model and we are open to exploring future methods that might enhance this diversity/quality tradeoff by better integrating trajectory-preserving techniques (such as the more efficient consistency distillation) with distribution matching methods. We believe this is a promising direction for future research.

We sincerely appreciate Reviewer uheV's constructive feedback. Below, we address the remaining concerns.

How to select the set of hyperparameters?

Thank you for your question. Our approach to selecting hyperparameters is straightforward and consistent across all datasets. We utilize the maximum batch size our compute setup allows. The learning rate is determined as the highest value that does not cause divergent loss within the first 500 iterations. We set the number of TTUR iterations to the minimum required for stability. The guidance scale is chosen based on what yields the best results for the teacher model. For the GAN weight, our method works well across a wide range of values, as demonstrated by the following ImageNet FID scores:

We will include these guidelines in our updated version.

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?

We are open to the possibility of a one-step generation that matches the performance of the teacher model. However, we want to emphasize that achieving this in a single-step generation is substantially more challenging than in a multi-step process, especially when distilling complex networks like SDXL. While our one-step generator can match the teacher model in quantitative metrics such as FID, qualitative aspects of image quality present greater challenges. Visually, the one-step process still produces some artifacts that are difficult to eliminate (See Figure 11). An exciting direction for future research could involve refining our training objectives to address these issues more directly, as seen in recent developments like those explored in HyperSD [1].

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?

While both models are trained on the LAION dataset, they likely utilize different subsets, leading to variations in the noise-to-data mapping. The higher quality and particularly the subtler details found in SDXL are indeed more challenging to capture, potentially requiring more advancements in loss design.

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?

We have observed improved performance with increased batch sizes and compute. We used a batch size of 128 for SDXL, as this is the maximum our current setup can accommodate within our compute budget of 3 days—especially considering that the SDXL model is three times larger than SDv1.5. Exploring a larger batch size, such as 2048, with additional resources in the future would indeed be an interesting experiment to potentially enhance performance further.

Some Writing Inconsistency

Thank you for pointing out these inconsistencies. Regarding the discrepancy between 50 and 100 steps in Figure 4, the teacher model uses 50 steps, but effectively has 100 forward passes due to the application of classifier-free guidance. For EDM, upon reevaluating the released model, we recorded a FID of 2.32, which we will update to 2.22 in our revised version to reflect the most accurate data. We appreciate your attention to detail and will correct the remaining issues. Thank you again for your valuable feedback.

Extra cost per training iteration of backward simulation

Thank you for raising the issue of training computational cost. The training iteration time for the model with backward simulation is approximately 9.2 seconds per iteration, compared to 7.8 seconds for the model without backward simulation. We also discovered that enabling backward simulation only during the last 20% of the training epochs yields comparable results. This approach offers a more computationally efficient option when resources are limited.

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?

From the data in Table 2, the difference in FID between the one-step and four-step distillation (0.3) is generally within the range of variability observed across different runs, indicating comparable performance. The significant improvement in Patch FID for the four-step distillation reflects better local image detail, aligning with qualitative improvements we observed. However, the CLIP score did decrease, which may suggest that the four-step method slightly sacrifices prompt alignment for image quality. This also happens for previous approaches with good text to image alignment like SDXL-Turbo.

[1] Ren, Yuxi, et al. "Hyper-sd: Trajectory segmented consistency model for efficient image synthesis." arXiv preprint arXiv:2404.13686 (2024).

We sincerely appreciate Reviewer aF5v's constructive feedback. Below, we address the remaining concerns.

It would be better to include a detailed training algorithm to clearly showcase the modifications over the original DMD training process.

Thank you for your suggestion. As you accurately summarized, our DMD2 model eliminates the need for the regression loss and incorporates a diffusion GAN loss along with backward simulation to mitigate training-inference mismatches. To clearly illustrate these modifications over the original DMD training process, we will include a detailed comparison of the training algorithms in the revised version of our paper.

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?

We acknowledge that DMD2 relies on access to real data to enhance diversity and image quality. However, it is important to note that the exact dataset used to train the teacher model is not required. For our training, we utilized a random set of 500,000 images from the LAION database, which is generally of lower quality than the curated aesthetic dataset of SDXL. This demonstrates that our GAN loss can be effectively applied using an alternative dataset, underscoring its versatility and universal applicability, regardless of the specific model being distilled.

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?

Thank you for your question. In DMD2, we demonstrate that much of the mode collapse observed in the original DMD method relates more to training issues than to an inherent inability of the DMD loss to support diverse generator training. Practically, we can assess the final performance under two scenarios:

• Class-Conditional Image Generation: We observed no mode collapse, as evidenced by the state-of-the-art FID scores for image generation (see Table 1).
• Text-to-Image Generation: This scenario is more complex. We utilize classifier-free guidance, which typically trades diversity for image quality. At high guidance settings, using the SDXL baseline, we achieved superior image quality with diversity comparable to or better than other distillation methods but slightly worse than the teacher model (see Table 6 and the new results in our Response to Reviewer aF5v).

Our current framework trains stably and produces excellent image quality along with comparable diversity. However, we are open to exploring future methods that might enhance this diversity/quality tradeoff by better integrating trajectory-preserving techniques (such as the more efficient consistency distillation) with distribution matching methods. We believe this is a promising direction for further research.

Why intermediate outputs of DMD2 follow a certain trajectory?

This observation was initially surprising to us as well. Although our generator does not follow the teacher diffusion model’s sampling trajectory, our training methods lead to few-step generators where the first output significantly shapes the general structure. Subsequent images tend to closely resemble this initial output. This effect is likely influenced by the relatively high signal-to-noise ratios in subsequent sampling steps, which preserve much of the structure even after noise injection.

We sincerely appreciate Reviewer aF5v's constructive feedback. Below, we address the remaining concerns.

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?

There may be some misunderstanding regarding our use of real data. In fact, our model does incorporate real data during training via the GAN loss component. This inclusion enhances diversity and helps to prevent mode collapse. The positive impact of integrating GAN loss with real data is evident in the improved performance detailed in Tables 3 and 4.

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.

Thank you for your suggestion! We have indeed conducted a diversity assessment using LPIPS diversity metrics, which is presented in Table 6 of the appendix. Our results indicate that our model's diversity is on par with other distillation approaches, such as the latent consistency model, and significantly exceeds that of purely GAN-based methods like SDXL-Turbo. While our model shows slightly worse diversity compared to the teacher model, it offers substantially better image quality, as illustrated in Figure 5. We believe this often represents a more advantageous trade-off for practical text-to-image applications.

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?

Thank you for your suggestion! In response, we retrained our SDXL model and monitored diversity metrics throughout the training process. The results from this new run showed a small diversity improvement over those reported in our paper. Specifically, the model displays higher diversity at the beginning, which gradually diminishes as image quality improves. Despite this trend, the final diversity scores of the model remain closely comparable to those of the teacher model (0.63 vs 0.64 for the teacher), indicating a well-maintained balance between diversity and image quality.

Setting without backward simulation

The reviewer’s interpretation is correct. In the setting without backward simulation, we add noise to real images and then feed these noisy images into our generator. The output from the generator is then supervised using both the DMD and GAN loss. Additionally, the fake score function is trained based on this generator output. We will clarify this process more thoroughly in the revised version of our paper.

how sensitive is alignment of fake score estimation to student's generator and how good is the quality of fake score at different stages of training and its implications?

Thank you for suggesting this set of further analyses. Due to time constraints, these will be included in the revised version of our paper. Regarding the first part, we believe achieving proper alignment between the fake score estimation and the generator’s output distribution is crucial. As shown in Figure 9 of the main paper, more frequent training of the fake score—aimed at enhancing its accuracy—leads to more stable generator training and overall improved performance. We will provide a more comprehensive analysis in the revised version. For the second part, we plan to retrain an ImageNet model and monitor the performance of the fake score model at various training stages by assessing the quality of its sample outputs. We appreciate the reviewers’ valuable suggestions and look forward to incorporating these insights!

We sincerely appreciate Reviewer Yut7's constructive feedback. Below, we address the remaining questions regarding numerical instability and human-related metrics.

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?

The GAN component in DMD2 does not introduce numerical instability. This stability can be attributed to our use of a diffusion GAN framework, where both the generators and discriminators are initialized from a pretrained diffusion model. Before classification, images are also treated with noise injection, enhancing stability. This method has proven more stable than traditional GANs that utilize pixel space discriminators without noise injection, as supported by several recent studies [1, 2]. Furthermore, DMD2 includes a two-timestep-scale update rule that bolsters stability. Additionally, DMD2 utilizes a weighted combination of DMD and GAN losses. The DMD loss consistently guides the overall structure of the images, preventing mode collapse and ensuring training stability even with the added GAN components.

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?

Thank you for this insightful suggestion! We have extended our evaluation of DMD2 and SDXL to include both ImageReward and aesthetic scores, using PartiPrompts [3] for consistency with our human evaluation. Below is a summary of our findings:

These results demonstrate that DMD2 consistently outperforms SDXL, corroborating the findings from our main paper's FID and human evaluation results. We will incorporate these additional metrics into the revised version of our paper.

[1] Wang, Zhendong, et al. "Diffusion-gan: Training gans with diffusion." ICLR 2023

[2] Xu, Yanwu, et al. "Ufogen: You forward once large scale text-to-image generation via diffusion gans." CVPR 2024.

[3] Yu, Jiahui, et al. "Scaling autoregressive models for content-rich text-to-image generation." arXiv preprint arXiv:2206.10789 2.3 (2022)