# Mapping Interstellar Dust with Gaussian Processes

acm, la, bl, jpc, dh, dmb  Andrew C. Miller
<EMAIL_ADDRESS>Work done while at the Data Science Institute at Columbia
University. Lauren Anderson
<EMAIL_ADDRESS>Boris Leistedt
<EMAIL_ADDRESS>John P. Cunningham
<EMAIL_ADDRESS>David W. Hogg
<EMAIL_ADDRESS>David M. Blei
<EMAIL_ADDRESS>
###### Abstract

Interstellar dust corrupts nearly every stellar observation, and accounting
for it is crucial to measuring physical properties of stars. We model the dust
distribution as a spatially varying latent field with a Gaussian process (GP)
and develop a likelihood model and inference method that scales to millions of
astronomical observations. Modeling interstellar dust is complicated by two
factors. The first is integrated observations. The data come from a vantage
point on Earth and each observation is an integral of the unobserved function
along our line of sight, resulting in a complex likelihood and a more
difficult inference problem than in classical GP inference. The second
complication is scale; stellar catalogs have millions of observations. To
address these challenges we develop ziggy, a scalable approach to GP inference
with integrated observations based on stochastic variational inference. We
study ziggy on synthetic data and the Ananke dataset, a high-fidelity
mechanistic model of the Milky Way with millions of stars. ziggy reliably
infers the spatial dust map with well-calibrated posterior uncertainties.

Keywords: Large-scale astronomical data, Gaussian processes, kernel methods,
scalable Bayesian inference, variational methods, machine learning

##  1 Introduction

The Milky Way galaxy is primarily comprised of dark matter, stars, and gas.
Within the gas, in its densest and coldest regions, dust particles form. The
stars of the Milky Way are embedded in this field of dust.

Back on Earth, astronomers try to map the stars, measuring the location,
apparent brightness, and color of each. But the dust between Earth and a star
obscures the light, corrupting the astronomers’ observation. Relative to the
star’s true brightness and color, the dust dims the brightness and reddens the
color. This corruption is called extinction, and it complicates inferences
about a star’s true distance and other true properties. Stellar extinction has
hindered many studies of stars in the Milky Way disk, which is where most dust
lies, and therefore the largest extinctions (Mathis, 1990).

To cope with these corruptions, astronomers need a map of interstellar dust,
an estimate of the density of dust at each location in the Galaxy. An accurate
dust map could be used to correct our astronomical measurements and sharpen
our knowledge of the Galaxy’s stars. Moreover, a dust map may be of
independent scientific interest—for example, it may reveal macroscopic
properties of the shape of the Milky Way, such as spiral arms.111Though hints
of spiral structures have been inferred from other observations (Georgelin and
Georgelin, 1976; Chen et al., 2019), whether the Milky Way galaxy is a grand
design spiral or if the spiral structures are more flocculent remains an open
question. An accurate dust map will shed light on the recent dynamical history
of the Milky Way.

But constructing such a map is difficult. We are embedded in our own dust
field, and so we cannot directly observe it. Rather, we can only observe a
noisy _integral_ of the field along the line of sight between Earth and a star
— the starlight extinction. (Here, the line of sight is the straight line
through space between the star and the earth.) Thus the inference problem is
to calculate a single, coherent spatial field of dust that can explain the
observed extinction of millions of spatially distributed stars. In this paper,
we use a large data set of astronomical measurements to infer a three-
dimensional map of interstellar dust.

The data comes from standard practice in astronomy, which is to estimate the
extinction of an individual star from its observed color and brightness using
a physical model of star formation and evolution.222The observed color and
brightness are directly derived from telescope images (e.g. photometry).
Extinction measurements are backed out from a theoretical distribution of
dust-free star colors and luminosities, where luminosity is the intrinsic
brightness of a star. This distribution can be based on isochrones or
empirically derived from a region of the sky known to have no dust (e.g. the
Milky Way halo). The extinction corresponds to how far the observed color and
brightness are from the set of theoretically plausible colors and
luminosities. The extinction uncertainty incorporates both noise in the
photometric measurement and prior uncertainty over the range of plausible
colors and luminosities.  This procedure results in an estimate of the
extinction and an approximate variance of the estimate about the true
extinction value. However, this estimate is derived from a single star’s color
and brightness—it ignores information about the spatial structure of the Milky
Way and the fact that all stars are observed through the same three-
dimensional density of dust.

The dust map is a spatial distribution of dust—an unobserved spatial
function—that can help refine these noisy measurements. Each measurement is
modeled conditional on the function and the three-dimensional location of the
star.333The spatial locations of stars can be derived from parallax
measurements. In this work we consider them fixed, but in future work will
additionally consider their uncertainty. Such a model shrinks estimates toward
their true values and reduces uncertainty about them. More formally, the
unknown dust map is a function ${\rho:\mathbb{R}^{D}\mapsto\mathbb{R}}$ from a
three-dimensional location in the universe $x$ to the density of dust at that
location. Our goal is to estimate this function from data.

We take a Bayesian nonparametric approach. We place a Gaussian process (GP)
prior on the dust map and posit a likelihood function for how astronomical
observations arise from it. We develop a scalable approximate posterior
inference algorithm to estimate the posterior dust map from large-scale
astronomical data. In addition to the scale of the data, the main challenge is
the likelihood function. Typical spatial analysis involves data that are noisy
evaluations of an unobserved function, and the likelihood is simple.
Astronomical data, however, comes from a limited vantage point—we only observe
the latent dust map as an _integrated process_ along a line of site to a star,
and this integral is baked into the likelihood of the observation.

More formally, let $e_{n}$ be the true extinction of starlight for star $n$ at
location $x_{n}$, let $a_{n}$ be the noisily measured extinction with
uncertainty $\sigma_{n}$. Given the unknown dust map, we model the noisy
measurement as a Gaussian whose mean is an integral from Earth to $x_{n}$
(Rezaei Kh. et al., 2017). With covariance function
$k_{{\boldsymbol{\mathbf{\theta}}}}(\cdot,\cdot)$, the full model is

$\displaystyle\rho(\cdot)$ $\displaystyle\sim GP(0,k_{\theta}(\cdot,\cdot))$
(1) $\displaystyle a_{n}$
$\displaystyle\sim\mathcal{N}\left(e_{n},\sigma_{n}^{2}\right),\quad\textrm{where
\,}e_{n}\triangleq\int_{x\in R_{n}}\rho(x)dx.$ (2)

Here $R_{n}$ is the set of points (i.e. the ray) from Earth to $x_{n}$; the
extinction $e_{n}$ is in an integral of the latent dust map $\rho(\cdot)$, and
it is through this integral that the map enters the likelihood. Figure 1
graphically depicts the distinction between pointwise and integrated
observations.

The data are $N$ locations, extinctions, and measurement errors, denoted
$\mathcal{D}=\\{a_{n},x_{n},\sigma_{n}^{2}\\}_{n=1}^{N}$. Conditional on the
data, we want to calculate the posterior dust map $\pi(\rho\mid\mathcal{D})$.
This posterior can estimate different properties of the latent dust, e.g., the
density of dust $\rho(\cdot)$ at a new location, the integral of dust
$\rho(\cdot)$ over new sets, and posterior uncertainty about these values.
Such inferences can aid many stellar studies.

But the posterior is difficult to compute, complicated by both the integrated
likelihood and the large scale of the data. (Modern catalogs of astronomical
data contain observations of millions of stars (Brown et al., 2018; Aguado et
al., 2019).) In theory the scale is helpful—each star provides information
about the unobserved dust map. But scaling a Gaussian process to millions of
observations is a significant computational challenge.

We overcome these challenges with a scalable algorithm for Gaussian process
inference with integrated observations, which we call ziggy.444ziggy is named
for Ziggy Stardust, David Bowie’s alter ego. In particular, ziggy is a
stochastic variational inference algorithm, one that uses stochastic
optimization, inducing points, and variational inference to approximate the
posterior. It handles general covariance functions and scales to millions of
data points. We study ziggy on both synthetic data and the Ananke data, which
comes from a high-fidelity mechanistic model of the Milky Way. We find that
ziggy accurately reconstructs the dust map and accuracy continues to improve
as the number of observations grows above a million.

In our applied setting, the integrated observations introduce technical
challenges. ziggy builds upon an existing scalable Bayesian inference
framework (Hoffman et al., 2013; Hensman et al., 2013), and introduces
additional approximations and computational techniques to address the
challenges created by integrated observations.

The paper is organized as follows. Section 2 describes related work for
estimating interstellar dust. Section 3.1 formally sets up the problem and
describes exact inference in GP models with integrated observations. Section 4
develops a stochastic variational inference algorithm that scales to millions
of stellar observations. Appendix F studies ziggy on a synthetic two-
dimensional example, comparing various settings of the algorithm that trade
computation for accuracy and flexibility. Section 5 studies ziggy with the
Ananke data set, showing that it recovers a well-calibrated three-dimensional
dust map that accurately predicts extinctions and global dust structure.
Section 6 concludes the paper and discusses directions for future research.

(a) Pointwise

(b) Integrated

Figure 1: Reconstructing an unobserved function from pointwise (left) and
integrated (right) observations. The latent function governs noisy pointwise
observations, $y_{n}$, and integrated observations, $a_{n}$. The task is to
reconstruct the unobserved function $\rho(x)$, depicted by the grey contours,
everywhere in the domain. Our perspective is limited to the origin (red dot);
our observations of this process are integrated along a compact set (i.e. a
ray).

##  2 Related research

This work builds on a foundation of research in both astronomy and statistics.

Estimating the latent dust map. The seminal work of Schlegel et al. (1998)
estimates the two-dimensional map of dust across the full sky using dust
emission (rather than dust absorption, which we use here). Emission is a more
direct estimate of the dust, but doesn’t allow for direct inferences of 3D
structure. Stars within the Milky Way, however, are embedded in the three-
dimensional dust field. Estimating per-star extinctions requires
characterizing the dust field as a function of distance, motivating the
construction of three-dimensional dust maps.

More recent approaches use hierarchical spatial models of noisy integrated
observations to reconstruct three-dimensional maps. These approaches, however,
rely on _discretizing space_ and modeling integrated observations as finite
sums. Various discretization strategies have been proposed, each with
different computational demands.

One approach models discretized lines of sight to every star (Rezaei Kh. et
al., 2017), but this approach can only accommodate a few thousand stellar
observations. Another approach uses a fine mesh of cubic voxels to describe
the dust distribution. This introduces a different computational trade off —
larger voxels introduce unrealistic spatial artifacts (Green et al., 2018,
2019), while smaller voxels significantly increase the computational burden
(Leike and Enßlin, 2019). Moreover, the theoretical length scale of the dust
distribution is small; discretizing three-dimensional space with fine enough
resolution to capture this small length scale will have to rely on additional
approximations to scale to the entire volume of the Milky Way. The approach we
develop here avoids discretization, works directly with continuous space, and
scales to millions of observations.

Spatial statistics and Gaussian processes. The field of spatial statistics has
developed many tools to estimate unobserved functions from noisy measurements
(Cressie, 1992). One ubiquitous and fundamental method is Gaussian process
regression or _kriging_ , which interpolates or smooths a latent function
given noisy observation at nearby locations (Krige, 1951; Matheron, 1963,
1973; Cressie, 1990; Rasmussen and Williams, 2006). A GP defines a probability
distribution over the unobserved function that encodes prior assumptions about
some properties — continuity, smoothness, and amplitude — while remaining
flexible. GPs admit analytically tractable inference routines in typical
settings, making them a useful prior for unobserved functions.

Scaling GPs to massive datasets is a more recent challenge. One approach is to
approximate the GP with inducing point methods (Quiñonero-Candela and
Rasmussen, 2005; Snelson and Ghahramani, 2006), where process values at
specific places in the space, the inducing points, are learned from the data,
and predictions are then produced using the inducing point process values.
Inference using this approximation scales better than exact GP inference.
Bolstering their use, recent theoretical work has characterized the error of
inducing point GP approximations (Burt et al., 2019). However, inducing point
methods still require analyzing the entire dataset simultaneously; this
requirement limits their applicability to datasets where the number of
observations is in the millions or billions.

Another thread of GP research in spatial statistics embeds them in more
complicated models and scales them (Cressie, 1992; Banerjee et al., 2014). For
example, nearest-neighbor GP’s are a scalable approximation for estimating the
posterior of a latent spatial field (Datta et al., 2016). Similarly, the
integrated nested Laplace approximations (INLA) framework is an approximate
inference methodology that computes posterior marginal uncertainty in latent
Gaussian models, including Gaussian process models with computationally
tractable precision matrices (Rue et al., 2009). Several scalable Gaussian
process approximations are reviewed in Heaton et al. (2019).

Gaussian processes have also been used within astronomy applications beyond
building spatial maps of interstellar dust. Methods to scale one-dimensional
GP regression to millions of observations have been developed and deployed
with success (Ambikasaran et al., 2014; Foreman-Mackey et al., 2017).

Scalable Bayesian inference. We build on methods for scaling Bayesian
inference to massive datasets. Variational methods (Jordan et al., 1999;
Wainwright et al., 2008; Blei et al., 2017) are a computationally efficient
alternative to Monte Carlo methods for posterior approximation. Variational
inference treats approximate inference as an optimization problem, fitting a
parameterized family to be close to the exact posterior. Stochastic
optimization (Robbins and Monro, 1951) scales variational inference to large
datasets, iteratively subsampling from the data to produce cheap, noisy
gradients of the objective; this strategy is called stochastic variational
inference (SVI) (Hoffman et al., 2013). Building off of SVI, stochastic
variational Gaussian processes (SVGP) bring in inducing point methods to scale
Gaussian processes to massive datasets (Hensman et al., 2013). We build on the
SVGP framework here, adapting it to fit the GP model of astronomical data.

##  3 A Gaussian process model of starlight extinctions

The data are $N$ stars, each one a tuple of its location $x_{n}$, a noisy
measurement of the extinction $a_{n}$, and the variance of the observation
$\sigma^{2}_{n}$; denote the data set
$\mathcal{D}\triangleq\\{x_{n},a_{n},\sigma^{2}_{n}\\}_{n=1}^{N}$. Given the
latent dust map $\rho(x)$, the likelihood of each observation $a_{n}$ is
defined in Equation 2, where the region of integration $R_{n}$ is the ray that
originates at the origin $O$ (i.e. the Earth) and ends at the spatial location
of the star, $x_{n}$,

$\displaystyle R_{n}$ $\displaystyle=\\{\alpha\cdot O+(1-\alpha)\cdot
x_{n}:\alpha\in[0,1]\\}\,.$ (3)

To complete the model, Equation 1 places a Gaussian process prior on the dust
map.

Given the data, the posterior distribution, $p(\rho\,|\,\mathcal{D})$,
summarizes the evidence about the latent dust map. Via the relationship
between $\rho(x_{n})$ and $e_{n}$ defined in Equation 2, we use this posterior
over $\rho$ to form estimates of the extinction for observed stars $e_{n}$,
the extinction at new locations $e_{*}$, and the dust map itself at new
locations $\rho(x_{*})$. Specifically, the compute the posterior expectations

$\displaystyle\hat{\rho}(x_{*})\triangleq\mathbb{E}\left[\rho(x^{*})\,|\,\mathcal{D}\right]\,,\quad\hat{e}_{n}\triangleq\mathbb{E}\left[e_{n}\,|\,x_{n},\mathcal{D}\right]\,,\quad\hat{e}_{*}\triangleq\mathbb{E}\left[e_{*}\,|\,x_{*},\mathcal{D}\right]\,.$
(4)

The posterior variance, e.g.,
$\mathbb{V}\left[e_{*}\,|\,x_{*},\mathcal{D}\right]$, describes the
uncertainty of the estimates.

The posterior dust map $\rho(\cdot)$ synthesizes information from nearby
sources to de-noise or shrink an individual extinction $e_{n}$, resulting in
more accurate inferences with smaller estimator variance. Moreover it enables
estimating the dust density at a new location $x_{*}$ or to estimate estimate
the extinction along a path to the new point.

In the following sections we discuss how to calculate the posterior and how to
approximate it with large data sets of stellar observations. Before that,
however, we discuss the choice of a Gaussian process prior.

Why use a Gaussian process? First, Gaussian processes can flexibly model
complex unobserved latent functions. The dust map $\rho(\cdot)$ is not easily
described by a parametric functional form, but GPs can adapt highly complex,
non-linear functions to describe data. Second, high-level properties of the
unobserved dust map can be captured by the choice of the GP’s covariance
function. In particular, we assume the dust map has specific structure—it is
continuous and two nearby locations are more likely to have similar dust
density values than two distant locations. GPs naturally capture such spatial
coherence. Third, all stellar observations are derived from the same
underlying dust map $\rho(x)$. Two noisy extinction measurements with nearby
integration paths must be coherently described by a single estimate of
$\rho(\cdot)$. The continuity, smoothness, and spatial coherence across
observations give traction in forming an accurate statistical estimate of the
unobserved dust map.

