# Additive Multi-Index Gaussian process modeling, with application to multi-
physics surrogate modeling of the quark-gluon plasma

Kevin Li111Department of Statistical Science, Duke University Simon Mak∗ J.-F.
Paquet222Department of Physics and Astronomy & Department of Mathematics,
Vanderbilt University Steffen A. Bass333Department of Physics, Duke University

###### Abstract

The Quark-Gluon Plasma (QGP) is a unique phase of nuclear matter, theorized to
have filled the Universe shortly after the Big Bang. A critical challenge in
studying the QGP is that, to reconcile experimental observables with
theoretical parameters, one requires many simulation runs of a complex physics
model over a high-dimensional parameter space. Each run is computationally
very expensive, requiring thousands of CPU hours, thus limiting physicists to
only several hundred runs. Given limited training data for high-dimensional
prediction, existing surrogate models often yield poor predictions with high
predictive uncertainties, leading to imprecise scientific findings. To address
this, we propose a new Additive Multi-Index Gaussian process (AdMIn-GP) model,
which leverages a flexible additive structure on low-dimensional embeddings of
the parameter space. This is guided by prior scientific knowledge that the QGP
is dominated by multiple distinct physical phenomena (i.e., multi-physics),
each involving a small number of latent parameters. The AdMIn-GP models for
such embedded structures within a flexible Bayesian nonparametric framework,
which facilitates efficient model fitting via a carefully constructed
variational inference approach with inducing points. We show the effectiveness
of the AdMIn-GP via a suite of numerical experiments and our QGP application,
where we demonstrate considerably improved surrogate modeling performance over
existing models.

Keywords: Bayesian nonparametrics, Gaussian processes, high energy physics,
surrogate modeling, uncertainty quantification, variational inference.

##  1 Introduction

The Quark-Gluon Plasma (QGP) is an exotic phase of nuclear matter whose
constituents, namely quarks and gluons, are the elementary building blocks of
protons, neutrons and nuclei. The QGP is theorized to have filled the Universe
shortly after the Big Bang, and the study of this plasma sheds light on the
conditions present in the early Universe. Recent promising work [27, 28, 58,
49] has focused on the use of virtual simulations from complex physics models,
to reconcile physical parameters $\mathbf{x}\in\mathbb{R}^{d}$ with
experimental data from particle colliders. Despite this progress, there
remains a critical bottleneck: such analysis requires many runs from the
expensive simulator $f(\cdot)$ at different parameters $\mathbf{x}$, each
requiring thousands of CPU hours [10]. This is exacerbated by the relatively
high dimension of $\mathbf{x}$, which may be on the order of 20 parameters for
full-scale QGP studies. Given a limited computing budget, one can afford only
hundreds of simulation runs over the high-dimensional parameter space, which
results in imprecise scientific findings [58].

Surrogate modeling [74, 37] provides a promising solution; see Figure 1 for a
visualization. The idea is simple but effective: simulation experiments are
first performed at carefully selected parameter points, then used as training
data to fit an surrogate model that efficiently emulates and quantifies
uncertainty on the simulation response surface $f(\cdot)$. There is a rich and
growing literature on probabilistic surrogate modeling, particularly using
Gaussian processes (GPs; [80, 37]), which offer a flexible Bayesian framework
for efficient prediction and uncertainty quantification (UQ). This includes
the seminal works on GP-based surrogates [72, 19, 87], as well as recent
extensions of such models for complex physics and engineering applications
[50, 14, 75, 81]. However, such models are known to suffer from a curse-of-
dimensionality; when the number of parameters $d$ becomes large, one may
require a sample size $n$ growing exponentially in $d$ to achieve satisfactory
predictive performance [85], which can easily become prohibitively expensive.
Given a tight budget on run size $n$, such models can thus yield poor
predictions with unacceptably high uncertainties in high dimensions.

Figure 1: Visualizing the role of probabilistic surrogate modeling for
studying the quark-gluon plasma using expensive multi-physics simulations
(figure adapted from images from the MADAI collaboration, [26], Brookhaven
National Laboratory and CERN).

One way to address this is to learn plausible low-dimensional structures in
$f$ and integrate such structures for predictive modeling. There has been
recent work in this direction for Bayesian surrogate modeling, particularly
using Gaussian processes. [57] explored the use of variable selection within
GPs for identifying sparsity (see also [4] and [25]). [39] proposed single-
index GP models, which can identify and quantify uncertainty on an active one-
dimensional linear embedding of the parameter space. [79, 77, 83] further
extended these models to identify higher-dimensional linear embeddings within
GPs. Such extensions, however, largely cannot quantify uncertainties in
estimating the underlying embedding, which as we show later may result in
highly overconfident predictions. [86, 56] investigated the use of random
embeddings within GPs. There is also a rich literature on leveraging additive
components of low-dimensional functions within GPs; see [23, 22, 53, 33, 21].

Despite this literature, there are two critical limitations of existing models
for our application. First, these approaches do not directly elicit (and thus
model for) the specific embedded low-dimensional structure guided by the
multi-physics nature of the physical system. Hearkening back to the well-known
Buckingham-$\pi$ theorem [9], it is accepted that complex systems have
structured low-dimensional manifold embeddings that represent a sparse number
of dominant physics. For the QGP, the simulated collision observables are
known to capture several types of dominant physics (e.g., geometry and quantum
fluctuations, relativistic hydrodynamic expansion), each of which depends on a
low-dimensional embedding of the parameter space. This notion of multi-
physics, where the simulator integrates multiple distinct types of physics, is
widely used for simulating complex phenomena, from reacting flows [60] to
particle collisions [55]. However, existing surrogate models do not elicit nor
model for this embedded structure, and thus may yield poor predictions with
high uncertainties when such multi-physics are present (as we show later in
Section 2.3). Second, existing models largely do not account for uncertainties
in estimating the underlying embedded structure (more on this in Section 2.2).
Particularly with a limited sample size $n$ in high dimensions $d$, neglecting
such uncertainties results in a wildly overconfident surrogate model, which
can easily lead to spurious scientific conclusions.

To address these limitations, we propose a novel Additive Multi-Index Gaussian
process (AdMIn-GP) model, which integrates the elicited low-dimensional
embedded structures from multi-physics for surrogate modeling. The AdMIn-GP
features a flexible additive model of GPs with each component active on
different low-dimensional linear embeddings of the parameter space, to model
the presence of multiple dominant physics within the complex simulation
system. We present a novel variational inference approach for efficient and
probabilistic prediction from the AdMIn-GP, leveraging carefully-constructed
shrinkage priors on the embedding matrices. We demonstrate the improved
predictive and uncertainty quantification performance of the AdMIn-GP over
existing models in a suite of numerical experiments and our motivating high-
energy physics application, thus showing that when multi-physics is present,
the integration of such structure can greatly improve surrogate modeling with
limited data. We then highlight how the AdMIn-GP sheds light on the extraction
and interpretation of dominant multi-physics in the quark-gluon plasma, which
can help guide scientific discoveries.

This paper is organized as follows. Section 2 describes the motivating quark-
gluon plasma application, and outlines key limitations of existing surrogate
models for this problem. Section 3 introduces the AdMIn-GP model, and presents
an efficient variational inference algorithm for model training and selection.
Section 5 compares the AdMIn-GP to the state-of-the-art in a suite of
numerical experiments. Section 6 then applies the AdMIn-GP for our application
on the surrogate modeling of the QGP, and discusses its potential for guiding
extraction of dominant multi-physics. Section 7 concludes the paper.

##  2 The quark-gluon plasma

We first outline key challenges underlying the surrogate modeling of the QGP,
then survey existing methods and investigate their limitations for our
application.

###  2.1 Simulating the QGP via heavy-ion collisions

As the theorized form of matter permeating the early Universe a few
microseconds after the Big Bang, the QGP is an important topic of study in
nuclear physics. This plasma is produced and explored in physical experiments
via the collisions of heavy nuclei at velocities close to the speed of light.
Such experiments are conducted at large particle colliders at Brookhaven
National Laboratory and the European Organization for Nuclear Research (CERN).
The temperature and pressure of these collisions converts the colliding nuclei
into a plasma of subatomic particles, namely, quarks and gluons. Due to its
rapid hydrodynamic expansion, the plasma quickly cools down and decays. The
decay products of the quark-gluon-plasma can subsequently be observed by
particle detectors.

To study properties of the QGP, one requires the coupling of data taken by
these particle detectors with complex nuclear physics simulation models, to
reconcile plasma properties (denoted by parameters
$\mathbf{x}\in\mathbb{R}^{d}$) with experimental data. These nuclear collision
simulations are typically performed in several successive stages (following
[27]) to faithfully capture the complex multi-physics phenomena. The first
stage is the impact of the colliding nuclei, and the deposition of energy from
this impact. This is followed by a pre-equilibrium phase, where the system of
quarks and gluons approaches local equilibrium. The third phase simulates the
evolution of plasma utilizing relativistic viscous fluid dynamics, where the
strongly-interacting plasma expands and rapidly cools down. The fourth stage
models particlization – the transition from the strongly-coupled quark-gluon
plasma to individual bound states of quarks and gluons (called “hadrons”),
whose interaction strength is insufficient to maintain a cohesive hydrodynamic
expansion. The final stage simulates the interaction and subsequent decay of
the hadronic bound states, until their density is so low that interactions
cease and the particles travel to the particle detectors. The behavior of this
collision system is controlled by parameters $\mathbf{x}$, which characterize
the energy deposition at the initial nuclear impact, its evolution towards
local equilibrium, the transport coefficients of the QGP (e.g., shear and bulk
viscosity), and the recombination of quarks and gluons into normal nuclear
bound states; details can be found in [27]. This complex simulator thus
captures the interplay between multiple types of physics, including finite
density quantum field theory, relativistic viscous hydrodynamics, relativistic
transport, and electromagnetic interactions.

Given parameters $\mathbf{x}$, we obtain a virtual simulation of an observable
$f(\mathbf{x})$ that one might observe from the decay of the plasma. The goal
is then to find parameters that best match the experimental observables from
particle colliders. In recent years, Bayesian inference with such simulators
have led to a deeper understanding of the properties of ultradense nuclear
matter [28, 27]. The computational cost of a full-scale study of nuclear
collisions is considerable, however, requiring thousands of CPU hours per
parameter. To address this critical bottleneck, surrogate models are
increasingly used to efficiently emulate the expensive simulator [68, 64, 5,
27, 58]; see Figure 1 for a visualization. Of course, its success hinges on
the ability to fit an accurate surrogate model with low predictive
uncertainty. This is a highly challenging task for the QGP: the parameter
space for $\mathbf{x}$ is quite high-dimensional ($d=17$), and many training
design points are thus needed to sufficiently populate this high-dimensional
space. However, each design point necessitates an expensive run of the
simulator $f(\cdot)$, which requires thousands of CPU hours. Given the
computational budget for our project, this limits us to $n\approx 500$ design
points. With such limited data, the careful construction of the surrogate
model is paramount for cost-efficient scientific discovery.

###  2.2 Existing GP surrogate models

Much of the literature on probabilistic surrogate modeling involves Gaussian
process models. We first provide below a brief review of GP modeling (see [74]
for further details), then discuss existing GP models that may apply for our
problem.

Let $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$ denote the black-box function
representing the observable simulated from the expensive simulator model.
Here, the simulated plasma observables are known to be noisy, with noise
corruption well-modeled via homoskedastic Gaussian noise (see, e.g., [27]).
Given simulated parameters $\mathbf{x}_{1},\cdots,\mathbf{x}_{n}$, we assume
the simulated observables follow:

$y_{i}=f(\mathbf{x}_{i})+\epsilon_{i},\quad\epsilon_{i}\mathbin{\overset{i.i.d.}{\kern
0.0pt\sim}}\mathcal{N}(0,\sigma^{2}),\quad i=1,\cdots,n.$ (1)

We then adopt a GP prior on the unknown response surface $f$, given as:

$f(\cdot)\sim\text{GP}\\{\mu,k(\cdot,\cdot)\\}.$ (2)

Here, $\mu$ is a mean parameter for the stochastic process, and
$k(\cdot,\cdot)$ is its covariance function that dictates smoothness of sample
paths. Given a lack of prior knowledge, $k(\cdot,\cdot)$ can be taken as the
squared-exponential or Matérn kernel [74]. Conditional on simulated data
$\mathbf{y}=(y_{1},\cdots,y_{n})$, the posterior predictive distribution of
$f$ at a new parameter $\mathbf{x}^{*}$ is:

$\displaystyle\small\begin{split}f(\mathbf{x}^{*})&|y_{1},\cdots,y_{n}\sim\mathcal{N}\\{\mu_{n}(\mathbf{x}^{*}),\sigma^{2}_{n}(\mathbf{x}^{*})\\},\\\
\mu_{n}(\mathbf{x}^{*})&=\mu+\mathbf{k}_{*,n}^{T}(\mathbf{K}_{n,n}+\beta^{-1}\mathbf{I}_{n\times
n})^{-1}(\mathbf{y}-\boldsymbol{\mu}),\;\sigma^{2}_{n}(\mathbf{x}^{*})=k(\mathbf{x}^{*},\mathbf{x}^{*})-\mathbf{k}_{*,n}^{T}(\mathbf{K}_{n,n}+\beta^{-1}\mathbf{I}_{n\times
n})^{-1}\mathbf{k}_{n,*}.\end{split}$ (3)

Here, $\mathbf{K}_{n,n}=[k(\mathbf{x}_{i},\mathbf{x}_{j})]_{i,j=1}^{n}$ is the
covariance matrix for the design points, and
$\mathbf{k}_{*,n}=[k(\mathbf{x}^{*},\mathbf{x}_{i})]_{i=1}^{n}$ is the
covariance vector between the design points and the new parameter. Equation
(3) highlights a key advantage of GP surrogates: it provides closed-form
expressions for prediction (emulation) and uncertainty quantification, which
can greatly speed-up downstream uses of the surrogate for optimization [52,
15] and design [7].

One limitation of standard GP surrogates, as mentioned earlier, is that it
suffers from a curse-of-dimensionality: as the number of parameters $d$ grows
large, one may require a sample size $n$ growing exponentially in $d$ to
achieve satisfactory predictions (see, e.g., minimax and posterior contraction
rates for GP models in [84, 85]). For the current application with $d=17$
parameters and highly expensive simulation runs, such a sample size is
unachievable with any reasonable budget! Given limited runs, one must then be
judicious in identifying low-dimensional structure for improving predictions.
We review existing literature on this below.

An important early work in this direction is [39], which proposed a fully
Bayesian single-index GP model (SIM-GP); this is further extended in [48] for
quantile regression. The key idea is in representing $f$ as
$f(\mathbf{x})=g(\mathbf{m}^{T}\mathbf{x})$, where
$c=\mathbf{m}^{T}\mathbf{x}$ represents a single latent variable that accounts
for variation over the response surface, and $g$ follows the aforementioned GP
model. The loading vector $\mathbf{m}$ is then inferred in a fully Bayesian
fashion via Markov chain Monte Carlo, using appropriate priors on $\mathbf{m}$
and GP length-scale parameters. However, as noted in [79], the SIM-GP and its
variants encounters difficulties when the embedded structures in $f$ are more
complex and have dimensionality greater than one; we shall see this later in
our application.

