# Efficient Prior Calibration From Indirect Data

O. Deniz Akyildiz  Mark Girolami  Andrew M. Stuart  Arnaud
Vadeboncoeur111Corresponding author (av537@cam.ac.uk).

###### Abstract

Bayesian inversion is central to the quantification of uncertainty within
problems arising from numerous applications in science and engineering. To
formulate the approach, four ingredients are required: a _forward model_
mapping the unknown parameter to an element of a solution space, often the
solution space for a differential equation; an _observation operator_ mapping
an element of the solution space to the data space; a _noise model_ describing
how noise pollutes the observations; and a _prior model_ describing knowledge
about the unknown parameter before the data is acquired. This paper is
concerned with learning the prior model from data; in particular, learning the
prior from multiple realizations of indirect data obtained through the noisy
observation process. The prior is represented, using a generative model, as
the pushforward of a Gaussian in a latent space; the pushforward map is
learned by minimizing an appropriate loss function. A metric that is well-
defined under empirical approximation is used to define the loss function for
the pushforward map to make an implementable methodology. Furthermore, an
efficient residual-based neural operator approximation of the forward model is
proposed and it is shown that this may be learned concurrently with the
pushforward map, using a bilevel optimization formulation of the problem; this
use of neural operator approximation has the potential to make prior learning
from indirect data more computationally efficient, especially when the
observation process is expensive, non-smooth or not known. The ideas are
illustrated with the Darcy flow inverse problem of finding permeability from
piezometric head measurements.

Keywords: inverse problems, generative models, prior learning, operator
learning, differential equations

##  1 Introduction

###  1.1 Set-Up

This paper is concerned with learning a generative model for unobserved
$\\{z^{(n)}\\}_{n=1}^{N}$ from indirect and noisy data
$\\{y^{(n)}\\}_{n=1}^{N}$ given by

(1) $\displaystyle y^{(n)}={\mathcal{G}}(z^{(n)})+\varepsilon^{(n)}_{\eta},$

where the noise $\varepsilon^{(n)}_{\eta}\sim\eta$ i.i.d. . The setting where
${\mathcal{G}}(\cdot)$ is the identity and the noise is zero is the standard
problem of generative modeling and is well studied; in the context of this
paper, Bayesian inversion, such a generative model may be used to construct a
prior measure from prior samples $\\{z^{(n)}\\}_{n=1}^{N}$. However, in many
science and engineering applications the prior samples
$\\{z^{(n)}\\}_{n=1}^{N}$ are not directly observed, but the
$\\{y^{(n)}\\}_{n=1}^{N}$, arising from (1) with choice ${\mathcal{G}}=g\circ
F^{\dagger}$ originating from multiple instantiations of physical systems are
available. Thus, we are interested in the setting where

(2) $\displaystyle y^{(n)}=g\circ
F^{\dagger}(z^{(n)})+\varepsilon^{(n)}_{\eta}.$

Here, $F^{\dagger}:Z\mapsto U$ maps between function spaces $Z,U$,
representing a PDE parameter-to-solution operator, and
$g:U\mapsto{\mathbb{R}}^{d_{y}}$ is a solution-to-data map. From the resulting
finite dimensional data we wish to construct a generative model for a prior
measure on a function space, giving rise to the unobserved prior samples
$\\{z^{(n)}\\}_{n=1}^{N}$.

To overview our approach to this problem we first describe it at a population
loss level. Let $\nu$ denote the law of the $y^{(n)}$ and $\mu$ the desired
law of the $z^{(n)}$. Letting $*$ denote convolution of measures and $\\#$
denote pushforward, introducing divergence $\mathsf{d}_{1}$ between
probability measures in data space and $\mathcal{H}(\cdot)$ the regularization
term on measures in the input space, we define population loss

