ODE-DPS: ODE-based Diffusion Posterior Sampling for Inverse Problems in Partial Differential Equation

Partial differential equations (PDEs) are fundamental concepts in various scientific and engineering fields, serving as the cornerstone of many disciplines. They are characterized by their dependence on parameters that govern their behavior, which allows them to describe a wide range of physical and operational scenarios. Solving partial differential equations under different parameter settings is meaningful because it enables predictions of physical laws under varying conditions. Various solution techniques, such as the finite element method (FEM), finite volumes, finite differences, and spectral methods, have been developed for tackling the problem [1, 2]. This process establishes a mapping from parameters to solutions, known as a parameter-to-solution map.

Conversely, when given experimental data, there is often a need to infer the unknown parameters associated with the observed phenomenon, thereby determining the parameter regimes of the observed system. This situation is known as the inverse problem. Inverse problems are generally more challenging than forward problems due to their ill-posed nature, where minor errors in observations can lead to significant inaccuracies in model parameters. This challenge is further exacerbated in scenarios involving observation noise, incomplete physics, and high-dimensional parameter spaces [3, 4].

So far, several regularization methods have been employed to address ill-posed inverse problems, including Tikhonov-type regularization methods, iterative regularization methods, and Bayesian inversion methods [3, 4]. The Tikhonov regularization method stabilizes the solution of the inverse problem by incorporating regularization terms and carefully selecting regularization parameters [5]. The iterative inversion method tackles the ill-posed nature of the inverse problem by implementing an iterative stopping criterion [5]. The Bayesian inversion method transforms ill-posed inverse problems into well-posed posterior distribution problems through the introduction of suitable prior distributions [3]. Although significant progress has been witnessed in various inverse problems, they are hindered by their inability to effectively utilize data from past records. This limitation leads to inaccurate inversion, making it challenging for the resulting parameters to meet practical requirements.

Recently, deep learning has emerged as a potent tool for addressing inverse problems related to partial differential equations (PDEs). For example, Raissi et al. [6] introduced physics-informed neural networks (PINNs) to tackle such problems. In this methodology, neural networks represent unknown parameters, leveraging physical equations to construct a loss function. Subsequently, stochastic gradient descent optimizes the neural network parameters, yielding numerical inversion results. In recent years, research in this field has flourished, incorporating various methodologies like weak adversarial networks for inverse problems [7], PDE-aware deep learning methods for inverse problems [8], implicit neural representations for seismic inversion [9], and inverse obstacle scatter [10]. These approaches operate on a case-by-case basis, necessitating complete retraining for each new set of unknown parameters, measured data, or modifications to other conditions in the inverse problem.

To address the aforementioned case-by-case scenario, researchers have proposed inverse neural operators as a solution for PDE inverse problems. Fan et al. [11] introduces compact neural network architectures for regularized inverse operators to solve the EIT inverse problem. Molinaro et al. [12] propose a novel architecture called Neural Inverse Operators (NIOs) for solving PDE inverse problems. NIOs utilize a composition of DeepONets and FNOs to approximate mappings from operators to functions. Ovadia et al. [13] combine a U-Net-based architecture with a vision transformer to design an inverse neural operator for solving various PDE inverse problems. While these inverse neural operators have demonstrated effective numerical inversion results for PDE inverse problems, acquiring pairs of data remains a challenging issue.

Presently, generative models [14, 15, 16, 17, 18, 19] have emerged as an alternative approach for solving inverse problems, offering the advantage of not necessitating paired unknown parameters and measurement data. This method only requires partial prior data about the unknown parameters and does not rely on their corresponding measurement data. Generative models constitute a category of machine learning models designed to learn the underlying data distribution and produce new data samples resembling the original. Typical examples of generative models encompass Autoencoders [16], Variational Autoencoders (VAEs) [17], and Generative Adversarial Networks (GANs) [18], among others. In recent years, diffusion models have gained prominence as mainstream generative models. Simultaneously, they have demonstrated effectiveness in addressing inverse problems within the visual field, as evidenced by studies such as [20, 21, 14, 22].

We present in this work an ODE-based Diffusion Posterior Sampling (ODE-DPS) algorithm for solving inverse problems in partial differential equations (PDEs). Specifically, by leveraging partial prior data of unknown parameters, combined with diffusion generative models and deep neural networks, we learn the prior distribution of the data, which is then applied in solving the Bayesian posterior distribution of the PDEs inverse problems. Then, we solve the posterior distribution using a conditional generation process achieved by solving a reverse-time stochastic differential equation. To enhance the inversion accuracy, we propose the ODE-DPS regularization inversion algorithm. This algorithm originates from an equivalent form of a noising process for the score-based diffusion model, with its equivalence guaranteed by the marginal probability density functions satisfying the same Fokker-Planck equation as in the forward noising process. Furthermore, based on experimental observations, we propose modifying the original algorithm by replacing the $L_{2}$-norm data residual term with an adaptive norm residual term. The purpose is to reduce the errors near the boundaries in the inversion results obtained by the score-based diffusion sampling algorithm. We summarize the main novelties of this paper as follows:

• •

  We propose an ODE-based diffusion sampling algorithm to solve inverse problems of PDEs. This algorithm utilizes a limited amount of prior data and integrates score-based generative models to learn the prior distribution of unknown parameters. Subsequently, it is applied in a conditional distribution-based backward generation process induced by Bayesian inversion.

• •

  By leveraging the Fokker-Planck equation satisfied by the marginal probability density function of the noising process in the generative model, we derived a new reverse ordinary differential equation. Subsequently, we discretized it numerically to develop our ODE-based DPS inversion algorithm.

• •

  Based on experimental observations, we propose modifying the diffusion posterior sampling algorithm by replacing the $L_{2}$-norm data residual term with an adaptive norm residual term to reduce the errors near the boundaries in the inversion results.

• •

  ODE-DPS algorithm requires only a small amount of prior data for unknown parameters and does not require observation data, also referred to as labeled data, during the training phase. Once trained, this generative model can be applied to various inverse problem-solving tasks. Compared to traditional inversion algorithms, it enhances the inversion accuracy with minimal reduction in efficiency.

• •

  Numerically, we compare our algorithm with traditional inversion methods across various inverse problems. It clearly shows that the ODE-DPS inversion algorithm can significantly improve the inversion accuracy.

The rest of this paper is organized as follows: In Section 2, we give an introduction to inverse problems. Then we elucidate the score-based diffusion model in Section 3. After that, we explain our method to solve the inverse problems in Section 4, apply it to some examples, and compare it with some traditional methods in Section 5. Finally, we end this paper with conclusions in Section 6.

In this section, we consider the inverse problem of a general form, aiming to estimate $\bar{x}_{0}\in D$ from the data represented by the operator equation:

where $F:D\subset X\rightarrow Y$ is an operator between reflexive Banach spaces $(X,\|\cdot\|_{X})$ and $(Y,\|\cdot\|_{Y})$ with domain $D$. Here, $\bar{y}_{\delta}$ denotes the observed data with noise $\xi_{\delta}$.

Generally speaking, the inverse problem (1) is ill-posed in the sense that the solution may not exist, be not unique, or not vary smoothly on the data due to the small noise caused by the measurement and recording process. Consequently, directly solving the operator equation (1) for ill-posed inverse problems is not feasible. To tackle these challenges, regularization techniques are commonly employed. The Tikhonov regularization method [23, 5] stands out as a well-established and effective approach for addressing inverse problems across various fields. This method replaces solving operator equation (1) by minimizing the following problem:

where $\|\bar{x}\|_{X}$ denotes the regularization term, and $\alpha$ is a regularization parameter. The efficacy of the Tikhonov regularization method in solving inverse problems is significantly influenced by the choice of regularization parameter and the regularization term norm $\|\cdot\|_{X}$. Optimal selection of a regularization parameter and regularization term can substantially enhance the accuracy of the inversion results. Nevertheless, determining the optimal parameter and norm are typically a challenging task.

Additionally, another method widely employed in solving inverse problems is the iterative method [24, 25]. However, due to the ill-posed nature of the inverse problem, the iterative inversion method often encounters semi-convergence. This phenomenon manifests as a decrease followed by an increase in the relative error of the numerical inversion result with iteration process. Therefore, the success of this inversion method hinges on choosing a suitable stopping criterion. Morozov’s discrepancy principle [26] is typically utilized for this purpose. Currently, integrating the discrepancy principle into iterative methods has emerged as another mainstream approach for addressing ill-posed inverse problems.

Although the inversion methods mentioned above have been applied to different inverse problems in the past, due to the ill-posedness of the inverse problem, we usually can only obtain low-precision inversion results. For example, let us consider the following equation:

where $a$ is a positive constant. We assume that the initial value $\phi$ and boundary $\psi$ and the constant $a$ are known (Here we set $a=1,~{}\psi(x,y,t)=0,~{}\phi(x,y)=20\sin 2\pi x\sin 2\pi y,~{}\Omega=[0,1]\times[0,1],T=1$), and use the measurement data $u(x,y,T)$ at the terminal time $T=1$ to recover the unknown source term $f$. We set the time step size to 0.01 and the grid size to 1/63 to discretizate $\Omega$, and denote the measurement data with noise as $u^{\delta}(x,y,T)$. We apply Landweber iteration regularization [27] and Tikhonov regularization method [5] for solving the ill-posed inverse source problem, and Morozov’s discrepancy principle [26] for stopping the Landweber iteration process. In the Tikhonov regularization method, we choose the norm of the regularization term to be the $l_{2}$ norm, with a regularization parameter of $\alpha=0.005$. The inversion results are shown in Figure 1. Compared with the exact solution in Figure 1a, the numerical results of the inversion in Figures 1b, 1c only captures large-scale features, and small-scale features are difficult to capture. In addition, the errors in Figures 1d, 1e also reflect that Landweber iteration regularization method and Tikhonov regularization method are difficult to give accurate inversion results. The main reasons why this inversion method struggles to achieve high-precision inversion results are primarily twofold:

1. 1.

   Although we only consider a linear inverse problem in this numerical example, their inherent ill-posed nature persists. Hence, it remains challenging to obtain an accurate solution for the ill-posed problem.

2. 2.

   While we have utilized the mainstream regularization methods to solve the aforementioned inverse problem, we have not incorporated any prior information about the inverse problem into the inversion process, apart from selecting the initial value and the norm of regularization term. In reality, leveraging prior information past-recorded of inverse problems can significantly enhance inversion accuracy.

Therefore, it is necessary for us to design a new regularized inversion algorithm to solve the ill-posed inverse problem. In recent years, deep learning-based inversion methods, as another form of regularization, have been applied across various fields of inverse problems, showing a trend of outperforming traditional inversion methods. Their advantage primarily stems from their ability to learn prior information about inverse problem solutions from previously recorded data, which can then be applied to the inversion process. Building upon this notion, this paper proposes a novel regularization inversion method. This method leverages partial prior information of unknown parameters, deep neural networks, score-based diffusion generative model, and the traditional Bayesian inversion framework to introduce a completely new regularization inversion approach. The following section will present the interpretation of our inversion method.

In this article, our main focus is on constructing new regularization algorithms using generative models and Bayesian inversion methods. Deep generative models are widely used in many subfields of AI and Machine Learning. These models, when combined with stochastic optimization methods, excel at capturing complex, high-dimensional data distributions such as images, text, and speech. Today, popular deep generative models include variational autoencoders (VAEs), generative adversarial networks (GANs), auto-regressive models, and so on. In particular, diffusion models, also known as diffusion probabilistic models or score-based generative models, represent a class of latent variable generative models that are found in applications for diverse inverse problem tasks, including image denoising [28], inpainting [29], super-resolution [30], and medical imaging [15]. These models have demonstrated state-of-the-art performance in these fields. Consequently, our objective in this study is to devise a novel regularization inversion method for addressing PDE inverse problems by integrating score-based diffusion models with traditional Bayesian inversion methods.

Score-based generative models represent a continuum of distributions evolving over time through a diffusion process. This process gradually transforms a data point into random noise, as described by a predetermined stochastic differential equation (SDE) independent of the data and devoid of trainable parameters. Reversing this process enables us to smoothly convert random noise into data for sample generation. Importantly, this reverse process adheres to a reverse-time SDE, derivable from the forward SDE based on the score of marginal probability densities over time. Consequently, we can approximate the reverse-time SDE by training a time-dependent neural network to estimate the scores, facilitating sample production using numerical SDE solvers. Below we will explain the score-based generative model in detail.

Consider the inverse problem mentioned in (1), and we assume that the unknown variable $\bar{x}_{0}$ and the observed data $\bar{y}_{\delta}$ are random variables, $\xi_{\delta}$ is the measurement noise. In the case of white Gaussian noise, $\xi_{\delta}\sim\mathcal{N}(0,\sigma^{2}I)$. In the Bayesian framework, one utilizes $p(\bar{x}_{0})$ as the prior, and denote the likelihood $p(\bar{y}_{\delta}|\bar{x}_{0})\sim\mathcal{N}(\bar{y}_{\delta}|F(\bar{x}_{0}),\sigma^{2}I)$. By Bayesian rules’, the posterior distribution $p(\bar{x}_{0}|\bar{y}_{\delta})$ denoted by

Then the inverse problem (1) is transformed into the problem of solving the posterior distribution $p(\bar{x}_{0}|\bar{y}_{\delta})$. To do it, we use the forward diffusion process (4) to obtain a family of diffused distribution $p(\bar{x}(t)|\bar{y}_{\delta})$ with the initial distribution $p(\bar{x}_{0}|\bar{y}_{\delta})$ and apply Anderson’s Theorem to derive the conditioned reverse time SDE

Then the solution to the inverse problem, that is, solving the posterior distribution $p(\bar{x}_{0}|\bar{y}_{\delta})$, is transformed into sampling from the conditional reverse time SDE (7). Since $p(\bar{x}(t)|\bar{y}_{\delta})\propto p(\bar{x}(t))p(\bar{y}_{\delta}|\bar{x}(t))$, the score $\nabla_{\bar{x}(t)}\log p(\bar{x}(t)|\bar{y}_{\delta})$ in (7) can be computed easily by

