Paper: Self-distillation for diffusion models (ICLR-2024)

Reviewer comment (Experiments):
Advantages of self-distillation while training diffusion models from scratch is unclear and additional results are needed to demonstrate its benefits while training from scratch. Figure 3 compares FID scores for distilled and the regular diffusion model while training for scratch only for the first 30k-60K steps. These FID scores are quite high for DDIM and longer training is needed. For instance, on CelebA, the original DDIM (with 10 sampling steps) achieves FID of 17.33 while in Figure 3, the lowest FID is around 100. Thus longer training will help improve FID scores, however it is unclear whether the gains seen in these plots due to self-distillation will continue for the entire duration of training, and whether the final FID scores of model that is trained from scratch with self-distillation is strictly better than the model that is trained without any self-distillation.

Author response:
* Your points regarding the high FID scores of distillation during training are justified. Given the comparatively pilot-like nature of DSD during training, previous FID scores were calculated with a relatively low number of samples generated for score calculation. We have re-run the experiments (with more epochs) and used our time since the last submission to optimize the hyperparameters for the training. Given the improved convergence of training overall, our latest results now demonstrate more clearly the effect of additional direct self-distillation during the standard process.

Reviewer comment (Evaluation):
The primary motivation of DSD is to use fewer parameter updates while achieving comparable image quality with progressive distillation, however the current FID scores in Table 2 are high across the datasets. Further, DSD does not result in high quality images, in terms of FID, in few step sampling (2-4 steps).

Author response:
* We feel that our relatively underwhelming scores for the DSD applications were at least in part due to the increase in generated images for calculation. For example, our imageNet results were calculated over 30K samples, against 5K of the existing PD variant. Lowering our samples to match the 5K shows Inception Scores similar to those of PD. We have therefore opted to instead use the evaluation code provided by the baseline models’ papers, along with the decision to calculate scores over 5K samples. Current scores now lie within the vicinity of those provided by the baseline papers.

Reviewer comment (Experiments):
In order for someone to strictly prefer DSD over PD, one would have to show at least one of the following: (1) DSD can achieve comparable or better FID than PD while using fewer parameter updates, in case we use fine-tuning DSD; (2) Using DSD while training from scratch simultaneously improves FID and reduces sampling steps than training a model from scratch and later doing PD. Currently, the paper does not include sufficient results to show either (1) or (2).

Author response:
* Table 2 represents only the results of the distillation of pre-trained models, not the results of distillation during training. To compare our distillation during training against any PD implementations would likely require a computational budget beyond our scope, given their resolutions of $256^2$ compared to our train-distilled $64^2$. However, implementing DSD in combination with PD to potentially achieve high quality, low DDIM samples more efficiently, is an approach our computational budget would allow as future research.

Reviewer comment (Experiments):
While DSD seems to outperform TSD at fewer parameter updates (from Table 2), it is unclear if DSD can match the FID of PD with more parameter updates. Currently, the models in the paper use 4K-5K parameter updates, and DSD quickly deteriorates in terms of FID while using 2 and 4 step sampling. The results would be much more convincing if FID scores could be matched with PD with slightly more parameter updates.

Author response:
* If the computation time allows for it, we will also attempt longer DSD runs to try and match the quality found at fewer sampling steps. We are currently limited in our ability to fine-tune the DSD process, which in part explains our focus on showing the working principle to be efficient and simple, rather than directly aiming to compete with matured existing methods. On a side-note, your idea of having DSD lower the sampling steps initially before introducing PD is one we have considered as well. We feel this idea is perhaps too much to introduce at this stage of the revision, but we will include it in our ideas for future work.

Reviewer comment (Reproducibility):
Does DSD maintain an exponentially moving average (EMA) copy of the model weights during training/fine-tuning?

Author response:
* We do indeed keep the EMA weights during distillation and train-distillation. Sampling for FID/IS calculation, as well as the samples shown in the paper, have been done with the EMA weights loaded.

Reviewer comment (Presentation):
Clarification is needed for Figure 2, specifically the X axis. The understanding is that this is \( N_i \), but the number of denoising steps along the X axis appears small. Shouldn't there be large step sizes \( N_i \) in self-distillation initially?

Author response:
* Regarding Figure 2, we have improved the explanation of the graph, which has received a dedicated section in the Appendix. This graph essentially shows during the entire distillation process, how often a specific $x^{’}(z_t)$ (for instance, the 3rd sampling step) was distilled into $x^{‘’}(z_{t-1})$ (then the 4th sampling step, in reverse order). For this reason, the x-axis spans 0-63. A conceptually simple example is DSDN, which loops from step 0-63 and updates each step against the consecutive sampling step’s output. Therefore, at each sampling step, an equal number of parameter updates has taken place. If we were to halve the total sampling steps at some point as done in DSDI, higher sampling steps (>32) are no longer seen, lowering their respective y-axis frequencies. Please let us know if this clarifies the graph somewhat.

Reviewer comment (Writing):
The readability of the paper would greatly improve if the notation is made more consistent. For example:  
- The loss in Eq. 1 uses the notation of \(\alpha \cdot \|\| \hat{z}_{t-1} - z_{t-1} \|\|^2 + \beta \cdot \|\| \hat{z}_{t-1} - \hat{z}_{t-2} \|\|^2\) but previous sections use \(\|\| \hat{x}_{t-1} - \hat{x}_{t-2} \|\|^2\) as loss for distillation. Further, the usual choice of loss objective for training diffusion models is \(\|\| \hat{x}_{t-1} - x_{gt} \|\|^2\) (\(x_{gt}\) is ground truth image) or \(\|\| \hat{\epsilon}_{t-1} - \epsilon \|\|^2\). Why was \(\|\| \hat{z}_{t-1} - z_{t-1} \|\|^2\) chosen instead for training (this is not about the second self-distillation loss term)?
- Minor: The results section uses S (see header for Table 1) without defining it in previous sections. It is presumed that this denotes the number of sampling steps in distillation.