###  3.1 GP Inference with Integrated Observations

In a GP, any finite set of function evaluations are multivariate normal
distributed with covariance defined by the covariance function,
$\text{cov}(\rho(x_{i}),\rho(x_{j}))=k_{{\boldsymbol{\mathbf{\theta}}}}(x_{i},x_{j})$.
When observations are corrupted with Gaussian noise—the typical setting for GP
models—the joint distribution between all observations and latent process
quantities remains multivariate normal and the posterior at a new point
$\rho(x^{*})$ can be expressed as a few matrix operations (Rasmussen and
Williams, 2006, Chapter 2). Inference in Gaussian process models with noisy
_integrated observations_ — defined in Equation 2—closely mirrors inference
with the more typical, noisy pointwise observations.

Consider the model in Equations 1 and 2. We derive the posterior distribution
by first characterizing the joint distribution over $\rho(x_{*})$ and
$\boldsymbol{\mathbf{a}}=(a_{1},\dots,a_{N})$ (Rasmussen and Williams, 2006).
Observe that linear operations of multivariate normal random variables remain
multivariate normal in distribution. As noted in Rasmussen and Williams (2006,
Section 9.8), because integration is a linear operation, the value of the
integral of $\rho$ over the domain $\mathcal{X}$ will remain Gaussian. This is
also true for any collection of definite integrals and the corresponding
(Gaussian) noisy observations.

The joint distribution over $\boldsymbol{\mathbf{a}}$ and
$\rho(x_{*})\triangleq\rho_{*}$ remains multivariate normal. Given the mean
and covariance of this joint distribution, we can compute the posterior
distribution over $\rho_{*}$ (or any set of $\rho_{*}$ points) given
observations $\boldsymbol{\mathbf{a}}$, $\boldsymbol{\mathbf{x}}$, and
$\boldsymbol{\mathbf{\sigma}}^{2}$. The joint distribution over
$\boldsymbol{\mathbf{a}}$ and $\rho_{*}$ mimics the standard Gaussian process
setup

$\displaystyle\begin{pmatrix}\boldsymbol{\mathbf{a}}\\\ \rho_{*}\end{pmatrix}$
$\displaystyle\sim\mathcal{N}\left(\boldsymbol{\mathbf{0}},\begin{pmatrix}{\boldsymbol{\mathbf{K}}}_{\boldsymbol{\mathbf{e}}}+\boldsymbol{\mathbf{\Sigma}}_{\boldsymbol{\mathbf{a}}}&\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{e}},*}\\\
{\boldsymbol{\mathbf{K}}}_{*,\boldsymbol{\mathbf{e}}}&\boldsymbol{\mathbf{K}}_{*}\\\
\end{pmatrix}\right)\,,$ (5)

where the covariance matrix blocks are defined

$\displaystyle({\boldsymbol{\mathbf{K}}}_{\boldsymbol{\mathbf{e}}})_{ij}$
$\displaystyle=\mathrm{Cov}(e_{i},e_{j})\,,\quad({\boldsymbol{\mathbf{K}}}_{*,\boldsymbol{\mathbf{e}}})_{*j}=\mathrm{Cov}(\rho_{*},e_{j})\,,\quad\boldsymbol{\mathbf{K}}_{*}=\mathrm{Cov}(\rho_{*},\rho_{*})\,\,,$
(6)

and the observation noise matrix
$\boldsymbol{\mathbf{\Sigma}}_{\boldsymbol{\mathbf{a}}}\equiv\texttt{diag}(\sigma_{1}^{2},\dots,\sigma_{N}^{2})$.
We note that because $\rho_{*}$ is scalar, $\boldsymbol{\mathbf{K}}_{*}$ is a
$1\times 1$ matrix. Computing these covariance entries enables us to populate
the joint covariance matrix and use standard multivariate normal conditioning
for posterior inference. This conditioning mimics the equations for standard
Gaussian process inference,

$\displaystyle\rho(x_{*})$ $\displaystyle\sim\mathcal{N}(\mu_{*},\Sigma_{*})$
(7) $\displaystyle\mu_{*}$
$\displaystyle=\boldsymbol{\mathbf{K}}_{*,\boldsymbol{\mathbf{e}}}\left(\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{e}}}+\boldsymbol{\mathbf{\Sigma}}_{\boldsymbol{\mathbf{a}}}\right)^{-1}\boldsymbol{\mathbf{a}}$
(8) $\displaystyle\Sigma_{*}$
$\displaystyle=\boldsymbol{\mathbf{K}}_{*}-\boldsymbol{\mathbf{K}}_{*,\boldsymbol{\mathbf{e}}}\left(\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{e}}}+\boldsymbol{\mathbf{\Sigma}}_{\boldsymbol{\mathbf{a}}}\right)^{-1}\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{e}},*}.$
(9)

To compute this posterior, we need to compute the covariances in Equation 6.

The chosen covariance function directly specifies
$\mathrm{Cov}(\rho_{*},\rho_{*})$ The other terms, $\mathrm{Cov}(e_{i},e_{j})$
and $\mathrm{Cov}(\rho_{*},e_{j})$ are slightly more complex, and require
integrating the covariance function in one or both of its arguments.

######  Claim 3.1 (Semi-integrated Covariance Function).

The covariance between a process value $\rho_{i}\triangleq\rho(x_{i})$ and an
integrated value $e_{j}\triangleq\int_{x\in R_{j}}\rho(x)dx$ takes the form

$\displaystyle\mathrm{Cov}(\rho_{i},e_{j})$ $\displaystyle=\int_{x\in
R_{j}}k(x_{i},x)dx=k^{(\mathrm{semi})}(x_{i},x_{j})\,.$ (10)

For consistency, we will write the integrated argument second.

The proof is in Appendix B.

######  Claim 3.2 (Doubly-integrated Covariance Function).

The covariance between two integrated process values $e_{i}$ and $e_{j}$ takes
the form

$\displaystyle\mathrm{Cov}(e_{i},e_{j})$
$\displaystyle=\int_{(x,x^{\prime})\in R_{i}\times
R_{j}}k(x,x^{\prime})dxdx^{\prime}=k^{(\mathrm{double})}(x_{i},x_{j})\,.$ (11)

The proof for the doubly-integrated kernel is similar to the semi-integrated
case.

To summarize, if we can compute these two types of covariance values then we
can plug the result into the distribution in Equation 5. As for standard GP
inference, we can then manipulate the joint distribution using Gaussian
conditioning to compute the posterior distribution over the value of the
unobserved function $\rho(x_{*})$ or functionals over test locations, for
example integrated values $e_{*}$. Algorithm 1 summarizes GP inference using
integrated measurements.

###  3.2 Choice of covariance function

Assumptions about the dust map $\rho(\cdot)$ can be encoded in the covariance
function $k_{{\boldsymbol{\mathbf{\theta}}}}(\cdot,\cdot)$. We focus on
kernels that are _stationary_ and _isotropic_ , which can be written
$k_{{\boldsymbol{\mathbf{\theta}}}}(x,y)=\sigma^{2}k\left(\frac{|x-y|}{\ell}\right)$
where ${\boldsymbol{\mathbf{\theta}}}\triangleq(\sigma^{2},\ell)$. The
parameter $\ell$ is the _length scale_ , and the parameter $\sigma^{2}$ is the
process marginal variance (Genton, 2001). The length scale controls how smooth
the function is; the marginal variance influences the function’s amplitude. We
will compare three families of kernel functions: squared exponential, Matérn,
and a kernel from Gneiting (Gneiting, 2002). The squared exponential kernel —
a commonly used covariance function for Gaussian process regression — admits
an analytically tractable form for the semi-integrated kernel, but not the
doubly integrated kernel. (See Appendix A.) Other kernels, such as the Matérn
(Matérn, 1960) or the kernel presented in Gneiting (2002) do not readily admit
a semi-integrated form, which complicates their use with integrated
observations. Similarly, none of the mentioned kernels, including the squared
exponential, admit a doubly-integrated form over two line segments. We develop
a method to overcome this technical limitation in the following section.
Furthermore, we use the analytic tractability of the semi-integrated kernel to
validate this method on the squared exponential kernel.

Beyond the choice of covariance function family, the covariance function
parameters (e.g. the length scale and process variance) ought to be tuned.
Tuning GP covariance parameters can be complex — the process values
$\boldsymbol{\mathbf{\rho}}$ can strongly depend on parameters $\sigma^{2}$
and $\ell$, making it difficult to explore their joint space. We discuss this
phenomenon within the context of our scalable approximate inference algorithm
in Section 4.

##  4 Scaling Integrated GP Inference

Gaussian process models scale poorly to large datasets (Quiñonero-Candela and
Rasmussen, 2005; Titsias, 2009). Computing the posterior distribution of the
latent function, the posterior predictive distribution of new observations, or
the marginal likelihood of observed data (e.g. for model comparison) all
require computing the inverse or the determinant of a $N\times N$ matrix — a
$O(N^{3})$ operation. As $N$ grows larger than a few thousand observations,
this computation becomes intractable. To address this bottleneck, a common
strategy is to approximate Gaussian process inference using $M<<N$ _inducing
points_ , or spatial locations in the input space that are used to approximate
the full Gaussian process. Inducing point approximations can be interpreted in
multiple ways. One interpretation is that the inducing points are used to
construct a rank $M$ approximation to the $N\times N$ matrix that can be
efficiently inverted (e.g., $O(N\cdot M^{3})$) (Quiñonero-Candela and
Rasmussen, 2005). Another interpretation (which we adopt below), is that the
inducing points are used to define a family of variational distributions for
approximating the posterior over $\rho$, and optimizing the variational
objective avoids direct manipulation of the $N\times N$ matrix (Titsias, 2009;
Hensman et al., 2013).

Integrated observations create an additional computational issue — the semi-
integrated or doubly-integrated covariance value between any two values may be
unavailable in closed form (except in some special cases) and may require
high-precision numerical approximation. With the integrated observations
considered here, we have an additional computational issue. Calculating the
semi-integrated or doubly-integrated covariance value between any two values
may be unavailable in closed form (except in some special cases) and may
require high-precision numerical approximation. The exact GP inference
algorithm in Algorithm 1 requires computing a $N\times N$ matrix of doubly-
integrated covariance values, which will be computationally prohibitive when
$N$ is a modest size. Furthermore, when tuning covariance function parameters
or comparing covariance functions, the $N\times N$ matrix of numerically
integrated covariance function values will need to be recomputed many times.

Thus, to use the model to analyze millions of integrated stellar observations,
we derive a scalable variational inference algorithm to perform approximate
posterior inference. Variational inference approximates a posterior
distribution by optimizing a parameterized family—the variational
distribution—to be close to the exact posterior (Jordan et al., 1999;
Wainwright et al., 2008; Blei et al., 2017). For the model here, our approach
builds on the stochastic variational Gaussian process (SVGP) framework
(Hensman et al., 2013). It uses stochastic optimization to approximate the
posterior (Hoffman et al., 2013), operating on small mini-batches of
observations.

Further, to accommodate the integrated observations, we construct a Monte
Carlo estimate of the semi-integrated kernel that generalizes to covariance
functions that do not admit a closed form semi-integrated version.
Additionally, we use the rotational invariance of stationary and isotropic
covariance functions to construct a fast approximation to the doubly
integrated covariance kernel.

This section describes the SVGP framework and our approach to adapt the
framework to incorporate integrated observations. Section 4.1 describes
relevant details of the SVGP framework, including inducing points, the
variational family, and forming efficient mini-batch estimates of the
variational objective. Section 4.2 details our approach to do inference with
integrated observations for a general class of semi-integrated kernels.

###  4.1 Stochastic Variational Gaussian Processes

SVGP casts inference in GP models as an optimization problem. Crucially, the
SVGP optimization objective is constructed such that it can be written as a
sum of $N$ terms, each depending on only one data point. Given this
construction, the objective can be optimized using stochastic gradients
computed with mini-batches of data. This makes inference more computationally
and memory efficient.

Here, we discuss the SVGP objective with the standard observations — noisy
versions of $\rho(x)$.. We then adapt this approximate inference framework to
integrated observations in Section 4.2. Our derivation of the variational
objective is slightly different from Hensman et al. (2013)—we first define a
structured approximating family and then directly derive the evidence lower
bound objective.

For this section’s presentation of inducing points, variational inference, and
stochastic variational Gaussian processes, consider the typical GP model for
$N$ observations with $\rho\sim GP\left(0,k_{\theta}(\cdot,\cdot)\right)$, and
standard (i.e., non-integrated) observations
$y_{n}\,|\,\boldsymbol{\mathbf{x}}_{n}\sim\mathcal{N}(\rho(\boldsymbol{\mathbf{x}}_{n}),\sigma_{n}^{2})$.
We write the observation vector $\boldsymbol{\mathbf{y}}=(y_{1},\dots,y_{N})$
and the corresponding latent process vector
$\boldsymbol{\mathbf{\rho}}=(\rho(\boldsymbol{\mathbf{x}}_{1}),\dots,\rho(\boldsymbol{\mathbf{x}}_{N})$).

Inducing points. Inducing points are a common tool used to scale Gaussian
process inference (Quiñonero-Candela and Rasmussen, 2005). An inducing point
is simply a location in the input space, $\bar{x}$, which has a corresponding
inducing point value $\rho(\bar{x})$. Inducing points are typically
distributed about the input space — in our experiments we place them on a
fixed 3-dimensional grid throughout space.

We augment the model above with $M$ inducing points and their corresponding
values

$\displaystyle\bar{\boldsymbol{\mathbf{x}}}$
$\displaystyle\triangleq\bar{\boldsymbol{\mathbf{x}}}_{1},\dots,\bar{\boldsymbol{\mathbf{x}}}_{M}$
inducing points (12) $\displaystyle\boldsymbol{\mathbf{u}}$
$\displaystyle\triangleq\rho(\bar{\boldsymbol{\mathbf{x}}}_{1}),\dots,\rho(\bar{\boldsymbol{\mathbf{x}}}_{M})$
inducing point values (13)

where $\boldsymbol{\mathbf{u}}$ is the $M$-length vector of values of $\rho$
evaluated at each of the $M$ inducing points. Note that the introduction of
$\bar{\boldsymbol{\mathbf{x}}}$ and $\boldsymbol{\mathbf{u}}$ has not altered
the original model.

The model above specifies two latent variables, $\boldsymbol{\mathbf{\rho}}$
and $\boldsymbol{\mathbf{u}}$, whose distributions we wish to infer given data
$\boldsymbol{\mathbf{y}}$. That is, we wish to characterize the posterior
distribution
$p(\boldsymbol{\mathbf{\rho}},\boldsymbol{\mathbf{u}}\,|\,\boldsymbol{\mathbf{y}},\boldsymbol{\mathbf{x}},\bar{\boldsymbol{\mathbf{x}}})$
in a computationally efficient way.555Note that we are considering a model
where the likelihood and prior are both Gaussian distributions — a conjugate
pair — which implies that the posterior distribution will also be Gaussian.
Indeed it will be, but manipulating Gaussian will scale cubically in $N$. The
SVGP approach is constructed to scale linearly in $N$. SVGP uses inducing
points and a specific variational approximation to avoid the manipulation of
the $N\times N$ observation covariance matrix.

Variational inference. Variational inference (VI) is an optimization-based
approach to approximate posterior inference (Jordan et al., 1999; Wainwright
et al., 2008; Blei et al., 2017). VI methods posit a _variational family_ of
distributions, $q\in Q$, and use optimization techniques to find the optimal
approximate distribution from the set $Q$. Here, each element of $Q$ is a
distribution over the latent quantities —
$q(\boldsymbol{\mathbf{\rho}},\boldsymbol{\mathbf{u}})$.

The variational inference framework defines an objective to be optimized — the
most common VI objective is the _evidence lower bound_ (ELBO)

$\displaystyle\mathcal{L}(q)$
$\displaystyle=\mathbb{E}_{q(\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{\rho}})}\left[\ln
p(\boldsymbol{\mathbf{\rho}},\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{y}}\,|\,\boldsymbol{\mathbf{x}},\bar{\boldsymbol{\mathbf{x}}})-\ln
q(\boldsymbol{\mathbf{\rho}},\boldsymbol{\mathbf{u}})\right]\leq\ln
p(\boldsymbol{\mathbf{y}}\,|\,\boldsymbol{\mathbf{x}},\bar{\boldsymbol{\mathbf{x}}})\,.$
(14)

Maximizing the ELBO minimizes the KL divergence between
$q(\boldsymbol{\mathbf{\rho}},\boldsymbol{\mathbf{u}})$ and the true posterior
$p(\boldsymbol{\mathbf{\rho}},\boldsymbol{\mathbf{u}}\,|\,\boldsymbol{\mathbf{y}},\boldsymbol{\mathbf{x}},\bar{\boldsymbol{\mathbf{x}}})$.