There has been much recent work on exploring a more flexible low-dimensional
embedding via active subspaces. Here, $f$ is modeled as the form
$f(\mathbf{x})=g(\mathbf{M}\mathbf{x})$, where
$\mathbf{M}\in\mathbb{R}^{p\times d}$, $p\ll d$, is a matrix that maps the
original parameters $\mathbf{x}$ onto an active lower-dimensional latent space
– the “active subspace”. Functions of this form are known as ridge functions
[69], and can be shown to naturally arise from simple physical laws (see,
e.g., the well-known Buckingham-$\pi$ theorem [9]), where $f$ is largely
controlled by one dominant physical process that depends on several key (but
latent) variables $\mathbf{c}=\mathbf{M}\mathbf{x}$. The challenge lies in
jointly estimating and quantifying uncertainty on the embedding matrix
$\mathbf{M}$ and the underlying function $g$ on the active subspace. Much of
the existing literature considers polynomial models on $g$ (see, e.g., [18,
34, 76, 40]), which facilitates efficient point estimation of $\mathbf{M}$ via
quadratic programming. Recent works, such as [77, 83], explore more flexible
forms of $g$ via GPs, with the goal of estimating the underlying active
subspace. These are also known as multi-index models in the literature (see
[41, 88]); however, existing work on such models are again largely restricted
to parametric models [44]. Another key limitation with the above existing
models is that they do not quantify uncertainty on the embedding matrix
$\mathbf{M}$ within predictive modeling. As we shall see next, this may result
in wildly overconfident surrogates with poor uncertainty quantification (see
[89] for similar observations).

Finally, there has been recent work (see, e.g., [33, 21, 13]) on identifying
active low-dimensional structure via the integration of projection pursuit
within a GP. Such models take the form
$f(\mathbf{x})=\sum_{l=1}^{L}g_{l}(\mathbf{m}_{l}^{T}\mathbf{x})$, i.e., $f$
can be represented as the sum of separate single-index models. Here, $g_{l}$
follows a GP, and the projection vectors $\mathbf{m}_{l}$ can either be
optimized or randomly sampled. While this provides a flexible extension of the
single-index GP, it suffers from a similar limitation as earlier models for
our application: it does not directly elicit (and thus does not model for) the
embedding structures dictated by the multi-physics present in the simulation
system. With limited data, this misspecification of the underlying embedded
structure may result in worse predictive performance compared to standard GPs,
as we shall see next.

###  2.3 Limitations for the QGP application

To motivate the AdMIn-GP, we investigate the above existing surrogates for our
application. We first simulate the training design points using Latin
hypercube designs [74] of sizes $n$ from the expensive physics model, which
has $d=17$ parameters. We then compare several existing models: (i) the
standard GP model using the squared-exponential kernel with automatic
relevance determination (ARD-GP), (ii) the fully Bayesian single-index GP
model (SIM-GP; [39]), (iii) the active subspace (or dimension reduced) GP
model (DR-GP; [77, 83, 79, 83]), which performs dimension reduction using
point estimates of the active subspaces, and (iv) the projection-pursuit-based
diverse projected additive GP (DPA-GP; [21]). Model hyperparameters are fitted
via maximization of the marginal likelihood, as recommended in [70, 82, 83].
The fitted models are then compared on its root-mean-squared-error (RMSE), the
continuous ranked probability score (CRPS; [35]), and the empirical coverage
of its 95% CIs over a testing set. The first measures the point prediction
accuracy, the second for probabilistic predictions, and the last measures
coverage performance. Further details on this set-up are provided in Section
6.

Figure 2: Boxplots of different predictive metrics (left: RMSE; middle: CRPS;
right: empirical coverage rate) for existing models and the proposed AdMIn-GP.

Figure 2 shows the predictive metrics for each model with varying sample sizes
$n=200-500$ for our QGP application. We see that, aside from the SIM-GP,
existing approaches that attempt to fit embedded low-dimensional structure
surprisingly perform worse than the standard GP, in terms of both predictions
and coverage! The SIM-GP yields comparable predictions to the standard GP (in
terms of RMSE and CRPS), but provides significantly worse uncertainty
quantification as its coverage rate is noticeably lower than the desired 95%.
There are two plausible reasons for this. First and foremost, when the
embedded low-dimensional structure (which we know to be present from prior
knowledge of multi-physics) is misspecified in the surrogate model, the fitted
model may in fact yield worse predictions and coverage over standard GPs that
do not leverage such structure. This highlights the need for a careful
elicitation of the embedded structure from prior scientific knowledge for
surrogate modeling. Second, the poor coverage of existing methods is not too
surprising, since such methods largely do not account for uncertainties in the
estimation of the embedding matrix $\mathbf{M}$. To foreshadow, Figure 2 shows
the performance of the proposed AdMIn-GP, which by modeling for the desired
embedding structure guided by the multi-physics of the QGP, appears to yield
significantly improved predictions with desired probabilistic coverage.

##  3 The AdMIn-GP model

We now present the proposed AdMIn-GP and justify how its modeled low-
dimensional structure can capture multi-physics in the simulator. We then
propose a variational inference approach using inducing points for efficient
posterior inference and prediction. Figure 3 visualizes our modeling framework
in plate diagram form; we elaborate on each part below.

###  3.1 Model specification

Figure 3: Plate diagram visualizing the inducing-points formulation for the
AdMIn-GP. White nodes represent latent (unobserved) variables, the shaded node
represents observed data, black points represent input variables or model
hyperparameters, and arrows indicate model dependencies.

Let us assume the same noisy observable model in (1), with
$y_{1},\cdots,y_{n}$ the noisy simulated outputs with i.i.d. noise
$\mathcal{N}(0,\beta^{-1})$ at input parameters
$\mathbf{x}_{1},\cdots,\mathbf{x}_{n}$. To capture the desired multi-physics
within $f$, the AdMIn-GP adopts the following model:

$f(\mathbf{x}_{i})=\sum_{l=1}^{L}f_{l}(\mathbf{x}_{i}):=\sum_{l=1}^{L}g_{l}(\mathbf{M}_{l}\mathbf{x}_{i}),\quad
i=1,\cdots,n.$ (4)

Here, for each term $l$, the embedding matrix
$\mathbf{M}_{l}\in\mathbb{R}^{p\times d}$ maps the original $d$ inputs in
$\mathbf{x}_{i}$ onto a lower $p$-dimensional subspace. This additive multi-
index form (4) is guided by the multi-physics structure present in the
simulator $f$, with each of the $L$ terms, namely
$g_{l}(\mathbf{M}_{l}\mathbf{x}_{i})$, modeling for the effect of distinct
dominant physics in the simulator. As mentioned in Section (2), each physics
is likely active on a low-dimensional (but latent) embedding of the $d$
parameters; see, e.g., the Buckingham-$\pi$ theorem [9]. This is reflected in
the fact that $g_{l}(\cdot)$ is active only on the lower $p$-dimensional
latent variables $\mathbf{c}_{l}:=\mathbf{M}_{l}\mathbf{x}$.

Of course, neither $g_{l}(\cdot)$ nor $\mathbf{M}_{l}$ are known in practice.
For the functions $\\{g_{l}(\cdot)\\}_{l=1}^{L}$, we then assign independent
zero-mean Gaussian process priors with kernel $k_{l}$, i.e.,

$g_{l}\mathbin{\overset{indep.}{\kern
0.0pt\sim}}\text{GP}\\{0,k_{l}(\cdot,\cdot)\\},\quad l=1,\cdots,L.$ (5)

This provides a flexible Bayesian nonparametric form for learning the effect
of the latent variables on the response of interest. For the embedding
matrices $\\{\mathbf{M}_{l}\\}_{l=1}^{L}$, we employ independent element-wise
priors of the form:

$(\mathbf{M}_{l})_{jk}\mathbin{\overset{indep.}{\kern
0.0pt\sim}}\textup{DExp}(0,\lambda_{l}),\quad
j=1,\cdots,p,\;k=1,\cdots,d,\;l=1,\cdots,L,$ (6)

where $\textup{DExp}(0,\lambda_{l})$ is the double exponential distribution
with mean 0 and scale parameter $\lambda_{l}$. Such priors have been widely
used for shrinkage in high-dimensional Bayesian linear regression [6]. They
are used here for two reasons. First, it is known from prior physical
knowledge that, for each dominant physics $l$, the active latent parameters
$\mathbf{c}_{l}:=\mathbf{M}_{l}\mathbf{x}$ depend on a sparse number of the
original $d$ parameters. For the QGP, such sparsity is expected from the
multi-stage nature of high-energy nuclear collisions, where groups of
parameters act at specific stages; it has also been observed experimentally in
prior QGP studies [73, 27]. Second, with limited $n$, some shrinkage of the
large number of matrix parameters can prevent model overfitting. Here, double
exponential priors are preferred over standard Gaussian priors since sub-
exponential tails are known to overshrink parameters empirically [71, 12]. We
will discuss the estimation of $\\{\lambda_{l}\\}_{l=1}^{L}$ and other
hyperparameters later.

For the kernels $\\{k_{l}\\}_{l=1}^{L}$, we found that the isotropic squared-
exponential form
$k_{l}(\mathbf{c},\mathbf{c}^{\prime})=\nu_{l}\exp\\{-(2\theta_{l})^{-1}\|\mathbf{c}-\mathbf{c}^{\prime}\|^{2}_{2}\\}$
works quite well in implementation, where different variance and length-scale
parameters are used for each of the $l$-th component. While such a kernel is
isotropic on the latent embedded space, the induced kernel on the original $d$
parameters is anisotropic; in the simple case of $L=1$ component,
$\text{Cov}\\{f(\mathbf{x},\mathbf{x}^{\prime})\\}=\nu_{l}\exp\\{-(2\theta_{l})^{-1}(\mathbf{x}-\mathbf{x}^{\prime})^{T}\mathbf{M}_{l}^{T}\mathbf{M}_{l}(\mathbf{x}-\mathbf{x}^{\prime})\\}$,
which is clearly anisotropic. With careful probabilistic estimation of the
embedding matrices $\mathbf{M}_{l}$ (see next), the AdMIn-GP thus facilitates
the identification of important combinations of parameters in $\mathbf{x}$ for
predictive modeling. The choice of an isotropic kernel in this embedded space
also has the key advantage of speeding up hyperparameter optimization and
providing numerical stability; more on this later. Our framework can also
accommodate other choices of isotropic kernels (e.g., the isotropic Matérn
[80]), with slight modifications on the following variational inference
procedure.

###  3.2 Inducing points and variational inference

While the above model is guided by the desired multi-physics embedding
structure, the many model parameters to estimate (e.g., parameters in
$\mathbf{M}_{l}$) can introduce significant computational complexities. To
address this, we extend the inducing points approach in [78, 82], which is
widely used for fitting “deep” variants of GPs [20], by facilitating tractable
marginalization in computing the posterior predictive distribution. Similar
approaches have been employed in the spatial statistics literature (see, e.g.,
the Gaussian predictive process models in [3]) for scaling up GP fitting for
massive datasets.

For the AdMIn-GP, the key idea is to make use of a set of $m$ inducing points
(or “pseudo-inputs”)
$\mathbf{z}_{1,l},\cdots,\mathbf{z}_{m,l}\in\mathbb{R}^{d}$ for each additive
component $l=1,\cdots,L$, with corresponding latent “pseudo-observations”
$\mathbf{u}_{l}=(f_{l}(\mathbf{z}_{1,l}),\cdots,f_{l}(\mathbf{z}_{m,l}))$, to
speed up computation. These points are not necessarily a subset of the
original data points $\mathbf{x}_{1},\cdots,\mathbf{x}_{n}$, nor its responses
a subset of the original responses $\mathbf{y}$. Let
$\mathbf{f}_{l}:=(f_{l}(\mathbf{x}_{1}),\cdots,f_{L}(\mathbf{x}_{n}))$, and
define the kernel matrices
$\mathbf{K}_{n,m}^{l}={[k_{l}(\mathbf{M}_{l}\mathbf{x}_{i},\mathbf{M}_{l}\mathbf{z}_{j,l})]_{i=1}^{n}}_{j=1}^{m}$,
$\mathbf{K}_{m,m}^{l}={[k_{l}(\mathbf{M}_{l}\mathbf{z}_{j,l},\mathbf{M}_{l}\mathbf{z}_{j^{\prime},l})]_{j=1}^{m}}_{j^{\prime}=1}^{m}$,
and
$\mathbf{K}_{n,n}^{l}={[k_{l}(\mathbf{M}_{l}\mathbf{x}_{i},\mathbf{M}_{l}\mathbf{x}_{i^{\prime}})]_{i=1}^{n}}_{i^{\prime}=1}^{n}$.
We can then write the generative distribution on the data $\mathbf{y}$ as:

$\displaystyle\begin{split}&\mathbf{y}|\mathbf{f}_{1},\dots,\mathbf{f}_{L}\sim\mathcal{N}\left(\sum_{l=1}^{L}\mathbf{f}_{l},\beta^{-1}\mathbf{I}_{n\times
n}\right),\quad l=1,\cdots,L,\\\
&\mathbf{f}_{l}|\mathbf{u}_{l},\mathbf{M}_{l}\sim\mathcal{N}\\{\mathbf{K}_{n,m}^{l}(\mathbf{K}_{m,m}^{l})^{-1}\mathbf{u}_{l},\mathbf{K}_{n,n}+\mathbf{K}_{n,m}^{l}(\mathbf{K}_{m,m}^{l})^{-1}\mathbf{K}_{m,n}^{l}\\},\\\
&\mathbf{u}_{l}|\mathbf{M}_{l}\sim\mathcal{N}(0,\mathbf{K}_{m,m}^{l}),\end{split}$
(7)

where each entry of $\mathbf{M}_{l}$ again follows the double-exponential
shrinkage priors in (6). Figure 3 shows a plate representation of this
inducing-points model formulation. The pseudo-observations
$\mathbf{u}=(\mathbf{u}_{1};\mathbf{u}_{2};\dots;\mathbf{u}_{L})$ are latent
variables that will be marginalized out later for posterior prediction. We
present next a variational inference (VI) approach that leverages this
inducing-points formulation for efficient (approximate) Bayesian predictions.

From (7), we wish to sample from the joint posterior distribution of the
function components $\mathbf{f}=(\mathbf{f}_{1};\cdots;\mathbf{f}_{L})^{T}$,
the pseudo-observations $\mathbf{u}$, and the embedding matrices
$\mathbf{M}=\\{\mathbf{M}_{1},\cdots,\mathbf{M}_{L}\\}$. We leverage a similar
variational bound as in [82], which was originally proposed for variational
inference of Gaussian process latent variable models. The desired posterior is
first approximated by the variational form:

$\displaystyle
p(\mathbf{f},\mathbf{u},\mathbf{M}|\mathbf{y})\approx\phi(\mathbf{u})\prod_{l=1}^{L}p(\mathbf{f}_{l}|\mathbf{u}_{l},\mathbf{M}_{l})\phi_{l}(\mathbf{M}_{l}),$
(VI)

