# The efficacy and generalizability of conditional GANs for posterior
inference in physics-based inverse problems

Deep Ray
Department of Aerospace and Mechanical Engineering
University of Southern California
Los Angeles, CA 90089, USA
<EMAIL_ADDRESS>&Harisankar Ramaswamy
Department of Aerospace and Mechanical Engineering
University of Southern California
Los Angeles, CA 90089, USA
<EMAIL_ADDRESS>&Dhruv V. Patel
Department of Mechanical Engineering
Stanford University
Stanford, CA 94305, USA
<EMAIL_ADDRESS>&Assad A. Oberai
Department of Aerospace and Mechanical Engineering
University of Southern California
Los Angeles, CA 90089, USA
<EMAIL_ADDRESS>
###### Abstract

In this work, we train conditional Wasserstein generative adversarial networks
to effectively sample from the posterior of physics-based Bayesian inference
problems. The generator is constructed using a U-Net architecture, with the
latent information injected using conditional instance normalization. The
former facilitates a multiscale inverse map, while the latter enables the
decoupling of the latent space dimension from the dimension of the
measurement, and introduces stochasticity at all scales of the U-Net. We solve
PDE-based inverse problems to demonstrate the performance of our approach in
quantifying the uncertainty in the inferred field. Further, we show the
generator can learn inverse maps which are local in nature, which in turn
promotes generalizability when testing with out-of-distribution samples.

##  1 Introduction

Inverse problems arise in various areas of science and engineering, such as
computerized tomography [1], seismology [2, 3], climate-modeling [4, 5], and
astronomy [6]. While the forward/direct problem is generally well-studied and
easier to solve, the inverse problem can be notoriously hard to tackle due to
the lack of well-posedness [7]. Reconstructing the inferred field from
measurements corrupted by noise, which is usually the case for most realistic
problems, makes the task even more challenging. Bayesian inference provides a
principled strategy to overcome these challenges by posing the problem in a
stochastic framework. Here Bayes’ rule is used to update the prior belief
about the inferred field, based on the knowledge gained from measurements [8].
The end result is the recovery of an expression of the posterior distribution
for the inferred field, or the ability to generate samples from such a
density. As opposed to deterministic methods of solving inverse problems, such
as regularization-based strategies [9], Bayesian inference can be used to
quantify the uncertainty in the reconstructed field. The knowledge of this
uncertainty can play a crucial role in applications where important decisions
need to be made based on the inferred field.

There has been growing interest in using deep learning tools to overcome some
of the computational bottlenecks that arise when using Bayesian inference in
practise. For instance, the Bayes’ update formula requires an explicit prior
to be constructed. Practitioners prefer to use simple densities, such as a
multivariate Gaussian distribution, which in some cases ensures an expression
for the posterior that is easy to sample from. However, such a density is not
suitable to represent complicated priors in most practical problems,
particularly in applications where the prior needs to be constructed from
snapshots of the inferred field.

Deep generative models have proven to be successful in learning the underlying
probability distribution from data arising from medical imaging, geophysical
processes and computer vision [10, 11, 12, 13]. In [14, 15], generative
adversarial networks (GANs) were used to learn the prior from data arising in
physics-based models. Further, this strategy helped in significantly reducing
the overall dimension of the problem, thus making posterior sampling
algorithms like MCMC computationally tractable. Physics informed GANs [16]
have also been designed to solve both forward and inverse problems, where the
underlying stochastic differential equations have been encoded into the GAN
architecture. In [17], a conditional GAN (cGAN) was designed to directly
approximate and sample from the posterior distribution arising in the image
reconstruction problem for computed tomography (CT) of the liver. At the same
time, a mathematical framework was developed to prove the weak convergence of
the learned distribution. More recently, cGANs have been used to solve inverse
problems associated with nearshore bathymetry [18].

Motivated by [17], we propose a deep learning based algorithm to solve
Bayesian inverse problems modelled by partial differential equations (PDEs).
In particular, we train a conditional Wasserstein GAN (cWGAN) on a dataset of
samples drawn from the joint probability distribution of the inferred field
and the measurement. Such a dataset can be either acquired from experiments,
or generated synthetically by assuming a suitable prior on the inferred field
and numerically solving the forward PDE model. Further, samples from the
posterior distribution can be generated by simply querying the trained network
instead of depending on MCMC algorithms, which provides a significant
computational boost over traditional Bayesian inference techniques. Using a
Wasserstein adversarial loss to train the cGAN allows us to take advantage of
the theory developed in [17] which states that the generated posterior
statistics converge to the true posterior statistics. We highlight below the
main contributions of the present work:

  1. 1.

We construct a cWGAN which uses a residual block-based U-Net architecture for
the generator subnetwork. While the general structure of the generator is
similar to [17], we use a much simpler architecture for the critic as compared
to the conditional mini-batch critic constructed in [17].

  2. 2.

We inject the latent variable at different scales of the U-Net using
conditional instance normalization (CIN) [19]. This helps introduce multi-
scale stochasticity, while giving the flexibility to choose the latent space
dimension independently of the dimension of the measurement or the inferred
field. This is another key distinction between the architecture proposed in
this paper as compared to [17] where the latent variable is stacked as an
additional channel to the generator input.

  3. 3.

In contrast to prior works that typically illustrate the performance of
conditional GANs on a single physics-based application, such as CT imaging of
the liver [20], nearshore bathymetry [18] or hydro-mechanical processes in
porous media [21], we consider multiple PDE-based inverse problems to
demonstrate the performance of the proposed algorithm. Specifically, we
consider the problem of inferring the initial condition of the heat conduction
problem, the coefficient of conductivity for the steady-state heat conduction
problem, and the shear modulus in elasticity imaging. Further, the cWGAN for
the elasticity imaging problem is trained on synthetic data and then tested on
experimentally measured data.

  4. 4.

Finally, we investigate the generalizablity of the cWGAN by testing the
performance of the trained network on out-of-distribution (OOD) measurements.
We also provide a theoretical justification for the generalization (Theorem
4.1) which is related to the local nature of the inverse map learned by the
cWGAN.

The rest of the paper is organized as follows. Section 2 describes the problem
setup and how a cWGAN can be used to solve inverse problems. Section 3 details
and motivates the specialized cWGAN architecture considered in this work.
Section 4 provides a theoretical explanation for the generalization observed
in the numerical results presented in Section 5. Concluding remarks are made
in Section 6.

##  2 Problem formulation

We begin by considering a forward/direct problem given by the map

$\bm{f}:\Omega_{\bm{x}}\rightarrow\Omega_{\bm{y}}$

which takes an input $\bm{x}\in\Omega_{\bm{x}}\subset\mathbb{R}^{{N_{x}}}$ to
produce an output/measurement
$\bm{y}\in\Omega_{\bm{y}}\subset\mathbb{R}^{{N_{y}}}$. For the examples
considered in the present work, the forward problem involves solving a PDE. In
this context, $\bm{x}$ denotes quantities such as the initial condition,
source term or material properties, $\bm{f}$ denotes the solution operator of
the PDE approximated by a suitable finite volume or finite element scheme,
while $\bm{y}$ denotes the (discretized) solution of the PDE. Further, the
measurement $\bm{y}$ might be corrupted by noise, which is typical of most
realistic scenarios.

We are interested in solving the inverse problem: given a measurement
$\bm{y}$, infer the associated $\bm{x}$. To solve this in the Bayesian
framework, we assume that the inferred field is modeled as realizations of the
vector-valued random variable $\bm{\mathcal{X}}$ with a (prior) density
$p_{\bm{\mathcal{X}}}$. Then, the forward map induces a marginal density
$p_{\bm{\mathcal{Y}}}$ for the random variable
$\bm{\mathcal{Y}}=\bm{f}(\bm{\mathcal{X}})$. The goal is to determine the
posterior density $p_{\bm{\mathcal{X}}|\bm{\mathcal{Y}}}$ given a measurement
$\bm{y}$. In particular, we want to effectively generate samples from the
posterior, which can then be utilized to determine posterior statistics.

We assume access to $M$ realizations $\\{\bm{x}^{(i)}\\}_{i=1}^{M}$ from
$p_{\bm{\mathcal{X}}}$, and construct the dataset of pairwise samples
$\mathcal{S}=\\{(\bm{x}^{(i)},\bm{y}^{(i)})\\}_{i=1}^{M}$ by solving the
forward model, exactly or approximately. Note that the samples in
$\mathcal{S}$ represent realizations from the joint distribution
$p_{\bm{\mathcal{X}}\bm{\mathcal{Y}}}$ which are supplied to a GAN.