Author response:
* Finally, we have included your improvements regarding notation.

Reviewer comment (Writing):
Clarity of writing. The manuscript suffers from significant issues with clarity, which impedes the reader's understanding of the methodological details. For example, on page 2, the latent variables  $\mathbf{z}_t$ are introduced without adequate context, leading to confusion. It is only in Section 3 that the reader learns these notations are adopted from Progressive Distillation (PD). Furthermore, the manuscript lacks a necessary introduction to PD and fails to clearly articulate the distinctions between Teacher-Student Distillation (TSD) and PD. The writing quality would benefit from thorough proofreading and reorganization.

Author response:
We will also aim to simplify the document in that the concept of PD naturally flows into our approach, which indeed is not the case currently. We will also try to reformulate TSD in an attempt to better distinguish this approach from both DSD and PD, since TSD is neither Direct, or dataset–driven like PD.

Reviewer comment (Writing):
The complexity of the method. The proposed method entails a three-level nested loop, as presented in Algorithm 2. Although it is not technically unacceptable, without a clear introduction of the notations, the writing issue complicates the understanding of the algorithm and scheduling. The recycling of previous latent variables in the loop, coupled with the scheduling, together add to the complexity. Despite eliminating the static teacher, the method does not appear to be a conceptual simplification of PD.

Author response:
The DSD algorithm's depiction is indeed over-complicated beyond necessity, where we aimed to include every possible scenario of the different DSD implementations (and their respective update schedules) into a single formulaic representation. Instead, we have removed those parts trying to encapsulate our step-scheduling, making the DSD algorithm instead a true, generalized application of the concept.

Reviewer comment (Evaluation):
The most severe weakness of this paper is the weak performance. It is well-known that the goal of distillation is to obtain more efficient student models with comparable performance to teacher models. However, in this paper, the proposed distilled model is not even close to the teacher model (LSUN results in Table 2). I strongly recommend the author to check the results and validity before submitting.

Author response:
Dear Reviewer, thank you for your insights. 

From your assessment it seems clear that our work would benefit from some improvements regarding clarity of results. To address your point regarding the number of parameters of Teacher vs Distilled models, the DSD distilled-models are exact copies of the teacher model in terms of structure and number of parameters. Given the nature of DSD, it might in fact be in error to state that there's a teacher model, since for all DSD approaches, this is in fact the distilled model before the first parameter update (only a single model is ever sampled from). We have clarified this in the latest version of our paper.

In similar need of clarification is the definition of our baseline models. I believe that, given your reference to Table 2, you may have interpreted the "PD" column as being our teacher model. Instead, our teacher models are displayed in Table 1. The "PD" column refers to the distillation results of other, non-DSD approaches, of which the checkpoints are sadly not available to us.

The great leap in image quality of these teacher/original models compared to ours was touched upon in the Results section, but it has come to our attention that our implementation of IS calculation was rather non-forgiving, as our decision to include many more samples for score-calculation increased robustness at a significant loss of maximum scores. Instead, we have now opted to use the models’ original authors’ evaluation code, leading to scores more in line with perceived image quality overall. For instance, our ImageNet results show higher Inception scores for 8 steps compared to PD (~140-200 vs 128.63), with DSDN narrowly managing to show improvements at just 4 steps (126.61 vs 126.36). We’ve also matched the PD papers’ choice of calculating the scores over 5000 samples instead.

Your other points regarding the writing/presentation are also very welcome and have been implemented in the latest revision. Large changes, excluding any minor improvements to readability, have been highlighted in our final document, which we feel improves the clarity of our results.

Thank you very much for your time.

Reviewer comment (Experiments):
I recommend the author compare the number of parameters in the distilled model to the original model. This information would demonstrate the potential benefits of the distilled model in terms of model size and resource requirements.

Author response:
Dear Reviewer, thank you for your insights. 

From your assessment it seems clear that our work would benefit from some improvements regarding clarity of results. To address your point regarding the number of parameters of Teacher vs Distilled models, the DSD distilled-models are exact copies of the teacher model in terms of structure and number of parameters. Given the nature of DSD, it might in fact be in error to state that there's a teacher model, since for all DSD approaches, this is in fact the distilled model before the first parameter update (only a single model is ever sampled from). We have clarified this in the latest version of our paper.

In similar need of clarification is the definition of our baseline models. I believe that, given your reference to Table 2, you may have interpreted the "PD" column as being our teacher model. Instead, our teacher models are displayed in Table 1. The "PD" column refers to the distillation results of other, non-DSD approaches, of which the checkpoints are sadly not available to us.

The great leap in image quality of these teacher/original models compared to ours was touched upon in the Results section, but it has come to our attention that our implementation of IS calculation was rather non-forgiving, as our decision to include many more samples for score-calculation increased robustness at a significant loss of maximum scores. Instead, we have now opted to use the models’ original authors’ evaluation code, leading to scores more in line with perceived image quality overall. For instance, our ImageNet results show higher Inception scores for 8 steps compared to PD (~140-200 vs 128.63), with DSDN narrowly managing to show improvements at just 4 steps (126.61 vs 126.36). We’ve also matched the PD papers’ choice of calculating the scores over 5000 samples instead.

Your other points regarding the writing/presentation are also very welcome and have been implemented in the latest revision. Large changes, excluding any minor improvements to readability, have been highlighted in our final document, which we feel improves the clarity of our results.

Thank you very much for your time.