Title: One-Step Diffusion Distillation through Score Implicit Matching
Venue: NeurIPS 2024
Accepted: yes
Keywords: Diffusion mode

Diffusion Distillation

Text-to-Image Generation

Generative Adversarial Network
Abstract: Despite their strong performances on many generative tasks, diffusion models require a large number of sampling steps in order to generate realistic samples.  This has motivated the community to develop effective methods to distill pre-trained diffusion models into more efficient models, but these methods still typically require few-step inference or perform substantially worse than the underlying model.  In this paper, we present Score Implicit Matching (SIM) a new approach to distilling pre-trained diffusion models into single-step generator models, while maintaining almost the same sample generation ability as the original model as well as being data-free with no need of training samples for distillation.  The method rests upon the fact that, although the traditional score-based loss is intractable to minimize for generator models, under certain conditions we \emph{can} efficiently compute the \emph{gradients} for a wide class of score-based divergences between a diffusion model and a generator. SIM shows strong empirical performances for one-step generators: on the CIFAR10 dataset, it achieves an FID of 2.06 for unconditional generation and 1.96 for class-conditional generation.  Moreover, by applying SIM to a leading transformer-based diffusion model, we distill a single-step generator for text-to-image (T2I) generation that attains an aesthetic score of 6.42 with no performance decline over the original multi-step counterpart, clearly outperforming the other one-step generators including SDXL-TURBO of 5.33, SDXL-LIGHTNING of 5.34 and HYPER-SDXL of 5.85. We will release this industry-ready one-step transformer-based T2I generator along with this paper.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a distillation-based accelerated sampling method for various score-based diffusion models, such as EDM and Stable Diffusion. The authors have named this method Score Implicit Matching (SIM), which is designed to compress information from a diffusion-based teacher into single-step generator models. This method is somewhat based on SID, and SIM can be regarded as an extended version of SID to some extent. In my opinion, the crucial contributions of this paper are its application of SIM to Stable Diffusion and the more general design choice of the distance function.\n\nSTRENGTHS:\n1. This paper applies SIM to text-to-image (t2i) generators such as Stable Diffusion, an important experiment validating the effectiveness of accelerated sampling algorithms.\n\n2. This paper enhances SID by generalizing the design choice of the distance function to six different versions, as shown in Table 4 of the Appendix.\n\n3. This paper provides proof of their method, although the contribution of this point is relatively small since most of the proofs are based on previous work. To be specific, (1) the proof in SID use the data distribution as the target but SIM utilize the score function as target; (2) SIM uses the derivative form for derivation, but SID draws its conclusions through \n$\\nabla_{x_t}\\log p_\\theta(x_t)=[x_g-x_t]/\\sigma^2_t$\n (In EDM $\\alpha_t \\equiv 1$); (3) the only difference is the difference in the distance function, but this does not affect the derivation in any way.\n\nWEAKNESSES:\n1. The performance of SIM seems to be not effective, although the authors' statement that SIDs are reproduced because they don't have publicly available code. The performance gap may originate from a few tricks, and it is not yet known whether these tricks can further improve the effectiveness of SIM. Thus. I hope the authors can retrain several experiments related to CIFAR10 to further evaluate the effectiveness of SIM through SID's official implementation [1].\n\n[1] https://github.com/mingyuanzhou/sid\n\n2. The author spends a lot of space in the ``Score Implicit Matching\" section to introduce the principle of Algorithm 1, but I think the biggest difference between SIM and SID is that SIM adopts the matching based on the score function while SID adopts the matching based on the sample. The rest of the logic is basically the same. Therefore, I think it is a bit unreasonable to take this part as the contribution of the paper.\n\n3. The authors of SID recently released a new paper named ``Long and Short Guidance in Score identity Distillation for One-Step Text-to-Image Generation'' (LSG) which is also an extended version of SID. I wish the authors would add LSG into related work (unfortunately, different metrics were used between SIM and LSG, so SIM cannot simply compare it with LSG).\n\nQUESTIONS:\nNo\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-17"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors study the distillation of score-based models in one-step generators. They propose a new objective function for score distillation coined Score-Based Divergence. This divergence measures the mean distance between the pre-trained score-based model and a score-based model learned on the distribution induced by the one-step generator. This method is then applied on image generation problems, on standard academic benchmarks and on large latent text-to-image models.\n\nSTRENGTHS:\nOverall, this is a strong paper with a sound and novel method that obtains solid experimental results. \n\n* The authors study the distillation of score-based models, which is a very hot research topic since inference time is the major bottleneck of such models. \n\n* As far as I know, the idea of score-based divergence is novel and original. Moreover, the method comes with state-of-the-art results in a very competitive area (distillation of score-based models to one-step generators).\n\n* The general idea of score-based divergences is sound and could lead to new research works, from applied ones to theoretical ones.\n\nWEAKNESSES:\n* Not a big weakness, but we could argue that the training is demanding from a computational viewpoint, since it involves three large neural networks: generator, score-based network for data distribution, score-based network for generator's distribution. However, this is not specific to this particular method. \n\n* Minor: typo in the TL;DR. \"The submission propose an general\" -> \"The submission proposes a general\"\n\nQUESTIONS:\n* Have you observed instabilities with the optimization procedure that alternates between online score and generator? Have you tried to explore other hyper-parameters, e.g. more steps or different learning rates for the online score network? \n\n* In ProlificDreamer, the authors use LORA to fine-tune the online score-based model. Have you tried such approach to reduce the computational cost? \n\n* Could your method be directly used for learning 3D models, such as ProlificDreamer? From what I understand, setting $\\theta$ to an ensemble of particles and generating 2D images through NERFs, this would be practicable. \n\n* Theoretically, your method is agnostic to the type of generator that is used, right? Thus, we could imagine training a model with low-dimensional latent space and different architecture, such as a StyleGAN. Have you explored this?\n\nLIMITATIONS:\n.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposed Score Implicit Matching (SIM) to distill diffusion models into a one-step generator. The core idea is to use the “score-gradient theorem” to transform the minimization of score-based divergences between generator and real score functions into an implicit and tractable minimization problem. The theory does not specify the underlying divergence or distance metric for score functions, and seems to generalize many distribution matching methods, such as SiD and DI. Experiments were conducted on CiFAR-10 and text-to-image generation to show the effectiveness and robustness of SIM, by distilling EDM and PixArt-$\\alpha$, respectively.\n\nSTRENGTHS:\n- Overall, the paper is well written and easy to read, although there is still a need to further polish the presentation (as I pointed in the Weaknesses section).\n- This work provides a good theoretical insight into how to transform the intractable score-based divergence into a tractable implicit minimization problem, where both share the same gradient. \n- It generalizes many previous works (SiD, DI) as its specific cases.\n- Experiments on CIFAR-10 and text-to-image generation showed convincing results to demonstrate the effectiveness of SIM.\n\nWEAKNESSES:\n- The reason why Pseudo-Huber distance is better than other choices (L2 Square in SiD or KL in DMD) is not quite clear to me. Can the authors provide some empirical evidence to support the hypothesis of the normalization effect? For example, we can use L2 distance (or setting c=0), does the method also work well as it also normalizes the difference between generator score and teacher score? Also, what $c$ value works the best for the Pseudo-Huber distance?\n- Can the authors show how SIM becomes a DMD objective if we set $d(\\cdot)$ to a reserve KL as in DMD? It is not quite clear to me how it works.\n- There are some other standard benchmarks for one-step distillation, including distilling EDM trained on ImageNet-64 and SD trained on text-image pairs. I wonder why not compare SIM with strong baselines on these two settings?\n- The presentation can be further improved. For example, the font size in Figure 2 is too small, and the legends in the right two figures are given without any explanation. Also, all the hyperlinks of references and citations are missing throughout the paper. Besides, there is a few typos: In Table 1, the FID of SiD ($\\alpha=1$) and SiD ($\\alpha=1.2$) can be switched. In line 24, “data synthetic” might be corrected as “data synthesis”. In line 254, “one generation steps” might be corrected as “one-step generation”. \n- This is not a major concern but there is an obvious gap between the reproduced SiD results and their reported ones. The authors admitted that it is because SiD code was not released and more importantly, the SIM results can be further improved if its hyperparameters can be tuned based on SiD’s setting. So I wonder that since SiD code has now been released: https://github.com/mingyuanzhou/SiD, maybe it is a good time to look into the reproduction issue and decide whether it can improve SIM’s performance.\n\nQUESTIONS:\nPlease see my questions in the Weaknesses section.\n\nLIMITATIONS:\nThe authors adequately addressed the limitations and potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

