# Normalising Flows for Bayesian Gravity Inversion

Henrietta Rakoczi
School of Physics and Astronomy
University of Glasgow
Glasgow, UK
<EMAIL_ADDRESS>
&Abhinav Prasad
School of Physics and Astronomy
University of Glasgow
Glasgow, UK
&Karl Toland
School of Physics and Astronomy
University of Glasgow
Glasgow, UK
&Christopher Messenger
School of Physics and Astronomy
University of Glasgow
Glasgow, UK
&Giles Hammond
School of Physics and Astronomy
University of Glasgow
Glasgow, UK

###### Abstract

Gravity inversion is a commonly applied data analysis technique in the field
of geophysics. While machine learning methods have previously been explored
for the problem of gravity inversion, these are deterministic approaches
returning a single solution deemed most appropriate by the algorithm. The
method presented here takes a different approach, where gravity inversion is
reformulated as a Bayesian parameter inference problem. Samples from the
posterior probability distribution of source model parameters are obtained via
the implementation of a generative neural network architecture known as
Normalising Flows. Due to its probabilistic nature, this framework provides
the user with a range of source parameters and uncertainties instead of a
single solution, and is inherently robust against instrumental noise. The
performance of the Normalising Flow is compared to that of an established
Bayesian method called Nested Sampling. It is shown that the new method
returns results with comparable accuracy 200 times faster than standard
sampling methods, which makes Normalising Flows a suitable method for real-
time inversion in the field. When applied to data sets with high
dimensionality, standard sampling methods can become impractical due to long
computation times. It is shown that inversion using Normalising Flows remains
tractable even at 512 dimensions and once the network is trained, the results
can be obtained in $O(10)$ seconds.

_Keywords_ gravity inversion, Bayesian, machine learning, normalising flows

## Acronyms

CDF

    cumulative distribution function
JS

    Jensen-Shannon
KL

    Kullback-Leibler
ML

    machine learning
NF

    Normalising Flow
NS

    Nested Sampling
PDF

    probability density function

##  1 Introduction

Gravimetry is a method that measures differences in the gravitational
acceleration on the Earth’s surface to a very high precision. Measuring
gravity can be a useful probe when the detection or characterisation of
underground features is desired as spatial or temporal changes in gravity can
indicate fluctuations in the density contrast between regions or objects of
interest (Torge (1989)). It can act as a complementary tool to other methods,
such as magnetic field or seismic measurements (Milano et al. (2021)), and can
offer an alternative when other methods are not viable. Gravimetry can be used
for civil engineering purposes, such as the detection of underground man-made
structures (Pearse et al. (2021), Eppelbaum (2010)), geophysical applications
such as volcano monitoring (Carbone et al. (2020)), or oil and gas exploration
(Gadirov et al. (2022)). In addition, gravimetry can not only be used for
detection and monitoring, but also for the deduction of the parameters of the
underground density distribution.

Modelling underground structures from geophysical measurements is an inverse
problem, where the computation of a mapping from the measured data to the
model governing the data is desired. In the case of gravimetry, this can be
formulated as an inverse source problem, such as

$\boldsymbol{\theta}=A^{-1}(\boldsymbol{d})$ (1)

where $\boldsymbol{\theta}$ is the source model or underground density
distribution, $\boldsymbol{d}$ is the measured gravity data, and $A^{-1}$ is
the inverse mapping (Zhdanov (2002)). In gravimetry it is often necessary that
the model is defined in higher dimensions than the data itself, which results
in an ill-posed problem. When measuring the gravitational acceleration of a
3-dimensional model on a flat surface, some information about the source is
inherently lost and the solution becomes non-unique. For instance, a massive
source buried deep underground might result in a similar gravity profile as a
much smaller, shallow object. These degeneracies cannot easily be broken, but
constraints can be introduced based on prior knowledge about the
characteristics of the source and local geology. Apart from this problem of
uniqueness, another concern in gravity inversion is the stability of the
solution in the presence of measurement noise. The solution is deemed stable
if it does not vary significantly within the noise range of the survey.

Traditional linear inversion algorithms apply the least squares method to
iteratively optimise the source model parameters until a solution is found
(Sen and Stoffa (2013), Lines and Treitel (1984)). The stability and
consistency of results from these methods are often susceptible to variations
in the measurement data and require the introduction of regularisation
functions (Tikhonov and Arsenin (1977)), which encode prior information and
need to be carefully chosen for the specific problem (Zhdanov (2009)). While
standard methods default to a single solution, we know that a family of
solutions exist due to non-uniqueness and noise. Therefore, additional
considerations are required to approximate uncertainties on the result.

Alternatively, gravity inversion can be reformulated as a Bayesian parameter
inference problem. Within this framework, the parameters of the source model
($\boldsymbol{\theta}$) are treated as random variables and a probability
density function (PDF) of these parameters is inferred to describe the
solution to the inversion. Due to its probabilistic nature, the method can
naturally account for instrumental noise and maintain stability. Any previous
knowledge of the source and characteristics of the site can also be encoded in
the prior, which is the assumed PDF of the parameters of interest before any
knowledge is obtained from the measurements. In addition, using this method we
obtain an ensemble of solutions in the form of samples from the Bayesian
posterior, instead of a single instance. This allows the user to have a better
understanding of the uncertainties associated with the problem, obtain
statistical measures, and visualise a whole range of possible models that
could explain the measured data. Bayes’ Theorem (Bayes (1763)) describes the
posterior PDF, $p(\boldsymbol{\theta}|\boldsymbol{d},A)$ as

$p(\boldsymbol{\theta}|\boldsymbol{d},A)=\frac{p(\boldsymbol{\theta}|A)p(\boldsymbol{d}|\boldsymbol{\theta},A)}{p(\boldsymbol{d}|A)}$
(2)

where $\boldsymbol{d}$ is the observed data, $A$ is the assumed forward model,
$\pi(\boldsymbol{\theta})\equiv p(\boldsymbol{\theta}|A)$ is the prior and
$L(\boldsymbol{\theta})\equiv p(\boldsymbol{d}|\boldsymbol{\theta},A)$ is the
likelihood. The denominator is a normalising constant, called the Bayesian
evidence. To obtain the evidence, the computation of the integral

$Z=p(\boldsymbol{d}|A)=\int
L(\boldsymbol{\theta})\pi(\boldsymbol{\theta})d\boldsymbol{\theta}$ (3)

over the whole parameter space is required, which can become intractable for
high-dimensional spaces. Therefore, various methods have been developed to
approximate this term or avoid computing it altogether.

Stochastic sampling methods (Sivia and Skilling (2006)) are commonly applied
to approximate the posterior. Some of these methods have been successfully
applied to gravity inversion (Brown (2015)), however the computation time
increases with the complexity and dimensionality of the posterior (Rodgers
(2017)). The application of machine learning (ML) can potentially improve on
sampling methods by reducing computation time and allowing the approximation
of posterior probability distributions even for high-dimensional data. In this
work, a specific type of generative neural network model implementing
Normalising Flow s (NF) (Rezende and Mohamed (2016)) is investigated for
applications in gravity inversion. Valentine and Sambridge (2022) identified
variational inference using generative models as a field of rapid development
and interest within the geophysics community. While various ML methods have
been applied to deterministic, single-solution gravity inversion (Yang et al.
(2022), Zhang et al. (2022)), the application of neural networks for
probabilistic gravity inversion is still in its early stages. More
specifically, NF s were recently explored for the inverse problem in seismic
tomography (Zhao et al. (2021)), but until now these not have been tested for
gravity data analysis.

##  2 Methodology

In this section different aspects of the methodology used in this work are
discussed. In Section 2.1, the Nested Sampling (NS) method is described, which
is a well-established algorithm used for probabilistic modelling and Bayesian
inference and is applied here as a benchmark algorithm. In Section 2.2, the
mathematical framework of the NF method is presented. Finally, in Section 2.3,
the applied forward model is discussed. The principles of operation of the
inversion algorithms are explained here, while the specifics of the
implementation are outlined in Section 4 of this paper.

###  2.1 Nested Sampling

In Bayesian analysis, evaluating the evidence and the posterior is trivial for
low dimensional data, described by a small number of parameters and by simple
probability distributions. For higher dimensions and more complex
distributions, the direct computation of the evidence and the posterior can
become intractable and other approaches are required. Algorithms such as
Markov-Chain Monte Carlo (Hastings (1970)) were designed to avoid the
evaluation of the evidence when only samples from the posterior distribution
are needed. Meanwhile, NS (Skilling (2006)), aims to efficiently compute an
approximation of the evidence from previously defined analytical likelihood
and prior functions. In addition, as is used in this work, samples from the
posterior can be produced as a by-product of the evidence evaluation. The idea
behind NS is to reduce the multi-dimensional integral defining the evidence
($Z$) in Equation 3 to a one-dimensional problem by defining the prior mass
($X$) via $dX=\pi(\boldsymbol{\theta})d\boldsymbol{\theta}$. The likelihood
can then be expressed in terms of the prior mass instead of the model
parameters, so that $L(X(\lambda))\equiv\lambda$, where

$X(\lambda)=\int_{L(\boldsymbol{\theta})>\lambda}\pi(\boldsymbol{\theta})d\boldsymbol{\theta}$
(4)

is the prior mass enclosed within the region of the parameter space bounded by
the likelihood contour $\lambda$. Therefore, Equation 3 becomes

$Z=\int_{0}^{1}L(X)dX,$ (5)

which can be computed using simple numerical integration methods.

The NS algorithm starts with sampling $N$ live points from the prior, and the
associated likelihood is computed for of each of these points. The point with
the lowest likelihood ($L_{1}$) is identified and removed from the set. A new
sample is drawn to replace it, which is only accepted if its associated
likelihood is $\geq L_{1}$. Setting $L_{1}$ to be a threshold shrinks the data
space volume that the samples are drawn from. This process is repeated in
iterations, tightening the likelihood bound each time, until the stopping
criteria is reached. For each iteration, the likelihood limit ($L_{i}$) as
well as the changes in the prior mass contained within this limit ($\Delta
X_{i}$) are stored. The integral in Equation 5 then can be approximated using
the trapezoid rule, such as $Z=\sum L_{i}\Delta X_{i}$ (Buchner (2023)). A
useful by-product of the process is that the removed live points can be used
to obtain posterior samples. At the end of the process, samples are drawn from
the assembly of all previous and final points with the $i^{th}$ point weighted
by the factor $L_{i}\Delta X_{i}$. These samples are representative of the
posterior.

