Paper: Physics Informed Distillation for Diffusion Models (ICLR-2024)

Reviewer comment (Novelty):
The literature of PINNs on ODE systems is very mature, and the paper did not propose innovations on how to better perform PINNs on ODE systems. Rather, it is simply applying the PINN techniques to ODE systems, limiting the contribution to developing a technique for distilling score based generative models.

Author response:
Probability flow ODEs can be viewed as a special type of ODE where the Lipschitz constant of such ODE is exploding near the origin[2]. To train a PINN for such an ODE system, 2 key innovations were used. Namely, the shifting of variable from Equation 8 to 9 was used to prevent the explosion of the loss value from one ieration to another when points near the singularity are sampled. Secondly, we adopt a time discretization similar to that in EDM[3]. In black box ODE solvers such as Euler Method, their truncation error is related with the Lipschitz constant of such ODE systems. As such, this error increases with the growing Lipschitz constant near the singularity point. To account for this, EDM time discretization was used where more points near the singularity was sampled. To the best of our knowledge, solving ODE systems on intervals near such Lipschitz singularity have never been considered in PINN literature. As such, we believe that our approach could also contribute to the PINN literature when such ODE systems are encountered.

Reviewer comment (Evaluation):
The paper does not achieve state of the art performance. As shown in Table 1, 2, 3, and Figure 7, the existing methods consistency model and EDM (which is heavily referred to when developing this paper) perform much better than the proposed method.

Author response:
With regards to the performance of PID, while our approach does not beat the current SOTA method, it is important to note that our method is able to achieve competitive performance, surpassing several recent works while resolving issues inherent in prior methodologies. Note that EDM model is a teacher Diffusion Model used in many recent distillation works as seen in Table 1 and 2, thus exhibiting a significantly slower sampling speed. Interestingly, our method exhibits a predictable trend with respect to hyperparameters. As demonstrated in Figure 3 in our main paper, PID performance consistently improves with increasing discretization, saturating when the discretization number is sufficiently large. In contrast, CM[4] is sensitive to the choice of discretization number, where high or low values from the optimal point results in performance deterioration. As such, our approach does not need to search for hyperparameters such as discretization numbers from dataset to dataset allowing us to set it at a sufficiently high number.

Reviewer comment (Experiments):
It would be great if the authors can demonstrate qualitatively and quantitatively how the proposed method is better than EDM and/or consistency model (such as distillation efficiency, inference time, etc.).

Author response:
Thank you for your clarification on this matter. I would like to elaborate on the distinction between the discretization approaches employed in Neural ODEs and our proposed method. In Neural ODE discretization, the utilisation of Neural ODE as a vector field involves sampling more steps, akin to diffusion. Thus, increasing discretization has been properly explored and has been shown to improve performance. However, in the context of knowledge distillation approaches for diffusion models, such as Consistency Models, it is noteworthy that increased discretization does not necessarily yield enhanced performance and necessitates careful tuning. In our approach, we learn the entire trajectory function, with the vector field being implicitly acquired. As such, increased discretization corresponds to querying points on this trajectory function, incurring no additional cost during training or inference. As our work is in the field of single step sampling, sampling efficiency is of great importance as iterative methods would defeat the purpose of distillation. Therefore, in the context of fast sampling methods, the comparison between discretization in the iterative Neural ODE and discretization in our method cannot be made on an apples-to-apples basis.

Reviewer comment (Theory):
Concern with the validity of the assumption that the score model is Lipschitz continuous, including how often this is the case and, if this property is enforced, how much it will hurt the performance.

Author response:
The main use of the Lipschitz condition assumption is to assume the existence and uniqueness of an ODE solution modelled by the diffusion model as well as the convergence of numerical ODE solvers. More specifically, the global truncation error proofs of many ODE solvers used in Diffusion literature involves the use of such Lipschitz continuity assumptions[1]. As such, this assumption is present in many diffusion model literature either explicitly as seen in Theorem 3.2 of DPM Solver[2] or implicitly when using ODE solvers to generate solutions where convergence of such solvers are assumed implicitly. Note that unlike GAN approaches, this assumption is not strong as we only require the model to be Lipschitz continuous to have valid ODE solutions and do not need to restrict the Lipschitz constant of such models.

Reviewer comment (Experiments):
Similarly, a first-order numerical approximation of the residual loss is proposed for the sake of efficient training, but no ablations are performed as to how much this discretization affects the performance.