Stochastic variational Gaussian processes The stochastic variational GP
framework (SVGP) (Hensman et al., 2013) uses a particular form for the
variational family. Given the set of inducing points
$\bar{\boldsymbol{\mathbf{x}}}$, the SVGP variational approximation is

$\displaystyle q(\boldsymbol{\mathbf{\rho}},\boldsymbol{\mathbf{u}})$
$\displaystyle=p(\boldsymbol{\mathbf{\rho}}\,|\,\boldsymbol{\mathbf{u}})\,q_{{\boldsymbol{\mathbf{\lambda}}}}(\boldsymbol{\mathbf{u}})\,,\quad
q_{{\boldsymbol{\mathbf{\lambda}}}}(\boldsymbol{\mathbf{u}})\triangleq\mathcal{N}(\boldsymbol{\mathbf{u}}\,|\,\boldsymbol{\mathbf{m}},\boldsymbol{\mathbf{S}})\,,\quad\text{and}\,\,{\boldsymbol{\mathbf{\lambda}}}\triangleq(\boldsymbol{\mathbf{m}},\boldsymbol{\mathbf{S}}),$
(15)

where $p(\boldsymbol{\mathbf{\rho}}\,|\,\boldsymbol{\mathbf{u}})$ is specified
by the Gaussian process prior and recall that $\boldsymbol{\mathbf{u}}$ and
$\boldsymbol{\mathbf{\rho}}$ are jointly multivariate normal. The variational
parameters ${\boldsymbol{\mathbf{\lambda}}}$ are fit by optimizing the ELBO
defined in Equation 14.

The SVGP approximation defined in Equation 15 is chosen because it has a very
specific property — when plugged into Equation 14, the objective decomposes
into the sum of $N$ separate terms. This allows us to create unbiased _mini-
batched_ estimators of the objective (and its gradients), enabling efficient
inference. We can see through a straightforward algebraic manipulation (we
suppress $\rho$ and $\boldsymbol{\mathbf{u}}$’s dependence on
$\boldsymbol{\mathbf{x}}$ and $\bar{\boldsymbol{\mathbf{x}}}$ to remove
clutter)

$\displaystyle\mathcal{L}({\boldsymbol{\mathbf{\lambda}}})$
$\displaystyle=\mathbb{E}_{q(\boldsymbol{\mathbf{\rho}},\boldsymbol{\mathbf{u}})}\left[\ln
p(\boldsymbol{\mathbf{\rho}},\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{y}})-\ln
q_{{\boldsymbol{\mathbf{\lambda}}}}(\rho,\boldsymbol{\mathbf{u}})\right]$ (16)
$\displaystyle=\mathbb{E}_{q_{{\boldsymbol{\mathbf{\lambda}}}}(\boldsymbol{\mathbf{u}})p(\boldsymbol{\mathbf{\rho}}\,|\,\boldsymbol{\mathbf{u}})}\left[\ln
p(\boldsymbol{\mathbf{y}}\,|\,\boldsymbol{\mathbf{\rho}})+\bcancel{\ln
p(\boldsymbol{\mathbf{\rho}}\,|\,\boldsymbol{\mathbf{u}})}+\ln
p(\boldsymbol{\mathbf{u}})-\ln
q_{{\boldsymbol{\mathbf{\lambda}}}}(\boldsymbol{\mathbf{u}})-\bcancel{\ln
p(\boldsymbol{\mathbf{\rho}}\,|\,\boldsymbol{\mathbf{u}})}\right]$ (17)
$\displaystyle=\underbrace{\mathbb{E}_{q_{{\boldsymbol{\mathbf{\lambda}}}}(\boldsymbol{\mathbf{u}})p(\boldsymbol{\mathbf{\rho}}\,|\,\boldsymbol{\mathbf{u}})}\left[\ln
p(\boldsymbol{\mathbf{y}}\,|\,\boldsymbol{\mathbf{\rho}})\right]}_{\text{(i)}}-KL(q_{{\boldsymbol{\mathbf{\lambda}}}}(\boldsymbol{\mathbf{u}})\,||\,p(\boldsymbol{\mathbf{u}}))\,\,,$
(18)

and we can write the term (i) as a sum over the $N$ observations

(i)
$\displaystyle\triangleq\mathbb{E}_{q_{{\boldsymbol{\mathbf{\lambda}}}}(\boldsymbol{\mathbf{u}})}\left[\mathbb{E}_{p(\boldsymbol{\mathbf{\rho}}\,|\,\boldsymbol{\mathbf{u}})}\left[\ln
p(\boldsymbol{\mathbf{y}}\,|\,\boldsymbol{\mathbf{\rho}})\right]\right]$ (19)
$\displaystyle=\mathbb{E}_{q_{{\boldsymbol{\mathbf{\lambda}}}}(\boldsymbol{\mathbf{u}})}\left[\mathbb{E}_{p(\boldsymbol{\mathbf{\rho}}\,|\,\boldsymbol{\mathbf{u}})}\left[\sum_{n=1}^{N}\ln
p(y_{n}\,|\,\rho_{n})\right]\right]$ (20)
$\displaystyle=\sum_{n=1}^{N}\underbrace{\mathbb{E}_{q_{{\boldsymbol{\mathbf{\lambda}}}}(\boldsymbol{\mathbf{u}})}\left[\mathbb{E}_{p(\boldsymbol{\mathbf{\rho}}\,|\,\boldsymbol{\mathbf{u}})}\left[\ln
p(y_{n}\,|\,\rho_{n})\right]\right]}_{\triangleq\mathcal{L}_{n}}\,.$ (21)

When the likelihood $p(y_{n}\,|\,\rho_{n})$ is Gaussian, the expectation in
each of the $\mathcal{L}_{n}$ terms can be computed analytically. Notice the
cancellation in Equation 17 eliminates the term that involves all $N$ data
points and a $N\times N$ matrix inversion, $\ln
p(\boldsymbol{\mathbf{\rho}}\,|\,\boldsymbol{\mathbf{u}})$ (Titsias, 2009).

To complete the derivation, we write a mini-batched estimator of the ELBO.
Given a set of randomly selected observations $B$,

$\displaystyle\hat{\mathcal{L}}({\boldsymbol{\mathbf{\lambda}}})$
$\displaystyle=\frac{N}{|B|}\sum_{b\in
B}\mathcal{L}_{b}-KL(q_{{\boldsymbol{\mathbf{\lambda}}}}(\boldsymbol{\mathbf{u}})\,||\,p(\boldsymbol{\mathbf{u}}))\,\,,$
(22)

which is an unbiased estimator of the full objective, and only touches
$|B|<<N$ observations. We can use gradients of this estimator to optimize the
lower bound with respect to variational parameters
${\boldsymbol{\mathbf{\lambda}}}=(\boldsymbol{\mathbf{m}},\boldsymbol{\mathbf{S}})$.
With $M$ inducing points, a mini-batched estimator can be computed in
$O(|B|\cdot M^{3})$ time.

When predicting, note that we only condition on the inducing point values
$\boldsymbol{\mathbf{u}}$, relying only on the prior covariance between
$\boldsymbol{\mathbf{u}}$ and $\rho_{*}$. The structure of this variational
approximation forces the inducing point values $\boldsymbol{\mathbf{u}}$ to
represent all of the information learned from the data. While this
approximation is no longer “nonparametric” in the traditional sense, it has
been shown that inducing point methods can provide good approximations to
exact GP inference provided enough inducing points are used (Burt et al.,
2019). Empirically, we find that we can adequately tune covariance parameters
and predict on held-out data using a modest number of inducing points (on the
order of $M=6{,}000$).

Natural gradients Given the objective in Equation 18, we are now tasked with
finding the optimal parameters ${{\boldsymbol{\mathbf{\lambda}}}}$. We use
stochastic optimization with mini-batches (Hoffman et al., 2013) along with
natural gradients (Amari, 1998), which are effective for variational inference
of Gaussian processes (Hensman et al., 2013).

For exponential family distributions (e.g. the multivariate normal), the
natural gradient can be computed by taking the standard gradient with respect
to the mean parameters (Hoffman et al., 2013). Following Hensman et al.
(2013), the natural gradient for variational parameters
${\boldsymbol{\mathbf{\lambda}}}=\boldsymbol{\mathbf{m}}_{{\boldsymbol{\mathbf{\lambda}}}},\boldsymbol{\mathbf{S}}_{{\boldsymbol{\mathbf{\lambda}}}}$
is straightforward to express using the natural parameterization of the
multivariate normal $q_{{\boldsymbol{\mathbf{\lambda}}}}$. See Appendix C for
more details.

Whitened parameterization. Draws of $\boldsymbol{\mathbf{u}}\sim
p(\boldsymbol{\mathbf{u}}\,|\,\bar{\boldsymbol{\mathbf{x}}},{\boldsymbol{\mathbf{\theta}}})$
are highly dependent on ${\boldsymbol{\mathbf{\theta}}}$. Consider a draw
$\boldsymbol{\mathbf{u}}$ conditioned on a fixed length scale parameter of
$\ell$. This same draw will be highly improbable under the prior with a
different length scale parameter, e.g. $2\ell$. The variational approximation
${\boldsymbol{\mathbf{\lambda}}}$ targets the posterior distribution
$p(\boldsymbol{\mathbf{u}}\,|\,\boldsymbol{\mathbf{y}},{\boldsymbol{\mathbf{\theta}}})$,
which is highly dependent on the structured Gaussian process prior. This
dependence frustrates the joint inference of ${\boldsymbol{\mathbf{\lambda}}}$
and ${\boldsymbol{\mathbf{\theta}}}$ — small changes in
${\boldsymbol{\mathbf{\theta}}}$ can make the approximate distribution over
$\boldsymbol{\mathbf{u}}$ suboptimal. This dependence forces gradient methods
to take small steps, which leads to extremely slow joint inference.

An effective strategy for coping with this dependence is to do posterior
inference in an alternative parameterization of the same model that decouples
variables $\boldsymbol{\mathbf{u}}$ and ${\boldsymbol{\mathbf{\theta}}}$. For
structured latent Gaussian models (e.g. Gaussian processes) this alternative
model uses the _whitened_ prior (Murray and Adams, 2010) or _non-centered_
parameterization (Bernardo et al., 2003). For the dust map model, we whiten
the prior by altering the target of posterior inference — instead of
approximating the posterior over inducing point values
$p(\boldsymbol{\mathbf{u}}\,|\,\boldsymbol{\mathbf{y}},{\boldsymbol{\mathbf{\theta}}})$,
we target the posterior over inducing point _noise_ values
$p(\boldsymbol{\mathbf{z}}\,|\,\boldsymbol{\mathbf{y}},{\boldsymbol{\mathbf{\theta}}})$,
where the relationship between $\boldsymbol{\mathbf{z}}$ and
$\boldsymbol{\mathbf{u}}$ is constructed to produce an equivalent model. For
notational clarity, we define
$\boldsymbol{\mathbf{K}}\triangleq\boldsymbol{\mathbf{K}}^{({\boldsymbol{\mathbf{\theta}}})}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{u}}}$
for a fixed value of ${\boldsymbol{\mathbf{\theta}}}$ and
$\boldsymbol{\mathbf{L}}$ such that
$\boldsymbol{\mathbf{K}}=\boldsymbol{\mathbf{L}}\boldsymbol{\mathbf{L}}^{\intercal}$
is a general matrix square root. The whitened model can be written

$\displaystyle\boldsymbol{\mathbf{z}}$
$\displaystyle\sim\mathcal{N}(0,I_{M})\,,\quad\boldsymbol{\mathbf{u}}\triangleq\boldsymbol{\mathbf{L}}\boldsymbol{\mathbf{z}}\sim\mathcal{N}(0,\boldsymbol{\mathbf{K}})\,.$
(23)

We now target the posterior distribution
$p(\boldsymbol{\mathbf{z}}\,|\,\boldsymbol{\mathbf{y}},{\boldsymbol{\mathbf{\theta}}})$
by fitting an approximation
$q_{{\boldsymbol{\mathbf{\lambda}}}}(\boldsymbol{\mathbf{z}})=\mathcal{N}(\boldsymbol{\mathbf{z}}\,|\,\tilde{\boldsymbol{\mathbf{m}}},\tilde{\boldsymbol{\mathbf{S}}})$.
The linear relationship between $\boldsymbol{\mathbf{z}}$ and
$\boldsymbol{\mathbf{u}}$ implies that
$\boldsymbol{\mathbf{m}}=\boldsymbol{\mathbf{L}}\tilde{\boldsymbol{\mathbf{m}}}$
and
$\boldsymbol{\mathbf{S}}=\boldsymbol{\mathbf{L}}\tilde{\boldsymbol{\mathbf{S}}}\boldsymbol{\mathbf{L}}^{\intercal}$.
We describe additional details of this parameterization in Appendix D.

###  4.2 Adapting to integrated observations

The previous section detailed a scalable approximate inference algorithm for
Gaussian processes in the typical noisy pointwise observation setting. Here,
we adapt this framework to integrated observations. The challenge is now
efficiently computing the appropriate semi- and doubly-integrated covariance
values to be used within each mini-batch update.

The natural gradient described in Equation 39 consists of two types of
covariance matrices —
$\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{u}}}$
and
$\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{\rho}}}$.
The former describes the prior covariance between inducing point values and
remains exactly the same in the integrated observation setting. The latter
describes the cross covariances between process locations and inducing point
values — this cross covariance must be adjusted to reflect the covariance
between the _integrated process_ and the inducing point values. This is
precisely what is defined by the semi-integrated covariance function. The
entries of
$\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{\rho}},\boldsymbol{\mathbf{u}}}$
have to simply be switched to the semi-integrated covariance defined in Claim
B.1

$\displaystyle(\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{\rho}},\boldsymbol{\mathbf{u}}})_{n,m}$
$\displaystyle=\int_{x\in R_{n}}k(x_{m},x)dx\,.$ (24)

In addition to the semi-integrated cross covariances, ziggy will need to
calculate doubly-integrated covariance terms. First, note that the gradient
defined in Equation 39 does not require doubly integrated terms. This is
convenient — the inducing point method does not require cross covariances
between integrated observations, and these cross covariances also happen to be
the most expensive to approximate with numerical methods. However, computing
the variational objective (Equation 18) and making predictions (Equation 78)
do require the doubly-integrated diagonal terms,
$\boldsymbol{\mathbf{K}}_{n,n}$ for each observation.

Given the above dependence on integrated covariance terms, the full adaptation
of the SVGP method to integrated observations runs into two technical hurdles.
First, the semi-integrated covariance is not available in closed form beyond
the simple squared exponential, which is considered too smooth (and thus of
limited capacity) in many settings. This limitation will make it difficult to
use custom covariance kernels constructed with astronomical theory (Leike and
Enßlin, 2019). Further, it will make it difficult to compare different models
and choose the best predictor among them. Second, even when the semi-
integrated covariance kernel is available in closed form, the doubly-
integrated covariance is not. The SVGP (and any inducing point) algorithm only
relies on _the diagonal_ of the covariance kernel, so we only need to
approximate this diagonal, scalar-valued function. We show that this scalar-
valued function can be easily approximated using a small grid of interpolated
values, where the “true” values are approximated using numerical integration.

We introduce two computationally efficient approximations to these operations
on the covariance function to overcome these technical issues: (i) Monte Carlo
estimates for semi-integrated covariance functions and (ii) linear
interpolation for the doubly-integrated diagonal covariance function.

Monte Carlo estimators for semi-integrated kernels The natural gradient
updates in Equation 39 require computing the covariance between observations
and inducing point values,
$\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{\rho}},\boldsymbol{\mathbf{u}}}$.
For integrated observations, the covariance between observations and inducing
points is given by the semi-integrated kernel in Equation 24. For kernels with
a closed form semi-integrated kernel (e.g. the squared exponential), we can
simply substitute $k^{(semi)}(\bar{x}_{m},x_{n})$ for $k(\bar{x}_{m},x_{n})$
in the loss and gradients and proceed within the typical SVGP framework.

For kernels without a closed form semi-integrated kernel, we run into a
computational issue — how do we compute the appropriate cross covariance
values? One approach is to use numerical integration — the semi-integrated
covariance is a one-dimensional integral that can be approximated with
quadrature techniques. However, each batch requires computing $M\times|B|$
semi-integrated covariance values — this can be computationally expensive,
which would force us to use small batches, and this leads to slower learning.

Stochastic natural gradient updates are already an estimate of the true
gradient — how _good_ does this estimate have to be to effectively find the
optimal solution? Numeric integration approximations to
$k^{(semi)}(\cdot,\cdot)$ are precise to nearly machine precision, but are
computationally expensive. Analogously we can ask, how _good_ do the semi-
integrated covariance approximations have to be to effectively find the
optimal solution?

