One-Step Diffusion Distillation through Score Implicit Matching

Thank you for your useful feedback. We will address your concerns one by one. Before that, we first give a summary of our work.

In this work, we introduce score implicit matching (SIM), a novel distillation algorithm that achieves competitive performances on CIFAR10 generation tasks, and lossless aesthetic performances on the one-step distillation of the DiT-based text-to-image diffusion model. Besides the strong empirical performances, the theoretical foundation of SIM is also sound: we prove that the SIM loss shares the same $\theta$ gradient as our proposed score-based divergence. Therefore using an SGD-type optimization algorithm to minimize SIM loss is secretly equivalent to minimizing the intractable score-based divergence. This distinguishes SIM from SiD's theoretical loss functions that are proved to be equal to Fisher divergence in value instead of parameter gradient when using an online learned score to replace the unknown generator score.

Next, we address your concerns one by one.

Q1. The relation between SIM and SiD

A1. We appreciate the SiD paper and have obtained inspiration from it. The SiD shows a solid step to minimize Fisher divergence other than KL. However, SIM and SiD are essentially different in the following aspects.

1) SIM loss has the same parameter gradient as the intractable score-based divergence. However, the SiD's theoretical losses only have equal value instead of an equal parameter gradient.

In Section 3.2, we have shown that SIM loss has the same parameter gradient (equation 3.5-3.7) as the target score-based divergence (equation 3.4) regardless of substituting the score function with an online learned score model (please check Theorem 3.1). This means that using an SGD-type optimization algorithm to minimize SIM loss is equivalent to minimizing the score-based divergence. On the contrary, the SiD theoretical loss functions (equation 13 and equation 20 in SiD paper) only have an equal value to the target Fisher divergence instead of the parameter gradient when using an online score function to replace the intractable generator score function. This means using a gradient-based optimization algorithm does not guarantee the minimization of the Fisher divergence. This might be the reason why SiD's theoretical-driven losses do not work well in practice as the authors of SiD have pointed out. After the theoretical argument, the authors of SiD empirically find that a linear combination of the two "bad" losses results in a strong loss function, which they call the "fused loss". Interestingly, we find that when taking SIM's distance function to be squared L2 distance, SIM recovers the empirical "fused loss" of SiD with $\alpha=1.0$. From this point of view, SIM, as well as our score-divergence gradient theorem brings a theoretical tool for analyzing SiD's empirical fused loss in a gradient view.

2) SIM is defined over a general family of score-based divergence with a flexible choice of distance function. As a comparison, the SiD only targets Fisher divergence, which can be viewed as a special case of SIM. Besides, trying to generalize SiD's theoretical loss (equation 20 in SiD) to support general distance functions is not direct. The proof of Equation 19 of the SiD paper relies on the expansion of the squared L2 difference. However, such an expansion may not hold for general distance functions such as L2 distance and the Pseudo-Huber distance.

3) SIM empirically shows faster convergence than SiD. This makes SIM favored when training large-scale one-step text-to-image generator models. Please see A2 for details.

4) SIM shows novel scalability to complicated text-to-image one-step distillation, however, the SiD did not show such a possibility before the NeurIPS submission.

In short conclusion, SIM has proven the equality of the loss function and the intractable divergence in a parameter gradient view, making it a solid theoretical contribution. Besides, SIM has shown success in scaling to train strong one-step text-to-image generators without performance losses, making it a solid empirical contribution.

Q2. Will SIM be further improved with the SiD codebase?

A2. We appreciate your constructive suggestion. After reading the SiD's official code, we find the official code of SiD has multiple tricks that are different from our original implementation. With some tests, we find that these tricks, such as masking out the nan samples, the weighting function, as well as the continuous-time Karras-$\rho$ noise sigma distribution help stabilize the training and improve the performance.

Due to the limited time of the rebuttal period and limited compute, we conduct experiments with two settings to compare SIM and SiD behavior at different training stages: (1) training from scratch, and (2) resuming from a nearly converged generator; Please check our global rebuttal and added one-page for technical details.

The new experiment shows that the faster convergence of SIM than SiD ($\alpha=1.0$) using the same implementation techniques. However, training to the end will take a very long time, and we have to acknowledge that we probably can not make a final conclusion that SIM will be better than SiD without sufficient computation. Besides, we appreciate SiD's high-quality implementations and solid contributions.

Q3. The authors of SID recently released a new paper named LSG, I wish the authors would add LSG into related work.