The algorithm requires the likelihood and prior functions associated with the
problem to be pre-defined, which encode some previously known information
about the noise associated with the gravity survey itself and the constraints
that can be placed on the source parameters before any knowledge is gained
from the measured data. Sampling methods have been successfully applied to
gravity (Brown (2015), Rodgers (2017)) and combined gravity and magnetic
(Ghalenoei et al. (2021)) inversion problems. However, these suffer from high
computational costs, which scales with the complexity of the source model.
This limits the complexity of the problems that are viable to compute and the
usability of the method in the field, where results need to be obtained within
minutes. In this work, a NS method is applied as a benchmark algorithm to
evaluate the relative accuracy and speed of the NF inversion presented.

###  2.2 Normalising Flows

NF s are a type of generative neural network that are ideally suited for the
modelling and sampling of complex data distributions. The basis of the method
described in Papamakarios et al. (2021) is to express the data vector
($\boldsymbol{\theta}$) as a transformation ($T$) of another vector
($\boldsymbol{u}$), which is sampled from an arbitrary base distribution
($p(\boldsymbol{u})$). The PDF of $\boldsymbol{\theta}$ can then be defined
via the change of variables formula

$p(\boldsymbol{\theta})=p(\boldsymbol{u})\det|J_{T}(\boldsymbol{u})|^{-1},$
(6)

where $J_{T}$ is the Jacobian matrix containing the partial derivatives of $T$
with respect to $\boldsymbol{u}$. To be able to model rich and complex PDF s,
$T$ can be a composite transformation made up of a number of invertible and
differentiable functions. These stacked transformations are often implemented
as a neural network with trainable parameters ($\boldsymbol{\phi}$). The
network parameters are iteratively tuned during training with the aim to
minimise the loss function, which is defined as the Kullback-Leibler (KL)
divergence between the network output ($q(\boldsymbol{\theta})$) and the
target distribution ($p(\boldsymbol{\theta})$). When the network is fully
trained, the output of the last layer of the network after forward
propagation, usually referred to as the latent space ($q(\boldsymbol{u})$), is
equal to the base distribution. Then, samples from the base distribution
($p(\boldsymbol{u})=q(\boldsymbol{u})$) can be passed through the inverted
transformations to obtain a sample from the real data space and the posterior
function can be evaluated. Since $p(\boldsymbol{u})$ can be freely chosen, it
is beneficial to opt for a distribution which can be cheaply sampled from, and
it is often set to be a multi-dimensional unit Gaussian.

The algorithm described above is a generative model which can be used to
interpolate a discrete data set and generate samples from the resulting
continuous distribution. However, to apply the network to parameter inference,
a conditional input needs to be introduced alongside the training data. Assume
that the aim is to infer the posterior PDF of some parameters
($\boldsymbol{\theta}$) conditioned on some observation ($\boldsymbol{d}$).
Then, the NF can be trained to approximate the conditional posterior PDF
($p(\boldsymbol{\theta}|\boldsymbol{d},A)$), such that

$p(\boldsymbol{\theta}|\boldsymbol{d},A)\approx
q(\boldsymbol{\theta}|\boldsymbol{d},A)=q(\boldsymbol{\theta};\boldsymbol{\phi}|\boldsymbol{d})=p(\boldsymbol{u|\boldsymbol{d}})|\det
J_{T}(\boldsymbol{u};\boldsymbol{\phi})|^{-1}$ (7)

where $q(\boldsymbol{\theta};\boldsymbol{\phi}|\boldsymbol{d})$ is the output
of the trained neural network, $T$ denotes the transformations parameterised
by the network parameters $\boldsymbol{\phi}$ and $A$ is an assumed model for
the data.

To apply this method to gravity inversion, a data set of source parameters
($\boldsymbol{\theta}$) needs to be generated from the chosen prior and the
corresponding gravity measurements ($\boldsymbol{d}$) can be simulated via the
gravity forward model ($A$), which forms the conditional part of the data set
once simulated noise is added. Once the network is trained on this data set,
samples can be generated from the posterior distribution of source parameters
conditioned on a specific gravimetry survey of interest, only returning
parameters that are consistent with the measurements. The NF method offers a
significant improvement in terms of computation time as compared to NS, as
once trained, inversion results can be obtained within seconds and the trained
network can be reused for multiple surveys.

###  2.3 Forward model

In gravity modelling, the forward model is the function that describes the
one-to-one mapping from the underground density distribution to the
theoretical surface gravity measurements. In traditional inversion methods
iterative forward calculations take place and the calculated surface gravity
is compared to the true values for optimization. In a NS algorithm iterative
calculations are also required, as the forward model forms part of the
likelihood function defined in Section 4. When applying a generative neural
network for inversion, we are also required to define a forward model, but
these are calculated in advance to form a part of the training data set. In
this work an analytical forward model is applied, which is specific to
rectangular prism anomalies (Li and Chouteau (1998)).

Figure 1: Diagram of the rectangular prism density anomaly for which the
forward model of the gravity measured at the point $(x,y,z)$ is defined in
Equation 8. (Li and Chouteau (1998))

The prism, visualised in Figure 1, is defined to have its faces aligned with
the $X$, $Y$ and $Z$ axes, with limits $[\xi_{1},\xi_{2}]$,
$[\eta_{1},\eta_{2}]$ and $[\zeta_{1},\zeta_{2}]$ correspondingly. The
vertical component of the gravitational acceleration ($g$) at the location
$R=(x,y,z)$ due to this body is

$g(R;\boldsymbol{\theta})=G\rho\sum_{i=1}^{2}\sum_{j=1}^{2}\sum_{k=1}^{2}\mu_{ijk}\left[x_{i}\ln(y_{j}+r_{ijk})+y_{j}\ln(x_{i}+r_{ijk})-z_{k}\arctan\frac{x_{i}y_{j}}{z_{k}r_{ijk}}\right]$
(8)

where

$x_{i}=x-\xi_{i},y_{j}=y-\eta_{j},z_{k}=z-\zeta_{k},\\\
r_{ijk}=\sqrt{x_{i}^{2}+y_{j}^{2}+z_{k}^{2}},\\\
\mu_{ijk}=(-1)^{i}(-1)^{j}(-1)^{k},$ (9)

$G$ is the gravitational constant and $\rho$ is the density of the object. The
source parameters ($\boldsymbol{\theta}$) that were chosen to describe each
instance of the rectangular prism anomalies are listed in Table 1, and it is
discussed in Section 3 how these are translated into the prism limits
($\xi,\eta,\zeta$).

##  3 Test Cases

The performance of the inversion algorithm was tested on a simulated data set
of source models consisting of rectangular prism voids located underneath the
surface. This data set was chosen as a simple test case that can be described
by a handful of parameters and for which the results are straight-forward to
visualise and interpret. While the modelling of natural geological features
usually requires more complex models, a simple data set as such can be applied
to the modelling of some human-made architecture, such as tunnels or bunkers.
The network was tested on two different data representations to showcase the
performance on data with varying dimensionality.

In the low-dimensional (parameterised) case, the location and size of the
rectangular prisms were described by 7 parameters, listed in Table 1 alongside
the chosen prior PDF s. To generate the training set, 1 million samples were
drawn from the prior and the forward model was calculated for each sample
following the method described in Section 2.3. For this, the source parameters
can be translated into the prism limits ($\xi,\eta,\zeta$) in Equation 9 via
eg., $\xi_{1}=c_{x}-l_{x}/2$ and $\xi_{2}=c_{x}+l_{x}/2$. Since for the
specified forward model the faces of the prism need to be aligned with the
$X,Y$ and $Z$ axes, the rotation parameterised by $\alpha$ is applied to the
survey coordinates in the opposite direction, resulting in the same effect in
the gravity values. The density contrast between the background and the object
itself ($\rho$) was set to 1500 kg/m3, approximated from the assumption that
it is a hollow prism encased in topsoil. The comparison between the NF and NS
algorithms is only done for this test case.

In the higher-dimensional (voxelised) case, the source parameters were first
drawn from the same priors as before and were translated into a
512-dimensional array of densities of cubic voxels on an $8\times 8\times 8$
grid covering a region of 120 m in each direction. This was done by
calculating the fraction of each voxel that overlaps with the rectangular
prism, and assigning a density value according to the average of the
background and the object densities weighted by this fraction. In addition,
Gaussian noise with a standard deviation of 50 kg/m3 was added to the density
values, to simulate fluctuations in topsoil density. Each voxel is then
considered as an instance of a rectangular prism, for which the analytical
forward model is calculated, and finally the results are summed for all
voxels. Similar to the parameterised case, the prism limits for Equation 9 can
be computed from the coordinates of the centre and the size of each voxel. The
rotation angle $\alpha$ is accounted for during the translation of the prism
parameters into density values, therefore the survey itself is not rotated.
While in this case swapping to this data representation seems impractical as
it reduces the resolution of the model while increasing the data volume and
introducing broader priors, it serves as a stepping-stone towards the
modelling of more complex non-rectangular bodies.

For the survey simulations, the stations were set to be arranged on an
8$\times$8 uniform grid with a separation of 10 m in each direction, and the
elevation kept at a constant. The survey was chosen to cover a smaller area
that the voxel grid to avoid edge effects, which means that the prism is
allowed to be located outside of the survey region. Since in real gravimetry
surveys taken in the field it is difficult to take measurements on an exact
grid, the voxelised data set was also combined with a noisy survey
configuration. Gaussian noise with a standard deviation of 2 $m$ was added to
the $x$ and $y$ coordinates of the stations. In the case where the survey
locations do not change between examples in the data set, only the gravity
values were provided as the conditional input to the NF. However, in the case
of varying, noisy survey locations, the $x$ and $y$ coordinates for each
measurement were also included. Overall three different test cases were
considered; the parameterised data representation with survey measurements
arranged on a uniform grid, and the voxelised data representation with both a
uniform survey grid and one with added noise to the grid locations. For all
cases Gaussian noise with a standard deviation of 10 $\mu$Gal was also added
to the survey measurements to simulate instrumental noise, which was chosen as
an arbitrary small value compared to the magnitudes of the expected relative
gravity signals.

##  4 Implementation

