Diff-Instruct: A Universal Approach for Transferring Knowledge From Pre-trained Diffusion Models
Due to the ease of training, ability to scale, and high sample quality, diffusion models (DMs) have become the preferred option for generative modeling, with numerous pre-trained models available for a wide variety of datasets. Containing intricate information about data distributions, pre-trained DMs are valuable assets for downstream applications. In this work, we consider learning from pre-trained DMs and transferring their knowledge to other generative models in a data-free fashion. Specifically, we propose a general framework called Diff-Instruct to instruct the training of arbitrary generative models as long as the generated samples are differentiable with respect to the model parameters. Our proposed Diff-Instruct is built on a rigorous mathematical foundation where the instruction process directly corresponds to minimizing a novel divergence we call Integral Kullback-Leibler (IKL) divergence. IKL is tailored for DMs by calculating the integral of the KL divergence along a diffusion process, which we show to be more robust in comparing distributions with misaligned supports. We also reveal non-trivial connections of our method to existing works such as DreamFusion, and generative adversarial training. To demonstrate the effectiveness and universality of Diff-Instruct, we consider two scenarios: distilling pre-trained diffusion models and refining existing GAN models. The experiments on distilling pre-trained diffusion models show that Diff-Instruct results in state-of-the-art single-step diffusion-based models. The experiments on refining GAN models show that the Diff-Instruct can consistently improve the pre-trained generators of GAN models across various settings.
Paper
Similar papers
Peer review
Summary
The paper discusses the distillation of pre-trained diffusion models. The paper introduces a framework called Diff-Instruct, which allows for the transfer of knowledge from pre-trained DMs to other generative models in a data-free manner. Diff-Instruct utilizes a mathematical foundation based on minimizing a divergence called Integral Kullback-Leibler (IKL) divergence. This divergence is specifically designed for DMs and is more robust in comparing distributions with misaligned supports. The text also highlights connections to existing works such as DreamFusion and generative adversarial training. The effectiveness of Diff-Instruct is demonstrated through two scenarios: distilling pre-trained diffusion models and refining existing GAN models. Experimental results show that Diff-Instruct achieves state-of-the-art performance in single-step diffusion-based models and consistently improves pre-trained generators of GAN models across various settings.
Strengths
* The paper tackles the problem of distilling diffusion models and addresses the important topic of accelerating inference for DMs. * The proposed method utilizes a novel formulation based on the Integral Kullback-Leibler (IKL) divergence, which adds further interest to the approach. The authors make an effort trying to explain why this divergence is a reasonable choice for distillation. * The experimental results showcased the effectiveness of this method, particularly in the context of 1-step distillation on datasets such as CIFAR and ImageNet64, where the results were highly promising. These findings demonstrate the potential of the proposed method for enhancing the efficiency and performance of diffusion models in small domains.
Weaknesses
* While the distillation results for CIFAR and ImageNet using the proposed method are good, they may not have practical utility in larger domains. Distillation becomes particularly crucial for such larger domains, and it remains unclear whether the method will be effective with larger images. * Additionally, although the method doesn't require additional data, it does necessitate the use of two extra models, namely the student model and an auxiliary model for gradient propagation. This requirement can be prohibitive when dealing with larger domains and bigger models. * I have some doubts about the notation (see questions). I would try to be consistent with the notation in the literature (see https://arxiv.org/abs/2011.13456).
Questions
* When discussing knowledge transfer, what is the primary objective? Is it knowledge distillation or transfer? It would be helpful if you could simplify and specifically tie your contributions to the results of your experiments. * Regarding Algorithm 1: when updating $\theta$, shouldn't $s$ be function of $\theta$? Could you clarify how the gradient has been formulated? * In Equation 1, does $q$ represent the forward or reverse process in a Diffusion Model? It appears you're using $q$ for the reverse process of a pre-trained Diffusion Model, which is conventionally represented as the approximate posterior. Could you clarify this point? * Furthermore, referencing Equation 1 in this paper (https://arxiv.org/abs/2011.13456), it seems that the notations in Equations 2.2 and 3.2 are inconsistent. Could you please provide some clarification? * After training the implicit model, what role does $s_{\phi}$ serve? Is it merely an auxiliary model, and if so, how do you initialize it? * How much time is required to distill a pre-trained diffusion model? * You've mentioned that the Inverse Kullback-Leibler (IKL) divergence is more robust. Could you explain how you've arrived at this conclusion?
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness
2 fair
Presentation
1 poor
Contribution
2 fair
Limitations
* Notation confusing or inconsistent. * Not clear what is the goal of the paper (transfer vs distillation).
Summary
This work distills diffusion models to a GAN-style generator. The author provides mathmatic connection to other works, as well as experimental comparisons between the distilled GAN-style generator with other (mostly) diffusion models.
Strengths
- This work introduce a way to distill diffusion model to GAN-style generator. - And shows the connection to other works. - The experiments show the proposed method works well.
Weaknesses
- Lack baseline method. Table1 and Table2 shows the performance many diffusion methods, but they are either diffusion model trained on different datasets or distilled diffusion model that may not use the same teather model as the author used (i.e. EDM). In this case, the numbers in table1 & 2 are not compariable. There should be at least one baseline method that use the same teather model (EDM), the same generator architecture (one-step GAN-style generator), but different distillation method. Here is a simple baseline idea, generating N images by the teather model (EDM), training a GAN on the generated images. - Lack time comparison. One big benefit of one-step generator (as well as distillation) is inference speed. Otherwise, we can use the teacher model directly.
Questions
Please check the weaknesses.
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness
3 good
Presentation
2 fair
Contribution
2 fair
Limitations
this work has no negative societal impact.
Summary
The paper introduces a method to transfer the knowledge of a pre-trained diffusion model to an implicit generator model, such that the distribution of the generated samples from the generator matches that of the pre-trained DM. The training is data free, and it can be divided into two alternating phases: training a diffusion model on generated samples from the generator model, and updating the generator model to produce samples that produces similar score with the pre-trained DM.
Strengths
1. The paper tackles a very interesting idea. While there are many papers targeting at improve the sampling speed of diffusion models, and multiple works have designed methods to distill diffusion models, the idea of distilling the knowledge to another generator model is, to the best of my knowledge, novel. 2. The method itself is solid and interesting. With the definition of IKL, the training objective becomes an explicit divergence minimization. 3. Results are comprehensive and strong. It shows results on both distilling diffusion models to a single step and enhancing existing GAN models. It also shows impressive performance on single step generation. The fact that it can further improve SOTA GANs is impressive.
Weaknesses
1. The presentation of the algorithm can be improved. Figure 1 isn't very informative, and a more detailed introduction should be included in the caption. Algorithm 1 isn't very clear, as it only shows the gradient formulation, without introducing how to compute it. For example, I am not very clear about how to estimate the integral in the gradient formulation. Did you just sample one random t as a very rough approximation? 2. As you mentioned, when we assume the generator output has dirac delta distribution, the formulation is essentially the same as SDS in dreamfusion. But we know that GAN is indeed a generator model with dirac delta distribution, so what is the real benefit of assuming non-dirac delta distribution for the generator? Did you actually find any benefit of training with algorithm 1 with an auxiliary DM over just training with 3.3? 3. I do not quite understand the distillation part. When doing the distillation (instead of GAN enhancement), at early stage the generator just output non-sense images, and matching the score on it can be meaningless. Is there any explanation on how it can be applied to distilling DM to a non-pretrained generator?
Questions
See above.
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
4 excellent
Presentation
2 fair
Contribution
3 good
Limitations
Limitations have been discussed.
Summary
The paper proposes a method for data free knowledge distillation from large scale diffusion models. It tries to distill the knowledge contained into a diffusion model, into another diffusion model or an explicit image generation model like a GAN. They propose a novel divergence measure obtained by integrating the KL divergence along the path of the diffusion model to match distilled model with the original instructor model. They show that the gradient update for the generator can simply be obtained as a function of the score functions as long as the generator is a differentiable function in its output. Using this the paper shows that they are able to distill knowledge from pre-trained diffusion models into generative models like GANs which enables faster inference. The paper provides empirical results to support their claims.
Strengths
1. The paper discusses an important problem of distilling pre-trained diffusion models, especially into single-shot prediction models which enables faster inference. 2. The proposed method is simple to understand and implement in practice. 3. Since the generator is a single-step predictor, it provides extremely fast inference at test time compared to standard diffusion models.
Weaknesses
1. The convergence of the proposed algorithm is not clear in the paper. The authors should elaborate as to why will the generator generate images which resemble the pre-trained diffusion model. In the absence of the generator, one can always sample random images and try to match a new diffusion model to a pre-trained model. Why will this approach not work? (One can use lesser number of reverse diffusion steps for this model) 2. This method requires training two models, which increases the space complexity of training
Questions
1. How does the generator converge to the distribution of the pre-trained diffusion model? 2. Can you reduce the number of reverse diffusion step in the first step of algorithm 1 so that its inference is comparable to the generator? Since the generator indirectly interacts with the pre-trained diffusion model. 3. How will simple score matching work, with lesser number of diffusion steps? What’s the need to introduce a generator model? Is it some form of regularization?
Rating
4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
1. Authors need to provide an elaborate discussion comparing the proposed method with its baselines for instance consistency models, and justify why the method is better compared to the baseline in terms of the theoretical convergence or inference
Summary
The paper proposes a general framework for distilling a pretrained diffusion model into an arbitrary one-shot latent-variable model. The objective is a reverse KL divergence integrated over time, and an extra score model is learned for the current model distribution. The generator model parameter and the score model parameter are updated alternatively. Empirical results were shown on distilling a diffusion model to a model of same architecture, and finetuning a pretrained GAN model.
Strengths
- This paper proposed a general framework for distilling pretrained diffusion model on a model-architecture agnostic fashion. Meaning that the distilled model doesn't have to have the same model architecture as many other distillation approaches for diffusion model. This can be useful for many potential applications such distilling 2D diffusion models for 3D. - Introducing a learnable score network for the generator model removes the need of learning an inference network as the usual variational inference requires. - Overall speaking I think this framework is nice and can potential have impact not only for 2D image distillation / refining, but also cross-modality knowledge transferring.
Weaknesses
- One piece missing and I'm eager to see is, if the implicit model is randomly initialized with an arbitrary architecture (e.g. a styleGAN generator), would this approach work? The experiments show that it works well if the implicit model is either initialized from a pretrained DM or a pretrained GAN model. It makes me doubt if this approach works only when the generator is already well-initialized. Adding this piece will make the paper much more stronger IMO. - The approach involves training two models jointly which inevitably introduces the issue of mismatching between two model classes and potentially results in instable training. I'd like to see more analysis towards this issue.
Questions
- Fig.1 is pretty low-res. Consider improving the figure quality. - There's a concurrent work (not a criticism!): ProlificDreamer: High-Fidelity and Diverse Text-to-3D Generation with Variational Score Distillation, which proposed essentially the same training objective in the context of text-to-3D and justified the objective from W gradient flow. Would be nice to compare with and discuss this work.
Rating
8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.
Confidence
5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
Yes.
Thank you for the clarifications. I read the rebuttal and it answers my questions.
Thanks for the rebuttal
Thanks for the rebuttal and providing additional results. My questions Q2-Q4 are well addressed by the authors. In terms of Q1, It is a bit surprising to me that this approach can work well with a randomly initialized generator model, as we know the reverse KL in Eq (3.1) is notorious for mode-seeking / mode-collapsing behavior, while seems not the case as shown in new Fig 2 right. Could you comment more on why in your framework, the mode-seeking bad property is avoided? Any special treatment in terms of hyperparamters, training procedure or modeling is required to make it work? Also, could you report numerical evaluations in that case?
Discussion on mode-seeking issue of random initialized generator: good intuition
Thank you for your response. We briefly summarize your question as Q1 and answer it in A1. **Q1**. More comment on the mode-seeking issue, training tricks, and numerical evaluation of training from scratch. **A1**: Thank you for your question, we are also interested in the training-from-scratch setting of Diff-Instruct. We agree that reverse KL can potentially lead to a mode-seeking issue, which is a well-known fact in the literature. In our train-from-scratch experiment, we find several "tricks" to help the training. (1) **noise level truncation**. We found that when training from scratch, direct use of all noise levels for Diff-Instruct will lead to failure (e.g. the generator fails to generate meaningful contents) because if the generator and teacher model is too dissimilar (as in training from scratch setting), score functions at small noise levels can be uninformative to “pull” the generator to teacher model. So we truncate the noise level by a minimal value (2.0 to 5.0) and find this helps. However, we think that truncating noise levels can lead to performance degradation because some noise levels are not used for Diff-Instruct. It is possible that more complex noise annealing tricks can further help but we leave the exploration of them in the revision. (2) **larger learning rate for $s_\phi$**. We find it better to use a larger learning rate for the $s_\phi$ than for the generator. Our intuition behind this is to let the generator update a little bit slower so that the diffusion model can be ready to follow the generator. **Analysis**: Overall, training from scratch works but faces a performance drop with an FID of 27.5 and an Inception Score of 9.88, which is significantly worse than the optimal initialized generator (sigma=5.0, FID=4.1). We appreciate your keen intuition on mode-seeking behavior for random-initialized generation. Figure 2 in our added rebuttal page shows some mode-seeking behaviors. For instance, all generated horses seem to turn their head to the right, and the color and shape of generated car images seem not diverse enough. This indicates that randomly initialized generators do demonstrate the **mode-seeking issue**. Therefore, it highlights the motivation and necessity of our contribution of **applying Tweedie's formula on the teacher diffusion model to obtain a better initialization** of the generator in the main text. We think this initialization technique is novel and easy to use for diffusion distillation, and it can potentially be helpful in other domains of knowledge transfer. In conclusion, incorporating knowledge from the teacher model through initialization shows better performance and more stable training than baseline (training from scratch). Developing new techniques, such as introducing a likelihood-ratio estimator, to improve the training from scratch may be an interesting research direction in the future. We thank you for your valuable suggestions which help to improve our work. We hope our answers have resolved all your concerns. If you still have any questions, please let us know.
Thank the authors for clarifying the training details of the training-from-scratch setting, and discussing about the mode-seeking issue. I do agree that this is an interesting and promising research direction to explore, and the results on distilling a well-initialized diffusion model are already pretty convincing. I'd like to increase my score to 8 to support the acceptance of this paper. Not necessary, but it'd be great if you could expand the discussion on potential new techniques on improving training-from-scratch a bit.
Thank you for liking our work. More discussions on training techniques.
We are glad that we have addressed your concerns. Thank you for the advice and we will incorporate them in our revision. We totally agree that training from scratch is an exciting direction. If we can avoid the mode-seeking behavior, training from scratch may potentially lead to stronger performance than the "pretraining + distillation" method. We think some inspiration can be drawn from two other domains: adversarial training (GAN) and text-to-3D generation. (1) The mode-seeking (or mode-hopping) behavior is a well-known issue for GAN. The community has developed many techniques to address the issue from either theoretical or empirical perspectives. For instance, some divergences, such as the one used in non-saturate GAN [1], have shown promising progress in alleviating mode-seeking behaviors. It would be interesting to explore whether they can help stabilize our Diff-Instruct. Besides, it is well-known that proper regularization of GAN's discriminator (the density ratio estimator) is crucial for the success of GAN training. There is a possibility that some forms of regularization of $s_\phi$ in Diff-Instruct could benefit the training from scratch (currently we do not assign regularization to $s_\phi$). (2) The text-to-3D content creation is another exciting scenario of knowledge transferring from teacher diffusion models to one-step generators. As we have shown in our main text, the Diff-Instruct has laid a solid mathematical foundation for many text-to-3D creation algorithms (such as DreamFusion and its variants) by minimizing IKL divergence. Therefore, techniques from text-to-3D generation algorithms, e.g., annealing the noise schedules of distillation, are potentially useful in other knowledge-transferring applications such as diffusion distillation and GAN improvement. [1] Generative Adversarial Networks [2] Towards Principled Methods for Training Generative Adversarial Networks
Thanks for you answers. Now table 1 and table 2 are better and clear. When comparing methods, it is better to fix other factors. For example, in table 1, it is hard to compare ReFlow (+distill)∗ and Diff-Instruct, because they use different Teachers (which are trained in different way and probabily on different data) and different data to "distill". So in table 1 and table 2, we can safely compare Diff-Instruct with PD, but not others. Table 3 and table 4, Cifar10 conditional class sampling performance, faces the same problem: some models use Cifar10 during training (e.g. Stylegan2+ADA+Tune + DI), some don't (e.g. Stylegan2), some may even use different pretraining dataset (e.g. BigGAN, FQ-GAN). I suggest author to provide a clear setting for the methods in Table 3 and Table 4.
Re-organize Table 3 and Table 4 to make them more clear.
Thank you for your response. Below, we summarize and answer your questions in the following paragraphs. **Q1**. More detailed setting of GAN improvement experiment (Table 3 and Table 4). **A1**. First, let us make some clarifications. By our understanding, we assume that you may have mistaken the "BigGAN+Tune" in Table 3 for our results and thought that we used a BigGAN pre-trained on larger datasets and then finetuned it on CIFAR10. We want to clarify that we have not experimented with BigGAN and the reported FID for "BigGAN + Tune" is directly taken from Table (b) in Figure 11 of the StyleGAN2-ADA paper [1], where they termed "BigGAN+Tuning". StyleGAN-ADA [1] is our main reference paper, and the FID values of ProGAN, AutoGAN, MultiHinge, FQ-GAN, and Stylegan2+ADA+Tune are also taken from the StyleGAN2-Ada paper. **So it is a misunderstanding that we use either different datasets or different methods**. Next, for the GAN improvement experiment itself, it is worth emphasizing our goal. We aim to show that by incorporating knowledge of a strong teacher diffusion model, our Diff-Instruct can further improve generators pre-trained with adversarial training (i.e. GAN). Therefore, taking a pre-trained GAN generator as a baseline, we evaluate whether Diff-Instruct can further improve its generative performance. To make our presentation more clear, we reorganized Table 3 and Table 4 to **Table 3** and **Table 4** to highlight the pre-training and post-training methods. **Table3**. Class-conditional sample quality on CIFAR10 with GAN generators. †Models that we implemented. DI represents for Diff-Instruct. | Method | Pre-trained Dataset | Pre-trained Method | Post-trained Method | FID | IS | | :------: | :------: | :------: | :------: | :------: | :------: | | BigGAN | CIFAR10-Cond | GAN | | 14.73 | 9.22 | | BigGAN+Tune | CIFAR10-Cond | GAN | | 8.47 | 9.07 ± 0.13 | | MultiHinge | CIFAR10-Cond | GAN | | 6.40 | 9.58 ± 0.09 | | FQ-GAN | CIFAR10-Cond | GAN | | 5.59 | 8.48 ± 0.09 | | Stylegan2 | CIFAR10-Cond | GAN | | 6.96 | 9.53 ± 0.06 | | Stylegan2† | CIFAR10-Cond | GAN | | 7.01 | 9.23 ± 0.07 | |**Stylegan2† + DI** | CIFAR10-Cond | GAN | DI (cond-EDM Teacher) | 6.62 | 9.40 ± 0.06 | | Stylegan2+ADA | CIFAR10-Cond | GAN | | 3.49 |**10.24 ± 0.07**| | Stylegan2+ADA+Tune | CIFAR10-Cond | GAN | | 2.42 | 10.14 ± 0.09 | |**Stylegan2+ADA+Tune + DI**| CIFAR10-Cond | GAN | DI (cond-EDM Teacher) | **2.27** | 10.11 ± 0.10 |
Re-organize Table 3 and Table 4 to make them more clear. (2)
**Table4**. Unconditional sample quality on CIFAR10 with GAN generators. †Models that we implemented. DI represents for Diff-Instruct. | Method | Pre-trained Dataset | Pre-trained Method | Post-trained Method | FID | IS | | :------: | :------: | :------: | :------: | :------: | :------: | | SNGAN | CIFAR10-Uncond | GAN | | 21.70 | 8.22 | | ProGAN | CIFAR10-Uncond | GAN | | 15.52 | 8.56 ± 0.06 | | AutoGAN | CIFAR10-Uncond | GAN | | 12.42 | 8.55 ± 0.10 | | SNGAN+DGflow | CIFAR10-Uncond | GAN | | 9.35 | 9.62 | | TransGAN | CIFAR10-Uncond | GAN | | 9.02 | 9.26 | | Stylegan2 | CIFAR10-Uncond | GAN | | 8.32 | 9.21 ± 0.09 | | Stylegan2† | CIFAR10-Uncond | GAN | | 8.21 | 9.09 ± 0.09 | |**Stylegan2† + DI** | CIFAR10-Uncond | GAN | DI (Uncond-EDM Teacher) | 7.56 | 9.16 ± 0.09 | | Stylegan2+ADA | CIFAR10-Uncond | GAN | | 5.33 |**10.02 ± 0.07**| | Stylegan2+ADA+Tune | CIFAR10-Uncond | GAN | | 2.92 | 9.83 ± 0.04 | |**Stylegan2+ADA+Tune + DI**| CIFAR10-Uncond | GAN | DI (Uncond-EDM Teacher) |**2.71** | 9.86 ± 0.04 | **Analysis**: As is shown in **Table3** and **Table4**, all generators are pre-trained on the same dataset (CIFAR10 cond/uncond), and **the StyleGAN2 generator performs the best among all generators with adversarial training**. So we decide to further improve the StyleGAN2 generator with Diff-Instruct as post-training to outperform adversarial training by incorporating knowledge of a strong EDM teacher diffusion model. It is worth noting that Diff-Instruct is **data-free**, so the post-training of the pre-trained generator does not need real data (we only need a teacher model). Besides, the Diff-Instruct does not need to sample from the teacher model. This makes Diff-Instruct more efficient than a simple baseline (generates samples from the teacher model and then trains the generator with adversarial training) because generating samples from diffusion models is costly, therefore the simple baseline is costly and can hardly be scaled to large-scale tasks such as text-to-image generation. The results in **Table 3** and **Table 4** show that Diff-Instruct can significantly improve the generators that are trained with adversarial training. For instance, the FID of StyleGAN2+ADA+Tune was improved from 2.42 (which is a previous SOTA for StyleGAN2 generator) to 2.27 for conditional generation, and from 2.92 to 2.71 for unconditional generation. In conclusion, in both the diffusion distillation and GAN improvement, the Diff-Instruct has shown its strong empirical performance over baseline methods (models). We really thank you for your valuable suggestions on writing and representation, which help us to improve our work. We hope our answers have resolved all your concerns. If you still have any questions, please let us know. [1] Training Generative Adversarial Networks with Limited Data
Thank you for your constructive suggestions.
We are glad that we have addressed your concerns. Thank you for the advice and we will incorporate them in our revision.
Thank your for your reviews, we are happy to provide clarifications.
We sincerely hope that our responses have adequately addressed the concerns you raised in your review. For any unresolved concerns or additional questions, please do not hesitate to let us know. We would be happy to provide further clarification and address any remaining issues.
Decision
Accept (poster)