A3. Thank you for your kind reminder. The LSG is publicly available after the NeurIPS submission deadline. After a careful reading, we find it technically interesting. We think the Long-and-short guidance technique proposed in LSG may further improve the SIM performance for text-to-image models. We will add the LSG to our related work in the revision.

We hope our response has resolved your concerns. If you still have any concerns, please let us know.

We are glad that you like our work. We appreciate your valuable suggestions, and we will incorporate them in our revision. In this paper, we introduce score implicit matching (SIM), a theoretically sound distillation method that secretly minimizes a general family of score-based divergences when distillation. Our experiments on CIFAR10 generation as well as the one-step DiT-based text-to-image generator model show the superior performances of SIM on industry-level applications. In the following paragraphs, we will address your concerns one by one.

Q1. Not a big weakness, but the training computation cost may be demanding.

A1. We acknowledge that the training cost is a little bit expensive because the SIM involves an additional online diffusion model. However, in practice, we find that the additional cost is acceptable by using Pytorch techniques. For example, in our text-to-image experiment, we use 4 Nvidia A100-80G GPUs for training the one-step generator. We use the Pytorch Distributed data-parallel (DDP) to support multi-GPU gradient syncing and use BF16 numerical formate for training. With BF16 training, we can effectively train the model with a total batch size of 256 even without the gradient accumulation. If we use gradient accumulation techniques, we can enlarge the batch size to 1024 with some speed decline.

Q2. Are there instabilities with the alternating optimization procedure? Would be better to explore other hyper-parameters.

A2. We appreciate your good intuition. We have observed that if the hyper-parameters are not set properly, the training may diverge. For instance, we find that when the weighting functions and the log-normal noise sigma distribution for training online diffusion models are not properly set, the alternating training will lead to a poor result. However, compared with Diff-Instruct and SiD, we find that the updating of a one-step generator using SIM is pretty robust to hyper-parameters. We find that using different loss scaling, different learning rates (in a reasonable interval), as well as several kinds of weighting functions, does not affect the SIM updates for the generator too much. We appreciate your constructive suggestions as a new direction to explore. We have tried to update the online diffusion twice, however, this slows down the training process by approximately 1.5 times and we do not observe significant performance improvement, therefore, we still use one diffusion update per generator update, with two updates using the learning learning rate. However, we will explore more of these settings in the revision.

Q3. Have you tried using Lora such as in ProlificDreamer [2] to reduce the computational cost?

A3. We appreciate your good intuition. We like the trick proposed in ProlificDreamer to use a Lora online diffusion, which can reduce the memory cost. However, as we have shown in A1, when we use BF16 training, we are not bothered by memory issues, therefore we do not use the Lora online diffusion. Another reason is that the Lora model may have worse modeling ability than a full-parameter model, which can potentially harm the performance of SIM.

Q4. Can SIM be used to train NeRF models using 2D diffusion? Can SIM be applied to other generator models such as GANs?

A4. We highly appreciate your good intuition. As you can see, the SIM is a general method that can be used for a wide range of applications other than one-step diffusion distillation. In the rebuttal period, we conduct two more experiments: (1) text-to-3D generation using text-to-2D diffusion; and (2) improving the GAN generator using diffusion models; Please check our Global rebuttal cell for details. On both new applications, SIM has shown strong performances, demonstrating its broad applicability.

We hope our response has resolved your concerns. If you still have any concerns, please let us know.

Thank you for your useful suggestions. We will address your concerns one by one. Before that, we first give a summary of the main contributions of our work.

In this work, we introduce score implicit matching (SIM), a novel distillation algorithm that achieves competitive performances on CIFAR10 generation tasks, and lossless aesthetic performances on the one-step distillation of the DiT-based text-to-image diffusion model. Besides the strong empirical performances, the theoretical foundation of SIM is also sound: we prove that the SIM loss has the same $\theta$ gradient as our proposed score-based divergence. Therefore using an SGD-type optimization algorithm to minimize SIM loss is secretly equivalent to minimizing the intractable score-based divergence. Besides, our proposed score-gradient theorem may bring new tools for future studies on generative models.

Next, we address your concerns one by one.

Q1. Clarify the reason for using Pseudo-Huber distance.