The benchmark NS method was implemented using built-in functionalities in the
bilby software package (Ashton et al. (2019)), using the dynesty algorithm
(Koposov et al. (2023)). This method requires the definition of prior PDF s,
which are defined in Table 1 for the 7 source parameters. An analytical
likelihood function also needs to be specified, which is used to model the
noise on the measurements. For the purposes of this work, Gaussian noise with
a standard deviation of $\sigma=$ 10 $\mu$Gal was added to the simulated
gravimetry data, therefore a Gaussian likelihood function needs to be
implemented. For a set of proposed source parameters ($\theta$) this is
described as

$\mathcal{L}(d|R,\boldsymbol{\theta})=\prod_{i}\frac{1}{\sqrt{2\pi\sigma^{2}}}\exp{\left(\frac{-(d_{i}-g(R_{i};\boldsymbol{\theta}))^{2}}{2\sigma^{2}}\right)},\\\
d_{i}=g(R_{i};\boldsymbol{\theta_{0}})+n_{i}$ (10)

where it is assumed that $d_{i}$ are independent data points measured at
locations $R_{i}$. The data is assumed to consist of a forward model
$g(R_{i};{\boldsymbol{\theta}})$, which is described in Section 2.3 and
computed from the true source parameters ($\boldsymbol{\theta_{0}}$), and
noise ($n_{i}$) sampled from a Gaussian distribution with standard deviation
$\sigma$. This is then compared to the forward model computed from the sampled
source parameters ($\boldsymbol{\theta}$) at each iteration of the algorithm
and the total likelihood is found from the product of each component at
locations $R_{i}$.

Table 1: The 7 parameters that define the location and size of the prism. The prism can only rotate around the z axis centred on $(c_{x},c_{y})$, and the rotation angle is defined as $\alpha$, which is constrained to $<\pi/2$. If the prism were allowed to rotate by more than $\pi/2$, the same configuration could be reproduced by taking a prism and rotating it by the angle $\pi/2+\alpha_{0}$ and by taking a prism with its $l_{x}$ and $l_{y}$ coordinates swapped and rotating it by $\alpha_{0}$, hence introducing a degeneracy in the posterior. In addition, the limits on the location and size in the $z$ dimension enforce that the prism remains below the surface. Parameter | Description | Prior
---|---|---
$c_{x}$ | Centre of the prism in the $x$ dimension. | Uniform(-60,60)
$c_{y}$ | Centre of the prism in the $y$ dimension. | Uniform(-60,60)
$c_{z}$ | Centre of the prism in the $z$ dimension. | Uniform(-60,20)
$l_{x}$ | Length of the prism in the $x$ dimension. | Uniform(0,120)
$l_{y}$ | Length of the prism in the $y$ dimension. | Uniform(0,120)
$l_{z}$ | Length of the prism in the $z$ dimension. | Uniform(0,80)
$\alpha$ | The rotation angle around the z axis. | Uniform(0,${\pi}/{2}$)

The NF method was implemented using glasflow (Williams et al. (2023)), a
python package based on pytorch (Paszke et al. (2019)) and nflows (Durkan et
al. (2020)). Before training, the data set consisting of 1 million data points
was split into 90% training and 10% validation data and was normalised using
the scikit-learn MinMaxScaler tool, which scales the data separately for each
dimension so that the values lie within the $[0,1]$ range. This was done for
both the data set of source models and the conditional data consisting of
survey measurements. The chosen architecture is built up of RealNVP (Dinh et
al. (2017)) transformations and is described by the parameters listed in Table
2. The final network parameters were selected based on rigorous optimisation
using Weights & Biases (Biewald (2020)) hyper-parameter optimisation tools.
The network architecture was optimised for the parameterised data set, which
was tested most rigorously, and was kept unaltered for the remaining test
cases to promote the flexibility and user-friendliness of the method.

The network was trained using the pytorch built-in Adam optimiser with an
initial learning rate of 0.001 and a ReduceLROnPlateau learning rate
scheduler, which reduces the learning rate by a factor of 0.05 whenever the
validation loss stops decreasing for 100 epochs. The training was completed
when the learning rate reached $10^{-6}$ and the validation loss reached a
plateau. Training can be done without the implementation of learning rate
scheduler, however this can make the process more efficient and allows the
user to utilise various learning rate values during a single training run. The
network architecture was constructed so that it can remain unchanged for the
three different test cases, however the required training time varied from
$\sim$4.5 hours for the parameterised case to $\sim$13 hours for the voxelised
case with a noisy survey grid using a NVIDIA Geforce RTX 3090 GPU. This
increase in training time is largely due to the increase in data volume. Once
the network is trained, 1000 samples can be obtained within 7 seconds as
compared to the $\sim$20 minute computation time associated with a single NS
run, and the trained network can be reused for multiple inversion cases. The
whole pipeline with examples is available on GitHub (Rakoczi (2023)).

Table 2: The parameters that define the RealNVP (Dinh et al. (2017)) neural network architecture that was implemented in this work. Only the parameters that differ from the default values implemented in glasflow (Williams et al. (2023)) are listed here. Parameter | Value
---|---
Transforms | 12
Blocks per transforms | 2
Neurons | 64
Batch size | 5000
Initial learning rate | 0.001

##  5 Results

In this section the results for the gravity inversion using the three
different data set discussed in Section 3 are discussed. The devised testing
methodologies are outlined and are used to scrutinise the performance and
accuracy of the presented NF inversion method, and comparisons are made with
the NS wherever relevant.

###  5.1 Survey consistency test

When using the trained network for inversion, samples are drawn from the
posterior probability distribution over the source parameters, conditioned on
the specific gravimetry survey of interest. The samples represent individual
instances of the source parameters that are all consistent with the survey
data within the limits of instrumental noise. To test this consistency, a
number of samples can be drawn from the posterior distribution of the NF and
forward modelled to compare to the original survey data. Summary statistics
based on 2000 samples are shown in Figure 2 for the three separate test cases.
The pixel-wise mean of the forward-modelled samples aligns well with the input
survey and the standard deviation is $\sim$1% of the amplitude of the relative
gravity signal for the parameterised case, and $\sim$5% for the voxelised
cases. This difference can be attributed to the relative difficulties of the
different test cases due to higher dimensionality and wider prior spaces, as
well as to the additional noise contributions from the background density
fluctuations, which are not present in the parameterised data. These results
demonstrate that the samples represent valid solutions to the inversion
problem.

Figure 2: Plot showing the comparison between the true gravimetry survey and
the simulated surveys from the NF output. Samples were drawn from the
posterior distribution of source parameters produced by the 3 different test
cases, and each were forward modelled using the method outlined in Section
2.3. Column 1 is the input gravimetry survey, column 2 is the pixel-wise mean
computed from 2000 samples and column 3 is the standard deviation.

###  5.2 Statistical consistency test

Another important test to conduct is to prove that the resulting posterior
from the sampling method is statistically consistent and free of bias. The
probability-probability (p-p) plot is a standard method developed to validate
Bayesian algorithms (Samantha R Cook and Rubin (2006), Talts et al. (2020)).
For this method, probability is considered as defined by frequentist
statistics, where the interest is the frequency ($F$) at which the true
parameter value lies within the $F\%$ credible interval of the posterior
probability distribution. In frequentist probability tests having a
sufficiently high number of test cases is essential, however a trade-off with
computational cost needs to be considered. In this work, 100 test cases were
randomly chosen from the prior to produce each p-p plot.

To construct the p-p plot shown in Figure 3, the following process was
applied; first 100 sets of source parameters were generated from the prior and
were forward modelled to compute the simulated gravity surveys with Gaussian
noise added to the measurements. Then, a number of samples were drawn
conditioned on each of these surveys separately from the posterior
distribution of source parameters output by the inference method. The
marginalised CDF at the location of the true parameter value in each dimension
was then calculated for each test case separately. A histogram of the values
of the obtained CDF s at the true parameter locations was computed, which
approximates to a uniform distribution when the method is consistent. Finally,
the CDF of this distribution was plotted, resulting in a series of lines
following the diagonal in the ideal case. However, due to the relatively low
number of test cases considered when constructing this plot, slight deviations
from the diagonal are expected. The method is considered consistent as long as
the curves lie within the 68%, 95% and 99.7% confidence intervals, which are
marked by the grey regions on the plot. When significant deviations from the
diagonal are observed, the shape of the curve can hint at specific bias in the
posterior. For example, if an S shape is observed, the posterior is over- or
under-constraining the parameter space, while in case of an bow shape, the
posterior is shifted from its ideal location.

The $p$-value was also computed from the maximum distance between the diagonal
line and the curves corresponding to the various parameter dimensions using
the Kolmogorov-Smirnov test. It is the probability of obtaining a value of a
distance statistic, or greater, under the null hypothesis that the underlying
cumulative distribution is diagonal. Extremely small $p$-values indicate that
the null hypothesis is unlikely. The combined value displayed on the top is
obtained using Fisher’s method to compute the $p$-value where the null
hypothesis is that each parameter obeys its own null hypothesis. The
$p$-values for each curve as well as the combined value are all expected to
appear to be random samples from a uniform distribution between 0 and 1.
Figure 3 shows evidence of the statistical consistency of the NS and the
various NF inversion test cases.

Figure 3: The various p-p plots obtained from the two different inversion
methods and three data sets; (a) Nested Sampling inversion of the
parameterised data set, (b) Normalising Flows inversion of the parameterised
data set, (c) Normalising Flows inversion of the voxelised data set and (d)
Normalising Flows inversion of the voxelised data set with added noise to the
survey locations. For the voxelised test case only a randomly picked 10
parameters {$q_{1}...q_{10}$} of the 512 dimensions were plotted for better
visibility. The p-values are displayed in the legend for each curve
separately, and the combined value is displayed above each graph. The method
of obtaining these plots is described in detail in Section 5.2.

###  5.3 Parameterised posterior comparison

We are also interested in visualising the posterior PDF of source parameters,
and testing the accuracy of the inference from the NF compared to the NS
method. In the case of a handful of parameters, the results of the inference
can be visualised on a so-called corner plot. This array of graphs depicts the
marginalised 1-dimensional probability distributions of individual parameters
on the diagonal panels, and the 2-dimensional combined distributions of pairs
of parameters in the panels below the diagonal. The locations of the true
parameter values were marked for reference, which are expected to be
consistent with the posterior contours. This plot can also be used to compare
probability distributions obtained from different inference methods applied to
the same input data under the same model and prior assumptions. If both
methods perform well, it is expected that these would constrain the parameter
space at the same level of accuracy, and ideally would be identical within the
uncertainties of finite sampling. The results for the inversion of the
gravimetry survey in the first row and first column of Figure 2 are shown in
Figure 4.

