Physics Informed Distillation for Diffusion Models

We appreciate your time in reviewing our paper. Below, we address your concerns:

Questions: The assumption that the score model is Lipschitz continuous.

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.

Adressing weaknesses raised

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.

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.

[1] Butcher, John Charles. Numerical methods for ordinary differential equations. John Wiley & Sons, 2016.

[2] Yang, Zhantao, et al. "Eliminating Lipschitz Singularities in Diffusion Models." arXiv:2306.11251 (2023).

[3] Karras, Tero, et al. "Elucidating the design space of diffusion-based generative models." NeuIPS(2022): 26565-26577.

[4] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. ICML2023

Question 1: Why a combination of soft and hard enforcement of boundary conditions is necessary?

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.

Question 2: The possible connections between PID and existing distillation[2,3].

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.

Question 3: Compare the first-order approximation with forward mode auto differentiation or a higher-order approximation.

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.

Question 4: PINN-based ODE solutions vs operator learning methods such as DSNO

In DSNO[4], operator learning is defined as learning the mapping between functions. In formulating it in such a manner, operator learning in a diffusion context involves mapping the constant function of noise to the trajectory function $x_t$. To learn this mapping, DSNO requires multiple functional evaluations of the different time steps in a parallel manner instead of the recursive one in diffusion models. As such, despite being able to perform single step inference to generate images, the parallel evaluations of the different time steps in DSNO causes it to be slower than standard single step inference as seen in Table 3 of their paper. In addition, the use of multiple time evaluations per element in their batch also incurs a significantly higher cost during training than posing it in a physics informed fashion where only single point evaluations are used. It is important to note however that the numerical gradients for training in a physics informed fashion incurs extra cause although not as much as multiple parallel evaluations in DSNO. Finally, operator learning as described in DSNO requires a trajectory dataset. This can be very memory consuming and expensive to produce as also mentioned in Consistency Models[2]. In contrast, proposing it in a physics informed fashion is data-free as only noise samples are needed, alleviating this problems entirely.

[1] Pao-Hsiung Chiu, Jian Cheng Wong, Chinchun Ooi, My Ha Dao, and Yew-Soon Ong. Can-pinn: A fast physics-informed neural network based on coupled-automatic–numerical differentiation method. Computer Methods in Applied Mechanics and Engineering, 2022

[2] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. ICML2023

[3] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models." arXiv preprint arXiv:2202.00512 (2022).

[4] Zheng, Hongkai, et al. "Fast sampling of diffusion models via operator learning." International Conference on Machine Learning. PMLR, 2023.

Thank you for taking the time to read and review our paper. Please see our responses below:

Weakness 1: Motivation of the parameterization used by PID (Eq.7).

We apologize for the lack of explanation on the paramaterization used in PID. We have added a detailed reasoning of this scheme in Appendix A.2 of the revised manuscript.

Weakness 2: Affect of numerical differentiation on performance.

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.

Weakness 3: Background for the soft enforcement of boundary conditions

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.

[1] Pao-Hsiung Chiu, Jian Cheng Wong, Chinchun Ooi, My Ha Dao, and Yew-Soon Ong. Can-pinn: A fast physics-informed neural network based on coupled-automatic–numerical differentiation method. Computer Methods in Applied Mechanics and Engineering, 2022

Thank you for your encouraging response. I would like to clarify that our submission, as it is within four months of the BOOT [1] paper's publication at the ICML workshop in June 20, can be regarded as concurrent work, which aligns with the ICLR Reviewer guidelines. We kindly request the reviewer's consideration of this timing and sincerely appreciate your understanding in this regard.

[1] Gu, Jiatao, et al. "Boot: Data-free distillation of denoising diffusion models with bootstrapping." ICML 2023 Workshop on Structured Probabilistic Inference {&} Generative Modeling. 2023.

We thank the reviewer for their detailed feedback. We address your concerns as below:

Weakness 2: Changing $L_{PINN}$ loss from Equation 9 to 8 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?

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.

Weakness 3: Theorem 1 is entirely based on L2 metric

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.

Question 1: FID using the Euler method with discretization setup.

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.

Question 2: Central difference for PID.

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.

[1]Yang, Zhantao, et al. "Eliminating Lipschitz Singularities in Diffusion Models." arXiv:2306.11251 (2023).

[2] Karras, Tero, et al. "Elucidating the design space of diffusion-based generative models." NeuIPS(2022): 26565-26577.

Thank you for your response. I hope to clarify that our submission, falling within four months of the BOOT [1] paper's publication in the ICML workshop in June 20 can be considered as concurrent work aligning with ICLR Reviewer guidelines. We kindly request the reviewer's consideration of this matter and thank you in advance for your understanding.

[1] Gu, Jiatao, et al. "Boot: Data-free distillation of denoising diffusion models with bootstrapping." ICML 2023 Workshop on Structured Probabilistic Inference {&} Generative Modeling. 2023.

We are grateful for taking the time to review our paper. We present our response to your concerns below.

Question: Similarities and Differences with concurrent work BOOT[1]

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.

Weaknesses: Manuscript update and minor typos.

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.

[1] Gu, Jiatao, et al. "Boot: Data-free distillation of denoising diffusion models with bootstrapping." ICML 2023 Workshop.

[2] Pao-Hsiung Chiu, Jian Cheng Wong, Chinchun Ooi, My Ha Dao, and Yew-Soon Ong. Can-pinn: A fast physics-informed neural network based on coupled-automatic–numerical differentiation method. Computer Methods in Applied Mechanics and Engineering, 2022

[3] Lu, Cheng, et al. "Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps." NeuIPS 2022: 5775-5787.

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.

[1] Gu, Jiatao, et al. "Boot: Data-free distillation of denoising diffusion models with bootstrapping." ICML 2023 Workshop on Structured Probabilistic Inference {&} Generative Modeling. 2023.

Thanks for pointing out the definition of concurrent work in the review guide. I think 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. But I respect and follow the review guide because it is not the reason for rejection. However, the final results of this work are still not promising. So I consider to raise my score to 6.