Pursuing this idea, we propose using a Monte Carlo approximation to the semi-
integrated covariance values — we sample uniformly along the integral path and
average the original covariance values. More formally, we introduce a uniform
random variable $\nu_{R_{n}}$ that takes values along the ray $R_{n}$ from the
origin to $x_{n}$. We can now write the semi-integrated covariance as

$\displaystyle k^{(semi)}(x_{m},x_{n})$ $\displaystyle=\int_{x\in
R_{n}}k(x_{m},x)dx=|R_{n}|\int_{\nu\in
R_{n}}\underbrace{\frac{1}{|R_{n}|}}_{=p(\nu)}k(x_{m},x)dx=|R_{n}|\,\mathbb{E}_{\nu}\left[k(x_{m},\nu)\right]\,.$

This form admits a simple Monte Carlo estimator

$\displaystyle\nu^{(1)},\dots,\nu^{(L)}$ $\displaystyle\sim
Unif(R_{n})\,,\quad\hat{k}^{(semi)}(x_{m},x_{n})=\frac{|R_{n}|}{L}\sum_{\ell}k(x_{m},\nu^{(\ell)})\,\,.$
(25)

It is straightforward to show that Equation 25 is an unbiased and consistent
estimator for the true semi-integrated covariance value evaluated at $x_{m}$
and $x_{n}$. While plugging this directly into the gradient formula results in
a biased estimator for the natural gradient, we find that with a modest number
of samples we can closely match the optimization performance of the true
natural gradient estimator. We investigate the relationship between samples
$L$ and the resulting optimization in Appendix Figure 9. Appendix E.1 contains
additional details and experimental results validating this approach.

Although conceptually similar, note that this Monte Carlo estimate of the
semi-integrated covariance is not equivalent to discretizing along each line
of site as in (Rezaei Kh. et al., 2017). The estimator can use a small number
of samples along the ray and still be useful within stochastic gradient
optimization.

Integrated Observation Variance. The SVGP algorithm requires an additional
covariance calculation — the marginal variance of an integrated observation at
point $x_{n}$, $k^{\mathrm{(doubly)}}(x_{n},x_{n})$. Note that we do not need
to compute the non-diagonal terms of the doubly integrated kernel, a
consequence of the separability of Equation 21. Similar to the semi-integrated
kernel function, the doubly integrated function can be approximated with
numerical methods — a double quadrature routine for each observation. But,
similar to the semi-integrated case, this double quadrature call becomes a
computational bottleneck within each batch. Furthermore, computing the
gradient of the ELBO with respect to the covariance parameters with
autodifferentiation tools would require differentiating through a double
quadrature call — prohibitive in our setting.

We solve this problem by observing that a rotationally invariant (and
stationary) kernel admits a doubly-integrated version that is a function only
of the distances to each of the two integral endpoints. That is,
$k^{\mathrm{(doubly)}}(x_{i},x_{j})$ can be written as a function
$f(|x_{i}|,|x_{j}|)$ for a fixed setting of covariance function parameters.
Further, we can write the diagonal of the doubly-integrated kernel as a
function of only the distance to the single point,
$k^{\mathrm{(doubly)}}(x_{n},x_{n})=f(|x_{n}|)$.

Given that this is a one-dimensional function shared by all $N$ of our
observations, we approximate the doubly-integrated kernel diagonal with linear
interpolation. The interpolation scheme allows us to cheaply compute a highly
accurate approximation to the doubly-integrated diagonal. It also enables us
to easily backpropagate gradients to the kernel function parameters, which
enables efficient tuning during the variational inference optimization
routine. Appendix E.2 describes the interpolation scheme for
$k^{\mathrm{(doubly)}}(x_{n},x_{n})$ in further detail.

###  4.3 Method summary and remarks

To summarize, ziggy scales Gaussian process inference with integrated
observations by fusing several strategies: (i) the explicit representation of
the inducing point posterior within the SVGP framework; (ii) a whitened
parameterization to decouple ${\boldsymbol{\mathbf{\lambda}}}$ and
${\boldsymbol{\mathbf{\theta}}}$ in the variational objective; (iii) on-the-
fly Monte Carlo estimates of the semi-integrated covariance function; (iv) and
fast kernel interpolation of the doubly integrated diagonal covariance
function that leverages the stationary and isotropic properties of the
covariance functions considered.

Algorithm 2 describes the main loop for updating based on stochastic gradients
of the variational parameters and the covariance function parameters $\theta$.
Maximizing the variational lower bound (Equation 18) with respect to prior
parameters $\theta$ is a similar strategy to empirical Bayes (Efron et al.,
2008; Efron, 2019), a method used to reduce the average error of posterior
estimates.

We find that the inducing point approach dovetails nicely with integrated
observations—most evaluations of the covariance function are between inducing
points and observations. This only requires computing the (easier) semi-
integrated covariance function. Further the doubly-integrated covariance
values between any two distinct observations never needs to be computed.
Rather, only the _diagonal_ of this covariance function needs to be computed,
but this conveniently reduces to a one-dimensional function that can be
efficiently approximated. ziggy circumvents computing the (prohibitive)
$N\times N$ doubly-integrated covariance.

We note that the SVGP framework for scaling Gaussian process is one approach
to build upon for our application — other frameworks for scaling kernel
methods could have been employed. One example is random (or carefully
selected) features for approximating the kernel function (Rahimi and Recht,
2008; Yang et al., 2012; Le et al., 2013). While these approximate kernel
methods could yield superior performance, one reason we reach for the SVGP
framework is extensibility. The variational inference framework can
incorporate more complex probabilistic graphical model structure. For this
application, additional types of observations can be incorporated into a more
sophisticated likelihood model to help identify the spatial dust distribution
and incorporate additional sources of uncertainty.

(a) True $\rho(x)$

(b) Integrated observations $a_{n}$

(c) Inferred $\mathbb{E}\left[\rho\,|\,\mathcal{D}\right]$.

(d) Inferred posterior scale
$\mathbb{V}\left[\rho(x)\,|\,\mathcal{D}\right]^{1/2}$.

Figure 2:  Integrated observation GPs can reveal spatial structure from a
severely limited vantage point. This figure summarizes domain simulated data
posterior inference using $N=1{,}000{,}000$ examples in a three-dimensional
space. Panel 2(a) depicts the true (unobserved) $\rho(x)$ that generates the
data. Panel 2(b) depicts the observed data, noisy integrated measurements of
$\rho$ from the origin to stellar positions in the domain simulation. 2(c)
depicts the inferred $\rho$ given these observations using ziggy with a
squared exponential kernel. 2(d) depicts the marginal posterior standard
deviation at each location in $\mathcal{X}$-space. All Panels are a slice at
$z=0$

##  5 Validation with a Domain Simulation

We developed ziggy with the goal of generating a three-dimensional dust map of
the Milky Way using 2 billion stars in the Gaia dataset (Brown et al., 2018).
As a step towards that goal—and to test ziggy in an applied environment—we
infer the dust distribution of a mechanistic and physically motivated domain
simulation. The domain dataset is a synthetic Gaia survey of a high resolution
Milky Way-like galaxy simulation.

The domain dust density $\rho(x)$ was generated in a Milky Way-like galaxy
simulation, one of many in the Latte suite of simulations of Milky-Way-mass
galaxies (Wetzel et al., 2016; Hopkins et al., 2018). It was run as part of
the Feedback In Realistic Environments (FIRE) simulation project, which self-
consistently models extinction and star formation in a cosmological context at
high resolution.

In astronomy, most observations prefer a model in which our Universe is
predominantly made of cold dark matter (Aghanim et al., 2018). On small
scales, however, there are significant challenges to this model (Klypin et
al., 1999), which this suite of simulations helps resolve (Wetzel et al.,
2016). This Latte simulation overcomes the computational challenge of
including gas, dust, and stars with high enough resolution to resolve these
tensions. It includes more complex physics of “normal” matter, as apposed to
just dark matter, including a high resolution disk of gas and dust. Because of
the results from these simulations, which resolved these tensions, astronomers
are more confident in a universal model of dark matter.

The simulation has enough resolution to resolve the formation of structures in
the dense gas that forms dust and stars, creating a realistic environment for
inference. The positions of the extinction observations, or the positions of
stars within $\rho(\cdot)$, was generated by the Ananke framework. The Ananke
framework generated a realistic, mock star catalog from the FIRE simulation,
and kept intact important observational relationships between gas, extinction,
and stellar populations (Sanderson et al., 2018). This catalogue was
specifically designed to resemble a Data Release 2 Gaia astrometric survey
(Brown et al., 2018), with a similar resolution of stellar density as Gaia. We
integrated $\rho(x)$ along lines of sight to 1 million stars in the Ananke
dataset within a $0.5\mathrm{kpc}\times 0.5\mathrm{kpc}\times 0.1\mathrm{kpc}$
region of the synthetic sun. So, we generate a set of N noisy integrated
observations

$\displaystyle x_{n}$ $\displaystyle\sim\text{Ananke}(\mathcal{X})$ stellar
locations (26) $\displaystyle e_{n}$ $\displaystyle=\int_{x\in
R_{n}}\rho(x)dx$ domain extinctions (27) $\displaystyle a_{n}$
$\displaystyle\sim\mathcal{N}(e_{n},\sigma_{n}^{2})$ noisy integrated
observation (28)

where the domain $\mathcal{X}=[(-0.25,0.25),(-0.25,0.25),(-0.05,0.05)]$, the
median extinction value is $0.01$, and noise variance is chosen to be
$\sigma_{n}=0.005$. So our median signal-to-noise value is $3$. The true
extinctions, $e_{n}$, are computed to high precision using numerical
quadrature, integrating from the origin to the Ananke stellar locations
$x_{n}$. The integrated observations are depicted in Figure 2(b). We estimate
this domain specific dust density $\rho(\cdot)$ from noisy integrated
observations using ziggy.

In Section 5.1, we compare posterior estimates of $\rho(\cdot)$ as a function
of training data set size. This highlights the importance of scaling inference
to more observations in order to obtain a more accurate statistical estimator.
In Section 5.2, we compare posterior estimates of $\rho(\cdot)$ as a function
of kernel choice, with a fixed data set size.

(a) Test $e$ RMSE

(b) Test $e$ log likelihood

Figure 3: Scaling to massive $N$ improves estimator performance. We compare
predictive RMSE (left) and log likelihood (right) as a function of data set
size $N$ on a held out sample of test stars.

(a) $\mathbb{E}[\rho\,|\,\mathcal{D}]$, $N=1{,}000$

(b) $\mathbb{E}[\rho\,|\,\mathcal{D}]$, $N=10{,}000$

(c) $\mathbb{E}[\rho\,|\,\mathcal{D}]$, $N=100{,}000$

(d) $\mathbb{E}[\rho\,|\,\mathcal{D}]$, $N=1{,}000{,}000$

Figure 4: Slices of the posterior mean function at $z=0$ for various dataset
sizes. More data leads to more accurate predictions of $\rho(\cdot)$.

###  5.1 The quality of the estimate

We study the quality of the estimate of $\rho(x)$ as a function of data set
size $N$. We fit the dust model to the Ananke dataset with data set sizes
$N\in\\{10^{3},10^{4},10^{5},10^{6}\\}$ for ${10^{5},10^{4},10^{3},100}$
epochs, respectively. We trained each model using mini-batches of size
$|B|=2{,}000$, an initial step size of 0.01, reducing it every epoch by
multiplying by a decay factor of 0.99. We measure model quality by computing
root mean squared error (RMSE) and log likelihood (LL) on a set of
$N_{\mathrm{test}}=2{,}000$ held out extinction values. Each model uses
$M=16\times 16\times 4$ inducing points, evenly spaced in a grid in the input
space.

Figure 2 summarizes the model’s fit using one million observations, the most
we tried in this example. Figure 2(c) displays the posterior mean for
$\rho(x)$ given the one million observations depicted in 2(b). Figure 2(d)
shows the posterior uncertainty (one standard deviation) about the estimated
mean. We see that the million-observation model recovers the true latent
function particularly well.

Figure 3 quantifies the improvement in the RMSE and log likelihood on a held
out sample of test stars as a function of dataset size $N$. As the dataset
size increases, both statistics drastically improve. Figure 4 visualizes a
direct comparison of the posterior estimate of the latent function as we
increase data set size $N$. We can clearly see that as we incorporate more
data into the nonparametric model, the form of the true underlying function
emerges. To accurately visualize and interpret large scale features of the
latent dust distribution, we will want to incorporate as many stellar
observations as possible. This is particularly true in the low signal to noise
regime of extinction estimates.

###  5.2 Comparing kernels

Similar to the synthetic dataset, we compare multiple models by fitting the
variational approximation and tuning the covariance function parameters for 5
different kernel choices.

In this domain setup, we ran each model for 100 epochs, saving the the model
with the best ELBO value. We used mini-batches of size $|B|=2{,}000$, and
started the step size at $0.01$, reducing it every epoch by multiplying by a
decay factor of $0.99$. For kernels that do not have a closed form semi-
integrated version, we used $L=50$ uniform grid Monte Carlo samples to
estimate the integrated covariance.

We validate the inferences against a set of $2{,}000$ held-out test
extinctions $e_{*}$. We quantify the quality of the inferences with mean
squared error, log likelihood, and the $\chi^{2}$-statistic on the test data.
We also visually validate ziggy with a $QQ$-plot and coverage comparison, also
on test data. We also visualize the behavior of ziggy as a function of
distance to the origin — specifically, we visualize how accurate and well-
calibrated test inferences are as synthetic observations are made farther
away.

In Figure 5 we compare the mean squared error (MSE), the log likelihood, and
the $\chi^{2}$ statistic on the held out test data. We compare five different
kernels – (i) Gneiting $\alpha=1$ from Rezaei Kh. et al. (2018) (ii-iv)
Matérn, $\nu\in\\{1/2,3/2,5/2\\}$, and (v) the squared exponential kernel. The
results match our expectations. The true $\rho(\cdot)$ is quite smooth, and we
see that the smoother kernels tend to have lower predictive error and higher
log likelihood. All kernels have similarly calibrated $\chi^{2}$ statistic,
except for Matérn, $\nu=3/2$, which under predicts it’s posterior variances.
This can also be seen in Figure 6. The test statistic comparison has chosen a
good match in the squared exponential kernel.

To test the calibration of posterior uncertainties, we inspect statistics of
test-sample $z$-scores for $e_{*}$. We test that the statistics of predictive
$z$-scores resemble a standard normal distribution using a QQ-plot. For a more
detailed description fo the QQ-plot, please see Section F.2. To compare
different models, we visualize QQ-plots for the different covariance functions
in Figure 6. As an additional summary of calibration, we compare the fraction
of predicted examples covered by $1/2,1,2$, and 3 posterior standard
deviations, summarized in Table 1. We see that there is room for improvement
in the posterior variances. The $z$-scores and empirical coverage show
posterior variances that are well calibrated for the squared exponential and
Matérn, $\nu=5/2$ kernels, but are too large for the other kernels. The models
for the less smooth kernels are compute bound, requiring thousands of inducing
points to tile the domain. We ran the models for 100 epochs, which maxed out
the memory and compute time. Miscalibration could be the result of
misspecification of the model itself — e.g., mismatch between the covariance
function $k$ and the true underlying $\rho$, or the inflexibility of the
variational approximation. To diagnose this shortcoming, methods that allow us
to scale the expressivity of the variational approximation — e.g., larger $M$
— would be necessary.

We also depict model fit statistics as a function of distance from the origin
in Figure 7. Figure 7(a) visualizes the test extinction values. ziggy is able
to accurately reconstruct the dust map. Figure 7(b) shows the posterior
standard deviation by distance, which shows a noise reduction of about
$10\times$ for most stars. In Figure 7(c) we see that the predictions have
posterior variances that are slightly too high, causing $z$-scores that lie
mostly within $2\sigma$ independent of distance, with only the most nearby
stars extending to $3\sigma$. Decent calibration, including at far distances,
is encouraging, as the goal is to form good estimates of this density far away
from the observer location.

(a) Test $e$ MSE

(b) Test $e$ log likelihood

(c) Test $e$ $\chi^{2}$

Figure 5: Predictive summaries across different models and inference schemes.
The smoother models have better predictive summary statistics which is
consistent with the underlying density field being very smooth.

Figure 6: QQ-plot for predicted distributions on 2,000 held out
integrated-$\rho$ test points $e_{*}$. Theoretical normal quantiles are on the
horizontal axis, and predictive $z$-scored quantiles are on the vertical axis.
Posterior for extinctions $e$ are moderately well calibrated, with the squared
exponential kernel having the best calibration.