Some interesting correlations can be observed in the arising distributions.
First of all, there is a negative correlation between the extent in the $z$
dimension ($l_{z}$) and depth of the centre of the prism ($c_{z}$). This can
be contributed to the fact that the same gravity anomaly can be linked to a
family of solutions ranging from a small and shallow object, to a larger
object deeper underneath the surface. From vertical gravitational acceleration
data only, these are not possible to distinguish. However, other less obvious
correlations appear as well. Both the $c_{z}$ and $l_{z}$ parameters are
correlated with the $l_{x}$ and $l_{y}$ parameters. This can be explained by
the angle between the gravity vector pointing to each point along the edge of
the prism in the $x-y$ plane and the vertical component of this vector
remaining constant. Since the density of the prism is constrained, the
contribution of each point within the prism to the total gravitational
acceleration vector at each survey location is also constrained. Since the
vertical component of gravity is directly measured, for the survey
measurements to remain unchanged for prisms located at different depths, the
angle between the gravity vector and its vertical component has to remain the
same. Therefore, when considering survey locations just outside the boundaries
of the prism, the edge of the prism has to lie closer to these points when
closer to the surface, and further away when deeper beneath the surface,
effectively shrinking the prism in the $x$ and $y$ dimensions. This effect can
be visualised using the analogy of shining a torch onto a screen, placing a
box in-between them and observing the shadow cast onto the screen. When the
box is closer to the screen (the prism is closer to the surface), its image
will be closely aligned with its size. However when the box is moved towards
the torch (the prism is deeper beneath the surface), its shadow becomes
larger.

While the corner plot in Figure 4 shows that the results from the two methods
agree and have comparable accuracy for a single test case, it is also
necessary to understand the general performance of the NF. In Figure 5 a
histogram of the Jensen-Shannon (JS) divergence between the 1-dimensional
marginalised posterior PDF s obtained from the two different methods is shown.
This provides a measure of distance between the distributions, with a minimum
value of 0 and a maximum of $\ln{2}$. It is shown here that 98% of the
distributions have a JS divergence $<0.1$ with a median of 0.0039. These
results indicate that the NF can consistently produce accurate posteriors
assuming that the NS approach produces the true posterior. One of the
potential reasons explaining the rare cases of higher JS divergence values is
the emergence of topologies in the true posterior PDF which are difficult to
model using NF s, such as multi-modality (Dinh et al. (2020)). In a few
specific cases when the $\alpha$ parameter lies close to the edge of the prior
space, the posterior of $\alpha$, $l_{y}$ and $l_{x}$ become multi-modal. This
is due to a specific degeneracy in the formulation of the prior, where a prism
rotated by $\alpha=\pi/2$ with specific $l_{x}$ and $l_{y}$ dimensions would
look the same as a prism with $\alpha=0$ with its dimensions swapped around.

Figure 4: The corner plot comparing the posterior probability distributions
obtained from the Nested Sampling (blue) and the Normalising Flows (orange)
for the specific test case. The 1-dimensional marginalised probability
distributions are shown on the diagonal for each parameter separately, with
the mean of the posterior distribution obtained from the Normalising Flows
displayed above. The dashed lines and the error bounds correspond to the
16$\%$ and 84$\%$ quantiles of the distribution. The black lines and markers
correspond to the true parameter values. The method of obtaining this plot is
explained in detail in Section 5.3. On the right hand side a table of the
Jensen-Shannon (JS) divergence between the 1-dimensional marginalised
posterior distributions produced by the Nested Sampling and Normalising Flows
methods are shown for this test case.  Figure 5:  A histogram of the Jensen-
Shannon (JS) divergence between the 1-dimensional marginalised posterior
distributions of source parameters produced by the Nested Sampling and
Normalising Flows methods for 100 test cases. This provides a measure for the
distance between the posterior distributions from the Normalising Flow and the
assumed ideal distribution, which is obtained from the Nested Sampling method.
The range of possible values for the JS divergence is $[0,\ln(2)]$, with
larger values indicating larger differences between the distributions. The
number of counts of JS divergence values within each bin for the different
parameter dimensions are marked by different colours and are stacked on top of
each other to show and overall distribution. Note the logarithmic scale on the
$x$-axis for better visibility.

###  5.4 Density distribution comparison

In case of the voxelised source model representation, each sample from the NF
is a 512-dimensional array of parameters representing the voxel densities. Due
to the high dimensionality, it is inconvenient to directly visualise the
posterior on a corner plot, and it is more informative to compute summary
statistics about the distribution. The mean, mode and standard deviation of
10000 samples from the posterior distribution of the source parameters are
shown in Figure 6 alongside the target source model, which is a single piece
of data the method is tested on. While in these test cases the source model
from which the gravimetry survey was simulated is known, it is not expected
that the inversion would return the true model when only the survey data is
provided due to the inherent degeneracies and noise added during the
simulation. It can however be expected that the resulting distribution of
source parameters includes the true model. To allow comparison between the
results from the parameterised and the voxelised inversion, the results from
the parameterised case were also translated into density values and the same
summary statistics were plotted in the topmost subplot of Figure 6.

The mean and standard deviation of the distribution of the source parameter
vectors were computed element-wise. The mean provides a visual representation
of all the locations where the prism is potentially located, and the standard
deviation is indicative of the uncertainty in the density predicted for the
specific voxel, which is highest around the predicted edges of the prism. The
mode is found by choosing the sample from the NF that has the highest
probability associated with it, hence it marks the single solution
approximating the peak of the multi-dimensional posterior. This is not
expected to closely agree with the true source model, it is only the most
probable solution based on the data. This phenomenon is illustrated in the
bottom subplot of Figure 6, where the mode of the distribution disagrees with
the true model in terms of its extent and location in the $z$ dimension as
well as its size in the $y$ dimension, while in the top subplot the mode
agrees with the truth more closely. In addition, the correlations observed in
Section 5.3 appear in Figure 6 as a cone-like feature in the mean of the
slices along the $x$ and $y$ dimensions, with the cone getting narrower
further from the surface. Therefore, the deeper the prism is, the smaller it
has to be in the $x$ and $y$ dimensions to remain consistent with the
gravimetry survey data. This is most apparent in the bottom panel presenting
results from an inversion with a noisy survey grid. In general, the
uncertainty in the results for this case appears lower than for the inversion
of the survey taken on a uniform grid, which suggests that a more randomised
survey configuration is more ideal for the NF inversion method.

Figure 6: Plot comparing the output of the Normalising Flows inversion to the
target density distribution. The three subplots correspond to the three
different test cases; the parameterised data set, the voxelised data set, and
the voxelised data set with added noise to the survey locations, from top to
bottom respectively. Each column corresponds to the $x$, $y$ and $z$
dimensions in order along which the volume is sliced, and the subcolumns in
order are the target distribution, and the mean, mode and standard deviation
of the samples from the posterior distribution of source parameters output by
the Normalising Flow. The rows correspond to the different slices, as marked
on the diagrams on top.

##  6 Discussion

There are various aspects that determine how well the parameter space can be
constrained for a specific inverse problem. Firstly, the distribution is
broadened by the inherent non-uniqueness which is ubiquitous in inverse
problems in gravimetry. The method presented here has the advantage over
conventional single-solution inversion that it does not ignore this non-
uniqueness and incorporates all possible solutions in the results. More
importantly, the resolution of the gravimetry survey and the instrumental
noise have a large effect on the accuracy with which the source parameters can
be inferred. These effects cannot be removed and can only be reduced by
intelligent survey design, good instruments and repeated measurements.
However, as the target distribution becomes narrower, the more challenging the
inference problem becomes and training time and network complexity increases.
The level of noise included in this work is relatively low compared to that of
currently available gravimeters on the market and only a simple Gaussian noise
profile was implemented here. When designing a model for surveys taken with a
specific gravimeter, more informed instrumental noise profiles can be applied.

One benefit of using machine learning for inversion is the speed at which
results can be obtained once the network is trained. However, this feature can
only be taken advantage of if the trained network is reusable, which means
that the chosen priors need to be general enough to be applicable to a range
of gravimetry surveys, but restrictive enough not to return irrelevant
results. A simple uniform prior space was implemented here, but for real-life
applications it is recommended to implement more restrictive priors with as
much information content as possible to make the inference faster and more
accurate. In addition, increasing model complexity can make the data set more
applicable to real-life situations. Any feature that can be parameterised and
its gravitational effect can be modelled, such as the geological background or
previous knowledge gained from other geophysical surveys, can be incorporated
in the source models. However, with increasing model complexity the volume of
the training data required to span the prior space is likely to increase too.
Depending on the speed of the forward model, data generation can prove to be a
bottleneck and a trade-off between the advantages of more complex models and
the computation time needs to be considered. The performance of the NF was
only tested for up to a few hundred data dimensions. Some inversion targets
might require even higher complexity to model, therefore exploring the
performance at higher dimensions and combining the inference with data
reduction methods are desirable avenues of research.

Another restriction in the current configuration of the method is that the
number of points in the gravimetry survey needs to be decided before the
measurements are taken and kept at a constant across different surveys if the
trained network is to be reused. For more flexibility, it would be beneficial
to develop the method to be able to take in a variable number of survey points
at unrestricted locations. This would also better allow the inversion to be
used in the field, for which the NF method is ideally suited due to fast
computation times. It could be used to immediately obtain inversion results,
to guide the user towards the best survey locations to be measured and to
indicate when the desired uncertainty threshold was found and the survey can
be concluded.

##  7 Conclusions

Bayesian inference methods have been successfully applied to gravity
inversion, however standard sampling methods suffer from long computation
times, which increases with the dimensionality of the data set. Here, a
gravity inversion method was implemented using Normalising Flows, which can
improve on standard sampling methods in terms of computation time and
flexibility. First, the method was tested with a data set consisting of
rectangular prism bodies described by 7 source parameters. The survey
measurements were simulated using the forward model described in Li and
Chouteau (1998) and Gaussian noise was added to the measurements to simulate
instrumental noise. The performance of Normalising Flows for this low-
dimensional example was compared to that of Nested Sampling, and it was shown
that the accuracy is comparable, while the Normalising Flow can produce
results within $\sim$ 7 seconds instead of $\sim$ 20 minutes. To test the
network with a higher-dimensional data set, the rectangular prisms were
translated into densities defined on a uniform 8$\times 8\times$8 cubic grid.
It was shown that the posterior distributions of the source parameters remain
tractable and accurate results are produced. The method was tested with two
different survey configurations, first with a simple uniform 8$\times$8 grid
followed by one with added Gaussian noise to the grid locations to more
closely simulate realistic survey configurations. As future work, we aim to
test the network with more complicated noise profiles and a more flexible
survey configurations.