where $p(\mathbf{f}_{l}|\mathbf{u}_{l},\mathbf{M}_{l})$ is the multivariate
normal density in (7). Here, $\phi(\mathbf{u})$ and $\phi_{l}(\mathbf{M}_{l})$
are variational distributions444We use the notation $\phi(\theta)$ to denote
the variational distribution of a parameter $\theta$, to distinguish it from
its generative distribution $p(\theta)$ in (7). that will be selected for a
computationally tractable variational bound for the log-marginal likelihood of
our observations, $\log p(\mathbf{y})$. The variational form (VI) and
resulting variational bound mimic those in [82] for the Bayesian GP latent
variable modeling. This form facilitates the analytical marginalization
desired for optimizing the evidence lower bound (see next), while also
providing flexibility for modeling dependencies between different additive
components, embedding matrices and pseudo-outputs.

Variational inference then proceeds via the maximization of a lower bound on
the marginal likelihood $\log p(\mathbf{y})$ via the variational approximation
(VI). Standard VI methods for mean-field and structured approximations (see,
e.g., [46]) are difficult to apply here due to the unstructured coupling of
variables in the multivariate normal likelihood
$p(\mathbf{f}_{1},\dots,\mathbf{f}_{L}|\mathbf{u},\mathbf{M}_{1},\cdots,\mathbf{M}_{L})$,
hence more established VI approaches (see [8]) cannot be used. We thus need to
explicitly derive below an analytic lower bound on the marginal likelihood for
maximizing variational parameters in $\phi$ and GP model parameters.

With (VI), the log-marginal likelihood can be bounded via the evidence lower
bound (ELBO; [8]):

$\displaystyle\begin{split}\log
p(\mathbf{y})&\geq\int\phi(\mathbf{u})\left(\prod_{l=1}^{L}p(\mathbf{f}_{l}|\mathbf{u}_{l},\mathbf{M}_{l})\phi_{l}(\mathbf{M}_{l})\right)\times\\\
&\quad\quad\quad\quad\log\left\\{\frac{p(\mathbf{y}|\mathbf{f}_{1}\dots\mathbf{f}_{L})\prod_{l=1}^{L}p(\mathbf{f}_{l}|\mathbf{u}_{l},\mathbf{M}_{l})p(\mathbf{u}_{l})p(\mathbf{M}_{l})}{\phi(\mathbf{u})\prod_{l=1}^{L}p(\mathbf{f}_{l}|\mathbf{u}_{l},\mathbf{M}_{l})\phi(\mathbf{M}_{l})}\right\\}d\mathbf{f}d\mathbf{M}d\mathbf{u}.\end{split}$
(8)

We can then marginalize out the latent function components $\mathbf{f}$,
yielding:

$\displaystyle\begin{split}&\int\phi(\mathbf{u})\left(\mathbb{E}_{\phi(\mathbf{M})}\left[\log\mathcal{N}\left(\mathbf{y};\sum_{l=1}^{L}\boldsymbol{\alpha}_{l},\beta^{-1}\mathbf{I}_{n\times
n}\right)\right]+\log\frac{p(\mathbf{u})}{\phi(\mathbf{u})}\right)d\mathbf{u}\\\
&\quad\quad-\sum_{l=1}^{L}\frac{\beta}{2}\left(\nu_{l}n-\text{tr}\left[(\mathbf{K}_{m,m}^{l})^{-1}\mathbb{E}_{\phi(\mathbf{M})}(\mathbf{K}_{m,n}^{l}\mathbf{K}_{n,m}^{l})\right]\right)-\sum_{l=1}^{L}\text{KL}\\{\phi(\mathbf{M}_{l})||p(\mathbf{M}_{l})\\},\end{split}$
(9)

where
$\boldsymbol{\alpha}_{l}=\mathbf{K}_{n,m}^{l}(\mathbf{K}_{m,m}^{l})^{-1}\mathbf{u}_{l}$,
$p(\cdot)$ denotes the generative distribution of a parameter from (7), and KL
denotes the Kullback-Liebler divergence. Examining the first integral, note
that if the variational distribution on $\mathbf{u}$ is chosen as:

$\phi(\mathbf{u})\propto\exp\left\\{\mathbb{E}_{\phi(\mathbf{M})}\left(\log\mathcal{N}\left(\mathbf{y};\sum_{l=1}^{L}\boldsymbol{\alpha}_{l},\beta^{-1}\mathbf{I}_{n\times
n}\right)\right)\right\\}p(\mathbf{u})$ (10)

then the bound in (9) reduces to:

$\displaystyle\small\begin{split}\log p(\mathbf{y})&\geq\int
p(\mathbf{u})\mathbb{E}_{\phi(\mathbf{M})}\left[\log\mathcal{N}\left(\mathbf{y};\sum_{l=1}^{L}\boldsymbol{\alpha}_{l},\beta^{-1}\mathbf{I}_{n\times
n}\right)\right]d\mathbf{u}\\\
&\quad-\sum_{l=1}^{L}\frac{\beta}{2}\left(\nu_{l}n-\text{tr}\left[(\mathbf{K}_{m,m}^{l})^{-1}\mathbb{E}_{\phi(\mathbf{M})}(\mathbf{K}_{m,n}^{l}\mathbf{K}_{n,m}^{l})\right]\right)-\sum_{l=1}^{L}\text{KL}\\{\phi(\mathbf{M}_{l})||p(\mathbf{M}_{l})\\}.\end{split}$
(11)

This lower bound can now be computed via simple integration of the pseudo-
observations $\mathbf{u}$ and the expected matrices
$\boldsymbol{\Psi}_{1}^{l}=\mathbb{E}_{\phi(\mathbf{M})}(\mathbf{K}_{n,m}^{l}),\boldsymbol{\Psi}_{2}^{l}=\mathbb{E}_{\phi(\mathbf{M})}(\mathbf{K}_{m,n}^{l}\mathbf{K}_{n,m}^{l})$.
We first choose the variational distribution on $\mathbf{M}$ as the matrix
normal distribution [16], i.e.,

$\phi_{l}(\mathbf{M}_{l})\sim\mathcal{MN}(\mathbf{A}_{l},\mathbf{H}_{l},\mathbf{U}_{l}),\quad
l=1,\cdots,L,$ (12)

where $\mathbf{A}_{l}$, $\mathbf{H}_{l}$ and $\mathbf{U}_{l}$ are the mean,
row-wise covariance and column-wise covariance matrices. With this, and
setting $k_{l}$ as the isotropic squared-exponential kernel, each element of
$\boldsymbol{\Psi}_{1}^{l}$ and $\boldsymbol{\Psi}_{2}^{l}$ can then be
computed in closed form. The derivations are, however, quite tedious and are
deferred to Appendix A. Integrating these closed-form expressions for
$\boldsymbol{\Psi}_{1}^{l}$ and $\boldsymbol{\Psi}_{2}^{l}$ into (11), we
arrive at the final variational bound for the log-marginal likelihood:

$\displaystyle\begin{split}\log
p(\mathbf{y})&\geq\log\left(\frac{\beta^{\frac{n}{2}}\prod_{l=1}^{L}\det(\mathbf{K}_{m,m}^{l})^{\frac{1}{2}}}{(2\pi)^{\frac{n}{2}}\det(\beta\mathbf{D}+\mathbf{P})^{\frac{1}{2}}}\exp\left\\{\frac{\beta^{2}}{2}\bar{\mathbf{y}}^{T}\mathbf{W}^{-1}\bar{\mathbf{y}}\right\\}\exp\left\\{-\frac{1}{2}\beta\mathbf{y}^{T}\mathbf{y}\right\\}\right)\\\
&\quad\quad-\sum_{l=1}^{L}\text{KL}\\{\phi_{l}(\mathbf{M}_{l})\\}||p(\mathbf{M}_{l}))-\sum_{l=1}^{L}\text{tr}(\mathbf{V}_{l}),\end{split}$
(13)

where
$\mathbf{V}_{l}=({\beta}/{2})(\nu_{l}n-\text{tr}\\{(\mathbf{K}_{m,m}^{l})^{-1}\boldsymbol{\Psi}_{2}^{l}\\})$,
$\bar{\mathbf{y}}=((\boldsymbol{\Psi}^{1}_{1})^{T}\mathbf{y},(\boldsymbol{\Psi}^{2}_{1})^{T}\mathbf{y}\dots(\boldsymbol{\Psi}^{L}_{1})^{T}\mathbf{y})^{T}$,
and $\mathbf{W}=\beta\mathbf{P}+\mathbf{D}$. Here,
$\mathbf{D}\in\mathbb{R}^{mL\times mL}$ is block diagonal with
$\mathbf{D}_{l,l}=\mathbf{K}_{m,m}^{l}$, and
$\mathbf{P}\in\mathbb{R}^{mL\times mL}$ is a block matrix with block diagonals
$\mathbf{P}_{l,l}=\boldsymbol{\Psi}^{l}_{2}$ and block off-diagonals
$\mathbf{P}_{s,t}=(\boldsymbol{\Psi}^{s}_{1})^{T}(\boldsymbol{\Psi}^{t}_{1})$.
Detailed derivation of this bound can be found in Appendix B

Finally, with (13), we then optimize this bound by maximizing with respect to
the variational parameters
$\Theta_{\phi}:=\\{(\mathbf{A}_{l},\mathbf{H}_{l},\mathbf{U}_{l})\\}_{l=1}^{L}$,
pseudo-inputs $\mathcal{Z}:={\\{\mathbf{z}_{j,l}\\}_{j=1}^{M}}_{l=1}^{L}$, and
AdMIn-GP model parameters $\Theta_{\rm GP}$, namely, its scale parameter
$\nu_{l}$, length-scale parameter $\theta_{l}$ and shrinkage parameter
$\lambda_{l}$ for each additive component $l=1,\cdots,L$. Details on this
optimization are provided in Section 4.1. The variational distribution (VI)
with these optimized parameters then serves as an approximation for the
complex posterior $p(\cdot|\mathbf{y})$, which facilitates efficient posterior
predictions (see Section 3.3). Since the derived bound (13) lower bounds the
log-marginal likelihood, the maximization of this bound with respect to model
parameters can be viewed as an empirical Bayes estimation [11] of such
parameters.

One bottleneck that arises is that, as each evaluation of the variational
bound (13) requires $\mathcal{O}(L^{3}(nm^{2}d^{3}p^{3}+m^{3}))$ operations,
the optimization over the full matrix normal variational family (12) can be
computationally quite costly. One solution is to restrict the variational
family (12) to the simpler form
$\phi_{l}(\mathbf{M}_{l})\sim\mathcal{MN}(\mathbf{A}_{l},\mathbf{I}_{p\times
p},\mathbf{I}_{d\times d}\mathbf{v}_{l})$, so that only the mean matrices
$\mathbf{A}_{l}$ and column-wise variances $\mathbf{v}_{l}$ need to be
optimized for variational inference. In addition to reducing the number of
variational parameters for optimization, this restriction facilitates the use
of rank-one matrix updates to achieve improved computational complexity per
evaluation of the bound (13), namely $\mathcal{O}(L^{3}(nm^{2}pd+m^{3}))$.
Technical details of this simplification are provided in Appendix A. In our
later experiments, this simplified variational form does not appear to
compromise much uncertainty quantification performance, particularly in
comparison to existing models.

###  3.3 Posterior predictions

Finally, we can make use of the variational form (VI) (with variational
parameters optimized via the ELBO) to approximate the desired posterior
predictive distribution $[f(\mathbf{x}_{\rm new})|\mathbf{y}]$, where
$\mathbf{x}_{\rm new}$ is a new input point. Note that our choice of
variational distributions in (10) and (12) induces the marginal distribution
on pseudo-observations:

$\phi(\mathbf{u})\sim\mathcal{N}(\mathbf{D}\mathbf{W}^{-1}\bar{\mathbf{y}},\mathbf{D}\mathbf{W}^{-1}\mathbf{D}).$
(14)

Integrating this within (VI), one can then show that $[f(\mathbf{x}_{\rm
new})|\mathbf{y},\mathbf{M}]$, the posterior predictive distribution
conditional on the embedding matrices $\mathbf{M}$, can be approximated as:

$\mathcal{N}\Big{(}\mathbf{K}_{{\mathbf{x}_{\rm
new}},\mathbf{M}}\mathbf{D}^{-1}\boldsymbol{\mu},\mathbf{K}_{\mathbf{x}_{\rm
new},\mathbf{x}_{\rm new}}-\mathbf{K}_{\mathbf{x}_{\rm
new},\mathbf{M}}\mathbf{D}^{-1}\mathbf{K}_{\mathbf{M},\mathbf{x}_{new}}+\mathbf{K}_{\mathbf{x}_{\rm
new},\mathbf{M}}\mathbf{D}^{-1}\mathbf{G}\mathbf{D}^{-1}\mathbf{K}_{\mathbf{M},\mathbf{x}_{\rm
new}}\Big{)}.$ (15)

Here, $\boldsymbol{\mu}=\mathbf{D}\mathbf{W}^{-1}\bar{\mathbf{y}}$,
$\mathbf{G}=\mathbf{D}\mathbf{W}^{-1}\mathbf{D}$, and
$\mathbf{K}_{\mathbf{x}_{\rm new},\mathbf{M}}$ is a vector of length $mL$ such
that its subvector (from indices $l$ to $m(l+1)-1$) is the cross-covariance
vector between the new point $\mathbf{x}_{\rm new}$ and the $m$ inducing
points for additive component $l$ given embedding matrices $\mathbf{M}$ (see
(7)). With this, we can then approximate the desired posterior distribution by
marginalizing over $\phi(\mathbf{M})=\prod_{l=1}^{L}\phi(\mathbf{M}_{l})$, the
variational distribution on $\mathbf{M}$ in (12):

$[f(\mathbf{x}_{\rm new})|\mathbf{y}]=\int[f(\mathbf{x}_{\rm
new})|\mathbf{M},\mathbf{y}]\;[\mathbf{M}|\mathbf{y}]\;d\mathbf{M}\approx\int\phi(f(\mathbf{x}_{\rm
new})|\mathbf{M})\phi(\mathbf{M})\;d\mathbf{M},$ (16)

where $\phi(f(\mathbf{x}_{\rm new})|\mathbf{M})$ is the multivariate normal
distribution in (15). One can thus sample from this approximate posterior
predictive distribution, by first sampling the embedding matrices $\mathbf{M}$
from $\phi(\mathbf{M})$ in (12), then sampling the prediction
$f(\mathbf{x}_{\rm new})$ from $\phi(f(\mathbf{x}_{\rm new})|\mathbf{M})$ in
(15). Both sampling steps are multivariate normal, and thus can be performed
efficiently after variational parameters are optimized via ELBO.

The above modeling approach for the AdMIn-GP highlights several advantages
over the state-of-the-art, particularly in addressing the limitations from
Section 2.3 for our application. First, the proposed model captures the
desired multi-physics structure in the computer simulator, via an additive
multi-index model with each additive component representing distinct physics
active on a low-dimensional manifold. With a careful selection of the additive
components and manifold ranks, the fitted model can potentially extract
interpretable physical features (see Section 6) that enable improved
predictions with limited data. Second, our approach quantifies uncertainty on
the underlying embedding matrices via a variational inference approach with
shrinkage priors on $\mathbf{M}$. This addresses a key weakness of the
existing models in Section 2.2, which largely do not account for uncertainty
on $\mathbf{M}$ and thus can yield poor coverage for our application (Section
2.3). Finally, the proposed approach enables efficient predictions via a
carefully-constructed variational approximation of the posterior predictive
distribution; we will see later in numerical experiments that this facilitates
quick and accurate surrogate modeling with rich uncertainty quantification.