A1. We appreciate your keen intuitions. Though the standard L2 distance seems a fair choice, using the Pseudo-Huber distance is better. The Phuber loss (equation 3.8) has a denominator of the form $\sqrt{\||y_t\||^2_2 + c^2}$. Here $y=s_g(x,t) - s_d(x,t)$ will be close to zero when the generator is converging. Therefore, if we set $c=0$ (the case of squared L2 distance), such a denominator is unstable and can lead to numerical issues. This means that Phuber with a small $c$ is better than the standard of L2 distance.

Another advantage of Phuber loss (equation 3.8) is that it has an adaptive self-normalization effect on the loss scale. This helps the training loss to be stable with a constant scale.

Q2. Can the authors show how SIM becomes a DMD objective if we set $d(\cdot)$ to be reversed KL?

A2. If we consider the EDM formulation of Algorithm 2 in the Appendix, the DMD loss can be viewed as a special case of SIM loss. Let $d_{\phi}(x_t,t)$ denotes the online EDM model, and $d_q(x_t,t)$ the teacher EDM model. In Algorithm 2, the SIM loss with Phuber distance writes: $L(\theta) = \frac{d\_{\phi}(x\_t,t) - d\_q(x\_t,t)}{\sqrt{\||d\_{\phi}(x\_t,t) - d\_q(x\_t,t)\||\_2^2 + c^2}} (d\_{\phi}(x\_t,t) - x\_0).$ If we cut off the $\theta$ gradient of $x_t$, then $d_{\phi}(x_t,t)$ and $d_q(x_t,t)$ will have no parameter dependence. So the SIM loss is equivalent to DMD loss with a weighting function $\frac{1}{\sqrt{\||d_{\phi}(x_t,t) - d_q(x_t,t)\||_2^2 + c^2}}$. However, the theory behind SIM and DMD is essentially different, because SIM aims to minimize a general family of score-based divergences, instead, DMD aims to minimize the KL divergence which is notorious for mode-collapse behavior.

There might be a reason why score-based divergence is a better choice than KL. Recall that the definition of KL divergence needs the likelihood ratio $\frac{p_g(x)}{p_d(x)}$, which will be ill-defined when $p_d$ and $p_d$ have misaligned density support. This can potentially lead to mode-collapse issues when using KL divergence as a target. However, the score-based divergence does not have such a "ratio", and therefore is safe in the case when $p_g$ and $p_d$ have misaligned support.

Q3. The motivation to distill PixelArt-$\alpha$ diffusion model.

A3. In Section 4.1, we evaluate the SIM on the CIFAR10 generation benchmark thoroughly and observe strong performances without using GAN losses. After that, we would like to scale SIM to challenging text-to-image generation. We choose PixelArt-$\alpha$ diffusion as a teacher model for three reasons. (1) We appreciate the great efforts of previous works on diffusion distillation of UNet-based t2i models. However, distilling DiT-based diffusion models, such as PixelArt-$\alpha$ diffusion, lack explorations. (2) The PixelArt-$\alpha$ model uses the T5-xxl text encoder, which is able to understand long prompts better than Stable Diffusion models with CLIP text encoders. We are motivated to explore high-quality one-step generator models that have such a long-prompt following ability. (3) We would like to explore the limits of a one-step model that is favored by human users in terms of Aesthetic Scores. The PixelArt-$\alpha$ is a good model that has been trained with high-quality aesthetic images.

Q4. Missing hyperlinks and some suggestions to improve the presentation.

A4. We feel sorry for the confusion. The hyperlinks are missing because we have separated the main pages and the appendix. We will incorporate your suggestions to improve our presentation in the revision.

Q5. Not a major concern, but I wonder if SIM can be further improved with the SiD codebase.

A5. We appreciate your constructive suggestion. After reading the SiD's official code, we find the official code of SiD has multiple tricks that are different from our original implementation. With some tests, we find that these tricks, such as masking out the nan samples, the weighting function, as well as the continuous-time Karras-$\rho$ noise sigma distribution help stabilize the training and improve the performance.

Due to the limited time of the rebuttal period and limited compute, we conduct experiments with two settings to compare SIM and SiD behavior at different training stages: training from scratch, and resuming from a nearly converged generator; Please check our global rebuttal and added one-page for technical details.

Our new experiment shows the faster convergence of SIM than SiD ($\alpha=1.0$) using the same implementation techniques. However, training to the end will take a very long time, and we have to acknowledge that we probably can not make a final conclusion that SIM will be better than SiD without sufficient computation. Besides, we appreciate SiD's high-quality implementations and solid contributions.

We hope our response has resolved your concerns. If you still have any concerns, please let us know.

