# To Sample or Not To Sample:
Retrieving Exoplanetary Spectra with Variational Inference and Normalising
Flows

Kai Hou Yip Department of Physics and Astronomy
University College London
Gower Street,WC1E 6BT London, United Kingdom  Quentin Changeat Department of
Physics and Astronomy
University College London
Gower Street,WC1E 6BT London, United Kingdom  European Space Agency (ESA),
ESA Office, Space Telescope Science Institute (STScI),
3700 San Martin Drive, Baltimore MD 21218, USA  Ahmed Al-Refaie Department of
Physics and Astronomy
University College London
Gower Street,WC1E 6BT London, United Kingdom  Ingo P. Waldmann Department of
Physics and Astronomy
University College London
Gower Street,WC1E 6BT London, United Kingdom

###### Abstract

Current endeavours in exoplanet characterisation rely on atmospheric retrieval
to quantify crucial physical properties of remote exoplanets from
observations. However, the scalability and efficiency of said technique are
under strain with increasing spectroscopic resolution and forward model
complexity. The situation becomes more acute with the recent launch of the
James Webb Space Telescope and other upcoming missions. Recent advances in
Machine Learning provide optimisation-based Variational Inference as an
alternative approach to perform approximate Bayesian Posterior Inference. In
this investigation we developed a Normalising Flow-based neural network,
combined with our newly developed differentiable forward model, Diff-$\tau$,
to perform Bayesian Inference in the context of atmospheric retrievals. Using
examples from real and simulated spectroscopic data, we demonstrate the
advantages of our proposed framework: 1) Training our neural network does not
require a large pre-comupted training set and can be trained with only a
single observation; 2) It produces high-fidelity posterior distributions in
excellent agreement with sampling-based retrievals and; 3) It requires up to
75% fewer forward model calls to converge to the same result. 4.) This
approach allows formal Bayesian model selection. We discuss the computational
efficiencies of Diff-$\tau$ in relation to TauREx3’s nominal forward model and
provide a lesson’s learned account of developing radiative transfer models in
differentiable languages. Our proposed framework contributes towards the
latest development of a neural-powered atmospheric retrieval. Its flexibility
and significant reduction in forward model calls required for convergence,
holds the potential to be an important addition to the retrieval tool box for
large and complex data sets along with sampling-based approaches.

††software: TauREx 3 (Al-Refaie et al., 2021, 2022), Tensorflow2 (Abadi et
al., 2015), tensorflow-probability (Dillon et al., 2017), scipy (Virtanen et
al., 2020), NumPy (Harris et al., 2020), Pandas (pandas development team,
2020),nestle (Barbary, 2021), corner (Foreman-Mackey, 2016), matplotlib
(Hunter, 2007), h5py (Collette, 2013)

##  1 Introduction

Atmospheric retrieval has become an indispensable tool for astronomers to
explain individual observations from transit, eclipse and phase curves
spectroscopy in both low (e.g. Tinetti et al., 2007; Line et al., 2013;
Kreidberg et al., 2014; Line et al., 2014; Lee et al., 2014; Haynes et al.,
2015; Evans et al., 2016; Tsiaras et al., 2016a; Line et al., 2016; Tsiaras et
al., 2016b; Sheppard et al., 2017; Stevenson et al., 2017; MacDonald &
Madhusudhan, 2017; Kreidberg et al., 2018; Tsiaras et al., 2019; Mikal-Evans
et al., 2019; MacDonald & Madhusudhan, 2019; Skaf et al., 2020; Pluriel et
al., 2020b; Zhang et al., 2020; Changeat & Al-Refaie, 2020; Chubb et al.,
2020; Alam et al., 2020; von Essen et al., 2020; Anisman et al., 2020; Carone
et al., 2021; Yip et al., 2021a; Edwards et al., 2021; Changeat & Edwards,
2021; Sheppard et al., 2021; Saba et al., 2022; Alam et al., 2021; Mugnai et
al., 2021; Swain et al., 2021; Changeat et al., 2021; Foote et al., 2022;
Mansfield et al., 2022; Changeat, 2022; Mikal-Evans et al., 2022, and
references therein) and high-resolution (e.g. Brogi & Line, 2019; Gibson et
al., 2020; Mollière et al., 2020; Seidel et al., 2020; Boucher et al., 2021;
MacDonald & Lewis, 2022; Challener & Rauscher, 2022; Meech et al., 2022;
Rasmussen et al., 2022; Harrington et al., 2022). Over the years, the
community has come up with a variety of retrieval frameworks, each coupled
with different modelling assumptions and sampling techniques (e.g. Irwin et
al., 2008; Madhusudhan & Seager, 2009; Line et al., 2013; Lavie et al., 2017;
Gandhi et al., 2019; Zhang et al., 2019; Lothringer & Barman, 2020; Min et
al., 2020; Cubillos & Blecic, 2021; Al-Refaie et al., 2021). As the number of
spectroscopic observations increases with the advent of new space and ground-
based observatories, the community has started to look at planetary
characterisation on a population level (Sing et al., 2016; Barstow et al.,
2017; Tsiaras et al., 2019; Pinhas et al., 2019; Mansfield et al., 2021;
Roudier et al., 2021; Changeat et al., 2022; Edwards et al., 2023).

At its core, atmospheric retrieval strives to find an atmospheric model that
can best explain a given observation. Most contemporary retrieval frameworks
formulated the inverse problem in terms of Bayesian statistics, where free
parameters of the physical model are framed as random variables. The
probability densities of these random variables ($\theta$) given the observed
data ($D$) are collectively referred to as the posterior distribution,
$p(\theta|D)$. The Bayes’ Theorem provides a way to calculate the posterior
distribution via the following relation:

$p(\theta|D)=\frac{p(D|\theta)p(\theta)}{p(D)}$ (1)

where $p(D|\theta)$ and $p(\theta)$ represent the likelihood and the prior
distributions respectively. However, the denominator, $p(D)$, or the Evidence
is intractable in most cases. The community has thus far relied on sampling
techniques such as Markov Chain Monte Carlo (MCMC) or Nested Sampling to
compute the approximate posterior distribution (see Madhusudhan, 2018, for a
recent review).

Nevertheless, sampling-based techniques are prohibitively slow when the
dimensionality of the problem and quantity of the observed data are large
(Zhang et al., 2019). This issue will become increasingly pressing with the
increased data volume originating from the recently launched James Webb Space
Telescope (JWST, Greene et al., 2016) and other future missions such as, Ariel
(Tinetti et al., 2021), Twinkle (Edwards et al., 2019) and ELTs from the
ground (e.g. Udry et al., 2014). These telescopes are designed to provide
hundreds of higher resolution spectroscopic measurements with wider wavelength
coverage. On the other hand, recent investigations have highlighted potential
biases associated with some commonly used modelling assumptions, such as
isothermal atmospheres, constant with altitude chemistry, or 1D plane parallel
approximation (e.g. Rocchetto et al., 2016; Changeat et al., 2019, 2020;
MacDonald et al., 2020; Pluriel et al., 2020a; Ih & Kempton, 2021), prompting
the need for more realistic treatment of the atmosphere (e.g. Irwin et al.,
2020; Feng et al., 2020; Changeat & Al-Refaie, 2020). This will inevitably
increase both the computational cost and complexity (Changeat et al., 2021) of
the forward model. The increase in both quantity of data and model complexity
signals the need for alternative approaches to computing the posterior
distribution.

Recent years have seen a surge in Machine Learning-based techniques being
applied to many areas within exoplanetary science, from data detrending (e.g.
Morvan et al., 2019, 2021a; Krick et al., 2020; Gebhard et al., 2020; Nikolaou
et al., 2020; Gebhard et al., 2022), to planet detection (e.g. Shallue &
Vanderburg, 2018; Yip et al., 2020; Yu et al., 2019; Valizadegan et al., 2021)
and to planet characterisation (e.g. Márquez-Neila et al., 2018; Zingales &
Waldmann, 2018; Cobb et al., 2019; Waldmann & Griffith, 2019; Oreshenko et
al., 2020; Hayes et al., 2020; Yip et al., 2021b; Haldemann et al., 2022;
Ardévol Martínez et al., 2022; Himes et al., 2022). In 2022, the topic of
planet characterisation has also been featured as a competition at the Neural
Information Processing Systems (NeurIPS, Changeat & Yip, 2022; Yip et al.,
2022) conference.

Variational Inference (VI) is a widely studied approach in the field of
machine learning (ML) to provide approximate posterior distributions for a
large and high dimensional data set with reduced computational demand compared
to Markovian sampling approaches (e.g. Blei et al., 2016; Fellows et al.,
2018; Shu et al., 2018; Buchholz et al., 2018; Zhang et al., 2019; Argelaguet
et al., 2020; Friston et al., 2020; Fortuin et al., 2020; Zhang et al., 2021;
Karchev et al., 2022; Lopez-Alvis et al., 2021, 2022). However, variational
methods require models that can provide their gradient with respect to some
(input) parameters. The field has recently explored different applications of
differentiable physical models. Differentiable models open up the possibility
to construct “physics-aware” neural networks, a type of network that is
explicitly constrained by physical laws (e.g. Raissi et al., 2019; Chen et
al., 2020; Amini Niaki et al., 2021; Viana & Subramaniyan, 2021; Cai et al.,
2021; Haghighat et al., 2021; Morvan et al., 2021b; Cuomo et al., 2022). For
instance, Kawahara et al. (2022) used Hamiltonian Monte Carlo (HMC), a
gradient-informed Monte Carlo sampling algorithm (Duane et al., 1987; Hoffman
& Gelman, 2011), to perform atmospheric retrieval of exoplanets on high
resolution spectroscopic data. Others have also applied HMC to speed up
lightcurve-fitting (e.g. Foreman-Mackey et al., 2021; Agol et al., 2021).

Here we present the following contributions:

  1. 1.

We present Diff-$\tau$, a Tensorflow-based fully differentiable atmospheric
forward model, based on the implementation of TauREx3(Al-Refaie et al., 2021,
2022).

  2. 2.

For the first time, we introduced Variational Inference as a more efficient
alternative to perform atmospheric retrieval.

  3. 3.

Our framework only requires a single data instance during training time, there
is no need for a large library of spectra for pre-training.

  4. 4.

We show that our framework formally takes into account the uncertainty
associated with the observation and is able to reproduce physically motivated
correlations between atmospheric parameters.

  5. 5.

Our Bayesian Neural Network is capable of producing posterior distributions on
par with distributions produced from sampling based approaches.

Figure 1: Overview of our proposed method. Left panel shows the structure of
Diff-$\tau$. It is almost identical to a conventional forward model, in the
sense that it takes in a set of physical parameters and outputs the
corresponding theoretical spectrum, but with an additional ability to provide
gradients. The right panel shows the schematic of our proposed VI-retrieval.
We provided a simple (user defined) seed distribution $\bm{\xi}$, observed
data and our prior bounds as inputs to the optimisation loop (Train). The
Normalising Flow (NF)-based neural network, i.e. NF, is tasked to transform
$\bm{\xi}$ to a surrogate distribution $q(\mathbf{z})$. We employed
variational inference to approximate the typically intractable posterior
distributions, where the network (NF) must learn the best transformation to
optimise ELBO (controlled by log($\mathcal{L}$) and $D_{\text{KL}}$). To
calculate log$(\mathcal{L})$, we used Diff-$\tau$ to transform samples from
$q(\mathbf{z})$ into spectra and compared with the observed data. Once the
training is completed, the trained model is able to provide a
$q^{*}(\mathbf{z})$ that best approximate $p(\mathbf{z}|\mathbf{x})$

.

##  2 Overview

In this investigation our core aim is to explore an alternative approach to
conventional, sampling approach with the use of modern deep learning
techniques. For simplicity we will denote conventional, sampling based
approach simply as NS-retrieval and our proposed approach as VI-retrieval. Our
approach involves three core components - a differentiable physical model,
formulation of Variational Inference and a Normalising Flow-based neural
network. Here we will provide a top-level overview of how the three components
interact with each other, also see Figure 1 for a schematic overview of the
VI-retrieval.

Instead of relying on sampling to map the unknown posterior distribution (as
one normally does with e.g. MCMC based approaches), VI-retrieval relies on
finding a best-fit surrogate distribution to the actual posterior distribution
through optimisation. The use of optimisation-based techniques means that VI-
retrieval can be orders of magnitudes faster than sampling-based approaches,
especially on high-dimensional problems. However, there are two implementation
difficulties that prevented the widespread use of variational methods in the
field of exoplanetary atmospheres. These are: (1) Gradients: many optimisation
procedures demand complete knowledge of the gradient flow within a
computational graph (i.e. a neural network), contemporary atmospheric forward
models break the flow as they are undifferentiable (2) Surrogate
distributions: the chosen distribution (usually a multinomal Gaussian
distribution) is often too simplistic to represent the actual, underlying
posterior distribution.