A GAN typically comprises two neural subnetworks, namely a generator $\bm{g}$
and a critic $d$. Let $\bm{z}\in\Omega_{\bm{z}}\subset\mathbb{R}^{{N_{z}}}$ be
a latent variable modeled using the random variable $\bm{\mathcal{Z}}$ with a
distribution $p_{\bm{\mathcal{Z}}}$, that is easy to sample from; for instance
a multivariate Gaussian distribution. Then the conditional GAN subnetworks are
given by the mappings

$\bm{g}:\Omega_{\bm{z}}\times\Omega_{\bm{y}}\rightarrow\Omega_{\bm{x}},\qquad
d:\Omega_{\bm{x}}\times\Omega_{\bm{y}}\rightarrow\mathbb{R}.$ (1)

For a given $\bm{y}$, the generator induces a distribution on
$\Omega_{\bm{x}}$. We denote this conditional distribution as
$p^{g}_{\bm{\mathcal{X}}|\bm{\mathcal{Y}}}$. The critic’s role is to
distinguish between true samples $(\bm{x},\bm{y})\sim
p_{\bm{\mathcal{X}}\bm{\mathcal{Y}}}$ and fake samples $(\bm{x}^{g},\bm{y})$,
where $\bm{x}^{g}\sim p^{g}_{\bm{\mathcal{X}}|\bm{\mathcal{Y}}}$.

To train the conditional GAN, we define the following objective function

$\displaystyle\mathcal{L}(d,\bm{g})$
$\displaystyle:=\underset{\begin{subarray}{c}(\bm{x},\bm{y})\sim
p_{\bm{\mathcal{X}}\bm{\mathcal{Y}}}\\\ \bm{z}\sim
p_{\bm{\mathcal{Z}}}\end{subarray}}{\mathbb{E}}\left[d(\bm{x},\bm{y})-d(\bm{g}(\bm{z},\bm{y}),\bm{y})\right]$
(2) $\displaystyle=\underset{\begin{subarray}{c}\bm{x}\sim
p_{\bm{\mathcal{X}}|\bm{\mathcal{Y}}}\\\ \bm{y}\sim
p_{\bm{\mathcal{Y}}}\end{subarray}}{\mathbb{E}}\left[d(\bm{x},\bm{y})\right]-\underset{\begin{subarray}{c}\bm{x}^{g}\sim
p^{g}_{\bm{\mathcal{X}}|\bm{\mathcal{Y}}}\\\ \bm{y}\sim
p_{\bm{\mathcal{Y}}}\end{subarray}}{\mathbb{E}}\left[d(\bm{x}^{g},\bm{y})\right].$

In (2), we begin with the definition of the loss term (also how it is
constructed in practise), and derive the final equality using the definition
of the joint distribution and that of
$p^{g}_{\bm{\mathcal{X}}|\bm{\mathcal{Y}}}$. The cWGAN is trained by solving
the following min-max problem

$(d^{*},\bm{g}^{*})=\underset{\bm{g}}{\arg\min}\ \underset{d}{\arg\max}\
\mathcal{L}(d,\bm{g}).$ (3)

Under the assumption that $d$ is 1-Lipschitz, it can be shown that
$\bm{g}^{*}$ in (3) minimizes the (mean) Wasserstein-1 distance between
$p_{\bm{\mathcal{X}}|\bm{\mathcal{Y}}}$ and
$p^{g}_{\bm{\mathcal{X}}|\bm{\mathcal{Y}}}$

$\bm{g}^{*}=\underset{\bm{g}}{\arg\min}\underset{\begin{subarray}{c}\bm{y}\sim
p_{\bm{\mathcal{Y}}}\end{subarray}}{\mathbb{E}}\left[W_{1}(p_{\bm{\mathcal{X}}|\bm{\mathcal{Y}}},p^{g}_{\bm{\mathcal{X}}|\bm{\mathcal{Y}}})\right],$
(4)

where $p_{\bm{\mathcal{Y}}}$ is the marginal distribution of
$\bm{\mathcal{Y}}$. A proof of this equivalence can be found in [17]. The
Lipschitz constraint on the critic can be weakly imposed using a gradient
penalty term [22] while training $d$ (also see Appendix A).

Convergence in the Wasserstein-1 metric (for $\bm{y}\sim
p_{\bm{\mathcal{Y}}}$) is equivalent to the weak convergence of the two
measures [23], i.e.,

$\underset{\begin{subarray}{c}\bm{x}\sim
p_{\bm{\mathcal{X}}|\bm{\mathcal{Y}}}\end{subarray}}{\mathbb{E}}\left[\ell(\bm{x})\right]=\underset{\begin{subarray}{c}\bm{x}^{g}\sim
p^{g}_{\bm{\mathcal{X}}|\bm{\mathcal{Y}}}\end{subarray}}{\mathbb{E}}\left[\ell(\bm{x}^{g})\right]=\underset{\begin{subarray}{c}\bm{z}\sim
p_{\bm{\mathcal{Z}}}\end{subarray}}{\mathbb{E}}\left[\ell(\bm{g}^{*}(\bm{z},\bm{y}))\right],\quad\forall\
\ell\in C_{b}(\Omega_{\bm{x}}),$ (5)

where $C_{b}(\Omega_{\bm{x}})$ is the space of continuous bounded functions
defined on $\Omega_{\bm{x}}$. This implies that for a given measurement
$\bm{y}$, any statistic of the posterior distribution can be evaluated by
sampling $\bm{z}$ from $p_{\bm{\mathcal{Z}}}$ and pushing it through the
trained generator $\bm{g}^{*}(\cdot,\bm{y})$.

Summarizing, the proposed algorithm to evaluate posterior statistics for a
given problem (characterized by the forward map $\bm{f}$) is:

  1. 1.

Generate or collect $\\{\bm{x}^{(i)}\\}_{i=1}^{M}$ sampled from
$p_{\bm{\mathcal{X}}}$.

  2. 2.

For each $\bm{x}^{(i)}$ solve the forward model to determine $\bm{y}^{(i)}$.
Thus, one constructs the pairwise dataset
$\mathcal{S}=\\{(\bm{x}^{(i)},\bm{y}^{(i)})\\}_{i=1}^{M}$.

  3. 3.

Train the cWGAN on the dataset $\mathcal{S}$ by solving the min-max problem
(3) using a stochastic gradient-based optimizer, such as Adam [24].

  4. 4.

Given a measurement $\bm{y}$, use Monte-Carlo to compute the desired
statistics of the posterior

$\underset{\begin{subarray}{c}\bm{x}\sim
p_{\bm{\mathcal{X}}|\bm{\mathcal{Y}}}\end{subarray}}{\mathbb{E}}\left[\ell(\bm{x})\right]\approx\frac{1}{K}\sum_{i=1}^{K}\ell(\bm{g}^{*}(\bm{z}^{(i)},\bm{y})),\quad\bm{z}^{(i)}\sim
p_{\bm{\mathcal{Z}}},\quad\ell\in C_{b}(\Omega_{\bm{x}}).$ (6)

##  3 The cWGAN architecture

The problems considered in this work assume the underlying two-dimensional
PDEs to be discretized on a uniform grid using a finite difference or finite
element solver. The corresponding $\bm{x}$ and $\bm{y}$ fields arising from
the discretization of the measured and inferred fields can be viewed as images
with a single channel for scalar-valued fields, or multiple channels for
vector-valued fields.

###  3.1 U-Net generator: a multi-scale map

The most popular network architecture to capture image-to-image
transformations is a U-Net. First proposed to efficiently segment biomedical
images [25], a U-Net extracts features from the input image at multiple scales
to construct the multi-scale features of the output field. As shown in Figure
1(a), it takes the input image $\bm{y}$ and pushes it through a contracting
path which extracts the multi-scale features as channels, while downsampling
(coarsening) the resolution of the image. This is followed by an expanding
path that slowly upscales the image to the desired resolution of the output
$\bm{x}$. Through skip connections, the U-Net pulls in the multi-scale
features of $\bm{y}$ learned in the contracting branch, which is concatenated
to the features (of the same spatial dimension) being learned in the expanding
branch. Details about the constitutive blocks of the U-Net are given in
Apendix A.

###  3.2 Conditional instance normalization (CIN): multi-scale stochasticity

A key difference between the generator architecture for the cWGAN we use
compared to one considered in [17] is the way the latent variable $\bm{z}$ is
fed into $\bm{g}$. In [17], $\bm{z}$ is stacked as an additional channel to
the input $\bm{y}$ of the generator. We instead use CIN [19] to inject the
latent information at various scales of the U-Net, as shown in Figure 1(a).
More precisely, an intermediate tensor $\bm{w}$ of shape $H^{\prime}\times
W^{\prime}\times C^{\prime}$ at a given scale of the U-Net, is transformed by
the following normalization for each channel $j=1,\cdots,C^{\prime}$

