Solution of physics-based inverse problems using conditional generative adversarial networks with full gradient penalty

Inverse problems arise in many fields of engineering and science. Typically, they involve recovering the input fields or parameters in a system given noisy measurements of its response. In contrast to direct problems, where one wishes to determine the response given the input parameters, inverse problems are notoriously hard to solve. This is because they are often ill-posed in the sense of Hadamard [1] and involve measurements that are corrupted by noise, while their solution often requires injecting some prior information or beliefs regarding the inferred parameters or fields.

The challenges described above can be addressed by working with a probabilistic version of the inverse problem where the input and response vectors/fields are treated as random vectors /processes. In this manuscript, we refer to the input vector as the vector of the parameters to be inferred, or the inferred vector for short. In the probabilistic framework, the solution of an inverse problem transforms to characterizing the conditional probability distribution of the inferred vector conditioned on a measurement of the response; this distribution is known as the posterior distribution. Using Bayes’ rule, the posterior distribution can be expressed (up to a multiplicative constant) as the product of a marginal distribution for the inferred vector with the conditional distribution for the measurement given the inferred parameters [2]. The former is replaced by a prior distribution for the parameters to be inferred, and the latter leads to the likelihood term, which is determined by the model for the noise in the measurement and the forward map that links the parameters to the response. Once the probabilistic inverse problem has been framed, a method like Markov Chain Monte Carlo (MCMC) [3, 4] may be applied to sample from the posterior distribution; these samples help approximate the posterior distribution that may otherwise be analytically intractable [5]. The steady state of the Markov chain leads to samples that are drawn from the desired conditional distribution for the input parameters.

While the MCMC method and its variants can be used to solve some probabilistic inverse problems, their efficacy diminishes under several conditions. These include problems where the dimension of the inferred vector becomes large (greater than $\approx$ 50, for example) — the so-called curse of dimensionality [6, 7, 8]. It also includes problems where prior information about the parameters to be inferred cannot be expressed in terms of a simple probability distribution and is determined through samples. Finally, some variants of the MCMC method with improved convergence properties, like the Hamiltonian Monte Carlo (HMC) method [9, 10], rely on computing the derivative of the forward map with respect to the parameters. In cases where the forward map is implemented through a complex legacy code, computing this derivative is not feasible and these methods cannot be used.

Recently, there has been a growing interest in utilizing novel machine and deep learning algorithms to address the challenges described above. Most of these methods first consider samples of the parameters to be inferred derived from a prior distribution and utilize the forward map and a model for noise to generate corresponding samples of the measurement vector. Thereafter, they use this dataset to train algorithms for sampling the conditional distribution for the parameters given the measured response. One such class of methods that utilize deep neural networks to map the response vector to the vector of parameters, is referred to as Bayesian Neural Networks (BNNs) [11, 12, 13, 14]. In these networks, in addition to the input and output vectors, the network weights are also treated as random variables. Thereafter, Bayes’ rule is applied to derive an expression for the posterior distribution of the network parameters (weights and biases) given the training data. Then, variational inference is used to derive an approximation for the posterior density of the network parameters. Once this density is known, multiple realizations of the network parameters are sampled and the resulting networks are applied to the same measured response to yield an ensemble of inferred vectors. The performance of BNNs is limited by the approximations that are typically incurred to make variational inference tractable, which include assuming mixture models for the weights and selecting a small subset of the network weights as stochastic [15, 16].

Another class of deep learning algorithms utilize Generative Adversarial Networks (GANs) [17] to learn and then sample from the posterior distribution. GANs are generative models that learn a reduced-order representation of a probability distribution from its samples and then efficiently generate samples from it. GANs comprise of two networks: a generator and a critic. The generator network maps low-dimensional vectors drawn from a simple probability distribution to the output domain, while the critic maps samples from the output domain to a scalar. For the Wasserstein GAN [18, 19], under suitable assumptions, it can be shown that the probability distribution of the samples produced by the trained generator converges to the true distribution (from which the training samples are drawn) in the Wasserstein-1 distance. This is achieved by imposing a 1-Lipschitz constraint on the critic. In [20], a WGAN was used to learn the prior distribution of the inferred vector, and then used in conjunction with the HMC method to learn the posterior conditional distribution. While this approach reduced the dimension of the underlying inverse problem by mapping it to the low-dimensional latent space of the WGAN, it still required the use of an HMC method to sample from the posterior distribution.