To circumvent the above difficulties, we built Diff-$\tau$ with the Tensorflow
library, and utilise its automatic differentiation capabilities to compute
gradients of the forward model. At the same time, we implemented a Normalising
Flow-based neural network, a deep learning approach that can transform a
simple, “seed” distribution (such as a multinomial Gaussian) to arbitrarily
complex distributions.

In the following sections we will explain the theoretical background behind
each technique, and in the latter part of the paper we will demonstrate how
these concepts are linked to each other in practise by providing retrieval
examples for three different scenarios.

###  2.1 Variational Inference

The mathematical theory of variation inference (VI) and its application in the
field of machine learning have been extensively discussed in (Blei et al.,
2016). There are on-going efforts to investigate the statistical implication
of using VI for parameter estimation (e.g. Pati et al., 2018; Chérief-
Abdellatif & Alquier, 2018). Here we provide a brief overview of the
methodology.

Given an observed spectrum $\mathbf{x}$ defined by the transit depths
($x_{i}$) and associate uncertainty ($\sigma_{i}$) at each spectral bin $i$ ,
the goal of an atmospheric retrieval is to find the posterior distribution
$p(\mathbf{z}|\mathbf{x})$ of the set of latent variables ($\mathbf{z}$) that
can best describe the observation under a specific atmospheric model
assumption $\mathcal{M}$. Instead of approximating the unnormalised
$p(\mathbf{z}|\mathbf{x})$ via sampling, variational methods approximate the
distribution by finding a best-matching surrogate distribution via
optimisation.

Suppose we have a family of probability distributions $\mathbb{Q}$
parameterised by some latent variables $\mathbf{z}$. The optimal distribution
(best-matching surrogate distribution to $p(\mathbf{z}|\mathbf{x})$) is one
that minimises the statistical distance to $p(\mathbf{z}|\mathbf{x})$. A
common choice is to compute the Kullback–Leibler (KL) divergence between the
two probability density functions (PDF) i.e.

$\displaystyle q^{*}(\mathbf{z})$
$\displaystyle=\operatorname*{arg\,min}_{q(\mathbf{z})\in\mathbb{Q}}\>D_{\text{KL}}[p(\mathbf{z}|\mathbf{x})||q(\mathbf{z})]$
(2)
$\displaystyle=\mathbb{E}\left[\text{log}\left(\frac{p(\mathbf{z}|\mathbf{x})}{q(\mathbf{z})}\right)\right]$
(3)

The KL-divergence measures the relative entropy between two PDFs with range
[0,$\infty$]. A score of 0 means that the two distributions contain identical
information, and any (positive) deviation from 0 means the two distributions
become increasingly different from one another (Kullback & Leibler, 1951). The
KL-divergence can be computed analytically if one knows the functional form of
both PDFs. In cases when the functional forms of one or both PDFs are unknown,
as we will see below, numerical approaches must be sought to approximate the
divergence.

However, Equation 3 itself cannot be the objective function for our
optimisation task, as we do _not_ have any knowledge of
$p(\mathbf{z}|\mathbf{x})$ to begin with. To negate the dependence on the
unknown true posterior distribution, variational inference provides an
alternative formulation, the Evidence Lower BOund (ELBO),

$\text{ELBO}=\mathbb{E}[\text{log}\,p(\mathbf{x}|\mathbf{z})]-D_{\text{KL}}[q(\mathbf{z})||p(\mathbf{z})]$
(4)

where the first term is the expected value of the log-likelihood
log($\mathcal{L}$) and the second term is the KL divergence between
$q(\mathbf{z})$ and the prior distribution $p(\mathbf{z})$. This formulation
can be understood, in terms of Bayesian Statistics, as a ‘tug of war’ between
the likelihood function and our prior belief on the distribution. We included
a detailed discussion on ELBO and its link to Equation 3 in Appendix A.

Unlike Equation 3, the ELBO does not require any knowledge on the intractable
evidence $p(\mathbf{x})$ and therefore it can be computed analytically. We can
hence find the optimal distribution $q^{*}(\mathbf{z})$ by minimising the
ELBO. This optimisation technique is well-studied in the field of deep
learning (Kingma & Welling, 2013; Kingma et al., 2016). Most contemporary
neural networks are trained to minimise a given loss function (e.g. ELBO in
our case) by relying on a combination of gradient descent and back-propagation
algorithms. Modern deep learning libraries such as Tensorflow (Abadi et al.,
2015), PyTorch (Paszke et al., 2019) and JAX (Bradbury et al., 2018) provide
easy access to model training and evaluation. In this study, our
implementation will be solely based on the Tensorflow framework, but other
deep learning frameworks can similarly be used (Kawahara et al., 2022).

Figure 2: Top panel: comparing outputs from Diff-$\tau$ in native (light grey)
and binned (Ariel) resolution (orange) and TauREx3 binned to the same
resolution (black dots and dashed line). The lower panel shows the residual
between the two binned spectra. The grey lines represent the mean (solid line)
and 1-$\sigma$ standard deviation (dashed lines).

###  2.2 Differentiable forward model, Diff-$\tau$

Diff-$\tau$ is an atmospheric forward model built entirely within the
Tensorflow framework (tensorflow). We followed the forward model formulation
as specified in Al-Refaie et al. (2021) to construct Diff-$\tau$, with minor
modifications to comply with the Tensorflow framework. As the code is largely
based on TauREx3, it provides excellent agreement between the two forward
models, see Figure 2 for an empirical comparison between the two. We leverage
the built-in automatic differentiation functionality (Baydin et al., 2018),
which has the ability to automatically differentiate (almost) any functions
without the need to explicitly specify the corresponding derivative form. This
is immensely helpful as atmospheric models are a mixture of different physical
processes, and deriving the respective derivative forms analytically can be a
time consuming and non-trivial task.

###  2.3 Normalising Flows

The stringent requirement of a pre-defined family of distributions,
$\mathbb{Q}$, presents a major limitation in using VI. For many real-life
scenarios, the desired posterior distributions are rarely Gaussian nor well-
defined. We implemented a Normalising Flow-based neural network to break the
limitations of Gaussianity in $q^{*}(\mathbf{z})$ in Equation 4 by
transforming it into an arbitrarily complex probability distribution.

Normalising Flow (NF) describes a mechanism to “craft” a complex, multi-model
distribution from a simple, “seed” distribution (Rippel & Adams, 2013; Rezende
& Mohamed, 2015; Kobyzev et al., 2021). This can usually be a distribution in
the exponential family (e.g. a Gaussian) or an Uniform distribution.

Suppose we have an invertible function $g$, such that we can transform a
random variable $\bm{\xi}\sim p_{\bm{\xi}}$ into another random variable
$\mathbf{y}\sim p_{\mathbf{y}}$ using $\mathbf{y}=g(\mathbf{\bm{\xi}})$, the
probability density $p_{\mathbf{y}}$ of the random variable $\mathbf{y}$ can
be computed using the change of variable formula, i.e.