$CIN(\bm{w},\bm{z})_{j}=\bm{\alpha}(\bm{z})_{j}\otimes\left(\frac{\bm{w}_{j}-\mu(\bm{w}_{j})}{\sigma(\bm{w}_{j})}\right)\oplus\bm{\beta}(\bm{z})_{j},$
(7)

where $\bm{\alpha}:\mathbb{R}^{N_{z}}\rightarrow\mathbb{R}^{C^{\prime}}$ and
$\bm{\beta}:\mathbb{R}^{N_{z}}\rightarrow\mathbb{R}^{C^{\prime}}$ are
(learnable) convolution layers that take $\bm{z}$ as the input. In (7) we have
used the notations $\otimes$ and $\oplus$ respectively to represent element-
wise multiplication and summation in the channel direction, while $\mu(.)$ and
$\sigma(.)$ respectively compute the mean and fluctuations along the spatial
directions. One can interpret the action of CIN in the following manner. Each
feature of the intermediate tensor is decomposed into average value and a
fluctuating component. The precise value of this average and the intensity of
the fluctuation is allowed to be stochastic, depending on the value of the
$\bm{z}$ through the (learned) maps $\bm{\alpha}$ and $\bm{\beta}$.

There are two fundamental advantages of injecting $\bm{z}$ using CIN:

  1. 1.

We can choose the dimension ${N_{z}}$ of the latent variable independently of
the dimension ${N_{y}}$ of $\bm{y}$.

  2. 2.

We introduce stochasticity at multiple scales/levels of the U-Net.

We remark that the CIN strategy has also been effectively used to learn many-
to-many maps with Augmented CycleGANs [26].

###  3.3 Simple critic architecture

In [17], the authors noticed the their generator was insensitive to
variability in $\bm{z}$ (mode collapse), and thus designed a specialized
conditional mini-batch critic to overcome this issue. However, we did not
notice such issues in our experiments when using the simpler critic
architecture shown in Figure 1(b). This can perhaps be attributed to how
$\bm{z}$ is being injected at multiple scales in $\bm{g}$, which is yet
another advantage of using CIN.

(a) Generator (b) Critic

Figure 1: Architecture (Type A) of generator and critic used in the
conditional GAN. The spatial dimension ${N_{x}}=H\times W$ of the input, the
channel parameter $C$, the latent dimension $N_{z}$ and the depth of generator
and critic vary for each experiment.

##  4 A discussion on generalization

In the numerical results presented in Section 5, we observe that the cWGAN
trained using a prior distribution of $\bm{x}$ of a certain type, produces
useful results even when the measurement corresponds to an $\bm{x}$ that is
not selected from this distribution. For example, in the inverse heat
conduction problem, the cWGAN trained using samples where $\bm{x}$ corresponds
to a single circular inclusion, produces reasonable results for measurements
where $\bm{x}$ represents two circular inclusions, or a single elliptical
inclusion. Similarly, for the inverse initial condition problem, the cWGAN
trained on MNIST dataset of handwritten digits produces reasonable results for
the notMNIST dataset which is comprised of digits, letters and other symbols
that are not handwritten. We believe that this “generalizability” learned by
the cWGAN emanates from the fact that the inverse map (the map from $\bm{y}$
to $\bm{x}$) and therefore the conditional distribution
$p_{\bm{\mathcal{X}}|\bm{\mathcal{Y}}}$ may depend only on a subset of the
components of the measurement vector. Then if the cWGAN also learns this
independence, and is perfectly trained for a single value of the complement of
this subset, it learns the true conditional distribution for all measurements.
We prove this assertion in the Theorem below. Thereafter, in the numerical
results section we demonstrate that for the problems we have selected, the
true (regularized) inverse map and the inverse map learned by the generator
satisfy these conditions related to the dependence of the solution only on a
subset of the components of the measurement vector.

######  Theorem 4.1.

Let $\overline{\bm{\mathcal{X}}}$ be a subset of the components of the random
vector $\bm{\mathcal{X}}$ and let $\bm{\mathcal{X}}^{\prime}$ be the
complement of this subset. For this choice of $\overline{\bm{\mathcal{X}}}$,
let $\overline{\bm{\mathcal{Y}}}$ be a subset of the components of the random
vector $\bm{\mathcal{Y}}$ and let $\bm{\mathcal{Y}}^{\prime}$ be the
complement of this subset. Thus, for a given $\overline{\bm{\mathcal{X}}}$, we
may partition
$\Omega_{\bm{y}}=\overline{\Omega}_{\bm{y}}\times\Omega_{\bm{y}}^{\prime}$ and
$\bm{\mathcal{Y}}=(\overline{\bm{\mathcal{Y}}},\bm{\mathcal{Y}}^{\prime})$.
Note that this decomposition depends on the choice of
$\overline{\bm{\mathcal{X}}}$. Let
$p_{\overline{\bm{\mathcal{X}}}|\bm{\mathcal{Y}}}$ denote the conditional
distribution for $\overline{\bm{\mathcal{X}}}$ obtained after marginalizing
$p_{\bm{\mathcal{X}}|\bm{\mathcal{Y}}}$ over $\bm{\mathcal{X}}^{\prime}$.
Further, assume the following conditions hold:

  1. 1.

$p_{\overline{\bm{\mathcal{X}}}|\bm{\mathcal{Y}}}$ is independent of
$\bm{\mathcal{Y}}^{\prime}$. That is for any
$\overline{\bm{y}}\in\overline{\Omega}_{\bm{y}}$ and
$\bm{y}^{\prime}_{1},\bm{y}^{\prime}_{2}\in\Omega_{\bm{y}}^{\prime}$,

$\displaystyle
p_{\overline{\bm{\mathcal{X}}}|\bm{\mathcal{Y}}}(\overline{\bm{x}}|(\overline{\bm{y}},\bm{y}^{\prime}_{1}))=p_{\overline{\bm{\mathcal{X}}}|\bm{\mathcal{Y}}}(\overline{\bm{x}}|(\overline{\bm{y}},\bm{y}^{\prime}_{2})).$
(8)

  2. 2.

The conditional density learned by the generator,
$p^{g}_{\overline{\bm{\mathcal{X}}}|\bm{\mathcal{Y}}}$, inherits this
independence. That is, for any
$\overline{\bm{y}}\in\overline{\Omega}_{\bm{y}}$ and
$\bm{y}^{\prime}_{1},\bm{y}^{\prime}_{2}\in\Omega_{\bm{y}}^{\prime}$,

$\displaystyle
p^{g}_{\overline{\bm{\mathcal{X}}}|\bm{\mathcal{Y}}}(\overline{\bm{x}}|(\overline{\bm{y}},\bm{y}^{\prime}_{1}))=p^{g}_{\overline{\bm{\mathcal{X}}}|\bm{\mathcal{Y}}}(\overline{\bm{x}}|(\overline{\bm{y}},\bm{y}^{\prime}_{2})).$
(9)

  3. 3.

For some fixed value of
$\bm{\mathcal{Y}}^{\prime}=\bm{y}^{\prime}_{a}\in\Omega_{\bm{y}}^{\prime}$,
the conditional density learned by the generator is equal to the true
conditional density. That is for any
$\overline{\bm{y}}\in\overline{\Omega}_{\bm{y}}$,

$\displaystyle
p^{g}_{\overline{\bm{\mathcal{X}}}|\bm{\mathcal{Y}}}(\overline{\bm{x}}|(\overline{\bm{y}},\bm{y}^{\prime}_{a}))=p_{\overline{\bm{\mathcal{X}}}|\bm{\mathcal{Y}}}(\overline{\bm{x}}|(\overline{\bm{y}},\bm{y}^{\prime}_{a})).$
(10)

Then, for any $\bm{y}\in\Omega_{\bm{y}}$,
$p^{g}_{\overline{\bm{\mathcal{X}}}|\bm{\mathcal{Y}}}(\overline{\bm{x}}|\bm{y})=p_{\overline{\bm{\mathcal{X}}}|\bm{\mathcal{Y}}}(\overline{\bm{x}}|\bm{y})$.

###### Proof.

For any $\overline{\bm{y}}\in\overline{\Omega}_{\bm{y}}$, and
$\bm{y}^{\prime}_{1}\in\Omega_{\bm{y}}^{\prime}$, we begin by using (9) with
$\bm{y}^{\prime}_{2}=\bm{y}^{\prime}_{a}$. This yields,

$\displaystyle
p^{g}_{\overline{\bm{\mathcal{X}}}|\bm{\mathcal{Y}}}(\overline{\bm{x}}|(\overline{\bm{y}},\bm{y}^{\prime}_{1}))$
$\displaystyle=$ $\displaystyle
p^{g}_{\overline{\bm{\mathcal{X}}}|\bm{\mathcal{Y}}}(\overline{\bm{x}}|(\overline{\bm{y}},\bm{y}^{\prime}_{a})).$

Thereafter, we use (10) in the equation above to arrive at