##  4 Model Implementation

We now discuss important implementation details of the AdMIn-GP that will be
employed in numerical experiments. We first present an approach for optimizing
parameters in the variational lower bound (13), then discuss efficient model
selection approaches for fitting the number of additive components $L$ and the
rank of embedding matrices $p$.

###  4.1 Parameter optimization

Recall that the optimization of the variational bound (13) facilitates the
estimation of variational parameters
$\Theta_{\phi}:=\\{(\mathbf{A}_{l},\mathbf{v}_{l})\\}_{l=1}^{L}$, pseudo-
inputs $\mathcal{Z}:={\\{\mathbf{z}_{j,l}\\}_{j=1}^{M}}_{l=1}^{L}$, and the
AdMIn-GP parameters $\Theta_{\rm
GP}:=\\{(\nu_{l},\theta_{l},\lambda_{l})\\}_{l=1}^{L}$. Letting
$\text{VLB}(\Theta_{\phi},\mathcal{Z},\Theta_{\rm GP})$ denote the variational
lower bound in (13), this optimization can be stated as:

$(\hat{\Theta}_{\phi},\hat{\mathcal{Z}},\hat{\Theta}_{\rm
GP})=\operatorname*{\arg\\!\max}_{\Theta_{\phi},\mathcal{Z},\Theta_{\rm
GP}}\text{VLB}(\Theta_{\phi},\mathcal{Z},\Theta_{\rm GP}).$ (17)

The optimized parameters
$(\hat{\Theta}_{\phi},\hat{\mathcal{Z}},\hat{\Theta}_{\rm GP})$ are then
plugged into (VI) for posterior approximation. In this formulation, both the
inputs and response should be standardized to zero mean and unit variance; we
found this to be essential for stable optimization. Due to the sheer number of
parameters to optimize, we further employed a common set of inducing points
$\\{\mathbf{z}_{j}\\}_{j=1}^{m}$ over all additive components, with common
shrinkage and scale parameters $\lambda$ and $\nu$. In our later experiments,
such simplifications appeared to yield significant computational speed-up
without noticeable compromise in predictive and UQ performance.

The high-dimensional optimization problem (17), however, can be highly
challenging. While analytic gradients can be derived from (VI), the sheer
number of parameters makes such gradients burdensome to compute and carry
around for optimization. We thus leverage recent tools on automatic
differentiation [67], which is widely used for efficient training of complex
machine learning models. The key idea is to exploit the computation of
$\text{VLB}(\cdot)$ as a sequence of elementary operations, to provide
accurate gradient estimates with a single forward evaluation of the objective.
In our implementation, we made use of the automatic differentiation
capabilities in PyTorch [67]. The ADAM optimizer [54], a popular optimizer in
machine learning, is then used for solving (17); we found that a learning rate
of 0.005 with 4000 optimization steps yielded acceptable performance, with a
smaller learning rate of 0.0025 if divergent behavior arises. We further made
use of GPU computing architecture to speed up matrix multiplication steps
involved in evaluating the objective function $\text{VLB}(\cdot)$.

A careful initialization of parameters is also important for successful
optimization. For the variational parameters $\Theta_{\phi}$, we found that a
random initialization
$\mathbf{M}_{l}\sim\mathcal{N}(\boldsymbol{0},\mathbf{I}_{p\times
p},\mathbf{I}_{d\times d})$ is quite effective for the embedding matrices. The
common set of inducing points is initialized as
$\mathbf{z}_{m}\mathbin{\overset{i.i.d.}{\kern
0.0pt\sim}}\mathcal{N}(\boldsymbol{0},\mathbf{I}_{m\times m})$, with column-
wise variances initialized at $\mathbf{v}_{l}=0.1$. For the GP model
parameters $\Theta_{\rm GP}$, we initialized the common shrinkage parameter
$\lambda$ at a relatively small value 0.01, and the common scale parameter
$\nu$ uniformly between $[1,2]$.

###  4.2 Model selection

Another important step is model selection, namely, the selection of the number
of additive components, $L$, and the rank of the embedded subspaces, $p$. A
standard approach would be to estimate such parameters via cross-validation
[42], but with the complexity of our model, this can be prohibitive in terms
of computational cost. Cross-validation may further induce larger biases in
the current limited data regime. We instead perform the model selection of $L$
and $p$ via the maximization of the variational bound (13). More formally,
letting $(\hat{\Theta}_{\phi}(L,p),\hat{\mathcal{Z}}(L,p),\hat{\Theta}_{\rm
GP}(L,p))$ denote the optimized parameters from (17) given fixed $L$ and $d$,
we select the optimal $L^{*}$ and $p^{*}$ via:

$(L^{*},p^{*})=\operatorname*{\arg\\!\max}_{L,p}\text{VLB}(\hat{\Theta}_{\phi}(L,p),\hat{\mathcal{Z}}(L,p),\hat{\Theta}_{\rm
GP}(L,p)).$ (18)

As VLB lower bounds the log-marginal likelihood of the data $\mathbf{y}$ (see
(13)), such a procedure can be viewed as an approximate empirical Bayes model
selection approach under uniform (flat) priors over feasible combinations of
$(L,p)$ [47]. In general, we find that selection of embedding dimension $p$ to
be less important than the number of additive components $L$; the shrinkage
priors on $\mathbf{M}_{l}$, coupled with known sparsity of the multi-physics
phenomena (see Section 3.1), results in similar empirical performance over
different choices of $p$.

In implementation, we found that to ensure numerical stability, the embedding
rank $p$ should be set above a lower limit of 4 for model selection. An upper
limit on $p$ does not appear necessary though, since the employed shrinkage
priors on $\mathbf{M}$ facilitate stable model fits for large $d$. For the
number of inducing points $m$, one should choose as many points as possible,
but this can be quite limited due to computation. For our later experiments,
we made use of $m=125$ inducing points, which seemed to provide adequate
performance.

##  5 Numerical Experiments

We now investigate the predictive and UQ performance of the proposed AdMIn-GP.
We first discuss the simulation set-up, including evaluation metrics and
compared models. We then explore the effectiveness of the AdMIn-GP in a suite
of simulation experiments.

###  5.1 Simulation set-up

We compare with the same four models tested in Section 2.3:

  * •

The standard GP model with the squared-exponential kernel and automatic
relevance determination (ARD-GP; see [37]), implemented in the GPy package
[36] in Python,

  * •

The fully Bayesian single-index model (SIM-GP; [39]), implemented in the R
package tgp [38],

  * •

The dimension-reduced (or active subspace) GP (DR-GP; [77, 79, 83]), with
point estimates for the active subspace, implemented in the Python package
GPyTorch [32].

  * •

The projection-pursuit-based diverse projected additive GP (DPA-GP; [21]).

The first is a standard baseline for surrogate modeling, and the latter
surrogate models integrate some form of dimension reduction. The AdMIn-GP is
fitted following the variational inference framework in Section 3 and the
optimization approach in Section 4.

We then generate the training design points $\\{\mathbf{x}_{i}\\}_{i=1}^{n}$
from a Latin hypercube design [61], with varying sample sizes of $n=300$,
$400$ and $500$. The number of input parameters $d$ is set to be 20 to reflect
our application (this will be increased in the next subsection). This
procedure is replicated 10 times to provide a quantification of simulation
variability. The embedded manifold dimension $p$ for DR-GP and DPA-GP are
selected in an empirical Bayes fashion to maximize the marginal likelihood (or
ELBO). Details for selection of $(L,p)$ for AdMIn-GP can be found in Section
5.3.

We then evaluate the performance of these models on several metrics. The first
is the standard root-mean-squared-error (RMSE) over a uniformly-sampled test
set of size 1000, which measures point prediction accuracy. The second is the
widely-used continuous ranked probability score (CRPS; [35]), which measures
accuracy of probabilistic predictions. For both measures, smaller values
indicate better predictive performance. Here, the CRPS is computed from 500
samples drawn from the posterior predictive distribution of each method [45].
The third metric measures the empirical coverage rate of the 95% posterior
predictive intervals over the same test set, and coverage close to the nominal
95% rate is desired.

###  5.2 Prediction and coverage

The base simulation case considers the following synthetic function on $f$,
composed of $L=3$ additive components that each depend on their separate
$p=2$-dimensional subspaces.

$\displaystyle\begin{split}f(\mathbf{x})&=0.4f_{1}(\mathbf{M}_{1}\mathbf{x})+0.3f_{2}(\mathbf{M}_{2}\mathbf{x})+0.3f_{3}(\mathbf{M}_{3}\mathbf{x}),\\\
f_{1}(\mathbf{c})&=\sin(c_{1})+\cos(0.5c_{2}),\quad
f_{2}(\mathbf{c})=\cos(0.5c_{1}c_{2}),\quad
f_{3}(\mathbf{c})=c_{2}\cos(0.75c_{1}).\end{split}$ (19)

We will use this to mimic the presence of dominant multi-physics in the
response surface, where each additive component captures a separate physics
that depends on a low-dimensional embedding (see [17]). To capture this, the
embedding matrices $\\{\mathbf{M}_{l}\\}_{l=1}^{3}$ are generated with
$(i,j)$-th entry sampled independently from $\mathcal{N}(3,1)$ if
$l\lfloor{d}/{3}\rfloor\leq j<(l+1)\lfloor{d}/{3}\rfloor$ and
$\mathcal{N}(0,1)$ otherwise, such that each additive component is dominantly
influenced by different subsets of parameters. Its rows are then normalized to
unit variance. We set the noise standard deviation $\sigma$ to be 15% of the
standard deviation for $f(\mathbf{u})$, where
$\mathbf{u}\sim\mathcal{N}(\boldsymbol{0},\mathbf{I})$.

Figure 4 shows the prediction metrics for the compared models in this base
simulation study. We see that the proposed AdMIn-GP yields the best prediction
metrics of all models: it provides notably lower RMSE and lower CRPS, thus
indicating more accurate point and probabilistic predictions. Similarly, the
empirical coverage rates for the AdMIn-GP is much closer to the desired 95%
rate, whereas existing models yield markedly lower coverage; this is not too
surprising, since much of these methods do not account for estimation
uncertainty of the embedded manifold structure. This confirms a previous
observation: when the underlying low-dimensional structure is misspecified
(i.e., in DR-GP and DPA-GP), the resulting model may yield noticeably poorer
predictions and coverage over the standard GP, which does not leverage such
structure! When this embedding is carefully elicited, the resulting model
(i.e., AdMIn-GP) can harness such structure to provide accurate predictions
and reliable uncertainties with limited training data – a key challenge in our
application.

We now investigate the effect of increasing the number of parameters $d$.
Figure 5 shows the predictive metrics for the compared models for two larger
choices of $d=25$ and $d=30$. We see that similar results hold: the AdMIn-GP
provide markedly more accurate point and probabilistic predictions, with
empirical coverage rates closer to the desired 95%. An interesting observation
is that, as $d$ increases, the models with misspecified low-dimensional
structure (DR-GP and DPA-GP) yield increasingly poorer predictive performance
over the standard GP (ARD-GP). This again suggests that, for complex systems
with many input parameters, the careful elicitation and integration of known
low-dimensional structures become increasingly crucial for effective surrogate
modeling with limited training data. The AdMIn-GP provides a means for
facilitating this integration for cost-efficient emulation of multi-physics
systems, as we show later in our motivating high-energy physics application.

Figure 4: Boxplots of the predictive metrics for the base simulation in $d=20$
dimensions over various sample sizes $n$.

Figure 5: Boxplots of the predictive metrics for the simulations in $d=25$ and
$d=30$ dimensions over various sample sizes $n$.

###  5.3 Model and variable selection

As the modeled embedded structure in the AdMIn-GP is guided by the underlying
multi-physics, the estimation of certain model parameters (e.g., the number of
additive components $L$ and the embedding matrices
$\\{\mathbf{M}_{l}\\}_{l=1}^{L}$) may in turn shed light on the underlying
dominant physics within the simulator. For example, in our QGP application,
the selection of $L$ may inform the degree of physics heterogeneity in the
heavy-ion collision system, and the estimation of
$\\{\mathbf{M}_{l}\\}_{l=1}^{L}$ can then help identify the nature of specific
dominant physics. With careful discussions with domain scientists, this may
serve as a useful tool for guiding scientific discovery, as we show later in
Section 6.

Consider first the model selection performance for the number of additive
components $L$ in the earlier simulations. Figure 6 shows the boxplots of the
“normalized” ELBO bound over different choices of $L$, in dimensions $d=20,25$
and 30. For each of $L=1,2,3$ and $4$, we set the embedding matrix dimension
$p$ for each component to be $p=12,6,4$ and 4, respectively, so that each
model has access to approximately the same number of features, namely
12.555For the $L=4$ case, we used $p=4$ since the setting of $p=3$ yielded
numerical instability issues. This “normalized” ELBO bound is the variational
bound (13) with a fixed number of components $L$, subtracted from the same
variational bound with $L=3$ (the true number of components). We thus select
the model with the largest ELBO, which lower bounds the marginal likelihood.
For small sample sizes ($n=300$), while the procedure tends to select the
correct model with $L=3$ components, it may err by selecting a smaller number
of components, particularly for larger $d$. This is not too surprising, since
there are simply insufficient data to justify the more complex model. As more
data are collected, we see that the correct model with $L=3$ is consistently
selected, which is as desired.

Figure 6: Boxplots of the “normalized” ELBO bound (13) for the AdMIn-GP with
$L$ additive components, over various sample sizes $n$. Negative values
suggest that the ELBO model selection procedure prefers the $L=3$ component
model.

Consider next the estimation of the embedding matrices
$\\{\mathbf{M}_{l}\\}_{l=1}^{L}$. An important question is the identifiability
of such matrices from the data $\mathbf{y}$. Indeed, since the GP length-scale
parameters $\\{\theta_{l}\\}_{l=1}^{L}$ are estimated from data, any scaling
of the embedding matrix $\mathbf{M}_{l}$ by a constant would be
unidentifiable, as it can be offset by a corresponding scaling of the length-
scale parameter. Fortunately, for the aforementioned goal of extracting
dominant physics, we require only the selection of non-zero entries in
$\\{\mathbf{M}_{l}\\}_{l=1}^{L}$ to pinpoint important variables in
$\mathbf{x}$ contributing to different multi-physics (see Section 6.2). While
we include no consistency results of this flavor for the AdMIn-GP (which are
beyond the scope of the paper), there is a rich literature on the theoretical
consistency and empirical performance of variable selection using GPs (see
[57, 66, 51]), which suggests that the proposed model might be useful for this
physics extraction goal. We investigate this empirically below.