$p_{\mathbf{y}}(\mathbf{y})=p_{\bm{\xi}}(f(\mathbf{y}))|\text{det
D}(f(\mathbf{y})|$ (5)

where $f$ is the inverse of $g$, i.e. $f\equiv g^{-1}$,
$\text{D}f(\mathbf{y})$ is the Jacobian of $f$, i.e.
$\text{D}f(\mathbf{y})=\frac{\partial f}{\partial\mathbf{y}}$ and
$\text{D}g(\bm{\xi})$ is the Jacobian of $g$,
$\text{D}g(\bm{\xi})=\frac{\partial g}{\partial\bm{\xi}}$. In terms of
generative models, the invertible function $g(.)$ is a generator that “pushes”
forward the seed distribution to a more complex distribution function (the
generative direction). On the other hand, the function $f(.)$ moves in the
opposite direction, transforming it back to a simple, “normalised”
distribution (normalising direction). It has been shown that one can generate
or craft any form of distribution $p_{\mathbf{y}}$ from any base distribution
$p_{\bm{\xi}}$, given that the generator $g$ can be arbitrarily complex
(Bogachev et al., 2005; Medvedev, 2008).

Up until now we have only shifted the problem from crafting an arbitrarily
complex density function to an arbitrarily complex generator function.
Fortunately, invertible functions (or bijections) have a nice property - the
composition of invertible functions is itself invertible, meaning that one can
build a successively more complicated function by chaining non-linear
invertible functions together, i.e.

$g=g_{N}\circ g_{N-1}...\circ g_{1}$ (6)

where $g_{1}...g_{N}$ is a set of $N$ bijective function. Similarly, $g$ has
an inverse,

$f=f_{1}\circ f_{2}...\circ f_{N}$ (7)

and conveniently, the determinant of the Jacobian $\text{D}f(\mathbf{y})$ is
the product of individual determinant of the Jacobians
$\text{D}f_{i}(\mathbf{y})$, i.e.

$\text{D}f(\mathbf{y})=\prod_{i}^{N}\text{D}f_{i}(\mathbf{\bm{\varphi}_{i}})$
(8)

we denote $\mathbf{\bm{\varphi}_{i}}$ as the resultant vector of the $i$-th
intermediate flow, i.e. $\mathbf{\bm{\varphi}_{i}}=g_{i}\circ\cdots\circ
g_{1}(\bm{\xi})=f_{i}\circ\cdots\circ f_{N}(\mathbf{y})$, where
$\mathbf{\bm{\varphi}_{N}}=\mathbf{y}$.

In the context of our investigation, we will transform our seed distribution
$\bm{\xi}\sim p_{\bm{\xi}}$ into $\mathbf{y}\sim p_{\mathbf{y}}$ and treat
$\mathbf{y}\sim p_{\mathbf{y}}$ as our surrogate distribution $\mathbf{z}\sim
q(\mathbf{z})$, i.e. $\mathbf{y}\sim p_{\mathbf{y}}\equiv\mathbf{z}\sim
q(\mathbf{z})$

However, these intermediate functions must -by definition- be diffeomorphic,
meaning they must be bijective and differentiable (including their inverses).
In recent years the field has put significant efforts in constructing
bijectors that conform with these restrictions but remain sufficiently
expressive and computationally efficient even in high-dimensional problems
such as images (Louizos & Welling, 2017; Rothfuss et al., 2019; Wu et al.,
2020; Zhang & Chen, 2021; Nielsen et al., 2020), notable bijector
architectures include MADE (Germain et al., 2015), Masked Autoregressive Flow
(Papamakarios et al., 2017), NICE (Dinh et al., 2015), RealNVP (Dinh et al.,
2016), Sylvester Normalising Flow (Berg et al., 2018), FFJORD (Grathwohl et
al., 2018), Glow (Kingma & Dhariwal, 2018) and NSF (Durkan et al., 2019a).
Normalising Flow has proven to be a highly successful approach in a wide range
of applications, including audio synthesis (van den Oord et al., 2018; Prenger
et al., 2019; Aggarwal et al., 2020), text translation (Jin et al., 2019;
Izmailov et al., 2020), anomaly detection (Gudovskiy et al., 2022; Rudolph et
al., 2021), time series forcasting (Feng et al., 2022; Rasul et al., 2020;
Schmidt & Simic, 2019) and image generation (Kingma & Dhariwal, 2018;
Grathwohl et al., 2018; Lugmayr et al., 2020).

##  3 Implementation

###  3.1 Flow-based model setup

We implemented the Inverse Autoregressive Flow (IAF, Kingma et al., 2016) as a
default bijector unit in our Normalising Flow based neural network. To perform
the transformation, we chained $N=10$ bijector units together, each controlled
by a 2-layer densely connected neural network with 64 hidden units and ReLU
activation (He et al., 2015) in each layer. To stabilise the network, we
followed Kingma et al. (2016) and added a Batch Normalisation layer (Ioffe &
Szegedy, 2015) after each bijector unit. We used Adam (Kingma & Ba, 2014) as
our optimiser with a scheduled learning rate (see below), the rest of the
settings are kept as default from Tensorflow. We used a multi-dimensional
uniform distribution as our seed distribution, $\bm{\xi}$.

###  3.2 ELBO Formulation

We begin by describing our implementation to the general ELBO formulation in
Equation 4. The ELBO consists of two terms - the log-likelihood term and the
prior term.

We defined the log-likelihood as an additive log-Gaussian PDF, i.e.

$\displaystyle\mathbb{E}[\text{log}p(\mathbf{x}|\mathbf{z})]$
$\displaystyle=\mathbb{E}[\text{log}(\mathcal{N}(\mathbf{x}|\bm{\hat{\mathbf{\mu}}},\bm{\sigma})]$
(9)
$\displaystyle=\mathbb{E}\left[\text{log}\left(\frac{1}{\sqrt{2\pi\bm{\sigma}^{2}}}\text{exp}\left(-\frac{1}{2}\frac{(\mathbf{x}-\bm{\hat{\mu}})^{2}}{\bm{\sigma}^{2}}\right)\right)\right]$
(10)

where $\bm{\hat{\mu}}$ = $\texttt{Diff-$\tau$}{}(\mathbf{z})$ is a forward
model generated by Diff-$\tau$ binned to the spectral resolution of the data
and $\bm{\sigma}$ is the observed uncertainty. This formulation is the same as
the likelihood formulation currently employed by many retrieval frameworks.

The second term, $D_{\text{KL}}[q(\mathbf{z})||p(\mathbf{z})$, computes the
KL-divergence between the surrogate distribution and the prior distribution
(as defined by the user), it can be expanded in a similar fashion as Equation
3:

$D_{\text{KL}}[q(\mathbf{z})||p(\mathbf{z})]=\mathbb{E}\left[\text{log}\left(\frac{q(\mathbf{z})}{p(\mathbf{z})}\right)\right].$
(11)

The functional form of $p(\mathbf{z})$ can be easily determined (as it is the
user defined priori distribution). The surrogate distribtion, $q(\mathbf{z})$
has undergone multiple transformations by the Flow-based neural network.
Instead of computing the KL divergence analyatical which requires knowledge of
the full analytical form of the distribution, we can approximate the KL
divergence via Monte Carlo sampling, as it is not costly to sample repeatedly
from the transformed surrogate distribution (Kingma et al., 2016). For this
investigation, we sampled both, the prior and surrogate distributions, 10,000
times to approximate the value of the KL-divergence. Consistent with existing
literature in exoplanetary retrieval, we imposed an uniform prior on all
physical parameters, but we note that any proper prior probability
distribution can be used.

Instead of optimising the ELBO objective in its actual formulation, we adopted
the weight annealing approach from Sun et al. (2022) to optimise a modified
ELBO objective, i.e.

$\text{loss}=\frac{1}{\beta}\mathbb{E}[\text{log}\,p(\mathbf{x}|\mathbf{z})]-D_{\text{KL}}[q(\mathbf{z})||p(\mathbf{z})]$
(12)

where $\beta=\text{max}(1,\beta_{0}\frac{epoch}{\tau})$. $\beta_{0}$ and
$\tau$ represent the weight constant and decay constant respectively. This
formulation prevents the neural network from converging to bad local minima at
the start of the training, and encourages it to explore different solutions
before converging. As training progresses, the objective function will slowly
converge back to the original ELBO formulation.

###  3.3 The role of Diff-$\tau$

The formulation of the log-likelihood function (Equation 10) prompts the need
to compare our input observation with theoretical spectra. The role of
Diff-$\tau$ can be seen simply as a deterministic transformation, i.e. from
physical parameters to spectra. As part of the optimisation process, the
network must learn to adhere to physical laws imposed by the forward model. In
the end, the network will produce outputs and correlations that are physically
constrained in order to achieve better scores. In other words, the network’s
behaviour becomes physically motivated and explainable.

###  3.4 Training Procedure

At training time, we define the seed (user defined) distribution $\bm{\xi}$ =
$U[\text{min},\text{max}]$111all bounded between [-1,1], our Flow-based neural
network (the generator) will then transform it into a surrogate distribution
$q(\mathbf{z})$ through our chain of bijections. We then sample from
$q(\mathbf{z})$ and generate the atmospheric forward model using Diff-$\tau$.
At each iteration, we will sample 5 times from the surrogate distribution and
compute the average (modified) ELBO objective, Equation 12. The entire process
is repeated until the optimisation has converged or is terminated. We have
implemented an Early Stopping procedure to stop the optimisation process if
the loss value (ELBO) does not improve over 50 epochs. As we have adopted a
weight annealing objective, the stopping criteria will only become effective
in the later stages of the optimisation, when the modified objective converges
back to the original ELBO formulation (weights stay unity)222 the weights
annealing strategy reduces the contribution from the much larger likelihood
function initially and slowly return it to its original value as training
progresses. The loss value will almost always increase as the weights slowly
increases back to unity). The best model is used to produce the results in
Section 4. As for the learning rate, we have implemented a cyclic learning
rate as suggested by Smith (2015) and Himes et al. (2022). Our experiments is
consistent with their claims of improved training performance and we show that
it outperforms constant learning rate and a step-wise decaying learning
rates333The learning rate will reduce if the loss value did not decrease after
certain number of epochs. in terms of loss value as well as speed of
convergence.

Once trained, the generator network is decoupled from the framework. At
inference time, we will initiate a seed distribution and pass it to used to
transform the seed distribution into our best matching surrogate distribution.

##  4 Application

In this section we will perform atmospheric retrieval with NS-retrieval and
VI-retrieval. We will demonstrate the technique through three cases:

  1. 1.

Real Hubble/WFC3 (HST/WFC3) observation of the hot-Jupiter HD209458 b.

  2. 2.

Simulated Ariel Tier-2 observation of the hot-Jupiter HD209458 b.

  3. 3.

Simulated observation (0.5µm -15µm) with N-point temperature profile (Waldmann
et al., 2015; Changeat et al., 2021) of the hot-Jupiter WASP-43 b.

Parameters | Ground Truth* | Case I | Case II | Priors | Reference
---|---|---|---|---|---
Tp (K) | 1449 | ✓ | ✓ |  $\mathcal{U}$(100,2500) | N/A
Rp (RJ) | 1.359 | ✓ | ✓ |  $\mathcal{U}$(0.5,2.5) | N/A
log10(H2O) | -5 | ✓ | ✓ |  $\mathcal{U}$(-12,-2) | Polyansky et al. (2018)
log10(CH4) | -5 |  | ✓ |  $\mathcal{U}$(-12,-2) | Yurchenko et al. (2017)
log10(CO2) | -5 |  | ✓ |  $\mathcal{U}$(-12,-2) | Yurchenko et al. (2020)
log10(NH3) | -8 |  | ✓ |  $\mathcal{U}$(-12,-2) | Yurchenko et al. (2011)
log10(CO) | -8 |  | ✓ |  $\mathcal{U}$(-12,-2) | Li et al. (2015)
log${}_{10}(\chi^{lee}_{mie}$) | N/A | ✓ |  |  $\mathcal{U}$(-40,-4) | Lee et al. (2013)
$q^{lee}_{mie}$ | N/A | ✓ |  |  $\mathcal{U}$(1,99) | Lee et al. (2013)
log${}_{10}(a^{lee}_{mie}$) | N/A | ✓ |  |  $\mathcal{U}$(-3,1) | Lee et al. (2013)
Table 1: Fitted parameters for Case I and II, along with their corresponding
ground truth, prior bounds, and molecular line list references for each trace
gas species. The checkmark (✓) indicates whether the parameter is included in
each case. The ground truth is only available for Case II.

###  4.1 Atmospheric model setup for Case I and II

We have taken values from Tsiaras et al. (2016a) as our reference values on
the HD209458 system. In both cases, we assumed a blackbody emission for the
host-star at Teff = 6065 K and stellar radius R∗ = 1.155 R⊙. We also assumed a
primary atmosphere (dominated by H and He, ratio = 0.175) at solar abundance.
We divided the atmosphere into 70 atmospheric layers over a 10-5 Pa and 106 Pa
pressure range (evenly spaced in logarithmic scale) and furthermore assumed an
iso-thermal T-P profile and an iso-abundance chemical profile. Table 1 shows
the fitted input parameters for each case and their corresponding ground
truth, prior bounds and formulation/linelist references. The atmospheric setup
of Case III can be found in Section 4.4.

###  4.2 Case I: HST/WFC3 observation of HD209458 b

Case I aims to demonstrate our method’s applicability to actual data. The
transmission spectrum is observed with the HST/WFC3 G141 grism and processed
by Iraclis (Tsiaras et al., 2016a, 2019). The detrending process is described
in details in Tsiaras et al. (2016a). The basic atmospheric setup follows
Section 4.1. For this case we are assuming a hydrogen/helium dominated
atmosphere, with opacities from trace gases absorption and Mie scattering
clouds444We note that the alternative and commonly used flat cloud model (i.e.
a constant in pressure opacity cut-off) is inherently not differentiable and
not physically viable, we therefore chose to include the differentiable Mie
scattering formulation in Diff-$\tau$. (Lee et al., 2013). We ran the
optimisation procedure for 2000 epochs, with the convergence parameters,
$\beta_{0}$ and $\tau$, set to 100 and 200 respectively. These values are
determined with a coarse hyperparameters search between $\beta_{0}$ = [100,
1000] and $\tau$ = [100, 1000].

###  4.3 Case II: Ariel Tier-2 observation on HD209458 b

In the second case we would like to understand the ability of our proposed
framework to retrieve the ground truth value and showcase the flexibility of
our framework to switch to different atmospheric assumptions and spectral
resolutions. We used ArielRad (Mugnai et al., 2020) to simulate the expected
noise level for each wavelength channel for HD209458 b at Ariel Tier 2
resolution. In terms of atmospheric chemistry, we adhered to the same setup as
described in Section 4.1 and include five trace gases H2O, CH4, CO and CO2 and
NH3. We chose this set of molecules due to their expected contribution in the
wavelength range considered, and because they have been successfully detected
in hot Jupiter’s atmospheres. We ran the optimisation procedure for 2000
epochs, with $\beta_{0}$ and $\tau$ set to 100 and 300 respectively.

###  4.4 Case III: N-Points TP profile retrieval with WASP-43 b

We simulated the atmosphere of a WASP43 b like planet observed at a customised
wavelength range, i.e. from 0.5 - 15 µm. The atmospheric model setup largely
followed Section 4.1, but with 50 atmospheric layers instead of 70 layers. As
for stellar and planetary parameters, we followed Hellier et al. (2011) and
set Teff = 4400 K, R∗= 0.6 R⊕, Mp = 1.78 MJ and Rp = 0.93 RJ. We included
log10(H2O) and log10(CH4) as trace gases in the atmosphere, with their
abundances set at -4 and -3 respectively. As for the temperature-pressure
profile, we adopted the dayside temperature profile of WASP-43 b retrieved by
Changeat et al. (2021). We then generated a (binned) atmospheric model with
the above settings, and added a 100 ppm Gaussian white noise to produce an
observation. The prior bounds of the parameters follows Table 1 apart from
planet temperature, where we expanded the bounds to [100, 3000]. The
observation remains the same for all retrievals performed in Case III.

As for the retrieval settings, we performed a N-Point retrieval with each
temperature point located at a fixed, pre-defined pressure point. The top and
bottom pressure points are fixed at $10^{-5}$ and $10^{6}$ respectively
throughout the investigation, while the (pressure) points in between them are
varying in accordance to the number of retrieved temperature (-pressure)
points. In all cases the pressure points are separated equally (in log-
pressure space). Other retrieved parameters included Rp, log10(H2O) and
log10(CH4). In summary, the total number of free parameters is N (temperature
points) + 3. The network setup mostly follows Section 3.4. We ran the
optimisation procedure for 5000 epochs, with $\beta_{0}$ and $\tau$ as 100 and
300 respectively.

##  5 Results

###  5.1 Case I and II

Figure 3: Posterior distribution of HD209458 b WFC3/G141 observation obtained
from NS-retrieval (blue) and VI-retrieval (yellow). The top right corner is an
empirical comparison between the two best fitted spectra as obtained by the
two methods. The shaded area shows the 1-$\sigma$ spread of the respective
retrieval approach.

In both cases (Figures 3 & 4), VI-retrieval (in yellow) is able to converge to
very similar results to NS-retrieval (in red). We observed an excellent
agreement between both methods and the percentile values555defined as the 16th
and 84th percentile, indicated by the dashed line in figures 3 & 4 are within
1-$\sigma$ of each other.

As for the shape of the posterior distribution, the surrogate distribution is
able to reproduce peaked, Gaussian-like distributions whenever there is
sufficient constrains from the observation (i.e. radius of the planet at 1 bar
pressure), and ,at the same time, produce a uniform-like distribution (with
upper limit) when the free parameters are unconstrained by the observation
(e.g. CO in Case II), conforming with our uniform prior. In addition to this,
the surrogate distribution provides a faithful reproduction of the covariances
between the free parameters. These correlations are resulted from the
interactions between the free parameters through the formulation of the
radiative transfer equations. They are most notable in Case II, when the high
S/N observation of Ariel allows better constraints on the free model
parameters.

Figure 4: Posterior distribution of HD209458 b in Ariel Tier2 resolution
obtained via NS-retrieval (blue) and VI-retrieval (yellow). The top right
corner is an empirical comparison between the two best fitted spectra as
obtained by the two methods. The shaded area shows the 1-$\sigma$ uncertainty
of the respective retrieved methods

Apart from the similarities, there are also noticeable differences between the
two retrievals, for instance VI-retrieval is less adept at capturing “cliff”
like distributions (e.g. log${}_{10}(\chi^{lee}_{mie}$)), when there is a
sharp upper or lower bound on a parameter. There are also instances (such as
Tp in Case I) when the surrogate distribution is only able to capture “part”
of the conditional distribution obtained via NS-retrieval. The difference
could be due to the imperfect optimisation process. The loss function (ELBO)
is a balance between the log-liklihood term and the prior term. Minimising one
term will always come at the expense of maximising the other. This situation
is aggravated by the fact that the variability of the log-likelihood term is
orders of magnitude higher than the prior term, which means that the training
will always be driven by the former term and produce sharply peaked
distributions. Our modified objective function (Equation 12) aims to alleviate
the imbalance between the two terms by explicitly lowering the contribution
from the log-likelihood term during the start of the training. However, the
underdispersed situation will likely return as the formulation slowly converge
back to the original ELBO form, as seen in both cases. An alternative remedy
is to increase permanently the contribution from the prior term, which will
inevitably increase the estimated uncertainty and makes it harder to compare
to the conventional approaches (i.e. NS-retrieval).

###  5.2 Case III

Figure 5 shows the outcome of the simulated WASP-43 b observation. Figure 6
shows the results of performing N-points retrievals from 6 free parameters to
12 free parameters. The first row compares the retrieved TP profiles from both
approaches (blue: NS-retrieval, yellow: VI-retrieval). The second row compares
the number of forward models calls (Ncalls) required by each approach. In all
cases VI-retrieval (in yellow) requires 6 times fewer Ncalls compared to NS-
retrieval (blue). The third row compares the log-evidence obtained by each
approaches. In all cases, the log-evidence obtained by VI-retrieval never goes
higher than the ones obtained via NS-retrieval, which is consistent with the
formulation of ELBO. Both approaches show a declining trend in log-evidence as
the number of retrieved points increases, with VI-retrieval declining more
rapidly than NS-retrieval. The faster decline in log-evidence may be affected
by the increase in dimensionality (free parameters), which makes the accurate
estimation of the log-evidence harder.

Figure 5: Left: Simulated observation of WASP-43 b from 0.5 µm to 15 µm;
Right: Dayside temperature profile used to generated this observation (adopted
from Changeat et al. (2021)), there are 50 pressure points separated evenly
(in logarithmic scale) between $10^{-5}$ to $10^{6}$. Figure 6: Comparing the
performance of VI-retrieval (yellow) and NS-retrieval(blue). First row: T-P
profiles retrieved by each approach, each subplot shows a different fixed-
point TP-profile, going from 3 points to 9 points at a step of 2. Second row:
The number of forward model calls requires for each N-points retrieval. Third
row: The log-evidence values retrieved by each approach.

##  6 Discussion

###  6.1 Limitations of Grid-based learning

Most contemporary machine learning based atmospheric retrievals are trained in
a supervised fashion with a large grid of simulated spectra produced by a
forward model. We broadly refer these models as grid-based models here. This
kind of training procedure takes away the computational burden of having to
generate thousands to millions of forward model on the fly during model
deployment and makes these models computationally fast for applications within
their trained domain. In other words, these models offload the computational
burden of atmospheric retrievals to the model training stage, with the fully
trained model being quick to run at inference. However, this approach comes
with three limitations: 1) model generalisability, 2.) lack of Bayesian
framework and 3.) a lack of model interpretability.