| Gneit-$\alpha=1$ | Mat-$\nu=1/2$ | Mat-$\nu=3/2$ | Mat-$\nu=5/2$ | SqExp | $\mathcal{N}(0,1)$
---|---|---|---|---|---|---
$\sigma$ |  |  |  |  |  |
0.5 | 0.894 | 0.760 | 0.451 | 0.730 | 0.625 | 0.383
1.0 | 0.965 | 0.900 | 0.677 | 0.900 | 0.869 | 0.683
2.0 | 0.997 | 0.969 | 0.878 | 0.991 | 0.986 | 0.955
3.0 | 1.000 | 0.994 | 0.947 | 1.000 | 0.999 | 0.997

Table 1: Coverage fractions at various levels of z-score standard deviation
$\sigma$ for test extinctions $e_{*}$. Compared with theoretical coverage
fractions, we see that our inferences generate predictions that are well
calibrated for the smoothest kernels

(a) Test $e$ values by distance

(b) Test $e$ posterior standard deviation

(c) Test $e$ z scores by distance

Figure 7: Posterior summaries as a function of distance to the origin
(observation point). The noise variance for the extinction estimates is
$\sigma_{n}=0.005$, so Figure a demonstrates our low signal to noise regime,
Figure b demonstrates our shrinkage by a factor of 10, and Figure c
demonstrates our decent calibration in this very low signal to noise regime.

##  6 Discussion

Mapping the three-dimensional distribution of dust in the Milky Way is
fundamental to astronomy. Star dust obscures our observations of star light,
and so an accurate dust map would help us more clearly observe the universe.
Such a dust map would be a new lens into the dynamics of the Milky Way,
providing a trace of the process by which it formed.

In this paper, we developed ziggy, a method to estimate the dust map from
millions of astronomical observations. ziggy incorporates noisy information
into a single, coherent spatial model. We developed technical innovations to
accommodate integrated observations into the Gaussian process framework, and
to scale such a model to millions of observations without spatial
discretization. We validated these algorithmic innovations with numerical
studies, and measured the performance of ziggy in both a synthetic setting and
a realistic, high-fidelity dataset used in astronomical study. We showed that
ziggy is able to incorporate millions of observations and recover accurate and
well-calibrated estimates of both stellar extinctions and pointwise values of
the three-dimensional dust map.

There are several avenues for future research and improvement. ziggy assumes
that distances are known, when in reality they are noisy. In practical
applications, we will only include stars with precise distance measurements,
lowering the spatial density of observations, which may degrade estimates. We
view ziggy as a step toward a probabilistic model of photometric measurements
that jointly infers brightnesses, colors, distance, and spatial dust. Jointly
modeling brightnesses, distances, and a spatial dust map using a Bayesian
hierarchical model will allow us to leverage the coherence of Bayesian
probabilistic modeling to form tighter estimates of all modeled quantities.

Additionally, a limitation of our treatment is the Gaussian error assumption.
Non-Gaussian observation errors could degrade extinction estimates — e.g.,
forcing $\rho$ to be more complex (i.e., with a smaller length scale) than it
truly is. A flexible hierarchical Bayesian model could incorporate a non-
Gaussian likelihood that could be more robust to heavier tailed observation
error.

A full map that spans the entirety of the Milky Way will require increasing
the capacity of the approximation. To reconstruct both global and local
features using an inducing point method, we need to ensure that small length
scales (relative to the input domain) can be resolved. When inducing points
are spatially distant from one another, the inducing point approximation will
not have the capacity to represent sharp changes in $\rho(x)$ at small scales.
This limitation can be overcome by including more, and closer, inducing
points. However, introducing more inducing points butts up against the
$O(M^{3})$ computational limitations. In this work, we assume we are able to
populate the space with enough inducing points; but the scaling concern is an
important avenue of future work. Incorporating more inducing points will be
crucial to resolving both global features and fine local features within the
Milky Way.

The ultimate goal of this line of work is to produce an accurate catalog of
the properties of stars, such as probabilistic brightnesses and distances, and
the focus of scale is motivated by the size of modern photometric catalogs.
PAN-STARRS 1 catalog includes 2.4 million detected stars (Flewelling, 2016),
the fifteenth data release of the Sloan Digital Sky Survey includes photometry
for over 260 million detected stars (Aguado et al., 2019), and the second GAIA
data release has 1.3 billion stars with parallax measurements that can be used
to estimate the interstellar dust distribution (Brown et al., 2018). As a
final avenue of research, we will scale ziggy to incorporate billions of
observations in a much larger spatial domain.

## References

  * Aghanim et al. (2018) Aghanim et al., N. (2018).  Planck 2018 results. VI. Cosmological parameters.  arXiv e-prints, page arXiv:1807.06209.
  * Aguado et al. (2019) Aguado, D. S., Ahumada, R., Almeida, A., Anderson, S. F., Andrews, B. H., Anguiano, B., Ortíz, E. A., Aragón-Salamanca, A., Argudo-Fernández, M., Aubert, M., et al. (2019).  The fifteenth data release of the sloan digital sky surveys: first release of manga-derived quantities, data visualization tools, and stellar library.  The Astrophysical Journal Supplement Series, 240(2):23.
  * Amari (1982) Amari, S.-I. (1982).  Differential geometry of curved exponential families-curvatures and information loss.  The Annals of Statistics, pages 357–385.
  * Amari (1998) Amari, S.-I. (1998).  Natural gradient works efficiently in learning.  Neural computation, 10(2):251–276.
  * Ambikasaran et al. (2014) Ambikasaran, S., Foreman-Mackey, D., Greengard, L., Hogg, D. W., and O’Neil, M. (2014).  Fast direct methods for Gaussian processes.  arXiv preprint arXiv:1403.6015.
  * Banerjee et al. (2014) Banerjee, S., Carlin, B. P., and Gelfand, A. E. (2014).  Hierarchical modeling and analysis for spatial data.  Chapman and Hall/CRC.
  * Bernardo et al. (2003) Bernardo, J., Bayarri, M., Berger, J., Dawid, A., Heckerman, D., Smith, A., and West, M. (2003).  Non-centered parameterisations for hierarchical models and data augmentation.  In Bayesian Statistics 7: Proceedings of the Seventh Valencia International Meeting, volume 307. Oxford University Press, USA.
  * 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.
  * Brown et al. (2018) Brown, A., Vallenari, A., Prusti, T., De Bruijne, J., Babusiaux, C., Bailer-Jones, C., Biermann, M., Evans, D. W., Eyer, L., Jansen, F., et al. (2018).  Gaia data release 2-summary of the contents and survey properties.  Astronomy & astrophysics, 616:A1.
  * Burt et al. (2019) Burt, D. R., Rasmussen, C. E., and van der Wilk, M. (2019).  Rates of Convergence for Sparse Variational Gaussian Process Regression.  arXiv preprint arXiv:1903.03571.
  * Chen et al. (2019) Chen, B.-Q., Huang, Y., Hou, L.-G., Tian, H., Li, G.-X., Yuan, H.-B., Wang, H.-F., Wang, C., Tian, Z.-J., and Liu, X.-W. (2019).  The Galactic spiral structure as revealed by O- and early B-type stars.  Monthly Notices of the Royal Astronomical Society, 487:1400–1409.
  * Cressie (1990) Cressie, N. (1990).  The origins of kriging.  Mathematical geology, 22(3):239–252.
  * Cressie (1992) Cressie, N. (1992).  Statistics for spatial data.  Terra Nova, 4(5):613–617.
  * Datta et al. (2016) Datta, A., Banerjee, S., Finley, A. O., and Gelfand, A. E. (2016).  Hierarchical nearest-neighbor Gaussian process models for large geostatistical datasets.  Journal of the American Statistical Association, 111(514):800–812.
  * Efron (2019) Efron, B. (2019).  Bayes, oracle bayes and empirical bayes.  Statistical science, 34(2):177–201.
  * Efron et al. (2008) Efron, B. et al. (2008).  Microarrays, empirical bayes and the two-groups model.  Statistical science, 23(1):1–22.
  * Fayed and Atiya (2014) Fayed, H. and Atiya, A. (2014).  An evaluation of the integral of the product of the error function and the normal probability density with application to the bivariate normal integral.  Mathematics of Computation, 83(285):235–250.
  * Flewelling (2016) Flewelling, H. (2016).  The Pan-STARRS 1 Medium Deep Field Variable Star Catalog.  In American Astronomical Society Meeting Abstracts #227, volume 227 of American Astronomical Society Meeting Abstracts, page 144.25.
  * Foreman-Mackey et al. (2017) Foreman-Mackey, D., Agol, E., Ambikasaran, S., and Angus, R. (2017).  Fast and scalable Gaussian process modeling with applications to astronomical time series.  The Astronomical Journal, 154(6):220.
  * Genton (2001) Genton, M. G. (2001).  Classes of kernels for machine learning: a statistics perspective.  Journal of machine learning research, 2(Dec):299–312.
  * Georgelin and Georgelin (1976) Georgelin, Y. M. and Georgelin, Y. P. (1976).  The spiral structure of our Galaxy determined from H II regions.  Astronomy & Astrophysics, 49:57–79.
  * Gneiting (2002) Gneiting, T. (2002).  Compactly supported correlation functions.  Journal of Multivariate Analysis, 83(2):493–508.
  * Green et al. (2018) Green, G. M., Schlafly, E. F., Finkbeiner, D., Rix, H.-W., Martin, N., Burgett, W., Draper, P. W., Flewelling, H., Hodapp, K., Kaiser, N., et al. (2018).  Galactic reddening in 3d from stellar photometry–an improved map.  Monthly Notices of the Royal Astronomical Society, 478(1):651–666.
  * Green et al. (2019) Green, G. M., Schlafly, E. F., Zucker, C., Speagle, J. S., and Finkbeiner, D. P. (2019).  A 3d dust map based on gaia, pan-starrs 1 and 2mass.  arXiv preprint arXiv:1905.02734.
  * Heaton et al. (2019) Heaton, M. J., Datta, A., Finley, A. O., Furrer, R., Guinness, J., Guhaniyogi, R., Gerber, F., Gramacy, R. B., Hammerling, D., Katzfuss, M., et al. (2019).  A case study competition among methods for analyzing large spatial data.  Journal of Agricultural, Biological and Environmental Statistics, 24(3):398–425.
  * Hensman et al. (2013) Hensman, J., Fusi, N., and Lawrence, N. D. (2013).  Gaussian processes for big data.  In Proceedings of the Twenty-Ninth Conference on Uncertainty in Artificial Intelligence, pages 282–290. AUAI Press.
  * Hoffman et al. (2013) Hoffman, M. D., Blei, D. M., Wang, C., and Paisley, J. (2013).  Stochastic variational inference.  The Journal of Machine Learning Research, 14(1):1303–1347.
  * Hopkins et al. (2018) Hopkins, P. F., Wetzel, A., Kereš, D., Faucher-Giguère, C.-A., Quataert, E., Boylan-Kolchin, M., Murray, N., Hayward, C. C., Garrison-Kimmel, S., and Hummels, C. (2018).  FIRE-2 simulations: physics versus numerics in galaxy formation.  Monthly Notices of the Royal Astronomical Society, 480(1):800–863.
  * Jordan et al. (1999) Jordan, M. I., Ghahramani, Z., Jaakkola, T. S., and Saul, L. K. (1999).  An introduction to variational methods for graphical models.  Machine learning, 37(2):183–233.
  * Klypin et al. (1999) Klypin, A., Kravtsov, A. V., Valenzuela, O., and Prada, F. (1999).  Where Are the Missing Galactic Satellites?  The Astrophysical Journal, 522(1):82–92.
  * Krige (1951) Krige, D. G. (1951).  A statistical approach to some basic mine valuation problems on the witwatersrand.  Journal of the Southern African Institute of Mining and Metallurgy, 52(6):119–139.
  * Le et al. (2013) Le, Q., Sarlós, T., and Smola, A. (2013).  Fastfood-computing hilbert space expansions in loglinear time.  In International Conference on Machine Learning, pages 244–252.
  * Leike and Enßlin (2019) Leike, R. and Enßlin, T. (2019).  Charting nearby dust clouds using gaia data only.  arXiv preprint arXiv:1901.05971.
  * Matérn (1960) Matérn, B. (1960).  Spatial variation: Stochastic models and their applications to some problems in forest surveys and other sampling investigations.  Meddelanden från Statens Skogsforskningsinstitut, 49:1–144.
  * Matheron (1963) Matheron, G. (1963).  Principles of geostatistics.  Economic geology, 58(8):1246–1266.
  * Matheron (1973) Matheron, G. (1973).  The intrinsic random functions and their applications.  Advances in applied probability, 5(3):439–468.
  * Mathis (1990) Mathis, J. S. (1990).  Interstellar dust and extinction.  Annual Review of Astronomy and Astrophysics, 28(1):37–70.
  * Murray and Adams (2010) Murray, I. and Adams, R. P. (2010).  Slice sampling covariance hyperparameters of latent Gaussian models.  In Advances in neural information processing systems, pages 1732–1740.
  * Quiñonero-Candela and Rasmussen (2005) Quiñonero-Candela, J. and Rasmussen, C. E. (2005).  A unifying view of sparse approximate Gaussian process regression.  Journal of Machine Learning Research, 6(Dec):1939–1959.
  * Rahimi and Recht (2008) Rahimi, A. and Recht, B. (2008).  Random features for large-scale kernel machines.  In Advances in neural information processing systems, pages 1177–1184.
  * Rasmussen and Williams (2006) Rasmussen, C. E. and Williams, C. K. I. (2006).  Gaussian Processes for Machine Learning.  The MIT Press.
  * Rezaei Kh. et al. (2017) Rezaei Kh., S., Bailer-Jones, C., Hanson, R., and Fouesneau, M. (2017).  Inferring the three-dimensional distribution of dust in the galaxy with a non-parametric method-preparing for gaia.  Astronomy & Astrophysics, 598:A125.
  * Rezaei Kh. et al. (2018) Rezaei Kh., S., Bailer-Jones, C. A. L., Hogg, D. W., and Schultheis, M. (2018).  Detection of the Milky Way spiral arms in dust from 3D mapping.  Astronomy & Astrophysics, 618:A168.
  * Robbins and Monro (1951) Robbins, H. and Monro, S. (1951).  A stochastic approximation method.  The Annals of Mathematical Statistics, 22(3):400–407.
  * Rue et al. (2009) Rue, H., Martino, S., and Chopin, N. (2009).  Approximate Bayesian inference for latent Gaussian models by using integrated nested Laplace approximations.  Journal of the royal statistical society: Series b (statistical methodology), 71(2):319–392.
  * Sanderson et al. (2018) Sanderson, R. E., Wetzel, A., Loebman, S., Sharma, S., Hopkins, P. F., Garrison-Kimmel, S., Faucher-Giguère, C.-A., Kereš, D., and Quataert, E. (2018).  Synthetic Gaia surveys from the FIRE cosmological simulations of Milky-Way-mass galaxies.  arXiv e-prints, page arXiv:1806.10564.
  * Schlegel et al. (1998) Schlegel, D. J., Finkbeiner, D. P., and Davis, M. (1998).  Maps of dust infrared emission for use in estimation of reddening and cosmic microwave background radiation foregrounds.  The Astrophysical Journal, 500(2):525.
  * Snelson and Ghahramani (2006) Snelson, E. and Ghahramani, Z. (2006).  Sparse Gaussian processes using pseudo-inputs.  In Advances in neural information processing systems, pages 1257–1264.
  * Titsias (2009) Titsias, M. (2009).  Variational learning of inducing variables in sparse Gaussian processes.  In Artificial Intelligence and Statistics, pages 567–574.
  * Wainwright et al. (2008) Wainwright, M. J., Jordan, M. I., et al. (2008).  Graphical models, exponential families, and variational inference.  Foundations and Trends® in Machine Learning, 1(1–2):1–305.
  * Wetzel et al. (2016) Wetzel, A. R., Hopkins, P. F., hoon Kim, J., Faucher-Giguère, C.-A., Kereš, D., and Quataert, E. (2016).  Reconciling Dwarf Galaxies with $\Lambda$CMD Cosmology: Simulating a Realistic Population of Satellites Around a Milky Way Mass Galaxy.  The Astrophysical Journal, 827(2):L23.
  * Yang et al. (2012) Yang, T., Li, Y.-F., Mahdavi, M., Jin, R., and Zhou, Z.-H. (2012).  Nyström method vs random fourier features: A theoretical and empirical comparison.  In Advances in neural information processing systems, pages 476–484.

Mapping Interstellar Dust with Gaussian Processes: Supplementary Material

##  Appendix A Covariance Function Definitions

We compare three families of covariance functions in this work: (i) squared
exponential, (ii) Matérn, (ii) and a kernel from Gneiting (Gneiting, 2002).
Each are parameterized with an amplitude scale $\sigma^{2}$ and length scale
parameter $\ell$.

The squared exponential kernel is defined

$\displaystyle k(x,y)$
$\displaystyle=\sigma^{2}\exp\left(-\frac{1}{2\ell^{2}}|x-y|_{2}^{2}\right)$
(29)