We now inspect how well our variational approach can identify important
variables within the three embedding matrices $\mathbf{M}_{1}$,
$\mathbf{M}_{2}$ and $\mathbf{M}_{3}$ in the earlier simulations. Figure 7
(left) shows the absolute values of the true matrices for the experiment in
$d=20$ dimensions, where each matrix is influenced by a sparse number of
parameters (as expected in multi-physics applications; see Section 3.1).
Figure 7 (right) shows the absolute values of the estimated matrices
$\hat{\mathbf{M}}_{1}$, $\hat{\mathbf{M}}_{2}$ and $\hat{\mathbf{M}}_{3}$ (via
posterior means) using the proposed variational inference approach. We see
that, with selected manifold rank $p=4$, each of the estimated matrices
roughly identified the correct block of active parameters. For
$\hat{\mathbf{M}}_{1}$, we see that the AdMIn-GP correctly identifies the
first block of six variables as active, with other variables largely set as
inert; similar conclusions hold for $\hat{\mathbf{M}}_{2}$ and
$\hat{\mathbf{M}}_{3}$. This identification of important variables in the
embeddings can help guide the discovery of interpretable dominant physics
embedded in the system, as we show next.

Figure 7:  (Left): Absolute values of the true (simulated) matrices
$\mathbf{M}_{1}$, $\mathbf{M}_{2}$ and $\mathbf{M}_{3}$ in $d=20$ dimensions.
(Right): Absolute values of the estimated matrices $\hat{\mathbf{M}}_{1}$,
$\hat{\mathbf{M}}_{2}$ and $\hat{\mathbf{M}}_{3}$ via posterior means using
the proposed variational inference approach.

##  6 Surrogate Modeling of the Quark-Gluon Plasma

We now return to our motivating application for the QGP, where the goal is to
train a probabilistic emulator for simulated heavy-ion collision observables.
Here, we investigate collisions of gold nuclei at near-light speeds (i.e.,
ultrarelativistic Au+Au collisions), which have been studied experimentally at
the Relativistic Heavy Ion Collider in Brookhaven National Laboratory. The
observables (responses) of interest are the so-called anisotropic flow
coefficients $v_{k}$’s, which quantify momentum anisotropy in the number of
particles emitted in the transverse plane. Letting $\phi$ be the angle in the
transverse plane, the particle number $dN/d\phi$ is expanded in a Fourier
series and the $v_{k}$’s are defined as:

$\frac{dN}{d\phi}=\left\langle\frac{dN}{d\phi}\right\rangle\left[1+2\sum_{k=1}^{\infty}v_{k}\cos\left(k\left(\phi-\Psi_{k}\right)\right)\right]$
(20)

where $\langle dN/d\phi\rangle$ is the average number of particles produced in
a given collision, $v_{k}$ are its Fourier coefficients and $\Psi_{k}$ are the
event plane angles [43]. Measurements of $v_{k}$ are averaged over thousands
of collisions, and calculations are averaged in a similar manner.

An important discovery in heavy-ion collisions is the correlation between the
initial spatial geometry of colliding nuclei and the final momentum asymmetry
of the number of particles at the end of the collisions (see [59] for a
review). As such, we use the second and third Fourier coefficients $v_{2}$ and
$v_{3}$, which are dominant contributions to this anisotropy [59]. The
measurement of these observables can further be categorized by “centrality
classes”, which describe the degree of overlap of the colliding nuclei at
impact. Figure 8 visualizes two collisions at low and high centrality. We make
use of $v_{2}$ and $v_{3}$ in the 0–5% centrality class (the 5% of collisions
with the most overlap of the colliding nuclei), and the 5–10% class (the next
5% of collisions with the most overlap), resulting in four observables of
interest. This is a subset of the observables employed in a recent Bayesian
analysis of the QGP [28], which yielded meaningful parameter constraints from
experiments.

Figure 8: Visualizing a central nucleus-nucleus collision (left) where the
nuclei undergo near head-on collisions, resulting in high centrality, and a
peripheral collision (right) where the nuclei collide off-center, resulting in
low centrality.

For inputs, the simulator has $d=17$ parameters: five describing the impact of
the nuclei and the resulting energy deposition, two related to the transition
of the systems into a quark-gluon plasma, nine for the temperature dependence
of the shear and bulk viscosity of the quark-gluon plasma and the associated
shear relaxation time, and one final parameter associated with the transition
of the plasma’s description from a fluid into a collection of individual
particles. The same parametrization of this collision system was used in prior
Bayesian analyses of the QGP (see [58, 24, 29]).

###  6.1 Prediction and coverage

Recall (see Section 2.1) that there are two key challenges for surrogate
modeling of the QGP. First, prediction over this high-dimensional ($d=17$)
space is challenging. Second, this is complicated by the costly simulation
runs needed to generate the training data, which requires thousands of CPU
hours per design point. With limited supercomputing budget for this project,
we can afford only $n^{*}=474$ training points, sampled from a Latin hypercube
design [61]. Given such limited data in high dimensions, the task of accurate
emulation with high predictive certainty is extremely challenging. We
investigate whether, by learning and integrating dominant multi-physics
structure for emulation, the AdMIn-GP can yield accurate and confident
surrogate modeling in this challenging setting.

As before, we compared the proposed AdMIn-GP with the same existing models in
our simulations (see Section 5.1). To evaluate predictive performance, we did
not use an additional testing set due to the highly expensive nature of
simulating such a test set. Rather, we made use of random splits of the full
$n^{*}=474$ into training sets of sample sizes $n=200$, $250$, $300$ and
$350$, with the remaining points used for testing. This random training-
testing split is replicated 10 times to measure error variability. The same
three metrics (RMSE, CRPS and empirical coverage rate) are used to gauge
predictive performance.

Figure 9 shows the RMSEs and CRPSs of the compared surrogate models, over
different sample sizes $n=200,250,300$ and $350$ for each observable. We see
that the proposed AdMIn-GP consistently yields significantly lower RMSEs to
existing models, over all sample sizes $n$. Existing models that leverage some
form of low-dimensional embedded structure (DR-GP, DPA-GP and SIM-GP) again
yield mediocre predictions, with errors noticeably larger than the standard GP
surrogate. From earlier simulations, this suggests that the underlying
embedded structure may indeed be misspecified by existing models; in
retrospect, this is not surprising, as such models do not account for the
presence of multi-physics. The standard GP surrogate, ARD-GP, yields
relatively high errors over all observables, and does not appear to decrease
as sample size $n$ increases. This confirms the limitations outlined in
Section 2.2: standard GPs suffer from the curse-of-dimensionality of requiring
an exponentially-growing sample size $n$ in dimension $d$ to achieve good
predictive performance. Here, with small $n$ and large $d$, such models thus
yield high errors that yield minimal decreases as $n$ increases. By eliciting
and integrating the desired embedded structures from multi-physics for
surrogate modeling, the AdMIn-GP can provide accurate predictions in this
challenging small $n$, large $d$ setting. This suggests that our model can
capture the true low-dimensional structure with much less misspecification. We
will show later that this learned structure can further guide the extraction
of interpretable multi-physics.

Figure 10 shows the corresponding empirical coverage rates of the compared
models in this application, over different sample sizes $n$. As in
simulations, we see that the proposed model provides much closer rates to the
desired nominal 95% rate for all observables and sample sizes. Existing
models, however, yield much lower coverage rates than 95%, particularly for
models that attempt to fit embedded low-dimensional structure. There may be
two reasons for this. First, as noted before, such existing models largely do
not consider uncertainties with estimating the underlying embedding structure
from data. With small sample sizes $n$ in high dimensions $d$, these
uncertainties can be large and may result in significant undercoverage if
ignored. Second, from earlier analysis, such existing models likely
misspecified the more complex embedding structure arising from the multi-
physics system. In doing so, the fitted model would be overconfident in
predicting the more complex response surface, resulting in large undercoverage
of predictive intervals.

Figure 9: Boxplots of predictive metrics (RMSE and CRPS) for the compared
models over various sample sizes $n$. Each of the four columns corresponds to
a different collision observable. Figure 10: Boxplots of the empirical
coverage rates (of 95% predictive intervals) for the compared models over
various sample sizes $n$. Each column corresponds to a different collision
observable.

###  6.2 Extracting dominant multi-physics

Figure 11: Boxplots of the “normalized” ELBO bound (13) for the AdMIn-GP with
$L$ additive components over various sample sizes $n$.. Negative values
suggest that the ELBO model selection procedure prefers the $L=2$ component
model.

Finally, we explore the model and variable selection performance of the AdMIn-
GP and discuss its implications for extracting multi-physics on the collision
system. We first investigate the selected number of additive components $L$,
then examine the selected important variables in the estimated embedding
matrices, and how this might guide the extraction of interpretable multi-
physics dominant within the QGP.

Consider first the selection of the number of additive components $L$ via the
ELBO bound (13). Here, we explored choices of $L$ from 2 to 4 components; the
setting of $L=1$ is excluded here since the QGP is known to be influenced by
multiple distinct types of physics [27]. Figure 11 shows the “normalized” ELBO
bound, i.e., the ELBO bound (13) subtracted from the baseline ELBO bound with
$L=2$, for different sample sizes $n$ and different observables. We see that,
for all observables, the AdMIn-GP selects the baseline $L=2$ setting, which
hints at two distinct dominant physics at play within the QGP. Further, as
sample size $n$ increases, we see a larger gap in the ELBO bound between the
baseline $L=2$ and larger choices of $L$, which again suggests the presence of
two dominant physics.

Figure 12: Absolute values of the estimated matrices $\hat{\mathbf{M}}_{1}$
and $\hat{\mathbf{M}}_{2}$ for each of the four QGP observables. The first two
columns are the estimated matrices for the first asymmetry observable at
centrality $0-5\%$ and $5-10\%$, and the last two columns are for the second
asymmetry observable at centrality $0-5\%$ and $5-10\%$. Parameters
corresponding to collision initial conditions and plasma shear viscosity are
highlighted in red and blue, respectively.

The nature of this dominant multi-physics within the QGP can further be
gleaned via a careful inspection of the estimated embedding matrices. Figure
12 shows the estimated matrices $\hat{\mathbf{M}}_{1}$ and
$\hat{\mathbf{M}}_{2}$ for the four observables. There are several
illuminating observations to note. Consider the first observable $v_{2}$,
i.e., the left two observables in Figure 12. We see that a majority of non-
zero entries selected by the first embedding matrix $\hat{\mathbf{M}}_{1}$
(top row) involves the parameters $x_{2}$, $x_{3}$ and $x_{4}$ (colored red on
plot). Upon closer inspection, these three parameters (Trento-$p$,
$\sigma_{k}$ and nucleon width) correspond specifically to the initial
conditions of the collision system [63]. It is well-accepted that the geometry
and fluctuations of the collision initial conditions are directly related to
momentum asymmetry [65, 62, 1, 2, 31], thus the first component of the AdMIn-
GP likely captures the effect of such dominant initial conditions. Similarly,
the second embedding matrix (bottom row) selects non-zero entries on the
parameters $x_{8}$ and, to a lesser extent, $x_{9}$ and $x_{11}$, which are
three of the four parameters (colored blue in plot) used to parametrize the
shear viscosity of the plasma (see [27] for a detailed discussion of this
parametrization). This dependence on shear viscosity affirms the sensitivity
analysis in [28], which found a strong dependence of the observable $v_{2}$ on
parameter $x_{8}$, but little dependence on the fourth shear viscosity
parameter $x_{10}$. The second component of the AdMIn-GP thus likely captures
the dependence of $v_{2}$ on plasma shear viscosity. One also observes a
slight dependence on initial conditions (i.e., the parameters
$x_{2},x_{3},x_{4}$) for the second matrix $\hat{\mathbf{M}}_{2}$, which
confirms the known correlation between shear viscosity with the initial
condition parameters of the QGP [43, 30].

Consider next the second observable $v_{3}$, i.e., the right two observables
in Figure 12. For the first embedding matrix $\hat{\mathbf{M}}_{1}$, we notice
a much larger dependence (compared to $v_{2}$) on the four shear viscosity
parameters ($x_{8}-x_{11}$). This is not surprising from a physics
perspective, as this asymmetric observable is related to shorter-scale
structures of the plasma fluid that are damped more strongly by viscosity [43,
30]. This stronger dependence on the shear viscosity is now captured in
$\hat{\mathbf{M}}_{1}$ as the dominant physics, with $\hat{\mathbf{M}}_{2}$
capturing a combination of initial condition parameters and their interaction
with plasma shear viscosity. This extraction of interpretable multi-physics in
observables $v_{2}$ and $v_{3}$ not only affirms the modeled embedding
structure in the AdMIn-GP, but also provides a data-driven approach for
discovering and verifying dominant physical phenomena within the QGP.

Finally, we note that, for each of the two momentum asymmetry observables
$v_{2}$ and $v_{3}$, the extracted dominant multi-physics (in the form of
important variables in $\hat{\mathbf{M}}_{1}$ and $\hat{\mathbf{M}}_{2}$) are
quite similar over different centralities; similar groups of parameters are
selected as active or inert between the two centrality bins $0-5\%$ and
$5-10\%$. This is consistent with the understanding that the centrality
dependence of $v_{2}$ and $v_{3}$ is moderate [43, 30], and similar results
should be expected for neighboring centralities.

##  7 Conclusion

The quark-gluon plasma (QGP) is a unique phase of nuclear matter, theoretized
to have filled the Universe shortly after the Big Bang. The study of this
plasma requires a probabilistic surrogate model, which can efficiently emulate
the complex relationship between many physical parameters $\mathbf{x}$ and its
observables $f(\mathbf{x})$ with expensive (and thus limited) training data.
We thus propose a new Additive Multi-Index Gaussian process (AdMIn-GP) model,
which features a flexible additive model of GPs with each component active on
different latent low-dimensional linear embeddings. This structure is directly
guided by prior knowledge that the QGP is controlled by several dominant
multi-physics. We present an efficient framework for model fitting via a
carefully constructed variational inference approach with inducing points. We
then demonstrate the improved predictive performance and uncertainty
quantification of the AdMIn-GP over competing methods in a suite of numerical
experiments and for surrogate modeling of the QGP, and show it may facilitate
the extraction of interpretable multi-physics for guiding scientific
discoveries.

Acknowledgements: The authors gratefully acknowledge funding from NSF CSSI
Frameworks grant 2004571 (KL, SM), NSF DMS 2210729, 2220496 (KL, SM) and U.S.
Department of Energy Grant DE-FG02-05ER41367 (SAB, JFP). The authors also
thank the JETSCAPE collaboration for useful comments and feedback.