Ncalls | NS-retrieval | VI-retrieval
---|---|---
WFC3/G141 | 105622 | 8781
Ariel | 65414 | 9570
Table 2: Comparing the number of forward model calls by each method for the
same atmospheric models (as in Case I) in WFC3/G141 and Ariel spectral
resolution. In both cases, the number of forward model calls by VI-retrieval
is significantly lower than NS-retrieval.

##### Generalisability

Any changes to the underlying model, spectral range or resolution, will hinder
the model’s performance, and in some cases, will require a full re-computation
of the training data from scratch and re-training the model (Márquez-Neila et
al., 2018; Zingales & Waldmann, 2018; Cobb et al., 2019; Yip et al., 2021b;
Haldemann et al., 2022; Ardévol Martínez et al., 2022). Such a scenario can be
triggered by anything as simple as adding an extra molecule that is previously
not present in the training data. These limitations can be alleviated to a
certain extent by training a surrogate forward model as done in Himes et al.
(2022). Their model is tasked to produce synthetic spectra at very high
resolution, and the output can subsequently be down-sampled to any appropriate
spectral range and resolution when required. Nevertheless, the model is not
immune to changes to the underlying atmospheric assumptions and will likely
need to be re-trained in those cases.

##### Lack of Bayesian framework

Most contemporary retrievals aim to map the Bayesian posterior distribution.
In contrast, most ML models applied in the field of atmospheric
characterisation are formulated to perform maximum likelihood estimation (e.g.
Márquez-Neila et al., 2018; Yip et al., 2021b; Haldemann et al., 2022; Ardévol
Martínez et al., 2022). The difference between these two objectives presents
an obstacle when trying to compare outputs from the two methodologies. Other
non-Bayesian approaches includes ExoGAN, in which case the generator is
trained to optimise the adversarial loss (Zingales & Waldmann, 2018), but this
loss function does not guarantee a robust computation of $P(\theta|D)$. There
are other ways to bypass this limitation. For instance, Cobb et al. (2019) and
Ardévol Martínez et al. (2022) use an ensemble of neural network to
approximate the conditional distribution. Himes et al. (2022) bypassed this
constraint by running a traditional retrieval with a ML-based surrogate
forward model.

##### Interpretability

Interpretability varies from one model to another. Learning algorithms such as
linear regression and Decision trees are some of the most transparent
algorithms, but with limited modelling capability. Deep learning algorithms
sits at the opposite end of the spectrum, with extensive learning
capabilities, but at the cost of very limited interpretability. Most grid-
based learning models (Márquez-Neila et al., 2018; Zingales & Waldmann, 2018;
Cobb et al., 2019; Yip et al., 2021b; Haldemann et al., 2022; Ardévol Martínez
et al., 2022) used more complex models to learn the covariance between
parameters from a large grid of spectral examples. However, the loss of
interpretability means that it becomes hard for users to understand when,
where and how the algorithm may break. With no explicit control on the
learning process (remember they are only asked to optimise the learning
objective), it is uncertain how the model may interpret the training data. By
relying only on the training data, it becomes important to make sure the
training data can adequately present the underlying forward model, which may
be computational expensive or in some cases difficult to ascertain (Fisher &
Heng, 2022).

###  6.2 A Flexible and Interpretable Bayesian Framework

Motivated by the above limitations. We presented an alternative approach to
train an end-to-end deep learning model for atmospheric retrievals. Instead of
focusing on a generalisable model (one that works on a wide range of spectra),
our framework is specific (or in other words, overfitted) to the observed
data, and similar to conventional Markovian sampling algorithms, it is not
generalisable and must be re-run for any changes in observed data and/or model
assumptions. By dropping the goal of training a generalisable model, we forego
the need to create a vast training sets as well as the need to pre-train our
model before any retrieval could take place. It furthermore affords us with
the flexibility of easily changing our input data and forward models. This
flexibility has allowed us to explore the performance of our model at
different spectral resolutions, wavelength ranges, observational uncertainties
and model assumptions with relative ease (as demonstrated above). However,
this alternative approach also comes with disadvantages, see Section 6.4.1.

To allow direct comparison with conventional retrievals, we formulated the
objective function to optimise the ELBO function, an alternative formulation
of the Bayes’ Theorem. This modification constrains the behaviour of our
trained model, and allows us to obtain results comparable to conventional
retrievals (as demonstrated above). Another advantage of this approach is the
drastic reduction of the number of forward model computation to a fraction of
its retrieval counterpart. Table 2 shows a comparison of the number of forward
model calls between NS- and VI-retrieval in Case I and II.

In the grid-based approaches, models have to learn implicitly the physical
relationships between atmospheric parameters from the training data. There is
no guarantee that physical laws are preserved or correctly represented in the
trained model. Our framework is purposely built to explicitly impose physical
laws. One can view our NF-based network as a generator and Diff-$\tau$ as a
corresponding decoder that transforms parameters into spectra. The
static666static in the sense that it remains invariant throughout the training
decoder acts as a physical regularizer to the generator, which explicitly
constrains the output from the generator to align with physical laws (any
misalignment will be reflected in the transformed spectra). This is
demonstrated through our examples, the surrogate distributions are able to
provide physically plausible correlations and are aligned with correlations
produced from standard Bayesian Sampling retrievals (NS-retrieval, blue
contour). Having a fully analytic model allows us to impose physical laws
directly, without going through a training data proxy.

###  6.3 Objective Function

The objective (loss) function is a crucial factor that governs the learning
behaviour of a deep learning model. Deep learning models in the literature are
usually trained to best-match the respective ground truth values of the
physical parameters777parameter values used to generate the forward
model(fitting-for-parameters). Here we opt to align our objective to that of
our retrieval counterpart, in other words, we are explicitly asking our model
to look for solutions that can best explain our observation (fitting-for-
observations).

Adopting an observation based likelihood not only allows us to align with the
objective function of conventional retrievals, it also has added advantage of
properly accounting for observational uncertainty. ML-based retrieval methods
often incorporate observational uncertainties through noise augmentation (Yip
et al., 2021b; Ardévol Martínez et al., 2022), this has resulted in over-
estimation of the error bound as compared to Nested Sampling based retrievals
(Ardévol Martínez et al., 2022).

In an ideal world, both approaches (fitting-for-parameters and fitting-for-
observations) will agree with each other. However, it is not the case for
inverse problems, where our observation is inherently corrupted888possible
sources include instrument and astrophysical noise sources as well as the
information loss from binning data and/or forward model and we may never be
able to recover the ground truth in some cases due to the loss of information
and the inverse processing being ill-defined. In such cases, differences in
the objective function may lead to different results. On one hand, the
fitting-for-parameter approach is asking the neural network to pursue
parameter values that may no longer be possible to retrieve (due to corruption
of the observed data), which will cause the neural network to exhibit
fictitious behaviour if it results in the lowering of the loss function values
(Yip et al., 2021b). On the other hand, the fitting-for-observation approach
explicitly ask for spectra that can explain the observation, not the
underlying ground truth. Of course, that will also mean that our approach is
not immune from the intrinsic retrieval biases induced by the atmospheric
forward itself (Rocchetto et al., 2016; Feng et al., 2020; MacDonald et al.,
2020). These biases can only be alleviated through increasing the complexity
of the atmospheric forward model to better represent the physical/chemical
processes leading to the observed spectra.

In terms of model development, our framework bypasses the need to train the
network with a large library of synthetic spectra before applying to actual
data, as we are training the network directly on actual observations. This
move avoids the problem of data shift (Quionero-Candela et al., 2009) \- where
our training distribution is different from our test distribution.

####  6.3.1 Computational cost of VI-retrieval

In this section we will discuss the computational cost associated with running
a retrieval using VI-retrieval compared to NS-retrieval. Under the same
forward model, the VI-retrieval will always beat the NS-retrieval, but the
matter becomes slightly more complicated as the two methods do not always
reuse the same forward model. We provide an initial assessment of the two
methodologies with two different forward models. Note that the following
figures are based on running both methods using a Macbook Pro with Apple M1
chip (8 CPU cores in total). These figures should be taken with caution as the
two forward models are at different developmental stages, TauREx3 is a highly
optimised code while Diff-$\tau$ is proof-of-concept forward model written
from scratch entirely in Tensorflow.

A single forward model call (FMcall) for Case I scenario takes about $0.89\pm
0.02$ seconds on Diff-$\tau$ and $0.26\pm 0.10$ seconds with TauREx3’s forward
model. For a single retrieval, VI-retrieval takes 8781 FMcall and NS-retrieval
takes 105,622 FMcall. Based on these numbers, VI-retrieval takes 7815s and NS-
retrieval takes 94,003s respectively. As for Case II, a single forward model
call takes about $5.08\pm 0.38$ seconds on Diff-$\tau$ and $0.51\pm 0.34$
seconds with TauREx3’s highly optimised forward model. For a single retrieval,
VI-retrieval takes 9,570 FMcall and NS-retrieval takes 65,414 FMcall. Based on
these numbers, VI-retrieval takes 48,615.2s and NS-retrieval takes 33,361s
respectively.

From our simple analysis above we can see that VI-retrieval’s computation time
is dependent on the speed of the forward model. It is empirically faster than
NS-retrieval if the computational time of the forward models are similar to
each other (Case I), but the advantage will elapse as the forward model takes
longer to compute . This difference in performance can be easily narrowed down
with model optimisation and maturity of the framework itself. In this vein,
VI-retrieval’s will play an important role in retrieval scenarios with
complex, and subsequently slow to run, forward models. In these situations the
optimisation objective becomes the number of forward model calls with VI-
retrieval’s being an order of magnitude more efficient.