Author response:
In section 4.2, we mention a study on PINNs that uses exact vs numerical approximations of the gradient during training[1]. One key observation in that paper is that with exact gradients, the model is vulnerable to converging to unphysical solutions, having low loss but modeling a significantly different solution than the exact solution. As such, the authors of that paper argue that leveraging local points via numerical approximations can alleviate this issue. Similarly, we observe a similar situation where using exact gradients performs poorly in practice and resort to using numerical gradients instead.

Reviewer comment (Writing):
A bit more background about PINNs, especially the soft enforcement of boundary conditions, could be helpful to better motivate the authors' choice of model parameterization.

Author response:
In section 3.2, we explain that PINN methods often either use soft or hard conditions to satisfy the boundary condition. However, as our method only exclusively uses hard conditions through skip connections, we did not include detailed explanations on how soft conditions in PINNs are enforced.

Reviewer comment (Theory):
Is there a reason why a combination of soft and hard enforcement of boundary conditions is necessary?

Author response:
In our approach, we only use hard conditions to satisfy the boundary conditions. In detail, we use skip connections to automatically satisfy the boundary conditions and thus do not explicitly need to train boundary conditions through soft conditioning. It is important to note however that despite not training for boundary conditions, the gradients at the boundary are still trained. To elaborate, we train our student model such that the gradient at the boundary corresponds well with the probability flow ODE.

Reviewer comment (Novelty):
Once a first-order discretization scheme for the residual loss is chosen, the student model training looks very similar in form to existing distillation techniques (1, 2). How is PID related to these techniques, e.g. can they be described as a special case of PID given a specific choice of time discretization of the probability flow ODE?

Author response:
While after the first-order discretization, the loss looks somewhat similar, one fundamental difference is present. Notably, in existing distillation techniques [2,3], the teacher model imparts knowledge to the student model without direct interaction. Specifically, the teacher model is fed with noised data, and the output of the student model is not incorporated back into the teacher model. In contrast, our approach, due to formulating it in a physics informed fashion, utilizes the output of the student model as the input for the teacher model instead of feeding in data directly. This distinction allows our method to learn in a data-free manner, as commonly observed in Physics-Informed Neural Network (PINN) literature. However, due to this difference, PID and existing distillation approaches unfortunately cannot be connected through considering different discretization schemes.

Reviewer comment (Experiments):
How does the first-order approximation of the residual error loss compare to training directly using the ODE (backpropagation), or with using a higher-order approximation?

Author response:
In section 4.2, we mention that in PINN literature, it has been observed that exact gradients using forward mode automatic differentiation leads to convergence with unphysical solutions[1]. This is primarily due to the fact that automatic differentiation relies only on singular point and thus can cause a form of overfitting in the PINN training paradigm while numerical methods rely on local points which do not exhibit this issue. Similarly, we observe that exact gradient for training performs poorly and therefore resort to first order gradient approximations. As for higher order numerical approximations, higher order approximations such as central difference method provides a further performance gain as it better approximates the gradients while still using local evaluation points. We present this results in the table below.

| Method  | FID | 
| -------- | -------|
| PID (1st order)     |  3.92     |
| PID (2nd order - Central difference)    |   3.68    |

Reviewer comment (Novelty):
Lack of an important related work: BOOT[1]. The proposed method seems almost exactly the same as BOOT, because they both distill the integral from time $T$ to time $t$, with the same integral and numerical differential method. Please compare with BOOT in details and discuss more about the own contributions.

Author response:
We apologize sincerely for not including BOOT as we were previously unaware of its existence. We have since added it in the revised manuscript in Table 1 and 2 for comparison. It is notable that both our proposed approach and BOOT[1] represent concurrent works in the field of distillation sampling with diffusion models. We present the similarities and differences in detail as follows. 

#### Similarities:

- **Same starting motivation**: They share the similarity in the fact that we both realize that distillation can be achieved by reducing the loss associated with the differential equation of the probability flow ODE system.

- **Same use of first order numerical gradient approximations**: In both our methods, we resort to using 1st order numerical gradient approximations instead of forward-mode auto differentiation. It is important to highlight that, unlike BOOT, our decision is not solely based on ease of implementation. Rather, our motivation stems from the insights in the Physics-Informed Neural Network (PINN) literature. Numerical gradient approximations have been demonstrated to prevent convergence to unphysical solutions in ODE systems, in contrast to auto-differentiation approaches[2] by relying on local points rather than a single point, as discussed in Section 4.2 of our main paper.

#### Differences:

- **Teacher Model Training: Signal ODE Distillation vs. EDM-based Approach**: A key distinction lies in the absence of signal ODE for distillation in our method. Notably, their approach involves distillation using the signal ODE system instead of the conventional ODE system. This difference holds significance as it leads them to train a new signal prediction teacher model (refer to Table 3 in their appendix, where a signal prediction model is employed as the teacher), rather than relying on pre-trained teacher models, as in our case. As such, they incur additional training cost associated with such teacher models.