##  8 Acknowledgments

The authors would like to extend their gratitude to Dr. Michael Williams and
Christian Chapman-Bird for their invaluable expert advice and insightful
ideas, which greatly enriched the contents of this paper. Special thanks go to
Mahijs Reinier Koymans and Dr. Craig Whitehill for their revisions and
constructive feedback, which played an instrumental role in enhancing the
quality of this manuscript. Their contributions have been indispensable in the
completion of this research project. The authors acknowledge the usage of the
analysis software nflows (Durkan et al. (2020)), glasflow (Williams et al.
(2023)), and bilby (Ashton et al. (2019)).

This work was primarily funded by UK Research and Innovation (Grant No.
ST/V005499/1). CM acknowledges partial support from Science and Technology
Research Council (Grant No. ST/L000946/1). For the purpose of open access, the
authors have applied a Creative Commons Attribution (CC BY) licence to any
Author Accepted Manuscript version arising.

## References

  * Torge [1989] Wolfgang Torge.  _Gravimetry_.  De Gruyter, 1989.  ISBN 9783110107029.
  * Milano et al. [2021] Maurizio Milano, Ramin Varfinezhad, Hamid Bizhani, Meysam Moghadasi, Ali Nejati Kalateh, and Hamidreza Baghzendani.  Joint interpretation of magnetic and gravity data at the golgohar mine in iran.  _Journal of Applied Geophysics_ , 195:104476, 2021.  ISSN 0926-9851.  doi:https://doi.org/10.1016/j.jappgeo.2021.104476.  URL https://www.sciencedirect.com/science/article/pii/S092698512100224X.
  * Pearse et al. [2021] Jillian Pearse, Andrés Cárdenas Contreras, Carol Vanessa Barrera Lopez, Nataly Castillo Ruiz, Henry Martínez Gómez, and Jean Baptiste Tary.  Gravity survey and modelling of the nemocón salt mine, colombia.  _Near Surface Geophysics_ , 19(3):365–376, 2021.  ISSN 1873-0604.  doi:https://doi.org/10.1002/nsg.12146.  URL https://www.earthdoc.org/content/journals/10.1002/nsg.12146.
  * Eppelbaum [2010] Lev Eppelbaum.  Application of microgravity at archaeological sites in israel: Some estimation derived from 3-d modeling and quantitative analysis of gravity field.  volume 1, 03 2010.  doi:10.4133/1.3176721.
  * Carbone et al. [2020] D. Carbone, L. Antoni-Micollier, G. Hammond, E. de Zeeuw van Dalfsen, E. Rivalta, C. Bonadonna, A. Messina, J. Lautier-Gaud, K. Toland, M. Koymans, K. Anastasiou, S. Bramsiepe, F. Cannavò, D. Contrafatto, C. Frischknecht, F. Greco, G. Marocco, R. Middlemiss, V. Ménoret, A. Noack, L. Passarelli, D. Paul, A. Prasad, G. Siligato, and P. Vermeulen.  The newton-g gravity imager: Toward new paradigms for terrain gravimetry.  _Frontiers in Earth Science_ , 8, 2020.  ISSN 2296-6463.  doi:10.3389/feart.2020.573396.  URL https://www.frontiersin.org/article/10.3389/feart.2020.573396.
  * Gadirov et al. [2022] Vagif Gadirov, Ekrem Kalkan, Adil Ozdemir, Yildiray Palabiyik, and Kamran Gadirov.  Use of gravity and magnetic methods in oil and gas exploration: Case studies from azerbaijan.  4:143–156, 07 2022.
  * Zhdanov [2002] Michael S. Zhdanov.  Chapter 1 - forward and inverse problems in geophysics.  In _Geophysical Inverse Theory and Regularization Problems_ , volume 36 of _Methods in Geochemistry and Geophysics_ , pages 3–28. Elsevier, 2002.  doi:https://doi.org/10.1016/S0076-6895(02)80038-5.  URL https://www.sciencedirect.com/science/article/pii/S0076689502800385.
  * Sen and Stoffa [2013] M.K. Sen and P.L. Stoffa.  _Global Optimization Methods in Geophysical Inversion_.  Cambridge University Press, 2013.  ISBN 9781107011908.  URL https://books.google.co.uk/books?id=GDazJPWr3K4C.
  * Lines and Treitel [1984] L.R. Lines and S. Treitel.  A review of least-squares inversion and its application to geophysical problems.  _Geophysical Prospecting_ , 32(2):159–186, 1984.  doi:https://doi.org/10.1111/j.1365-2478.1984.tb00726.x.  URL https://onlinelibrary.wiley.com/doi/abs/10.1111/j.1365-2478.1984.tb00726.x.
  * Tikhonov and Arsenin [1977] A.N. Tikhonov and V.I.A. Arsenin.  _Solutions of Ill-posed Problems_.  Halsted Press book. Winston, 1977.  ISBN 9780470991244.  URL https://books.google.co.uk/books?id=ECrvAAAAMAAJ.
  * Zhdanov [2009] Michael S. Zhdanov.  New advances in regularized inversion of gravity and electromagnetic data.  _Geophysical Prospecting_ , 57(4):463–478, 2009.  doi:https://doi.org/10.1111/j.1365-2478.2008.00763.x.  URL https://onlinelibrary.wiley.com/doi/abs/10.1111/j.1365-2478.2008.00763.x.
  * Bayes [1763] T. Bayes.  An essay towards solving a problem in the doctrine of chances.  _Philosophical Transactions_ , 53:370–418, 1763.
  * Sivia and Skilling [2006] D. S. Sivia and J. Skilling.  _Data analysis: a Bayesian tutorial_.  Oxford University Press, Oxford, 2nd edition, 2006.
  * Brown [2015] G Brown.  Bayesian mass anomaly estimation with measurements of gravity.  In _2nd IET International Conference on Intelligent Signal Processing 2015 (ISP)_ , pages 1–6, 2015.  doi:10.1049/cp.2015.1754.
  * Rodgers [2017] Anthony David Rodgers.  _Bayesian inference from terrestrial gravimetry measurements of near-surface anomalies using a bespoke reversible-jump Markov chain Monte Carlo algorithm_.  PhD thesis, University of Birmingham, 2017.
  * Rezende and Mohamed [2016] Danilo Jimenez Rezende and Shakir Mohamed.  Variational inference with normalizing flows, 2016.
  * Valentine and Sambridge [2022] Andrew Valentine and Malcolm Sambridge.  Emerging directions in geophysical inversion, 2022.
  * Yang et al. [2022] Qianguo Yang, Xiangyun Hu, Shuang Liu, Qu Jie, Huaijiang Wang, and Qiuhua Chen.  3-d gravity inversion based on deep convolution neural networks.  _IEEE Geoscience and Remote Sensing Letters_ , 19:1–5, 2022.  doi:10.1109/LGRS.2020.3047131.
  * Zhang et al. [2022] Lianzhi Zhang, Guibin Zhang, Yao Liu, and Zhenyu Fan.  Deep learning for 3-d inversion of gravity data.  _IEEE Transactions on Geoscience and Remote Sensing_ , 60:1–18, 2022.  doi:10.1109/TGRS.2021.3110606.
  * Zhao et al. [2021] Xuebin Zhao, Andrew Curtis, and Xin Zhang.  Bayesian seismic tomography using normalizing flows.  _Geophysical Journal International_ , 228(1):213–239, 07 2021.  ISSN 0956-540X.  doi:10.1093/gji/ggab298.  URL https://doi.org/10.1093/gji/ggab298.
  * Hastings [1970] W. K. Hastings.  Monte Carlo sampling methods using Markov chains and their applications.  _Biometrika_ , 57(1):97–109, 04 1970.  ISSN 0006-3444.  doi:10.1093/biomet/57.1.97.  URL https://doi.org/10.1093/biomet/57.1.97.
  * Skilling [2006] John Skilling.  Skilling, j.: Nested sampling for general bayesian computation. bayesian anal. 1(4), 833-860.  _Bayesian Analysis_ , 1:833–860, 12 2006.  doi:10.1214/06-BA127.
  * Buchner [2023] Johannes Buchner.  Nested sampling methods, 2023.
  * Ghalenoei et al. [2021] Emad Ghalenoei, Jan Dettmer, Mohammed Y Ali, and Jeong Woo Kim.  Gravity and magnetic joint inversion for basement and salt structures with the reversible-jump algorithm.  _Geophysical Journal International_ , 227(2):746–758, 07 2021.  ISSN 0956-540X.  doi:10.1093/gji/ggab251.  URL https://doi.org/10.1093/gji/ggab251.
  * Papamakarios et al. [2021] George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan.  Normalizing flows for probabilistic modeling and inference, 2021.
  * Li and Chouteau [1998] Xiong Li and Michel Chouteau.  Three-dimensional gravity modeling in all space.  _Surveys in Geophysics_ , 19, 1998.  doi:10.1023/A:1006554408567.  URL https://doi.org/10.1023/A:1006554408567.
  * Ashton et al. [2019] Gregory Ashton, Moritz Hübner, Paul D. Lasky, Colm Talbot, Kendall Ackley, Sylvia Biscoveanu, Qi Chu, Atul Divakarla, Paul J. Easter, Boris Goncharov, Francisco Hernandez Vivanco, Jan Harms, Marcus E. Lower, Grant D. Meadors, Denyz Melchor, Ethan Payne, Matthew D. Pitkin, Jade Powell, Nikhil Sarin, Rory J. E. Smith, and Eric Thrane.  Bilby: A user-friendly bayesian inference library for gravitational-wave astronomy.  _The Astrophysical Journal Supplement Series_ , 241(2):27, apr 2019.  doi:10.3847/1538-4365/ab06fc.  URL https://doi.org/10.3847%2F1538-4365%2Fab06fc.
  * Koposov et al. [2023] Sergey Koposov, Josh Speagle, Kyle Barbary, Gregory Ashton, Ed Bennett, Johannes Buchner, Carl Scheffler, Ben Cook, Colm Talbot, James Guillochon, Patricio Cubillos, Andrés Asensio Ramos, Ben Johnson, Dustin Lang, Ilya, Matthieu Dartiailh, Alex Nitz, Andrew McCluskey, Anne Archibald, Christoph Deil, Dan Foreman-Mackey, Danny Goldstein, Erik Tollerud, Joel Leja, Matthew Kirk, Matt Pitkin, Patrick Sheehan, Phillip Cargile, Ruskin Patel, and Ruth Angus.  joshspeagle/dynesty: v2.1.2, June 2023.  URL https://doi.org/10.5281/zenodo.7995596.
  * Williams et al. [2023] Michael J. Williams, jmcginn, federicostak, and John Veitch.  uofgravity/glasflow: v0.2.0, February 2023.  URL https://doi.org/10.5281/zenodo.7598678.
  * Paszke et al. [2019] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  Pytorch: An imperative style, high-performance deep learning library.  In _Advances in Neural Information Processing Systems 32_ , pages 8024–8035. Curran Associates, Inc., 2019.  URL http://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf.
  * Durkan et al. [2020] Conor Durkan, Artur Bekasov, Iain Murray, and George Papamakarios.  nflows: normalizing flows in PyTorch, November 2020.  URL https://doi.org/10.5281/zenodo.4296287.
  * Dinh et al. [2017] Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio.  Density estimation using real nvp, 2017.
  * Biewald [2020] Lukas Biewald.  Experiment tracking with weights and biases, 2020.  URL https://www.wandb.com/.  Software available from wandb.com.
  * Rakoczi [2023] Henrietta Rakoczi.  G.I.Flow, 11 2023.  URL https://github.com/rakoczi-h/flowinv.
  * Samantha R Cook and Rubin [2006] Andrew Gelman Samantha R Cook and Donald B Rubin.  Validation of software for bayesian models using posterior quantiles.  _Journal of Computational and Graphical Statistics_ , 15(3):675–692, 2006.  doi:10.1198/106186006X136976.  URL https://doi.org/10.1198/106186006X136976.
  * Talts et al. [2020] Sean Talts, Michael Betancourt, Daniel Simpson, Aki Vehtari, and Andrew Gelman.  Validating bayesian inference algorithms with simulation-based calibration, 2020.
  * Dinh et al. [2020] Laurent Dinh, Jascha Sohl-Dickstein, Hugo Larochelle, and Razvan Pascanu.  A rad approach to deep mixture models, 2020.