$\displaystyle
p^{g}_{\overline{\bm{\mathcal{X}}}|\bm{\mathcal{Y}}}(\overline{\bm{x}}|(\overline{\bm{y}},\bm{y}^{\prime}_{1}))$
$\displaystyle=$ $\displaystyle
p_{\overline{\bm{\mathcal{X}}}|\bm{\mathcal{Y}}}(\overline{\bm{x}}|(\overline{\bm{y}},\bm{y}^{\prime}_{a})).$

Finally, using (8) with $\bm{y}^{\prime}_{2}=\bm{y}^{\prime}_{a}$ in the
equation above, we have

$\displaystyle
p^{g}_{\overline{\bm{\mathcal{X}}}|\bm{\mathcal{Y}}}(\overline{\bm{x}}|(\overline{\bm{y}},\bm{y}^{\prime}_{1}))$
$\displaystyle=$ $\displaystyle
p_{\overline{\bm{\mathcal{X}}}|\bm{\mathcal{Y}}}(\overline{\bm{x}}|(\overline{\bm{y}},\bm{y}^{\prime}_{1})).$

Since any $\bm{y}\in\Omega_{\bm{y}}$ can be written as
$\bm{y}=(\overline{\bm{y}},\bm{y}^{\prime})$, this completes the proof. ∎

In the following section, we consider $\overline{\bm{\mathcal{X}}}$ to be the
subset of components of $\bm{\mathcal{X}}$ centered around a given spatial
location (see Figure 2). For the inverse initial condition and the inverse
thermal conductivity problems, we demonstrate that the true (regularized)
inverse map and the inverse map learned by the cWGAN are local. More
precisely, the value of $\bm{\mathcal{X}}$ at a given spatial location depends
only on the values of the measurements in its vicinity. Therefore, for the
choice of $\overline{\bm{\mathcal{X}}}$ displayed in Figure 2, the
corresponding subset of measurements, $\overline{\bm{\mathcal{Y}}}$, that
influences $\overline{\bm{\mathcal{X}}}$ is comprised of components that are
in the spatial vicinity of $\overline{\bm{\mathcal{X}}}$. Then, using the
result of Theorem 4.1, we conclude that the cWGAN trained using a given
dataset will yield reasonable results for another dataset as long as the
spatially local features in the two sets are similar. This will hold even if
the global features in these sets are different. The generalizability results
obtained in Section 5 are consistent with this conclusion. In each instance,
the cWGAN is trained on data in which the local spatial distribution of the
inferred field is of the form of a sharp transition along a smooth curve. Once
the cWGAN is trained on this data, and then applied to instances where this
local feature is retained, it continues to perform well. This is despite the
fact that the the global features among the two datasets are different.

Figure 2: A schematic representation of the subsets
$\overline{\bm{\mathcal{X}}}$ and $\overline{\bm{\mathcal{Y}}}$ in our
problems.

##  5 Numerical results

We consider inverse problems arising in the context of three distinct PDE
models to demonstrate the performance of the proposed algorithm. The first two
models correspond to the time-dependent and steady-state heat conduction
equations, for which the training and test data are synthetically generated.
The third model solves the inverse elastography problem, where the training
data is generated synthetically but the test data comes from real-life
experiments. A separate cWGAN is trained for each problem, with specific
details about the architecture and hyper-parameters given in Appendix A. We
also show the performance of the trained networks on out-of-distribution (OOD)
test samples, to highlight the generalization capabilities of such an
algorithm. Further, based on the discussion in the previous section, we
explain this generalization by presenting numerical and analytical evidence
that the true and learned inverse maps are local.

###  5.1 Heat conduction: inferring initial condition

Consider the two-dimensional time-dependent heat conduction problem on a
bounded domain $\Omega\subset\mathbb{R}^{2}$ with Dirichlet boundary
conditions

$\displaystyle\frac{\partial u(\bm{s},t)}{\partial
t}-\nabla\cdot(\kappa(\bm{s})\nabla u(\bm{s},t))$
$\displaystyle=b(\bm{s}),\qquad$ $\displaystyle\forall\
(\bm{s},t)\in\Omega\times(0,T)$ (11) $\displaystyle u(\bm{s},0)$
$\displaystyle=u_{0}(\bm{s}),\qquad$ $\displaystyle\forall\ \bm{s}\in\Omega$
(12) $\displaystyle u(\bm{s},t)$ $\displaystyle=0,\qquad$
$\displaystyle\forall\ (\bm{s},t)\in\partial\Omega\times(0,T).$ (13)

Here $u$ denotes the temperature field, $\kappa$ denotes the conductivity
field, $b$ denotes the heat source and $u_{0}$ denotes the initial temperature
field. Given a noisy temperature field at final time $T$, we wish to infer the
initial condition $u_{0}$. This is a severely ill-posed problem as significant
information is lost via the diffusion process when moving forward in time.

We set the spatial domain to be $\Omega=[0,2\pi]^{2}$ and the final time to be
$T=1$. We denote the discretized initial and final (noisy) temperature field
as $\bm{x}$ and $\bm{y}$, respectively, evaluated on a $28\times 28$ Cartesian
grid. The training and test data samples pairs are generated by i) sampling
$\bm{x}$ based on some prior, ii) computing the final temperature field for
the given $\bm{x}$ by solving (11)-(13) using a central-space-backward-time
finite difference scheme, and iii) adding uncorrelated Gaussian noise to
obtain the noisy $\bm{y}$.

####  5.1.1 Rectangular prior

We assume the following rectangular prior on the initial condition

$u_{0}(\bm{s})=\begin{cases}2+2\frac{(s_{1}-\xi_{1})}{(\xi_{3}-\xi_{1})}&\quad\text{if
}s_{1}\in[\xi_{1},\xi_{3}],\ s_{2}\in[\xi_{4},\xi_{2}],\\\
0&\quad\text{otherwise},\end{cases}$ (14)

with $\xi_{1},\xi_{4}\sim\mathcal{U}[0.2,0.4]$ and
$\xi_{2},\xi_{3}\sim\mathcal{U}[0.6,0.8]$. This leads to an initial condition
which has a uniform background value of zero and a linear profile in the
rectangular inclusion whose top-left and bottom-right corners are (randomly)
located at $(\xi_{1},\xi_{2})$ and $(\xi_{3},\xi_{4})$, respectively. The
final temperature fields are obtained by setting $b=0$ and $\kappa=0.64$ in
the PDE model, and adding noise sampled from $\mathcal{N}(\bm{0},\bm{I})$. In
Figure 3, we show a few samples of the discretized $\bm{x}$, the clean and
noisy $\bm{y}$ used to construct the datasets.

(a) Sample 1 (b) Sample 2 (c) Sample 3 (d) Sample 4

Figure 3: Samples from rectangular dataset used to train the cWGAN. The clean
measurements are also shown to contextualize the amount of noise added.

The cWGAN is trained on a dataset with 10,000 training sample pairs, and
tested on a measurement $\bm{y}$ not in the training set. The corresponding
$\bm{x}$, along with reference pixel-wise statistics of posterior (given
$\bm{y}$) are shown in Figure 4. The reference statistics have been computed
by performing Monte Carlo sampling directly in the $\bm{\xi}$ space. Note that
this is computationally feasible only because the dimension of the $\bm{\xi}$
is small (4-dimensional). We observe that the standard deviation (SD) peaks
along the edges of the rectangular region indicating that these are the
regions of highest uncertainty.

A priori, it is not clear what ${N_{z}}$ (the dimension of the latest space)
should be. Since it represents the dimension of the posterior distribution, it
is reasonable to expect ${N_{z}}$ to be much smaller than ${N_{x}}$. We train
different cWGANs by varying ${N_{z}}$ and compute the GAN-based statistics
using (6) with $K=800$. As can be seen in Figure 5, the statistics with all
${N_{z}}$ considered are qualitatively similar to the reference values. In
fact, even ${N_{z}}=1$ gives accurate results, which leads to the maximal
dimension reduction. We also make a quantitative assessment by considering the
the $L^{1}$ error plotted in Figure 6. Although we cannot see a trend as
${N_{z}}$ increases, the error appears to be in the narrow range $0.097\pm
0.019$ for the mean and $0.089\pm 0.014$ for standard deviation.

Figure 4: Test sample with reference mean and SD.

We remark here that the proposed algorithm can approximate the mean of the
posterior, which can be quite different from the ground truth $\bm{x}$ used to
generate the (clean) $\bm{y}$. This difference is quite evident in Figure 4.
In most realistic scenarios, the $\bm{x}$ may not be available for a given
noisy measurement. However, we can use techniques to find the most important
samples (amongst $K$) generated by $\bm{g}$. For instance, we use the reduced
rank QR (RRQR) decomposition of the $K$-snapshot matrix to determine the most
important samples [27]. We show the four most important samples (ranked left
to right) in Figure 7 for ${N_{z}}=3$. This reduced set of four samples
captures a significant amount of the variance observed in the $K=800$ samples
obtained from the generator. We also observe that these samples are quite
distinct from each other, with the first sample being closest to the true
sample.