## References

  * Alver, [2007] Alver, B. (2007).  System size, energy, pseudorapidity, and centrality dependence of elliptic flow.  Phys. Rev. Lett., 98:242302.
  * Alver and Roland, [2010] Alver, B. and Roland, G. (2010).  Collision geometry fluctuations and triangular flow in heavy-ion collisions.  Phys. Rev. C, 81:054905.  [Erratum: Phys.Rev.C 82, 039903 (2010)].
  * Banerjee et al., [2008] Banerjee, S., Gelfand, A. E., Finley, A. O., and Sang, H. (2008).  Gaussian predictive process models for large spatial data sets.  Journal of the Royal Statistical Society. Series B (Statistical Methodology), 70(4):825–848.
  * Ben Salem et al., [2019] Ben Salem, M., Bachoc, F., Roustant, O., Gamboa, F., and Tomaso, L. (2019).  Sequential dimension reduction for learning features of expensive black-box functions.  Preprint at https://hal.science/hal-01688329v2/file/main.pdf.
  * Bernhard et al., [2019] Bernhard, J. E., Moreland, J. S., and Bass, S. A. (2019).  Bayesian estimation of the specific shear and bulk viscosity of quark–gluon plasma.  Nature Physics, 15(11):1113–1117.
  * Bhattacharya et al., [2015] Bhattacharya, A., Pati, D., Pillai, N. S., and Dunson, D. B. (2015).  Dirichlet–Laplace priors for optimal shrinkage.  Journal of the American Statistical Association, 110(512):1479–1490.
  * Binois et al., [2018] Binois, M., Huang, J., Gramacy, R. B., and Ludkovski, M. (2018).  Replication or exploration? Sequential design for stochastic simulation experiments.  Technometrics, 61(1):7–23.
  * Blei et al., [2017] Blei, D. M., Kucukelbir, A., and McAuliffe, J. D. (2017).  Variational inference: A review for statisticians.  Journal of the American Statistical Association, 112(518):859–877.
  * Buckingham, [1914] Buckingham, E. (1914).  On physically similar systems; illustrations of the use of dimensional equations.  Physical Review, 4(4):345.
  * Cao et al., [2021] Cao, S., Chen, Y., Coleman, J., Mulligan, J., Jacobs, P., Soltz, R., Angerami, A., Arora, R., Bass, S., Cunqueiro, L., et al. (2021).  Determining the jet transport coefficient $\hat{q}$ from inclusive hadron suppression measurements using Bayesian parameter estimation.  Physical Review C, 104(2):024905.
  * Carlin and Louis, [1997] Carlin, B. P. and Louis, T. A. (1997).  Bayes and Empirical Bayes Methods for Data Analysis.  Springer.
  * Castillo and van der Vaart, [2012] Castillo, I. and van der Vaart, A. (2012).  Needles and Straw in a Haystack: Posterior concentration for possibly sparse sequences.  The Annals of Statistics, 40(4):2069 – 2101.
  * Chen and Tuo, [2022] Chen, G. and Tuo, R. (2022).  Projection pursuit Gaussian process regression.  IISE Transactions.  To appear.
  * Chen et al., [2020] Chen, J., Mak, S., Joseph, V. R., and Zhang, C. (2020).  Function-on-function kriging, with applications to three-dimensional printing of aortic tissues.  Technometrics, 63:384–395.
  * Chen et al., [2023] Chen, Z., Mak, S., and Wu, C. F. J. (2023).  A hierarchical expected improvement method for Bayesian optimization.  Journal of the American Statistical Association, to appear.
  * Chikuse, [2003] Chikuse, Y. (2003).  Statistics on Special Manifolds, volume 1.  Springer.
  * Constantine et al., [2014] Constantine, P. G., Dow, E., and Wang, Q. (2014).  Active subspace methods in theory and practice: applications to kriging surfaces.  SIAM Journal on Scientific Computing, 36(4):A1500–A1524.
  * Constantine et al., [2015] Constantine, P. G., Zaharatos, B., and Campanelli, M. (2015).  Discovering an active subspace in a single-diode solar cell model.  Statistical Analysis and Data Mining: The ASA Data Science Journal, 8(5):264–273.
  * Currin et al., [1991] Currin, C., Mitchell, T., Morris, M., and Ylvisaker, D. (1991).  Bayesian prediction of deterministic functions, with applications to the design and analysis of computer experiments.  Journal of the American Statistical Association, 86(416):953–963.
  * Damianou and Lawrence, [2013] Damianou, A. and Lawrence, N. D. (2013).  Deep Gaussian processes.  In Carvalho, C. M. and Ravikumar, P., editors, Proceedings of the Sixteenth International Conference on Artificial Intelligence and Statistics, pages 207–215.
  * Delbridge et al., [2020] Delbridge, I., Bindel, D., and Wilson, A. G. (2020).  Randomly projected additive Gaussian processes for regression.  Proceedings of the 37th International Conference on Machine Learning, 119:2453–2463.
  * Durrande et al., [2012] Durrande, N., Ginsbourger, D., and Roustant, O. (2012).  Additive covariance kernels for high-dimensional Gaussian process modeling.  Annales de la Faculté des sciences de Toulouse : Mathématiques, 21(3):481–499.
  * Duvenaud et al., [2011] Duvenaud, D. K., Nickisch, H., and Rasmussen, C. (2011).  Additive Gaussian processes.  In Advances in Neural Information Processing Systems, volume 24.
  * Ehlers et al., [2022] Ehlers, R., Angerami, A., Arora, R., Bass, S., Cao, S., Chen, Y., Du, L., Dai, T., Elfner, H., Fan, W., et al. (2022).  Bayesian analysis of QGP jet transport using multi-scale modeling applied to inclusive hadron and reconstructed jet data.  arXiv preprint arXiv:2208.07950.
  * Eriksson and Jankowiak, [2021] Eriksson, D. and Jankowiak, M. (2021).  High-dimensional Bayesian optimization with sparse axis-aligned subspaces.  In Proceedings of the Thirty-Seventh Conference on Uncertainty in Artificial Intelligence, pages 493–503. PMLR.  ISSN: 2640-3498.
  * Everett, [2021] Everett, D. (2021).  Quantifying the Quark Gluon Plasma.  The Ohio State University.
  * [27] Everett, D., Ke, W., Paquet, J.-F., Vujanovic, G., Bass, S., Du, L., Gale, C., Heffernan, M., Heinz, U., Liyanage, D., et al. (2021a).  Multisystem Bayesian constraints on the transport coefficients of QCD matter.  Physical Review C, 103(5):054904.
  * [28] Everett, D., Ke, W., Paquet, J.-F., Vujanovic, G., Bass, S., Du, L., Gale, C., Heffernan, M., Heinz, U., Liyanage, D., et al. (2021b).  Phenomenological constraints on the transport properties of QCD matter with data-driven model averaging.  Physical Review Letters, 126(24):242301.
  * Everett et al., [2022] Everett, D., Oliinychenko, D., Luzum, M., Paquet, J.-F., Vujanovic, G., Bass, S., Du, L., Gale, C., Heffernan, M., Heinz, U., et al. (2022).  Role of bulk viscosity in deuteron production in ultrarelativistic nuclear collisions.  Physical Review C, 106(6):064901.
  * Gale et al., [2013] Gale, C., Jeon, S., and Schenke, B. (2013).  Hydrodynamic modeling of heavy-ion collisions.  Int. J. Mod. Phys. A, 28:1340011.
  * Gardim et al., [2012] Gardim, F. G., Grassi, F., Luzum, M., and Ollitrault, J.-Y. (2012).  Mapping the hydrodynamic response to the initial geometry in heavy-ion collisions.  Phys. Rev. C, 85:024908.
  * Gardner et al., [2018] Gardner, J., Pleiss, G., Weinberger, K. Q., Bindel, D., and Wilson, A. G. (2018).  GPytorch: Blackbox matrix-matrix Gaussian process inference with GPU acceleration.  In Advances in Neural Information Processing Systems, volume 31.
  * Gilboa et al., [2013] Gilboa, E., Saatçi, Y., Cunningham, J., and Gilboa, E. (2013).  Scaling multidimensional Gaussian processes using projected additive approximations.  In International Conference on Machine Learning, pages 454–461. PMLR.
  * Glaws et al., [2017] Glaws, A., Constantine, P. G., Shadid, J. N., and Wildey, T. M. (2017).  Dimension reduction in magnetohydrodynamics power generation models: Dimensional analysis and active subspaces.  Statistical Analysis and Data Mining: The ASA Data Science Journal, 10(5):312–325.
  * Gneiting and Raftery, [2007] Gneiting, T. and Raftery, A. E. (2007).  Strictly proper scoring rules, prediction, and estimation.  Journal of the American Statistical Association, 102(477):359–378.
  * GPy, [2012] GPy (since 2012).  GPy: A Gaussian process framework in Python.  http://github.com/SheffieldML/GPy.
  * Gramacy, [2020] Gramacy, R. (2020).  Surrogates: Gaussian Process Modeling, Design and Optimization for the Applied Sciences.  Chapman Hall/CRC.
  * Gramacy, [2007] Gramacy, R. B. (2007).  tgp: an R package for Bayesian nonstationary, semiparametric nonlinear regression and design by treed Gaussian process models.  Journal of Statistical Software, 19:1–46.
  * Gramacy and Lian, [2012] Gramacy, R. B. and Lian, H. (2012).  Gaussian process single-index models as emulators for computer experiments.  Technometrics, 54(1):30–41.
  * Guan et al., [2020] Guan, C., Lu, M., Zeng, W., Yang, D., and Han, D. (2020).  Prediction of standard enthalpies of formation based on hydrocarbon molecular descriptors and active subspace methodology.  Industrial & Engineering Chemistry Research, 59(10):4785–4791.
  * Hall and Li, [1993] Hall, P. and Li, K.-C. (1993).  On almost linearity of low dimensional projections from high dimensional data.  The Annals of Statistics, 21(2):867 – 889.
  * Hastie et al., [2009] Hastie, T., Tibshirani, R., and Friedman, J. H. (2009).  The Elements of Statistical Learning: Data Mining, Inference, and Prediction, volume 2.  Springer.
  * Heinz and Snellings, [2013] Heinz, U. and Snellings, R. (2013).  Collective flow and viscosity in relativistic heavy-ion collisions.  Ann. Rev. Nucl. Part. Sci., 63:123–151.
  * Hemant and Cevher, [2012] Hemant, T. and Cevher, V. (2012).  Active learning of multi-index function models.  In Pereira, F., Burges, C., Bottou, L., and Weinberger, K., editors, Advances in Neural Information Processing Systems, volume 25.
  * Hersbach, [2000] Hersbach, H. (2000).  Decomposition of the continuous ranked probability score for ensemble prediction systems.  Weather and Forecasting, 15:559–570.
  * Hoffman and Blei, [2015] Hoffman, M. and Blei, D. (2015).  Stochastic structured variational inference.  In Lebanon, G. and Vishwanathan, S. V. N., editors, International Conference on Artificial Intelligence and Statistics, pages 361–369. PMLR.
  * Hoffman et al., [2013] Hoffman, M. D., Blei, D. M., Wang, C., and Paisley, J. (2013).  Stochastic variational inference.  Journal of Machine Learning Research, 14(40):1303–1347.
  * Hu et al., [2013] Hu, Y., Gramacy, R. B., and Lian, H. (2013).  Bayesian quantile regression for single-index models.  Statistics and Computing, 23:437–454.
  * Ji et al., [2021] Ji, Y., Mak, S., Soeder, D., Paquet, J., and Bass, S. A. (2021).  A graphical Gaussian process model for multi-fidelity emulation of expensive computer codes.  arXiv preprint arXiv:2108.00306.
  * Ji et al., [2022] Ji, Y., Yuchi, H. S., Soeder, D., Paquet, J.-F., Bass, S. A., Joseph, V. R., Wu, C., and Mak, S. (2022).  Conglomerate multi-fidelity Gaussian process modeling, with application to heavy-ion collisions.  arXiv preprint arXiv:2209.13748.
  * Jiang and Tokdar, [2021] Jiang, S. and Tokdar, S. T. (2021).  Variable selection consistency of Gaussian process regression.  The Annals of Statistics, 49(5):2491–2505.
  * Jones et al., [1998] Jones, D. R., Schonlau, M., and Welch, W. J. (1998).  Efficient global optimization of expensive black-box functions.  Journal of Global Optimization, 13(4):455–492.
  * Kandasamy et al., [2015] Kandasamy, K., Schneider, J., and Poczos, B. (2015).  High dimensional Bayesian optimisation and bandits via additive models.  In Bach, F. and Blei, D., editors, International Conference on Machine Learning, pages 295–304. PMLR.
  * Kingma and Ba, [2017] Kingma, D. P. and Ba, J. (2017).  Adam: A method for stochastic optimization.  In International Conference on Learning Representations.
  * Kumar et al., [2023] Kumar, A., Tachibana, Y., Sirimanna, C., Vujanovic, G., Cao, S., Majumder, A., Chen, Y., Du, L., Ehlers, R., Everett, D., et al. (2023).  Inclusive jet and hadron suppression in a multistage approach.  Physical Review C, 107(3):034911.
  * Letham et al., [2020] Letham, B., Calandra, R., Rai, A., and Bakshy, E. (2020).  Re-examining linear embeddings for high-dimensional Bayesian optimization.  In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H., editors, Advances in Neural Information Processing Systems, pages 1546–1558.
  * Linkletter et al., [2006] Linkletter, C., Bingham, D., Hengartner, N., Higdon, D., and Ye, K. Q. (2006).  Variable selection for Gaussian process models in computer experiments.  Technometrics, 48(4):478–490.
  * Liyanage et al., [2022] Liyanage, D., Ji, Y., Everett, D., Heffernan, M., Heinz, U., Mak, S., and Paquet, J.-F. (2022).  Efficient emulation of relativistic heavy ion collisions with transfer learning.  Physical Review C, 105(3):034910.
  * Luzum and Petersen, [2014] Luzum, M. and Petersen, H. (2014).  Initial state fluctuations and final state correlations in relativistic heavy-ion collisions.  J. Phys. G, 41:063102.
  * Mak et al., [2018] Mak, S., Sung, C.-L., Wang, X., Yeh, S.-T., Chang, Y.-H., Joseph, V. R., Yang, V., and Wu, C. F. J. (2018).  An efficient surrogate model for emulation and physics extraction of large eddy simulations.  J Amer Statist Assoc, 113(524):1443–1456.
  * McKay et al., [2000] McKay, M. D., Beckman, R. J., and Conover, W. J. (2000).  A comparison of three methods for selecting values of input variables in the analysis of output from a computer code.  Technometrics, 42(1):55–61.
  * Miller and Snellings, [2003] Miller, M. and Snellings, R. (2003).  Eccentricity fluctuations and its possible effect on elliptic flow measurements.  arXiv preprint nucl-ex/0312008.
  * Moreland et al., [2015] Moreland, J. S., Bernhard, J. E., and Bass, S. A. (2015).  Alternative ansatz to wounded nucleon and binary collision scaling in high-energy nuclear collisions.  Phys. Rev. C, 92(1):011901.
  * Novak et al., [2014] Novak, J., Novak, K., Pratt, S., Vredevoogd, J., Coleman-Smith, C., and Wolpert, R. (2014).  Determining fundamental properties of matter created in ultrarelativistic heavy-ion collisions.  Physical Review C, 89(3):034917.
  * Ollitrault, [1992] Ollitrault, J.-Y. (1992).  Anisotropy as a signature of transverse collective flow.  Phys. Rev. D, 46:229–245.
  * Paananen et al., [2019] Paananen, T., Piironen, J., Andersen, M. R., and Vehtari, A. (2019).  Variable selection for Gaussian processes via sensitivity analysis of the posterior predictive distribution.  In International Conference on Artificial Intelligence and Statistics, pages 1743–1752.
  * Paszke et al., [2017] Paszke, A., Gross, S., Chintala, S., Chanan, G., Yang, E., DeVito, Z., Lin, Z., Desmaison, A., Antiga, L., and Lerer, A. (2017).  Automatic Differentiation in PyTorch.  NIPS 2017 Workshop Autodiff.
  * Petersen et al., [2011] Petersen, H., Coleman-Smith, C., Bass, S. A., and Wolpert, R. (2011).  Constraining the initial state granularity with bulk observables in Au+Au collisions at GeV.  Journal of Physics G: Nuclear and Particle Physics, 38(4):045102.
  * Pinkus, [2015] Pinkus, A. (2015).  Ridge Functions, volume 205.  Cambridge University Press.
  * Rasmussen and Ghahramani, [2000] Rasmussen, C. and Ghahramani, Z. (2000).  Occam’s razor.  In Advances in Neural Information Processing Systems. MIT Press.
  * Ray and Szabó , [2021] Ray, K. and Szabó , B. (2021).  Variational Bayes for high-dimensional linear regression with sparse priors.  Journal of the American Statistical Association, 117(539):1270–1281.
  * Sacks et al., [1989] Sacks, J., Welch, W. J., Mitchell, T. J., and Wynn, H. P. (1989).  Design and analysis of computer experiments.  Statistical Science, 4(4):409–423.
  * Sangaline and Pratt, [2016] Sangaline, E. and Pratt, S. (2016).  Toward a deeper understanding of how experiments constrain the underlying physics of heavy-ion collisions.  Physical Review C, 93(2):024908.
  * Santner et al., [2003] Santner, T. J., Williams, B. J., Notz, W. I., and Williams, B. J. (2003).  The Design and Analysis of Computer Experiments.  Springer.
  * Sauer et al., [2022] Sauer, A., Cooper, A., and Gramacy, R. B. (2022).  Vecchia-approximated deep Gaussian processes for computer experiments.  Journal of Computational and Graphical Statistics.  To appear.
  * Scillitoe et al., [2021] Scillitoe, A., Seshadri, P., and Wong, C. (2021).  Polynomial ridge flowfield estimation.  Physics of Fluids, 33(12):127110.
  * Seshadri et al., [2019] Seshadri, P., Yuchi, S., and Parks, G. T. (2019).  Dimension reduction via Gaussian ridge functions.  SIAM/ASA Journal on Uncertainty Quantification, 7(4):1301–1322.
  * Snelson and Ghahramani, [2005] Snelson, E. and Ghahramani, Z. (2005).  Sparse Gaussian processes using pseudo-inputs.  In Advances in Neural Information Processing Systems. MIT Press.
  * Snelson and Ghahramani, [2012] Snelson, E. and Ghahramani, Z. (2012).  Variable noise and dimensionality reduction for sparse Gaussian processes.  arXiv preprint arXiv:1206.6873.
  * Stein, [1999] Stein, M. L. (1999).  Interpolation of Spatial Data: Some Theory for Kriging.  Springer Science & Business Media.
  * Sung et al., [2020] Sung, C.-L., Hung, Y., Rittase, W., Zhu, C., and Wu, C. F. J. (2020).  A generalized gaussian process model for computer experiments with binary time series.  Journal of the American Statistical Association, 115(530):945–956.
  * Titsias and Lawrence, [2010] Titsias, M. and Lawrence, N. D. (2010).  Bayesian Gaussian process latent variable model.  In International Conference on Artificial Intelligence and Statistics, pages 844–851.
  * Tripathy et al., [2016] Tripathy, R., Bilionis, I., and Gonzalez, M. (2016).  Gaussian processes with built-in dimensionality reduction: Applications to high-dimensional uncertainty propagation.  Journal of Computational Physics, 321:191–223.
  * Van de Geer and van de Geer, [2000] Van de Geer, S. A. and van de Geer, S. (2000).  Empirical Processes in M-Estimation, volume 6.  Cambridge University Press.
  * van der Vaart and van Zanten, [2008] van der Vaart, A. W. and van Zanten, J. H. (2008).  Rates of contraction of posterior distributions based on Gaussian process priors.  The Annals of Statistics, 36(3):1435–1463.
  * Wang et al., [2016] Wang, Z., Hutter, F., Zoghi, M., Matheson, D., and De Freitas, N. (2016).  Bayesian optimization in a billion dimensions via random embeddings.  Journal of Artificial Intelligence Research, 55(1):361–387.
  * Welch et al., [1992] Welch, W. J., Buck, R. J., Sacks, J., Wynn, H. P., Mitchell, T. J., and Morris, M. D. (1992).  Screening, predicting, and computer experiments.  Technometrics, 34(1):15–25.
  * Xia, [2008] Xia, Y. (2008).  A multiple-index model and dimension reduction.  Journal of the American Statistical Association, 103(484):1631–1640.
  * Yuchi et al., [2023] Yuchi, H. S., Mak, S., and Xie, Y. (2023).  Bayesian uncertainty quantification for low-rank matrix completion.  Bayesian Analysis, 18(2):491–518.