#  _Zeroth-Order SciML_ : Non-intrusive Integration of Scientific Software
with Deep Learning

Ioannis Tsaknakis  Bhavya Kailkhura  Sijia Liu  Donald Loveland  James
Diffenderfer  Anna Maria Hiszpanski  Mingyi Hong

###### Abstract

Using deep learning (DL) to accelerate and/or improve scientific workflows can
yield discoveries that are otherwise impossible. Unfortunately, DL models have
yielded limited success in complex scientific domains due to large data
requirements. In this work, we propose to overcome this issue by integrating
the abundance of scientific knowledge sources (SKS) with the DL training
process. Existing knowledge integration approaches are limited to using
differentiable knowledge source to be compatible with first-order DL training
paradigm. In contrast, our proposed approach treats knowledge source as a
black-box in turn allowing to integrate virtually any knowledge source. To
enable an end-to-end training of SKS-coupled-DL, we propose to use zeroth-
order optimization (ZOO) based gradient-free training schemes, which is non-
intrusive, i.e., does not require making any changes to the SKS. We evaluate
the performance of our ZOO training scheme on two real-world material science
applications. We show that proposed scheme is able to effectively integrate
scientific knowledge with DL training and is able to outperform purely data-
driven model for data-limited scientific applications. We also discuss some
limitations of the proposed method and mention potentially worthwhile future
directions.

Machine Learning, ICML

##  1 Introduction

Deep learning (DL) provides incredible opportunities to answer some of the
most important and difficult questions in a wide range of scientific
applications. However, many of the scientific use cases such as materials
discovery do not have access to large amounts of data required to train DL
models. In such data-limited complex scientific domains, purely data-driven DL
models have yielded limited success or led to unsatisfactory results. At first
glance, the task of training an accurate deep neural network (DNN) from a
small amount of data appears impossible. Coming to the rescue is a wealth of
prior domain knowledge in the form of scientific software codes – most of
which is currently not being utilized as purely-data driven models
inefficiently try to learn already known chemistry from scratch.

The biggest challenge in integrating scientific knowledge into the learning
process is formulating it in a manner that is compatible with the current
training paradigm. Given the fact that nearly all DNN training approaches are
based on differentiable programming, all existing efforts in integrating
scientific knowledge(von Rueden et al., 2019) are either (a) restricted to
using a very limited subset of knowledge that is already in a differentiable
form, or (b) spending significant effort on modifying knowledge sources to
support automatic differentiability (autodiff). As an example of (a), Physics-
aware machine learning (ML) area (Karniadakis et al., 2021) uses analytical
equations, which only capture a very limited aspect of otherwise complex
system. As an example of (b), Differentiable Physics/Chemistry area (Kasim et
al., 2022; Belbute-Peres et al., 2020) requires updating, or altogether
rewriting, existing software codes to provide autodiff functionality, which is
extremely costly. This points to a crucial roadblock in using DL in scientific
domains: we do not have methods that can exploit all available scientific
knowledge without being intrusive, i.e., requiring a rewrite of existing
software codes.

To overcome the challenge of effortless integration of scientific knowledge
with DL, we propose making changes to the ML paradigm to be compatible with
scientific software, rather than the other way around. Specifically, we
propose to leverage gradient-free optimization approaches to train scientific
software coupled deep neural nets. Note that resulting “sciML” system does not
provide gradients rendering popular first-order methods (e.g., stochastic
gradient descent) inapplicable. Unlike the current gradient-based training
paradigm, i.e., backprop, our proposed technique aims to train sciML models
accurately using only function queries (or forward passes), thereby bypassing
the stringent requirement of gradient availability. To achieve this, we adopt
a subset of gradient-free optimization, Zeroth-Order Optimization (ZOO) (Liu
et al., 2020, 2018), and train sciML system in an end-to-end manner. To show
the feasibility of our approach, we apply our technique on a popular material
discovery pipeline on two datasets. This pipeline is used to produce novel
molecules with desired properties by searching for molecular graphs in the
graph generative model’s latent space. The sciML system for this problem is
comprised of the following three components: 1) Generative model: junction
tree variational autoencoder (JT-VAE), 2) Molecular Featurizer: RDKit
cheminformatics software, and 3) Predictive model: message passing neural
network (MPNN) (see Figure 1). Note that the reasoning behind using a
generative model is to transform the original discrete space where molecules
lie to a continuous space that is amenable for search and optimization.
Unfortunately, the presence of RDKit software (which is effectively a black-
box function) renders the sciML system non-differentiable. This precludes the
use of gradient-based methods for end-to-end training. On the other hand,
existing approaches resort to poor workaround of using purely data-driven
models. These approaches pretrain the JT-VAE and MPNN separately using first-
order methods. However, in this approach, the JT-VAE training if not aware of
the property of interest predicted by MPNN. Thus, JT-VAE latent space often is
not amenable to maximizing the property predicted by MPNN. More importantly
purely data-driven approaches require a large amount of data for learning
generalizable patterns.

Figure 1: Scientific machine learning pipeline for molecular property
prediction and generation.

In this paper, on the other hand, we aim to perform an end-to-end training of
sciML sysyem in Figure 1. With a clever application of zeroth-order methods,
we manage to overcome this obstacle and develop gradient-free end-to-end
training algorithms. We show that the proposed zeroth-order training approach
can outperform purely data-driven first-order methods by exploiting the
knowledge encoded in the scientific software RDKit. This is the key
contribution of this paper, and we hope that its application extends beyond
the material science use case considered in this work.

##  2 Problem Formulation

Although our formulation of integrating black-box knowledge with DNN is quite
general, we concretize our discussion with a specific real-world material
science application.

###  2.1 Molecular Property Prediction System

Here we provide details on our molecular property prediction system (Figure 1)
that serves as a fundamental building block for molecular screening and
discovery problems.

Let us define a dataset $D=\\{(\mathbf{x}_{i},\mathbf{y}_{i})\\}_{i=1}^{N}$ of
$N$ molecules, where $\mathbf{x}_{i}$ is the graph $G=(V,E)$ of a single
molecule with vertex features $\\{\mathbf{s}_{u}\\}_{u\in V}$ and edge
features $\\{\mathbf{e}_{uv}\\}_{uv\in E}$; the vertex and edge features
describe certain atom and bond characteristics, respectively. The vector
$\mathbf{y}_{i}$ denotes certain chemical properties, such as density and heat
of formation.

The predictive system of interest consists of the following three distinct
components.

  1. 1.

Generative model (JT-VAE)

We use the Junction Tree Variational Auto-Encoder (JT-VAE) (Jin et al., 2018)
tailored for working on molecular graphs. It consists of: 1) an encoder which
is a neural network $F_{E}(\cdot;\mathbf{u})$ that maps the input $\mathbf{x}$
(molecular graph) to a latent representation $\mathbf{z}$ where $\mathbf{u}$
denotes the encoder parameters, and 2) a decoder which is a neural network
$F_{D}(\cdot;\mathbf{v})$ that maps (reconstructs) the latent representation
$\mathbf{z}$ back to a molecular graph $\tilde{\mathbf{x}}$ where $\mathbf{v}$
denotes the decoder parameters.

  2. 2.

Molecular Featurizer (RDKit)

JT-VAE output is followed by a cheminformatics software that takes a molecule
(in the SMILES string form) as input, computes certain molecular descriptors,
and outputs a molecular graph where these descriptors are then incorporated as
node and edge features. We use the RDKit library (Landrum, 2016) as our
cheminformatics knowledge base. For each problem, any feature that has a
constant value across the whole dataset is removed prior to training. From a
modelling perspective we treat this component as a black box $G(\cdot)$ in
which we can send queries (molecules) $\tilde{\mathbf{x}}$ and get molecular
graphs $\hat{\mathbf{x}}$ with certain edge and node features as outputs.
However, note that the black-box nature of $G$ prevents us from having direct
access to its gradients, and thus the pipeline is effectively rendered non-
differentiable.

  3. 3.

Predictive model (MPNN)

We use the Message Passing Neural Net (MPNN) as our predictive model, which we
denote with $H(\cdot;\mathbf{w})$ where $\mathbf{w}$ denotes the model
parameters. MPNN takes a molecular graph $\hat{\mathbf{x}}$ as an input and
predicts certain molecular properties. Note that MPNN is a graph neural
network (GNN); thus, it offers the advantage of leveraging node/edge features
and connectivity available in the network structure of a molecule.