The Matern class of kernels

$\displaystyle k(x,y)$
$\displaystyle=\sigma^{2}\frac{2^{1-\nu}}{\Gamma(\nu)}\left(\sqrt{2\nu}\frac{d}{\ell}\right)^{\nu}K_{\nu}\left(\sqrt{2\nu}\frac{d}{\ell}\right)$
(30)

And the Gneiting kernel is defined

$\displaystyle k(\tau)$
$\displaystyle=\sigma^{2}(1+\tau^{\alpha})^{-3}\left((1-\tau)\text{cos}(\tau\cdot\pi)+(1/\pi)\text{sin}(\tau\cdot\pi)\right)$
(31)

for $\tau=|x-y|/\ell$, $\tau<1$.

We graphically depict these covariance functions in Figure 8.

##  Appendix B Integrated Covariance Functions

Below is a simple derivation of the semi-integrated covariance function.

######  Claim B.1 (Semi-integrated Covariance Function).

The covariance between a process value $\rho_{i}\triangleq\rho(x_{i})$ and an
integrated value $e_{j}\triangleq\int_{x\in R_{j}}\rho(x)dx$ takes the form

$\displaystyle\text{Cov}(\rho_{i},e_{j})$ $\displaystyle=\int_{x\in
R_{j}}k(x_{i},x)dx$ (32) $\displaystyle=k^{(semi)}(x_{i},x_{j})\,.$ (33)

For consistency, we will write the integrated argument second.

###### Proof.

Without loss of generality consider a mean zero process $\rho$. For positive
and integrable covariance functions $k(\cdot,\cdot)$ (e.g. the Matérn and
squared exponential) the double integral will be finite. By Fubini’s theorem
we can reverse the order of integration, yielding

$\displaystyle\text{Cov}(\rho_{i},e_{j})$
$\displaystyle=\mathbb{E}\left[\rho_{i}\int_{x\in R_{j}}\rho(x)dx\right]$
$\displaystyle=\mathbb{E}\left[\int_{x\in R_{j}}\rho_{i}\rho(x)dx\right]$
$\displaystyle=\int_{x\in R_{j}}\mathbb{E}\left[\rho_{i}\rho(x)\right]dx$
$\displaystyle=\int_{x\in R_{j}}\text{Cov}(\rho_{i},\rho(x))dx\,=\int_{x\in
R_{j}}k(x_{i},x)dx\,.$

∎

Figure 8: Correlation function for squared exponential, Matérn, and Gneiting
kernels. Here the squared exponential and Matérn kernels have $\sigma^{2}=1$
and $\ell=1$; the Gneiting kernel has $\ell=2$. Note that the Gneiting kernel
goes to zero above $\ell$, while the other two kernels. Further, the steepness
of the function at 0 determines how smooth the GP sample paths will be — the
squared exponential kernel is the smoothest, while the Gneiting $\alpha=.5$ is
the most jagged.

###  B.1 Semi-integrated Squared exponential kernel

The _squared exponential_ covariance function is commonly used for Gaussian
process regression (Rasmussen and Williams, 2006). For $D$-dimensional inputs
$x_{i},x_{j}$, it is

$\displaystyle k_{sqe}(x_{i},x_{j})$
$\displaystyle=\sigma^{2}\exp\left(-\frac{1}{2}(x_{i}-x_{j})^{\intercal}L^{-1}(x_{i}-x_{j})\right),$
(34)

where ${\boldsymbol{\mathbf{\theta}}}=(L,\sigma^{2})$ are the $D\times D$-size
length scale matrix $L$ and the process marginal variance $\sigma^{2}$. For an
isotropic Gaussian process, the length scale matrix $L$ is $\ell\cdot I_{D}$
for a scalar length scale $\ell$. The length scale controls the smoothness of
the function; smaller length scales create more wiggly functions and longer
length scales create more slowly-changing functions. This kernel puts support
on functions with infinitely differentiable sample paths.

The squared exponential kernel admits an analytically tractable form for the
semi-integrated kernel

$\displaystyle k_{\mathrm{sqe}}^{(\mathrm{semi})}(x_{i},x_{j})$
$\displaystyle=\int_{[0,1]}k_{\mathrm{sqe}}(x_{i},\alpha\cdot
x_{j})||x_{j}||_{2}d\alpha$ (35)
$\displaystyle=\frac{1}{2}||x_{j}||_{2}\sigma^{2}\sqrt{\frac{2\pi}{a}}\exp\left(\frac{b^{2}}{2a}-\frac{c}{2}\right)\left[\mathrm{erf}\left(\frac{1-b/a}{\sqrt{2/a}}\right)-\mathrm{erf}\left(\frac{-b/a}{\sqrt{2/a}}\right)\right]$
(36)

where $\text{erf}(\cdot)$ is the error function, and $a$, $b$, and $c$ are
defined

$\displaystyle a=x_{j}^{\intercal}L^{-1}x_{j}\,,\quad\quad
b=x_{i}^{\intercal}L^{-1}x_{j}\,,\quad\quad c=x_{i}^{\intercal}L^{-1}x_{i}$

and the origin, or observer, is at 0.

The doubly-integrated squared exponential covariance function is not as
straightforward to express. Deriving the doubly-integrated squared exponential
covariance function requires integrating the expression in Equation 36.
Integrating the error function multiplied by a quadratic requires either
numerical approximation or a more sophisticated approach (Fayed and Atiya,
2014). We note that, while it is straightforward to analytically derive the
integral of the squared exponential kernel over the _unbounded region_
$[-\infty,\infty]^{2}$ (e.g. to compute a Gaussian normalizing constant), it
is not straightforward to integrate over two-dimensional compact sets, e.g.
$[0,x_{i}]\times[0,x_{j}]$ — this is exactly what is required to compute the
doubly-integrated kernel.

Data: data set $\mathcal{D}\triangleq\\{x_{n},\sigma_{n},a_{n}\\}_{n=1}^{N}$,
test location $x_{*}$, covariance functions: standard $k(\cdot,\cdot)$, semi-
integrated $k^{(\textrm{semi})}(\cdot,\cdot)$, and doubly-integrated
$k^{(\textrm{double})}(\cdot,\cdot)$

Result: $\mu_{*}$, $\sigma_{*}^{2}$ for posterior
$p(\rho(x_{*})|\mathcal{D})=\mathcal{N}(\mu_{*},\sigma_{*}^{2})$

$(\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{e}}})_{i,j}\leftarrow
k^{(\textrm{double})}(x_{i},x_{j})$ for $i,j=1,\dots,N$ ;

// doubly-integrated Gram matrix

$(\boldsymbol{\mathbf{K}}_{*,\boldsymbol{\mathbf{e}}})_{*,j}\leftarrow
k^{(\textrm{semi})}(x_{*},x_{j})$ for $j=1,\dots,N$ ;

// semi-integrated cross covariance

$K_{*,*}\leftarrow k(x_{*},x_{*})$ ;

// prior variance from standard covariance function

$\Sigma\leftarrow\text{diag}(\sigma_{1}^{2},\dots,\sigma_{N}^{2})$

$\mu_{*}\leftarrow\boldsymbol{\mathbf{K}}_{*,\boldsymbol{\mathbf{e}}}(\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{e}}}+\Sigma)^{-1}\boldsymbol{\mathbf{a}}$

$\sigma_{*}^{2}\leftarrow
K_{*}-\boldsymbol{\mathbf{K}}_{*,\boldsymbol{\mathbf{e}}}(\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{e}}}+\Sigma)^{-1}\boldsymbol{\mathbf{K}}_{*,\boldsymbol{\mathbf{e}}}^{\intercal}$

return _$\mu_{*},\sigma_{*}^{2}$_

Algorithm 1 Exact Gaussian process inference with noisy integrated
observations.

##  Appendix C Stochastic Natural Gradients for SVGPs

Following Hensman et al. (2013) and Hoffman et al. (2013), the natural
gradient for variational parameters
${\boldsymbol{\mathbf{\lambda}}}=\boldsymbol{\mathbf{m}}_{{\boldsymbol{\mathbf{\lambda}}}},\boldsymbol{\mathbf{S}}_{{\boldsymbol{\mathbf{\lambda}}}}$
is straightforward to express using the natural parameterization of the
multivariate normal $q_{{\boldsymbol{\mathbf{\lambda}}}}$

$\displaystyle{\boldsymbol{\mathbf{\eta}}}_{1}=\boldsymbol{\mathbf{S}}_{{\boldsymbol{\mathbf{\lambda}}}}^{-1}\boldsymbol{\mathbf{m}}_{{\boldsymbol{\mathbf{\lambda}}}}\,\,\,,\quad{\boldsymbol{\mathbf{\eta}}}_{2}=-\frac{1}{2}\boldsymbol{\mathbf{S}}_{\lambda}^{-1}\,\,\,.$
(37)

The natural gradient of $\mathcal{L}$ with respect to
${\boldsymbol{\mathbf{\eta}}}_{1}$ and ${\boldsymbol{\mathbf{\eta}}}_{2}$ can
be written

$\displaystyle\tilde{\nabla}_{{\boldsymbol{\mathbf{\eta}}}_{2}}\mathcal{L}$
$\displaystyle=-\frac{1}{2}\Lambda+{\boldsymbol{\mathbf{\eta}}}_{2}$ (38)
$\displaystyle\tilde{\nabla}_{{\boldsymbol{\mathbf{\eta}}}_{1}}\mathcal{L}$
$\displaystyle=-\frac{1}{2}\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{u}}}^{-1}\tilde{\boldsymbol{\mathbf{K}}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{\rho}}}\boldsymbol{\mathbf{y}}+{\boldsymbol{\mathbf{\eta}}}_{1}$
(39)

where we define

$\displaystyle\Lambda$
$\displaystyle=\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{u}}}^{-1}\tilde{\boldsymbol{\mathbf{K}}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{\rho}}}\tilde{\boldsymbol{\mathbf{K}}}_{\boldsymbol{\mathbf{\rho}},\boldsymbol{\mathbf{u}}}\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{u}}}^{-1}+\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{u}}}^{-1}$
(40)
$\displaystyle=\sum_{n}\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{u}}}^{-1}\left(\tilde{\boldsymbol{\mathbf{K}}}_{\boldsymbol{\mathbf{u}},n}\tilde{\boldsymbol{\mathbf{K}}}_{n,\boldsymbol{\mathbf{u}}}\right)\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{u}}}^{-1}+\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{u}}}^{-1}\,,$
(41)

and

$\displaystyle\tilde{\boldsymbol{\mathbf{K}}}_{\boldsymbol{\mathbf{\rho}},\boldsymbol{\mathbf{u}}}$
$\displaystyle\triangleq\Sigma^{-1/2}\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{\rho}},\boldsymbol{\mathbf{u}}}$
(42)
$\displaystyle(\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{\rho}},\boldsymbol{\mathbf{u}}})_{n,m}$
$\displaystyle=k(x_{n},\bar{x}_{m})$ (43) $\displaystyle\Sigma$
$\displaystyle\triangleq\text{diag}(\sigma_{1}^{2},\dots,\sigma_{N}^{2})\,.$
(44)

Note that Equation 41 shows that the term $\Lambda$ can be decomposed into a
sum of $N$ data-specific terms. Further, the gradient for
${\boldsymbol{\mathbf{\eta}}}_{1}$ includes the term
$\tilde{\boldsymbol{\mathbf{K}}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{\rho}}}\boldsymbol{\mathbf{y}}=\sum_{n}\tilde{\boldsymbol{\mathbf{K}}}_{\boldsymbol{\mathbf{u}},n}y_{n}$
that also decomposes into a sum over the $N$ data-specific terms. These two
decompositions enable us to compute unbiased mini-batched natural gradients
and optimize the ELBO using stochastic gradient ascent. An additional benefit
of using the natural gradient is that a stochastic gradient update to
${\boldsymbol{\mathbf{\eta}}}_{2}$ will stay in the positive semi-definite
cone. This allows us to directly parameterize
${\boldsymbol{\mathbf{\eta}}}_{2}$ and not its Cholesky factor or symmetric
square root, simplifying our implementation.

##  Appendix D Whitened Parameterization Gradients

The relationship between the whitened model over $\boldsymbol{\mathbf{z}}$ and
the standard model over $\boldsymbol{\mathbf{u}}$ implies the following
relationship between the Gaussian variational parameters

$\displaystyle\boldsymbol{\mathbf{m}}$
$\displaystyle=\boldsymbol{\mathbf{L}}\tilde{\boldsymbol{\mathbf{m}}}$ (45)
$\displaystyle\boldsymbol{\mathbf{S}}$
$\displaystyle=\boldsymbol{\mathbf{L}}\tilde{\boldsymbol{\mathbf{S}}}\boldsymbol{\mathbf{L}}^{\intercal}.$
(46)

Similarly, the natural parameters have the relationship

$\displaystyle{\boldsymbol{\mathbf{\eta}}}_{1}$
$\displaystyle=\boldsymbol{\mathbf{S}}^{-1}\boldsymbol{\mathbf{m}}=(\boldsymbol{\mathbf{L}}^{\intercal})^{-1}\tilde{\boldsymbol{\mathbf{S}}}^{-1}\boldsymbol{\mathbf{L}}^{-1}\boldsymbol{\mathbf{L}}\tilde{\boldsymbol{\mathbf{m}}}=(\boldsymbol{\mathbf{L}}^{\intercal})^{-1}\tilde{{\boldsymbol{\mathbf{\eta}}}}_{1}$
(47) $\displaystyle{\boldsymbol{\mathbf{\eta}}}_{2}$
$\displaystyle=-\frac{1}{2}\boldsymbol{\mathbf{S}}^{-1}=-\frac{1}{2}(\boldsymbol{\mathbf{L}}^{\intercal})^{-1}\tilde{\boldsymbol{\mathbf{S}}}^{-1}\boldsymbol{\mathbf{L}}^{-1}=(\boldsymbol{\mathbf{L}}^{\intercal})^{-1}\tilde{{\boldsymbol{\mathbf{\eta}}}}_{2}\boldsymbol{\mathbf{L}}^{-1}.$
(48)

The natural gradient is invariant to parameterization (Amari, 1982), which
implies that the whitened natural gradient transformation can be derived by
multiplying the standard update by $\boldsymbol{\mathbf{L}}^{\intercal}$

$\displaystyle{\boldsymbol{\mathbf{\eta}}}_{1}^{(t+1)}$
$\displaystyle={\boldsymbol{\mathbf{\eta}}}_{1}^{(t)}+({\boldsymbol{\mathbf{\mu}}}-{\boldsymbol{\mathbf{\eta}}}_{1}^{(t)})$
natural gradient updates (49)
$\displaystyle{\boldsymbol{\mathbf{\eta}}}_{2}^{(t+1)}$
$\displaystyle={\boldsymbol{\mathbf{\eta}}}_{2}^{(t)}+(-\frac{1}{2}\Lambda-{\boldsymbol{\mathbf{\eta}}}_{1}^{(t)})$
(50) $\displaystyle\implies\tilde{{\boldsymbol{\mathbf{\eta}}}}_{1}^{(t+1)}$
$\displaystyle=\tilde{{\boldsymbol{\mathbf{\eta}}}}_{1}^{(t)}+(\boldsymbol{\mathbf{L}}^{\intercal}{\boldsymbol{\mathbf{\mu}}}-\tilde{{\boldsymbol{\mathbf{\eta}}}}_{1}^{(t)})$
whitened natural gradient updates (51)
$\displaystyle\tilde{{\boldsymbol{\mathbf{\eta}}}}_{2}^{(t+1)}$
$\displaystyle=\tilde{{\boldsymbol{\mathbf{\eta}}}}_{2}^{(t)}+(-\frac{1}{2}\boldsymbol{\mathbf{L}}^{\intercal}\Lambda\boldsymbol{\mathbf{L}}-\tilde{{\boldsymbol{\mathbf{\eta}}}}_{2}^{(t)}).$
(52)

This yields a natural interpretation — the whitened sufficient statistics must
undo the inverse covariance structure from the prior before updating the
variational parameters.666Note that the natural parameters are in “information
form” — pre- and post-multiplying by the covariance Cholesky will whiten a
precision matrix.

When it comes time to predict, we will simply transform the whitened
parameters into standard parameters using Equation 46, and proceed with the
usual multivariate Gaussian conditioning. In Section G.1 we show that
optimizing in the whitened space can yield better tuned covariance function
parameters. This empirical evidence supports the practical conclusion that
jointly tuning the variational approximation over $\boldsymbol{\mathbf{z}}$
and ${\boldsymbol{\mathbf{\theta}}}$ is an easier optimization problem than
jointly tuning parameters over $\boldsymbol{\mathbf{u}}$ and
${\boldsymbol{\mathbf{\theta}}}$ due to the stronger dependence between
$\boldsymbol{\mathbf{u}}$ and ${\boldsymbol{\mathbf{\theta}}}$.