comments: SUMMARY:
This paper proposes a new distribution matching loss between the one-step generator and the pre-trained diffusion model. The reverse KL divergence proposed in Diff-Instruct is generalized through the Score-divergence gradient Theorem.

STRENGTHS:
1. The Score-divergence theorem is well adapted to generalize Diff-Inst.

2. The empirical results with generalized objective function show better with expanded hyperparameter space.

3. A method like CTM utilizes GAN loss (w/ true data) to boost the performance. Besides the CTM, this paper's performance seems good.

4. The distillation training cost seems cheap.

WEAKNESSES:
1. More analysis on $\alpha$ seems required.

2. The score network for the student model is still required. I think the existence of an auxiliary score network is expensive. For example, a recent paper in [1] does not utilize an auxiliary score network. 

3. Adding recall metric in cifar-10 experiments if you want to claim Diff-Instruct suffers from mode-collapse. 

4. Since Diff-Instruct uses reverse-KL divergence, their mode collapse is intuitive for me. Why does the objective function you propose not to suffer from mode collapse? I know why forward KL divergence can cover the mode and reverse KL can not do it. Please explain it similar sense. There may be a difference depending on the $\alpha$.

[1] Multistep Distillation of Diffusion Models via moment-matching

QUESTIONS:
1. Is the score-divergence theorem mentioned in the diffusion model community? It seems applicable to diffusion model training.

