

\subsubsection*{#1}}

\pagestyle{headings}
\markright{Reference sheet: \texttt{natbib}}

\usepackage{shortvrb}
\MakeShortVerb{\|}

\begin{document}
\thispagestyle{plain}

\newcommand{\textsc{Bib}\TeX}{\textsc{Bib}\TeX}
\newcommand{\texttt{#1}\def\filedate{#2}\def\fileversion{#3}}}{\texttt{#1}\def\filedate{#2}\def\fileversion{#3}}}
\begin{center}{\bfseries\Large
    Reference sheet for \texttt{#1}\def\filedate{#2}\def\fileversion{#3}}\ usage}\\
    \large(Describing version \fileversion\ from \filedate)
\end{center}

\begin{quote}\slshape
For a more detailed description of the \texttt{#1}\def\filedate{#2}\def\fileversion{#3}}\ package, \LaTeX\ the
source file \texttt{#1}\def\filedate{#2}\def\fileversion{#3}}\texttt{.dtx}.
\end{quote}
\head{Overview}
The \texttt{#1}\def\filedate{#2}\def\fileversion{#3}}\ package is a reimplementation of the \LaTeX\ |\cite| command,
to work with both author--year and numerical citations. It is compatible with
the standard bibliographic style files, such as \texttt{plain.bst}, as well as
with those for \texttt{harvard}, \texttt{apalike}, \texttt{chicago},
\texttt{astron}, \texttt{authordate}, and of course \texttt{#1}\def\filedate{#2}\def\fileversion{#3}}.

\head{Loading}
Load with |\usepackage[|\emph{options}|]{|\texttt{#1}\def\filedate{#2}\def\fileversion{#3}}|}|. See list of
\emph{options} at the end.

\head{Replacement bibliography styles}
I provide three new \texttt{.bst} files to replace the standard \LaTeX\
numerical ones:
\begin{quote}\ttfamily
 plainnat.bst \qquad abbrvnat.bst \qquad unsrtnat.bst
\end{quote}
\head{Basic commands}
The \texttt{#1}\def\filedate{#2}\def\fileversion{#3}}\ package has two basic citation commands, |\citet| and
|\citep| for \emph{textual} and \emph{parenthetical} citations, respectively.
There also exist the starred versions |\citet*| and |\citep*| that print
the full author list, and not just the abbreviated one.
All of these may take one or two optional arguments to add some text before
and after the citation.
\begin{quote}
\begin{tabular}{l@{\quad$\Rightarrow$\quad}l}
  |\citet{jon90}| & Jones et al. (1990)\\
  |\citet[chap.~2]{jon90}| & Jones et al. (1990, chap.~2)\\[0.5ex]
  |\citep{jon90}| & (Jones et al., 1990)\\
  |\citep[chap.~2]{jon90}| & (Jones et al., 1990, chap.~2)\\
  |\citep[see][]{jon90}| & (see Jones et al., 1990)\\
  |\citep[see][chap.~2]{jon90}| & (see Jones et al., 1990, chap.~2)\\[0.5ex]
  |\citet*{jon90}| & Jones, Baker, and Williams (1990)\\
  |\citep*{jon90}| & (Jones, Baker, and Williams, 1990)
\end{tabular}
\end{quote}
\head{Multiple citations}
Multiple citations may be made by including more than one
citation key in the |\cite| command argument.
\begin{quote}
\begin{tabular}{l@{\quad$\Rightarrow$\quad}l}
  |\citet{jon90,jam91}| & Jones et al. (1990); James et al. (1991)\\
  |\citep{jon90,jam91}| & (Jones et al., 1990; James et al. 1991)\\
  |\citep{jon90,jon91}| & (Jones et al., 1990, 1991)\\
  |\citep{jon90a,jon90b}| & (Jones et al., 1990a,b)
\end{tabular}
\end{quote}

\head{Numerical mode}
These examples are for author--year citation mode. In numerical mode, the
results are different.
\begin{quote}
\begin{tabular}{l@{\quad$\Rightarrow$\quad}l}
  |\citet{jon90}| & Jones et al. [21]\\
  |\citet[chap.~2]{jon90}| & Jones et al. [21, chap.~2]\\[0.5ex]
  |\citep{jon90}| & [21]\\
  |\citep[chap.~2]{jon90}| & [21, chap.~2]\\
  |\citep[see][]{jon90}| & [see 21]\\
  |\citep[see][chap.~2]{jon90}| & [see 21, chap.~2]\\[0.5ex]
  |\citep{jon90a,jon90b}| & [21, 32]
\end{tabular}
\end{quote}
\head{Suppressed parentheses}
As an alternative form of citation, |\citealt| is the same as |\citet| but
\emph{without parentheses}. Similarly, |\citealp| is |\citep| without
parentheses. Multiple references, notes, and the starred variants
also exist.
\begin{quote}
\begin{tabular}{l@{\quad$\Rightarrow$\quad}l}
  |\citealt{jon90}| & Jones et al.\ 1990\\
  |\citealt*{jon90}| & Jones, Baker, and Williams 1990\\
  |\citealp{jon90}| & Jones et al., 1990\\
  |\citealp*{jon90}| & Jones, Baker, and Williams, 1990\\
  |\citealp{jon90,jam91}| & Jones et al., 1990; James et al., 1991\\
  |\citealp[pg.~32]{jon90}| & Jones et al., 1990, pg.~32\\
  |\citetext{priv.\ comm.}| & (priv.\ comm.)
\end{tabular}
\end{quote}
The |\citetext| command
allows arbitrary text to be placed in the current citation parentheses.
This may be used in combination with |\citealp|.
\head{Partial citations}
In author--year schemes, it is sometimes desirable to be able to refer to
the authors without the year, or vice versa. This is provided with the
extra commands
\begin{quote}
\begin{tabular}{l@{\quad$\Rightarrow$\quad}l}
  |\citeauthor{jon90}| & Jones et al.\\
  |\citeauthor*{jon90}| & Jones, Baker, and Williams\\
  |\citeyear{jon90}|   & 1990\\
  |\citeyearpar{jon90}| & (1990)
\end{tabular}
\end{quote}
\head{Forcing upper cased names}
If the first author's name contains a \textsl{von} part, such as ``della
Robbia'', then |\citet{dRob98}| produces ``della Robbia (1998)'', even at the
beginning of a sentence. One can force the first letter to be in upper case
with the command |\Citet| instead. Other upper case commands also exist.
\begin{quote}
\begin{tabular}{rl@{\quad$\Rightarrow$\quad}l}
  when & |\citet{dRob98}| & della Robbia (1998) \\
  then & |\Citet{dRob98}| & Della Robbia (1998) \\
     &   |\Citep{dRob98}| & (Della Robbia, 1998) \\
     &   |\Citealt{dRob98}| & Della Robbia 1998 \\
     &   |\Citealp{dRob98}| & Della Robbia, 1998 \\
     &   |\Citeauthor{dRob98}| & Della Robbia
\end{tabular}
\end{quote}
These commands also exist in starred versions for full author names.

\head{Citation aliasing}
Sometimes one wants to refer to a reference with a special designation,
rather than by the authors, i.e. as Paper~I, Paper~II. Such aliases can be
defined and used, textual and/or parenthetical with:
\begin{quote}
\begin{tabular}{lcl}
  |\defcitealias{jon90}{Paper~I}|\\
  |\citetalias{jon90}| & $\Rightarrow$ & Paper~I\\
  |\citepalias{jon90}| & $\Rightarrow$ & (Paper~I)
\end{tabular}
\end{quote}
These citation commands function much like |\citet| and |\citep|: they may
take multiple keys in the argument, may contain notes, and are marked as
hyperlinks.
\head{Selecting citation style and punctuation}
Use the command |\bibpunct| with one optional and 6 mandatory arguments:
\begin{enumerate}
\item the opening bracket symbol, default = (
\item the closing bracket symbol, default = )
\item the punctuation between multiple citations, default = ;
\item the letter `n' for numerical style, or `s' for numerical superscript
      style, any other letter for
      author--year, default = author--year;
\item the punctuation that comes between the author names and the year
\item the punctuation that comes between years or numbers when common author
      lists are suppressed (default = ,);
\end{enumerate}

The optional argument is the character preceding a post-note, default is a
comma plus space. In redefining this character, one must include a space if
one is wanted.

Example~1, |\bibpunct{[}{]}{,}{a}{}{;}| changes the output of
\begin{quote}
  |\citep{jon90,jon91,jam92}|
\end{quote}
into [Jones et al. 1990; 1991, James et al. 1992].

Example~2, |\bibpunct[; ]{(}{)}{,}{a}{}{;}| changes the output of
\begin{quote}
  |\citep[and references therein]{jon90}|
\end{quote}
into (Jones et al. 1990; and references therein).

\head{Other formatting options}
Redefine |\bibsection| to the desired sectioning command for introducing
the list of references. This is normally |\section*| or |\chapter*|.

Define |\bibpreamble| to be any text that is to be printed after the heading but
before the actual list of references.

Define |\bibfont| to be a font declaration, e.g.\ |\small| to apply to
the list of references.

Define |\citenumfont| to be a font declaration or command like |\itshape|
or |\textit|.

Redefine |\bibnumfmt| as a command with an argument to format the numbers in
the list of references. The default definition is |[#1]|.

The indentation after the first line of each reference is given by
|\bibhang|; change this with the |\setlength| command.

The vertical spacing between references is set by |\bibsep|; change this with
the |\setlength| command.

\head{Automatic indexing of citations}
If one wishes to have the citations entered in the \texttt{.idx} indexing
file, it is only necessary to issue |\citeindextrue| at any point in the
document. All following |\cite| commands, of all variations, then insert
the corresponding entry to that file. With |\citeindexfalse|, these
entries will no longer be made.

\head{Use with \texttt{chapterbib} package}

The \texttt{#1}\def\filedate{#2}\def\fileversion{#3}}\ package is compatible with the \texttt{chapterbib} package
which makes it possible to have several bibliographies in one document.

The package makes use of the |\include| command, and each |\include|d file
has its own bibliography.

The order in which the \texttt{chapterbib} and \texttt{#1}\def\filedate{#2}\def\fileversion{#3}}\ packages are loaded
is unimportant.

The \texttt{chapterbib} package provides an option \texttt{sectionbib}
that puts the bibliography in a |\section*| instead of |\chapter*|,
something that makes sense if there is a bibliography in each chapter.
This option will not work when \texttt{#1}\def\filedate{#2}\def\fileversion{#3}}\ is also loaded; instead, add
the option to \texttt{#1}\def\filedate{#2}\def\fileversion{#3}}.

Every |\include|d file must contain its own
|\bibliography| command where the bibliography is to appear. The database
files listed as arguments to this command can be different in each file,
of course. However, what is not so obvious, is that each file must also
contain a |\bibliographystyle| command, \emph{preferably with the same
style  argument}.
\head{Sorting and compressing citations}
Do not use the \texttt{cite} package with \texttt{#1}\def\filedate{#2}\def\fileversion{#3}}; rather use one of the
options \texttt{sort} or \texttt{sort\&compress}.

These also work with author--year citations, making multiple citations appear
in their order in the reference list.

\head{Long author list on first citation}
Use option \texttt{longnamesfirst} to have first citation automatically give
the full list of authors.

Suppress this for certain citations with |\shortcites{|\emph{key-list}|}|,
given before the first citation.

\head{Local configuration}
Any local recoding or definitions can be put in \texttt{#1}\def\filedate{#2}\def\fileversion{#3}}\texttt{.cfg} which
is read in after the main package file.

\head{Options that can be added to \texttt{\char`\\ usepackage}}
\begin{description}
\item[\ttfamily round] (default) for round parentheses;
\item[\ttfamily square] for square brackets;
\item[\ttfamily curly] for curly braces;
\item[\ttfamily angle] for angle brackets;
\item[\ttfamily colon] (default) to separate multiple citations with
     colons;
\item[\ttfamily comma] to use commas as separaters;
\item[\ttfamily authoryear] (default) for author--year citations;
\item[\ttfamily numbers] for numerical citations;
\item[\ttfamily super] for superscripted numerical citations, as in
     \textsl{Nature};
\item[\ttfamily sort] orders multiple citations into the sequence in
     which they appear in the list of references;
\item[\ttfamily sort\&compress] as \texttt{sort} but in addition multiple
     numerical citations are compressed if possible (as 3--6, 15);
\item[\ttfamily longnamesfirst] makes the first citation of any reference
     the equivalent of the starred variant (full author list) and subsequent
     citations normal (abbreviated list);
\item[\ttfamily sectionbib] redefines |\thebibliography| to issue
     |\section*| instead of |\chapter*|; valid only for classes with a
     |\chapter| command; to be used with the \texttt{chapterbib} package;
\item[\ttfamily nonamebreak] keeps all the authors' names in a citation on
     one line; causes overfull hboxes but helps with some \texttt{hyperref}
     problems.
\end{description}
\end{document}

\section{Introduction}
Atmospheric retrieval has become an indispensable tool for astronomers to explain individual observations from transit, eclipse and phase curves spectroscopy in both low \citep[e.g.][and references therein]{Tinetti2007, Line_2013, Kreidberg_2014, Line_2014, Lee_2014, Haynes_Wasp33b_spectrum_em,Evans_2016, Tsiaras_2016, Line_2016, Tsiaras_2016_55cnc, Sheppard_2017, Stevenson_2017, MacDo_2017, Kreidberg_2018, Tsiaras2019,Evans_2019, MacDo_2019,Skaf_2020,Pluriel_2020,Zhang_2020,Changeat_2020, Chubb_2020, Alam_2020, Von_essen_2020, Anisman_2020, Carone_2021, Yip_2021, Edwards_2021,Changeat_Edwards_2021, Sheppard_2021, Saba_2021,Alam_2021,Mugnai_2021, Swain_2021, Changeat_w43,Foote_2022, Mansfield_2022_w77, Changeat_2022, evans_2022} and high-resolution \citep[e.g.][]{Brogi_2019,Gibson_2020,Molliere_2020,Seidel_2020,Boucher_2021}. Over the years, the community has come up with a variety of retrieval frameworks, each coupled with different modelling assumptions and sampling techniques \citep[e.g.][]{irwin2008,Madhu_retrieval_method,Line_2013,helios,Gandhi_2019,Zhang_2019,Lothringer_2020,Min_2020,Cubillos_2021,taurex3}. As the number of spectroscopic observations slowly increases, the community has started to look at planetary characterisation on a population level~\citep[][]{Sing_2016, Barstow_2017,Tsiaras2019,Pinhas_2019,Mansfield_2021, Roudier_2021, Changeat_fivekeys}. 


At its core, a 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:
\begin{equation}
    P(\theta|D) = \frac{P(D|\theta)P(\theta)}{P(D)}
\end{equation}
\noindent 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 \citep[see ][for a recent review]{Madhusudhan_retreival_review}. 

Nevertheless, sampling-based techniques are prohibitively slow when the dimensionality of the problem and quantity of the observed data are large \citep[][]{VI_survey}. This issue will become increasingly pressing with the increased data volume originating from the recently launched James Webb Space Telescope  \citep[JWST,][]{greene_jwst_exo} and other future missions such as, Ariel \citep[][]{Tinetti_2021}, Twinkle \citep[][]{edwards_exo} and ELTs from the ground \citep[e.g.][]{Udry14}. These telescopes are designed to provide hundreds of higher resolution spectroscopic measurements with wide 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 \citep[e.g.][]{Rocchetto_2016, Changeat_2019, alfnoor, MacDonald_2020,Pluriel_2020_bias,Ih_2021}, prompting the need for more realistic treatment of the atmosphere \citep[e.g.][]{Irwin_2020,Feng_2020,Changeat_2020}. This will inevitably increase both the computational cost and model complexity \citep[][]{Changeat_w43}. 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~\citep[e.g.][]{morvan,morvan_2,Krick_rf_irac,Gebhard_2020}, to planet detection~\citep[e.g.][]{Shallue_2018,Direct,Yu_tess,exominer} and  characterisation~\citep[e.g.][]{Marquez_rf,Zingales_2018,Cobb_2019,himes2020,Oreshenko_2020,Yip_sensitivity,Haldemann2022,Ardevol_CNN}. 

Variational Inference (VI) is a widely studied and scalable approach in the field of machine learning (ML) to provide approximate posterior distributions for a large and high dimensional data set within a reasonable amount of time \citep[][]{Blei_VI,VI_survey}. However, VI requires a fully differentiable model. 
The field has recently explored different applications of differentiable physical models. Differentiable models open up the possibility to construct ``physics-aware" neural network, a type of network that respects physical laws \citep[][]{RAISSI2019686}. For instance, \cite{Kawahara_2021} used Hamiltonian Monte Carlo (HMC), a gradient-informed Monte Carlo sampling algorithm \citep[][]{Duane_1987,NUTS}, to perform atmospheric retrieval of exoplanets on high resolution spectroscopic data. Others have also applied HMC to speed up lightcurve-fitting \citep[e.g.][]{Foreman_2021,Agol_2021}. 

Here we present the following contributions:
\begin{enumerate}
    \item We present \texttt{Diff-$\tau$}{}, a \texttt{Tensorflow}-based fully differentiable atmospheric forward model, based on the implementation of \texttt{TauREx3} \citep{taurex3, taurex3.1}.
    \item For the first time, we introduced Variational Inference as a more efficient alternative to perform atmospheric retrieval. 
    \item Our framework only requires a \textit{single} data instance during training time, there is no need for a large library of spectra for pre-training.
    \item We show that our framework takes into account the uncertainty associated with the observation and is able to reproduce physically motivated correlations between atmospheric parameters.
    \item Our Bayesian Neural Network is capable of producing posterior distributions on par with distributions produced from sampling process. 
\end{enumerate}

\begin{figure*}[t]
\centering
\includegraphics[width=0.8\textwidth]{updated_plot.drawio.pdf}
\caption{Overview of our investigation. Left panel shows the structure of \texttt{Diff-$\tau$}{}. It takes in a set of physical parameters and outputs the corresponding theoretical spectrum. Right panel shows the schematic of our proposed \texttt{VI-retrieval}{}. We provided a simple noise distribution $\bm{\xi}$, observed data and our prior bounds as inputs to the optimisation loop (Train). The Normalising Flow-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 \texttt{Diff-$\tau$}{} to transform samples from $q(\mathbf{z})$ into spectra and compared with the observed data. Once the training is completed, the trained NF is able to provide a $q^*(\mathbf{z})$ that best approximate $p(\mathbf{z}|\mathbf{x})$}.
\label{fig:schematic}
\end{figure*}



\section{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 \texttt{NS-retrieval}{} and our proposed approach as \texttt{VI-retrieval}{}. Our approach involves three core components - a differentiable physical model, Variational Inference and Normalising Flow. Here we will provide a top-level overview of how the three components interact with each other, also see \Cref{fig:schematic} for a schematic overview of the \texttt{VI-retrieval}{}. 

Instead of relying on sampling to map the actual, underlying posterior distribution (as one normally does with e.g. MCMC based approaches), \texttt{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 \texttt{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 wide-spread use of variational inference in the field of exoplanetary atmospheres. These are: (1) \textbf{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 in-differentiable (2) \textbf{Surrogate distributions:} the chosen distribution (usually a multinomal Gaussian distribution) is often too simplistic to represent the actual, underlying posterior distribution.

To complete the ``circuit" we built \texttt{Diff-$\tau$}{} with \texttt{Tensorflow} libraries, 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 multinomal 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 demonstrate how these concepts are linked to each other in practise by providing retrieval examples in two different scenarios. 



\subsection{Variational Inference}
The mathematical theory of variation inference (VI) and its application in the field of machine learning have been extensively discussed in \citep{Blei_VI}. There are on-going efforts to investigate the statistical implication of using VI for parameter estimation \citep[e.g.][]{Pati_pmlr,cherief_VI_param}. 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 estimate 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, VI approximates 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}$. We can measure the statistical distance, e.g. the Kullback-Leibler  divergence ($D_{\text{KL}}$), between $q(\mathbf{z})$ and $p(\mathbf{z}|\mathbf{x})$, and find an optimal distribution $q^*(\mathbf{z})$ that minimises the divergence from $p(\mathbf{z}|\mathbf{x})$, i.e.

\begin{equation}
    q^*(\mathbf{z}) = \argmin_{q(\mathbf{z}) \in \mathbb{Q}} \: D_{\text{KL}}[p(\mathbf{z}|\mathbf{x})||q(\mathbf{z})]
    \label{eqn:KL_obj}
\end{equation}

\noindent However, this optimisation procedure is impossible as we do \emph{not} have any knowledge of $p(\mathbf{z}|\mathbf{x})$ to begin with. To negate the dependence on the unknown true posterior distribtion, we can instead optimise the Evidence Lower BOund (ELBO), 
\begin{equation}
    \text{ELBO} = \mathbb{E}[\text{log}\,p(\mathbf{x}|\mathbf{z})] - D_{\text{KL}}[q(\mathbf{z})||p(\mathbf{z})]
    \label{eqn:elbo}
\end{equation}
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 \Cref{eqn:KL_obj} in Appendix \ref{sec:app}.

Unlike \Cref{eqn:KL_obj}, 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 \citep{vae,IAF}. 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 \texttt{Tensorflow} \citep[][]{tensorflow}, \texttt{PyTorch} \citep[][]{pytorch} and \texttt{JAX} \citep[][]{jax2018github} provide easy access to model training and evaluation. Our implementation will be solely based on the \texttt{Tensorflow} framework,  other deep learning frameworks can also be used \citep{Kawahara_2021}. 

\begin{figure*}
\includegraphics[width=1\textwidth]{FM_compare.pdf}
\caption{Top panel: comparing outputs from \texttt{Diff-$\tau$}{} in native (pink) and binned (Ariel) resolution (orange) and \texttt{TauREx3}{} binned to the same resolution (blue 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).  }
\label{fig:agreement}
\end{figure*}


\subsection{Differentiable forward model, \texttt{Diff-$\tau$}{}}
\texttt{Diff-$\tau$}{} is an atmospheric forward model built entirely within the \texttt{Tensorflow} framework \citep[][]{tensorflow}. We followed the forward model formulation as specified in \cite{taurex3} to construct \texttt{Diff-$\tau$}{}, with only very minor modifications to comply with the \texttt{Tensorflow} framework. As the code is largely based on \texttt{TauREx3}, it provides excellent agreement between the two forward models, see \Cref{fig:agreement} for an empirical comparison between the two. The advantage of having it entirely in a deep learning framework is to leverage the built-in automatic differentiation functionality \citep[][]{Gunes_2015}, 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 for each can be a time consuming and non-trivial task. As discussed above, having a differentiable models opens up many interesting avenues to explore, in this investigation we aim to showcase how differentiable models enable us to perform Variational Inference to approximate Bayesian posterior distribution.

While the \texttt{Autograd} library has lessened the overhead to translate the forward model from one framework to another, translating between the two frameworks remains a non-trivial task. Section \ref{sec:comments} will discuss the authors' comments on implementing physical models within the Tensorflow deep learning frameworks. 

\subsection{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 normality in $q^*(\mathbf{z})$ in \Cref{eqn:elbo} 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~\citep[][]{NF_foundation,NF_VI,Kobyzev_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. 
\begin{equation}
    p_{\mathbf{y}}(\mathbf{y}) = p_{\bm{\xi}}(f(\mathbf{y}))|\text{det D} (f(\mathbf{y})|
\end{equation}
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 density (the generative direction). On the other hand, the function $f(.)$ moves in the opposite direction, transforming a complex density 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$ is arbitrarily complex ~\citep[][]{Bogachev_2005,Med_proof_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. 
\begin{equation}
    g = g_N \circ g_{N-1} ... \circ g_{1} 
\end{equation}
where $g_{1} ... g_N $ is a set of $N$ bijective function. Similarly, $g$ has an inverse,
\begin{equation}
    f = f_1 \circ f_{2} ... \circ f_{N} 
\end{equation}
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.
\begin{equation}
    \text{D}f(\mathbf{y}) = \prod_i^N \text{D}f_i(\mathbf{\bm{\varphi}_i})
\end{equation}
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}$.

However, these intermediate functions must by definition be diffeomorphic, meaning they must be bijective and differentiable (including their inverses). There has been a significant focus on constructing non-linear bijectors that conform with these restrictions but remain sufficiently expressive and computationally efficient even in high-dimensional problems (such as images), notable bijector architectures include MADE \citep[][]{MADE}, Masked Autoregressive Flow \citep[][]{MAF}, NICE \citep[][]{NICE}, RealNVP \citep[][]{realnvp}, FFJORD \citep[][]{ffjord}, Glow \citep[][]{kingma_glow_2018} and NSF \citep[][]{NSF}. Normalising Flow has met with critical success in a wide range of applications, including audio \citep[][]{oord2018parallel}, text \citep[][]{jin-unsupervised} and image generation \citep[][]{kingma_glow_2018,ffjord}. 

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 it as our surrogate distribution $\mathbf{z} \sim q(\mathbf{z})$.

\section{Implementation}

\subsection{ELBO Formulation}
We will begin by describing our implementation to the general ELBO formulation in \Cref{eqn:elbo}. The ELBO consists of the log-likelihood term and the prior  term. The latter term can be computed analytically, i.e.
\begin{equation}
    D_{\text{KL}}[q(\mathbf{z})||p(\mathbf{z})] = \mathbb{E}\left[\text{log}\left(\frac{q(\mathbf{z})}{p(\mathbf{z})}\right)\right]
    \label{eqn:prior_term}
\end{equation}
However, an analytical form of  $q(\mathbf{z})$ may not always be possible (i.e. a distribution crafted by Normalising Flow). Here we have approximated the KL divergence via Monte Carlo sampling. Consistent with existing literature in exoplanetary retrieval, we imposed a uniform prior on the physical parameters, but it can also be any other kind of prior. 


Additionally, we define our log-likelihood as an additive log-Gaussian probability density function, i.e. 
\begin{align}
    \mathbb{E}[\text{log} p(\mathbf{x}|\mathbf{z})] &= \mathbb{E}[\text{log}(\mathcal{N}(\mathbf{x}|\bm{\hat{\mathbf{\mu}}}, \bm{\sigma})] \\
     &= \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]
    \label{eqn:llh}
\end{align}
where $\bm{\hat{\mu}}$ = $\texttt{Diff-$\tau$}{}(\mathbf{z})$ is a atmospheric forward model generated by \texttt{Diff-$\tau$}{} and $\bm{\sigma}$ is the observed uncertainty. This formulation is the same as the likelihood formulation currently employed by many retrieval frameworks.

Instead of optimising the ELBO objective in its actual formulation, we adopted the weight annealing approach from \cite{Sun2022} to optimise a modified ELBO objective, i.e. 
\begin{equation}
    \text{loss} = \frac{1}{\beta}\mathbb{E}[\text{log}\,p(\mathbf{x}|\mathbf{z})] - D_{\text{KL}}[q(\mathbf{z})||p(\mathbf{z})]
    \label{eqn:mod_elbo}
\end{equation}
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. During the training stage, the objective function will slowly converge back to the original ELBO formulation. 

\subsection{Training Procedure}


At training time, our Flow-based neural network (the generator, for details on the architecture please refer to Section \ref{sec:NF_setup}) is tasked to transform a simple noise distribution (we used Uniform noise here) $\bm{\xi}$ into a latent surrogate distribution $q(\mathbf{z})$ that can best approximate $p(\mathbf{z}|\mathbf{x})$. To calculate the ELBO loss, samples are drawn and transformed via \texttt{Diff-$\tau$}{} into synthetic spectra. These spectra are then compared with the observed data to estimate the ELBO values. The entire process is repeated until the optimisation procedure is completed or terminated. 

The role of \texttt{Diff-$\tau$}{} can be seen simply as a deterministic transformation, i.e. from physical parameters to spectra. However, 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 other words, the outputs and the correlations are no longer black boxes but are physically motivated. 

Once trained, the generator is now capable of transforming simple distributions in to the desired posterior distribution. 



\section{Application}
\label{sec:application}
In this section we will perform atmospheric retrieval with \texttt{NS-retrieval}{} and \texttt{VI-retrieval}{}. We will demonstrate the technique through two cases:
\begin{enumerate}
    \item Real Hubble/WFC3 (HST/WFC3) observation of the hot-Jupiter HD209458\,b.
    \item Simulated Ariel Tier-2 observation of the same planet.
    
\end{enumerate}
Note that, while we have demonstrate our results using HD209458\,b, \texttt{VI-retrieval}{} is not confined to any planet nor planet classes.

\subsection{Normalising Flow model setup}
\label{sec:NF_setup}
We implemented Inverse Autoregressive Flow \citep[IAF,][]{IAF} as a default bijector unit in our Normalising Flow based neural network. To perform the transformation, we chained 10 bijector units together, each controlled by a 2-layer densely connected neural network with 64 hidden units and ReLU activation \citep[][]{relu} in each layer. To stabilise the network, we followed \cite{IAF} and added a Batch Normalisation layer \citep[][]{BatchNorm} after each bijector unit. We used Adam as our optimiser with the learning rate kept at 0.001, the rest of the settings are kept as default. We used a multi-dimensional uniform distribution as our ``seed" distribution. At each iteration, we will sample 5 times from the surrogate distribution and compute the average (modified) ELBO objective, \Cref{eqn:mod_elbo}. In both cases we set $\beta_0$ and $\tau$ as 500 and 200 respectively.

\begin{table*}[]
\centering
\begin{tabular}{ccccccc}
\toprule
Parameters          & Ground Truth*  &     Case I           & Case II           & Priors  & Scale & Reference \\ \midrule
T$_p$ (K)           & 1449          & \checkmark            &    \checkmark      &   $\mathcal{U}$(100,2500)    & Uniform & N/A \\
R$_p$ (R$_J$)         & 1.359        & \checkmark              &   \checkmark       &   $\mathcal{U}$(0.5,2.5)     & Uniform & N/A \\
log$_{10}$(H$_2$O)   & -5           & \checkmark            &  \checkmark        &   $\mathcal{U}$(-12,-2)     &  Log-Uniform & \cite{polyansky_h2o} \\
log$_{10}$(CH$_4$)   & -5           &             &    \checkmark      &   $\mathcal{U}$(-12,-2)      & Log-Uniform  & \cite{ExoMol_CH4_new} \\
log$_{10}$(CO$_2$)   & -5           &             &   \checkmark       &   $\mathcal{U}$(-12,-2)      & Log-Uniform & \cite{Yurchenko_2020}\\
log$_{10}$(NH$_3$)   & -8           &           &  \checkmark        &   $\mathcal{U}$(-12,-2)      & Log-Uniform  & \cite{ExoMol_NH3} \\
log$_{10}$(CO)      & -8           &          &  \checkmark        &      $\mathcal{U}$(-12,-2)   & Log-Uniform & \cite{Li_CO}\\
log$_{10}(\chi^{lee}_{mie}$)   & N/A          &  \checkmark              &          &   $\mathcal{U}$(-40,-4)      & Log-Uniform  & \cite{lee_scattering} \\
$q^{lee}_{mie}$   & N/A          &  \checkmark             &          &   $\mathcal{U}$(1,99)      & Uniform  & \cite{lee_scattering} \\
log$_{10}(a^{lee}_{mie}$)   & N/A          &   \checkmark            &          &   $\mathcal{U}$(-3,1)      & Log-Uniform  & \cite{lee_scattering} \\


\bottomrule
\end{tabular}
\caption{Fitted parameters in each scenario, along with their corresponding ground truth, prior bounds, scales and molecular line list references for each trace gas species. The checkmark (\checkmark) indicates whether the parameter is included in Case I and II. The ground truth is only available for Case II.  }
\label{tab:GT_and_priors}

\end{table*}

\subsection{Atmospheric model setup}
\label{sec:atmo_setup}
We have taken values from \cite{Tsiaras_2016} as our reference values on the HD209458 system. In both cases, we assume a blackbody emission for the host-star at T$_{eff}$ = 6065\,K and stellar radius R$_*$ = 1.155\,R$_{\odot}$. We also assume a primary atmosphere (dominated by H and He, ratio = 0.175) at solar abundance. We divide the atmosphere into 70 atmospheric layers over a 10$^{-5}$\,Pa and 10$^6$\,Pa pressure range and furthermore assume an iso-thermal T-P profile and an iso-abundance chemical profile. \Cref{tab:GT_and_priors} shows the fitted input parameters for each case and their corresponding ground truth, prior bounds and formalism/linelist references.


\subsection{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 \texttt{Iraclis} \citep{Tsiaras_2016, Tsiaras2019}. The detrending process is described in details in \cite{Tsiaras_2016}. The basic atmospheric setup follows Section\,\ref{sec:atmo_setup}. For this case we are assuming a hydrogen/helium dominated atmosphere, with trace water and absorption from Mie scattering clouds\footnote{We 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 formalism in \texttt{Diff-$\tau$}{}.} \citep{lee_scattering}. We ran the optimisation procedure for 1500 epochs. 



\begin{figure*}
    \centering
    \includegraphics[width=0.8\textwidth]{figures/wfc3_combined.pdf}
    \caption{Posterior distribution of HD209458\,b WFC3/G141 observation obtained from \texttt{NS-retrieval}{} (red) and \texttt{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.}
    \label{fig:hst_retrieval}
\end{figure*}

\subsection{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 \texttt{ArielRad} \citep[][]{mugnai_Arielrad} 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\,\ref{sec:atmo_setup} and include five trace gases H$_2$O, CH$_4$, CO and CO$_2$ and NH$_3$. 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 atmospheres. We ran the optimisation procedure for 3000 epochs.


\section{Results}
In both cases (see Figure~\ref{fig:hst_retrieval} and Figure~\ref{fig:ariel_retrieval} and, the \texttt{VI-retrieval}{} (in yellow) is able to converge to similar results to \texttt{NS-retrieval}{} (in red), with . Table compares the quartile values return from both retrievals. We can see a very good agreement between both methods and the median values are within 1-$\sigma$ to each other. As for the shape of the posterior distribution, the surrogate distribution is both able to reproduce peaked distributions whenever there is sufficient constrain 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, aligning with our uniform prior.

The surrogate distribution also displayed similar covariances between the free parameters, it is most notable in Case II, when the high S/N observation of Ariel simulated observations allows better constraints of the free model parameters. 

There are also notable differences between the two retrievals, for instance \texttt{VI-retrieval}{} is less adept at capturing ``cliff" like distribution (e.g. log$_{10}(\chi^{lee}_{mie}$)), when there is a sharp transition from presence to absence of a physical phenomenon. 

Our objective 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 worsen 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 (\Cref{eqn:mod_elbo}) aims to alleviate the imbalance between the two terms by lowering the contribution from the log-likelihood term during the start of the training. However, the underdispersed situation will likely to return as the formulation slowly converge back to the ELBO formulation, as seen from both cases. An alternative remedy is to increase permanently the contribution from the prior term, which will inevitably increase the estimated uncertainty. 


\begin{figure*}
    \centering
    \includegraphics[width=0.8\textwidth]{figures/ariel.pdf}
    \caption{Posterior distribution of HD209458\,b in Ariel Tier2 resolution obtained via \texttt{NS-retrieval}{} (red) and \texttt{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}
    \label{fig:ariel_retrieval}
\end{figure*}


\section{Discussion}
\subsection{A Flexible and Interpretable Framework}
Despite numerous recent developments in applying ML techniques to infer atmospheric properties from spectroscopic data, the applicability of modern deep learning models in atmospheric retrieval remains limited. This is mainly due to two limiting factors: 1) their generalisability and training requirements and 2) their interpretability. 




\begin{table}[]
\centering
\begin{tabular}{rcc}
\toprule
Ncalls                & \texttt{NS-retrieval}{} & \texttt{VI-retrieval}{} \\ \midrule
WFC3/G141       & 105622      &  7500    \\
Ariel & 65414   & 15000 \\
\bottomrule
\end{tabular}
\label{tab:Ncalls}
\caption{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 \texttt{VI-retrieval}{} is significantly lower than \texttt{NS-retrieval}{}.}
\end{table}


Models are currently trained on a large data sets, encapsulating a fixed set of model assumptions and pre-defined input grids. Any changes to the spectral range, resolution or model assumptions will effectively hinder the model's performance (assuming the change is small) or in most cases, require re-training ~\citep{Marquez_rf,Zingales_2018,Cobb_2019,himes2020,Yip_sensitivity,Haldemann2022,Ardevol_CNN}. The method presented here presents an alternative approach to train deep learning model. Instead of generating a big data set for model training, our network is specific (or in other words, overfitted) to the observed data, and similar to Markovian sampling algorithms, it is not generalisable and must be re-run for any changes in observed data and/or model assumptions. However, using an optimisation algorithm (i.e. VI) over a sampling one drastically reduces the number of forward model computations required and results in significant run-time savings when forward models are computationally expensive. Table \ref{tab:Ncalls} shows a comparison of the number of forward model calls between \texttt{NS-} and \texttt{VI-retrieval}{} on the same atmospheric model (as Case I) in two different spectral resolutions. By dropping the goal of training a generalisable model, we forego the need to create vast training sets and expensive pre-training of our model. 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). On the other hand, generalisable models must be re-trained for any changes in observed data and/or model assumptions. 

The second major limitation of ML and in particular deep models is their inherently un-interpretable architecture. For instance, conventional Variational Autoencoders~\citep[made up of neural network based encoder and decoder,][]{vae}, do not produce naturally interpretable latent variables. Our implementation, on the other hand, naturally produce interpretable latent variables by using a physical, deterministic forward model as our ``decoder". Here we constrain the neural network to produce physically viable solutions, in other words, \emph{the network is constrained explicitly by the physics}. 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 (\texttt{NS-retrieval}{}, red contour). 

\subsection{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 parameters (parameter values used to generate the forward model). 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. 

In an ideal world, both approaches will agree with each other. However, it is not the case for inverse problems, where our observation is inherently corrupted\footnote{possible sources include limitations from instruments or other astrophysical noises} 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. The conventional approach will therefore be asking the neural network to pursue ground truth(s) that may no longer be possible to retrieve, which will cause the neural network to exhibit fictitious behaviour if it results in the lowering of the loss function values \citep[][]{Yip_sensitivity}. Our proposed objective, on the other hand, naturally takes into account these observational uncertainties and is inherently constrained by our deterministic physical model. 

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 \citep[][]{datashift} - where our training distribution is different from our test distribution. 



Our approach proposed in this paper grants us better interpretability of the results, compared to other machine learning retrival architectures, due to the inclusion of the physical forward model \texttt{Diff-$\tau$}. However it does of course not guard from the intrinsic retrieval biases induced by the atmospheric forward itself \citep[][]{Rocchetto_2016,Feng_2020, MacDonald_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.

\subsubsection{Model Selection}
\begin{table}[]
\centering
\begin{tabular}{rcccc}
\toprule
Model             & ELBO & Ref  & log$_{10}(\mathcal{B})$\\
\midrule
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\\
\bottomrule
\end{tabular}
\label{tab:model_selection}
\caption{Retrieved ELBO from \texttt{VI-retrieval}{} and Bayesian (log-) Evidence from \texttt{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. }
\end{table}
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 \texttt{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. 
\begin{equation}
    \mathcal{B}_{10} = \frac{P(\textbf{x}|\mathcal{M}_1)P(\mathcal{M}_1)}{P(\textbf{x}|\mathcal{M}_0)P(\mathcal{M}_0)}
\label{eqn:bayesfactor}
\end{equation}
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 \texttt{VI-retrieval}{} and \texttt{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 \ref{tab:model_selection} compares the corresponding ELBO from \texttt{VI-retrieval}{} and Bayesian Evidence from \texttt{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 \ref{tab:model_selection} allows us to differentiate different models. Following Jeffrey's guideline scale \citep[][]{jeffreys_theory_1998, Hobson2002, Padilla_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 \cite{Changeat_w43}.

\subsection{A word on differentiable frameworks}
\label{sec:comments}
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 \texttt{Tensorflow}, \texttt{PyTorch} and \texttt{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 \citep[][]{Kawahara_2021} 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. 

\begin{enumerate}
    \item \textbf{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 \emph{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.
    \item \textbf{Differentiability is not guaranteed:}  Not all operations are naturally differentiable. 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.
    \item \textbf{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 \citep[][]{hu2019difftaichi,hu2019taichi}. 
    \item \textbf{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-compatible\footnote{one good example is the switch from \texttt{Tensorflow} 1 to \texttt{Tensorflow} 2.} and may impact the long term sustainability of the developed model. 
\end{enumerate}

\section{Conclusion}

In this paper we have introduced the differentiable forward model (\texttt{Diff-$\tau$}{}) based on \texttt{TauREx3}{} and implemented in \texttt{Tensorflow}. We further introduced a Variational Inference neural network combined with a density-alternating  (Normalising Flow) approach. We showed that it is possible to compute the approximate Bayesian posterior distribution without resorting to computationally more expensive sampling based technique. 
Through our examples we have demonstrated three advantages of our framework:
\begin{enumerate}
    \item \textbf{Speed}: Our \texttt{VI-retrieval}{} only requires 10\% or less forward model calls compared to \texttt{NS-retrieval}{} to converge.
    \item \textbf{Flexibility}: Unlike many deep learning based framework, which relies on a large pre-computed library of spectra and somewhat obscure loss objectives, 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.
    \item \textbf{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. 
\end{enumerate}

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.  
Our framework is \emph{not} meant to replace sampling based framework, but to \emph{complement} existing frameworks. The approximate retrieval performed by variational inference can act as a precursor before the more computationally heavy Nested Sampling retrievals are required. Its rapid nature 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. 


\begin{acknowledgments}
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. 

\end{acknowledgments}


\vspace{5mm}



\software{TauREx 3~\citep[][]{taurex3, taurex3.1}, Tensorflow2 ~\citep[][]{tensorflow2015-whitepaper}, tensorflow-probability ~\citep[][]{dillon2017tensorflow}, 
scipy~\citep[][]{2020SciPy-NMeth}, NumPy~\citep[][]{harris2020array}, Pandas ~\citep[][]{reback2020pandas},nestle~\citep[][]{barbary_2021}, corner~\citep[][]{corner}, matplotlib~\citep[][]{matplotlib}, h5py~\citep[][]{collette_python_hdf5_2014}
          }



