PIDM] Diffusion Model With Physical Regularization

Diffusion Model 30/09/2024

3 main points
✔️ Proposes a theoretical method to introduce physical constraints into diffusion models
✔️ Confirms regularization capability for overfitting
✔️ Can expand not only equations but also inequalities and optimal functions

Physics-Informed Diffusion Models

written byJan-Hendrik Bastek, WaiChing Sun, Dennis M. Kochmann
(Submitted on 21 Mar 2024)
Comments: 15 pages, 4 figures; added further theoretical motivation, new residual estimation mechanism and additional experimental study
Subjects: Machine Learning (cs.LG); Computational Engineering, Finance, and Science (cs.CE)

code：

The images used in this article are from the paper, the introductory slides, or were created based on them.

Summary

Diffusion models have extremely high performance and versatility in approximating very complex data distributions, and their application in the natural sciences has spread rapidly in recent years. Given their application in a scientific context, it is often the case that for a particular class of problems, the governing equations that the data follow are explicitly known, and one wishes to provide a regularization that the data follow these governing equations. To date, many scientific applications of diffusion models have been purely data-driven, and it has been nontrivial whether the samples generated obey the physical laws. In this study, we investigate theoretical methods to provide physical regularization to diffusion models so that the generated samples obey the governing equations, and demonstrate their usefulness through numerical experiments.

Background

An important foundation in current AI is the diffusion model. Diffusion models are extremely powerful in approximating the complex distributions that data follows, and their use has spread rapidly. In recent years, diffusion models have also found applications in the natural sciences. However, most scientific applications of diffusion models are purely aimed at approximating the distribution followed by the data from the data, and there has been insufficient consideration of their integration with scientific knowledge. In certain cases, where the governing equations followed by the data are explicitly known from scientific knowledge, it is necessary to establish methods to physically constrain the diffusion model so that the generated samples follow the governing equations. Against this backdrop, this study theoretically investigates the use of physical regularization in diffusion models.

Proposed Method

Here we describe several important elements during this study.

Diffusion Model

Diffusion models are a type of state-of-the-art generative model. Briefly, the purpose of a diffusion model is to approximate and model the distribution $q(x_0)$ that the data $x_0$ follows. To achieve this, the diffusion model considers a $T$-step series of data from $x_0$ to $x_T$, adding Gaussian noise at each step so that $x_T$ is pure Gaussian noise. This process is called the forward difussion process and is defined below.

where $\{\beta\in (0,1)\}_{t=1}^T$ is the parameter that determines the diffusion process. We also define below the series that generates samples from the Gaussian process as the inverse operation of this process.

Here, the unknown distribution, $q(x_{t-1}|x_t)$,is approximated as$p_θ$$(x_{t-1}|x_t)$ bya neural network. Finally, following the simplification by Ho et al. we train by the following loss function

In other words, conceptually, the model is trained so that the error between $\hat{x}_0$, which is the result of adding noise and then removing the noise, and the original data $x_0$ is small. The above is a brief explanation of the diffusion model.

Governing Equation

In general, governing equations can be expressed abstractly as follows

In addition, consider the following general boundary conditions

However, where is the abstract differential operator, is the boundary conditions, and is the solution to the governing equations. During this study, we assume that the sample $x_0$ generated by the diffusion model satisfies the governing equations up to and including the boundary conditions described above. In this study, $x_0$ is considered to be data in the form of an image, for example, an image representing the distribution of stress in a mechanics problem. According to this governing equation, we define the following residuals

So, simply put, this residual measures whether the governing equation is satisfied.

Physics-Informed Diffusion Model

Objective function design

In this study, a hypothetical residual $\hat{r}$ is introduced to avoid compromising the stochastic perspective of the diffusion model and is assumed to follow the following distribution

Using this,we consider thehypothetical likelihood $p_θ$($\hat{r}$)as follows.

Using these, the physical regularization objective can be expressed as follows

This means that the parameters are adjusted so that the probability of the residuals being zero is maximized. The authors point out that this is a probabilistic interpretation of the error function of a physics-informed neural network.

Furthermore, during this study, in addition to the error function described above, an error function due to observed data is also considered. As the authors point out, this has the effect of preventing the function to be estimated from collapsing. That is, a function that is zero everywhere, for example, is a solution that satisfies a particular type of governing equation but is not physically meaningful. We need an error from the observed data as some sort of regularization term so that we do not get stuck with such a nontrivial solution when searching for the first floor. Adding the error from the observed data, we define the objective function as follows

Simplification of the objective function

Imposing a zero residual on the diffusion model for all series of denoising is too strict for regularization and risks compromising the flexibility of the diffusion model. Therefore, the authors designed the regularization to scale at each step as appropriate. That is, we designed the regularization to become stronger as the denoising is applied from $T$ to step 0. The following is the definition for this. Figure 1 shows a schematic of the PIDM for the denoising process.

Figure 1: Schematic of PIDM in the noise reduction process.

Furthermore, by simplifying the objective function, taking into account that is designed to approximate , we can finally obtain the following objective function

The first term is the error function for the observed data and the second term is the error function for the governing equations. The above is a brief description of PIDM.

Experimental Results

In this study, several numerical experiments are performed to demonstrate the effectiveness of PIDM. In this commentary, we will focus on an example of a two-dimensional Darcy flow that represents flow in a porous medium as one of the more physically relevant examples. The governing equations are as follows

However,

It is. In this study, several models were prepared for comparison. (i) a diffusion model that learns from an objective function with data, (ii) a diffusion model that learns from an objective function with data only but considers information on residuals as guidance (Phisics-guided diffusion), and (iii) a model that but with a first-order modification of the residuals during inference (CoCoGen). This study also considers models that generate the permittivity and pressure in the above equations.

Figure 2. Visualization of the process of learning each model.

As Figure 2, we visualize the history of the error function in the process of learning each model and the PIDM proposed in this study. From Figure 2(a), we can see that PIDM improves by about two orders of magnitude as the error against the residuals. Furthermore, Figure 2(b) shows that the conventional method is expected to overfit the test data as the learning process progresses, but PIDM prevents this by physically regularizing the test data. These results suggest that the regularization term not only improves accuracy, but also prevents overfitting.

Figure 3. (a) Permittivity and pressure given by the normal diffusion model, (b) (c) Permittivity and pressure given by PIDM.

Also shown as Figure 3 is a comparison of the permittivity and pressure given by the normal diffusion model and PIDM. Also shown is the spatial distribution of the residuals in that case. Consistent with the results in Figure 2, this result also confirms that the estimation is more accurate than that of the ordinary diffusion model. In addition, (b) and (c) in Figure 3 represent different physical states, suggesting that the PIDM can represent a wide variety of states without falling into a single solution. These results support the high potential of PIDM.

Summary and Conclusion

In this study, we theoretically derived the diffusion model PIDM with the addition of a physical regularization term and demonstrated its performance through numerical experiments. In particular, we were able to impose physical regularization directly on the diffusion model, rather than by "modifying" latent variables during inference, as in previous studies. Numerical experiments also suggest that PIDM is robust to overfitting, in addition to improving the error on residuals by about two orders of magnitude compared to ordinary diffusion models. The theoretical foundation provided by this study will be important as the demand for diffusion models in the natural sciences continues to increase. It is expected to be applied to a wide variety of specific, concrete natural science problems in the future and to serve as the basis for a general-purpose tool.