In contrast to the approach described above, conditional GANs [21] can be used to approximate the posterior distribution directly. Adler and Öktem [22] proposed a conditional WGAN (cWGAN) that aims to directly learn the posterior distribution, and used it to improve the results of an inverse problem in medical imaging. Ray et al. [23] improved the architecture of cWGAN and applied it to solve complex inverse problems motivated by physics-driven applications. Kovachki et al. [24] proposed the Monotone GAN (MGAN) to solve inverse problems. This method differs from the cWGAN in several ways. First, it uses a latent space whose dimension is the same as that of the inferred vector, whereas the cWGAN typically uses a latent space of reduced dimension. Second, it relies on block triangular transformations to map the latent space to the space of the inferred vector, whereas a cWGAN uses a U-NET [25] architecture. Finally, it requires the generator that transforms the latent space to the space of inferred parameters to be monotonic, whereas the cWGAN imposes a 1-Lipschitz condition on the critic and not the generator.

The cWGAN proposed by Adler and Öktem [22] uses a critic network that is 1-Lipschitz with respect to the inferred vector only. The 1-Lipschitz constraint is implemented weakly through a gradient penalty term, and the gradient is computed with respect to the inferred vector. Herein, we refer to this as Partial-GP. As a result of the partial 1-Lipschitz constraint being imposed on the critic, the Kantorovich-Rubinstein duality (see [26] for example) can be invoked to show that the critic for the cWGAN with Partial-GP belongs to the space of functions that can only be optimized to approximate the Wassertstein-1 distance between distributions of the inferred vector. Based on this, Adler and Öktem [22] develop a training objective function for the cWGAN that aims to reduce the averaged Wasserstein-1 distance between the true posterior distribution and the one induced by the cWGAN; the average is taken over all possible values of the measurement vector. This does not necessarily mean that the distribution of the inferred vector conditioned on any one particular value of the measurement vector will converge to the true posterior distribution even if the critic and generator networks are sufficiently expressive, and are perfectly trained.

In this work we improve upon the cWGAN with Partial-GP [22], with a novel cWGAN formulation which differs from the original formulation in one critical aspect. In our approach, we require the critic network to be a 1-Lipschitz function of all its inputs — the inferred and measurement vector. Therefore, the crucial algorithmic difference that arises is that the gradient must be computed with respect to both arguments of the critic as opposed to just the inferred vector for Partial-GP. Remarkably, this small change in the implementation has a significant effect on the properties of the final network. Due to the critic being 1-Lipschitz with respect to its entire input argument, we can show, via the Kantorovich-Rubinstein duality, that the critic in the cWGAN belongs to the space of functions that can minimize the Wasserstein-1 between distributions of the joint random vector containing both the inferred variable and the measurements. As a result, we show that the generator of a cWGAN that is perfectly trained with the same training objective as [22] will converge to the true conditional distribution for the inferred vector for any measurement vector.

Additionally, the convergence guarantees of the proposed cWGAN with Full-GP allow us to develop a new evaluation strategy where the cWGAN is evaluated at measurement values obtained by perturbing the actual measurement vector with samples from a zero-mean Gaussian distribution that has a minute variance. We show that this new evaluation strategy may lead to improved performance of cWGANs. Further, we study the behavior of the proposed cWGAN with Full-GP on a suite of illustrative examples and its ability to approximate the joint distribution and the posterior distribution. We also apply the proposed cWGANs to two challenging inverse problems arising in inverse heat conduction and elastography [27] that help establish its efficacy in solving physics-based inverse problems.

The format of the remainder of this manuscript is as follows. In the following section we introduce the mathematical background required to develop and analyze the new cWGAN with Full-GP. In Section 3, we describe how the Full-GP cWGAN is trained and then used to solve a probabilistic inverse problem. In Section Section 4, we perform a detailed analysis of the convergence of the conditional density generated by Full-GP cWGAN with increasing number of trainable generator parameters. We also compare it with the corresponding analysis for the cWGAN with Partial-GP. In Section Section 5, we apply the Full-GP and Partial-GP cWGAN networks to inverse problems with increasing complexity, and evaluate their ability to recover the true posterior distribution for cases where it is known. We observe that the Full-GP cWGAN performs better. We end with conclusions in Section 6.

In this work, we propose a cWGAN model with Full-GP which, unlike the original cWGAN proposed by Adler and Öktem [22], imposes a 1-Lipschitz constraint on the critic $d$ with respect to its entire argument $({\bm{x}},{\bm{y}})$. Let us denote by Lip the 1-Lipschitz functions on $\mathbb{R}^{{N_{X}}}\times\mathbb{R}^{{N_{Y}}}$. If we assume that the maximization in (5) is over $d\in\text{Lip}$, then we can show by the Kantorovich–Rubinstein duality argument [26] that

where $W_{1}$ is the Wasserstein-1 distance defined on the space of joint probability measures on $\mathbb{R}^{{N_{X}}}\times\mathbb{R}^{{N_{Y}}}$. In practice, the 1-Lipschitz constraint on the critic can be imposed by augmenting the loss function by a gradient penalty term [19] when training the critic, i.e., by solving

where $\lambda>0$ is a hyper-parameter which appropriately weights the penalty term. We use the following penalty term:

where $\mathcal{U}(0,1)$ denotes the uniform distribution on $[0,1]$, $\nabla$ denotes the gradient with respect to both its arguments, and

The gradient penalty term in Eq. 8 imposes a 1-Lipschitz constraint on the critic with respect to both its arguments ${\bm{x}}$ and ${\bm{y}}$. Using this modified penalty term is crucial to ensure that solving minmax problem is equivalent to minimizing the $W_{1}$ distance between the true joint measure $\mu_{{\bm{X}}{\bm{Y}}}$ and the generated joint $\mu^{\bm{g}}_{{\bm{X}}{\bm{Y}}}$. This in turn leads to an alternative analysis of (weak) convergence to the true conditional measure $\mu_{{\bm{X}}|{\bm{Y}}}$ which is described in Section 4.

In this section, we analyze the convergence of the cWGAN with Full-GP and show how the optimal generator can be used to approximate the posterior statistics from Eq. 11 in accordance to Eq. 12. This proof, which forms the main result of this section, is contained in Theorem 4.1 for $\sigma=0$ and in Theorem 4.2 for $\sigma>0$. The proofs of these theorems rely on three distinct assumptions. The first (Assumption 4.1) is a statement of the convergence of a standard (not conditional) WGAN with increasing number of learnable parameters. This is also implicitly assumed in the original WGAN papers [18, 19]. The second assumption (Assumption 4.2) assumes that the true and learned conditional expectations are bounded. Finally, the third (Assumption 4.3) invokes the absolutely continuity of the marginal measure of measurement vector and ensuring the existence of a bounded marginal density.

Let $\mathcal{G}_{n}$ be the family of generator networks such that any network ${\bm{g}}\in\mathcal{G}_{n}$ has $n$ trainable parameters. We begin with our first assumption about the convergence of a standard WGAN with increasing capacity of the networks.

Under Assumption 4.1, the sequence of measures $\mu^{{\bm{g}}^{*}_{n}}_{{\bm{X}}{\bm{Y}}}$ will converge to $\mu_{{\bm{X}}{\bm{Y}}}$ in the $W_{1}$ metric, and hence weakly [26], i.e.,

Next, we show how an appropriate choice of $\ell\in C_{b}(\mathbb{R}^{{N_{X}}}\times\mathbb{R}^{{N_{Y}}})$ leads to approximating the conditional expectation (11). Given a generator ${\bm{g}}$, we define the generated conditional expectations for any $q\in C_{b}(\mathbb{R}^{{N_{X}}})$ as

We require the conditional expectations to be finite and computable.

We also make the following assumption about the true marginal $\mu_{\bm{Y}}$ to simplify the exposition.

Now, let $\hat{{\bm{y}}}\sim\mu_{\bm{Y}}$ be given for which $p_{\bm{Y}}(\hat{{\bm{y}}})\neq 0$. Let $\mu_{{\bm{Y}}_{\sigma}}:=\mathcal{N}(\hat{y},\sigma^{2}\bm{I})$ be the Gaussian measure on $\mathbb{R}^{N_{Y}}$ with density

The following standard result shows how the Dirac measure can be approximated using a Gaussian.

For the rest of the theoretic discussion, we omit the superscript $*$ and it is understood that ${\bm{g}}_{n}$ denotes the optimal generator in accordance to Assumption 4.1. We are now ready to state and prove the main results about approximating Eq. 11.

We construct a sequence of joint measures with elements

which is different from $\mu_{{\bm{X}}{\bm{Y}}}$ or $\mu^{{\bm{g}}_{n}}_{{\bm{X}}{\bm{Y}}}$. We now prove a robustness result which shows that we can obtain error estimates similar to Theorem 4.1 even if we inject a controlled amount of noise into the given measurement ${\bm{y}}$ in accordance to $\mu_{{\bm{Y}}_{\sigma}}$.

Note that

Taking the Monte Carlo approximation of Eq. 31 leads to the approximation formula Eq. 12 in our algorithm, with the choice of $\sigma$ motivated by the proof of Theorem 4.2. In particular, to get a reasonable approximation of the conditional expectation, we need to pick a $\sigma=\bar{\sigma}$ small enough to ensure that inequalities 26, 27 and 29 are satisfied and a generator large (expressive) enough to ensure the inequality 30 is satisfied.

In this section, several numerical experiments are presented to analyze the performance of the proposed cWGANs with Full-GP. The experiments include: one-dimensional problems where the target measure is known; an inverse heat conduction problem where the reference statistics of the target measure are available [23]; and a complex inverse problem in elastography where the target measure is unknown. For the first two problems, where the target statistics are available and we have a reliable reference to estimate error, we use both the Partial-GP [22] (see Eq. 10), and Full-GP (see Eq. 8) approaches to perform a comparison of the two methods.