##  Appendix E Scaling Integrated Observation GPs: Algorithm Details

###  E.1 Monte Carlo estimators for semi-integrated kernels

For tractable semi-integrated covariance functions (e.g. the squared
exponential kernel) we can easily substitute a closed form computation for
$\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{\rho}},\boldsymbol{\mathbf{u}}}$
in Equation 39. However, when the semi-integrated covariance term is
intractable, directly using natural gradient updates becomes complicated. One
approach is to use numeric integration — the semi-integrated covariance is a
one-dimensional integral that can be approximated with quadrature techniques.
However, when we have $N$ and $M$ very large, solving $N\times M$ numerical
integration problems to populate the cross covariance matrix can be too
expensive to be practical.

Stochastic natural gradient updates employ computationally cheap estimates of
the true gradient to solve the original optimization problem. The operative
question is, how _good_ do these estimates have to be to effectively find the
optimal solution?

Similarly, numeric integration approximations to $k^{(semi)}(\cdot,\cdot)$ are
precise to nearly machine precision, but are computationally expensive.
Analogously we can ask, how _good_ do the semi-integrated covariance
approximations have to be to effectively find the optimal solution?

Pursuing this idea, we propose using a Monte Carlo approximation to the semi-
integrated covariance values — we sample uniformly along the integral path and
average the original covariance values. More formally, we introduce a uniform
random variable $\nu_{R_{n}}$ that takes values along the ray $R_{n}$ from the
origin to $x_{n}$. We can now write the semi-integrated covariance as

$\displaystyle k^{(semi)}(x_{m},x_{n})$ $\displaystyle=\int_{x\in
R_{n}}k(x_{m},x)dx$ (53) $\displaystyle=|R_{n}|\int_{\nu\in
R_{n}}\underbrace{\frac{1}{|R_{n}|}}_{=p(\nu)}k(x_{m},x)dx$ (54)
$\displaystyle=|R_{n}|\,\mathbb{E}_{\nu}\left[k(x_{m},\nu)\right]\,.$ (55)

This form admits a simple Monte Carlo estimator

$\displaystyle\nu^{(1)},\dots,\nu^{(L)}$ $\displaystyle\sim Unif(R_{n})$ (56)
$\displaystyle\hat{k}^{(semi)}(x_{m},x_{n})$
$\displaystyle=\frac{|R_{n}|}{L}\sum_{\ell}k(x_{m},\nu^{(\ell)})\,\,.$ (57)

It is straightforward to show that Equation 57 is an unbiased and consistent
estimator for the true semi-integrated covariance value evaluated at $x_{m}$
and $x_{n}$.

Uniformly distributed grids The unbiased Monte Carlo estimator
$\hat{k}^{(semi)}$ will have some variance about the mean. The quality of an
unbiased estimator is often measured by how low its variance is — the more
precise the better. Independent and identically distributed samples
$\nu^{(\ell)}$ will yield the correct expectation, but may be lower variance
than a more clever sampling scheme. For our application, we uniformly sample
along the ray that traces from the origin $O$ to the observation location
$x_{n}$. Intuitively, drawing two uniform random variables that are very close
to one another will carry redundant information (for smooth functions
$k(\cdot,\cdot)$). Spreading our samples out to cover more of the length of
the ray may yield more precise estimates, provided that we can construct a
sample that yields the correct expectation.

One inexpensive way to generate a set of identically distributed (but not
independent samples) is to lay down a random grid of points. Consider the
sampling procedure for $L$ samples along the unit interval

$\displaystyle\nu^{(1)}$ $\displaystyle\sim Unif(0,1/L)$ (58)
$\displaystyle\nu^{(\ell)}$
$\displaystyle=\nu^{(\ell-1)}+\tfrac{1}{L}\quad\text{ for }\ell=2,\dots,L\,.$
(59)

This sampling scheme will generate a set of $L$ correlated samples that
constructs a randomly placed grid within the unit interval. If we randomly
permute the variables $\nu_{1},\dots,\nu_{L}$ (or, equivalently, randomly
select one of them) then the marginal distribution of the random variable will
be $Unif(0,1)$ — this ensures that the expectation of the estimator in
Equation 25 will remain unbiased. However, because the samples are anti-
correlated, the variance of the estimator can be substantially reduced.

Biased gradient estimates The Monte Carlo estimate of
$\hat{k}^{(semi)}(x_{m},x_{n})$ is unbiased, but non-linear functions of an
unbiased estimator, in general, may jettison this property. In our
application, the goal is to compute cheap, unbiased estimates of the gradient
of the variational parameters
$\boldsymbol{\mathbf{m}},\boldsymbol{\mathbf{S}}$ (or their information form
parameterization
${\boldsymbol{\mathbf{\eta}}}_{1},{\boldsymbol{\mathbf{\eta}}}_{2}$).

For a single term $n$, the gradient for $\eta_{1}$ is a linear function of the
semi-integrated covariance vector
$\boldsymbol{\mathbf{K}}_{n,\boldsymbol{\mathbf{u}}}=\left[k^{(semi)}(x_{n},x_{1}),\dots,k^{(semi)}(x_{n},x_{M})\right]$.777Note
that the order of the subscripts indicates the vector dimensions, e.g.
$\boldsymbol{\mathbf{K}}_{n,\boldsymbol{\mathbf{u}}}\in\mathbb{R}^{1\times M}$
and
$\boldsymbol{\mathbf{K}}_{n,\boldsymbol{\mathbf{u}}}^{\intercal}=\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},n}\in\mathbb{R}^{M\times
1}$. Due to this linearity, plugging in an unbiased estimator,
$\hat{\boldsymbol{\mathbf{K}}}_{n,\boldsymbol{\mathbf{\rho}}}$, for
$\boldsymbol{\mathbf{K}}_{n,\boldsymbol{\mathbf{u}}}$ will form an unbiased
gradient estimate for $\tilde{\nabla}_{\eta_{1}}\mathcal{L}$. However,
plugging in $\hat{\boldsymbol{\mathbf{K}}}_{n,\boldsymbol{\mathbf{u}}}$ into
the gradient term for $\eta_{2}$ will form a non-linear function of an
unbiased estimator, resulting in a biased estimate of
$\tilde{\nabla}_{\eta_{2}}$.

To see this, consider the vector of unbiased estimates of the semi-integrated
covariance

$\displaystyle\hat{\boldsymbol{\mathbf{K}}}_{n,\boldsymbol{\mathbf{\rho}}}$
$\displaystyle=\left[\hat{k}^{(semi)}(x_{1},x_{n}),\dots,\hat{k}^{(semi)}(x_{m},x_{n})\right]\,.$
(60)

The stochastic component of the gradient for $\eta_{2}$ can be written

$\displaystyle\mathbb{E}\left[\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{u}}}^{-1}\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},n}\boldsymbol{\mathbf{K}}_{n,\boldsymbol{\mathbf{u}}}\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{u}}}^{-1}\right]$
$\displaystyle=\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{u}}}^{-1}\mathbb{E}\left[\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},n}\boldsymbol{\mathbf{K}}_{n,\boldsymbol{\mathbf{u}}}\right]\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{u}}}^{-1}$
(61)

where
$\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{u}}}^{-1}$
can be compute exactly, leaving the expectation of the outer product of the
semi-integrated covariance terms
$\mathbb{E}\left[\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},n}\boldsymbol{\mathbf{K}}_{n,\boldsymbol{\mathbf{u}}}\right]$
Given an unbiased estimator for
$\boldsymbol{\mathbf{K}}_{n,\boldsymbol{\mathbf{u}}}$ the naive plugin
estimate for
$\boldsymbol{\mathbf{K}}_{n,\boldsymbol{\mathbf{u}}}\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},n}$
is

$\displaystyle\hat{O}^{(plugin)}=\hat{\boldsymbol{\mathbf{K}}}_{\boldsymbol{\mathbf{u}},n}\hat{\boldsymbol{\mathbf{K}}}_{n,\boldsymbol{\mathbf{u}}}$
(62)

which is _biased_ in general

$\displaystyle\mathbb{E}\left[\hat{\boldsymbol{\mathbf{K}}}_{\boldsymbol{\mathbf{u}},n}\hat{\boldsymbol{\mathbf{K}}}_{n,\boldsymbol{\mathbf{u}}}\right]$
$\displaystyle\neq\mathbb{E}\left[\hat{\boldsymbol{\mathbf{K}}}_{\boldsymbol{\mathbf{u}},n}\right]\mathbb{E}\left[\hat{\boldsymbol{\mathbf{K}}}_{n,\boldsymbol{\mathbf{u}}}\right]=\boldsymbol{\mathbf{K}}_{n,\boldsymbol{\mathbf{u}}}\boldsymbol{\mathbf{K}}_{n,\boldsymbol{\mathbf{u}}}^{\intercal}\,.$
(63)

Fortunately, we can construct an unbiased estimator for this outer product.
Note that the outer product of an expectation can be expressed as

$\displaystyle\mathbb{E}\left[\boldsymbol{\mathbf{x}}\right]\mathbb{E}\left[\boldsymbol{\mathbf{x}}\right]^{\intercal}$
$\displaystyle=\mathbb{E}\left[\boldsymbol{\mathbf{x}}\boldsymbol{\mathbf{x}}^{\intercal}\right]-\text{Cov}(\boldsymbol{\mathbf{x}})$
(64)

Given an $L$-size sample
$\boldsymbol{\mathbf{x}}_{1},\dots,\boldsymbol{\mathbf{x}}_{L}$ we can compute
unbiased estimators for $\mu=\mathbb{E}[\boldsymbol{\mathbf{x}}]$,
$M\triangleq\mathbb{E}\left[\boldsymbol{\mathbf{x}}\boldsymbol{\mathbf{x}}^{\intercal}\right]$,
and $\Sigma\triangleq\text{Cov}(\boldsymbol{\mathbf{x}})$

$\displaystyle\hat{\mu}$
$\displaystyle=\frac{1}{L}\sum_{\ell}\boldsymbol{\mathbf{x}}_{\ell}$ (65)
$\displaystyle\hat{\Sigma}$
$\displaystyle=\frac{1}{L-1}\sum_{\ell}(\boldsymbol{\mathbf{x}}_{\ell}-\hat{\mu})(\boldsymbol{\mathbf{x}}_{\ell}-\hat{\mu})^{\intercal}$
(66) $\displaystyle\hat{M}$
$\displaystyle=\frac{1}{L}\sum_{\ell}\boldsymbol{\mathbf{x}}_{\ell}\boldsymbol{\mathbf{x}}_{\ell}^{\intercal}$
(67)

With these estimators, $\hat{M}-\hat{\Sigma}$ forms an unbiased estimate for
$\mathbb{E}[\boldsymbol{\mathbf{x}}]\mathbb{E}[\boldsymbol{\mathbf{x}}]^{\intercal}$

$\displaystyle\mathbb{E}\left[\hat{M}-\hat{\Sigma}\right]$
$\displaystyle=\mathbb{E}[\boldsymbol{\mathbf{x}}]\mathbb{E}[\boldsymbol{\mathbf{x}}]^{\intercal}\,.$
(68)

We can use this fact to construct a de-biased estimator for $O$

$\displaystyle\hat{O}^{(unbiased)}$
$\displaystyle=\hat{M}\left(\\{\hat{\boldsymbol{\mathbf{K}}}^{(\ell)}_{n,\boldsymbol{\mathbf{u}}}\\}_{\ell=1}^{L}\right)-\hat{\Sigma}\left(\\{\hat{\boldsymbol{\mathbf{K}}}^{(\ell)}_{n,\boldsymbol{\mathbf{u}}}\\}_{\ell=1}^{L}\right)\,.$
(69)

Though $\hat{M}-\hat{\Sigma}$ forms an unbiased and symmetric estimator of the
outer product, this estimator may not be positive semi-definite. If this
estimator is not positive semi-definite, then a gradient step along this
direction may take the optimization routine outside of the PSD cone. To ensure
safe gradient updates, we will have to meticulously track the magnitude of the
lowest eigenvalues of this matrix.

Though biased, the naive plugin estimator in Equation 62 will always be non-
negative definite, resulting in gradient updates that are guaranteed to remain
in the PSD cone. Proceeding with this estimator, we must ask, how much bias is
too much bias? We explore this question empirically in Section F, comparing
the optimization traces of algorithms using various sample sizes $L$ to the
exact $k^{(semi)}(\cdot,\cdot)$ values in the squared exponential kernel case.
We see that, though biased, this gradient estimator can closely approximate
the performance of the exact semi-integrated kernel given a modest number of
samples ($\approx 20$).

Data: $\mathcal{D}\triangleq\\{a_{n},x_{n},\sigma^{2}_{n}\\}_{n=1}^{N}$
(data), $k^{({\boldsymbol{\mathbf{\theta}}})}(\cdot,\cdot)$ (covariance
function);
$\left({\boldsymbol{\mathbf{\lambda}}}^{(0)},{\boldsymbol{\mathbf{\theta}}}^{(0)}\right)$
(initial values); $k(\cdot,\cdot)$ (covariance function);
$(\Delta_{{\boldsymbol{\mathbf{\lambda}}}},\Delta_{{\boldsymbol{\mathbf{\theta}}}})$
(step-sizes); $B$ (batch size); $L$ (Monte Carlo semi-integrated estimator
sample size);

Result: ${\boldsymbol{\mathbf{\lambda}}},{\boldsymbol{\mathbf{\theta}}}$
variational and covariance function parameters

$\hat{k}^{(dd)}(\cdot)\leftarrow\text{Interpolate-Doubly-
Diag}(k^{({\boldsymbol{\mathbf{\theta}}}_{0})}(\cdot,\cdot))$ ;

// construct diagonal interpolator

for  _$t\leftarrow 1$ to $t_{max}$_ do

$\mathcal{D}_{t}\leftarrow\text{Batch}(\mathcal{D},B)$ ;

// mini-batch of data

${\boldsymbol{\mathbf{\lambda}}}^{(t)}\leftarrow{\boldsymbol{\mathbf{\lambda}}}^{(t-1)}+\Delta_{{\boldsymbol{\mathbf{\lambda}}}}\cdot\tilde{\nabla}^{(MC)}_{{\boldsymbol{\mathbf{\lambda}}}}\mathcal{L}({\boldsymbol{\mathbf{\lambda}}}^{(t-1)},{\boldsymbol{\mathbf{\theta}}}^{(t-1)},\mathcal{D}_{t},L)$
;

// MC natural grad update

${\boldsymbol{\mathbf{\theta}}}^{(t)}\leftarrow{\boldsymbol{\mathbf{\theta}}}^{(t-1)}+\Delta_{{\boldsymbol{\mathbf{\theta}}}}\cdot{\nabla}^{(MC)}_{{\boldsymbol{\mathbf{\theta}}}}\mathcal{L}({\boldsymbol{\mathbf{\lambda}}}^{(t-1)},{\boldsymbol{\mathbf{\theta}}}^{(t-1)},\mathcal{D}_{t},L,\hat{k}^{(dd)}(\cdot))$
;

// MC kernel grad update

$\text{sgd-
update}({\boldsymbol{\mathbf{\lambda}}}^{(t)},\tilde{\nabla}{{\boldsymbol{\mathbf{\lambda}}}}^{(MC)})$
;

// gradient-based update

end for

return
${\boldsymbol{\mathbf{\lambda}}}^{(t_{max})},{\boldsymbol{\mathbf{\theta}}}^{(t_{max})}$
;

// optimal parameters

Algorithm 2 Scalable integrated observation GP inference with mini-batches and
Monte Carlo semi-integrated covariance estimates. The natural gradient
estimator uses $L$ Monte Carlo samples to approximate the semi-integrated
covariance and $B$ data observations to approximate the complete data
objective.

Monte Carlo SVGP comparison