Using the above notations, the predicted property of molecule $i$ is denoted
as

$\displaystyle\widetilde{\mathbf{y}}_{i}=H\left(G\left(F_{D}\left(F_{E}\left(\mathbf{x}_{i};\mathbf{u}\right);\mathbf{v}\right)\right);\mathbf{w}\right).$
(1)

###  2.2 RDKit Integrated Training of JT-VAE and MPNN

Note that material science applications usually are data-limited, which makes
it difficult for purely data-driven models to learn generalizable patterns. In
this work, we aim to improve the performance of the aforementioned prediction
system by leveraging end-to-end training with RDKit as an additional knowledge
source. The objective of the training procedure is to jointly optimize the
parameters of JT-VAE (i.e., encoder and decoder) and MPNN. This goal is
achieved by solving the following optimization problem:

$\displaystyle\min_{\mathbf{u},\mathbf{v},\mathbf{w}}L(\mathbf{x},\mathbf{y};\mathbf{u},\mathbf{v},\mathbf{w}),$
(2) $\displaystyle
L=L_{mpnn}(\mathbf{x},\mathbf{y};\mathbf{u},\mathbf{v},\mathbf{w})+\lambda\;L_{jtvae}\left(\mathbf{x};\mathbf{u},\mathbf{v}\right),$
$\displaystyle
L_{mpnn}=\frac{1}{N}\sum\limits_{i=1}^{N}\|\widetilde{\mathbf{y}}_{i}-\mathbf{y}_{i}\|^{2},$
$\displaystyle
L_{jtvae}=\ell\left(F_{D}\left(F_{E}\left(\\{\mathbf{x}_{i}\\}_{i=1}^{N};\mathbf{u}\right);\mathbf{v}\right),\\{\mathbf{x}_{i}\\}_{i=1}^{N}\right),$

where $\ell(\cdot,\cdot)$ denotes the JT-VAE loss (for more details refer to
(Jin et al., 2018)), $\lambda$ is a regularization parameter, and
$\widetilde{\mathbf{y}}_{i}$ is defined in (1). Note that $L_{mpnn}$ is a
function of RDKit output, which makes it necessary to use gradient-free
optimization methods.

##  3 Proposed Zeroth-Order Training Algorithm

To solve (2), we adopt a variant of gradient-free optimization, Zeroth-Order
Optimization (ZOO), which operates similarly to iterative methods: gradient
estimation (via function evaluations), descent direction computation, and
solution update. Thus, our ZOO end-to-end training replaces gradients with
their estimates (whose exact form is defined in the next section) and proceeds
as traditional training.

###  3.1 ZOO for the sciML System in Figure 1

Specifically, the gradient of loss $L$ (w.r.t. to the model parameters
$\mathbf{u},\mathbf{v},\mathbf{w}$) are as follows:

$\displaystyle\nabla_{\mathbf{w}}L=\nabla_{\mathbf{w}}L_{mpnn}(\\{(\mathbf{x}_{i},\mathbf{y}_{i})\\}_{i=1}^{N};\mathbf{u},\mathbf{v},\mathbf{w})$
$\displaystyle\widetilde{\nabla}_{\mathbf{u},\mathbf{v}}L=\widetilde{\nabla}_{\mathbf{u},\mathbf{v}}L_{mpnn}(\\{(\mathbf{x}_{i},\mathbf{y}_{i})\\}_{i=1}^{N};\mathbf{u},\mathbf{v},\mathbf{w})$
$\displaystyle\quad\quad\quad+\lambda\nabla_{\mathbf{u},\mathbf{v}}L_{jtvae}\left(\\{\mathbf{x}_{i}\\}_{i=1}^{N};\mathbf{u},\mathbf{v}\right)$
$\displaystyle=\frac{1}{N}\sum\limits_{i=1}^{N}2\left(\widetilde{\mathbf{y}}_{i}-\mathbf{y}_{i}\right)\widetilde{\nabla}_{\mathbf{u},\mathbf{v}}\widetilde{\mathbf{y}}_{i}+\lambda\nabla_{\mathbf{u},\mathbf{v}}L_{jtvae},$

where the symbol $\widetilde{\nabla}$ is used to denote a gradient estimate.

Using the chain rule we have

$\displaystyle\widetilde{\nabla}_{\mathbf{u}}\widetilde{\mathbf{y}}_{i}=\nabla_{\mathbf{u}}F_{E}(\mathbf{x}_{i};\mathbf{u})\nabla_{1}F_{D}\left(F_{E}\left(\mathbf{x}_{i};\mathbf{u}\right);\mathbf{v}\right)\times$
$\displaystyle\widetilde{\nabla}G\left(F_{D}\left(F_{E}\left(\mathbf{x}_{i};\mathbf{u}\right);\mathbf{v}\right)\right)\nabla_{1}H\left(G\left(F_{D}\left(F_{E}\left(\mathbf{x}_{i};\mathbf{u}\right);\mathbf{v}\right)\right);\mathbf{w}\right)$
(3)
$\displaystyle\widetilde{\nabla}_{\mathbf{v}}\widetilde{\mathbf{y}}_{i}=\nabla_{\mathbf{v}}F_{D}\left(F_{E}\left(\mathbf{x}_{i};\mathbf{u}\right);\mathbf{v}\right)\widetilde{\nabla}G\left(F_{D}\left(F_{E}\left(\mathbf{x}_{i};\mathbf{u}\right);\mathbf{v}\right)\right)$
$\displaystyle\quad\quad\quad\times\nabla_{1}H\left(G\left(F_{D}\left(F_{E}\left(\mathbf{x}_{i};\mathbf{u}\right);\mathbf{v}\right)\right);\mathbf{w}\right),$
(4)

where $\nabla_{1}$ denotes the gradient w.r.t. the first argument (i.e., the
model’s input) and $\widetilde{\nabla}G$ some gradient estimate of the black-
box function $G(\cdot)$.

#### Challenges with sciML System in Figure 1.

Upon a closer examination of aforementioned ZOO steps, notice some major
technical hurdles arise due to subtle application level details. Next, we
explain these issues in detail.

  1. 1.

The gradients involved in (3), (4) are not well-defined for the application of
interest. Specifically, notice that

     * •

The term
$\widetilde{\nabla}G\left(F_{D}\left(F_{E}\left(\mathbf{x};\mathbf{u}\right);\mathbf{v}\right)\right)$
denotes the gradient estimate of graph w.r.t. SMILES string.

     * •

The term
$\nabla_{1}H\left(G\left(F_{D}\left(F_{E}\left(\mathbf{x};\mathbf{u}\right);\mathbf{v}\right)\right);\mathbf{w}\right)$
denotes the gradient of a scalar w.r.t. a graph.

  2. 2.

The gradient estimates we aim to form for black-box function $G(\cdot)$ will
involve a quantity of the form $G(\mathbf{x}+\mu\mathbf{d})-G(\mathbf{x})$,
(which typically appears in zeroth-order gradient estimators) where
$\mathbf{x}$ is a molecule and $\mathbf{d}$ is a direction in the molecules
space (the exact representation, e.g., SMILES string or graph, is not relevant
at this point); the output of $G(\cdot)$ is a graph. Then, it is not clear
what the meaning of the above difference is.

Figure 2: We change the boundaries of the three components for our algorithm
design compared to the original one in Fig. 1.

###  3.2 Reparameterization to the Rescue

From the above discussion, it is clear that we need to use reparameterization
in order to resolve these issues. Our approach is to adjust the boundaries of
the three components that comprise the system (see Figure 2). Specifically, we
perform the following changes:

  * •

The boundary between the 1st component and the 2nd component becomes the
intermediate latent layer of the JT-VAE. We maintain the same notation as
previously and denote the input-output relation with
$F_{E}\left(\mathbf{x};\mathbf{u}\right)$.

  * •

The 2nd component includes the layers of the MPNN which performs the embedding
of the graph into a continuous latent space. Specifically, the output of the
2nd component is the readout function of the MPNN. Note that the presence of
the RDKit module in this component renders the pipeline incompatible to first-
order methods. We denote this component with
$\widetilde{G}(\cdot;\mathbf{v},\mathbf{w}_{1})$, where $\mathbf{w}_{1}$
includes the parameters of the MPNN inside this component.

  * •

The third component now includes only the part of the MPNN after the readout
function. We denote the parameters involved in this part with $\mathbf{w}_{2}$
and the input-output relation as $\widetilde{H}(\cdot;\mathbf{w}_{2})$.

The output of the reparameterized system can be written as
$\widetilde{\mathbf{y}}=\widetilde{H}(\widetilde{G}(F_{E}(\mathbf{x}_{i},\mathbf{u});\mathbf{v},\mathbf{w}_{1});\mathbf{w}_{2})$.
As a result, the chain rule is rearranged as:

$\displaystyle\widetilde{\nabla}_{\mathbf{u}}\widetilde{\mathbf{y}}_{i}=\nabla_{\mathbf{u}}F_{E}(\mathbf{x}_{i};\mathbf{u})\widetilde{\nabla}_{1}\widetilde{G}(F_{E}(\mathbf{x}_{i};\mathbf{u});\mathbf{v},\mathbf{w}_{1})$
$\displaystyle\quad\quad\quad\times\nabla_{1}\widetilde{H}(\widetilde{G}(F_{E}(\mathbf{x}_{i};\mathbf{u});\mathbf{v},\mathbf{w}_{1});\mathbf{w}_{2})$
$\displaystyle\widetilde{\nabla}_{\mathbf{v}}\widetilde{\mathbf{y}}_{i}=\widetilde{\nabla}_{\mathbf{v}}\widetilde{G}(F_{E}(\mathbf{x}_{i};\mathbf{u});\mathbf{v},\mathbf{w}_{1})$
$\displaystyle\quad\quad\quad\times\nabla_{1}\widetilde{H}(\widetilde{G}(F_{E}(\mathbf{x}_{i};\mathbf{u});\mathbf{v},\mathbf{w}_{1});\mathbf{w}_{2}).$

Zeroth-order estimators. One aspect we have not addressed so far is the
specific choice for the gradient estimator of the black-box function
$\widetilde{\nabla}_{1}\widetilde{G}(\cdot),\widetilde{\nabla}_{\mathbf{v}}\widetilde{G}(\cdot)$.
In this work, we employ a coordinate-wise zeroth-order (ZO) gradient estimator
(Liu et al., 2020):