####  6.3.2 Model Selection

Model | ELBO | Ref | log${}_{10}(\mathcal{B})$
---|---|---|---
Flat line | 62.74 | 62.83 | 315.66
No Methane | 345.37 | 347.18 | 33.03
Complete | 378.40 | 380.20 | N/A
Overspecified Model | 374.00 | 377.74 | 4.4
Table 3: Retrieved ELBO from VI-retrieval and Bayesian (log-) Evidence from
NS-retrieval (Ref) for each scenario. log${}_{10}(\mathcal{B})$ shows the log-
difference between the ELBO of our complete model and other competing models.

Model selection is a key part of model evaluation cycle. So far, the machine
learning retrieval literature has largely omitted the issue by training
networks under one or several fixed atmospheric assumptions. Our flexible
framework and similar objective function to NS-retrieval means that we can,
for the first time, utilise some of the tools frequently used by sampling-
based retrievals to compare the retrieval results from different models.

Given that our surrogate distribution is a good approximation of the
underlying posterior distribution, our ELBO, despite being a lower bound,
should closely approximate the Bayesian evidence. We can therefore use the
ELBO as a proxy of the evidence, and compare our models by estimating the
Bayes’ factor.

$\mathcal{B}_{10}=\frac{P(\textbf{x}|\mathcal{M}_{1})P(\mathcal{M}_{1})}{P(\textbf{x}|\mathcal{M}_{0})P(\mathcal{M}_{0})}$
(13)

where P$(\textbf{x}|\mathcal{M}_{k})$ represents the Bayesian Evidence
attained from model $\mathcal{M}_{k}$, $P(\mathcal{M}_{k})$ represents our
prior belief on a particular model.

As an empirical example, we took our example from Case II and performed VI-
retrieval and NS-retrieval with different atmospheric assumptions, including a
flat line model, an incomplete model (without methane), a complete model (as
specified in Case II) and an overspecified model (Case II + TiO and VO). Table
3 compares the corresponding ELBO from VI-retrieval and Bayesian Evidence from
NS-retrieval. The ELBO retrieved from each model closely follows, but always
smaller than, the corresponding Bayesian Evidence, as set out from the
definition of ELBO. The Bayes’ factor displayed in Table 3 allows us to
differentiate different models. Following Jeffrey’s guideline scale (Jeffreys,
1998; Hobson et al., 2002; Padilla et al., 2019), the Bayes’ factor strongly
favours our complete model over other competing models. For more discussion on
using Bayesian evidence as a model selection tool, please refer to the
Appendix in Changeat et al. (2021).

###  6.4 Limitations

####  6.4.1 Limitations on VI-retrieval

In this section we will focus on the framework of VI-retrieval.

  1. 1.

Generalisbility: As set out from the design goal, VI-retrieval is targeted to
work on a single observation. This approach has earned us flexibility, where
we can freely change our model assumptions, spectral range and resolution with
relative ease. However, the same advantage has also limited the
generalisbility of our framework. As opposed to other deep learning
approaches, which can be rapidly deployed to a data set within its training
set range, it will have to be re-trained for each observation.

  2. 2.

Convergence: Gradient descent helps to converge to an global optimum if
function is convex. However in the presence of model degeneracy, the function
becomes non-convex and convergence to global minima is not guaranteed.

  3. 3.

Flow Bijections: The expressiveness of the surrogate distribution depends
highly on the architecture of neural network and the type of Flow bijectors
(IAF in our case). While our investigation has demonstrated the flexibility if
these bijectors, it does not mean that any distribution (Durkan et al., 2019b)
can be mimicked. It remains an ongoing effort to design bijectors that are
both flexible and computational efficient.

  4. 4.

Approximation: The formulation of variational inference, especially the use of
ELBO in the objective function, means that the retrieved surrogate
distribution will always be an approximation to the ground truth. The use of
Flow-based neural networks has of course improved the fidelity of the
approximation.

  5. 5.

Tendency to produce underdispersed solution: The implementation of ELBO and
its modified form in Section 3.2 hints the tendency for the network to produce
underdispersed solution. The negative log-liklihood term (1st term) is almost
always going to be larger than the prior term (2nd term), meaning that the
network will tend to produce a sharp distribution first before widening itself
to comply with the prior. This effect is also related to the expressiveness of
the bijections, their flexibility is ultimately bounded by the range of
‘action’ that they can perform to transform the distribution. We note here
that this effect is minor in our case as the shape of the target distributions
are relatively simple.

  6. 6.

Hyperparameter tuning: Similar to conventional neural networks,
hyperparameters do not always stay optimal from one setup to another, which
means some degree of hyperparameter search should be performed to identify a
good setup However, this is alleviated by identifying the key hyperparameters
that may influence the optimization procedure, in our case we realised the
annealing weights plays a role in the optimisation procedure, owing to its
close relationship with the objective function. Other hyperparameters are
generally fixed and do not differ much.

####  6.4.2 Limitations on differentiable frameworks

The ability to differentiate and provide gradients with respect to some
quantities is central to modern deep learning algorithms. The growing
popularity of Machine Learning in recent years has accelerated the development
of differentiable frameworks, among them Tensorflow, PyTorch and JAX have
amassed a substantial user bases.

There have been several recent attempts within the field of exoplanets to
develop differentiable physical models within these frameworks. Initial
results show that these differentiable models may hold the keys to overcome
the curse of dimensionality brought by our increasingly complex models (Morvan
et al., 2021b; Kawahara et al., 2022) and may one day enable us to perform
population study without placing significant demand on computational
resources.

However, these frameworks are not without issues. Here we would like to
provide our “user experience” for readers who are interested and/or would like
to implement their own models.

  1. 1.

Significant overhead: Current development in differentiable programming
mandates that any implementation must be written entirely in terms of a chosen
framework. While it has become relatively straightforward to translate _most_
operations from one framework to another, it is not a trivial and error-free
process. These frameworks come with their own programming restrictions and
conventions. Users are expected to adhere to these conventions or otherwise
they might risk losing the ability to differentiate.

  2. 2.

Differentiability is not guaranteed: Not all operations are differentiable.
Some operations may be mathematically non-differentiable999such as a
discontinuity in the function, e.g. logarithmic function or they are not
designed to have a gradient, such as a LookUpTable or interpolation. Depending
on the algorithmic structure of the forward (atmospheric) model, this obstacle
may present difficulties in obtaining a valid gradient. One good example is
the grey cloud model. This model assumes the planet becomes completely opaque
below a certain pressure level/altitude. This model, while easy to implement
in any computational language, is not differentiable. A possible way around is
to implement their own gradient like in Kawahara et al. (2022), but this
approach is only possible if one knows the differentiable form.

  3. 3.

Sub-optimal Performance: Many modern differentiable languages are built and
designed around deep learning based applications. They are not designed to
handle complex computational models. In other words, a differentiable model
may suffer from reduced performance compared to its un-differentiable
counterpart (Hu et al., 2020, 2019).

  4. 4.

Rapidly evolving language: Differentiable languages are under constant
development and rapid release cycles in response to latest researches. Some of
these developments may not be backward-compatible101010one good example is the
switch from Tensorflow 1 to Tensorflow 2. and may impact the long term
sustainability of the developed model.

##  7 Conclusion

In this paper we introduced the differentiable forward model (Diff-$\tau$)
based on TauREx3 and implemented in Tensorflow. We combined our newly
developed model with our density-alternating neural network and showed that it
is possible to compute the approximate Bayesian posterior distribution that
are in excellent agreement with the ones produces from computationally more
expensive sampling based techniques.  Through our examples we have
demonstrated three advantages of our framework:

  1. 1.

Fewer forward model calls: Our VI-retrieval requires 25% or fewer forward
model calls compared to NS-retrieval to converge, which opens up opportunities
for more rapid retrieval with more complex (i.e. slower) forward models.

  2. 2.

Flexibility: Unlike many deep learning based frameworks, which rely on a large
pre-computed libraries of spectra, our proposed framework resembles more
closely a traditional retrieval set up by explicitly including the physical
forward model in the deep learning architecture. Consequently, it retains many
of the advantages of a traditional retrieval codes, such as the freedom to
choose the model assumptions, number of free parameters, spectral wavelength
range and observed uncertainty, without having to produce a separate training
dataset each time.

  3. 3.

Error propagation: By incorporating observational uncertainty directly into
the likelihood function, we demonstrated the capability of our network to
produce uncertainty bounds that are on par with conventional atmospheric
retrievals.

  4. 4.

Model Selection: We demonstrated, for the first time, that we can compare the
adequacy of our neural network model to a given observation by computing the
Bayesian Evidence and Bayes Factor.

Our proposed framework presents a major step towards the wider adoption of
neural network powered atmospheric retrievals. With significantly higher
spectral resolutions and signal-to-noise ratios afforded by JWST and Ariel
data, atmospheric forward models must consequently increase in complexity to
model these data accurately. Such an increase in the dimensionality of the
problem will result in significant strain on traditional sample-based
retrieval approaches. While the results shines light to rapid AI-assisted
Bayesian inference, we must stress that our framework is _not_ meant to
replace sampling based framework, but to _complement_ existing frameworks. VI-
retrieval, similar to most MCMC based approach, are susceptible to non-global
minima, NS-retrieval on the other hand, tends to allow more thorough
exploration of the parameter space. This difference becomes important when
multiple solutions are equally plausible for a given observation. The
approximate retrieval performed by VI-retrieval can act as a precursor before
the more computationally heavy Nested Sampling retrievals are required. Its
rapid convergence means that one can scan through multiple candidate
atmospheric models before converging into a few promising models for further,
more detailed evaluations. The parameter bounds extracted by our framework can
also act as an informative prior to speed up the sampling process of
conventional retrievals.

This project has received funding from the European Research Council (ERC)
under the European Union’s Horizon 2020 research and innovation programme
(grant agreement No 758892, ExoAI) and the European Union’s Horizon 2020
COMPET programme (grant agreement No 776403, ExoplANETS A). Furthermore, we
acknowledge funding by the Science and Technology Funding Council (STFC)
grants: ST/K502406/1, ST/P000282/1, ST/P002153/1 and ST/S002634/1.

##  Appendix A ELBO derivation

In this section we will describe the mathematical derivation that led to the
formulation of the Evidence Lower BOund, ELBO. We will start from Equation 3,
i.e.

$q^{*}(\mathbf{z})=\operatorname*{arg\,min}_{q(\mathbf{z})\in\mathbb{Q}}\>D_{\text{KL}}[p(\mathbf{z}|\mathbf{x})||q(\mathbf{z})]$
(A1)

We start by expanding the left hand side of the above expression:

$D_{\text{KL}}[p(\mathbf{z}|\mathbf{x})||q(\mathbf{z})]=\mathbb{E}[\text{log}q(\mathbf{z})]-\mathbb{E}[\text{log}p(\mathbf{z}|\mathbf{x})]\\\
$ (A2)

and expand the latter term with the multiplication rule and chain rule of
probability,