Combing (7) with (8), we have

Leveraging (9), scored-based generative models provide a way to solve the Bayesian inverse problem, i.e., to solve the posterior distribution $p(\bar{x}_{0}|\bar{y}_{\delta})$. In (9), we have two terms that should be computed: the score function $\nabla_{\bar{x}(t)}\log p(\bar{x}(t))$ and the likelihood function $\nabla_{\bar{x}(t)}\log p(y_{\delta}|\bar{x}(t))$. To compute the former term $\nabla_{\bar{x}(t)}\log p(\bar{x}(t))$, we can simply use the pre-trained score function $s_{\theta^{*}}(\bar{x}(t),t)$. However, the latter term is hard to acquire in closed-form due to the dependence on the time $t$, as there only exists dependence between $\bar{y}_{\delta}$ and $\bar{x}_{0}$.

In this section, we show the performance of our method by three different inverse problems. Experimental performance is evaluated using both qualitative analysis and quantitative metric. We define the relative $l_{2}$ error as follows:

where $\bar{x}_{ref}$ is the ground truth and $\bar{x}_{rec}$ is the reconstructed parameter. The relative $l_{2}$ error measures the degree of accuracy in reconstructing the unknown parameters.

To show the advantages of our proposed inversion algorithm, we will contrast it with the traditional regularization inversion methods, namely the Landweber iterative regularization method [27, 5] and the Tikhonov regularization method [5], in the paper. We illustrate the methods as follows:

1. 1.

   Method 1: Landweber iteration regularization. It is a widely used regularization inversion method applied in solving ill-posed inverse problems, and various extensions of this method have been developed for different inverse problems. In this paper, we only focus on the most basic version [27], i.e., given $\bar{x}_{0}$ is an initial guess which may incorporate a prior knowledge of an exact solution for inverse problem, the iteration is defined via the adjoint $F^{{}^{\prime}}(\cdot)^{*}$ of $F^{{}^{\prime}}(\cdot)$ as follows:

   $$\bar{x}_{k+1}^{\delta}=\bar{x}_{k}^{\delta}+\zeta_{k}F^{{}^{\prime}}(\bar{x}_{k}^{\delta})^{*}(\bar{y}^{\delta}-F(\bar{x}_{k}^{\delta})),~{}k=0,1,2,\cdots,$$

   where $F^{{}^{\prime}}(\cdot)$ is the Fréchet derivative of forward operator $F(\cdot)$ and $\zeta_{k}$ is the step size of $k$th step. To obtain a stable approximation of the method, the iteration process need to be stopped after $k_{*}=k_{*}(\delta)$ steps according to a generalized discrepancy principle, i.e.,

   $$\|\bar{y}^{\delta}-F(\bar{x}_{k_{*}}^{\delta})\|\leq\tau\delta<\|\bar{y}^{\delta}-F(\bar{x}_{k}^{\delta})\|,~{}0\leq k\leq k_{*},$$

   where $\tau>1$ is a constant. In general, we set $\tau=1.01$. Additionally, in Lemma 2.4 of reference [5], we obtain that the Landweber iteration step is the steepest descent step with a stepsize $\zeta_{k}$ of a quadratic functional $\Psi(\bar{x}^{\delta})=\frac{1}{2}\|\bar{y}^{\delta}-F(\bar{x}^{\delta})\|_{2}^{2}$. We denote the descent direction by $\nabla_{\bar{x}^{\delta}}\Psi(\bar{x}^{\delta})$.

2. 2.

   Method 2: An improved Landweber iteration regularization. Following the principles of our algorithm, we modify the descent direction using $\nabla_{\bar{x}^{\delta}}\Psi_{\rho}(\bar{x}^{\delta})$, where $\Psi_{\rho}(\bar{x}^{\delta})=\frac{1}{2}\|\bar{y}^{\delta}-F(\bar{x}^{\delta})\|_{\rho}^{2}$, with $\|\cdot\|_{\rho}$ defined in (17).

3. 3.

   Method 3: Tikhonov regularization method. Our objective is modified to solve the following optimization problem:

   $$\bar{x}^{\delta}:=\arg\min\limits_{\bar{x}}\|\bar{y}^{\delta}-F(\bar{x})\|_{2}^{2}+\alpha\|\bar{x}\|_{2}^{2}.$$

   (19)

   Here we set $\alpha=0.005$ or $\alpha=0.1$ in different experiments, and we find that this choice yields the best inversion results for the regularization method.

Unless otherwise specified, in all numerical examples, for the three traditional regularization methods mentioned above, we fix the maximum number of iterations at 2000, and set the step size $\zeta_{k}$ in iterations equal to our method ODE-DPS.