$\displaystyle\widetilde{\nabla}_{1}\widetilde{G}(F_{E}(\mathbf{x};\mathbf{u});\mathbf{v},\mathbf{w}_{1})=$
$\displaystyle\frac{\widetilde{G}(F_{E}(\mathbf{x};\mathbf{u})+\mu_{1}\mathbf{d}_{1};\mathbf{v},\mathbf{w}_{1})-\widetilde{G}(F_{E}(\mathbf{x};\mathbf{u});\mathbf{v},\mathbf{w}_{1})}{\mu_{1}}\mathbf{d}_{1}$
$\displaystyle\widetilde{\nabla}_{\mathbf{v}}\widetilde{G}(F_{E}(\mathbf{x};\mathbf{u});\mathbf{v},\mathbf{w}_{1})=$
$\displaystyle\frac{\widetilde{G}(F_{E}(\mathbf{x};\mathbf{u});\mathbf{v}+\mu_{2}\mathbf{d}_{2},\mathbf{w}_{1})-\widetilde{G}(F_{E}(\mathbf{x};\mathbf{u});\mathbf{v},\mathbf{w}_{1})}{\mu_{2}}\mathbf{d}_{2},$

where $\mu_{1},\mu_{2}$ are the smoothing parameters, and
$\mathbf{d}_{1},\mathbf{d}_{2}$ are directions in the (continuous) latent
space; the elements of these direction vectors are Gaussian random variables
of mean $0$ and variance $\sigma^{2}$.

##  4 Experiments

#### Goal:

Our experiments try to find: can the ZO training approach integrate black-box
knowledge source in sciML learning process to learn more generalizable models
despite of data scarcity?

#### Datasets:

We consider the following two datasets:

  1. 1.

ZINC (Sterling & Irwin, 2015): We have $5500$ molecules in the training set
and $5000$ molecules in the test set. Similar to (Kusner et al., 2017; Jin et
al., 2018), the property we predict is a penalized version of the water-
octanol partition coefficients ($\log P$), specifically the quantity $\log
P-SA-RS$, where $SA$ is the synthetic accessibility score and $RS$ is the ring
size.

  2. 2.

HOF (Griessen & Riesterer, 1988) dataset: This dataset contains $5600$
molecules in the training set and $7000$ in the test set. The property of
interest is the heat of formation.

#### Models and Hyperparameters:

For each dataset we perform the following set of experiments:

  1. 1.

Stage1: We train a JT-VAE and a MPNN model independently (or in a purely data-
driven manner) using the DGL-LifeSci package111https://github.com/awslabs/dgl-
lifesci.

  2. 2.

Stage2: For Stage2 models, we inject RDKit knowledge into the purely data-
driven models by using zeroth-order end-to-end training.

Below we provide details on our hyperparameter selection, which remain the
same across both experiments.

  1. 1.

JT-VAE: We use learning rate $=0.0007$, network depth$=3$, size of hidden
layers$=450$ (for both the encoder and the decoder), and dimension of latent
space$=56$.

  2. 2.

MPNN: We use learning rate$=0.0005$, size of node input features$=39$, size of
edge input features$=50$, size of output node representations$=128$, size of
hidden edge representations$=64$, and number of message passing steps$=6$.

Note that the batch size was $16$ for Stage1 experiments, and $8$ for Stage2.
Moreover, in Stage2 experiments we employ gradient clipping in order to deal
with exploding gradients. We also set $\lambda=1$ and freeze the parameters of
the decoder (this is a reasonable choice as the decoder was already pretrained
sufficiently well in Stage1). Finally, zeroth-order gradient estimators uses
$\mu_{1}=\mu_{2}=0.001$.

#### Metrics:

We compare the performance of our approach w.r.t. baselines in terms of
following metrics:

  1. 1.

Reconstruction ability of the generative model.

The molecule is passed through the JT-VAE, where it is first decoded (i.e., we
obtain its continuous representation in the latent space), and then encoded
back into a (graph) molecular form. We consider the reconstruction successful
if the molecule we get in the output (after the decoder) is the same as the
molecule we provide as input (before the encoder). We report the
reconstruction accuracy, which is defined as

$\text{accuracy}=\frac{\\#\text{successfully reconstructed
molecules}}{\text{total number of molecules evaluated}}.$

  2. 2.

Predictive ability of the supervised (predictive) model

We compute the Root Mean Square Error (RMSE) over the dataset, between the
true value of the molecular property and the value predicted by the MPNN,

$\text{RMSE}=\left[\frac{1}{N}\sum\limits_{i=1}^{N}\left(y_{i}-\widetilde{y}_{i}\right)^{2}\right]^{1/2}.$

#### Results:

Table 1 presents the results of the experiments described above. Note that
MPNN is trained in Stage 1 for a maximum number of epochs such that we can no
longer improve its performance (i.e., decrease RMSE). Then, notice that the
jointly training the MPNN (along with the JT-VAE) leads to better prediction
performance, under both datasets. Although we observe a decrease in the
reconstruction error, note that reconstruction is not our main objective. In
fact the main motivation behind using JT-VAE is to obtain a continuous latent
space. Regardless, a successful training by simply using forward passes (i.e.,
without gradients) is exciting and motivates more research in this direction.

Dataset | Model (Measure) | Stage1 | Stage2
---|---|---|---
ZINC | JT-VAE (Recon. accuracy) | $4.59\%$ | $3.97\%$
MPNN (RMSE) with JT-VAE | $1.83$ | 1.66
HOF | JT-VAE (Recon. accuracy) | $22.5\%$ | $15.5\%$
MPNN (RMSE) with JT-VAE | $53$ | 52.2
Table 1: Performance comparison of purely data-driven models (trained with
first-order method) with RDKit integrated models (trained with proposed
zeroth-order method).

We also mention some limitation of the proposed ZOO scheme. First, ZOO
requires large runtime mainly due to application related factors that we
cannot completely control. For instance, the decoding phase of the JT-VAE,
which is involved in the computation of
$\widetilde{G}(F_{E}(\mathbf{x};\mathbf{u})+\mu_{1}\mathbf{d}_{1};\mathbf{v},\mathbf{w}_{1})$
and
$\widetilde{G}(F_{E}(\mathbf{x};\mathbf{u});\mathbf{v}+\mu_{2}\mathbf{d}_{2},\mathbf{w}_{1})$
(in the ZO estimates) can be very slow, especially for novel molecules.
Second, JT-VAE requires pretraining in order to decode molecules reasonably
fast. Finally, training requires carefully balancing prediction error vs.
reconstruction accuracy. Choosing right hyperparameters is critical to a
stable training.

##  5 Conclusion

In this paper, we developed gradient-free training approaches for enabling
integration of black-box scientific knowledge sources with deep learning.
Using two materials datasets, we showed that proposed knowledge-integrated
models outperform purely data-driven models in small-data regime. In the
future we will explore the use of more sophisticated zeroth-order gradient
estimators, e.g., average several gradient estimates across different
directions. Another worthwhile direction is to apply the proposed approach to
additional applications where scientific knowledge is available in the form of
a black-box. Our results provide a new perspective on the integration of
knowledge sources with deep learning. We hope that our positive results will
encourage deep learning community to pay more attention to ZOO-based knowledge
integration.

## Acknowledgement

This work was performed under the auspices of the U.S. Department of Energy by
Lawrence Livermore National Laboratory under Contract DE-AC52-07NA27344. James
Diffenderfer’s effort was supported by the LLNL-LDRD Program under Project No.
22-FS-019.

## References

  * Belbute-Peres et al. (2020) Belbute-Peres, F. D. A., Economon, T., and Kolter, Z.  Combining differentiable pde solvers and graph neural networks for fluid flow prediction.  In _International Conference on Machine Learning_ , pp. 2402–2411. PMLR, 2020.
  * Griessen & Riesterer (1988) Griessen, R. and Riesterer, T.  Heat of formation models.  _Hydrogen in intermetallic compounds I_ , pp. 219–284, 1988.
  * Jin et al. (2018) Jin, W., Barzilay, R., and Jaakkola, T.  Junction tree variational autoencoder for molecular graph generation.  In _International Conference on Machine Learning_ , pp. 2323–2332. PMLR, 2018.
  * Karniadakis et al. (2021) Karniadakis, G. E., Kevrekidis, I. G., Lu, L., Perdikaris, P., Wang, S., and Yang, L.  Physics-informed machine learning.  _Nature Reviews Physics_ , 3(6):422–440, 2021\.
  * Kasim et al. (2022) Kasim, M. F., Lehtola, S., and Vinko, S. M.  Dqc: A python program package for differentiable quantum chemistry.  _The Journal of chemical physics_ , 156(8):084801, 2022.
  * Kusner et al. (2017) Kusner, M. J., Paige, B., and Hernández-Lobato, J. M.  Grammar variational autoencoder.  In _International conference on machine learning_ , pp. 1945–1954. PMLR, 2017.
  * Landrum (2016) Landrum, G.  Rdkit: Open-source cheminformatics software.  2016\.  URL https://github.com/rdkit/rdkit/releases/tag/Release_2016_09_4.
  * Liu et al. (2018) Liu, S., Kailkhura, B., Chen, P.-Y., Ting, P., Chang, S., and Amini, L.  Zeroth-order stochastic variance reduction for nonconvex optimization.  _Advances in Neural Information Processing Systems_ , 31, 2018.
  * Liu et al. (2020) Liu, S., Chen, P.-Y., Kailkhura, B., Zhang, G., Hero III, A. O., and Varshney, P. K.  A primer on zeroth-order optimization in signal processing and machine learning: Principals, recent advances, and applications.  _IEEE Signal Processing Magazine_ , 37(5):43–54, 2020.
  * Sterling & Irwin (2015) Sterling, T. and Irwin, J. J.  Zinc 15 – ligand discovery for everyone.  _Journal of Chemical Information and Modeling_ , 55(11):2324–2337, 2015.  doi: 10.1021/acs.jcim.5b00559.  URL https://doi.org/10.1021/acs.jcim.5b00559.  PMID: 26479676.
  * von Rueden et al. (2019) von Rueden, L., Mayer, S., Beckh, K., Georgiev, B., Giesselbach, S., Heese, R., Kirsch, B., Pfrommer, J., Pick, A., Ramamurthy, R., et al.  Informed machine learning–a taxonomy and survey of integrating knowledge into learning systems.  _arXiv preprint arXiv:1903.12394_ , 2019.