- **Divergent Boundary Approaches: Our Method vs. BOOT**: A key discrepancy that arises between the 2 methods is in our approach towards the boundary conditions. While BOOT opts to solve the signal ODE, they encounter challenges in easily satisfying boundary conditions through simple skip connections. Consequently, they resort to training with soft boundary conditions. In contrast, due to our use of the base ODE system modeled by the diffusion model, such boundary conditions can easily be satisfied through the use of skip connections. The benefit of this is twofold. Firstly, the need for soft conditions in BOOT incurs additional cost during training through the introduction of the boundary condition loss. In addition, due to the use of soft conditions, such boundary conditions are not always satisfied. As a result, reducing the loss associated with the differential equation may not lead to valid solutions. We observe this in the revised Table 1 and 2 of our manuscript where we outperform BOOT in both datasets. 

In summary, while the beginning motivations of both ours and BOOT are similar, notable differences in the training schemes are evident that provides the following benefit: 
1) We are able to use pretrained teacher models for distillation while BOOT trains its own signal prediction teacher model.
2) We do not need to train for boundary conditions, while BOOT has to due to its soft conditions, incurring a slight additional training cost.
3) Our boundary condition is always satisfied, unlike BOOT, causing the reduction of the PINN loss to provide valid solutions to the probability flow ODE, resulting in improved FID values across all datasets.

Reviewer comment (Evaluation):
The results in Table 1 is unfair. Some of the results are based on the checkpoint of the VPSDE in ScoreSDE[2], but some of the results are based on the checkpoint of EDM[3]. The authors should at least split them into two parts.

Author response:
Thank you for meticulously reviewing our paper. We have partitioned the result comparison in the main result table based on the distillation teacher model in our revised manuscript. Additionally, we have incorporated the results from DPM-Solver-fast into the main result table and fixed the typos accordingly.

Reviewer comment (Experiments):
The results of DPM-Solver in Table 1 seems not the best results. e.g., please see Table 1 in [4], where it involves "dpm-solver-fast".

Author response:
Thank you for meticulously reviewing our paper. We have partitioned the result comparison in the main result table based on the distillation teacher model in our revised manuscript. Additionally, we have incorporated the results from DPM-Solver-fast into the main result table and fixed the typos accordingly.

Reviewer comment (Writing):
A small typo: please use $\mathcal{O}(\Delta t)$ instead of $\mathcal{O}\Delta t$.

Author response:
Thank you for meticulously reviewing our paper. We have partitioned the result comparison in the main result table based on the distillation teacher model in our revised manuscript. Additionally, we have incorporated the results from DPM-Solver-fast into the main result table and fixed the typos accordingly.

Reviewer comment (Novelty):
Although the comparison between BOOT and the proposed method is fair after the rebuttal, there are still many similarities because BOOT is the first work which combines the PINN loss with diffusion distillation.

Author response:
We apologize sincerely for not including BOOT as we were previously unaware of its existence. We have since added it in the revised manuscript in Table 1 and 2 for comparison. It is notable that both our proposed approach and BOOT[1] represent concurrent works in the field of distillation sampling with diffusion models. We present the similarities and differences in detail as follows. 

#### Similarities:

- **Same starting motivation**: They share the similarity in the fact that we both realize that distillation can be achieved by reducing the loss associated with the differential equation of the probability flow ODE system.

- **Same use of first order numerical gradient approximations**: In both our methods, we resort to using 1st order numerical gradient approximations instead of forward-mode auto differentiation. It is important to highlight that, unlike BOOT, our decision is not solely based on ease of implementation. Rather, our motivation stems from the insights in the Physics-Informed Neural Network (PINN) literature. Numerical gradient approximations have been demonstrated to prevent convergence to unphysical solutions in ODE systems, in contrast to auto-differentiation approaches[2] by relying on local points rather than a single point, as discussed in Section 4.2 of our main paper.

#### Differences:

- **Teacher Model Training: Signal ODE Distillation vs. EDM-based Approach**: A key distinction lies in the absence of signal ODE for distillation in our method. Notably, their approach involves distillation using the signal ODE system instead of the conventional ODE system. This difference holds significance as it leads them to train a new signal prediction teacher model (refer to Table 3 in their appendix, where a signal prediction model is employed as the teacher), rather than relying on pre-trained teacher models, as in our case. As such, they incur additional training cost associated with such teacher models.