(a) Mean

(b) SD

Figure 5: Mean and SD computed with 800 samples of $z$, for cWGANs trained
with varying latent space dimension $({N_{z}})$.

(a) Mean (b) SD

Figure 6: $L_{1}$ error in computed statistics compared to the reference
statistics as a function of the number ${N_{z}}$.  Figure 7: Most important
samples ranked left to right using RRQR algorithm on 800 samples for
${N_{z}}=3$.

####  5.1.2 MNIST prior

Next, we assume a more complex prior on $\bm{x}$. We use linearly scaled MNIST
[28] handwritten digits as $\bm{x}$, so that it has a background value of zero
and takes the value $4.0$ on the digit. The final temperature fields are
obtained by setting $b=0$ and $\kappa=0.2$ and adding noise sampled from
$\mathcal{N}(\bm{0},0.3\bm{I})$. In Figure 8, we show a few samples used to
construct the datasets.

(a) Sample 1 (b) Sample 2 (c) Sample 3 (d) Sample 4

Figure 8: Samples from MNIST dataset used to train the cWGAN. The clean
measurements are also shown to contextualize the amount of noise added.

We train a cWGAN with ${N_{z}}=100$ on a dataset with 10,000 training sample
pairs. The trained generator is first used on test samples from the same
distribution as the training set. As can be seen in Figure 9, the mean
computed using the cWGAN ($K$=800) correctly captures the ground truth
$\bm{x}$. The SD is the highest at the boundaries of the digits, indicating a
higher uncertainty in predicting the sharp transition region of initial
temperature field. Furthermore, certain regions where the mean deviates from
$\bm{x}$ (marked by red rectangles) also have high values of SD. This
highlights the utility of computing the SD in that it quantifies the
uncertainty in the inference, and points to the regions in the domain where
the uncertainty is high and the inference less trustworthy.

(a) Sample 1 (b) Sample 2 (c) Sample 3

Figure 9: Inferring initial condition for test samples chosen from the same
distribution as the training set (MNIST prior).

Next, we use the trained generator on OOD test samples where the $\bm{x}$ is
chosen from the linearly scaled notMNIST dataset111Available at:
http://yaroslavvb.blogspot.com/2011/09/notmnist-dataset.html. We observe from
the results shown in Figure 10, that the generator struggles to capture high-
temperature regions closer to the boundary and the broader zones in the
interior. Note that MNIST digits in the training set are more spatially
centered and have narrower features. Nonetheless, the generator is able to
visibly reconstruct the underlying notMINST characters. This indicates the
trained cWGAN is able to generalize reasonably well to OOD samples. Based on
the discussion following Theorem 4.1, we conclude that this surprising
performance of the cWGAN can be attributed to the fact that (a) the true and
learned inverse maps are local (shown below), and (b) the local spatial
features of the OOD and in-distribution datasets are similar (sharp variations
in the initial condition field across smooth spatial curves), even though the
global features are very different.

(a) Sample 1 (b) Sample 2 (c) Sample 3

Figure 10: Inferring initial condition for OOD test samples (notMNIST prior).

To demonstrate the locality of the true inverse map, we assume the solution at
time $T=1.0$ on the domain $[0,2\pi]^{2}$ to be given by a Gaussian bump

$u(\bm{s},T)=\frac{1}{\sqrt{2\pi}\sigma}\exp{\left(-\frac{|\bm{s}-\bm{s}^{0}|^{2}}{2\sigma^{2}}\right)},\quad\sigma=0.7,$
(15)

centered at $\bm{s}^{0}\in[0,2\pi]^{2}$. We then solve the inverse problem
(with $\kappa=0.2$) using an FFT algorithm to obtain the solution at $t=0$.
Since the inverse heat equation is highly ill-posed, we solve a regularized
version of the inverse problem, which involves truncating the high frequency
modes (hyper-diffusion) before the evolving (backwards-in-time) the FFT of
$u(\bm{s},T)$. Further, to avoid instabilities when the bump (15) is centered
close to the domain boundary, we solve the inverse problem on an extended
domain $[-2\pi,4\pi]^{2}$ with a uniform discretization that ensures the
target domain has $28\times 28$ nodes. We found that retaining the first 25
modes (on the extended domain) was sufficient to solve the problem. The
solution profiles at $t=1.0$ and $t=0.0$ for $\bm{s}^{0}=(0,0)$ on the target
domain are shown in Figure 11. As expected we observe that the Gaussian bump
“tightens” as we march backwards in time. To see the local dependence of
$u(\bm{s},0)$ on $u(\bm{s},T)$, we move the bump by considering all $28\times
28$ locations of $\bm{s}^{0}$ and visualize $u(\bm{s},0)$ at fixed locations
$\bm{s}=\bm{s}^{1}$. This is shown in Figure 12, where the red markers
indicate the spatial location $\bm{s}^{1}$ being considered. Clearly, the
$u(\bm{s},T)$ has a weaker influence on $u(\bm{s},0)$ as the center of the
bump $\bm{s}^{0}$ moves further away from $\bm{s}^{1}$. In other words, the
regularized inverse map is local in nature.

(a) $u(\bm{s},T)$ (b) $u(\bm{s},0)$ with FFT

Figure 11: The profiles of $u(T)$ and the corresponding $u(0)$ obtained with
an FFT algorithm, for $\bm{s}^{0}=(0,0)$.

(a) Location 1  (b) Location 2  (c) Location 3  (d) Location 4  (e) Location 5
(f) Location 6

Figure 12: The value of $u(\bm{s}^{1},T)$ at select locations/components
$\bm{s}^{1}$ (marked in red) as $\bm{s}^{0}$ varies in $[0,2\pi]^{2}$.

Next, we demonstrate the the inverse map learned by the generator is also
local. We consider the trained $\bm{g}$ and compute the gradient of the $k$-th
component of the prediction with respect to the network input $\bm{y}$. We
evaluate the magnitude of gradient and average over 100 distinct samples of
$\bm{y}$ and 10 realizations of $\bm{z}$, i.e.,

$\overline{\text{grad}}_{k}=\frac{1}{1000}\sum_{i=1}^{100}\sum_{j=1}^{10}\left|\frac{\partial
g_{k}}{\partial\bm{y}}(\bm{z}^{(j)},\bm{y}^{(i)})\right|,\ \bm{y}^{(i)}\sim
p_{\bm{\mathcal{Y}}},\quad\bm{z}^{(j)}\sim p_{\bm{\mathcal{Z}}},\ 1\leq
k\leq{N_{x}}.$ (16)

The averaged gradients for a few components of $\bm{g}$ are shown in Figure
13. Note that the gradient for each component is concentrated in the
neighbourhood of the corresponding component in $\bm{y}$. In other words, the
domain of influence of $k$-th component of $\bm{g}$’s output is a
neighbourhood of the $k$-th component of $\bm{y}$, but not all of $\bm{y}$.
The local nature of $\bm{g}$ is not unexpected, as most of the operations in
the U-Net architecture are local. Further, the non-local conditional instance
normalization in the intermediate layers does not seem to substantially alter
the local influence $\bm{y}$ on the prediction.

(a) Location 1 (b) Location 2 (c) Location 3

(d) Location 4 (e) Location 5 (f) Location 6

Figure 13: Average component-wise gradient of $\bm{g}$ trained on MNIST data.
The red marker in each tile denotes the component/location of $x$ under
consideration.

###  5.2 Steady-state heat conduction: inferring conductivity

Here we apply our method to infer the conductivity for a steady-state heat
conduction problem when the noisy temperature distribution is known a priori.
We modify (11)-(13) by dropping the time-dependent term and ignoring the
initial condition:

$\displaystyle-\nabla\cdot(\kappa(\bm{s})\nabla u(\bm{s}))$
$\displaystyle=b,\qquad$ $\displaystyle\forall\ \bm{s}\in\Omega$ (17)
$\displaystyle u(\bm{s})$ $\displaystyle=0,\qquad$ $\displaystyle\forall\
\bm{s}\in\partial\Omega.$ (18)

The non-linear relation between conductivity $\kappa$ and the temperature
field $u$ is now dependent only on the spatial variable $\bm{s}$. Further, we
pose this problem on a unit square domain $\Omega=(0,1)^{2}$ and set the
source $b=10$.