##  Appendix A Closed-form expressions for $\boldsymbol{\Psi}_{1}^{l}$ and
$\boldsymbol{\Psi}_{2}^{l}$

We provide in the following a detailed derivation of the closed-form
expectations for $\boldsymbol{\Psi}_{1}^{l}=[\phi_{1,i,j}^{l}]$ and
$\boldsymbol{\Psi}_{2}^{l}=[\phi_{2,i,j}^{l}]$, which then provides tractable
bounds for variational inference. The integrals to be computed are:

$\displaystyle\begin{split}&\phi_{1,i,j}^{l}=\int\nu_{l}\exp\left\\{-\frac{||\mathbf{M}_{l}\mathbf{x}_{i}-\mathbf{z}_{j}||^{2}}{2}\right\\}\pi_{\mathcal{MN}}(\mathbf{M}_{l}|\mathbf{A}_{l},\mathbf{H}_{l},\mathbf{U}_{l})d\mathbf{M}_{l}\\\
&\phi_{2,i,j}^{l}=\sum_{k=1}^{n}\int\nu_{l}^{2}\exp\left\\{-\frac{||\mathbf{z}_{i}-\mathbf{M}_{l}\mathbf{x}_{k}||^{2}}{2}\right\\}\exp\left\\{-\frac{||\mathbf{z}_{j}-\mathbf{M}_{l}\mathbf{x}_{k}||^{2}}{2}\right\\}\pi_{\mathcal{MN}}(\mathbf{M}_{l}|\mathbf{A}_{l},\mathbf{H}_{l},\mathbf{U}_{l})d\mathbf{M}_{l},\end{split}$
(21)

where $\pi_{\mathcal{MN}}(\cdot|\mathbf{A},\mathbf{H},\mathbf{U})$ denotes the
density for the matrix normal distribution with mean matrix $\mathbf{A}$, row-
wise covariance matrix $\mathbf{H}$, and column-wise covariance matrix
$\mathbf{U}$.

We will derive these integrals first in the above general setting (i.e., with
general forms for $\mathbf{H}_{l}$ and $\mathbf{U}_{l}$), then show how to
simplify these expression under a simplified variational form with
$\mathbf{H}_{l}=\mathbf{I}_{p\times p}$ and
$\mathbf{U}_{l}=\mathbf{v}_{l}\mathbf{I}_{d\times d}$. For the general
setting, first observe that, since
$\mathbf{M}_{l}\sim\mathcal{MN}(\mathbf{A}_{l},\mathbf{H}_{l},\mathbf{U}_{l})$,
it follows that
$\mbox{vec}(\mathbf{M}_{l})\sim\mathcal{MVN}(\mbox{vec}(\mathbf{A}_{l}),\mathbf{U}_{l}\otimes\mathbf{H}_{l})$.
Thus, the required computation reduces to one involving a multivariate
Gaussian integral. From the first expression in (21), we then have:

$\displaystyle\phi_{1,i,j}^{l}$
$\displaystyle=C_{j}\int\exp\left\\{-\frac{1}{2}\left(\text{tr}(\mathbf{M}_{l}^{T}\mathbf{M}_{l}\mathbf{x}_{i}\mathbf{x}_{i}^{T})-2\text{tr}(\mathbf{M}_{l}^{T}\mathbf{z}_{j}\mathbf{x}_{i}^{T}\right)\right\\}$
$\displaystyle\quad\times\exp\left\\{-\frac{1}{2}\left(\mbox{vec}(\mathbf{M}_{l})^{T}(\mathbf{U}_{l}\otimes\mathbf{H}_{l})^{-1}\mbox{vec}(\mathbf{M}_{l})-2\mbox{vec}(\mathbf{M}_{l})^{T}(\mathbf{U}_{l}\otimes\mathbf{H}_{l})^{-1}\mbox{vec}(\mathbf{A}_{l})\right)\right\\}d\mbox{vec}(\mathbf{M}_{l})$

where
$C_{j}=\nu_{j}\cdot{\exp\Big{(}-\frac{1}{2}\mbox{vec}(\mathbf{A}_{l})^{T}(\mathbf{U}_{l}\otimes\mathbf{H}_{l})^{-1}\mbox{vec}(\mathbf{A}_{l})-\frac{1}{2}\mathbf{z}_{j}^{t}\mathbf{z}_{j}\Big{)}}/((2\pi)^{\frac{pd}{2}}|\mathbf{U}_{l}\otimes\mathbf{H}_{l}|^{\frac{1}{2}})$.
Recall that
$\text{tr}(\mathbf{M}_{l}^{T}\mathbf{M}_{l}\mathbf{x}_{i}\mathbf{x}_{i}^{T})=\mbox{vec}(\mathbf{M}_{l})^{T}(\mathbf{x}_{i}\mathbf{x}_{i}^{T}\otimes\mathbf{I}_{d\times
d})\mbox{vec}(\mathbf{M}_{l})$ and
$\text{tr}(\mathbf{M}_{l}^{T}\mathbf{z}_{j}\mathbf{x}_{i}^{T})=\mbox{vec}(\mathbf{M}_{l})^{T}\mbox{vec}(\mathbf{x}_{j}\mathbf{x}_{i}^{T})$.
We can now complete the square, such that:

$\displaystyle\begin{split}\phi_{1,i,j}^{l}&=\nu\cdot\frac{\exp\left\\{-\frac{1}{2}\mbox{vec}(\mathbf{A}_{l})^{T}(\mathbf{U}_{l}\otimes\mathbf{H}_{l})^{-1}\mbox{vec}(\mathbf{A}_{l})-\frac{1}{2}\mathbf{z}_{j}^{t}\mathbf{z}_{j}\right\\}}{(2\pi)^{\frac{pd}{2}}|\mathbf{U}_{l}\otimes\mathbf{H}_{l}|^{\frac{1}{2}}|(\mathbf{x}_{i}\mathbf{x}_{i}^{T}\otimes\mathbf{I}_{p\times
p})+(\mathbf{U}_{l}\otimes\mathbf{H}_{l})^{-1}|^{\frac{1}{2}}}\\\
&\quad\times\exp\left\\{\frac{1}{2}(\mbox{vec}(\mathbf{z}_{j}\mathbf{x}_{i}^{T})+(\mathbf{U}_{l}\otimes\mathbf{H}_{l})^{-1}\mbox{vec}(\mathbf{A}_{l}))^{T}\big{(}(\mathbf{x}_{i}\mathbf{x}_{i}^{T}\otimes\mathbf{I}_{p\times
p})+(\mathbf{U}_{l}\otimes\mathbf{H}_{l})^{-1}\big{)}^{-1}\right.\\\
&\left.\quad\times(\mbox{vec}(\mathbf{z}_{j}\mathbf{x}_{i}^{T})+(\mathbf{U}_{l}\otimes\mathbf{H}_{l})^{-1}\mbox{vec}(\mathbf{A}_{l}))\right\\}.\end{split}$
(22)

This provides a closed-form expression for computing the entries of
$\boldsymbol{\Psi}_{1}^{l}$.

We can perform the exact same procedure for $\phi_{2,i,j}^{l}$ with little
modification. From the second expression in (21), define the $k$-th component
as:

$\displaystyle\begin{split}&\phi_{2,i,j,k}^{l}=\int\nu_{l}^{2}\exp\left\\{-\frac{||\mathbf{z}_{i}-\mathbf{M}_{l}\mathbf{x}_{k}||^{2}}{2}\right\\}\exp\left\\{-\frac{||\mathbf{z}_{j}-\mathbf{M}_{l}\mathbf{x}_{k}||^{2}}{2}\right\\}\pi_{\mathcal{MN}}(\mathbf{M}_{l}|\mathbf{A}_{l},\mathbf{H}_{l},\mathbf{U}_{l})d\mathbf{M}_{l}\\\
&=\nu_{l}^{2}\cdot\frac{\exp\left\\{-\frac{1}{2}\mbox{vec}(\mathbf{A}_{l})^{T}(\mathbf{U}_{l}\otimes\mathbf{H}_{l})^{-1}\mbox{vec}(\mathbf{A}_{l})-\frac{1}{2}(\mathbf{z}_{j}+\mathbf{z}_{i})^{T}(\mathbf{z}_{j}+\mathbf{z}_{i})\right\\}}{(2\pi)^{\frac{pd}{2}}|\mathbf{U}_{l}\otimes\mathbf{H}_{l}|^{\frac{1}{2}}|(2\mathbf{x}_{k}\mathbf{x}_{k}^{T}\otimes\mathbf{I}_{p\times
p})+(\mathbf{U}_{l}\otimes\mathbf{H}_{l})^{-1}|^{\frac{1}{2}}}\\\
&\times\exp\left\\{\frac{1}{2}(\mbox{vec}((\mathbf{z}_{j}+\mathbf{z}_{i})\mathbf{x}_{k}^{T})+(\mathbf{U}_{l}\otimes\mathbf{H}_{l})^{-1}\mbox{vec}(\mathbf{A}_{l}))^{T}\big{(}(2\mathbf{x}_{k}\mathbf{x}_{k}^{T}\otimes\mathbf{I}_{p\times
p})+(\mathbf{U}_{l}\otimes\mathbf{H}_{l})^{-1}\big{)}^{-1}\right.\\\
&\left.\times(\mbox{vec}((\mathbf{z}_{j}+\mathbf{z}_{i})\mathbf{x}_{k}^{T})+(\mathbf{U}_{l}\otimes\mathbf{H}_{l})^{-1}\mbox{vec}(\mathbf{A}_{l}))\right\\}.\end{split}$
(23)

Together with (21), this provides a closed-form expression for computing the
entries of $\boldsymbol{\Psi}_{2}^{l}$.

One disadvantage with the above general setting (i.e., with general forms for
$\mathbf{H}_{l}$ and $\mathbf{U}_{l}$) is that the closed-form expressions
(22) and (23) require the inverting of $pd\times pd$ matrices, which requires
$\mathcal{O}(p^{3}d^{3})$ work and thus can be costly. We show next that we
can greatly reduce this computational complexity if we restrict the
variational form to $\mathbf{U}_{l}=\mathbf{v}_{l}\mathbf{I}_{d\times
d},\mathbf{H}_{l}=\mathbf{I}_{p\times p}$. From this restricted variational
form, one can show that:

$\displaystyle\phi_{1,i,j}^{l}$
$\displaystyle=C_{j}\int\exp\left\\{-\frac{1}{2}\big{(}\mbox{tr}(\mathbf{M}_{l}^{T}\mathbf{M}_{l}\mathbf{x}_{i}\mathbf{x}_{i}^{T})-2\mbox{tr}(\mathbf{M}_{l}^{T}\mathbf{z}_{j}\mathbf{x}_{i}^{T}\big{)}\big{)}\right\\}$
$\displaystyle\quad\times\exp\left\\{-\frac{1}{2}\big{(}\mbox{tr}(\mathbf{M}_{l}^{T}\mathbf{M}_{l}\mathbf{U}_{l}^{-1})-2\mbox{tr}(\mathbf{M}_{l}^{T}\mathbf{A}_{l}\mathbf{U}_{l}^{-1})\right\\}d\mathbf{M}_{l},$

where
$C_{j}={\exp\big{(}-\frac{1}{2}\mbox{tr}(\mathbf{U}_{l}^{-1}\mathbf{A}_{l}^{T}\mathbf{A}_{l})-\frac{1}{2}\mathbf{z}_{j}\mathbf{z}_{j}^{T}\big{)}}/({|\mathbf{U}_{l}|^{\frac{d}{2}}2\pi^{\frac{pd}{2}}})$.
We can then use the matrix-normal version of completing the square to obtain:

$\displaystyle\phi_{1,i,j}^{l}=\frac{C_{j}}{|\mathbf{U}_{l}^{-1}+\mathbf{x}_{i}\mathbf{x}_{i}^{T}|^{\frac{d}{2}}}\exp\left\\{\frac{1}{2}\mbox{tr}\big{(}(\mathbf{z}_{j}\mathbf{x}_{i}^{T}+\mathbf{A}_{l}\mathbf{U}_{l}^{-1})^{T}(\mathbf{z}_{j}\mathbf{x}_{i}^{T}+\mathbf{A}_{l}\mathbf{U}_{l}^{-1})(\mathbf{U}_{l}^{-1}+\mathbf{x}_{i}\mathbf{x}_{i}^{T})^{-1}\big{)}\right\\}.$
(24)

The above closed-form expression for $\phi_{1,i,j}^{l}$ (under the restricted
variational form) can be evaluated in significantly less work that the earlier
$\mathcal{O}(p^{3}d^{3})$. First, note that $C_{j}$ can be computed (for all
$j$) in $\mathcal{O}(d^{2}p+mp)$ work. Next, we can make use of careful linear
algebra simplifications to avoid expensive matrix inversions and
multiplications in the exponent term in (24). Note that the below expression
admit rank-one updates:

$\displaystyle(\mathbf{U}^{-1}+\mathbf{x}_{i}\mathbf{x}_{i}^{T})^{-1}=\mathbf{U}-\frac{\mathbf{U}\mathbf{x}_{i}\mathbf{x}_{i}^{T}\mathbf{U}}{1+\mathbf{x}_{i}^{T}\mathbf{U}\mathbf{x}_{i}},\quad|\mathbf{U}^{-1}+\mathbf{x}_{i}\mathbf{x}_{i}^{T}|=|\mathbf{U}^{-1}|(1+\mathbf{x}_{i}^{T}\mathbf{U}\mathbf{x}_{i}).$

Thus, these quantities can be computed in $\mathcal{O}(d)$ because
$\mathbf{U}$ is diagonal. Now examining the exponent term in (24), we can
distribute the product and re-arrange using the permutation invariance of
$\text{tr}(\cdot)$ to obtain:

$\displaystyle\mbox{tr}\big{(}(\mathbf{z}_{j}\mathbf{x}_{i}^{T}+\mathbf{A}_{l}\mathbf{U}_{l}^{-1})^{T}(\mathbf{z}_{j}\mathbf{x}_{i}^{T}+\mathbf{A}_{l}\mathbf{U}_{l}^{-1})(\mathbf{U}_{l}-\frac{\mathbf{U}_{l}\mathbf{x}_{i}\mathbf{x}_{i}^{T}\mathbf{U}_{l}}{1+\mathbf{x}_{i}^{T}\mathbf{U}_{l}\mathbf{x}_{i}})\big{)}$
$\displaystyle=\mbox{tr}(\mathbf{z}_{j}^{T}\mathbf{z}_{j}\mathbf{x}_{i}^{T}\mathbf{U}_{l}\mathbf{x}_{i})-\frac{\mbox{tr}(\mathbf{z}_{j}^{T}\mathbf{z}_{j}\mathbf{x}_{i}^{T}\mathbf{U}_{l}\mathbf{x}_{i}\mathbf{x}_{i}^{T}\mathbf{U}_{l}\mathbf{x}_{i})}{1+\mathbf{x}_{i}^{T}\mathbf{U}_{l}\mathbf{x}_{i}}$
$\displaystyle\quad+2\mbox{tr}(\mathbf{z}_{j}^{T}\mathbf{A}_{l}\mathbf{x}_{i})-\frac{\mbox{tr}(2\mathbf{z}_{j}^{T}\mathbf{A}_{l}\mathbf{x}_{i}\mathbf{x}_{i}^{T}\mathbf{U}_{l}\mathbf{x}_{i})}{1+\mathbf{x}_{i}^{T}\mathbf{U}_{l}\mathbf{x}_{i}}$
$\displaystyle\quad+\mbox{tr}(\mathbf{U}_{l}^{-1}\mathbf{A}_{l}^{T}\mathbf{A}_{l})-\frac{\mbox{tr}(\mathbf{x}_{i}^{T}\mathbf{A}_{l}^{T}\mathbf{A}_{l}\mathbf{x}_{i})}{1+\mathbf{x}_{i}^{T}\mathbf{U}_{l}\mathbf{x}_{i}}.$

All these terms then decomposes into the product of scalars resulting from
vector-matrix products. Combining these steps, we then see that
$\boldsymbol{\Psi}_{1}^{l}$ can be computed in $\mathcal{O}(mnpd)$ work under
the above restricted variational form.

The calculation for $\boldsymbol{\Psi}_{2}^{l}$ proceeds in an analogous
manner. We write

$\displaystyle\phi_{2,i,j,k}^{l}=\frac{C_{j}}{|\mathbf{U}_{l}^{-1}+2\mathbf{x}_{k}\mathbf{x}_{k}^{T}|^{\frac{d}{2}}}\exp\left\\{\frac{1}{2}\text{tr}\big{(}((\mathbf{z}_{i}+\mathbf{z}_{j})\mathbf{x}_{k}^{T}+\mathbf{A}_{l}\mathbf{U}_{l}^{-1})^{T}((\mathbf{z}_{j}+\mathbf{z}_{i})\mathbf{x}_{k}^{T}+\mathbf{A}_{l}\mathbf{U}_{l}^{-1})(\mathbf{U}_{l}^{-1}+2\mathbf{x}_{k}\mathbf{x}_{k}^{T})^{-1}\big{)}\right\\}.$

A similar decomposition (as in $\phi_{1,i,j}^{l}$) of the above exponent term
can be performed, allowing $\boldsymbol{\Psi}_{2}^{l}$ to be computed in
$\mathcal{O}(nm^{2}pd)$ work if $\boldsymbol{\Psi}_{1}^{l}$ has already been
computed.

##  Appendix B Derivation of ELBO bound

We show next the detailed derivation of the variational lower bound (13).
Recall that, for the AdMIn-GP, its generative process can be written as:

$\displaystyle\mathbf{y}|\mathbf{f}_{1}\dots\mathbf{f}_{L}\sim\mathcal{MVN}\left(\sum_{l=1}^{L}\mathbf{f}_{l},\beta^{-1}\mathbf{I}_{n\times
n}\right),$
$\displaystyle\mathbf{f}_{l}|\mathbf{u}_{l},\mathbf{M}_{l}\sim\mathcal{MVN}(\mathbf{K}_{n,m}^{l}(\mathbf{K}_{m,m}^{l})^{-1}\mathbf{u}_{l},\mathbf{K}_{n,n}^{l}+\mathbf{K}_{n,m}^{l}(\mathbf{K}_{m,m}^{l})^{-1}\mathbf{K}_{m,n}^{l}),\quad
l=1,\dots,L,$
$\displaystyle\mathbf{u}_{l}\sim\mathcal{MVN}(\mathbf{0},(\mathbf{K}_{m,m}^{l})^{-1}),\quad
l=1,\dots,L,$ $\displaystyle\mathbf{M}_{l}\sim\mathcal{MN}(\mathbf{0}_{p\times
d},\mathbf{I}_{p\times p},\mathbf{v}_{l}\mathbf{I}_{d\times d}),\quad
l=1,\dots,L.$

The variational approximation made is:

$\displaystyle
p(\mathbf{f}_{1}\dots\mathbf{f}_{L},\mathbf{u}_{1}\dots\mathbf{u}_{L},\mathbf{M}_{1}\dots\mathbf{M}_{L}|\mathbf{y})\approx\phi(\mathbf{u})\prod_{l=1}^{L}p(\mathbf{f}_{l}|\mathbf{M}_{l},\mathbf{u}_{l})\phi(\mathbf{M}_{l}),$

where, as mentioned in the main paper, we choose variational distributions of
the form:

$\displaystyle\phi(\mathbf{u})\sim\mathcal{MVN}(\mathbf{u}|\mathbf{D}\mathbf{W}^{-1}\bar{\mathbf{y}},\mathbf{D}\mathbf{W}^{-1}\mathbf{D})$
$\displaystyle\phi_{l}(\mathbf{M}_{l})\sim\mathcal{MN}(\mathbf{A}_{l},\mathbf{I}_{p\times
p},\mathbf{I}_{d\times d}\mathbf{v}_{l}).$

With this, we can then write out the ELBO bound, then simplify via the chosen
variational distribution on $\phi(\mathbf{u})$:

$\displaystyle log\left\\{p(\mathbf{y})\right\\}$
$\displaystyle\geq\int\prod_{l=1}^{L}p(\mathbf{f}_{l}|\mathbf{u}_{l},\mathbf{M}_{l})\phi(\mathbf{u}_{l})\phi(\mathbf{M}_{l})\log\Big{(}\frac{p(\mathbf{y}|\mathbf{f}_{1}\dots\mathbf{f}_{L})\prod_{l=1}^{L}p(\mathbf{f}_{l}|\mathbf{u}_{l},\mathbf{M}_{l})p(\mathbf{u}_{l})p(\mathbf{M}_{l})}{\prod_{l=1}^{L}p(\mathbf{f}_{l}|\mathbf{u}_{l},\mathbf{M}_{l})\phi(\mathbf{u}_{l})\phi(\mathbf{M}_{l})}\Big{)}d\mathbf{f}d\mathbf{M}d\mathbf{u}$
$\displaystyle=\int\phi(\mathbf{u})\phi(\mathbf{M})\Big{(}\int\prod_{l=1}^{L}p(\mathbf{f}_{l}|\mathbf{u}_{l},\mathbf{M}_{l})\log\big{(}p(\mathbf{y}|\mathbf{f}_{1}\dots\mathbf{f}_{L})\big{)}d\mathbf{f}+\log\left(\frac{p(\mathbf{u})}{\phi(\mathbf{u})}\right)\Big{)}d\mathbf{M}d\mathbf{u}$
$\displaystyle\quad\quad-\sum_{l=1}^{L}KL(\phi(\mathbf{M}_{l})||p(\mathbf{M}_{l}))$
$\displaystyle=\log\Bigg{(}\int
p(\mathbf{u})\exp\left\\{\mathbb{E}_{\phi(M)}\left[\log\big{(}N(\mathbf{y}|\sum_{l=1}^{L}\boldsymbol{\alpha}_{l},\beta^{-1}\mathbf{I}_{n\times
n})\big{)}\right]\right\\}d\mathbf{u}\Bigg{)}$
$\displaystyle\quad\quad-\sum_{l=1}^{L}\frac{\beta}{2}\Big{(}\nu_{l}n-\mbox{tr}\big{(}(\mathbf{K}_{m,m}^{l})^{-1}\mathbb{E}_{\phi(\mathbf{M})}(\mathbf{K}_{m,n}^{l}\mathbf{K}_{n,m}^{l})\Big{)}-\sum_{l=1}KL(\phi(\mathbf{M}_{l})||p(\mathbf{M}_{l})),$

where
$\boldsymbol{\alpha}_{l}=\mathbf{K}_{n,m}^{l}(\mathbf{K}_{m,m}^{l})^{-1}\mathbf{u}_{l}$.

The key integral in the above expression is

$\displaystyle\int
p(\mathbf{u})\exp\left\\{\mathbb{E}_{\phi(\mathbf{M})}\left[\log\big{(}N(\mathbf{y}|\sum_{l=1}^{L}\boldsymbol{\alpha}_{l},\beta^{-1}\mathbf{I}_{n\times
n})\big{)}\right]\right\\}d\mathbf{u}$
$\displaystyle=\left(\frac{\beta}{2\pi}\right)^{\frac{n}{2}}\exp\left\\{-\frac{\beta}{2}\mathbf{y}^{T}\mathbf{y}\right\\}$
$\displaystyle\times\int
p(\mathbf{u})\exp\left\\{-\frac{\beta}{2}\left(\sum_{l=1}^{L}\sum_{s=1}^{L}\mathbf{u}_{l}(\mathbf{K}_{m,m}^{l})^{-1}\mathbb{E}_{\phi(\mathbf{M})}\left[(\mathbf{K}_{m,n}^{l}\mathbf{K}_{n,m}^{s})\right](\mathbf{K}_{m,m}^{s})^{-1}\mathbf{u}_{s}-2\sum_{l=1}\mathbf{y}^{T}E_{\phi(\mathbf{M}_{l})}\left[(\mathbf{K}_{n,m}^{l})(\mathbf{K}_{m,m}^{l})^{-1}\mathbf{u}_{l}\right]\right)\right\\}d\mathbf{u}$
$\displaystyle=\left(\frac{\beta}{2\pi}\right)^{\frac{n}{2}}\exp\left(-\frac{\beta}{2}\mathbf{y}^{T}\mathbf{y}\right)\int
p(\mathbf{u})\exp\left(-\frac{\beta}{2}\mathbf{u}^{T}\mathbf{D}^{-1}\mathbf{P}\mathbf{D}^{-1}\mathbf{u}-\frac{\beta}{2}\bar{\mathbf{y}}^{T}\mathbf{D}^{-1}\mathbf{u}\right)d\mathbf{u}$

Finally, recalling that
$p(\mathbf{u})\sim\mathcal{MVN}(\mathbf{0},\mathbf{D})$, this reduces to a
simple multivariate Gaussian integral, and we can then write the final
variational bound as:

$\displaystyle\log(p(\mathbf{y}))$
$\displaystyle\geq\log\left(\frac{\beta^{\frac{n}{2}}\prod_{l=1}^{L}|\mathbf{K}_{m,m}^{l}|^{\frac{1}{2}}}{(2\pi)^{\frac{n}{2}}|\beta\mathbf{D}+\mathbf{P}|^{\frac{1}{2}}}\exp\left(\frac{1}{2}\beta^{2}\bar{\mathbf{y}}^{T}\mathbf{W}^{-1}\bar{\mathbf{y}})\right)\exp\left(-\frac{1}{2}\beta\mathbf{y}^{T}\mathbf{y}\right)\right)$
$\displaystyle\quad\quad-\sum_{l=1}^{L}KL(\phi_{l}(\mathbf{M}_{l})||p(\mathbf{M}_{l}))-\sum_{l=1}^{L}\mathbf{V}_{l},$

where
$\mathbf{V}_{l}=({\beta}/{2})(\nu_{l}n-\text{tr}\\{(\mathbf{K}_{m,m}^{l})^{-1}\boldsymbol{\Psi}_{2}^{l}\\})$
and $\mathbf{W}=\beta\mathbf{P}+\mathbf{D}$.