- **Divergent Boundary Approaches: Our Method vs. BOOT**: A key discrepancy that arises between the 2 methods is in our approach towards the boundary conditions. While BOOT opts to solve the signal ODE, they encounter challenges in easily satisfying boundary conditions through simple skip connections. Consequently, they resort to training with soft boundary conditions. In contrast, due to our use of the base ODE system modeled by the diffusion model, such boundary conditions can easily be satisfied through the use of skip connections. The benefit of this is twofold. Firstly, the need for soft conditions in BOOT incurs additional cost during training through the introduction of the boundary condition loss. In addition, due to the use of soft conditions, such boundary conditions are not always satisfied. As a result, reducing the loss associated with the differential equation may not lead to valid solutions. We observe this in the revised Table 1 and 2 of our manuscript where we outperform BOOT in both datasets. 

In summary, while the beginning motivations of both ours and BOOT are similar, notable differences in the training schemes are evident that provides the following benefit: 
1) We are able to use pretrained teacher models for distillation while BOOT trains its own signal prediction teacher model.
2) We do not need to train for boundary conditions, while BOOT has to due to its soft conditions, incurring a slight additional training cost.
3) Our boundary condition is always satisfied, unlike BOOT, causing the reduction of the PINN loss to provide valid solutions to the probability flow ODE, resulting in improved FID values across all datasets.

Reviewer comment (Novelty):
Although they can be considered as concurrent work, the two works are too similar, and the author needs to discuss them more in detail, instead of only comparing the FID results.

Author response:
Thank you for your kind response. We would like to draw your attention to a point regarding the ICLR Reviewer guidelines. In it, works only in peer reviewed conferences such as online conferences and workshops within 4 months of submission are considered as concurrent work. As the paper BOOT [1] was published at the ICML workshop in June 20, this places it within the four-month timeframe leading up to our submission on September 28. It is our understanding that, in accordance with this timeline, our work should be regarded as concurrent. We kindly request the reviewer's consideration of this matter and thank you in advance for your understanding.

Reviewer comment (Theory):
Equation 9 and 8 are equivalent up to a scaling factor for L2 metric, but not for arbitrary distance metric such as LPIPS, which is used for the main results. Changing $L_{PINN}$ from equation 8 to 9 will change the loss landscape. However, this step is not justified or explained in the paper. Why not use the original PINN loss given by equation 8?

Author response:
We apologize for the lack of clarification when moving from Equation 8 to 9. We have added a detailed explanation in the revised manuscript. The primary motivation for this change is the Lipschitz explosion of the probability flow ODE, where the ODE system goes to infinity for small time values. As training in physics-informed fashion trains to match the gradients of the network to the values in the ODE system, directly training with an exploding system will lead to poor performance. To resolve this issue, we transform to a stabler loss system where the values do not explode near this singularity point.

Reviewer comment (Theory):
The authors choose LPIPS metric in the paragraph before Theorem 1. However, Theorem 1 is fully based on L2 metric, which is a bit confusing. I do not see how Theorem 1 can extend to LPIPS metric.

Author response:
In Theorem 1, we assumed that the metric used was a valid distance metric, thus possessing specific properties. As such, the proof in Theorem 1 is not specific to the L2 metric but rather any valid distance metric. It is important to note that despite the distance metric used, the truncation error is still expressed in L2. This is not due to the metric used for training but rather the assumption of the Lipschitz continuity of the Diffusion Model about the L2 metric.

Reviewer comment (Experiments):
Theorem 1 shows that PID will be equivalent to the Euler method with the same number of discretization steps N if the PID loss is zero. Can you also report the FID of the corresponding Euler method with the same $N$? This may help us understand the underlying gap between the learned model and the actual Euler method.

Author response:
The FID value of EDM[2] when sampling with an Euler solver in 250 time steps is given in the following table. Additionally, we have also incorporated this table in the updated manuscript.

| Sampler        | CIFAR     |  ImageNet | 
| --------       | -------   | -------   |
| EDM - Base     |   2.04    |  2.44     |
| EDM - Euler   |   2.10    |  2.41     |

Reviewer comment (Experiments):
The central difference method is more accurate than the forward difference that is used in the paper. Would it be beneficial to use the central difference for PID?

Author response:
The table below presents the results of using the central difference approach instead of the 1st order approach proposed in the main paper. Indeed, since second-order gradient approximations such as major difference methods provide better gradient approximations, this leads to a further gain in performance.

| Method  | FID | 
| -------- | -------|
| PID (1st order)     |  3.92     |
| PID (2nd order - Central difference)    |   3.68    |