The training dataset for this problem consists of 8,000 samples of
($\bm{x}$,$\bm{y}$) pairs. In order to construct each pair, we sample $\kappa$
from a prior distribution and then solve for $u$, from (17) and (18), using
the standard Bubnov-Galerkin approach implemented in FEniCS [29]. We use
triangular elements to discretize the domain $\Omega$ and first-order Lagrange
shape functions to approximate the trial solutions and weighting functions.
Both $\kappa$ and the computed field $u$ are then projected onto a square grid
with $64\times 64$ nodes, to obtain $\bm{x}$ and $\bm{y}$ (without noise)
respectively. Finally, an uncorrelated Gaussian noise
$\eta\sim\mathcal{N}(0,\sigma^{2}\mathbb{I})$ is added to $\bm{y}$. We choose
$\sigma$ to be 2.5% of the maximum temperature value from the entire dataset.
The prior distribution is constructed such that each $\kappa$ field sampled
from this distribution consists of a flat circular inclusion on a constant
background value of unity. The coordinates of the center of the circle,
$(\xi_{1},\xi_{2})$, its radius, $\xi_{3}$, and the contrast $\xi_{4}$ are
sampled from the uniform distributions,

$\xi_{1}\sim\mathcal{U}(0,1),\quad\xi_{2}\sim\mathcal{U}(0,1),\quad\xi_{3}\sim\mathcal{U}(0.05,0.3),\quad\xi_{4}\sim\mathcal{U}(2,10).$

Few samples from the dataset are shown in Figure 14.

The predictions of the trained network (refer appendix A for the architecture
and hyper-parameters) on a few “in-distribution” test samples are shown in
Figure 15 where the mean, SD, and important samples are computed for $K=800$.
In general, we see that the mean is able to capture the location and size of
the target inclusion well. The SD values peak at the edges of the circles
where there is a sharp transition in $\kappa$ (the same behavior was noted in
Section 5.1). We also observe an increase in SD near the corners of the square
domain (see Figure 15(b)) indicating a region of greater uncertainty. This is
attributed to the fact that the measured temperature field along the edges,
and especially at the corners, is determined by the boundary conditions alone
and is not influenced by variations in the thermal conductivity. Once again,
we observe that the important samples determined through RRQR capture the
variations observed in the samples, especially in the contrast between the
inclusion and the background. Of the four samples considered for each case,
the first appears to be the closest to the target field.

(a) Sample 1 (b) Sample 2 (c) Sample 3 (d) Sample 4

Figure 14: Samples from the dataset that was used to train the network for
inferring the conductivity. Each x sample consists of a circular inclusion
with a randomly chosen contrast value.

(a) Sample 1 (b) Sample 2 (c) Sample 3

Figure 15: Inferring conductivity for test samples generated from circular
priors (same distribution as the training set).

Next we investigate the performance of the trained generator on OOD samples
from two other distributions: elliptical inclusions (Figure 16) and two
circular inclusions (Figure 17). For the elliptical inclusions we observe that
the generator correctly locates the centroids in all cases and predicts the
contrast between the inclusion and the background accurately. It also
generates inclusions that are elliptical, however they appear to be more
“circular” when compared to the target field. For the two circular inclusions
we observe that generator is able to generate two distinct inclusions (a case
that is never included in the training set) as long as the inclusions are
distant from each other (Figure 17(b)) or are of comparable size (Figure
17(c)). In the case of two unequal inclusions of very different sizes it
misses the smaller inclusion (Figure 17(a)). Further, it incorrectly infers
that the magnitude of conductivity in the two inclusions is roughly the same.
Overall, similar to the inverse initial condition problem, we conclude that
cWGAN generalizes reasonably well to these OOD datasets.

(a) Sample 1 (b) Sample 2 (c) Sample 3

Figure 16: Inferring conductivity for OOD samples generated with elliptical
priors.

(a) Sample 1 (b) Sample 2 (c) Sample 3

Figure 17: Inferring conductivity for OOD samples involving two circles.

(a) Location 1

(b) Location 2

(c) Location 3

(d) Location 4

(e) Location 5

(f) Location 5

Figure 18: Average component-wise gradient data from the network used for
conductivity inference. The red marker denotes the component/location (of $x$)
under consideration.

In order to better understand this generalizability of the cWGAN we return to
the discussion following Theorem 4.1 in Section 4. Similar to the inverse
problem for the initial condition, in the paragraphs below we demonstrate that
the true and learned inverse maps for the inverse conductivity problem are
local. Thus, Theorem 4.1 tells us that the cWGAN will generalize well to
datasets that have the same local features as the one used for training. It is
clear that this is the case for the OOD and in-distribution datasets
considered in this problem, since at a local level they involve sharp changes
in the thermal conductivity along smooth spatial curves.

To examine the locality of the true inverse map, we perturb the temperature by
a spatially confined perturbation, $\delta u$, and determine the corresponding
perturbation in the inferred thermal conductivity. In order to eliminate the
effect of boundary conditions, we consider the free-space problem. The
perturbed temperature and conductivity fields together satisfy

$\displaystyle\nabla\cdot\big{(}(\kappa+\delta\kappa)\nabla(u+\delta
u)\big{)}=f.$ (19)

Using (17), and retaining only the first order terms, this yields

$\displaystyle\nabla\cdot\big{(}\delta\kappa\nabla
u\big{)}=-\nabla\cdot\big{(}\kappa\nabla\delta u\big{)}.$ (20)

Given the pair $(\kappa,u)$, and a perturbation in the measurement, $\delta
u$, this equation determines the corresponding change in the conductivity,
$\delta\kappa$. A regularized (by diffusion) version of this equation reads,

$\displaystyle\nabla\cdot\big{(}\delta\kappa\nabla
u\big{)}+\rho\nabla^{2}(\delta\kappa)=-\nabla\cdot\big{(}\kappa\nabla\delta
u\big{)},$ (21)

which is an advection-diffusion equation for $\delta\kappa$ with a “velocity
field” given by $\nabla u$ and a diffusion coefficient $\rho$. It is driven by
the source term involving $\delta u$. The formal solution to this problem,
obtained after utilizing the free-space Green’s function,
$g(\bm{s},\bm{s}^{\prime})$, and performing integration-by-parts twice, is

$\displaystyle\delta\kappa(\bm{s})=-\int_{\mathbb{R}^{2}}\nabla\cdot\big{(}\kappa(\bm{s}^{\prime})\nabla
g(\bm{s},\bm{s}^{\prime})\big{)}\delta u(\bm{s}^{\prime})d\bm{s}^{\prime}.$
(22)

Now assuming that $\delta u$ is concentrated around $\bm{s}^{0}$ and
integrates to $\delta U$, the integral above may be approximated as

$\displaystyle\delta\kappa(\bm{s})\approx-\delta
U\big{(}\nabla\kappa(\bm{s}^{0})\cdot\nabla
g(\bm{s},\bm{s}^{0})+\kappa(\bm{s}^{0})\nabla^{2}g(\bm{s},\bm{s}^{0})\big{)}.$
(23)

In two-dimensions, $g(\bm{s},\bm{s}^{0})\propto\ln|\bm{s}-\bm{s}^{0}|$, which
implies that the two terms on the right hand side of (23) have singularities
of the type $|\bm{s}-\bm{s}^{0}|^{-1}$ and $|\bm{s}-\bm{s}^{0}|^{-2}$, and
therefore decay rapidly away from $\bm{s}^{0}$. This implies that
$\delta\kappa$ also decays rapidly away from $\bm{s}^{0}$, and therefore the
true regularized inverse map for this problem is local.

Next, we demonstrate that the inverse map learned by the generator is local.
We compute the gradient of the $k$-th component of the prediction with respect
to the network input $\bm{y}$ (see (16)). The averaged gradients for several
components are shown in Figure 18. Once again we observe that the gradient for
each component is concentrated in the neighbourhood of the corresponding
component in $\bm{y}$ indicating that this map is also spatially local.

###  5.3 Elastography: inferring shear modulus

Elastography is a promising medical imaging technique where ultrasound is used
to image tissue as it is deformed using an external force. The sequence of
ultrasound images thus obtained are used to determine the displacement field
inside the tissue, which can be related to the spatial distribution of the
shear modulus via the equations of equilibrium for an elastic solid,

$\displaystyle\nabla\cdot\bm{\sigma}(\bm{u}(\bm{s}))$ $\displaystyle=0,\qquad$
$\displaystyle\forall\ \bm{s}\in\Omega,$ (24)

and appropriate boundary conditions. Here
$\bm{\sigma}(\bm{(}u))=2\mu(\nabla^{s}\bm{u}+(\nabla\cdot\bm{u})\bm{I})$ is
the Cauchy stress tensor for an incompressible linear isotropic elastic solid
in a state of plane stress, $\bm{u}$ is the displacement field, and $\mu$ is
the shear modulus. For the specific experimental configuration considered in
this example, the domain $\Omega$ is a $34.608\times 26.297$ mm rectangle. The
boundary conditions are traction free in both directions on the left and right
edges, and traction free along the horizontal direction on the top and bottom
edges. In addition, the compression of the specimen is modeled by setting the
vertical component of displacement to 0.084 mm on the top edge and to 0.392 mm
along the bottom edge.