Thank you for your constructive feedback. We will address your concerns one by one. Before that, we first give a summary of the main contributions of our work.

In this work, we introduce score implicit matching (SIM), a novel diffusion distillation algorithm that achieves competitive performances on CIFAR10 generation tasks, and lossless aesthetic performances on the one-step distillation of the DiT-based text-to-image diffusion model. Besides the strong empirical performances, the theoretical foundation of SIM is also sound: we prove that the SIM loss shares the same $\theta$ gradient as our proposed score-based divergence. Therefore using an SGD-type optimization algorithm to minimize SIM loss is secretly equivalent to minimizing the intractable score-based divergence. Besides, our proposed score-gradient theorem may bring new theoretical tools for future studies on generative models.

Next, we address your concerns one by one.

Q1. More analysis on $\alpha$.

A1. We are not sure what you mean by $\alpha$. In the SIM algorithm (Algorithm 1), we do not have an $\alpha$ hyperparameter. We guess that you mentioned $\alpha$ to mean the SiD algorithm, which uses an empirical $\alpha$ hyperparameter to fuse its loss functions. SIM is different from SiD in theory. As we have shown in Sections 3.3 and 3.4, when taking the distance function to be squared L2 distance, SiD's empirical fused loss with an $\alpha=1.0$ is a special case of SIM. Besides, SIM supports a flexible choice of distance functions to define and instantiate different score-based loss functions. This distinguishes SIM from previous methods such as SiD and Diff-Instruct.

Q2. The existence of an auxiliary score network is expensive. A recent paper in [1] does not utilize an auxiliary score network.

A2. We appreciate your keen intuition. We acknowledge that SIM needs an additional online diffusion model, which brings more memory costs. However, in practice, we find that the additional cost is acceptable if we use Pytorch techniques. For example, in our text-to-image experiment, we use 4 Nvidia A100-80G GPUs for training the one-step generator. We use the Pytorch Distributed data-parallel (DDP) to support multi-GPU gradient syncing and use the BF16 numerical format for training. We can effectively train the model with a total batch size of 256 even with no gradient accumulation. If we use gradient accumulation techniques, we can enlarge the batch size to 1024 with some speed decline.

However, we highly appreciate your suggestions to refer to [1] to figure out the possibility of getting rid of the online diffusion model. After a careful read of [1], we like the idea of moment matching to drop the online diffusion model. We will involve a discussion on SIM and [1] in the revision and leave the study that drops the online diffusion while maintaining strong one-step generation performance in our future work.

Q3. Why does SIM not suffer from mode collapse?

A3. We appreciate your good question! SIM is essentially different from the Diff-Instruct algorithm in the divergence that the generator aims to minimize. If we use $p_d$ to represent data distribution and $p_g$ the generator distribution, the Diff-Instruct minimizes the integral of the KL divergence between the generator and the teacher diffusion model with a form $\mathcal{D}\_{KL} (p\_g, p\_d) = \mathbb{E}\_{x\sim p\_g}\log \frac{p\_g(x)}{p\_d(x)},$ while the SIM minimizes the score-based divergence with a form $\mathcal{D}\_{SD}(p\_g,p\_d) = \mathbb{E}\_{x\sim \pi} d(\nabla\_x \log p\_g(x) - \nabla\_x \log p\_d(x)).$ KL divergence is notorious for mode-collapse issues because the likelihood ratio $\frac{p_g(x)}{p_d(x)}$ will be ill-defined if $p_g$ and $p_d$ have misaligned density support. However, the score-based divergence does not have such a "ratio", and therefore is safe in the case when $p_g$ and $p_d$ have misaligned support. Besides, our SIM with Phuber distance (equation 3.8) has a self-normalization form, which helps stabilize the training loss to a constant scale which potentially addresses the mode-collapse issue.

Q4. Is the score-divergence theorem mentioned in the diffusion model community?

A4. As we introduced in Section 3.2, the Score-divergence gradient Theorem is proved by taking the parameter gradient to both sides of the so-called score-projection identity. The score-projection identity is well-known for proving the equivalence of denoising score matching and classical score matching in [2]. Recently, [3] also used the projection identity to prove value equality when deriving SiD loss.

However, as far as we know, our score-divergence gradient theorem is the first time it appeared in the community, marking it as an important theoretical contribution in our paper. In this paper, we propose and use the score-divergence gradient theorem to prove the gradient equality of SIM loss and intractable score-based divergence instead of value equality. This distinguishes SIM from previous methods such as SiD that only prove value equality that is insufficient in theory.