$\displaystyle D_{\text{KL}}[p(\mathbf{z}|\mathbf{x})||q(\mathbf{z})]$
$\displaystyle=\mathbb{E}[\text{log}(q(\mathbf{z}))]-\mathbb{E}[\text{log}p(\mathbf{z},\mathbf{x})]+\text{log}p(\mathbf{x})$
(A3)
$\displaystyle=\mathbb{E}[\text{log}(q(\mathbf{z}))]-\mathbb{E}[\text{log}p(\mathbf{x}|\mathbf{z})]+\mathbb{E}[\text{log}p(\mathbf{z})]+\text{log}p(\mathbf{x})$
(A4)
$\displaystyle=-(\mathbb{E}[\text{log}p(\mathbf{x}|\mathbf{z})]-(\mathbb{E}[\text{log}(q(\mathbf{z}))]-\mathbb{E}[\text{log}p(\mathbf{z})])+\text{log}p(\mathbf{x})$
(A5)
$\displaystyle=-(\mathbb{E}[\text{log}p(\mathbf{x}|\mathbf{z})]-D_{\text{KL}}[q(\mathbf{z})||p(\mathbf{z})])+\text{log}p(\mathbf{x}).$
(A6)

Here we will substitute the expression for ELBO as defined in Equation 4,

$D_{\text{KL}}[p(\mathbf{z}|\mathbf{x})||q(\mathbf{z})]=-\text{ELBO}+\text{log}p(\mathbf{x})$
(A8)

From the above expression we can see that the KL divergence is determined by
the interaction between the ELBO and the Bayesian Evidence term. While the
former term depends on the quality of the surrogate distribution
$q(\mathbf{z})$, the latter term is a constant as it depends only on the data
$\mathbf{x}$. This, combined with the fact that $D_{\text{KL}}\geq 0$ , the
ELBO term is inherently constrained by $\text{log}p(\mathbf{x})$, i.e.

$\text{ELBO}\leq\text{log}p(\mathbf{x})$ (A9)

Hence the term is known as the Evidence Lower BOund (ELBO). We can therefore
use ELBO as our objective function, since maximising the ELBO is equivalent to
minimising the KL divergence. In our implementation, the objective function is
set to minimise the negative ELBO.

## References

  * Abadi et al. (2015) Abadi, M., Agarwal, A., Barham, P., et al. 2015, TensorFlow: Large-Scale Machine Learning on Heterogeneous Systems.  https://www.tensorflow.org/
  * Aggarwal et al. (2020) Aggarwal, V., Cotescu, M., Prateek, N., Lorenzo-Trueba, J., & Barra-Chicote, R. 2020, in ICASSP 2020 - 2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), 6179–6183, doi: 10.1109/ICASSP40776.2020.9053678
  * Agol et al. (2021) Agol, E., Dorn, C., Grimm, S. L., et al. 2021, The Planetary Science Journal, 2, 1, doi: 10.3847/PSJ/abd022
  * Al-Refaie et al. (2022) Al-Refaie, A. F., Changeat, Q., Venot, O., Waldmann, I. P., & Tinetti, G. 2022, ApJ, 932, 123, doi: 10.3847/1538-4357/ac6dcd
  * Al-Refaie et al. (2021) Al-Refaie, A. F., Changeat, Q., Waldmann, I. P., & Tinetti, G. 2021, ApJ, 917, 37, doi: 10.3847/1538-4357/ac0252
  * Alam et al. (2020) Alam, M. K., López-Morales, M., Nikolov, N., et al. 2020, AJ, 160, 51, doi: 10.3847/1538-3881/ab96cb
  * Alam et al. (2021) Alam, M. K., López-Morales, M., MacDonald, R. J., et al. 2021, ApJ, 906, L10, doi: 10.3847/2041-8213/abd18e
  * Amini Niaki et al. (2021) Amini Niaki, S., Haghighat, E., Campbell, T., Poursartip, A., & Vaziri, R. 2021, Computer Methods in Applied Mechanics and Engineering, 384, 113959, doi: 10.1016/j.cma.2021.113959
  * Anisman et al. (2020) Anisman, L. O., Edwards, B., Changeat, Q., et al. 2020, AJ, 160, 233, doi: 10.3847/1538-3881/abb9b0
  * Ardévol Martínez et al. (2022) Ardévol Martínez, F., Min, M., Kamp, I., & Palmer, P. I. 2022, A&A, 662, A108, doi: 10.1051/0004-6361/202142976
  * Argelaguet et al. (2020) Argelaguet, R., Arnol, D., Bredikhin, D., et al. 2020, Genome biology, 21, 1
  * Barbary (2021) Barbary, K. 2021, nestle: Nested sampling algorithms for evaluating Bayesian evidence, Astrophysics Source Code Library, record ascl:2103.022.  http://ascl.net/2103.022
  * Barstow et al. (2017) Barstow, J. K., Aigrain, S., Irwin, P. G. J., & Sing, D. K. 2017, ApJ, 834, 50, doi: 10.3847/1538-4357/834/1/50
  * Baydin et al. (2018) Baydin, A. G., Pearlmutter, B. A., Radul, A. A., & Siskind, J. M. 2018, Journal of Marchine Learning Research, 18, 1
  * Berg et al. (2018) Berg, R. v. d., Hasenclever, L., Tomczak, J. M., & Welling, M. 2018, arXiv preprint arXiv:1803.05649
  * Blei et al. (2016) Blei, D. M., Kucukelbir, A., & McAuliffe, J. D. 2016, arXiv e-prints, arXiv:1601.00670.  https://arxiv.org/abs/1601.00670
  * Bogachev et al. (2005) Bogachev, V. I., Kolesnikov, A. V., & Medvedev, K. V. 2005, Sbornik: Mathematics, 196, 309, doi: 10.1070/SM2005v196n03ABEH000882
  * Boucher et al. (2021) Boucher, A., Darveau-Bernier, A., Pelletier, S., et al. 2021, AJ, 162, 233, doi: 10.3847/1538-3881/ac1f8e
  * Bradbury et al. (2018) Bradbury, J., Frostig, R., Hawkins, P., et al. 2018, JAX: composable transformations of Python+NumPy programs, 0.2.5.  http://github.com/google/jax
  * Brogi & Line (2019) Brogi, M., & Line, M. R. 2019, AJ, 157, 114, doi: 10.3847/1538-3881/aaffd3
  * Buchholz et al. (2018) Buchholz, A., Wenzel, F., & Mandt, S. 2018, in Proceedings of Machine Learning Research, Vol. 80, Proceedings of the 35th International Conference on Machine Learning, ed. J. Dy & A. Krause (PMLR), 668–677.  https://proceedings.mlr.press/v80/buchholz18a.html
  * Cai et al. (2021) Cai, S., Mao, Z., Wang, Z., Yin, M., & Karniadakis, G. E. 2021, Acta Mechanica Sinica, 37, 1727, doi: 10.1007/s10409-021-01148-1
  * Carone et al. (2021) Carone, L., Mollière, P., Zhou, Y., et al. 2021, A&A, 646, A168, doi: 10.1051/0004-6361/202038620
  * Challener & Rauscher (2022) Challener, R. C., & Rauscher, E. 2022, AJ, 163, 117, doi: 10.3847/1538-3881/ac4885
  * Changeat (2022) Changeat, Q. 2022, AJ, 163, 106, doi: 10.3847/1538-3881/ac4475
  * Changeat & Yip (2022) Changeat, Q., & Yip, K. H. 2022, arXiv e-prints, arXiv:2206.14633.  https://arxiv.org/abs/2206.14633
  * Changeat & Al-Refaie (2020) Changeat, Q., & Al-Refaie, A. 2020, ApJ, 898, 155, doi: 10.3847/1538-4357/ab9b82
  * Changeat et al. (2020) Changeat, Q., Al-Refaie, A., Mugnai, L. V., et al. 2020, AJ, 160, 80, doi: 10.3847/1538-3881/ab9a53
  * Changeat et al. (2021) Changeat, Q., Al-Refaie, A. F., Edwards, B., Waldmann, I. P., & Tinetti, G. 2021, ApJ, 913, 73, doi: 10.3847/1538-4357/abf2bb
  * Changeat & Edwards (2021) Changeat, Q., & Edwards, B. 2021, ApJ, 907, L22, doi: 10.3847/2041-8213/abd84f
  * Changeat et al. (2019) Changeat, Q., Edwards, B., Waldmann, I. P., & Tinetti, G. 2019, ApJ, 886, 39, doi: 10.3847/1538-4357/ab4a14
  * Changeat et al. (2022) Changeat, Q., Edwards, B., Al-Refaie, A. F., et al. 2022, ApJS, 260, 3, doi: 10.3847/1538-4365/ac5cc2
  * Chen et al. (2020) Chen, Y., Lu, L., Karniadakis, G. E., & Dal Negro, L. 2020, Optics express, 28, 11618
  * Chubb et al. (2020) Chubb, K. L., Min, M., Kawashima, Y., Helling, C., & Waldmann, I. 2020, A&A, 639, A3, doi: 10.1051/0004-6361/201937267
  * Chérief-Abdellatif & Alquier (2018) Chérief-Abdellatif, B.-E., & Alquier, P. 2018, Electronic Journal of Statistics, 12, doi: 10.1214/18-ejs1475
  * Cobb et al. (2019) Cobb, A. D., Himes, M. D., Soboczenski, F., et al. 2019, AJ, 158, 33, doi: 10.3847/1538-3881/ab2390
  * Collette (2013) Collette, A. 2013, Python and HDF5 (O’Reilly)
  * Cubillos & Blecic (2021) Cubillos, P. E., & Blecic, J. 2021, MNRAS, 505, 2675, doi: 10.1093/mnras/stab1405
  * Cuomo et al. (2022) Cuomo, S., Schiano di Cola, V., Giampaolo, F., et al. 2022, arXiv e-prints, arXiv:2201.05624.  https://arxiv.org/abs/2201.05624
  * Dillon et al. (2017) Dillon, J. V., Langmore, I., Tran, D., et al. 2017, arXiv preprint arXiv:1711.10604
  * Dinh et al. (2015) Dinh, L., Krueger, D., & Bengio, Y. 2015, in 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Workshop Track Proceedings, ed. Y. Bengio & Y. LeCun.  http://arxiv.org/abs/1410.8516
  * Dinh et al. (2016) Dinh, L., Sohl-Dickstein, J., & Bengio, S. 2016, arXiv preprint arXiv:1605.08803
  * Duane et al. (1987) Duane, S., Kennedy, A. D., Pendleton, B. J., & Roweth, D. 1987, Physics Letters B, 195, 216, doi: 10.1016/0370-2693(87)91197-X
  * Durkan et al. (2019a) Durkan, C., Bekasov, A., Murray, I., & Papamakarios, G. 2019a, Advances in neural information processing systems, 32
  * Durkan et al. (2019b) —. 2019b, Advances in neural information processing systems, 32
  * Edwards et al. (2019) Edwards, B., Rice, M., Zingales, T., et al. 2019, Experimental Astronomy, 47, 29, doi: 10.1007/s10686-018-9611-4
  * Edwards et al. (2021) Edwards, B., Changeat, Q., Mori, M., et al. 2021, AJ, 161, 44, doi: 10.3847/1538-3881/abc6a5
  * Edwards et al. (2023) Edwards, B., Changeat, Q., Tsiaras, A., et al. 2023, ApJS, 269, 31, doi: 10.3847/1538-4365/ac9f1a
  * Evans et al. (2016) Evans, T. M., Sing, D. K., Wakeford, H. R., et al. 2016, ApJ, 822, L4, doi: 10.3847/2041-8205/822/1/L4
  * Fellows et al. (2018) Fellows, M., Mahajan, A., Rudner, T. G. J., & Whiteson, S. 2018, arXiv e-prints, arXiv:1811.01132.  https://arxiv.org/abs/1811.01132
  * Feng et al. (2022) Feng, S., Xu, K., Wu, J., et al. 2022, arXiv preprint arXiv:2205.07493
  * Feng et al. (2020) Feng, Y. K., Line, M. R., & Fortney, J. J. 2020, AJ, 160, 137, doi: 10.3847/1538-3881/aba8f9
  * Fisher & Heng (2022) Fisher, C., & Heng, K. 2022, ApJ, 934, 31, doi: 10.3847/1538-4357/ac7801
  * Foote et al. (2022) Foote, T. O., Lewis, N. K., Kilpatrick, B. M., et al. 2022, AJ, 163, 7, doi: 10.3847/1538-3881/ac2f4a
  * Foreman-Mackey (2016) Foreman-Mackey, D. 2016, The Journal of Open Source Software, 1, 24, doi: 10.21105/joss.00024
  * Foreman-Mackey et al. (2021) Foreman-Mackey, D., Luger, R., Agol, E., et al. 2021, The Journal of Open Source Software, 6, 3285, doi: 10.21105/joss.03285
  * Fortuin et al. (2020) Fortuin, V., Baranchuk, D., Raetsch, G., & Mandt, S. 2020, in Proceedings of Machine Learning Research, Vol. 108, Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, ed. S. Chiappa & R. Calandra (PMLR), 1651–1661.  https://proceedings.mlr.press/v108/fortuin20a.html
  * Friston et al. (2020) Friston, K. J., Parr, T., Zeidman, P., et al. 2020, Wellcome open research, 5
  * Gandhi et al. (2019) Gandhi, S., Madhusudhan, N., Hawker, G., & Piette, A. 2019, AJ, 158, 228, doi: 10.3847/1538-3881/ab4efc
  * Gebhard et al. (2020) Gebhard, T. D., Bonse, M. J., Quanz, S. P., & Schölkopf, B. 2020, arXiv e-prints, arXiv:2010.05591.  https://arxiv.org/abs/2010.05591
  * Gebhard et al. (2022) —. 2022, A&A, doi: https://doi.org/10.1051/0004-6361/202142529
  * Germain et al. (2015) Germain, M., Gregor, K., Murray, I., & Larochelle, H. 2015, in International Conference on Machine Learning, PMLR, 881–889
  * Gibson et al. (2020) Gibson, N. P., Merritt, S., Nugroho, S. K., et al. 2020, MNRAS, 493, 2215, doi: 10.1093/mnras/staa228
  * Grathwohl et al. (2018) Grathwohl, W., Chen, R. T., Bettencourt, J., Sutskever, I., & Duvenaud, D. 2018, arXiv preprint arXiv:1810.01367
  * Greene et al. (2016) Greene, T. P., Line, M. R., Montero, C., et al. 2016, ApJ, 817, 17, doi: 10.3847/0004-637X/817/1/17
  * Gudovskiy et al. (2022) Gudovskiy, D., Ishizaka, S., & Kozuka, K. 2022, in Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (WACV), 98–107
  * Haghighat et al. (2021) Haghighat, E., Raissi, M., Moure, A., Gomez, H., & Juanes, R. 2021, Computer Methods in Applied Mechanics and Engineering, 379, 113741
  * Haldemann et al. (2022) Haldemann, J., Ksoll, V., Walter, D., et al. 2022, arXiv e-prints, arXiv:2202.00027.  https://arxiv.org/abs/2202.00027
  * Harrington et al. (2022) Harrington, J., Himes, M. D., Cubillos, P. E., et al. 2022, , 3, 80, doi: 10.3847/PSJ/ac3513
  * Harris et al. (2020) Harris, C. R., Millman, K. J., van der Walt, S. J., et al. 2020, Nature, 585, 357, doi: 10.1038/s41586-020-2649-2
  * Hayes et al. (2020) Hayes, J. J. C., Kerins, E., Awiphan, S., et al. 2020, MNRAS, 494, 4492, doi: 10.1093/mnras/staa978
  * Haynes et al. (2015) Haynes, K., Mandell, A. M., Madhusudhan, N., Deming, D., & Knutson, H. 2015, The Astrophysical Journal, 806, 146.  http://stacks.iop.org/0004-637X/806/i=2/a=146
  * He et al. (2015) He, K., Zhang, X., Ren, S., & Sun, J. 2015, Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification, arXiv, doi: 10.48550/ARXIV.1502.01852
  * Hellier et al. (2011) Hellier, C., Anderson, D. R., Collier Cameron, A., et al. 2011, A&A, 535, L7, doi: 10.1051/0004-6361/201117081
  * Himes et al. (2022) Himes, M. D., Harrington, J., Cobb, A. D., et al. 2022, , 3, 91, doi: 10.3847/PSJ/abe3fd
  * Hobson et al. (2002) Hobson, M. P., Bridle, S. L., & Lahav, O. 2002, MNRAS, 335, 377, doi: 10.1046/j.1365-8711.2002.05614.x
  * Hoffman & Gelman (2011) Hoffman, M. D., & Gelman, A. 2011, arXiv e-prints, arXiv:1111.4246.  https://arxiv.org/abs/1111.4246
  * Hu et al. (2020) Hu, Y., Anderson, L., Li, T.-M., et al. 2020, ICLR
  * Hu et al. (2019) Hu, Y., Li, T.-M., Anderson, L., Ragan-Kelley, J., & Durand, F. 2019, ACM Transactions on Graphics (TOG), 38, 201
  * Hunter (2007) Hunter, J. D. 2007, Computing in Science & Engineering, 9, 90, doi: 10.1109/MCSE.2007.55
  * Ih & Kempton (2021) Ih, J., & Kempton, E. M. R. 2021, AJ, 162, 237, doi: 10.3847/1538-3881/ac173b
  * Ioffe & Szegedy (2015) Ioffe, S., & Szegedy, C. 2015, arXiv e-prints, arXiv:1502.03167.  https://arxiv.org/abs/1502.03167
  * Irwin et al. (2020) Irwin, P. G. J., Parmentier, V., Taylor, J., et al. 2020, MNRAS, 493, 106, doi: 10.1093/mnras/staa238
  * Irwin et al. (2008) Irwin, P. G. J., Teanby, N. A., de Kok, R., et al. 2008, J. Quant. Spec. Radiat. Transf., 109, 1136, doi: 10.1016/j.jqsrt.2007.11.006
  * Izmailov et al. (2020) Izmailov, P., Kirichenko, P., Finzi, M., & Wilson, A. G. 2020, in Proceedings of Machine Learning Research, Vol. 119, Proceedings of the 37th International Conference on Machine Learning, ed. H. D. III & A. Singh (PMLR), 4615–4630.  https://proceedings.mlr.press/v119/izmailov20a.html
  * Jeffreys (1998) Jeffreys, H. 1998, The Theory of Probability, Oxford Classic Texts in the Physical Sciences (OUP Oxford).  https://books.google.co.uk/books?id=vh9Act9rtzQC
  * Jin et al. (2019) Jin, L., Doshi-Velez, F., Miller, T., Schwartz, L., & Schuler, W. 2019, in Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics (Florence, Italy: Association for Computational Linguistics), 2442–2452, doi: 10.18653/v1/P19-1234
  * Karchev et al. (2022) Karchev, K., Coogan, A., & Weniger, C. 2022, MNRAS, 512, 661, doi: 10.1093/mnras/stac311
  * Kawahara et al. (2022) Kawahara, H., Kawashima, Y., Masuda, K., et al. 2022, ApJS, 258, 31, doi: 10.3847/1538-4365/ac3b4d
  * Kingma & Ba (2014) Kingma, D. P., & Ba, J. 2014, arXiv e-prints, arXiv:1412.6980.  https://arxiv.org/abs/1412.6980
  * Kingma & Dhariwal (2018) Kingma, D. P., & Dhariwal, P. 2018, arXiv:1807.03039 [cs, stat].  http://arxiv.org/abs/1807.03039
  * Kingma et al. (2016) Kingma, D. P., Salimans, T., Jozefowicz, R., et al. 2016, Improving Variational Inference with Inverse Autoregressive Flow, arXiv, doi: 10.48550/ARXIV.1606.04934
  * Kingma & Welling (2013) Kingma, D. P., & Welling, M. 2013, arXiv e-prints, arXiv:1312.6114.  https://arxiv.org/abs/1312.6114
  * Kobyzev et al. (2021) Kobyzev, I., Prince, S. J., & Brubaker, M. A. 2021, IEEE Transactions on Pattern Analysis and Machine Intelligence, 43, 3964, doi: 10.1109/tpami.2020.2992934
  * Kreidberg et al. (2014) Kreidberg, L., Bean, J. L., Désert, J.-M., et al. 2014, ApJ, 793, L27, doi: 10.1088/2041-8205/793/2/L27
  * Kreidberg et al. (2018) Kreidberg, L., Line, M. R., Parmentier, V., et al. 2018, AJ, 156, 17, doi: 10.3847/1538-3881/aac3df
  * Krick et al. (2020) Krick, J. E., Fraine, J., Ingalls, J., & Deger, S. 2020, AJ, 160, 99, doi: 10.3847/1538-3881/aba11f
  * Kullback & Leibler (1951) Kullback, S., & Leibler, R. A. 1951, The annals of mathematical statistics, 22, 79
  * Lavie et al. (2017) Lavie, B., Mendonça, J. M., Mordasini, C., et al. 2017, AJ, 154, 91, doi: 10.3847/1538-3881/aa7ed8
  * Lee et al. (2013) Lee, J.-M., Heng, K., & Irwin, P. G. J. 2013, ApJ, 778, 97, doi: 10.1088/0004-637X/778/2/97
  * Lee et al. (2014) Lee, J.-M., Irwin, P. G. J., Fletcher, L. N., Heng, K., & Barstow, J. K. 2014, ApJ, 789, 14, doi: 10.1088/0004-637X/789/1/14
  * Li et al. (2015) Li, G., Gordon, I. E., Rothman, L. S., et al. 2015, ApJS, 216, 15, doi: 10.1088/0067-0049/216/1/15
  * Line et al. (2014) Line, M. R., Knutson, H., Wolf, A. S., & Yung, Y. L. 2014, ApJ, 783, 70, doi: 10.1088/0004-637X/783/2/70
  * Line et al. (2013) Line, M. R., Wolf, A. S., Zhang, X., et al. 2013, ApJ, 775, 137, doi: 10.1088/0004-637X/775/2/137
  * Line et al. (2016) Line, M. R., Stevenson, K. B., Bean, J., et al. 2016, AJ, 152, 203, doi: 10.3847/0004-6256/152/6/203
  * Lopez-Alvis et al. (2021) Lopez-Alvis, J., Laloy, E., Nguyen, F., & Hermans, T. 2021, Computers & Geosciences, 152, 104762
  * Lopez-Alvis et al. (2022) Lopez-Alvis, J., Nguyen, F., Looms, M. C., & Hermans, T. 2022, Journal of Geophysical Research (Solid Earth), 127, e22581, doi: 10.1029/2021JB022581
  * Lothringer & Barman (2020) Lothringer, J. D., & Barman, T. S. 2020, AJ, 159, 289, doi: 10.3847/1538-3881/ab8d33
  * Louizos & Welling (2017) Louizos, C., & Welling, M. 2017, in Proceedings of Machine Learning Research, Vol. 70, Proceedings of the 34th International Conference on Machine Learning, ed. D. Precup & Y. W. Teh (PMLR), 2218–2227.  https://proceedings.mlr.press/v70/louizos17a.html
  * Lugmayr et al. (2020) Lugmayr, A., Danelljan, M., Gool, L. V., & Timofte, R. 2020, in European conference on computer vision, Springer, 715–732
  * MacDonald et al. (2020) MacDonald, R. J., Goyal, J. M., & Lewis, N. K. 2020, ApJ, 893, L43, doi: 10.3847/2041-8213/ab8238
  * MacDonald & Lewis (2022) MacDonald, R. J., & Lewis, N. K. 2022, ApJ, 929, 20, doi: 10.3847/1538-4357/ac47fe
  * MacDonald & Madhusudhan (2017) MacDonald, R. J., & Madhusudhan, N. 2017, MNRAS, 469, 1979, doi: 10.1093/mnras/stx804
  * MacDonald & Madhusudhan (2019) —. 2019, MNRAS, 486, 1292, doi: 10.1093/mnras/stz789
  * Madhusudhan (2018) Madhusudhan, N. 2018, Atmospheric Retrieval of Exoplanets, ed. H. J. Deeg & J. A. Belmonte, 104, doi: 10.1007/978-3-319-55333-7_104
  * Madhusudhan & Seager (2009) Madhusudhan, N., & Seager, S. 2009, ApJ, 707, 24, doi: 10.1088/0004-637X/707/1/24
  * Mansfield et al. (2021) Mansfield, M., Line, M. R., Bean, J. L., et al. 2021, Nature Astronomy, 5, 1224, doi: 10.1038/s41550-021-01455-4
  * Mansfield et al. (2022) Mansfield, M., Wiser, L., Stevenson, K. B., et al. 2022, AJ, 163, 261, doi: 10.3847/1538-3881/ac658f
  * Márquez-Neila et al. (2018) Márquez-Neila, P., Fisher, C., Sznitman, R., & Heng, K. 2018, Nature Astronomy, 2, 719, doi: 10.1038/s41550-018-0504-2
  * Medvedev (2008) Medvedev, K. 2008, Theory of Stochastic Processes, 1
  * Meech et al. (2022) Meech, A., Aigrain, S., Brogi, M., & Birkby, J. L. 2022, MNRAS, 512, 2604, doi: 10.1093/mnras/stac662
  * Mikal-Evans et al. (2019) Mikal-Evans, T., Sing, D. K., Goyal, J. M., et al. 2019, MNRAS, 488, 2222, doi: 10.1093/mnras/stz1753
  * Mikal-Evans et al. (2022) Mikal-Evans, T., Sing, D. K., Barstow, J. K., et al. 2022, Nature Astronomy, 6, 471, doi: 10.1038/s41550-021-01592-w
  * Min et al. (2020) Min, M., Ormel, C. W., Chubb, K., Helling, C., & Kawashima, Y. 2020, A&A, 642, A28, doi: 10.1051/0004-6361/201937377
  * Mollière et al. (2020) Mollière, P., Stolker, T., Lacour, S., et al. 2020, A&A, 640, A131, doi: 10.1051/0004-6361/202038325
  * Morvan et al. (2019) Morvan, M., Nikolaou, N., Tsiaras, A., & Waldmann, I. 2019, ApJ
  * Morvan et al. (2021a) Morvan, M., Tsiaras, A., Nikolaou, N., & Waldmann, I. P. 2021a, PASP, 133, 034505, doi: 10.1088/1538-3873/abe6e8
  * Morvan et al. (2021b) —. 2021b, PASP, 133, 034505, doi: 10.1088/1538-3873/abe6e8
  * Mugnai et al. (2020) Mugnai, L. V., Pascale, E., Edwards, B., Papageorgiou, A., & Sarkar, S. 2020, Experimental Astronomy, 50, 303, doi: 10.1007/s10686-020-09676-7
  * Mugnai et al. (2021) Mugnai, L. V., Modirrousta-Galian, D., Edwards, B., et al. 2021, AJ, 161, 284, doi: 10.3847/1538-3881/abf3c3
  * Nielsen et al. (2020) Nielsen, D., Jaini, P., Hoogeboom, E., Winther, O., & Welling, M. 2020, Advances in Neural Information Processing Systems, 33, 12685
  * Nikolaou et al. (2020) Nikolaou, N., Waldmann, I. P., Tsiaras, A., et al. 2020, arXiv e-prints, arXiv:2010.15996.  https://arxiv.org/abs/2010.15996
  * Oreshenko et al. (2020) Oreshenko, M., Kitzmann, D., Márquez-Neila, P., et al. 2020, AJ, 159, 6, doi: 10.3847/1538-3881/ab5955
  * Padilla et al. (2019) Padilla, L. E., Tellez, L. O., Escamilla, L. A., & Vazquez, J. A. 2019, arXiv e-prints, arXiv:1903.11127.  https://arxiv.org/abs/1903.11127
  * pandas development team (2020) pandas development team, T. 2020, pandas-dev/pandas: Pandas, latest, Zenodo, doi: 10.5281/zenodo.3509134
  * Papamakarios et al. (2017) Papamakarios, G., Pavlakou, T., & Murray, I. 2017, Advances in neural information processing systems, 30
  * Paszke et al. (2019) Paszke, A., Gross, S., Massa, F., et al. 2019, in Advances in Neural Information Processing Systems 32, ed. H. Wallach, H. Larochelle, A. Beygelzimer, F. d Alché-Buc, E. Fox, & R. Garnett (Curran Associates, Inc.), 8024–8035.  http://papers.neurips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf
  * Pati et al. (2018) Pati, D., Bhattacharya, A., & Yang, Y. 2018, in Proceedings of Machine Learning Research, Vol. 84, Proceedings of the Twenty-First International Conference on Artificial Intelligence and Statistics, ed. A. Storkey & F. Perez-Cruz (PMLR), 1579–1588.  https://proceedings.mlr.press/v84/pati18a.html
  * Pinhas et al. (2019) Pinhas, A., Madhusudhan, N., Gandhi, S., & MacDonald, R. 2019, MNRAS, 482, 1485, doi: 10.1093/mnras/sty2544
  * Pluriel et al. (2020a) Pluriel, W., Zingales, T., Leconte, J., & Parmentier, V. 2020a, A&A, 636, A66, doi: 10.1051/0004-6361/202037678
  * Pluriel et al. (2020b) Pluriel, W., Whiteford, N., Edwards, B., et al. 2020b, AJ, 160, 112, doi: 10.3847/1538-3881/aba000
  * Polyansky et al. (2018) Polyansky, O. L., Kyuberis, A. A., Zobov, N. F., et al. 2018, Monthly Notices of the Royal Astronomical Society, 480, 2597
  * Prenger et al. (2019) Prenger, R., Valle, R., & Catanzaro, B. 2019, in ICASSP 2019-2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), IEEE, 3617–3621
  * Quionero-Candela et al. (2009) Quionero-Candela, J., Sugiyama, M., Schwaighofer, A., & Lawrence, N. D. 2009, Dataset Shift in Machine Learning (The MIT Press)
  * Raissi et al. (2019) Raissi, M., Perdikaris, P., & Karniadakis, G. 2019, Journal of Computational Physics, 378, 686, doi: https://doi.org/10.1016/j.jcp.2018.10.045
  * Rasmussen et al. (2022) Rasmussen, K. C., Brogi, M., Rahman, F., et al. 2022, AJ, 164, 35, doi: 10.3847/1538-3881/ac6bfa
  * Rasul et al. (2020) Rasul, K., Sheikh, A.-S., Schuster, I., Bergmann, U., & Vollgraf, R. 2020, arXiv preprint arXiv:2002.06103
  * Rezende & Mohamed (2015) Rezende, D. J., & Mohamed, S. 2015, Variational Inference with Normalizing Flows, arXiv, doi: 10.48550/ARXIV.1505.05770
  * Rippel & Adams (2013) Rippel, O., & Adams, R. P. 2013, High-Dimensional Probability Estimation with Deep Density Models, arXiv, doi: 10.48550/ARXIV.1302.5125
  * Rocchetto et al. (2016) Rocchetto, M., Waldmann, I. P., Venot, O., Lagage, P. O., & Tinetti, G. 2016, ApJ, 833, 120, doi: 10.3847/1538-4357/833/1/120
  * Rothfuss et al. (2019) Rothfuss, J., Ferreira, F., Boehm, S., et al. 2019, arXiv preprint arXiv:1907.08982
  * Roudier et al. (2021) Roudier, G. M., Swain, M. R., Gudipati, M. S., et al. 2021, AJ, 162, 37, doi: 10.3847/1538-3881/abfdad
  * Rudolph et al. (2021) Rudolph, M., Wandt, B., & Rosenhahn, B. 2021, in 2021 IEEE Winter Conference on Applications of Computer Vision (WACV) (Los Alamitos, CA, USA: IEEE Computer Society), 1906–1915, doi: 10.1109/WACV48630.2021.00195
  * Saba et al. (2022) Saba, A., Tsiaras, A., Morvan, M., et al. 2022, AJ, 164, 2, doi: 10.3847/1538-3881/ac6c01
  * Schmidt & Simic (2019) Schmidt, M., & Simic, M. 2019, arXiv preprint arXiv:1906.06904
  * Seidel et al. (2020) Seidel, J. V., Ehrenreich, D., Pino, L., et al. 2020, A&A, 633, A86, doi: 10.1051/0004-6361/201936892
  * Shallue & Vanderburg (2018) Shallue, C. J., & Vanderburg, A. 2018, AJ, 155, 94, doi: 10.3847/1538-3881/aa9e09
  * Sheppard et al. (2017) Sheppard, K. B., Mandell, A. M., Tamburo, P., et al. 2017, ApJ, 850, L32, doi: 10.3847/2041-8213/aa9ae9
  * Sheppard et al. (2021) Sheppard, K. B., Welbanks, L., Mandell, A. M., et al. 2021, AJ, 161, 51, doi: 10.3847/1538-3881/abc8f4
  * Shu et al. (2018) Shu, R., Bui, H. H., Zhao, S., Kochenderfer, M. J., & Ermon, S. 2018, Advances in Neural Information Processing Systems, 31
  * Sing et al. (2016) Sing, D. K., Fortney, J. J., Nikolov, N., et al. 2016, Nature, 529, 59, doi: 10.1038/nature16068
  * Skaf et al. (2020) Skaf, N., Bieger, M. F., Edwards, B., et al. 2020, AJ, 160, 109, doi: 10.3847/1538-3881/ab94a3
  * Smith (2015) Smith, L. N. 2015, arXiv e-prints, arXiv:1506.01186.  https://arxiv.org/abs/1506.01186
  * Stevenson et al. (2017) Stevenson, K. B., Line, M. R., Bean, J. L., et al. 2017, AJ, 153, 68, doi: 10.3847/1538-3881/153/2/68
  * Sun et al. (2022) Sun, H., Bouman, K. L., Tiede, P., et al. 2022, arXiv e-prints, arXiv:2201.08506.  https://arxiv.org/abs/2201.08506
  * Swain et al. (2021) Swain, M. R., Estrela, R., Roudier, G. M., et al. 2021, AJ, 161, 213, doi: 10.3847/1538-3881/abe879
  * Tinetti et al. (2007) Tinetti, G., Vidal-Madjar, A., Liang, M.-C., et al. 2007, Nature, 448, 169, doi: 10.1038/nature06002
  * Tinetti et al. (2021) Tinetti, G., Eccleston, P., Haswell, C., et al. 2021, arXiv e-prints, arXiv:2104.04824.  https://arxiv.org/abs/2104.04824
  * Tsiaras et al. (2016a) Tsiaras, A., Waldmann, I. P., Rocchetto, M., et al. 2016a, ApJ, 832, 202, doi: 10.3847/0004-637X/832/2/202
  * Tsiaras et al. (2019) Tsiaras, A., Waldmann, I. P., Tinetti, G., Tennyson, J., & Yurchenko, S. N. 2019, Nature Astronomy, 3, 1086, doi: 10.1038/s41550-019-0878-9
  * Tsiaras et al. (2016b) Tsiaras, A., Rocchetto, M., Waldmann, I. P., et al. 2016b, ApJ, 820, 99, doi: 10.3847/0004-637X/820/2/99
  * Udry et al. (2014) Udry, S., Lovis, C., Bouchy, F., et al. 2014, arXiv e-prints, arXiv:1412.1048.  https://arxiv.org/abs/1412.1048
  * Valizadegan et al. (2021) Valizadegan, H., Martinho, M., Wilkens, L. S., et al. 2021, arXiv e-prints, arXiv:2111.10009.  https://arxiv.org/abs/2111.10009
  * van den Oord et al. (2018) van den Oord, A., Li, Y., Babuschkin, I., et al. 2018, in Proceedings of Machine Learning Research, Vol. 80, Proceedings of the 35th International Conference on Machine Learning, ed. J. Dy & A. Krause (PMLR), 3918–3926.  https://proceedings.mlr.press/v80/oord18a.html
  * Viana & Subramaniyan (2021) Viana, F. A., & Subramaniyan, A. K. 2021, Archives of Computational Methods in Engineering, 28, 3801
  * Virtanen et al. (2020) Virtanen, P., Gommers, R., Oliphant, T. E., et al. 2020, Nature Methods, 17, 261, doi: 10.1038/s41592-019-0686-2
  * von Essen et al. (2020) von Essen, C., Mallonn, M., Hermansen, S., et al. 2020, A&A, 637, A76, doi: 10.1051/0004-6361/201937169
  * Waldmann & Griffith (2019) Waldmann, I. P., & Griffith, C. A. 2019, Nature Astronomy, 3, 620, doi: 10.1038/s41550-019-0753-8
  * Waldmann et al. (2015) Waldmann, I. P., Rocchetto, M., Tinetti, G., et al. 2015, ApJ, 813, 13, doi: 10.1088/0004-637X/813/1/13
  * Wu et al. (2020) Wu, H., Köhler, J., & Noé, F. 2020, Advances in Neural Information Processing Systems, 33, 5933
  * Yip et al. (2021a) Yip, K. H., Changeat, Q., Edwards, B., et al. 2021a, AJ, 161, 4, doi: 10.3847/1538-3881/abc179
  * Yip et al. (2021b) Yip, K. H., Changeat, Q., Nikolaou, N., et al. 2021b, AJ, 162, 195, doi: 10.3847/1538-3881/ac1744
  * Yip et al. (2020) Yip, K. H., Nikolaou, N., Coronica, P., et al. 2020, in Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2019, Würzburg, Germany, September 16–20, 2019, Proceedings, Part III, Springer, 322–338
  * Yip et al. (2022) Yip, K. H., Waldmann, I. P., Changeat, Q., et al. 2022, arXiv e-prints, arXiv:2206.14642.  https://arxiv.org/abs/2206.14642
  * Yu et al. (2019) Yu, L., Vanderburg, A., Huang, C., et al. 2019, AJ, 158, 25, doi: 10.3847/1538-3881/ab21d6
  * Yurchenko et al. (2017) Yurchenko, S. N., Amundsen, D. S., Tennyson, J., & Waldmann, I. P. 2017, A&A, 605, A95, doi: 10.1051/0004-6361/201731026
  * Yurchenko et al. (2011) Yurchenko, S. N., Barber, R. J., & Tennyson, J. 2011, MNRAS, 413, 1828, doi: 10.1111/j.1365-2966.2011.18261.x
  * Yurchenko et al. (2020) Yurchenko, S. N., Mellor, T. M., Freedman, R. S., & Tennyson, J. 2020, MNRAS, 496, 5282, doi: 10.1093/mnras/staa1874
  * Zhang et al. (2019) Zhang, C., Butepage, J., Kjellstrom, H., & Mandt, S. 2019, IEEE Transactions on Pattern Analysis & Machine Intelligence, 41, 2008, doi: 10.1109/TPAMI.2018.2889774
  * Zhang et al. (2019) Zhang, M., Chachan, Y., Kempton, E. M. R., & Knutson, H. A. 2019, PASP, 131, 034501, doi: 10.1088/1538-3873/aaf5ad
  * Zhang et al. (2020) Zhang, M., Chachan, Y., Kempton, E. M. R., Knutson, H. A., & Chang, W. H. 2020, ApJ, 899, 27, doi: 10.3847/1538-4357/aba1e6
  * Zhang & Chen (2021) Zhang, Q., & Chen, Y. 2021, Advances in Neural Information Processing Systems, 34, 16280
  * Zhang et al. (2021) Zhang, X., Nawaz, M. A., Zhao, X., & Curtis, A. 2021, in Advances in Geophysics, Vol. 62 (Elsevier), 73–140
  * Zingales & Waldmann (2018) Zingales, T., & Waldmann, I. P. 2018, AJ, 156, 268, doi: 10.3847/1538-3881/aae77c