In a typical elastography problem only the displacement component along the
axis of the ultrasound transducer is measured accurately. This corresponds to
the vertical direction in our configuration. Thus the inverse problem is:
given appropriate boundary conditions and the vertical component of the
displacement field, determine the spatial distribution of the shear modulus.
For the mathematical foundations of elasatography and an analaysis of the
uniqueness of this problem, the reader is referred to [30, 31].

The training dataset for this problem consists of 8,000 samples of
($\bm{x}$,$\bm{y}$), which correspond to the discretized values of $\mu$ and
$u_{2}$ (vertical component of displacement), respectively. These samples are
generated by first sampling $\mu$ from a prior distribution and then solving
for $u_{2}$ using the finite element method with linear triangular elements.
Both $\mu$, and the computed field $u_{2}$, are then projected onto a
$56\times 56$ grid, to obtain $\bm{x}$ and $\bm{y}$ (without noise)
respectively. An uncorrelated Gaussian noise
$\eta\sim\mathcal{N}(0,\sigma^{2}\mathbb{I})$, with $\sigma=0.001$ is added to
$\bm{y}$. The prior distribution for $\mu$ is constructed such that each field
consists of a stiff inclusion on a uniform background of 4.7 kPa. The
coordinates of the center of the circle, $(\xi_{1},\xi_{2})$, its radius,
$\xi_{3}$, and the contrast $\xi_{4}$ are sampled from the uniform
distributions:

$\xi_{1}\sim\mathcal{U}(7.1,19.2),\quad\xi_{2}\sim\mathcal{U}(7.1,27.6),\quad\xi_{3}\sim\mathcal{U}(3.5,7),\quad\xi_{4}\sim\mathcal{U}(1,8).$

Some of the resulting training samples are shown in Figure 19.

The experimental data used for this inference problem was obtained from
ultrasound scans on a tissue-mimicking phantom [32]. The phantom was
manufactured using a mixture of gelatin, agar, and oil and consisted of a
stiff inclusion embedded within a softer substrate. The phantom was gently
compressed and the vertical displacement field within the phantom was
measured. This displacement field is shown in the top-left panel of Figure 20.

The results from the shear modulus distribution generated by the cWGAN
conditioned on the measured displacement, are shown in Figure 20. These
include the mean, the SD, and the three important samples as determined by the
RRQR factorization of the matrix of $K=800$ samples generated by the cWGAN. We
observe that the highest uncertainty occurs in a thin layer along the boundary
of the inclusion. We also observe that there is very little variability among
the three important samples. A comparison of the quantitative metrics
calculated from the mean image, and independent experimental measurements is
as follows [32]:

  1. 1.

Vertical distance from the bottom edge to the center of the inclusion:
estimated = 12.0 mm; measured = 13.0 mm; error = 7.7%.

  2. 2.

Diameter of the inclusion: estimated = 10.3 mm; measured = 10.0 mm; error =
3%.

  3. 3.

Shear modulus of the inclusion: estimated = 13.3 kPa; measured = 10.7 kPa;
error = 24.3 %.

These errors are within the range of errors obtained from other deterministic
approaches used to solve this inverse problem [33]. The advantage of the
Bayesian approach is that it provides additional estimates of uncertainty.

(a) Sample 1 (b) Sample 2 (c) Sample 3 (d) Sample 4

Figure 19: Samples from the dataset (circular priors) that was used to train
the network for inferring shear modulus. Figure 20: Inference on an
experimentally measured displacement field.

##  6 Conclusion

In this manuscript, we proposed a deep learning algorithm to solve Bayesian
inference problems where the forward map is modelled using a PDE. Conditional
WGANs are trained on paired samples of the inferred field and measurements,
with the final generator used to generate samples from the posterior
distribution given a noisy (test) measurement. The U-Net architecture of the
generator is able to capture multi-scale features of the measurement to
reconstruct the inferred field. An important component of the proposed deep
learning model is the use of CIN to inject the latent information into the
generator. This has a threefold benefit i) the dimension of the latent space
can be chosen independently of the dimension of the measurement or inferred
field, ii) stochasticity is introduced at multiple scales of the generator,
and iii) the added stochasticity allows the use of a simple critic
architecture without suffering from mode-collapse.

The proposed cWGAN is used to solve inverse problems arising from three
different PDE models. The numerical experiments show that the local values of
SD is a good measure of the uncertainty in the reconstructed fields (see
Figure 9, for example). It is shown that the trained cWGANs can lead to
reasonable predictions when used with OOD measurements. A theoretical
reasoning is provided to explain the generalizability of the networks, which
hinges on the local nature of the learned inverse map and the true
(regularized) inverse map. This locality is demonstrated, either analytically
or empirically, when presenting the numerical results. Finally, we show for
the inverse elasticity imaging problem that the network trained on synthetic
data yields satisfactory reconstructions when tested on experimental
measurements.

The results presented in this work provide a first look into the
generalizability of cWGANs. However, a deeper analysis is warranted to unravel
the full potential of such an approach. Further, it worth exploring the
benefits of embedding the underlying physics into the network architecture, or
the objective function. These will be the focus of future work.

## Acknowledgments

The authors acknowledge support from ARO grant W911NF2010050, NRL grant
N00173-19-P-1119, and AIER (USC), and resources from the Center for Advanced
Research Computing (USC).