(Functional 1)
$\displaystyle\mathsf{J}_{1}(\mu)=\mathsf{d}_{1}\Bigl{(}\nu,\eta*(g\circ
F^{\dagger})_{\\#}\mu\Bigr{)}+{\mathcal{H}}(\mu).$

To develop algorithms to find $\mu$ we will represent $\mu$ as pushforward
under map $T^{\alpha}$ of Gaussian measure $\mu_{0}$ on a latent
space.222Other generative models replacing the Gaussian with different, but
also straightforward to sample, measures can easily be accommodated; we choose
a Gaussian in the latent space to make the presentation explicit. Here
$\alpha\in{\mathbb{R}}^{d_{\alpha}}$ represents a finite dimensional
parameterization of the pushforward from $\mu_{0}$ to $\mu$. We replace the
regularization term $\mathcal{H}(\cdot)$ on measure $\mu$ by regularization
term $h:{\mathbb{R}}^{d_{\alpha}}\to{\mathbb{R}}$ on $\alpha.$ We then
consider the modified population loss

(Functional 2)
$\displaystyle\mathsf{J}_{2}(\alpha)=\mathsf{d}_{1}\Bigl{(}\nu,\eta*(g\circ
F^{\dagger}\circ T^{\alpha})_{\\#}\mu_{0}\Bigr{)}+h(\alpha).$

We also observe that $F^{\dagger}$ may be expensive to compute and it may be
desirable to replace it by a neural operator $F^{\phi}$ whose parameters
$\phi\in{\mathbb{R}}^{d_{\phi}}$ need to be learned so that $F^{\phi}\approx
F^{\dagger}$ in sets of high probability under $\mu.$ But we do not know
$\mu$, indeed we are trying to find it; thus the optimal parameters $\phi$
will depend on $\alpha$ and be defined by $\phi=\phi^{*}(\alpha)$. We thus
introduce loss function

(Functional 3)
$\displaystyle\mathsf{J}_{3}(\alpha)=\mathsf{d}_{1}\Bigl{(}\nu,\eta*(g\circ
F^{\phi^{*}(\alpha)}\circ T^{\alpha})_{\\#}\mu_{0}\Bigr{)}+h(\alpha).$

At the heart of all these loss functions is a matching of distributions. In
practice both $\nu$ and the pushforward of $\mu_{0}$ will only be available
empirically and so it is necessary that the divergence $\mathsf{d}_{1}$ can be
readily evaluated on empirical measures. Empiricalized versions of the
functionals $\mathsf{J}_{2},\mathsf{J}_{3}$ will form the basis of
computational methodology proposed in this paper. The mapping
$\phi^{*}(\alpha)$ will also be learned using minimization of a loss function,
involving matching of distributions and evaluated empirically. The functional
$\mathsf{J}_{1}$ provides a theoretical underpinning of our approach; and in
the case $N=1$ will be linked, in the empirical setting, to Bayes Theorem.

Subsection 1.2, which follows, summarizes our contributions and outlines the
paper. In Subsection 1.3 we review relevant literature in the area and in
Subsection 1.4 we overview the notation and model problem (Darcy) used in the
paper.

###  1.2 Contributions and Outline

The proposed novel methodology allows the construction of a calibrated measure
over the parameters underlying a PDE model, given data from a collection of
physical systems. Furthermore, the methodology may be combined with a novel
concurrent neural operator approximation of the PDE. Together these ideas hold
the potential to improve the accuracy and efficiency of Bayesian inversion and
of generative modeling for physical systems. The work can be broken down into
five primary contributions:

  1. 1.

We introduce a suitable choice for divergence $\mathsf{d}_{1}$ based on sliced
Wasserstein-2 distance and demonstrate that it leads to computationally
feasible objective $\mathsf{J}_{2}$ (Functional 2).

  2. 2.

We introduce a residual-based probabilistic loss function to define choice of
parameters $\phi^{*}(\alpha)$ in the neural operator approximation.

  3. 3.

With this definition of $\phi^{*}(\alpha)$ we demonstrate a computationally
feasible objective $\mathsf{J}_{3}$ (Functional 3).

  4. 4.

We show that, with our choice of $\mathsf{d}_{1}$, minimization of Functional
1 may be linked to Bayes Theorem when $N=1$.

  5. 5.

In order to be concrete we describe our methodology in the context of the
Darcy flow model of porous medium flow which may be viewed as a mapping from
the permeability field ($z$) to linear functionals of the piezometric head
($y$). Numerical experiments with Darcy flow, for two different choices of
pushforward families $T^{\alpha}$, are used to demonstrate feasibility and
consistency of the proposed methodology.

In Section 2 we describe the efficient residual-based approach to operator
learning that we adopt in this paper, addressing contribution 2. Section 3
introduces specific divergences for definition of Functionals 1–3, and for the
residual-based learning, addressing contributions 1. and 3. ; we also address
contribution 4, in Theorem 3.2. Section 4 discusses algorithmic details,
giving further detail on contributions 1. and 3. whilst Section 5 implements
the algorithms on the Darcy flow problem, contribution 5. We conclude and
discuss future works in Section 6.

###  1.3 Literature Review

We overview relevant literature, First, we discuss the learning of priors from
data. Secondly we describe the surrogate modeling literature. And finally, we
discuss related transport-based inference methods.

####  1.3.1 Learning Priors

The task of selecting a best prior has received much attention in the Bayesian
statistics community, with many objectives and ideals in mind [32]. Certain
efforts concentrate on the careful formulation of uninformative priors [33, 5,
6]. Others focus on mathematical tractability through conjugacy [51] or
eliciting priors from domain experts [48]. Some see the prior as an
opportunity to share information from observations with possibly different
underlying parameters, that are assumed to be drawn from the same
distribution. Taking this point of view, are methods related to hierarchical
Bayes [28], empirical Bayes [54, 55] and parametric empirical Bayes [15, 44].
In many ways, the idea of using a set of data to explicitly target an
unconditional distribution is the basis of many modern generative modeling
methods in the field of machine learning (ML). These include score-based
models [64], diffusion models [62, 30], variational auto-encoders [36],
energy-based models [63, 67], normalizing flows [49], gradient-flows [10], and
more. Data-based unconditional distributions like these have been shown to be
not only expressive generative models, but also powerful priors for Bayesian
inversion [25, 11]. These ideas appear in a number of recent papers focussed
on inverse-problems [50, 3, 2]. A key difference in our work and the basis for
our contribution is that, while above works focus on learning (and using)
priors from which unconditional samples are available, our methodology here
extends this idea to learn priors from indirect data. This idea of indirect
knowledge of priors is exploited in [27] for linear operators where the set of
Bayesian inverse problems defined individually by (1), for each $n$, is solved
in the case where ${\mathcal{G}}(\cdot)=A\cdot$ for some linear operator $A$;
the collection of inverse problems is used to both learn prior information and
learn the dependence of each individual posterior on data.

####  1.3.2 Surrogate Modeling and Operator Learning

In many engineering applications related to design optimization, parameter
inversion, and forward uncertainty quantification tasks, it is necessary to
evaluate numerical models many times. When using a classical numerical scheme,
there is no information carry-over from one numerical solve to the next.
Hence, there is room for improvement in the form of somehow interpolating
information from the one numerical solution to the next. The idea of replacing
computer code with a cheap-to-evaluate statistical interpolation model is well
explored [56, 34]. Such methods have found their place as viable model order
reduction techniques for multi-query problems. Modern advances have now begun
to pose the task of surrogate modeling directly in function-space, resulting
in the field of operator learning [8, 42, 38, 39]. These methods are based on
gathering input-output datasets of PDE parameters and PDE solutions obtain via
classical numerical schemes, such as finite element models (FEM) and spectral
methods. Physics informed surrogate modeling attempts to directly incorporate
PDE information into the learning task [53, 73, 40]. The methods may then be
data-less or semi-data-informed. Certain classes of methods also pose prior
distributions over PDE parameters and attempt to learn a surrogate trained on
random draws from that prior [68, 69].

####  1.3.3 Pushforwards and Wasserstein Losses

Minimization of regularized loss functions over the space of probability
measures, as exemplified by Functional 1, is central to modern computational
statistics and machine learning: it lies at the heart of variational inference
[71] and in many other emerging inference problems for probability measures
[18, 41]. In statistical inference and ML, there are many learning objective
functions to make use of. A well known method is maximizing the marginal
likelihood. However, the task we are interested in for this work, is
distributional learning. Hence, we must use a statistical divergence between
measures. A common choice is the KL divergence. However this is not useful
when both measures being compared are empirical, due to the non-overlapping
support of sampled Diracs, as is the case in this paper.

Noting this restriction on suitable divergences we proceed to identify
appropriate choices. The (MMD) maximum mean divergence [61, 23] and (closely
related) energy distances [65, 59] are one possible class metrics that could
be used. In this work, we focus on the use of computationally tractable
optimal transport based metrics on the space of measures; in particular we use
the sliced-Wasserstein metric [9]. Many other works have explored the use of
these optimal transport (OT) based metrics for ML inference tasks [37, 21, 45,
47, 41]. There is also work on combining pushforward measures and Bayesian
inference in [14, 13, 43, 24]. Solving inverse problems with Wasserstein loss
is also a partially explored topic in [1] and conditional flow matching [17].
Some works have also looked at Wasserstein metrics between pushforward
measures [57]. Wasserstein and sliced-Wasserstein metrics have found use in
approximate Bayesian computation (ABC) [66, 4, 46] where the recovered measure
have been shown to converge to the Bayesian posterior. Parameter estimation
with Wasserstein metrics in a purer form is explored in [7]. In discussing OT
based learning objective, it is important to mention entropy regularization,
as in the Sinkhorn algorithm [19]. In these works, entropy regularization is
put on the finite-dimensional observational space. We will see that entropy
regularization comes up differently in our proposed methodology.

###  1.4 Notation

Let $D\subset\mathbb{R}^{d}$ be bounded and open. We denote the boundary of
the set $D$ by $\partial D$. We use the $L^{p}(D)$ classes of $p-$th power
Lebesgue integrable functions, $1\leq p<\infty$, extending to $p=\infty$ in
the usual way via the essential supremum. We denote by $C^{\infty}(D)$ the set
of infinitely differentiable functions, and by $H^{1}_{0}(D)$ the Sobolev
space of functions with one square-integrable weak derivative and homogeneous
Dirichlet boundary conditions; we denote by $H^{-1}(D)$ the dual of
$H^{1}_{0}(D)$ with respect to the canonical pairing through Lebesgue
integration over $D$.

Let ${\mathcal{P}}(X)$ denote the space of probability measures on measurable
space $X$. Divergences on the space of probability measures are denoted by
$\mathsf{d}$, sometimes with subscript $i\in\\{1,2\\}.$ We denote by
$f_{\\#}\mu$ the pushforward measure given by $f_{\\#}\mu(A)=\mu(f^{-1}(A))$
for all $\mu$ measurable sets $A$. We denote indexing over different instances
of variables in a collection (such as a dataset or set of randomly sampled
variables) with a superscript in parenthesis, accessing elements of a vector
are done through a subscript in parenthesis, and incrementing (such as in
summations) is done with a plain subscript. We denote by
$\langle\cdot,\cdot\rangle_{A}=\langle\cdot,A^{-1}\cdot\rangle$ the covariance
weighted inner-product, for any positive self-adjoint $A$, with induced norm
$\|\cdot\|_{A}.$

Consider (2) and assume that $z^{(n)}\sim\mu^{\dagger}\in\mathcal{P}(Z)$ where
$Z\subset L^{\infty}(D)$ is separable. Let $F^{\dagger}:Z\to H^{1}_{0}(D)$ and
let $g:H_{0}^{1}(D)\to\mathbb{R}^{d_{y}}$ denote a set of functionals. If we
assume that $\varepsilon^{(n)}_{\eta}\sim\eta$ for $n=1,\ldots,N$ are i.i.d
noise variables then $y^{(n)}\sim\eta*(g\circ
F^{\dagger})_{\\#}\mu^{\dagger}.$ The problem of interest is to recover from
the observations $\\{y^{(n)}\\}_{n=1}^{N}$ the law $\mu^{\dagger}$ of the
parameter field.

To be concrete we will work with inverse problems defined by the Darcy
equation

(3a) $\displaystyle\nabla\cdot(z\nabla u)+f$ $\displaystyle=0,\quad\forall
x\in\text{D},$ (3b) $\displaystyle u$ $\displaystyle=0,\quad\forall
x\in\partial D.$

Here $z$ denotes permeability and $u$ the piezometric head. The mapping
$z\mapsto u$ may be viewed as mapping $F^{\dagger}:Z\to H^{1}_{0}(D),$ for
appropriately defined $Z$. To be concrete we fill focus on this setting, and
the problem of determining a prior on $z$ from noisy linear functionals of $u$
defined by mapping
$g:H^{1}_{0}(D)\to{\mathbb{R}}^{d_{y}}\cong\bigl{(}H^{-1}(D)\bigr{)}^{d_{y}}.$
The reader will readily see that the ideas in the paper apply more generally,
and that consideration of the Darcy problem is simply for expository purposes.

##  2 Residual-Based Neural Operator

This section is devoted to defining $\phi^{*}(\alpha)$ which appears in
Functional 3. Recall that $\phi$ are the parameters of the neural network
surrogate PDE model, and that their optimization depends on the underlying
input measure on which the surrogate needs to be accurate; for this reason
they depend on $\alpha$ the set of parameters characterizing the prior we seek
to determine. To define function $\phi^{*}(\alpha)$ we proceed as follows. Let
$F^{\phi}:Z\to H^{1}_{0}(D)$ be a parametric family of maps approximating
$F^{\dagger}$. Now define residual operator $R:Z\times H_{0}^{1}(D)\to
H^{-1}(D)$. Intuitively, we can express our PDE as $R(z,u)=0.$ In the case of
the Darcy equation (3) we may write

$R(z,u)=\nabla\cdot(z\nabla u)+f,$

and we have $u\in H_{0}^{1}(D)$, $z\in Z\subset L^{\infty}(D)$, and $f\in
L^{2}(D)$. Note next that we can write $u=F^{\dagger}(z)$ by the definition of
forward map, hence we have $R(z,F^{\dagger}(z))=0$ for all $z\in Z$. In order
to incorporate this information into our loss functional, we define
$R^{\phi}:Z\to H^{-1}(D)$ as

(4) $\displaystyle R^{\phi}(z)=R(z,F^{\phi}(z)),$

where $F^{\phi}$ is the parametric family of maps. In order to obtain a
computable loss, we introduce a discretization operator
$\mathcal{O}:H^{-1}(D)\to\mathbb{R}^{d_{o}}$, where $d_{o}$ is the dimension
of the output of $\mathcal{O}$. In particular, given a set of basis functions
$\\{v_{i}\\}_{i=1}^{d_{o}}$ for $H_{0}^{1}(D)$, we can write333Here
$\langle\cdot,\cdot\rangle$ denotes the duality pairing between $H^{1}_{0}(D)$
and $H^{-1}(D)$

(5) $\displaystyle\mathcal{O}(R)_{i}=\langle
v_{i},R\rangle=\int_{D}v_{i}R(z,u)(x)\mathrm{d}x,$

for $i=1,\ldots,d_{o}$, for any given $v_{i}\in H_{0}^{1}(D)$, noting that
integration by parts may be used to to show well-definedness in the given
function space setting. To compactly represent the discretization process with
our emulator $F^{\phi}$, we define $\mathcal{O}^{\phi}:X\to\mathbb{R}^{d_{o}}$
as

(6) $\displaystyle\mathcal{O}^{\phi}(z)=\mathcal{O}(R^{\phi}(z)).$

We are now in a position to define our final loss functional using the
constructions above. To learn the optimal parameter $\phi$ as $\alpha$ varies,
we define the following coupled loss functional and associated minimization
problem:

(Functional 4) $\displaystyle\mathsf{J}_{4}(\phi;\alpha)$
$\displaystyle=\mathsf{d}_{2}\Bigl{(}\delta_{0},(\mathcal{O}^{\phi}\circ
T^{\alpha})_{\\#}\mu_{0})\Bigr{)},$ $\displaystyle\phi^{\star}(\alpha)$
$\displaystyle=\underset{\phi}{\mathrm{argmin}}\;\mathsf{J}_{4}(\phi;\alpha);$

here $\mathsf{d}_{2}$ is a divergence term between probability measures on
$\mathbb{R}^{d_{o}}$ and $\delta_{0}$ is the Dirac measure at zero. Minimizing
$\mathsf{J}_{4}(\cdot;\alpha)$ for given $\alpha$ determines a residual-based
approximation of $F^{\dagger}$, accurate with respect to
$(T^{\alpha})_{\\#}\mu_{0}$. Using this expression shows that Functional 3,
Functional 4 define a bilevel optimization scheme [60]. This scheme is the
heart of our proposed methodology. Using the bilevel approach avoids balancing
the contributions of Functional 3 and Functional 4 that arise from an additive
approach. A similar bilevel optimization scheme is employed in [72], for
similar reasons, to solve a different problem. In the next section we provide
specific examples $\mathsf{d}_{1}$ and $\mathsf{d}_{2}$ and discuss some of
their properties.

##  3 Choice of Divergences

In order to instantiate the bilevel optimzation scheme defined by Functional
3, Functional 4 to obtain an implementable algorithm, we now define specific
choices of the divergence terms $\mathsf{d}_{1}$ and $\mathsf{d}_{2}$
(Subsection 3.1) and discuss empirical approximation of the input measures
(Subsection 3.2) required to evaluate these divergences in practice. And, to
further establish a context for our work on learning priors, we make a
connection to Bayesian inversion in the setting of equation (1) when $N=1$
(Subsection 3.3.)

###  3.1 Divergences

In this work, to effectively and efficiently compare empirical measures, we
will use the sliced-Wasserstein (SW) distance to define the divergence term
for $\mathsf{d}_{1}$ and Wasserstein distance for $\mathsf{d}_{2}$. To define
precisely what we do, we first introduce the weighted Wasserstein distance

(7)
$\displaystyle\mathsf{W}_{2,\mathrm{B}}^{2}(\nu,\mu)=\inf_{\gamma\in\Pi(\nu,\mu)}\int_{\mathbb{R}^{d}\times\mathbb{R}^{d}}\|x-y\|_{\mathrm{B}}^{2}{\mathrm{d}}\gamma(x,y),$

where coupling $\Pi(\nu,\mu)$ is the set of all joint probability measures on
$\mathbb{R}^{d}\times\mathbb{R}^{d}$ with marginals $\nu$ and $\mu$, and
$\mathrm{B}$ is positive and self-adjoint. We let
$\mathsf{W}_{2}:=\mathsf{W}_{2,\mathrm{I}}.$ The following lemma shows that
the weighted squared Wasserstein-2 distance can be seen as the squared
Wasserstein-2 distance of pushforwards of the original measures.

######  Lemma 3.1.

For $P_{\mathrm{B}}(\cdot)=\mathrm{B}^{-1/2}\,\cdot$ it follows that

(8)
$\displaystyle\mathsf{W}^{2}_{2,\mathrm{B}}(\nu,\mu)=\mathsf{W}^{2}_{2}(P_{\mathrm{B}\\#}\nu,P_{\mathrm{B}\\#}\mu).$

###### Proof.

See Appendix A.1. ∎

We now define the weighted and sliced-Wasserstein distance by

(9)
$\displaystyle\mathsf{SW}^{2}_{2,\mathrm{B}}(\nu,\mu)=\int_{\mathbb{S}^{d-1}}\mathsf{W}_{2}^{2}(P^{\theta}_{\mathrm{B}\\#}\nu,P^{\theta}_{\mathrm{B}\\#}\mu){\mathrm{d}}\theta,$

where
$P^{\theta}_{\mathrm{B}}(\cdot)=\langle\mathrm{B}^{-\frac{1}{2}}\cdot,\theta\rangle.$
The sliced-Wasserstein distance leads to a computationally efficient
alternative to the Wasserstein distance because it may be implemented by Monte
Carlo approximation of integration over $\theta$ and then each slice,
resulting from a randomly chosen $\theta$, involves only evaluation of a
Wasserstein distance between probabilities in ${\mathbb{R}}.$

Assuming that $\eta=\mathcal{N}(0,\Gamma)$ we then consider the losses
Functional 3, Functional 4 with
$\mathsf{d}_{1}(\cdot,\cdot)=\frac{d_{y}}{2}\mathsf{SW}^{2}_{2,\Gamma}(\cdot,\cdot)$
and $\mathsf{d}_{2}(\cdot,\cdot)=\mathsf{W}^{2}_{2}(\cdot,\cdot),$ so that

(10a) $\displaystyle\mathsf{J}_{3}(\alpha)$
$\displaystyle=\frac{d_{y}}{2}\mathsf{SW}^{2}_{2,\Gamma}\Bigl{(}\nu,\eta*(g\circ
F^{\phi^{*}(\alpha)}\circ T^{\alpha})_{\,\\#}\mu_{0}\Bigr{)}+h(\alpha),$ (10b)
$\displaystyle\mathsf{J}_{4}(\phi;\alpha)$
$\displaystyle=\mathsf{W}_{2}^{2}\Bigl{(}\delta_{0},(\mathcal{O}^{\phi}\circ
T^{\alpha})_{\,\\#}\mu_{0}\Bigr{)},$ (10c) $\displaystyle\phi^{\star}(\alpha)$
$\displaystyle=\underset{\phi}{\mathrm{argmin}}\;\mathsf{J}_{4}(\phi;\alpha).$

###  3.2 Empiricalization

The optimization problem in (10) forms the basis of our computational
methodology in this paper. However, to implement it, we need to use empirical
approximations of the two input measures that define the loss function
$\mathsf{J}_{3}.$ There are two ways in which evaluation of $\mathsf{J}_{3}$
defined by (10) must be empiricalized to make a tractable algorithm:

  * •

measure $\nu$ is replaced by measure

$\nu^{N}=\frac{1}{N}\sum_{n=1}^{N}\delta_{y^{(n)}},$

reflecting the fact that working with $\nu$ is not computationally tractable,
but samples from $\nu$ are available;

  * •

measures $\nu^{N}$ and $(g\circ F^{\phi^{*}(\alpha)}\circ
T^{\alpha})_{\,\\#}\mu_{0}$ are replaced by their empirical approximations
using, in each case, $N_{s}$ independent samples, reflecting the fact that
working with $(g\circ F^{\phi^{*}(\alpha)}\circ T^{\alpha})_{\,\\#}\mu_{0}$ is
not computationally tractable, together with the computational simplicity
arising from using the same number of samples in each argument of the
divergence;

  * •

measure $(\mathcal{O}^{\phi}\circ T^{\alpha})_{\,\\#}\mu_{0}$ is empiricalized
using $N_{r}$ independent samples from $\mu_{0}$.

###  3.3 Connection to Bayes Theorem

Consider the inverse problem defined by (1) in setting $N=1:$

(11) $\displaystyle y={\mathcal{G}}(z)+\varepsilon,$

where $\varepsilon\sim\eta:=\mathcal{N}(0,\Gamma)$ and
${\mathcal{G}}:Z\to{\mathbb{R}}^{d}.$ Consider Functional 1 with
$\nu=\delta_{y}$, ${\mathcal{G}}(\cdot)$ replacing $(g\circ
F^{\dagger})(\cdot)$444There is nothing intrinsic to the factorization
${\mathcal{G}}=g\circ F^{\dagger}$ in this subsection; results are expressed
purely in terms of ${\mathcal{G}}$. and choice of $\mathsf{d}_{1}$ as in
Subsection 3.1:

(12)
${\mathsf{J}}(\mu)=\frac{d}{2}\mathsf{SW}^{2}_{2,\Gamma}\Bigl{(}\delta_{y},\eta*({\mathcal{G}})_{\\#}\mu\Bigr{)}+{\mathcal{H}}(\mu).$

Thus we have returned to the population loss description of the problem in the
setting where optimization to determine the prior is over all probability
measures, not the parameterized family that we will use for computation. Thus
the regularization is on the space of probability measures. The message of the
following theorem is that, with this problem formulation on the space of
measures, in the case $N=1$ and with appropriate choice of regularization,
Bayes theorem is recovered.

######  Theorem 3.2.

Define $\mathcal{H}(\mu):=D_{\mathrm{KL}}(\mu||\mu^{\mathrm{prior}})$ for some
probability measure $\mu^{\mathrm{prior}}$ $\in$ ${\mathcal{P}}(Z).$ Consider
the Bayesian inverse problem defined by (11) with $z\sim\mu^{\mathrm{prior}}$
independent of $\varepsilon\sim\eta:=\mathcal{N}(0,\Gamma).$ Then the
minimizer of (12) over the set of probability measures
{$\mu\in{\mathcal{P}}(Z)$ :
${\mathbb{E}}_{y^{\prime}\sim\eta*({\mathcal{G}}_{\\#})\mu}\|y^{\prime}\|_{\Gamma}^{2}<\infty\\}$
is the Bayesian posterior given by

(13a) $\displaystyle{\mathrm{d}}\mu^{y}(A)$
$\displaystyle=\frac{1}{\mathcal{Z}}\int_{A}\exp\left(-\frac{1}{2}\|y-{\mathcal{G}}(z)\|^{2}_{\Gamma}\right){{\mathrm{d}}\mu^{\mathrm{prior}}(z)},$
(13b) $\displaystyle\mathcal{Z}$
$\displaystyle=\int_{Z}\exp\left(-\frac{1}{2}\|y-{\mathcal{G}}(z)\|^{2}_{\Gamma}\right){{\mathrm{d}}\mu^{\mathrm{prior}}(z)}.$

######  Remark 3.3.

As discussed in the introduction to this paper, our method learns a data-
informed probability measure – a generative model – for $z$, which may be used
as a prior for downstream inference tasks. The theorem shows that, in the
special case of $N=1$, the learned prior actually coincides with the Bayesian
posterior for the inverse problem defined by (1), if regularizer $\mathcal{H}$
is chosen appropriately. This is entirely consistent with our broader agenda
when $N>1$ as the posterior distribution is the natural prior for downstream
tasks in Bayesian inference.

To prove this theorem we establish a sequence of lemmas. The first shows that
the weighted Wasserstein-2 distance may be simplified when one of its argument
is a Dirac:

######  Lemma 3.4.

For any $y\in\mathbb{R}^{d}$ and $\mu\in\mathcal{P}(\mathbb{R}^{d})$

(14)
$\displaystyle\mathsf{W}^{2}_{2,\Gamma}(\delta_{y},\mu)=\mathbb{E}_{x\sim\mu}\|y-x\|_{\Gamma}^{2}.$

###### Proof.

See Appendix A.2. ∎

Using the definition of pushforward, it follows from Lemma 3.4 that

(15)
$\displaystyle\mathsf{W}_{2,\Gamma}^{2}(\delta_{y},{\mathcal{G}}_{\\#}\mu)=\mathbb{E}_{z\sim\mu}\|y-{\mathcal{G}}(z)\|_{\Gamma}^{2}=\int_{\mathbb{R}^{d}}\|y-{\mathcal{G}}(z)\|_{\Gamma}^{2}{\mathrm{d}}\mu(z).$

The following lemma shows that the sliced-Wasserstein metric also simplifies
when one of its argument is a Dirac:

######  Lemma 3.5.

Let $y\in\mathbb{R}^{d}$ and $\mu\in\mathcal{P}(\mathbb{R}^{d})$, and assume
${\mathbb{E}}_{z\sim\mu}\|z\|_{\Gamma}^{2}<\infty$. Then

(16)
$\displaystyle\mathsf{SW}_{2,\Gamma}^{2}(\delta_{y},\mu)=\frac{1}{d}\mathsf{W}_{2,\Gamma}^{2}(\delta_{y},\mu).$

###### Proof.

See Appendix A.3 ∎

The final lemma concerns convolution:

######  Lemma 3.6.

Let $\mathsf{I}:{\mathcal{P}}(Z)\to{\mathbb{R}}$ be defined by
$\mathsf{I}(\cdot)=\mathsf{W}_{2,\Gamma}^{2}(\delta_{y},\cdot)$. Let $\eta$ be
a centred random variable and finite second moment. Then

$\displaystyle\mathsf{I}(\eta*\mu)=\mathsf{I}(\mu)+{\mathbb{E}}_{\varepsilon\sim\eta}\|\varepsilon\|_{\Gamma}^{2}.$

###### Proof.

See Appendix A.4. ∎

###### Proof of Theorem 3.2.

By equation (15) and Lemmas 3.5 and 3.6 it follows that the loss function (12)
can be written as

(17)
$\displaystyle\mathsf{J}(\mu)=\frac{1}{2}\mathbb{E}_{z\sim\mu}\|y-{\mathcal{G}}(z)\|^{2}_{\Gamma}+\mathcal{H}(\mu)+C_{1},$

where $C_{1}$ is independent of $\mu.$ With the choice
$\mathcal{H}(\mu):=D_{\mathrm{KL}}(\mu||\mu^{\mathrm{prior}})$ we see that
$\mathsf{J}(\mu)=D_{\mathrm{KL}}(\mu||\mu^{y})+C_{2}$ where $C_{2}$ is
independent of $\mu$, and is determined by $C_{1}$ and $\mathcal{Z}.$ The
result follows since $D_{\mathrm{KL}}(\cdot||\mu^{y})$ is minimized at
$\mu^{y}.$ ∎

######  Remark 3.7.

In order to use Lemmas 3.5 and 3.6 in the proof of Theorem 3.2, we assume the
pushforward measure $\eta*({\mathcal{G}}_{\\#}\mu)$ has
$\|\cdot\|_{\Gamma}^{2}$-moment finite as stated in Theorem 3.2. This will
follow from second moments of ${\mathcal{G}}_{\\#}\mu.$ Such results can often
be established using the Fernique theorem; see for example [16, 20].

######  Remark 3.8.

Now consider functional $\mathsf{J}(\cdot)$ given by (12), replacing
${\mathcal{G}}$ by the specific choice $g\circ F^{\dagger}$ arising from the
inverse problem defined by $\eqref{eq:2}$ in the case where $N=1.$ If we now
parameterize target measure $\mu$ by $\mu=T^{\alpha}_{\\#}\mu_{0}$, make the
choice of KL divergence for $\mathcal{H}$ and view $\mathsf{J}(\cdot)$ as
parameterized by $\alpha\in{\mathbb{R}}^{d_{\alpha}}$ rather than by
$\mu\in{\mathbb{P}}(Z),$ we obtain

$\displaystyle{\mathsf{J}}(\alpha)=\frac{d}{2}\mathsf{SW}^{2}_{2,\Gamma}\Bigl{(}\delta_{y},\eta*(g\circ
F^{\dagger}\circ
T^{\alpha})_{\\#}\mu_{0}\Bigr{)}+D_{\mathrm{KL}}((T^{\alpha})_{\\#}\mu_{0}||\mu^{\mathrm{prior}}).$

By the same reasoning used in the proof of Theorem 3.2 we may rewrite this as

$\displaystyle{\mathsf{J}}(\alpha)=\frac{1}{2}\mathbb{E}_{z\sim(T^{\alpha})_{\\#}\mu_{0}}\|y-(g\circ
F^{\dagger})(z)\|^{2}_{\Gamma}+D_{\mathrm{KL}}((T^{\alpha})_{\\#}\mu_{0}||\mu^{\mathrm{prior}}).$

From this form it is clear that minimizing ${\mathsf{J}}(\alpha)$ over
$\alpha\in{\mathbb{R}}^{d_{\alpha}}$ is simply the variational Bayes
methodology applied to approximate the Bayesian posterior $\mu^{y}$ given by
equation (13). Furthermore this provides motivation for the consideration of
Functional 2 to determine the prior on $z$, in the case where $N>1$, noting
that it reduces to variational Bayes when $N=1$ with appropriate choices of
$\mathsf{d}_{1}(\cdot,\cdot)$ and $h(\cdot).$

##  4 Algorithms

Algorithm 1 Prior Calibration

1: Initialize $\alpha_{0}$, $T$ (number of iterations), $N_{s}$ (number of
samples for $\mathsf{J}_{2}$), $F^{\dagger}$ (forward model)

2: for $t=1,\ldots,T$ do

3: for $i=1,\ldots,N_{s}$ do

4: Sample $z^{(i)}\sim\mu^{\alpha_{t-1}}$

5: Sample $\varepsilon_{\eta}^{(i)}\sim\eta$

6: Compute ${y^{\prime}}^{(i)}=g\circ
F^{\dagger}(z^{(i)})+\varepsilon_{\eta}^{(i)}$

7: Sample $y^{(i)}\sim\nu^{N}$

8: Compute $\mathsf{J}_{2}(\alpha_{t-1})$ from the $N_{s}$ samples
$\\{y^{(i)},{y^{\prime}}^{(i)}\\}_{i=1}^{N_{s}}$.

9: $\alpha_{t}=\text{OPTIMISER}(\alpha_{t-1},\mathsf{J}_{2}(\alpha_{t-1}))$

10: return $\alpha^{\star}\leftarrow\alpha_{T}$.

In this section we discuss practical aspects pertaining to the implementation
of the proposed methodologies. Firstly, although Remark 3.8 suggest a specific
choice for regularization $h(\alpha)$ in Functional 2 or Functional 3; in
practice it is not computationally straightforward to work with this choice
and simpler choices are made. Secondly, the evaluation of the objective
functions is carried out using the empiricalization described in Subsection
3.2. Thirdly the physics-based residual can be evaluated using Lemma 3.4 and
the considerations of Subsection 3.2. Fourthly, once loss functions are
evaluated we obtain gradients with respect to prior (and operator) parameters
through back-propagation with standard ML library tools, such as JAX [12].

In Algorithm 1 we show how to implement the proposed inference methodology of
Functional 2 for the task of prior calibration given a forward model
$F^{\dagger}$. In Algorithm 2 we show how to implement Functional 3 for joint
prior calibration and operator learning. We note that to correctly implement
the bilevel optimization scheme proposed to minimize Functional 3 we must
differentiate through the lower-level optimization steps given by Functional 4
to take into account the dependence of $\phi^{\star}(\alpha)$ on $\alpha$.
This incurs additional computational costs; efficient approximation methods
for this task will be explored in future works.

We note, only the outer-most parameter update loops must be computed
sequentially. All other loops can be efficiently computed in parallel in a GPU
efficient manner. In all experiments we make use of the Adam optimizer [35].
The sliced-Wasserstein implementation is based on that of [26]. If one is
interested in applying this methodology to nonlinear PDEs, computing
Functional 2 requires the solution of $N_{s}$ nonlinear systems of equations
at every parameter update step. However, this is not the case for the method
using Functional 3, Functional 4 as we only need to compute $N_{r}$ residuals,
which is done the same manner for linear or nonlinear PDEs. (See, also,
Subsection 3.2 for definitions of $N_{s},N_{r}.$)

Algorithm 2 Prior Calibration and Operator Learning

1: Initialize $\alpha_{0},\phi_{0}$, $T$ (number of iterations of outer-loop
$\alpha$ updates), $L$ (number of inner-loop $\phi$ updates), $N_{s}$ (number
of samples for $\mathsf{J}_{3}$), $N_{r}$ (number of samples for
$\mathsf{J}_{4}$)

2: for $t=1,\ldots,T$ do

3: for $l=1,\ldots,L$ do

4: for $j=1,\ldots,N_{r}$ do

5: Sample $z^{(j)}\sim\mu^{\alpha_{t-1}}$

6: Compute $r^{(j)}=\mathcal{O}^{\phi_{l-1}(\alpha_{t-1})}(z^{(j)})$

7: Compute $\mathsf{J}_{4}(\alpha_{t-1},\phi_{l-1}(\alpha_{t-1}))$ from the
$N_{r}$ samples $\\{r^{(j)}\\}_{j=1}^{N_{r}}$.

8:
$\phi_{l}(\alpha_{t-1})=\text{OPTIMISER}(\phi_{l-1}(\alpha_{t-1}),\mathsf{J}_{4}(\alpha_{t-1},\phi_{l-1}(\alpha_{t-1}))$

9: $\phi^{\star}(\alpha_{t-1})\leftarrow\phi_{L}(\alpha_{t-1})$

10: for $i=1,\ldots,N_{s}$ do

11: Sample $z^{(i)}\sim\mu^{\alpha_{t-1}}$

12: Sample $\varepsilon_{\eta}^{(i)}\sim\eta$

13: Compute ${y^{\prime}}^{(i)}=g\circ
F^{\phi^{\star}(\alpha_{t-1})}(z^{(i)})+\varepsilon_{\eta}^{(i)}$

14: Sample $y^{(i)}\sim\nu^{N}$

15: Compute $\mathsf{J}_{3}(\alpha_{t-1},\phi^{\star}(\alpha_{t-1}))$ from the
$N_{s}$ samples $\\{y^{(i)},{y^{\prime}}^{(i)}\\}_{i=1}^{N_{s}}$.

16:
$\alpha_{t}=\text{OPTIMISER}(\alpha_{t-1},\mathsf{J}_{3}(\alpha_{t-1},\phi^{\star}(\alpha_{t-1})))$

17: return
$\\{\alpha^{\star}\leftarrow\alpha_{T},\phi^{\star}(\alpha^{\star})\leftarrow\phi^{\star}(\alpha_{T})\\}$.

##  5 Numerical Results for Darcy Flow

We now illustrate the performance of our methodology to learn generative
models for priors, based on indirect observations. All our numerical examples
are in the setting of equation (3). The presented PDE is to be interpreted in
the weak sense; see Appendix B for a short discussion on the weak form and
approximate computational methods exploiting GPU-efficient array shifting
operations. We consider two classes of coefficient function $z:$ the first is
a class of piecewise constant functions with discontinuity sets defined as
level sets of a smooth field, in Subsection 5.1; the second is a class of
functions defined as the pointwise exponential of a smooth field, in
Subsection 5.2. In both cases we will write the prior as pushforward of a
Gaussian measure, on the underlying smooth field, under a parameter-dependent
map $T^{\alpha}$; we refer to the first class as _level set priors_ and the
second as _lognormal priors_. We attempt to learn the parameters $\alpha$
through minimization of either Functional 2 (using numerical PDE solver) or
Functional 3 (using residual-based operator learning); in all cases we use the
framework of Section 3. We denote the true measure which we wish to recover by
$\mu^{\dagger}.$

Throughout we use regularization
$h(\alpha)=\frac{1}{2\sigma_{h}^{2}}\|\log(\alpha)-m_{h}\|^{2}_{2}$ in which
the log is applied component-wise; we refer to $\sigma_{h}$ and $m_{h}$ as the
standard-deviation and mean of the regularizer term. For all experiments, we
use 1000 slicing directions to evaluate the sliced-Wasserstein term, and we
minimize the relevant loss functions using Adam [35] with a learning rate of
$10^{-2}$ decayed by half four times on the parameters $\alpha$, other than
where explicitly stated. All experiments are run on a 24GB NVIDIA RTX 4090
GPU.

The numerical results we present substantiate the following conclusions:

  * •

The proposed methodology using Functional 2 can effectively recover the true
parameters of a level set prior and a lognormal prior.

  * •

The proposed methodology using Functional 2 can effectively recover the true
parameters of a level set prior despite using a smoothening of the level set
formulation, so that there is prior mis-specification.

  * •

We highlight the impact of dataset size ($N$), and the number of samples used
to estimate the loss ($N_{s}$), on the quality of the prior parameter
estimation; this demonstrates the strength of distributional inference for
prior calibration.

  * •

The proposed methodology using Functional 3, Functional 4 can achieve
comparable parameter estimation accuracy to that under Functional 2; thus
jointly estimating the operator approximation and the parameters of the prior
is both feasible and, for reasons of efficiency, will be desirable in settings
where $F^{\dagger}$ is expensive to evaluate, is not differentiable or is not
available.

  * •

We show how, under certain circumstances, prior parameters may be
unidentifiable, and we show limitations of the proposed methodology in this
setting.

###  5.1 Level Set Priors

(a) Normalized Spectrum, $\beta=4$ (b) Smoothed level set function

Figure 1: (a) Spectrum decay of the square root of the eigen values for the
covariance operator (20) (which corresponds to the standard deviaiton of the
basis expansion coefficient of each mode $\varphi_{j}(x)$) for $\beta=4$ and
different choices of lengthscale parameter $\lambda$. As $\lambda$ increases,
more modes play a significant role. (b) A comparison of smooth level set
function $\tilde{z}$ and sharp level set function $z$ transitioning at $x=0.5$
for $\tau=100$.

Prior construction here is based on the methodology introduced in [22, 31]. We
define $\alpha=\\{\kappa^{-},\kappa^{+},\lambda\\}$ where
$\kappa^{\pm}\in(0,\infty)$ and $\lambda>0.$ We now define a function
$z:C^{\infty}(D)\times\mathbb{R}^{2}\to L^{\infty}(D)$ which will take a
$\lambda-$dependent function in $C^{\infty}(D)$, and the two values
$\kappa^{\pm}\in\mathbb{R}^{2}$ to create a function in $L^{\infty}(D).$ We
introduce parameterized family of measures $(T^{\alpha})_{\\#}\mu_{0}$ on
$L^{\infty}(D)$ by pushing forward a Gaussian measure on $a\in C^{\infty}(D)$
under map $T^{\alpha}$ to define a measure on $z.$ To this end we define $z$
by

(18) $\displaystyle z(x)$
$\displaystyle=\kappa^{-}\mathbbm{1}_{D^{-}(a)}(x)+\kappa^{+}\mathbbm{1}_{D^{+}(a)}(x),$
(19) $\displaystyle D^{-}(a)$ $\displaystyle=\\{x\in D|a(x)<0\\},\quad
D^{+}(a)=\\{x\in D|a(x)\geq 0\\}.$

To complete the description of $T^{\alpha}$ we construct $a$ as a
$\lambda-$dependent Gaussian random field. To do this we fix a scalar
$\beta>0.$ The construction of $a$ differs in details between dimensions one
and two. For a 1D physical domain $D$, we define the $\lambda-$dependent
Gaussian measure on $a$ through the Karhunen-Loève expansion

$\displaystyle a(x;\lambda,\beta)$
$\displaystyle=\sum_{j=1}^{J}\Bigl{(}j^{2}\pi^{2}+\lambda^{2}\Bigr{)}^{-\beta/2}\varepsilon_{j}\varphi_{j}(x),\quad\varepsilon_{j}\sim{\mathcal{N}}(0,1)\;\mathrm{i.i.d,}$

where $\varphi_{j}(x)=\cos(j\pi x).$ In dimension two we generalize to obtain

$\displaystyle a(x;\lambda,\beta)$
$\displaystyle=\sum_{j=1,k=1}^{J,K}\Bigl{(}(j^{2}+k^{2})\pi^{2}+\lambda^{2}\Bigr{)}^{-\beta/2}\varepsilon_{j,k}\varphi_{j,k}(x),\quad\varepsilon_{j,k}\sim{\mathcal{N}}(0,1)\;\mathrm{i.i.d},$

with $\varphi_{jk}(x)=\cos(j\pi x_{(1)})\cos(k\pi x_{(2)}).$ For all
experiments, we fix $\beta=4$.

(a) several $u$ and $y$ (b) Random fields $\bar{a}$ (c) Sharp $z$ (d) Smooth
$z$, $\tau=10$

Figure 2: (a)-(c) Data Generated from sharp prior $\mu^{\dagger}$, which are
the PDE solution fields with the observational data at points, the normalized
GRF $\bar{a}$ samples, and one example of a level set function resulting from
one of the $\bar{a}$ fields. (d) The same diffusivity field $\tilde{z}$ as in
(c), but smoothed with (21) and $\tau=10$.

Note that, for $J=\infty$ in dimension one (resp. $(J,K)=\infty$ in dimension
two) $a\sim{\mathcal{N}}(0,C_{\lambda,\beta})$ where

(20) $\displaystyle
C^{-1}_{\lambda,\beta}=\bigl{(}-\Delta+\lambda^{2}\mathrm{I}\bigr{)}^{\beta},$

and $-\Delta$ is the Laplacian equipped with homogeneous Neumann boundary
conditions. Draws from this measure have Hölder regularity up to exponent
$\beta-d/2$ where $d$ is dimension of domain $d.$ This statement about the
probability measure from which $a$ is drawn is approximate when $J$ is (resp.
$(J,K)$ are) finite. Because of the countable nature of the construction of
the level set prior, it lies in a separable subspace $Z$ of $L^{\infty}(D).$
The derivative of the objective functional with respect to $\lambda$ is not
well behaved for the sharp level set setup as $\partial_{a}{z}$ is zero almost
everywhere. Hence we introduce a smoothened level set parametrization of
$z(\alpha)$ as

(21) $\displaystyle\tilde{z}(x;\tau)$
$\displaystyle=\frac{1}{2}\tanh(\tau\,\bar{a})(\kappa^{+}-\kappa^{-})+\frac{1}{2}(\kappa^{+}-\kappa^{-})+\kappa^{-},$
(22) $\displaystyle\bar{a}$ $\displaystyle=a/{\|a\|_{L^{2}(D)}},$

where the parameter $\tau$ controls the sharpness of transition from
$\kappa^{-}$ to $\kappa^{+}$. In Fig. 1(a) we show the spectrum decay of the
GRF generated from the covariance operator (20) in dimension one. Figure 1(b)
shows the smoothing of the sharp level set function $z$ and its spatial
derivative. In Fig. 2 we show 10 sampled random fields $\bar{a}$ and the
associated PDE solutions for the 1D Darcy problem. We only show one diffusion
field $z$ for clarity.

(a) Convergence $\mu^{\dagger}$ (b) Loss $\mu^{\dagger}$

(c) Convergence $\tilde{\mu}^{\dagger}$ (d) Loss $\tilde{\mu}^{\dagger}$

Figure 3: Comparison of convergence of $\alpha=\\{\kappa^{\pm},\lambda\\}$ for
data generated from the physically realistic sharp $\mu^{\dagger}$ and the
smoothed $\tilde{\mu}^{\dagger}$ for 1D Darcy with Functional 2. We plot in
(a) and (c) the convergence of the parameters themselves, and in (b) and (d)
the loss function values.

We note that there is a symmetry in inference between $\kappa^{+}$ and
$\kappa^{-}$, so after convergence we sort the $\kappa^{\pm}$ to associate
$\kappa^{+}$ with the larger value and vice versa for $\kappa^{-}$. With this
experimental setup established, we can test the proposed methodology.
Subsections 5.1.1 and 5.1.2 concern the setting where we use a PDE solve for
the forward model, and hence minimize Functional 2; we consider dimensions
$d=1$ and $d=2$ respectively. In Subsections 5.1.3 and 5.1.4 we combine
learning of $\alpha$ with operator learning to replace the PDE solve, and
hence minimize Functional 3, Functional 3; we again consider dimensions $d=1$
and $d=2$ respectively.

####  5.1.1 Prior Calibration: 1D Darcy

In this section we focus our attention on using Functional 2 to infer
parameters of a prior $(T^{\alpha})_{\\#}\mu_{0}$ given an empirical measure
$\nu^{N}$ and a forward operator $F^{\dagger}$.

(a) Convergence for $\lambda$ (b) Convergence for $\kappa^{+}$

(c) Convergence for $\kappa^{-}$

Figure 4: Converged parameter estimation for $\lambda,\kappa^{+},\kappa^{-}$
individually, for different dataset sizes $N$, and number of samples $N_{s}$
for Functional 2 on the 1D Darcy problem.

We use $N=1000$ data $y$ each of $d_{y}=50$ noisy pointwise observations of
the solution with $\Gamma=0.01^{2}\mathrm{I}$ observational noise covariance.
The prior KL expansion is truncated at 20 terms. We use $N_{s}=1000$ samples
to evaluate Functional 2. The true parameters for data generation are
$\lambda=8,\kappa^{+}=2,\kappa^{-}=1$ and we set the level set smoothing
parameter $\tau=10$. We use a finite element mesh of 100 nodes with a weak
form residual computed through array shifting and solved with conjugate
gradients and the forcing function $f$ is set to a constant value of 10 for
all examples. No regularizer $h(\alpha)$ is used for the 1D example. We
minimize using the smoothed model without sharp interfaces but we derive data
both with the sharp and smoothed interfaces. In both cases we are able to
accurately recover all three hyper parameters in $\alpha$: see Fig. 3. The
runtime for these experiments is 41 seconds for the full 1k iterations. The
relative error of the predicted parameters for the smoothed level set prior
are $0.56\%$, $0.28\%$, and $0.96\%$ for $\kappa^{+}$, $\kappa^{-}$, and
$\lambda$ respectively. We observe that the mis-specification between
$\mu^{\dagger}$ (draws are discontinuous) and $(T^{\alpha})_{\\#}\mu_{0}$
(draws are smooth) is reflected in a small inference bias, as would be
expected. In Fig. 4 we plot the prior parameter estimation mean and standard
deviations for 100 random initializations, varying the number of samples
$N_{s}$ used for estimating the loss functions and the size of the dataset
($N$). The $\lambda$ parameters are initialized from
$\log\lambda\sim\mathrm{Unif}(\log 0.5,\log 4)$,
$\log\kappa^{-}\sim\mathrm{Unif}(\log 0.5,\log 4)$, and
$\log\kappa^{-}\sim\mathrm{Unif}(\log 6,\log 10)$. We observe that the mean of
the recovered parameters roughly converges for this problem setup after a
dataset size $N=100$ and number of samples $N_{s}=100$.

####  5.1.2 Prior Calibration: 2D Darcy

In this section we generalize the setting of the previous section two the two-
dimensional domain $D=[0,1]\times[0,1]$. As in the 1D example, we use $N=1000$
data $y$ each of $d_{y}=50$ noisy pointwise observations of the solution with
$\Gamma=0.01^{2}\mathrm{I}$ observational noise covariance. The prior KL
expansion is truncated at 20 terms in each dimension, for a total of 400
terms. We use $N_{s}=100$ samples to evaluate Functional 2. The true
parameters for data generation are $\lambda=5,\kappa^{+}=2,\kappa^{-}=1$ and
we set the level set smoothing parameter $\tau=5$. We use a finite element
mesh of $100\times 100$ nodes. The regularizer $h(\alpha)$ in Functional 2 has
means $m_{h,\lambda}=\log(10)$ and $m_{h,\kappa^{\pm}}=\log(3)$, with
standard-deviations $\sigma_{h}=2$. In Fig. 5, two permeability fields
associated with their sets of observations from the dataset are shown.

(a) $z$ field, 3

(b) Smoothed $z$ field, 3

(c) $u$ and 50 observations, 3

Figure 5: Plots of a sharp permeability fields and the smoothed field and PDE
solutions with observation locations in the dataset for the 2D Darcy problems
We not that the solutions are computed from the sharp level sets for the
dataset, whereas $(T^{\alpha})_{\\#}\mu_{0}$ in the loss function uses the
smoothed level sets to make use of gradient-based optimization.

(a) $\alpha$ convergence (b) Losses

Figure 6: Convergence of prior parameters $\alpha$ for 2D Darcy trained with
Functional 2.

Fig. 6 shows the convergence plots for prior parameter calibration. Inference
time was 1.52 minutes for the 2k iterations. The relative error of the
predicted parameters are $0.39\%$, $0.03\%$, and $1.96\%$ for $\kappa^{+}$,
$\kappa^{-}$, and $\lambda$ respectively.

####  5.1.3 Prior Calibration and Operator Learning: 1D Darcy

(a) $\alpha$ convergence (b) Losses

Figure 7: Convergence of prior parameters $\alpha$, with learned operator
$F^{\phi}$ on 1D Darcy with Functional 3 and Functional 4. We take 10
parameter updates steps on $\phi$ in the lower objective of the bilevel
optimization for every step on $\alpha$.

We now turn our attention to the task of using Functional 3 and Functional 4
to jointly estimate $(T^{\alpha})_{\\#}\mu_{0}$ and learn a parametrized
operator $F^{\phi}$. We choose this operator to be a 4-layer Fourier Neural
Operator [39] with 64 dimensional channel space and 8 Fourier modes and Silu
activation functions [29] ($\mathbf{\sigma}(x)=x\,\mathrm{sigmoid}(x)$). The
last layer performs a point-wise multiplication with a function that is zero
on the boundary ($\sin(\pi x)$). Functional 4 is evaluated with $N_{r}=20$
samples and Functional 3 is evaluated with $N_{s}=1000$ samples. We perform 10
update steps on $\phi$ per $\alpha$ update in the lower-level optimization
routine. We use the same data setup, prior, and regularizer as in Subsection
5.1.1. In Fig. 7 we show that we are able to jointly learn the neural operator
approximation of the forward model and the unknown parameters. The runtime is
2.44 minutes. The relative error for the FNO predictions against the solver
for the data samples $z\sim\mu^{\dagger}$ from the true sharp level set prior
is 0.40%. The relative error of the predicted parameters are $1.13\%$,
$0.71\%$, and $4.18\%$ for $\kappa^{+}$, $\kappa^{-}$, and $\lambda$
respectively.

####  5.1.4 Prior Calibration and Operator Learning: 2D Darcy

(a) $\alpha$ convergence with $\phi^{\star}(\alpha)$ (b) Losses

Figure 8: Convergence of prior parameters $\alpha$ for the level set prior on
the 2D Darcy problem with jointly learned operator $F^{\phi}$. For every step
of prior update $\alpha$, 10 parameter updates are performed on $\phi$.

We now apply the joint prior-operator learning methodology to the 2D setting,
keeping the same FNO parameters as in 1D case (Subsection 5.1.3) but for 2
dimensional input fields. The last layer of the FNO is now $\sin(\pi
x_{(1)})\sin(\pi x_{(2)})$. We use the same data setup, prior, and regularizer
as in Subsection 5.1.2 and $N_{r}=20$, $N_{s}=100$. Figure 8 shows the
convergence for the learning of both $(T^{\alpha})_{\\#}\mu_{0}$ and
$F^{\phi}$. The model converged accurately after a runtime of 56.46 minutes.
The relative error for the FNO predictions against the solver for the data
samples $z\sim\mu^{\dagger}$ from the true sharp level set prior is of 0.73%.
The relative error of the predicted parameters are $0.10\%$, $0.98\%$, and
$3.19\%$ for $\kappa^{+}$, $\kappa^{-}$, and $\lambda$ respectively.

###  5.2 Lognormal Priors

In this subsection we focus on learning parameters of a lognormal prior. We
introduce a Matérn-like field $a$ with regularity $\nu$, amplitude $\sigma$,
and length-scale $\ell$; to be specific we generate a centred Gaussian random
field (GRF) with covariance operator

$\displaystyle
C_{\sigma,\ell,\nu}=\gamma\ell^{d}(-\ell^{2}\Delta+\mathrm{I})^{-\nu-d/2},$

setting

$\displaystyle\gamma=\sigma^{2}\frac{2^{d}\pi^{d/2}\Gamma(\nu+d/2)}{\Gamma(\nu)},$

where $\Gamma(\cdot)$ is the Gamma function (not to be confused with our use
of $\Gamma$ as the observational noise covariance). In practice we truncate a
Karhunen-Loeve expansion in the form

$\displaystyle
a(x;\sigma,\ell,\nu)=\sum_{j=1,k=1}^{J,K}\sqrt{\gamma\ell^{d}\left(\ell^{2}(j^{2}+k^{2})\pi^{2}+1\right)^{-\nu-d/2}}\varepsilon_{j,k}\varphi_{j,k}(x),\;\;\varepsilon_{j,k}\sim{\mathcal{N}}(0,1)\mathrm{i.i.d.}$

We then set $z=\exp(a).$ Again, because of the countable nature of the
construction of the lognormal prior, it lies in a separable subspace $Z$ of
$L^{\infty}(D).$ We focus on estimating $\alpha=\\{\nu,\ell\\}$ which
represent the regularity and length-scale of the lognormal permeability field
$z$. We do not attempt to jointly infer the amplitude $\sigma$ as this induces
a lack of identifiability of parameters at the level of the Karhunen-Loève
expansion spectrum. We note that, as the amplitude parameter does not contain
much spatially dependent information, it can be accurately estimated by other
means, hence here it set to $\sigma=1$. We attempt to jointly learn
$\alpha=\\{\nu,\ell\\}$ and $\phi$ the parameters of a residual-based neural
operator approximation. In Subsection 5.2.1 we consider a setting where the
regularity and lengthscale parameters are identifiable and show successful
joint learning of $(\alpha,\phi).$ However it is intuitive that the regularity
and lengthscale may not be separately identifiable; we show in Subsection
5.2.2 that in such situations the entanglement of jointly learning
$(\alpha,\phi)$ can cause convergence problems with the proposed methodology.

(a) $a$

(b) $z$

(c) $u$ and 50 observations

Figure 9: (a) A sampled GRF $a$, (b) the exponentiated random field $z$, (c)
the solution field associated to $z$ with the 50 observation locations.

(a) $\alpha$ convergence (b) Loss (c) $\alpha$ convergence with
$\phi^{\star}(\alpha)$ (d) Losses

Figure 10: (a)-(b) Convergence and loss for $\alpha=\\{\nu,\ell\\}$ on the 2D
Darcy problem for the lognormal prior. (b)-(c) Convergence of combinations of
prior parameters $\alpha=\\{\nu,\ell\\}$ estimated jointly with $\phi$ for 2D
Darcy on the lognormal prior with 10 $\phi$ updates in the lower-level
optimization for every $\alpha$ step.

####  5.2.1 Identifiable Setting

We use $N=1000$ data $y$ each of $d_{y}=50$ noisy pointwise observations of
the solution with $\Gamma=0.01^{2}\mathrm{I}$ observational noise covariance.
The prior KL expansion is truncated at 20 terms in each dimension for a total
of 400 expansion terms. We set $N_{s}=100$. The true parameters for data
generation are $\nu=1.5,\ell=0.5$. We use a finite element mesh of $100\times
100$ nodes. The regularizer $h(\alpha)$ in Functional 3 has means
$m_{h,\nu}=\log(3.5)$ and $m_{h,\ell}=\log(1)$, with standard-deviations
$\sigma_{h}=2$. In Fig. 10(a, b) we show the convergence plots for the prior
only learning. The relative error on parameter estimation is $1.26\%$ and
$0.66\%$ on $\nu$ and $\ell$ respectively. The runtime is of 4.26 mins.

We then test the joint estimation of $\alpha$ and $\phi$ with Functional 3,
Functional 4. The data setup, prior, and regularizers are the same as in the
$\alpha$-only case described in Subsection 5.1.2 with $N_{r}=20$, $N_{s}=100$.
The operator setup is the same as in Subsection 5.1.4. In Fig. 10(c, d) we
show the joint prior learning with operator learning. The achieved relative
error for the learned $F^{\phi}$ on $z\sim\mu^{\dagger}$ is $0.128\%$. The
relative error on parameter estimation is $0.44\%$ and $3.13\%$ on $\nu$ and
$\ell$ respectively. The runtime is of $112$ mins for the 10k iterations.

####  5.2.2 Unidentifiable Setting

(a) $\alpha$ convergence (b) Loss

(c) $\alpha$ convergence with $\phi^{\star}(\alpha)$ (d) Losses

Figure 11: (a)-(b) Convergence and loss for $\alpha$ learning. (c)-(d)
Convergence and losses for $\alpha$ and $\phi$ learning. Both sets of figures
are for the 2D Darcy on log normal prior for unidentifiable $\alpha$ parameter
values.

For certain parameter values of $\alpha$, from which the true data is
generated, there can be a lack of identifiability. In the previous section,
the data was generated from a relatively rough random field with a relatively
long length-scale. Hence the smoothness and length-scale parameters each have
a distinct impact on the prior’s spectrum. We will now look at a smooth random
field setting $\nu=4$, keeping $\ell=0.5$. Now, both parameters will have a
similar influence on the spectrum of the prior. Other parameters for the data,
solver/operator, and regularizer are kept the same as in Subsection 5.2.1. In
Fig. 11 we show the convergence of the $\alpha$ parameters for the
unidentifiable parameter regime for the $\alpha$ only learning and the joint
$\alpha,\phi$ prior and operator learning tasks. We can see the lack of
identifiability causes more issues in the operator learning as it struggles to
gain any inferential traction. For these examples, it was necessary to run the
learning models for more iterations (20k) and use a 6-time decaying learning
rate on $\alpha$ (as opposed to the usual 4 in previous examples).
Furthermore, due to training instability, the $\phi$ parameters had to be
trained with AMSGrad [52](with the same settings) instead of Adam. As we can
see from the Fig. 11, the proposed methodology applied to this particular
unidentifiable setup does not fail silently, the $\alpha$ only learning has a
very long convergence time, and the $\alpha$ and $\phi$ learning does not
reach equilibrium. Other possible avenues to identifying non-uniqueness in
prior parameter estimation include assessing the sensitivity of the estimated
parameters on the choice of regularizer mean and standard-deviation, or the
use of hierarchical parametric priors where a distribution on $\alpha$ is
learned. There, a wide (or multimodal) inferred distribution on $\alpha$ could
indicate non-identifiability of the prior parameters as well as the covariance
between entangled prior parameters.

##  6 Conclusions

In this work we propose a novel methodology for learning a generative model
for a prior on function space, based on indirect and noisy observations
defined through solution of a PDE. The learning scheme is based on the
minimization of a loss functional computing divergences at a distributional
level. We prove our methodology recovers the Bayesian posterior when
observations originate from a single physical system ($N=1$). We demonstrate
the accuracy of our methodology on a series of example pertinent for
practitioners. These are 1D and 2D steady-state Darcy flown problems for two
different parametric priors. Namely, a level set prior and a lognormal prior.
Furthermore, we show how the proposed framework can be augmented to jointly
learn an operator, mapping samples of PDE parameter fields drawn from the
estimated prior, and the solution space of the PDE, through a bilevel
optimization strategy. The operator learning takes place on-the-fly during the
optimization of the prior parameters, and is residual-based. Finally, we
indicate the possible pitfalls of unidentifiable priors and show the emergent
behavior of the methodology under this setting. A selection of avenues for
future work include, the use of different metrics on the space of measures,
testing on different PDE systems (nonlinear, coupled, time-evolving etc.),
other choices of parametrized prior measures, studying identifiability and
non-uniqueness of inferred prior parameters, and testing the downstream use of
such learned priors and operators as in Bayesian inversion and generative
modeling of physical systems.

## Funding

MG is grateful for support through EPSRC grants EP/X037770/1, EP/Y028805/1,
EP/W005816/1, EP/V056522/1 , EP/V056441/1, EP/T000414/1, and EP/R034710/1. AMS
is grateful for support through a Department of Defense Vannevar Bush Faculty
Fellowship, from the Air Force Office of Scientific Research under MURI award
number FA9550-20-1-0358 (Machine Learning and Physics-Based Modeling and
Simulation) and by the SciAI Center, funded by the Office of Naval Research
(ONR), under Grant Number N00014-23-1-2729. AV is supported through the EPSRC
ROSEHIPS grant EP/W005816/1.

##  Appendix A Proofs

###  A.1 Proof of Lemma 3.1

###### Proof.

We first remember our definition

(23)
$\displaystyle\mathsf{W}^{2}_{2,\mathrm{B}}(\nu,\mu)=\inf_{\gamma\in\Pi(\nu,\mu)}\int_{\mathbb{R}^{d}\times\mathbb{R}^{d}}\|x-y\|_{\mathrm{B}}^{2}{\mathrm{d}}\gamma(x,y).$

By the definition of pushforward

(24)
$\displaystyle\mathsf{W}^{2}_{2}(P_{\mathrm{B}\\#}\nu,P_{\mathrm{B}\\#}\mu)$
$\displaystyle=\inf_{\gamma\in\Pi(\nu,\mu)}\int_{\mathbb{R}^{d}\times\mathbb{R}^{d}}\|\mathrm{B}^{-1/2}x-\mathrm{B}^{-1/2}y\|_{2}^{2}{\mathrm{d}}\gamma(x,y)$
(25)
$\displaystyle=\inf_{\gamma\in\Pi(\nu,\mu)}\int_{\mathbb{R}^{d}\times\mathbb{R}^{d}}\|x-y\|_{\mathrm{B}}^{2}{\mathrm{d}}\gamma(x,y).$

∎

###  A.2 Proof of Lemma 3.4

###### Proof.

To show the desired result we determine a point at which the infimum over
couplings $\Pi(\delta_{y},\mu)$ is achieved. It is known from [58][Section
§1.4] that when one of the measures in the argument of a Kantorovich problem
(KP) (of which Wasserstein metrics are a special case) is a Dirac, i.e.
$\delta_{y}$ for any $y\in{\mathbb{R}}^{d}$, then the set of couplings
$\Pi(\delta_{y},\mu)$ with marginals $\delta_{y}$, $\mu$, contains a single
element, namely $\delta_{y}\otimes\mu$. It follows that

$\displaystyle\mathsf{W}^{2}_{2,\Gamma}(\delta_{y},\mu)$
$\displaystyle=\inf_{\gamma\in\Pi(\delta_{y},\mu)}\int_{\mathbb{R}^{d}\times\mathbb{R}^{d}}\|y^{\prime}-x\|_{\Gamma}^{2}\mathrm{d}\gamma(y^{\prime},x)$
$\displaystyle=\int_{\mathbb{R}^{d}\times\mathbb{R}^{d}}\|y^{\prime}-x\|_{\Gamma}^{2}({\mathrm{d}}\delta_{y}(y^{\prime})\otimes{\mathrm{d}}\mu(x))$
(26) $\displaystyle=\mathbb{E}_{x\sim\mu}\|y-x\|_{\Gamma}^{2}.$

∎

###  A.3 Proof of Lemma 3.5

###### Proof.

We want to show, for $y\in\mathbb{R}^{d}$ and
$\mu\in\mathcal{P}(\mathbb{R}^{d})$,

$\displaystyle\mathsf{SW}_{2,\Gamma}^{2}(\delta_{y},\mu)=\frac{1}{d}\mathsf{W}_{2,\Gamma}^{2}(\delta_{y},\mu).$

Thus, we have

$\displaystyle\mathsf{SW}^{2}_{2,\Gamma}(\delta_{y},\mu)$
$\displaystyle=\int_{\mathbb{S}^{d-1}}\mathsf{W}_{2}^{2}(P^{\theta}_{\Gamma\\#}\delta_{y},P^{\theta}_{\Gamma\\#}\mu){\mathrm{d}}\theta$
$\displaystyle=\int_{\mathbb{S}^{d-1}}\mathsf{W}_{2}^{2}(\delta_{\langle\Gamma^{-1/2}y,\theta\rangle},P^{\theta}_{\Gamma\\#}\mu){\mathrm{d}}\theta$
(27a)
$\displaystyle=\int_{\mathbb{S}^{d-1}}\mathbb{E}_{z\sim\mu}(\langle\Gamma^{-1/2}y,\theta\rangle-\langle\Gamma^{-1/2}z,\theta\rangle)^{2}{\mathrm{d}}\theta$
(27b)
$\displaystyle=\mathbb{E}_{z\sim\mu}\int_{\mathbb{S}^{d-1}}(\langle\Gamma^{-1/2}y,\theta\rangle-\langle\Gamma^{-1/2}z,\theta\rangle)^{2}{\mathrm{d}}\theta$
$\displaystyle=\mathbb{E}_{z\sim\mu}\int_{\mathbb{S}^{d-1}}\bigl{(}\langle\Gamma^{-1/2}(y-z),\theta\rangle\bigr{)}^{2}{\mathrm{d}}\theta,$

where, in order to move from (27a) to (27b), we use the fact that
${\mathbb{E}}_{z\sim\mu}\|z\|_{\Gamma}^{2}<\infty$. Furthermore,

$\displaystyle\langle\Gamma^{-1/2}a,\theta\rangle\langle\Gamma^{-1/2}b,\theta\rangle=(\Gamma^{-1/2}a)^{\top}\theta\theta^{\top}(\Gamma^{-1/2}b),$

and noting, from [70][Section §3.3.1], the covariance of the uniform on the
sphere, $\mathrm{Cov\,Unif}(\mathbb{S}^{d-1})$
$=\int_{\mathbb{S}^{d-1}}\theta\theta^{\top}{\mathrm{d}}\theta=\frac{1}{d}\mathrm{I}$.
Then

$\displaystyle\int_{\mathbb{S}^{d-1}}\langle\Gamma^{-1/2}a,\theta\rangle\langle\Gamma^{-1/2}b,\theta\rangle{\mathrm{d}}\theta$
$\displaystyle=\frac{1}{d}(\Gamma^{-1/2}a)^{\top}(\Gamma^{-1/2}b)$
$\displaystyle=\frac{1}{d}\langle a,b\rangle_{\Gamma}.$

In particular

$\displaystyle\int_{\mathbb{S}^{d-1}}\bigl{(}\langle\Gamma^{-1/2}(y-z),\theta\rangle\bigr{)}^{2}{\mathrm{d}}\theta$
$\displaystyle=\frac{1}{d}\|y-z\|^{2}_{\Gamma}.$

And so, from (27b) and Lemma 3.4 we obtain

$\displaystyle\mathsf{SW}^{2}_{2,\Gamma}(\delta_{y},\mu)$
$\displaystyle=\frac{1}{d}\mathbb{E}_{z\sim\mu}\|y-z\|^{2}_{\Gamma}$
$\displaystyle=\frac{1}{d}\mathsf{W}^{2}_{2,\Gamma}(\delta_{y},\mu).$

∎

###  A.4 Proof of Lemma 3.6

###### Proof.

We have

$\displaystyle\mathsf{I}(\eta*\mu)$
$\displaystyle={\mathbb{E}}_{x^{\prime}\sim\eta*\mu}\|y-x^{\prime}\|_{\Gamma}^{2},$
$\displaystyle\mathsf{I}(\mu)$
$\displaystyle={\mathbb{E}}_{x\sim\mu}\|y-x\|_{\Gamma}^{2}.$

Thus

$\displaystyle\mathsf{I}(\eta*\mu)$
$\displaystyle={\mathbb{E}}_{(x,\varepsilon)\sim\mu\otimes\eta}\|y-x-\varepsilon\|_{\Gamma}^{2},$
$\displaystyle={\mathbb{E}}_{(x,\varepsilon)\sim\mu\otimes\eta}\Bigl{(}\|y-x\|_{\Gamma}^{2}-2\langle
y-x,\varepsilon\rangle_{\Gamma}+\|\varepsilon\|_{\Gamma}^{2}\Bigr{)},$
$\displaystyle={\mathbb{E}}_{x\sim\mu}\|y-x\|_{\Gamma}^{2}+{\mathbb{E}}_{\varepsilon\sim\eta}\|\varepsilon\|_{\Gamma}^{2},$
$\displaystyle=\mathsf{I}(\mu)+{\mathbb{E}}_{\varepsilon\sim\eta}\|\varepsilon\|_{\Gamma}^{2},$

using that $\eta$ is centred, independent of $\mu$ and has finite second
moments.

∎

##  Appendix B Weak Form Residuals with Array Shifting

We now show how to compute a residual in the weak form for a Darcy problem on
domain $D$, with homogenous boundary conditions using array shifting. Taking
(3) and testing the domain residual against as set of test functions
$\\{v_{i}\\}_{i=1}^{d_{o}}$, we obtain the discretized variational formulation

$\displaystyle\mathcal{O}(R(z,u))_{i}=\langle
v_{i},R(z,u)\rangle=\int_{D}v_{i}(\nabla\cdot(z\nabla
u)){\mathrm{d}}x+\int_{D}v_{i}f{\mathrm{d}}x=0,\forall v_{i}\in V.$

Through integration by parts we obtain the weak form. Here, we use as a
shorthand $\mathsf{r}_{i}=\mathcal{O}(R(z,u))_{i}$. In 2D

(28) $\displaystyle\mathsf{r}_{i}$
$\displaystyle=-\int_{D}z\partial_{x_{(1)}}v_{i}\,\partial_{x_{(1)}}u\,{\mathrm{d}}x-\int_{D}z\partial_{x_{(2)}}v_{i}\,\partial_{x_{(2)}}u\,{\mathrm{d}}x+\int_{D}v_{i}f\,{\mathrm{d}}x.$

$j,k$$j,k_{-1}$$j,k_{+1}$$j_{+1},k$$j_{+1},k_{+1}$$j_{-1},k_{-1}$$j_{-1},k$$1$$2$$3$$4$$5$$6$
Figure 12: 2D FEM mesh centred at node $jk$. We assume equal spacing, $h$, of
the nodes.

Traditional finite element solvers would assemble a system of sparse linear
equations of the form $\mathsf{Au=f}$ where
$\mathsf{A}\in\mathbb{R}^{d_{o}\times d_{o}}$ is sparse, and
$\mathsf{u,f}\in\mathbb{R}^{d_{o}}$. The matrix vector product $\mathsf{Au}$
is comprised of the first two terms in (28) and $\mathsf{f}$ is the third
term. To compute residuals using array shifting in 2D we use double indexing,
denoted by $jk$, to represent the mesh nodes as shown in Fig. 12. We denote
the set of indices
$\\{i\\}_{i=1}^{d_{o}}=\mathsf{ravel}(\\{j,k\\}_{j,k=1}^{j,k=\sqrt{d_{o}}})$.
Then

$\displaystyle\mathsf{r}_{jk}=-\mathsf{r}_{x_{(1)},jk}-\mathsf{r}_{x_{(2)},jk}+\mathsf{r}_{f,jk},$

where each of the three terms corresponds the ones in (28). Assuming $z$ is
given at the nodes and is piece-wise constant from the top left of an element,
we have

$\displaystyle\mathsf{r}_{x_{(1)},jk}=$
$\displaystyle\frac{1}{2}((z_{j_{-1},k_{-1}}+z_{j,k_{-1}})(u_{j,k}-u_{j,k_{-1}})-(z_{j,k}+z_{j_{-1},k})(u_{j,k_{+1}}-u_{j,k})),$
$\displaystyle\mathsf{r}_{x_{(2)},jk}=$
$\displaystyle\frac{1}{2}((z_{j,k_{-1}}+z_{j,k})(u_{j,k}-u_{j_{+1},k})-(z_{j_{-1},k}+z_{j_{-1},k_{-1}})(u_{j_{-1},k}-u_{j,k})).$

We compute the tested inhomogeneous term as
$\int_{D}vf\,{\mathrm{d}}x=\sum_{e=1}^{6}\int_{D_{e}}vf\,{\mathrm{d}}x$
integrating using 1 point Gauss integration at $(1/3h,1/3h)$. Thus

$\displaystyle\mathsf{r}_{f,jk}=\frac{h^{2}}{9}\left(3f_{j,k}+f_{j_{-1},k_{-1}}+f_{j,k_{-1}}+f_{j_{+1},k}+f_{j_{+1},k_{+1}}+f_{j,k_{+1}}+f_{j_{-1},k}\right).$

All operations can be rapidly computed using array shifting. To solve the
linear differential equation we use an iterative linear solver such as
conjugate gradient (CG), for positive definite matrices $\mathsf{A}$, or GMRES
for more general problems. To solve nonlinear differential equations, one
would use Newton’s method. The Newton update step can be seen as solving a
linear system with a Jacobian-vector product
$\mathsf{u_{n+1}}=\mathsf{{u}_{n}-{\mathrm{J}_{\mathcal{O}}(u_{n})}^{-1}\mathcal{O}(u_{n})}$.
Deep learning libraries are designed to compute Jacobian-vector products with
high efficiency, resulting in a GPU efficient Newton-Krilov method.

## References

  * [1] J. Adler, A. Ringh, O. Öktem, and J. Karlsson, Learning to solve inverse problems using wasserstein loss, preprint arXiv:1710.10898, (2017).
  * [2] S. Arridge, P. Maass, O. Öktem, and C.-B. Schönlieb, Solving inverse problems using data-driven models, Acta Numer., 28 (2019), pp. 1–174.
  * [3] M. Asim, M. Daniels, O. Leong, A. Ahmed, and P. Hand, Invertible generative models for inverse problems: mitigating representation error and dataset bias, ICML, (2020), pp. 399–409.
  * [4] M. A. Beaumont, W. Zhang, and D. J. Balding, Approximate bayesian computation in population genetics, Genetics, 162 (2002), pp. 2025–2035.
  * [5] J. M. Bernardo, Reference posterior distributions for bayesian inference, J. R. Stat. Soc. Ser. B. Stat. Methodol., 41 (1979), pp. 113–128.
  * [6] J. M. Bernardo and A. F. Smith, Bayesian theory, vol. 405, John Wiley & Sons, 2009.
  * [7] E. Bernton, P. E. Jacob, M. Gerber, and C. P. Robert, On parameter estimation with the wasserstein distance, Inf. Inference, 8 (2019), pp. 657–676.
  * [8] K. Bhattacharya, B. Hosseini, N. B. Kovachki, and A. M. Stuart, Model reduction and neural networks for parametric pdes, SMAI J. Comput. Math., 7 (2021), pp. 121–157.
  * [9] N. Bonneel, J. Rabin, G. Peyré, and H. Pfister, Sliced and radon wasserstein barycenters of measures, Journal of Mathematical Imaging and Vision, 51 (2015), pp. 22–45.
  * [10] N. Bonnotte, Unidimensional and evolution methods for optimal transportation, PhD thesis, Université Paris Sud-Paris XI; Scuola normale superiore (Pise, Italie), 2013.
  * [11] B. Boys, M. Girolami, J. Pidstrigach, S. Reich, A. Mosca, and O. D. Akyildiz, Tweedie moment projected diffusions for inverse problems, preprint arXiv:2310.06721, (2023).
  * [12] J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang, JAX: composable transformations of Python+NumPy programs, (2018).
  * [13] T. Butler, J. Jakeman, and T. Wildey, A consistent bayesian formulation for stochastic inverse problems based on push-forward measures, preprint arXiv:1704.00680, (2017).
  * [14] T. Butler, J. Jakeman, and T. Wildey, Combining push-forward measures and bayes’ rule to construct consistent solutions to stochastic inverse problems, SIAM J. Sci. Comput., 40 (2018), pp. A984–A1011.
  * [15] G. Casella, An introduction to empirical bayes data analysis, Amer. Statist., 39 (1985), pp. 83–87.
  * [16] J. Charrier, Strong and weak error estimates for elliptic partial differential equations with random coefficients, SIAM J. Numer. Anal., 50 (2012), pp. 216–246.
  * [17] J. Chemseddine, P. Hagemann, C. Wald, and G. Steidl, Conditional wasserstein distances with applications in bayesian ot flow matching, preprint arXiv:2403.18705, (2024).
  * [18] F. R. Crucinio, V. De Bortoli, A. Doucet, and A. M. Johansen, Solving a class of fredholm integral equations of the first kind via wasserstein gradient flows, Stochastic Process. Appl., (2024), p. 104374.
  * [19] M. Cuturi, Sinkhorn distances: Lightspeed computation of optimal transport, Advances in neural information processing systems, 26 (2013).
  * [20] M. Dashti and A. M. Stuart, Uncertainty quantification and weak approximation of an elliptic inverse problem, SIAM J. Numer. Anal., 49 (2011), pp. 2524–2542.
  * [21] I. Deshpande, Y.-T. Hu, R. Sun, A. Pyrros, N. Siddiqui, S. Koyejo, Z. Zhao, D. Forsyth, and A. G. Schwing, Max-sliced wasserstein distance and its use for gans, Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, (2019), pp. 10648–10656.
  * [22] M. M. Dunlop, M. A. Iglesias, and A. M. Stuart, Hierarchical bayesian level set inversion, Stat. Comput., 27 (2017), pp. 1555–1584.
  * [23] G. K. Dziugaite, D. M. Roy, and Z. Ghahramani, Training generative neural networks via maximum mean discrepancy optimization, preprint arXiv:1505.03906, (2015).
  * [24] T. A. El Moselhy and Y. M. Marzouk, Bayesian inference with optimal maps, J. Comput. Phys., 231 (2012), pp. 7815–7850.
  * [25] B. T. Feng, J. Smith, M. Rubinstein, H. Chang, K. L. Bouman, and W. T. Freeman, Score-based diffusion models as principled priors for inverse imaging, Proceedings of the IEEE/CVF International Conference on Computer Vision, (2023), pp. 10520–10531.
  * [26] R. Flamary, N. Courty, A. Gramfort, M. Z. Alaya, A. Boisbunon, S. Chambon, L. Chapel, A. Corenflos, K. Fatras, N. Fournier, L. Gautheron, N. T. Gayraud, H. Janati, A. Rakotomamonjy, I. Redko, A. Rolet, A. Schutz, V. Seguy, D. J. Sutherland, R. Tavenard, A. Tong, and T. Vayer, Pot: Python optimal transport, JMLR, 22 (2021), pp. 1–8.
  * [27] A. F. Gao, O. Leong, H. Sun, and K. L. Bouman, Image reconstruction without explicit priors, ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), (2023), pp. 1–5.
  * [28] A. Gelman, J. B. Carlin, H. S. Stern, and D. B. Rubin, Bayesian data analysis, Chapman and Hall/CRC, 1995.
  * [29] D. Hendrycks and K. Gimpel, Gaussian error linear units (gelus), preprint arXiv:1606.08415, (2016).
  * [30] J. Ho, A. Jain, and P. Abbeel, Denoising diffusion probabilistic models, Advances in neural information processing systems, 33 (2020), pp. 6840–6851.
  * [31] M. A. Iglesias, K. Lin, and A. M. Stuart, Well-posed bayesian geometric inverse problems arising in subsurface flow, Inverse Problems, 30 (2014), p. 114001.
  * [32] E. T. Jaynes, Prior probabilities, IEEE Transactions on systems science and cybernetics, 4 (1968), pp. 227–241.
  * [33] H. Jeffreys, An invariant form for the prior probability in estimation problems, Proceedings of the Royal Society of London. Series A. Mathematical and Physical Sciences, 186 (1946), pp. 453–461.
  * [34] M. C. Kennedy and A. O’Hagan, Bayesian calibration of computer models, J. Roy. Statist. Soc. Ser. A, 63 (2001), pp. 425–464.
  * [35] D. P. Kingma and J. Ba, Adam: A method for stochastic optimization, preprint arXiv:1412.6980, (2014).
  * [36] D. P. Kingma and M. Welling, Auto-encoding variational bayes, preprint arXiv:1312.6114, (2013).
  * [37] S. Kolouri, K. Nadjahi, U. Simsekli, R. Badeau, and G. Rohde, Generalized sliced wasserstein distances, Advances in neural information processing systems, 32 (2019).
  * [38] N. Kovachki, Z. Li, B. Liu, K. Azizzadenesheli, K. Bhattacharya, A. Stuart, and A. Anandkumar, Neural operator: Learning maps between function spaces with applications to pdes, JMLR, 24 (2023), pp. 1–97.
  * [39] Z. Li, N. Kovachki, K. Azizzadenesheli, B. Liu, K. Bhattacharya, A. Stuart, and A. Anandkumar, Fourier neural operator for parametric partial differential equations, preprint arXiv:2010.08895, (2020).
  * [40] Z. Li, H. Zheng, N. Kovachki, D. Jin, H. Chen, B. Liu, K. Azizzadenesheli, and A. Anandkumar, Physics-informed neural operator for learning partial differential equations, ACM/JMS Journal of Data Science, (2021).
  * [41] A. Liutkus, U. Simsekli, S. Majewski, A. Durmus, and F.-R. Stöter, Sliced-wasserstein flows: Nonparametric generative modeling via optimal transport and diffusions, ICML, (2019), pp. 4104–4113.
  * [42] L. Lu, P. Jin, G. Pang, Z. Zhang, and G. E. Karniadakis, Learning nonlinear operators via deeponet based on the universal approximation theorem of operators, Nature machine intelligence, 3 (2021), pp. 218–229.
  * [43] Y. Marzouk, T. Moselhy, M. Parno, and A. Spantini, An introduction to sampling via measure transport, arXiv preprint arXiv:1602.05023, (2016).
  * [44] C. N. Morris, Parametric empirical bayes inference: theory and applications, J. Amer. Statist. Assoc., 78 (1983), pp. 47–55.
  * [45] K. Nadjahi, Sliced-wasserstein distance for large-scale machine learning: theory, methodology and extensions, Institut polytechnique de Paris, (2021).
  * [46] K. Nadjahi, V. De Bortoli, A. Durmus, R. Badeau, and U. Şimşekli, Approximate bayesian computation with the sliced-wasserstein distance, ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), (2020), pp. 5470–5474.
  * [47] K. Nguyen and N. Ho, Energy-based sliced wasserstein distance, Advances in Neural Information Processing Systems, 36 (2024).
  * [48] A. O’Hagan, Expert knowledge elicitation: subjective but scientific, Amer. Statist., 73 (2019), pp. 69–81.
  * [49] G. Papamakarios, E. Nalisnick, D. J. Rezende, S. Mohamed, and B. Lakshminarayanan, Normalizing flows for probabilistic modeling and inference, JMLR, 22 (2021), pp. 1–64.
  * [50] D. V. Patel and A. A. Oberai, GAN-Based Priors for Quantifying Uncertainty in Supervised Learning, SIAM/ASA J. Uncertain. Quantif., 9 (2021), pp. 1314–1343.
  * [51] H. Raiffa and R. Schlaifer, Applied statistical decision theory, vol. 78, John Wiley & Sons, 2000.
  * [52] S. J. Reddi, S. Kale, and S. Kumar, On the convergence of adam and beyond, preprint arXiv:1904.09237, (2019).
  * [53] M. Rixner and P.-S. Koutsourelakis, A probabilistic generative model for semi-supervised training of coarse-grained surrogates and enforcing physical constraints through virtual observables, J. Comput. Phys., 434 (2021), p. 110218.
  * [54] H. Robbins, The empirical bayes approach to statistical decision problems, The Annals of Mathematical Statistics, 35 (1964), pp. 1–20.
  * [55] H. E. Robbins, An empirical bayes approach to statistics, in Breakthroughs in Statistics: Foundations and basic theory, Springer, 1992, pp. 388–394.
  * [56] J. Sacks, W. J. Welch, T. J. Mitchell, and H. P. Wynn, Design and analysis of computer experiments, Statistical Science, 4 (1989), pp. 409–423.
  * [57] A. Sagiv, The wasserstein distances between pushed-forward measures with applications to uncertainty quantification, Commun. Math. Sci., 18 (2020), pp. 707–724.
  * [58] F. Santambrogio, Optimal transport for applied mathematicians, Birkäuser, NY, 55 (2015), p. 94.
  * [59] D. Sejdinovic, B. Sriperumbudur, A. Gretton, and K. Fukumizu, Equivalence of distance-based and rkhs-based statistics in hypothesis testing, The annals of statistics, (2013), pp. 2263–2291.
  * [60] A. Sinha, P. Malo, and K. Deb, A review on bilevel optimization: From classical to evolutionary approaches and applications, IEEE transactions on evolutionary computation, 22 (2017), pp. 276–295.
  * [61] A. Smola, A. Gretton, L. Song, and B. Schölkopf, A hilbert space embedding for distributions, International conference on algorithmic learning theory, (2007), pp. 13–31.
  * [62] Y. Song and S. Ermon, Generative Modeling by Estimating Gradients of the Data Distribution, Advances in Neural Information Processing Systems, 32 (2019).
  * [63] Y. Song and D. P. Kingma, How to train your energy-based models, preprint arXiv:2101.03288, (2021).
  * [64] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole, Score-based generative modeling through stochastic differential equations, preprint arXiv:2011.13456, (2020).
  * [65] G. J. Székely and M. L. Rizzo, Energy statistics: A class of statistics based on distances, Journal of statistical planning and inference, 143 (2013), pp. 1249–1272.
  * [66] S. Tavaré, D. J. Balding, R. C. Griffiths, and P. Donnelly, Inferring coalescence times from dna sequence data, Genetics, 145 (1997), pp. 505–518.
  * [67] Y. W. Teh, M. Welling, S. Osindero, and G. E. Hinton, Energy-based models for sparse overcomplete representations, JMLR, 4 (2003), pp. 1235–1260.
  * [68] A. Vadeboncoeur, Ö. D. Akyildiz, I. Kazlauskaite, M. Girolami, and F. Cirak, Fully probabilistic deep models for forward and inverse problems in parametric pdes, J. Comput. Phys., 491 (2023).
  * [69] A. Vadeboncoeur, I. Kazlauskaite, Y. Papandreou, F. Cirak, M. Girolami, and O. D. Akyildiz, Random grid neural processes for parametric partial differential equations, ICML, (2023).
  * [70] R. Vershynin, High-dimensional probability: An introduction with applications in data science, vol. 47, Cambridge university press, 2018.
  * [71] M. J. Wainwright, M. I. Jordan, et al., Graphical models, exponential families, and variational inference, Foundations and Trends in Machine Learning, 1 (2008), pp. 1–305.
  * [72] R. Z. Zhang, X. Xie, and J. Lowengrub, Bilo: Bilevel local operator learning for pde inverse problems, preprint arXiv:2404.17789, (2024).
  * [73] Y. Zhu, N. Zabaras, P.-S. Koutsourelakis, and P. Perdikaris, Physics-constrained deep learning for high-dimensional surrogate modeling and uncertainty quantification without labeled data, J. Comput. Phys., 394 (2019), pp. 56–81.