Q5. Will you release the code?

A5. We will release the code if the paper gets published.

Q6. Add a recall metric to show that Diff-Instruct suffers from mode-collapse.

Q6. Thank you for your suggestion. In the rebuttal period, we compute the precision and recall of the DI model. It has a recall of 0.52 which is pretty low, indicating its mode-collapse issue.

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

[2] A Connection Between Score Matching and Denoising Autoencoders

[3] Score identity Distillation: Exponentially Fast Distillation of Pretrained Diffusion Models for One-Step Generation

We hope our answers have resolved your concerns, and if you still have any concerns, please do let us know.

We are glad that you like our work. We appreciate your valuable suggestions, and we will take them in the revision. In the following paragraphs, we will address your concerns one by one.

Q1. Doing a broader comparison with other distillation and generative modeling techniques would provide a more comprehensive evaluation of SIM's relative performance and applicability.

A1. In this paper, we have compared SIM with Diff-Instruct and SiD, showing its soundness on both the theoretical and empirical side. Another leading diffusion distillation method might be the consistency trajectory model (CTM) [1], which uses self-consistency to distill the model. Though CTM has shown strong performances, its performances rely on using additional LPIPS losses and GAN losses, which need ground truth data and are also tricky. Besides, the CTM model requires running multiple steps of reversed diffusion sampling, which is more computationally inefficient than SIM without running multiple-step sampling with teacher diffusion. We will compare with other leading distillation methods in the revision.

Q2. Including results on more diverse datasets, such as medical imaging, and satellite imagery would strengthen the claim of SIM's broad applicability.

A2. We highly appreciate your suggestions to use SIM on broader applications. In the rebuttal period, we conduct two more experiments: (1) text-to-3D generation using text-to-2D diffusion; and (2) improving the GAN generator using diffusion models; Please check our Global rebuttal cell for details. On both new applications, SIM has shown strong performances, showing its broad applicability.

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

A3. We appreciate your keen intuition. One advantage of SIM is the image-data-free property. However, we agree that incorporating training data with SIM will potentially strengthen the SIM in real-world applications. A possible way to incorporate data is to use additional GAN discriminator and GAN losses together with the SIM loss. This can inject knowledge from new data into the generator model. Another direction that is worth exploring is to extend the one-step generator to multi-step generative models using variants of SIM. This potentially will combine techniques with consistency-based models. However, both the GAN loss and the multi-step generalization are more like a combination of different helpful techniques in generative modeling. In this submission, our goal is to thoroughly study the theory and practical performance of SIM. We will leave the study that incorporates data with SIM in the future.

Q4. Discussing scenarios where SIM would not perform well would help in understanding the boundaries and limitations of this approach and guide future improvements.

A4. We appreciate your constructive suggestion. In our text-to-image experiment, we find our SIM-DiT-600M sometimes fails to generate high-quality tiny human faces, hands, as well as legs. However, we find that the teacher PixelArt-$\alpha$ model also has such issues. We believe that a stronger teacher diffusion model will possibly address such failure cases. Please check Figure 2 in our added one-page in global rebuttal for some failure cases.

Q5. The human preference study.

A5. We feel sorry for the confusion, we put the human preference comparison of SIM-DiT-600M and PixelArt-$\alpha$ diffusion model in Table 3 in our submission. The result shows that the teacher PixelArt-$\alpha$ model slightly outperforms the SIM one-step model with a winning rate of 54.88%. We will polish the presentation in the revision.

Q6. Implementation details.

A6. We will release the code if the paper gets published. In this response, we add some more implementation details for the one-step text-to-image generator model. When distilling the generator, we first reformulate the PixelArt-$\alpha$ diffusion with the so-called "data-prediction" formulation as proposed in the EDM paper [2]. We then initialize our generator using the "data-prediction" model with a fixed noise sigma of 2.5, following the same setting of Diff-Instruct and SiD. When distillation, we use a fixed classifier-free guidance scale of 4.5 for teacher diffusion. For the auxiliary diffusion model, we do not use the CFG. This setting is inspired by the DMD paper [3]. We will add more implementation details in the Appendix in the revision.

Q7. Some typos.

A7. We will address the typos and improve our writing in the revision.

[1] Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion

[2] Elucidating the Design Space of Diffusion-Based Generative Models

[3] One-step Diffusion with Distribution Matching Distillation

We hope our response has resolved your concerns. If you still have any concerns, please let us know.