Figure 9: Monte Carlo estimates of semi-integrated covariance functions enable
efficient ELBO inference competitive with the analytic semi-integrated
covariance. Above is a comparison of ELBO traces by mini-batch ($x$-axis) for
the squared exponential kernel. The analytic semi-integrated kernel was used
in blue (hidden behind the $L=50$ purple line. The optimal value is depicted
with the dotted gray line.

The squared exponential kernel admits an analytic form for the semi-integrated
covariance function. This gives us a way to directly compare our Monte Carlo
semi-integrated approximation within SVGP to the same algorithm using exact
semi-integrated covariance values. In this experiment, we compare the progress
of the exact variational objective throughout optimization using different
gradient estimators — the analytic version, and the biased Monte Carlo
gradients using sample size $L=5,7,10,20,$ and $50$.

We fit variational parameters using stochastic natural gradient descent with
mini-batches of size $1{,}000$ and step size $0.1$. We fix the covariance
function parameters in this experiment to report a more direct comparison
between exact and Monte Carlo gradients — in subsequent experiments we apply
gradient updates to these two parameters as well. Inducing points are placed
on a fixed grid of size $20x20$ ($M=400$), tiling the 2-d space.

At each batch, we compute the complete-data ELBO objective (using the exact
squared exponential semi-integrated covariance function across methods).
Additionally, we compute the optimal value of the ELBO (by solving for the
optimal values of $\boldsymbol{\mathbf{m}}_{{\boldsymbol{\mathbf{\lambda}}}}$
and $\boldsymbol{\mathbf{S}}_{{\boldsymbol{\mathbf{\lambda}}}}$ in closed form
— an operation that is too costly when $N$ is on the order of millions of
observations) which we depict alongside optimization traces.

Figure 9 compares the progress of optimizers using different gradient
estimators. The “analytic” version uses the closed form semi-integrated
covariance, while the Monte Carlo runs use sample sizes of varying values. We
see that with as few as $L=20$ samples the resulting optimization routine very
nearly finds the optimal value (depicted in grey), and almost exactly matches
the path of the “analytic” routine. This confirms that these Monte Carlo semi-
integrated covariance estimators are suitable plug-in estimators for fitting
models with covariance functions that do not admit an analytic form for the
semi-integrated covariance, (e.g. the Gneiting and Matérn kernels).

###  E.2 Fast doubly-integrated kernel diagonal interpolation

Approximating doubly-integrated diagonal terms To make predictions within the
SVGP framework, we need to compute the marginal variance of each observation
according to the Gaussian process prior. Concretely, this amounts to computing
the $\boldsymbol{\mathbf{K}}_{*,*}$ term in Equation 81. In the integrated
observation setting, this requires computing the variance of an integrated
observation

$\displaystyle(\tilde{}\boldsymbol{\mathbf{K}}_{N})_{n,n}=Cov(a_{n},a_{n})=Var(e_{n})=Var\left(\int_{[0,1]}\rho(\alpha\cdot
x_{n})d\alpha\right)\,.$ (70)

This marginal variance corresponds to the diagonal of the doubly-integrated
kernel and may not be analytically tractable. For the squared exponential
kernel, this would require integrating the semi-integrated kernel expressed in
Equation 36 over a compact set, which does not yield a closed form (Fayed and
Atiya, 2014).

Fortunately, for a one-dimensional compact set, this integral can be
numerically computed relatively efficiently. Unfortunately, if we are scaling
to millions of observations, we will have to perform millions of numerical
integrations to simply evaluate the likelihood of the data under one setting
of the covariance function parameters. Further, we will have to recompute all
of these numerical integrations at each iteration when fitting covariance
function parameters (or computing their posterior).

The inducing point variance correction method requires only computing the
marginal variance of each integrated observation under the prior. For
stationary kernels this is a function only of the distance of the observation
to the origin (or more generally, the length of the integrated ray). Similar
to the semi-integrated kernel, we can see that this is true by expanding the
stationary kernel

$\displaystyle\int_{0}^{1}k(|x-\alpha x|)d\alpha$
$\displaystyle=\int_{0}^{1}k((x^{\intercal}x+\alpha^{2}x^{\intercal}x-2\alpha
x^{\intercal}x)^{1/2})d\alpha$ (71)
$\displaystyle=\int_{0}^{1}k((|x|^{2}+\alpha^{2}|x|^{2}-2\alpha|x|^{2})^{1/2})d\alpha$
(72) $\displaystyle=f(|x|)\,.$ (73)

Due to this stationarity, we can view each numerical computation as evaluating
the function $k_{\theta}(\boldsymbol{\mathbf{d}}_{n})=\tilde{K}_{n}n$ where
$\boldsymbol{\mathbf{d}}_{n}$ is the distance of observation
$\boldsymbol{\mathbf{x}}_{n}$ from the origin. This perspective suggests a
numerically cheap approximation — simply estimate this function along a fixed
grid of distances (that spans the range of observed distances) and interpolate
the value for each observation.

This one-dimensional function can be estimated independent of the length scale
parameter for isotropic and stationary covariance functions. The effect of the
length scale parameter is to effectively redefine distance, or the magnitude
of $|x|$. We can account for this change by re-scaling the $|x|$ distance and
multiplying the result by $\ell$. Specifically for the diagonal, this becomes
simply a function of a one-dimensional distance, which can easily be
interpolated with a small number of knots.

Making predictions. To estimate the process value $\rho_{*}$ at new location
$x_{*}$, we use the posterior predictive distribution
$p(\rho_{*}\,|\,\mathcal{D})$. Within the SVGP framework, we approximate this
distribution using the posterior
$q_{{\boldsymbol{\mathbf{\lambda}}}}(\boldsymbol{\mathbf{u}})$

$\displaystyle p(\rho_{*}\,|\,\mathcal{D})$ $\displaystyle=\int
p(\rho_{*}\,|\,\boldsymbol{\mathbf{u}})p(\boldsymbol{\mathbf{u}}\,|\,\mathcal{D})d\boldsymbol{\mathbf{u}}$
(74) $\displaystyle\approx\int
p(\rho_{*}\,|\,\boldsymbol{\mathbf{u}})q_{{\boldsymbol{\mathbf{\lambda}}}}(\boldsymbol{\mathbf{u}})d\boldsymbol{\mathbf{u}}\,\,.$
(75)

Under the approximation, $\boldsymbol{\mathbf{u}}$ and $\rho_{*}$ are jointly
normal, making this predictive distribution available in closed form

$\displaystyle p(\rho_{*}\,|\,\mathcal{D})$
$\displaystyle=\mathcal{N}(\mu_{*},\sigma^{2}_{*})$ (76)
$\displaystyle\mu_{*}$
$\displaystyle=\boldsymbol{\mathbf{K}}_{*,\boldsymbol{\mathbf{u}}}\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{u}}}^{-1}\boldsymbol{\mathbf{m}}$
(77) $\displaystyle\sigma_{*}^{2}$
$\displaystyle=\tilde{\boldsymbol{\mathbf{K}}}_{*,*}+\boldsymbol{\mathbf{K}}_{*,\boldsymbol{\mathbf{u}}}\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{u}}}^{-1}\boldsymbol{\mathbf{S}}\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{u}}}^{-1}\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},*}$
(78)

where

$\displaystyle(\boldsymbol{\mathbf{K}}_{*,\boldsymbol{\mathbf{u}}})_{m}$
$\displaystyle=\text{Cov}(\rho_{*},\boldsymbol{\mathbf{u}}_{m})$ (79)
$\displaystyle\boldsymbol{\mathbf{K}}_{*,*}$
$\displaystyle=\text{Var}(\rho_{*})$ (80)
$\displaystyle\tilde{\boldsymbol{\mathbf{K}}}_{*,*}$
$\displaystyle\triangleq\boldsymbol{\mathbf{K}}_{*,*}-\boldsymbol{\mathbf{K}}_{*,\boldsymbol{\mathbf{u}}}\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},\boldsymbol{\mathbf{u}}}^{-1}\boldsymbol{\mathbf{K}}_{\boldsymbol{\mathbf{u}},*}\,.$
(81)

Prediction for integrated process values, ${e_{*}=\int_{x\in
R_{*}}\rho(x)dx}$, can be computed similarly by substituting semi-integrated
covariance values into the vector
$\boldsymbol{\mathbf{K}}_{*,\boldsymbol{\mathbf{u}}}$ and the doubly
integrated covariance value $\boldsymbol{\mathbf{K}}_{*,*}$.

(a) True $\rho(x)$

(b) Integrated observations $a_{n}$

(c) Inferred $\mathbb{E}\left[\rho\,|\,\mathcal{D}\right]$.

(d) Inferred posterior scale
$\mathbb{V}\left[\rho(x)\,|\,\mathcal{D}\right]^{1/2}$.

Figure 10: Integrated observation GPs can reveal spatial structure from a
severely limited vantage point. This figure summarizes synthetic data
posterior inference using $N=1{,}000{,}000$ examples in a two-dimensional
space. Panel 10(a) depicts the true (unobserved) $\rho(x)$ that generates the
data. Panel 10(b) depicts the observed data, noisy integrated measurements of
$\rho$ from the origin to uniformly random points in the square. 10(c) depicts
the inferred $\rho$ given these observations using ziggy with a squared
exponential kernel. 10(d) depicts the marginal posterior standard deviation at
each location in $\mathcal{X}$-space.

##  Appendix F Validation on Synthetic Data

In this section we study the performance of ziggy on synthetic data under
varying conditions. To empirically evaluate ziggy we measure the quality of
its statistical inferences on a two-dimensional synthetic example. To generate
synthetic data, we first define the unobserved
${\rho:\mathbb{R}^{2}\mapsto\mathbb{R}}$ and generate noise with known
variance. The form we chose for $\rho$ is

$\displaystyle\rho(x)$
$\displaystyle=4+\sum_{d=1}^{d}x_{d}\cdot\sin(2x_{d}^{2})\,,$ (82)

which is visualized in Figure 10(a). This test function has clear regions of
high density far from the observation point (e.g. around $x=(-1.5,-1.5)$) and
it is positive everywhere. We generate a dataset of $N$ noisy integrated
observations

$\displaystyle x_{n}$ $\displaystyle\sim\text{Unif}(\mathcal{X})$ simulated
star locations (83) $\displaystyle e_{n}$ $\displaystyle=\int_{x\in
R_{n}}\rho(x)dx$ simulated extinctions (84) $\displaystyle a_{n}$
$\displaystyle\sim\mathcal{N}(e_{n},\sigma_{n}^{2})$ noisy integrated
observation (85)

where the domain is $\mathcal{X}=[-2,2]^{2}$, and the noise variance is chosen
to be $\sigma_{n}^{2}=4$ to roughly approximate the noise level of extinction
measurements. The true extinctions, $e_{n}$, are computed to high precision
using numerical quadrature, integrating from the origin to the random location
$x_{n}$. The integrated observations are depicted in Figure 10(b). We estimate
this synthetic $\rho(\cdot)$ from noisy integrated observations using the
approach developed in the previous section. In this experiment, we fix a grid
of $20\times 20$ inducing points, evenly spaced from $[-2,2]$ in each of the
two dimensions.

In Section F.1, the first set of simulations compares posterior estimates of
$\rho(\cdot)$ as a function of training data set size. Posterior estimates
continue to become more accurate and well-calibrated even as the number of
integrated observation $N$ grows from one hundred thousand to one million.
This accuracy highlights the importance of scaling inference to more
observations in order to obtain a more accurate statistical estimator.

In Section F.2, we compare posterior estimate quality with a fixed dataset
over different kernels, each encoding different assumptions about the
underlying function. This experiment tests the Monte Carlo approach that can
incorporate kernels that do not admit an analytic form for the semi-integrated
kernel. Finally, Appendix G describes additional synthetic experiments that
examine the effect of two algorithm hyper parameter choices on variational
objective optimization, the whitened parameterization and the number of Monte
Carlo semi-integrated covariance samples.

A python and PyTorch implementation is available.888Available at [REDACTED TO
PRESERVE AUTHOR BLINDNESS]. Code to reproduce experiments are included in the
repository.

###  F.1 The quality of the estimate

We study the quality of the estimate of $\rho(x)$ as a function of data set
size $N$. We fit the dust model to synthetic observations with data set sizes
$N\in\\{10^{3},10^{4},10^{5},10^{6}\\}$. We measure model quality by computing
root mean squared error (RMSE) and log likelihood (LL) on a set of
$N_{\mathrm{test}}=2{,}000$ held out extinction values

$\displaystyle\mathrm{RMSE}$
$\displaystyle=\left(\frac{1}{N_{\mathrm{test}}}\sum_{n=1}^{N_{\mathrm{test}}}\left(\mathbb{E}\left[e_{*}\,|\,\mathcal{D}\right]-e_{*}\right)^{2}\right)^{1/2}$
(86) $\displaystyle\mathrm{LL}$
$\displaystyle=\frac{1}{N_{\mathrm{test}}}\sum_{n=1}^{N_{\mathrm{test}}}\ln
p(e_{n}\,|\,\mathcal{D})\,\,.$ (87)

Each model uses $M=20\times 20$ inducing points, evenly spaced in a grid in
the input space. We trained each model until convergence, using batches of
size $1{,}000$, a step size of $1e-3$ and a kernel parameter step size of
$1e-6$. (Additional optimization details can be found in the implementation.)

Figure 10 summarizes the model’s fit using one million observations, the most
we tried in this example. Figure 10(c) displays the posterior mean for
$\rho(x)$ given the one million observations depicted in 10(b). Figure 10(d)
shows the posterior uncertainty (one standard deviation) about the estimated
mean. The million-observation model recovers the true latent function well.

Figure 11 summarizes predictive quality of ziggy as a function of data set
size on held out test data using RMSE and LL. This illustrates the benefit of
including more training data — root mean squared error (RMSE) decreases (and
log likelihood increases) significantly as data set size grows. For example,
the RMSE for $N=100{,}000$ is $2.6$ times bigger than the RMSE for
$N=1{,}000{,}000$, (.082 vs. .031) and the average log likelihood is
significantly better for the million-observation model.

Figure 12 makes a more direct graphical comparison of the posterior estimate
of the latent function as we increase data set size $N$. As we incorporate
more data into the nonparametric model, the form of the true underlying
function emerges. To accurately visualize and interpret large scale features
of the latent dust distribution, we will want to incorporate as many stellar
observations as possible.

(a) Test $e$ RMSE

(b) Test $e$ log likelihood

Figure 11: Scaling to massive $N$ improves estimator performance. Above, we
compare predictive RMSE (left) and log likelihood (right) as a function of
data set size $N$ on a held out sample of test stars.

(a) $\mathbb{E}[\rho\,|\,\mathcal{D}]$, $N=1{,}000$

(b) $\mathbb{E}[\rho\,|\,\mathcal{D}]$, $N=10{,}000$

(c) $\mathbb{E}[\rho\,|\,\mathcal{D}]$, $N=100{,}000$

(d) $\mathbb{E}[\rho\,|\,\mathcal{D}]$, $N=1{,}000{,}000$

Figure 12: More data leads to more accurate function posterior predictions.

###  F.2 Comparing kernels

(a) Test $e$ MSE

(b) Test $e$ log likelihood

(c) Test $e$ $\chi^{2}$

(d) Test $\rho(x)$ MSE

(e) Test $\rho(x)$ log likelihood

(f) Test $\rho(x)$ $\chi^{2}$

Figure 13: Predictive summaries across different models and inference schemes.

We compared multiple models by fitting the variational approximation and
tuning the covariance function parameters for five different kernel choices.
In this synthetic setup, we ran each model for 50 epochs, saving the the model
with the best ELBO value. We used mini-batches of size $|B|=1{,}000$, and
started the step size at $.1$, reducing it every epoch by multiplying by a
decay factor of .95.

For kernels that do not have a closed form semi-integrated version, we used
$L=30$ uniform grid Monte Carlo samples to estimate the integrated covariance.
We chose this on the conservative end — in Appendix E.1 shows that on the
order of $L=10$ to $20$ samples is often just as effective as the analytic
kernel in the squared exponential case.

We validate the inferences against a set of $2{,}000$ held-out test
observations, which include ground truth extinctions $e_{*}$ and pointwise
values, $\rho_{*}$. We quantify the quality of the inferences with mean
squared error, log likelihood, and the $\chi^{2}$-statistic on the test data.
We also visually validate ziggy with a $QQ$-plot and coverage comparison, also
on test data. We also visualize the behavior ziggy as a function of distance
to the origin — specifically, we visualize how accurate and well-calibrated
test inferences are as synthetic observations are made farther away.

In Figure 13 we compare the mean squared error (MSE), the log likelihood, and
the $\chi^{2}$ statistic on the held out test data. We compare five different
kernels – (i) Gneiting $\alpha=1$, (ii-iv) Matérn, $\nu\in\\{1/2,3/2,5/2\\}$,
and (v) the squared exponential kernel. The results match our expectations.
The true $\rho(\cdot)$ is smooth, and the smoother kernels tend to have lower
predictive error, higher log likelihood, and better calibrated $\chi^{2}$
statistic (closer to one). The test statistic comparison has chosen a good
match in the squared exponential kernel.

We also note that the MSE for the process evaluations $\rho_{*}$ is much
higher than the test MSE for extinctions $e_{*}$. This is also expected —
intuitively, if the $\rho(\cdot)$ predictions are well-calibrated, then over a
larger distance the pointwise errors can be averaged out, resulting in a
better predictor for extinctions.

A crucial property of a good posterior approximation is well-calibrated
uncertainty. To evaluate the quality of the uncertainty measurements for
unobserved extinctions and pointwise process values, we inspect statistics of
test-sample $z$-scores

$\displaystyle z_{*}$
$\displaystyle=\frac{e_{*}-\mathbb{E}[e_{*}\,|\,x_{*},\mathcal{D}]}{\sqrt{\mathbb{V}\left[e_{*}\,|\,x_{*},\mathcal{D}\right]}}$
(88)