## References

  * [1] F. Natterer.  The Mathematics of Computerized Tomography.  Society for Industrial and Applied Mathematics, 2001.
  * [2] Wences P Gouveia and John A Scales.  Resolution of seismic waveform inversion: Bayes versus Occam.  Inverse Problems, 13(2):323–349, apr 1997.
  * [3] Öz Yilmaz.  Seismic Data Analysis: Processing, Inversion, and Interpretation of Seismic Data.  Society of Exploration Geophysicists, 01 2001.
  * [4] Charles Jackson, Mrinal K. Sen, and Paul L. Stoffa.  An efficient stochastic Bayesian approach to optimal parameter and uncertainty estimation for climate model predictions.  Journal of Climate, 17(14):2828 – 2841, 2004.
  * [5] Sixun Huang, Jie Xiang, Huadong Du, and Xiaoqun Cao.  Inverse problems in atmospheric science and their application.  Journal of Physics: Conference Series, 12:45–57, jan 2005.
  * [6] I. Craig and John C. Brown.  Inverse Problems in Astronomy, A guide to inversion strategies for remotely sensed data.  Boston: A. Hilger, 1986.
  * [7] Jacques Hadamard.  Sur les problèmes aux dérivées partielles et leur signification physique.  Princeton university bulletin, pages 49–52, 1902.
  * [8] Masoumeh Dashti and Andrew M. Stuart.  The Bayesian Approach to Inverse Problems, pages 311–428.  Springer International Publishing, Cham, 2017.
  * [9] Heinz Werner Engl, Martin Hanke, and Andreas Neubauer.  Regularization of inverse problems, volume 375.  Springer Science & Business Media, 1996.
  * [10] Jay Whang, Erik Lindgren, and Alex Dimakis.  Composing normalizing flows for inverse problems.  In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 11158–11169. PMLR, 18–24 Jul 2021.
  * [11] Hwan Goh, Sheroze Sheriffdeen, Jonathan Wittmer, and Tan Bui-Thanh.  Solving Bayesian inverse problems via variational autoencoders, 2021\.
  * [12] Gabrio Rizzuti, Ali Siahkoohi, Philipp A Witte, and Felix J Herrmann.  Parameterizing uncertainty by deep invertible networks: An application to reservoir characterization.  In SEG Technical Program Expanded Abstracts 2020, pages 1541–1545. Society of Exploration Geophysicists, 2020.
  * [13] Gregory Ongie, Ajil Jalal, Christopher A Metzler, Richard G Baraniuk, Alexandros G Dimakis, and Rebecca Willett.  Deep learning techniques for inverse problems in imaging.  IEEE Journal on Selected Areas in Information Theory, 1(1):39–56, 2020.
  * [14] Dhruv V. Patel and Assad A. Oberai.  Gan-based priors for quantifying uncertainty in supervised learning.  SIAM/ASA Journal on Uncertainty Quantification, 9(3):1314–1343, 2021.
  * [15] Dhruv V Patel, Deep Ray, and Assad A Oberai.  Solution of physics-based Bayesian inverse problems with deep generative priors, 2021.
  * [16] Liu Yang, Dongkun Zhang, and George E.M. Karniadakis.  Physics-informed generative adversarial networks for stochastic differential equations.  SIAM Journal on Scientific Computing, 42(1):A292–A317, feb 2020\.
  * [17] Jonas Adler and Ozan Öktem.  Deep Bayesian inversion.  arXiv preprint arXiv:1811.05910, 2018.
  * [18] Yizhou Qian, Mojtaba Forghani, Jonghyun Harry Lee, Matthew Farthing, Tyler Hesser, Peter Kitanidis, and Eric Darve.  Application of deep learning-based interpolation methods to nearshore bathymetry, 2020.
  * [19] Vincent Dumoulin, Jonathon Shlens, and Manjunath Kudlur.  A learned representation for artistic style, 2017.
  * [20] Jonas Adler and Ozan Öktem.  Solving ill-posed inverse problems using iterative deep neural networks.  Inverse Problems, 33(12):124007, dec 2017.
  * [21] T Kadeethum, D O’Malley, J N Fuhg, Y Choi, J Lee, H S Viswanathan, and N Bouklas.  A framework for data-driven solution and parameter estimation of PDEs using conditional generative adversarial networks.  Nature Computational Science, 1, 12 2021.
  * [22] Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville.  Improved training of wasserstein gans.  In Advances in neural information processing systems, pages 5767–5777, 2017.
  * [23] Cédric Villani.  Optimal Transport: Old and New.  Grundlehren der mathematischen Wissenschaften. Springer Berlin Heidelberg, 2008.
  * [24] Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization, 2017.
  * [25] Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In Nassir Navab, Joachim Hornegger, William M. Wells, and Alejandro F. Frangi, editors, Medical Image Computing and Computer-Assisted Intervention – MICCAI 2015, pages 234–241, Cham, 2015. Springer International Publishing.
  * [26] Amjad Almahairi, Sai Rajeshwar, Alessandro Sordoni, Philip Bachman, and Aaron Courville.  Augmented cyclegan: Learning many-to-many mappings from unpaired data.  In International Conference on Machine Learning, pages 195–204. PMLR, 2018.
  * [27] Tony F Chan and Per Christian Hansen.  Some applications of the rank revealing qr factorization.  SIAM Journal on Scientific and Statistical Computing, 13(3):727–741, 1992.
  * [28] Yann LeCun, Corinna Cortes, and CJ Burges.  Mnist handwritten digit database.  ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist, 2, 2010.
  * [29] Martin Sandve Alnæs, Jan Blechta, Johan E. Hake, August Johansson, Benjamin Kehlet, Anders Logg, Chris N. Richardson, Johannes Ring, Marie E. Rognes, and Garth N. Wells.  The FEniCS project version 1.5.  Archive of Numerical Software, 3, 2015.
  * [30] Paul E Barbone and Assad A Oberai.  Elastic modulus imaging: some exact solutions of the compressible elastography inverse problem.  Physics in Medicine & Biology, 52(6):1577, 2007.
  * [31] Paul E Barbone and Assad A Oberai.  A review of the mathematical and computational foundations of biomechanical imaging.  Computational Modeling in Biomechanics, pages 375–408, 2010.
  * [32] Theo Z Pavan, Ernest L Madsen, Gary R Frank, Jingfeng Jiang, Antonio AO Carneiro, and Timothy J Hall.  A nonlinear elasticity phantom containing spherical inclusions.  Physics in Medicine & Biology, 57(15):4787, 2012.
  * [33] Jean Francois Dord, Sevan Goenezen, Assad A Oberai, Paul E Barbone, Jingfeng Jiang, Timothy J Hall, and Theo Pavan.  Validation of quantitative linear and nonlinear compression elastography.  Ultrasound Elastography for Biomedical Applications and Medicine, pages 129–142, 2018.

##  Appendix A cWGAN loss and architecture

The cWGANs in this work were trained using the following loss function

$\displaystyle\mathcal{L}(d,\bm{g})$
$\displaystyle=\underset{\begin{subarray}{c}(\bm{x},\bm{y})\sim
p_{\bm{\mathcal{X}}\bm{\mathcal{Y}}}\\\ \bm{z}\sim
p_{\bm{\mathcal{Z}}}\end{subarray}}{\mathbb{E}}\left[d(\bm{x},\bm{y})-d(\bm{g}(\bm{z},\bm{y}),\bm{y})+\mathcal{G}\mathcal{P}\right],$
$\displaystyle\mathcal{G}\mathcal{P}$
$\displaystyle=\lambda\underset{\begin{subarray}{c}\epsilon\sim\mathcal{U}(0,1)\end{subarray}}{\mathbb{E}}\left[(\|\partial_{1}d(\bm{h}(\bm{x},\bm{y},\bm{z},\epsilon),\bm{y})\|_{2}-1)^{2}\right],$

which has been augmented with a gradient penalty term to constrain the critic
to be 1-Lipschitz on $\Omega_{\bm{x}}$. Here, $\mathcal{U}(0,1)$ denotes the
uniform distribution on $[0,1]$, $\partial_{1}d(.,.)$ denotes the derivative
with respect to the first argument of the critic, and

$\bm{h}(\bm{x},\bm{y},\bm{z},\epsilon)=\epsilon\bm{x}+(1-\epsilon)\bm{g}(\bm{z},\bm{y}).$

For all experiments considered in the present work, the gradient penalty
parameter is set as $\lambda=10$.

As shown in Figures 1, the architecture of the generator and critic comprises
various network blocks. We describe the key components below.

  * •

Conv($n,s,k$) denotes a 2D convolution with $k$ filters of size $n$ and stride
$s$. If $n>1$, reflective padding of width 1 is applied in the spatial
dimensions of the input before applying the convolution. Furthermore, if the
third argument is absent, the number of filter is taken as equal to the number
of channels in the input tensor.

  * •

The CIN block is used to inject the latent variable $\bm{z}$ into different
levels of the generator’s U-Net architecture. Its action (channel-wise) is
given by (7).

  * •

Res block denotes a residual block, which is shown in Figure 21(c).

    * –

When appearing in the generator, it takes as input an intermediate tensor
$\bm{w}$ and the latent variable $\bm{z}$. The latent variable is injected
into the residual block using conditional instance normalization (CIN).

    * –

When appearing in the critic, it takes as input only the intermediate tensor
$\bm{w}$, with CIN replaced by layer normalization.

    * –

In certain cases, when explicitly mentioned, no normalization is used. This
typically happens towards the beginning of the generator and critic networks.

  * •

Down($k$) denotes the down-sampling block shown in Figure 21(a), where $k$
denotes the factor by which the number of input channels increases. This block
makes use of 2D average pooling to reduce the spatial resolution by a factor
of 2. When appearing in the critic, the block only takes as input the
intermediate tensor $\bm{w}$.

  * •

Up($k$) denotes the up-sampling block shown in Figure 21(b). It receives the
output $\bm{w}$ from the previous block and concatenates it (unless specified)
with the output $\widetilde{\bm{w}}$ of a down-sampling block of the same
spatial size through a skip connection. The Up($k$) block reduces the number
of output channels to $C^{\prime}/k$, where $C^{\prime}$ denotes the number of
channels in the input $\bm{w}$. Unless specified, 2D nearest neighbour
interpolation is used to increase the spatial resolution by a factor of 2.

  * •

Dense($k$) denotes a fully connected layer of width $k$

The size of the generator and critic used for the various problems are as
follows:

  * •

Inferring initial condition: A smaller architecture is used, compared to what
is shown in Figure 1. More precisely, the U-Net is shallower, with one less
Down(2) block and Up(2) block. The critic has one less Down(2) block. For both
datasets, we set $H=W=28$ and $C=32$.

  * •

Inferring conductivity: The architecture is as shown in Figure 1 with the only
difference being that CIN was not used in the first Up(1) layer. Further, we
set $H=W=64$ and $C=32$.

  * •

Inferring shear modulus: The architecture is as shown in Figure 1, with
$H=W=56$ and $C=32$.

(a) Down(k) (b) Up(k) (c) Res block

Figure 21: Key components used to build the generator and critic networks.

The networks are trained on TensorFlow, using the Adam optimizer for 1000
epochs with $\beta_{1}=0.5$, $\beta_{2}=0.9$ and learning-rates $10^{-3}$. The
remaining hyper-parameters are listed in Table 1. Each network takes less than
20 hours to train on a single NVIDIA Tesla P100 GPU.

Table 1: Hyper-parameters for cWGAN Inferred field |  | Initial condition
---
| Conductivity
---
| Sheer modulus
---
Training Data | Rectangular | MNIST | Circles | Circles
Training samples | 10,000 | 10,000 | 8000 | 8000
${N_{x}}$ | $28\times 28$ | $28\times 28$ | $64\times 64$ | $56\times 56$
${N_{z}}$ | Multiple | 100 | 50 | 50
Batch size | 50 | 50 | 64 | 64
Activation param. | 0.1 | 0.1 | 0.2 | 0.2
$n_{\text{critic}}/n_{\text{gen}}$ | 4 | 4 | 5 | 5