2. Will you release the code?

LIMITATIONS:
None.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-17

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper ‘One-Step Diffusion Distillation through Score Implicit Matching’ introduces a novel framework Score Implicit Matching (SIM), to distill pre-trained diffusion models into single-step generator models. This approach achieves almost no performance loss compared to the teacher diffusion model while being data-free i.e. this approach does not require ground truth data. \n\n**Key Insights**\n- The paper proposes a flexible class of score-based divergences between the generator model (student) and the diffusion model (teacher). - Though such divergences cannot be computed explicitly, the gradients of these divergences can be computed via the score-gradient theorem, allowing for efficient training.\n- The Pseudo-Huber distance function in SIM results in stronger robustness to hyper-parameters such as learning rate and faster convergence.\n- The paper provides a detailed algorithm for implementing SIM, which involves alternating phases of learning the marginal score function and updating the generator model.\n\n**Empirical Performance**\n- On CIFAR10, SIM achieves an FID of 2.17 for unconditional generation and 1.96 for class-conditional generation.\n- For text-to-image generation, a single-step generator distilled using SIM from a leading diffusion-transformer-based model attains an aesthetic score of 6.42, outperforming other one-step generators like SDXL-TURBO, SDXL-LIGHTNING, and HYPER-SDXL.\n\n**Comparative Analysis**\n- SIM is compared to other methods like Diff-Instruct (DI) and Score Identity Distillation (SiD), demonstrating better performance in terms of robustness, convergence speed, and generative quality.\n- The empirical results highlight SIM’s superiority, particularly in maintaining the performance of the original multi-step models in a one-step generation framework.\n\n**Practical Applications**\n- The data-free nature and robust convergence make SIM a highly efficient method for distilling diffusion models.\n- The method shows strong potential for scaling to more complex tasks and larger neural networks.\n\n**Limitations and Future Work**\n- The applicability of SIM to other generative models, such as flow-matching models, is yet to be unexplored.\n- While SIM is data-free, incorporating new data might further enhance performance, which is a potential area for future research.\n\nOverall, the paper presents a significant advancement in the field of diffusion model distillation, offering a practical and efficient method for transforming pre-trained multi-step diffusion models into one-step generators without compromising performance.\n\nSTRENGTHS:\n**Originality**\n- The paper presents a novel approach - Score Implicit Matching (SIM) to distill pre-trained diffusion models into one-step generator models. This approach stands out by being data-free, which is a significant deviation from conventional distillation methods that often require extensive training data.\n- The paper’s key technical insight, the score-gradient theorem, allows computation of gradients for score-based divergences, enabling efficient training. This is a novel contribution that enhances the feasibility of implicit minimization of these divergences.\n- Pseudo-Huber Distance Function normalizes the distance vector y_t, which in turn stabilizes training loss and results in robustness to hyper-params and faster convergence.\n\n**Quality**\n- The paper provides compelling empirical evidence of SIM's effectiveness on CIFAR10 image generation. The approach also outperforms other one-step generators on aesthetic score on COCO-2017 validation set. \n- The authors conduct thorough comparisons with existing methods, such as Diff-Instruct (DI) and Score Identity Distillation (SiD), showing that SIM not only outperforms these methods but also converges faster and is more robust to hyper-parameters.\n\n**Significance**\n- SIM makes a significant contribution to the practical deployment of diffusion models in real-world applications, particularly where computational efficiency is critical.\n- The robustness and fast convergence of SIM suggest that it can scale to more complex tasks and larger neural networks, making it a valuable tool for future research and application in various domains, including image and video generation.\n- The potential of SIM to generalize to other generative models, such as flow-matching models, opens up new avenues for research and application, further enhancing its significance in the field of generative modeling.\n\n**Clarity**\n- The paper is well-organized, with a logical flow from problem formulation to the introduction of the SIM method, followed by empirical evaluations and discussions of results.\n- Key concepts, such as the score-gradient theorem and the Pseudo-Huber distance function, are explained in detail for reproducibility. \nThe tables and figures that present empirical results and algorithmic steps help in better understanding the performance and implementation details of SIM.\n\n\nThe paper presents a substantial advancement in the field of diffusion model distillation, offering a robust and efficient method for converting multi-step diffusion models into one-step generators without compromising performance. The combination of theoretical innovation with strong empirical results underscores the potential of SIM for widespread application and further research in generative modeling.\nOverall, \"One-Step Diffusion Distillation through Score Implicit Matching\" is a well-executed study that addresses a key challenge in diffusion models and provides a promising solution with broad applicability.\n\nWEAKNESSES:\n- **Narrow Scope of Comparisons**: The empirical comparisons focus primarily on Diff-Instruct (DI) and Score Identity Distillation (SiD). While these are relevant baselines, doing a broader comparison with other state-of-the-art distillation and generative modeling techniques would provide a more comprehensive evaluation of SIM's relative performance and applicability.\n- **Dataset Limitation**: The experiments are mainly conducted on the CIFAR-10, SAM-LLaVA-Caption10M and COCO-2017 validation dataset. Including results on more diverse datasets, such as higher resolution images, different image domains (e.g., medical imaging, satellite imagery), and more complex text-to-image datasets, would strengthen the claim of SIM's broad applicability.\n- **Addressing Failure Modes**: The paper does not thoroughly discuss the potential failure modes or limitations of the SIM method. Identifying scenarios where SIM might underperform, and providing insights or hypotheses on why these failures might occur, would be valuable for understanding the boundaries of the method’s applicability and for guiding future improvements.\n- **Exploration of Data Incorporation**: While the data-free nature of SIM is highlighted as a strength, the potential benefits of incorporating new data during distillation are mentioned but not explored. Providing preliminary experiments or a more detailed discussion on how incorporating data might enhance the quality of the distilled models would be a constructive addition.\n\nThe paper presents significant contributions to the field of diffusion models. Addressing the above weaknesses would further enhance its impact and applicability.\n\nQUESTIONS:\nSome of these are already mentioned under weakness. \n- Evaluating the performance of SIM on different image domains would broaden the applicability of this approach\n- Discussing scenarios where SIM would not perform well would help in understanding the boundaries and limitations for this approach and guide future improvements.\n\nAdditionally, it would be helpful if the authors can provide mode implementation details such as hyper-parameter tuning strategies and potential pitfalls during the training process. \n\nMisc\n- In the appendix, it is mentioned that a human preference study was conducted and the authors collected 30 responses in total. It would be valuable to know the results from the human preference study.\n- Grammar/Typos\n   - Line 130 - ‘... if we choose the sampling distribution to be the diffused ..’\n   - Line 220 appears to be grammatically incorrect - ‘It is on par with the CTM and the SiD’s official implementation has yet to be released’\n\nLIMITATIONS:\nThe authors have acknowledged limitations of the proposed approach. \n\n- The applicability of SIM has not been explored in the space of other generative models such as flow-matching models.\n- Although the approach is data-free, the potential benefit of incorporating data is yet to be explored. \n\nI do not see potential negative societal impacts of this work.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-28"}