Parameterizing pressure-temperature profiles of exoplanet atmospheres
  with neural networks
Atmospheric retrievals (AR) of exoplanets typically rely on a combination of
a Bayesian inference technique and a forward simulator to estimate atmospheric
properties from an observed spectrum. A key component in simulating spectra is
the pressure-temperature (PT) profile, which describes the thermal structure of
the atmosphere. Current AR pipelines commonly use ad hoc fitting functions here
that limit the retrieved PT profiles to simple approximations, but still use a
relatively large number of parameters. In this work, we introduce a
conceptually new, data-driven parameterization scheme for physically consistent
PT profiles that does not require explicit assumptions about the functional
form of the PT profiles and uses fewer parameters than existing methods. Our
approach consists of a latent variable model (based on a neural network) that
learns a distribution over functions (PT profiles). Each profile is represented
by a low-dimensional vector that can be used to condition a decoder network
that maps $P$ to $T$. When training and evaluating our method on two publicly
available datasets of self-consistent PT profiles, we find that our method
achieves, on average, better fit quality than existing baseline methods,
despite using fewer parameters. In an AR based on existing literature, our
model (using two parameters) produces a tighter, more accurate posterior for
the PT profile than the five-parameter polynomial baseline, while also speeding
up the retrieval by more than a factor of three. By providing parametric access
to physically consistent PT profiles, and by reducing the number of parameters
required to describe a PT profile (thereby reducing computational cost or
freeing resources for additional parameters of interest), our method can help
improve AR and thus our understanding of exoplanet atmospheres and their
habitability.
Timothy D. Gebhard, Daniel Angerhausen, Björn S. Konrad, Eleonora Alei, Sascha P. Quanz, Bernhard Schölkopf
# Parameterizing pressure-temperature profiles

###### Abstract

Context:Atmospheric retrievals (AR) of exoplanets typically rely on a combination of a Bayesian inference technique and a forward simulator to estimate atmospheric properties from an observed spectrum. A key component in simulating spectra is the pressure-temperature (PT) profile, which describes the thermal structure of the atmosphere. Current AR pipelines commonly use ad hoc fitting functions here that limit the retrieved PT profiles to simple approximations, but still use a relatively large number of parameters.

Aims:In this work, we introduce a conceptually new, data-driven parameterization scheme for physically consistent PT profiles that does not require explicit assumptions about the functional form of the PT profiles and uses fewer parameters than existing methods.

Methods:Our approach consists of a latent variable model (based on a neural network) that learns a distribution over functions (PT profiles). Each profile is represented by a low-dimensional vector that can be used to condition a decoder network that maps \(P\) to \(T\).

Results:When training and evaluating our method on two publicly available datasets of self-consistent PT profiles, we find that our method achieves, on average, better fit quality than existing baseline methods, despite using fewer parameters. In an AR based on existing literature, our model (using two parameters) produces a tighter, more accurate posterior for the PT profile than the five-parameter polynomial baseline, while also speeding up the retrieval by more than a factor of three.

Conclusions:By providing parametric access to physically consistent PT profiles, and by reducing the number of parameters required to describe a PT profile (thereby reducing computational cost or freeing resources for additional parameters of interest), our method can help improve AR and thus our understanding of exoplanet atmospheres and their habitability.

## 1 Introduction

With now over 5000 confirmed planet detections outside the solar system (see, e.g., Christiansen, 2022), the exoplanet science community is increasingly expanding from detecting new planets to also characterizing them. One important tool for this are atmospheric retrievals (AR), that is, "the inference of atmospheric properties of an exoplanet given an observed spectrum" (Madhusudhan, 2018). The properties of interest here include, for example, the chemical composition of the atmosphere (i.e., abundances of different chemical species) or the presence of clouds or haze. Deriving these properties from a spectrum is a classic example of an inverse problem. As with many such problems, the standard approach is to combine Bayesian inference methods, such as nested sampling (Skilling, 2006), with a simulator for the "forward" direction--in our case, turning a set of atmospheric parameters \(\mathbf{\theta}\) into the corresponding spectrum. Given the spectrum \(y\) of a planet of interest, one can then estimate a posterior distribution \(p(\mathbf{\theta}\,|\,y)\) by iteratively drawing a sample of parameter values from a prior \(\pi(\mathbf{\theta})\), passing them to a simulator, and comparing the simulated spectrum \(\hat{y}\) to the observed one through a likelihood function \(L\) to guide the next parameter sample that is drawn. In practice, the simulator is often deterministic, in which case \(L\) is essentially equivalent to the noise model: For example, optimizing \(\mathbf{\theta}\) to minimize \(\|y-\hat{y}(\mathbf{\theta})\|^{2}\) implies the assumption that the noise in \(y\) comes from an uncorrelated Gaussian.

Traditionally, due to computational limitations, much of the existing work on atmospheric retrievals has made the simplifying assumption of treating atmospheres as one-dimensional. More recent studies, however, have also begun to explore 2D and 3D approaches (see, e.g., Chubb and Min, 2022; Nixon and Madhusudhan, 2022; Zingales et al., 2022).

Pressure-temperature profilesOne of the key factors that determines the observable spectrum of a planet's atmosphere is its thermal structure, that is, the temperature of each atmospheric layer. The thermal structure is described by the pressure-temperature profile, or PT profile, which is a function \(f:\mathbb{R}^{+}\rightarrow\mathbb{R}^{+}\) that maps an atmospheric pressure to its corresponding temperature. In reality, of course, this function also depends on the location (e.g., in the case of the Earth, the air over the equator is warmer than over the poles), and such variations can be accounted for when using general circulation models (GCMs). However, as mentioned above, it is common to limit oneself to one-dimensional approximations. The implications of such an approximation are discussed, for example, in Blecic et al. (2017).

There are, in principle, two different ways to determine the PT profile during an AR: One may either solve a set of (simpli
fied) radiative transfer equations, or approximate the PT profile through an ad hoc fitting function (see, e.g., Seager, 2010; Line et al., 2013; Heng, 2017). In the latter case, which is the one that we consider in this work, one has to choose a parameterization for the PT profile, that is, one needs to assume a parametric family of functions \(f_{\mathbf{z}}:P\to T\) which map pressure values onto corresponding temperature values and have tunable parameters \(\mathbf{z}\in\mathbb{R}^{d}\) that control how exactly this mapping looks like (see Figure 1 for an illustration). These parameters \(\mathbf{z}\) are a subset of the full set of parameters \(\mathbf{\theta}\) for which the retrieval procedure estimates a posterior distribution (see above). Throughout this work, we commonly refer to \(\mathbf{z}\) as a (latent) representation of a PT profile. The choice of \(f_{\mathbf{z}}\) will, in general, be determined by a trade-off between different (and possibly conflicting) requirements and desiderata, including the following:

* For every profile that we can expect to find during the retrieval, there should exist a \(\mathbf{z}\) such that \(f_{\mathbf{z}}\) is a good approximation of this profile (surjectivity).
* For every value of \(\mathbf{z}\) from some given domain (e.g., a subset of \(\mathbb{R}^{d}\)), \(f_{\mathbf{z}}\) is a valid PT profile. Together with surjectivity, this implies that the image of \(\mathbf{z}\mapsto f_{\mathbf{z}}\) is equal to the set of physically sensible (and relevant) PT profiles.
* The mapping \(\mathbf{z}\mapsto f_{\mathbf{z}}\) should be smooth in the sense that small changes in \(\mathbf{z}\) should only result in small changes in \(f_{\mathbf{z}}\) (i.e., continuity); otherwise, the retrieval loop may not converge.
* The dimensionality \(d=\dim(\mathbf{z})\) should be as small as possible, since methods such as nested sampling often do not scale favorably in the number of retrieval parameters.
* We should be able to define a simple prior \(p(\mathbf{z})\) for \(\mathbf{z}\), and for \(\mathbf{z}\sim p(\mathbf{z})\), the induced distribution over functions \(f_{\mathbf{z},\mathbf{z}-p(\mathbf{z})}\) should constitute an appropriate prior distribution over the PT profiles of the population of planets that we are considering.
* Ideally, the dimensions of \(\mathbf{z}\) are interpretable as physical quantities, or correspond to physical processes (e.g., "\(z_{1}\) is the mean temperature," or "\(z_{2}\) controls the presence of an inversion").

Of course, as suggested above, no choice of parameterization will perfectly satisfy all requirements simultaneously, and some sort of compromise is required.

Related workBarstow and Heng (2020), who have identified the parameterization of PT profiles as one of the "outstanding challenges of exoplanet atmospheric retrievals," find the approaches of Madhusudhan and Seager (2009) and Guillot (2010) to be the most popular among the atmospheric retrieval community. These parameterizations, as well as the ones proposed by Hubeny et al. (2003), Hansen (2008), Heng et al. (2011), and Robinson and Catling (2012), all consist of (semi)-analytic expressions that seek to describe the relationship between pressure and temperature mainly in terms of physically interpretable quantities, such as the opacity or optical depth. Simpler (yet less physically motivated) approaches include the use of splines (e.g., Zhang et al., 2021) as well as low-order polynomials (e.g., Konrad et al., 2022). Finally, a data-driven approach was presented by Schreier et al. (2020), who compute a singular value decomposition of a representative set of PT profiles, and parameterize new profiles as a combination of the first \(k\) singular vectors. All these approaches typically use four to six parameters to describe a PT profile, which is often a substantial fraction of the total number of retrieved parameters. To the best of our knowledge, no studies to date have used machine learning (especially neural networks) to parameterize pressure-temperature profiles of exoplanet atmospheres.

ContributionsIn this work, we explore a conceptually new approach to parameterizing PT profiles that is based on the idea of using neural networks to learn efficient (i.e., low-dimensional) representations of PT profiles, as well as the corresponding decoding mechanisms, from simulated data.1 Using two different datasets of PT profiles obtained with full climate models, we show that our approach can achieve better fit quality with fewer parameters than our two baseline methods, while still easily integrating into an existing atmospheric retrieval framework.

Footnote 1: We presented early versions of this work at AbSciCon 2022 and the “ML and the Physical Sciences” workshop at NeurIPS 2022.

## 2 Method

Our goal is the following: We want to learn a model (in our case: a neural network) that takes two inputs: (1) a pressure value \(p\), and (2) a vector \(\mathbf{z}\in\mathbb{R}^{d}\), also called "representation", which contains a highly compressed description of a PT profile. The output of the model is the temperature \(t\) at pressure \(p\), for the profile described by \(\mathbf{z}\). During an atmospheric retrieval, the Bayesian inference routine can then propose values of \(\mathbf{z}\), which can be converted to a PT profile by conditioning the network on the proposed \(\mathbf{z}\) (that is, fixing the input \(\mathbf{z}\)) and evaluating it at different pressures \(p_{i}\) to get the corresponding temperatures \(t_{i}\). Thus, we will infer a posterior distribution over \(\mathbf{z}\), which corresponds directly to a posterior over PT profiles.

To learn such a network, we assume that we have a dataset \(\mathcal{D}\) of PT profiles which were generated by a full climate model or radiative transfer code to ensure they are physically consistent. These models typically simulate an atmosphere in a layer-by-layer fashion, meaning that each PT profile is given in the form

Figure 1: Schematic illustration of the problem setting considered in this paper: “Parameterizing PT profiles” means finding a way to represent PT profiles as vectors \(\mathbf{z}\in\mathbb{R}^{d}\), together with a decoding mechanism \(f_{\mathbf{z}}:P\to T\) that converts these \(\mathbf{z}\)-vectors back into a function that maps pressure values onto temperature values. Our goal in this work is to learn both the representations \(\mathbf{z}\) and the mechanism \(f_{\mathbf{z}}\) from data obtained with simulations using full climate models.
of a set support points, \(\{(p_{i},t_{i})\}_{i=1,\ldots,N}\), where \(N\) is the number of layers, and \(p_{i}\) and \(t_{i}\) are the pressure and temperature in the \(i\)-th layer. Importantly, our method does not require that \(p_{i}\) be the same for all profiles in \(\mathcal{D}\): Different profiles can use different pressure grids, and this is indeed the case for the datasets we work with. Moreover, in principle, even the number of layers \(N\) can vary between profiles. (We do not explicitly consider this case here, but it implies only minor technical complications at training time.) This flexibility also makes it easy, for example, to combine PT profiles from different models or pre-existing atmosphere grids into a single training dataset.

The approach that we propose in the following is inspired by the idea of a (conditional) neural process (NP; Garnelo et al. 2018, 2018). A neural process is a type of machine learning model that learns a distribution over functions and combines the advantages of Gaussian processes with the flexibility and scalability of neural networks. In astrophysics, the usage of (conditional) neural processes is still relatively rare; examples include Park & Choi (2021), Covorici-Hajdinjak et al. (2021), and Jankov et al. (2022), all of which use conditional NPs for interpolation.

Our method works as follows (see also Figure 2). We employ two neural networks: an encoder \(E\) and a decoder \(D\). The encoder is only required during training and will not be used during an AR. At training time, we take a PT profile--consisting of a vector \(\mathbf{p}=(p_{1},...,p_{N})\) of pressure values and a vector \(\mathbf{t}=(t_{1},...,t_{N})\) of corresponding temperatures--from our training data set and pass it to \(E\), which outputs a representation \(\mathbf{z}\in\mathbb{R}^{d}\) of the profile:

\[\mathbf{z}=E(\mathbf{p},\mathbf{t})\,. \tag{1}\]

The dimensionality \(d\) of \(\mathbf{z}\) can be chosen freely (but changing \(d\) will require re-training \(E\) and \(D\)). Generally, we want to keep \(d\) as low as possible (e.g., \(d=2\)), but of course, higher values also imply more expressive or informative representations.

In the next step, the representation \(\mathbf{z}\) is then used to condition the decoder network \(D\). Once conditioned on \(\mathbf{z}\), the decoder network \(k\) is a PT profile, that is, a function that takes in a single value (a pressure) and outputs a single value (a temperature). We then pass all pressure values \(p_{i}\) of our input profile through \(D\) to get the respective predicted temperature values \(\hat{t}_{i}\):

\[\hat{t}_{i}=D(p_{i}\,|\,\mathbf{z})\,, \tag{2}\]

where we use \(D(\cdot\,|\,\mathbf{z}):\mathbb{R}^{+}\to\mathbb{R}^{+}\) to denote the function that is given by conditioning \(D\) on \(\mathbf{z}\) (i.e., fixing \(\mathbf{z}\) as an input). We train the encoder and decoder jointly (i.e., we update the weights of the neural networks) to minimize the difference between the true and predicted temperatures. More specifically, we minimize the following mean squared error (MSE) reconstruction loss, which is a standard choice in machine learning for regression problems:

\[\mathcal{L}_{\text{rec}}(\mathbf{t},\hat{\mathbf{t}})=\frac{1}{N}\sum_{i=1}^{N}(t_{i }-\hat{t}_{i})^{2}\,. \tag{3}\]

In practice, we do this not only for a single PT profile, but minimize the average of \(\mathcal{L}_{\text{rec}}\) over our entire training dataset in a batched fashion. We also note that this choice of reconstruction loss assigns the same weight to all \(N\) atmospheric layers, and refer to Section 5.4 for a brief discussion of possible alternatives.

To ensure that we can define a simple prior for \(\mathbf{z}\) during a retrieval, we furthermore want to constrain the distribution of the \(\mathbf{z}\)'s produced by the encoder. There are various potential approaches for this (see Appendix A). For this work, we adopt an idea from Zhao et al. (2017) and add a second term to the loss function that encourages \(\mathbf{z}\) to follow a standard Gaussian:

\[\mathcal{L}_{\text{MMD}}=\text{MMD}^{2}\left(\{\mathbf{z}_{i}\}_{i=1,\ldots,b}, \{\mathbf{s}_{i}\}_{i=1,\ldots,b}\right)\,, \tag{4}\]

where:

\[\mathbf{s}_{i}\sim\mathcal{N}_{d}(0,1)\,,\,\,\,i=1,\ldots,b\,.\]

Here, \(\mathcal{N}_{d}(0,1)\) denotes a \(d\)-dimensional standard Gaussian distribution (i.e., mean zero and identity matrix as covariance matrix). Furthermore, \(b\) is the batch size of our training, and MMD stands for Maximum Mean Discrepancy (Borgwardt et al. 2006; Gretton et al. 2012), which is a kernel-based metric that measures whether two samples come from the same distribution (see Appendix A.2 for more details). By minimizing this MMD term, we encourage the model to produce values of \(\mathbf{z}\) that follow the same distribution as the \(\mathbf{s}_{i}\), that is, a \(d\)-dimensional standard Gaussian.

Experimentally, we found that despite this MMD term, there are sometimes a few profiles that are mapped to a \(\mathbf{z}\) very far from the origin, which is undesirable if we want to define a simple prior for \(\mathbf{z}\) during a retrieval. To prevent this problem, we introduce a third loss term with a softplus barrier function on the norm of \(\mathbf{z}\):

\[\mathcal{L}_{\text{norm}}(\mathbf{z})=\nicefrac{{1}}{{k}}\cdot\log\left(1+\exp\{ k\cdot\left(\|\mathbf{z}\|_{2}-\tau\right)\right)}\,. \tag{5}\]

This is a smooth version of the ReLU\((x):=\max(0,x)\) function, where \(k\) controls the amount of smoothing: Larger values of \(k\) increase the similarity to ReLU. For this work, we have fixed \(k\) at 100 (ad hoc choice), after observing that using a standard ReLU function destabilized the training for us. The parameter \(\tau\) defines the threshold on \(\|\mathbf{z}\|_{2}\) and was set to \(\tau=3.5\) (ad hoc choice).

Figure 2: Schematic illustration of our model: During training, the encoder network \(E\) maps a PT profile (consisting of a vector \(\mathbf{p}=(p_{1},...,p_{N})\) of pressure values and a vector \(\mathbf{t}=(t_{1},...,t_{N})\) of corresponding temperature values) onto a latent representation \(\mathbf{z}\in\mathbb{R}^{d}\). This latent code is then used to condition a decoder network \(D\), and \(D(\cdot\,|\,\mathbf{z})\) is evaluated on each \(p_{i}\) to get the corresponding predicted temperatures \(\hat{t}_{i}\). During an atmospheric retrieval, \(\mathbf{z}\) is proposed by the Bayesian inference method (e.g., nested sampling).
The total loss that we minimize is then just a weighted sum (with \(\beta,\gamma\in\mathbb{R}^{+}\)) of our three loss terms:

\[\mathcal{L}=\mathcal{L}_{\text{rec}}+\beta\cdot\mathcal{L}_{\text{MMD}}+\gamma \cdot\mathcal{L}_{\text{norm}}\,. \tag{6}\]

The hyperparameters \(\beta\) and \(\gamma\) control the tradeoff between the different loss terms: (1) Increasing \(\beta\) corresponds to placing a stronger prior on the distribution of \(\mathbf{z}\) in the latent space, usually at the cost of a higher reconstruction error. We have tried different values of \(\beta\) and found that the performance is relatively similar over a range of about 0.1 to 1. Too high values of \(\beta\) can lead to model collapse, where the prior on \(\mathbf{z}\) becomes too strong and prevents the encoder from learning meaningful representations. For this work, we set \(\beta=0.1\). This was an ad hoc choice and not the result of systematic optimization, which would likely yield a different value for each dataset and value of \(\dim(\mathbf{z})\). (2) The value of \(\gamma\) should be chosen such that \(\gamma\cdot\mathcal{L}_{\text{norm}}\) dominates the other loss terms when \(\|\mathbf{z}\|_{2}>t\). We set \(\gamma=10\) for this work (again without optimization).

To use our trained models during atmospheric retrieval, we provide a convenient Python wrapper that allows the model to be loaded as a normal function, taking as input a value for \(\mathbf{z}\) and a numpy array of pressure values, and returning an array of corresponding temperatures. We note again that neither the lengths of the pressure grids nor the exact values need to match the pressure grid(s) used during training. (However, the range of pressure values should be approximately the same; otherwise, evaluating \(D(\cdot\,|\mathbf{z})\) at pressure values it has never encountered during training may result in nonphysical outputs).

## 3 Datasets

We validate our proposed approach using two different, publicly available datasets of PT profiles: one for terrestrial planets, and one for hot Jupiters. A visual comparison of the two datasets and the respective \(P\)-\(T\) space that they cover is shown in Figure 3.

### PvATMOS

The PvATMOS dataset was generated by Chopra et al. (2023) and consists of 124 314 stable atmospheres of Earth-like planets around solar-type stars generated with Atmos, a 1D coupled photochemistry-climate model (Arney et al., 2016; Meadows et al., 2018). The atmospheres in this dataset were simulated using \(N=101\) layers covering an altitude range of 0-80 km. Consequently, each PT profile is given by two 101-dimensional vectors for the pressures and temperatures. In addition, about 80 other atmospheric quantities are available, such as fluxes and abundance profiles of chemical species. We manually removed one PT profile from the dataset because it was completely out-of-distribution and caused stability problems during training.

Furthermore, it has been brought to our attention that the simulations of the PvATMOS dataset suffered from a couple of issues that affect the physical correctness of the simulated PT profiles. In particular, the ClimMain.f specified a H\({}_{2}\)O \(k\)-coefficient file that assumed pressure broadening by H\({}_{2}\)O itself, which is not appropriate for a \(N_{2}\)-dominated atmosphere. Further, the model version also did not allow for convection above 40 km. Consequently, the simulated PT profiles tend to underpredict stratospheric and mesospheric temperatures compared to the expectation from a real Earth twin. Newer versions of the Atmos code are not affected by these issues and use updated H\({}_{2}\)O and CO\({}_{2}\)\(k\)-coefficients (Teal et al., 2022; Vidaurri et al., 2022). While this may limit the scientific usefulness of the PT profiles, it should not affect our ability to use them to demonstrate the applicability of our method as such, that is, to show how we can learn efficient, low-dimensional parameterizations of PT profiles from simulated data.

### Govl-2020

This dataset was published by Goyal et al. (2020) and consists of 11 293 PT profiles, corresponding to 89 hot Jupiters, each with four recirculation factors, six metallicities and six C/O ratios.2 Due to missing data, we manually removed three profiles. The data were simulated with Atmo (Amundsen et al., 2014; Tremblin et al., 2015, 2016; Drummond et al., 2016),3 a radiative-convective equilibrium model for atmospheres, using \(N=51\) atmospheric layers. In addition to the PT profiles, the dataset also includes transmission and emission spectra for each atmosphere, as well as over 250 abundance profiles of different chemical species.

Footnote 2: The authors report that around 12 % of the models did not converge; hence the dataset is smaller than the expected 12 816 combinations.

Footnote 3: We point out that Atmos and Atmo are not the same code.

## 4 Experiments and results

In this section, we describe the experiments that we have performed to test our proposed approach and show our results.

Figure 3: These plots show all the PT profiles in our two datasets. Each segment of a line—that is, the connection between the points \((p_{i},t_{i})\) and \((p_{i+1},t_{i+1})\)—is color-coded by the density of the profiles at its respective \((p,t)\) coordinate, which was obtained through a 2D KDE. The green line in the PvATMOS plot shows the one PT profile that we manually removed for being out-of-distribution. \(\vartriangle\)
### Training and evaluation procedure

We use stochastic gradient descent to train our models (i.e., one pair of neural networks \(E\) and \(D\) for each combination of a dataset and a value of \(\text{dim}(\mathbf{z})\)) on a random subset of the datasets described in Section 3, and evaluate their performance on another random subset that is disjoint from the training set. For more detailed information about our network architectures, their implementation, and our training procedure, see Appendix B, or take a look at our code, which is available online.

### Reconstruction quality

As a first experiment, we study how well our trained model can approximate previously unseen PT profiles from the test set (cf. Appendix B.3), and compare the performance with two baseline methods: (1) low-order polynomials, and (2) a PCA-based approach similar to the method of Schreier et al. (2020).

Setup:For each \(\text{dim}(\mathbf{z})\in\{1,2,3,4\}\), we train three instances of our model (using three different random seeds to initialize the neural networks and control the split between training and validation) on the training set. We do not consider higher values of \(\text{dim}(\mathbf{z})\) here mainly because one of the goals of this work was to demonstrate that our method requires fewer fitting parameters than the baseline methods. Then, we take the learned decoder and, for each profile in the test set, use nested sampling to find the optimal \(\mathbf{z}\) (denoted \(\mathbf{z}^{*}\)) to reproduce that profile. In this case, "optimal" means that \(\mathbf{z}^{*}\) minimizes the mean squared error (MSE):

\[\mathbf{z}^{*}=\operatorname*{arg\,min}_{\mathbf{z}}\frac{1}{N}\cdot\sum_{i=1}^{N}(t_ {i}-D(p_{i}\,|\,\mathbf{z}))^{2}\,. \tag{7}\]

We add this additional optimization step to the evaluation to remove the effect of the encoder: Ultimately, we are only interested in how well \(D\) can approximate a given profile for some \(\mathbf{z}\)--which does not necessarily have to match the output of \(E\) for that profile perfectly. We have chosen nested sampling over gradient-based optimization--which is possible given that \(D\) is fully differentiable both with respect to \(p\) and \(\mathbf{z}\)--because the latter is generally not available during an AR, unless one uses a differentiable forward simulator (see, e.g., Diff-r from Yip et al., 2022). Nested sampling, therefore, should give us a more realistic idea of which profiles we can find during a retrieval. Our optimization procedure is based on _UltraNest_(Buchner, 2021), with 400 live points and a truncated standard normal prior, limiting each \(z_{i}\) to \([-4,4]\) (because \(\mathcal{L}_{\text{norm}}\) limits \(\|\mathbf{z}\|\) to \(\tau=3.5\)).

For the polynomial baseline, we simply fit each profile in the test set using a polynomial with degree \(n_{\text{poly}}-1\) for \(n_{\text{poly}}\in\{2,3,4,5\}\). The minimization objective of the fit is again the MSE.

Finally, for the PCA baseline, we compute a PCA on all the temperature vectors in the training set, and then fit the temperature vectors in the test set as a linear combination of the principal components (PCs), for \(n_{\text{PC}}\in\{2,3,4,5\}\), once again minimizing the MSE. We note that this PCA baseline must be

Figure 4: Examples of pressure-temperature profiles from our test set (black) together with the best approximation using our trained model (red). Columns show different values of \(\text{dim}(\mathbf{z})\in\{1,2,3,4\}\); rows show the best, median, and worst example (in terms of the RMSE). Results are shown for the PvATMOS dataset; additional plots for the Goyal-2020 are found in Figure D.1 in the appendix.
taken with a grain of salt, for two reasons. First, it requires us to project all profiles onto a single common pressure grid to ensure that the \(i\)-th entry of the temperature vectors always has the same interpretation (i.e, temperature at pressure \(p_{i}\)). We do this by linear interpolation. Consequently, the common pressure grid is determined by the intersection of the pressure ranges of all profiles, meaning that profiles may get clipped both at high and low pressures. In practice, this affects in particular the Govat-2020 data where, for example, the pressure in the deepest layer varies by more than an order of magnitude between different profiles, whereas for PvATMOS, the profiles all cover a similar pressure range. Second, combining principal components only returns vectors, not functions, meaning that if we want to evaluate a profile at an arbitrary pressure value, we again have to use interpolation.

Results: We are showing some example profiles and the best-fit reconstruction using our model in Figure 4. The main results of this experiment--the distribution of the (root) mean square error (RMSE) on the test set--are found in Figure 5. The RMSE is simply defined as the square root of the MSE defined in Equation (3), that is, the root of the mean squared difference between a given temperature vector and its best-fit reconstruction. As Figure 5 shows, our method achieves a significantly lower reconstruction error than both baseline methods, both for the PvATMOS dataset and for the Govat-2020 dataset (see Figure 4 in the appendix). We find that on the PvATMOS dataset, even the one-dimensional version of our model (i.e., each profile is parameterized as just a single number) outperforms the polynomial baseline with five fitting parameters, and is approximately on-par with the PCA solution using five components.

To account for the absolute scale of the data (i.e., Earth-like temperatures vs. hot Jupiters) and allow comparisons across the two datasets, we also compute the mean relative error (MRE):

\[\text{MRE}=\frac{1}{N}\cdot\sum_{i=1}^{N}\frac{|\hat{t}_{i}-t_{i}|}{t_{i}}\,. \tag{8}\]

We compute this quantity for every profile in the test set, and then take the median of this distribution. For PCA and our method, the medians are then also mean-averaged over the different runs. The final results are given in Figure 6. Looking at the mean relative errors, we note our errors for the Goyal-2020 dataset are systematically larger than for the PvATMOS dataset. We suspect that can be explained as the result of two factors: First, the PvATMOS training dataset is ten times bigger than the Goyal-2020 training dataset, and, second, the Goyal-2020 dataset covers a more diverse set of functions compared to the PvATMOS dataset.

Interpreting the latent representations of PT profiles in the context of the corresponding atmospheres

In this section, we take a closer look at the latent representations of PT profiles that our models learn and the extent to which they lend themselves to interpretation. We begin with Figure 7, where we present an illustration of our entire method using a model with \(\text{dim}(\mathbf{z})=2\) that we trained on the Goyal-2020 dataset. The upper part of the figure shows how the encoder \(E\) maps PT profiles--which are given by vectors \(\mathbf{p}\) and \(\mathbf{t}\), which in this example are 51-dimensional--onto 2-dimensional latent representations \(\mathbf{z}\). The three colored dots in the latent space correspond to the three PT profiles shown on top, which are examples from our test set

Figure 5: Distributions of the reconstruction error (RMSE) for a polynomial baseline, a PCA-based baseline, and our method, for different number of fitting parameters. For our model, as well as the PCA baseline, each distribution is the average of three runs using different random seeds controlling the initialization and the train/validation split. The dashed lines indicate the respective medians. Results are shown for the PvATMOS dataset; additional plots for the Goyal-2020 are found in Figure 4 in the appendix.
(i.e., profiles that were not used for training). The gray dots in the latent space show the representations of all PT profiles in our test set. Their distribution approximately follows a 2D standard Gaussian (cf. the marginal distributions), which is precisely what we would expect, as this was the reason we introduced the \(\mathcal{L}_{\text{MMO}}\) term to the loss function. The bottom part of the figure illustrates what happens if we place a regular grid on the latent space and, for each \(\mathbf{z}\) on the grid, evaluate \(D(\cdot\,|\,\mathbf{z})\) on a fixed pressure vector \(\mathbf{p}^{\prime}\), which does not have to match the \(\mathbf{p}\) of any profile in the training or test dataset. For example, \(\mathbf{p}^{\prime}\) can have a higher resolution (i.e., more atmospheric layers) than the training data.

In both parts of the figure, we observe that PT profiles with different shapes--implying different physical processes--are mapped to different parts of the latent space, and that similar PT profiles are grouped together in the latent space. Here, "similar" refers not only to the shape of the PT profiles, but also to the properties of the corresponding atmospheres. To see this, we turn to Figure 8. As mentioned above, both the PvATMOS and Goyal-2020 datasets contain not only PT profiles, but also various other properties of the corresponding atmospheres, such as concentrations or fluxes of chemical species. In Figure 8, we show several scatter plots of the latent representations of the test set (i.e., the \(\mathbf{z}\) we obtained in the first experiment), color-coded according to some of these atmospheric properties--for example, the mean concentration of CO\({}_{2}\) in the atmosphere corresponding to a given PT profile. For ease of visualization, we limit ourselves to the case \(\dim(\mathbf{z})=2\). We observe that there are clear patterns in these scatter plots that show a strong correlation between our latent representations and the properties of the corresponding atmospheres, again demonstrating that PT profiles from physically similar atmospheres are grouped together.

With this in mind, we return to the lower half of Figure 7, where we evaluate \(D\) on a grid of \(\mathbf{z}\) values. We observe that not only do different shapes of PT profiles correspond to different parts of the latent space, but the PT profiles given \(D(\cdot\,|\,\mathbf{z})\) also vary smoothly as a function of \(\mathbf{z}\). This means that the decoder not only reproduces the real PT profiles that were used to train it, but also allows smooth interpolation between them. Along the axes of the latent space, one can identify specific behaviors: For example, fixing \(z_{2}=0\) and increasing \(z_{1}\) leads to cooling in the upper atmosphere, while fixing \(z_{1}=2\) and increasing \(z_{2}\) leads to heating that extends further into the upper layers. Profiles around the center of the latent space (i.e., \(z_{1}\approx z_{2}\approx 0\)) are almost isothermal, indicating a relatively uniform temperature distribution throughout the atmosphere. A negative value of \(z_{1}\) is generally associated with the presence of a thermal inversion, and for these cases, \(z_{2}\) seems to encode the altitude where the inversion happens. For instance, for \(z_{1}=-2\) the profile shows an inversion in the mid-atmosphere (around \(\log_{10}(P)=-2\)) for \(z_{2}=2\), whereas for \(z_{2}=-2\), the inversion only happens at very high altitudes/low pressures around \(\log_{10}(P)=-6\).

Finally, we can also draw direct connections between the behavior of the PT profiles in latent space and the correlations of \(\mathbf{z}\) with some of the atmospheric parameters that we show in Figure 8. For example, a high concentration of TiO (panel 3 of Figure 8) corresponds to the hot, inverted atmospheres in the upper left of the latent space. This is consistent with expectations from the literature, where TiO is indeed one of the canonical species proposed to cause temperature inversions in hot Jupiters (e.g., Hubeny et al., 2003; Fortney et al., 2008; Piette et al., 2020).

Overall, this suggests that the latent space is to some extent interpretable: While the latent variables may not correspond directly to known physical parameters--as is the case, for example, for the analytical parameterizations of Madhusudhan and Seager (2009) and Guillot (2010)--their relationships with other variables may provide insight into the behavior of exoplanet atmospheres.

Finally, a small word of caution: While we have just shown that the latent space is _in principle_ interpretable, any _specific_ interpretation--that is, a statement along the lines of "PT profiles with hot upper layers are found in the top left corner of the latent space"--will only hold for one specific model. If we re-train that same model using a different random seed to initialize the neural network, or using a different dataset, we will end up with an equivalent but different latent space. This is because our model is, in general, not identifiable, and there exists no preferred system of coordinates for \(\mathbf{z}\): For example, we can rotate or mirror the latent space without losing information, and unless we add additional constraints to the objective function, there is no reason why the model should favor any particular orientation. However, this simply means that the value of \(\mathbf{z}\) can only be interpreted for a given trained model. It does not, in any way, affect the ability of our model to be used for atmospheric retrievals.

### Usage in an atmospheric retrieval

A major motivation for the development of our method is the possibility of using it for atmospheric retrievals. In this experiment, we show how we can easily plug our trained model into an existing atmospheric retrieval pipeline. The target PT profile we use for this experiment is taken from the PvATMOS test set. We made

Figure 6: Comparison of the median test set MRE for the different methods and datasets (lower is better). Reminder: For each profile, we compute the _mean_ relative error (over all atmospheric layers), and then aggregate by computing the _median_ over all profiles. Finally, for the PCA baseline as well as our method, the results are also _mean_-averaged over different random seeds.
## 6 Conclusion

Fig. 7: A visual illustration of the encoder \(E\), the decoder \(D\), and their connection via the latent space. See Section 4.3 in the main text for a detailed explanation. This figure is best viewed in color.
sure that it is covered by the training data (i.e., there are similar PT profiles in the training set); however, in order not to make it too easy for our method, we chose a PT profile shape that is comparatively rare in the training data (see below). In Appendix C, we consider an even more challenging setting and show two additional results for PT profiles that were generated with a different code and that are not covered by our training distribution.

Setup: For our demonstration, we add our trained model to the atmospheric retrieval routine first introduced in Konrad et al. (2022) and developed further in Alei et al. (2022) and Konrad et al. (2023), which is based on petitRADTRANS (Molliere et al., 2019), LIFEsim (Dannert et al., 2022) and PyMultiNest (Buchner et al., 2014). The goal of this experiment is to compare the PT profile parameterization capabilities of a version of our model (trained on the PyATMOS dataset using \(\dim(\mathbf{z})=2\)) with the fourth-order polynomial (i.e., five fitting parameters) used in Konrad et al. (2022) and Alei et al. (2022). We perform two cloud-free atmospheric retrievals (using both our PT model or the polynomial one) on the simulated thermal emission spectrum of an Earth-like planet orbiting a Sun-like star at a distance of \(10\,\mathrm{pc}\), assuming a photon noise signal-to-noise ratio of \(10\) at a wavelength of \(\lambda=$11.2\,\mathrm{\SIUnitSymbolMicro m}$\). The wavelength range was chosen as \([3\,\mathrm{\SIUnitSymbolMicro m},20\,\mathrm{\SIUnitSymbolMicro m}]\), and the spectral resolution was \(R=\nicefrac{{\lambda}}{{\lambda}}=200\). The number of live points for PyMultiNest was set to \(700\). In addition to the parameters required for the PT profile, we retrieve \(10\) parameters of interest; see Table 1 for an overview. This setup closely matches Konrad et al. (2022) and Alei et al. (2022).

For the retrieval with our model, we use a 2-dimensional uniform prior, \(\mathcal{U}_{2}(-4,4)\) for \(\mathbf{z}\). The reason for choosing this uniform prior instead of a Gaussian is that the shape of the ground truth PT profile is relatively under-represented in the PyATMOS dataset that we use for training. This means that our trained model, together with a Gaussian prior on \(\mathbf{z}\), assigns it a very small prior probability, making it difficult for the retrieval routine to find the correct PT profile. Using a uniform prior is an easy way to give more weight to "rare" profiles. We discuss this choice, as well as potential alternatives, in more detail Section 5.3.

Results: We show the main results of this experiment in the form of the retrieved PT profile and spectrum residuals in Figure 9. First, we find that the spectrum residuals demonstrate that both retrievals reproduce the simulated spectrum with sufficient accuracy, despite the differences in the PT parametrization. The residual's quantiles are centered on the truth, and are significantly smaller than the photon noise-level. Second, for the retrieved PT profiles, we visually find that the result obtained with our model--unlike the polynomial baseline--is in good agreement with the ground truth, except at the highest layers of the atmosphere, where it tends to underestimate the true temperature. We believe this can be explained by the fact that the upper layers of the atmosphere (i.e., low pressures) have little to no effect on the simulated exoplanet spectrum (cf. the emission contribution functions in the right panel of Section 4.4), making it hard for the retrieval routine to constrain the PT structure in the upper atmosphere. Third, the retrieved constraints for the surface pressure \(P_{0}\) and temperature \(T_{0}\), are much tighter and more accurate for our model than the polynomial baseline. Of course, this is also partly because our model represents a much narrower prior over the PT profiles that we are willing to accept as an explanation of the data (i.e., the spectrum) compared to the polynomial. This is an assumption, and in this case, it is justified by construction. In general, however, the decision of what is an appropriate prior for a given atmospheric retrieval problem is beyond the scope of the method as such. We discuss this in more detail in Section 5.2. Finally, we note that the retrieval using our model was significantly faster than the baseline: by decreasing the number of parameters required to fit the PT profile from five (polynomial) to two (our model), we reduced the total runtime by a factor of 3.2.

## 5 Discussion

In this section, we discuss the advantages and limitations of our method and suggest directions for future work.

### Advantages

Our approach for parameterizing PT profiles has multiple potential benefits compared to existing approaches. First, our model makes realistic, physically consistent PT profiles--which re

Figure 8: Examples of 2D latent spaces, where each \(\mathbf{z}\) from the test set is color-coded by a property of the corresponding atmosphere. Results are shown for the Goyal-2020 dataset; additional plots for the PyATMOS are found in Figure D.3 in the appendix.
quire a computationally expensive solution of radiative transfer equations--available as an ad hoc fitting function. It works, in principle, for all types of planets (Earth-like, gas giants, etc.) as long as we can create a suitable training dataset through simulations. In this context, "suitable" refers to, for example, the total size of the dataset and the density in parameter space (cf. the single out-of-distribution profile for PvATMOS). Our method then basically provides a tool which allows using the distribution of the PT profiles in that training set as a prior for the PT profile during a retrieval. Second, during an atmospheric retrieval, the parameterization of our model uses a simple prior for \(\mathbf{z}\) that does not require hard-to-interpret fine-tuning (see also Section 5.3). Third, as we have shown in Section 4, our method can fit realistic PT profiles with fewer parameters than the baseline and still achieve a lower fit error. Limiting the number of parameters needed to express a PT profile during retrieval can lead to faster processing times, or conversely, allow the retrieval of additional parameters when operating within a fixed computational budget. We would like to emphasize at this point that accelerating atmospheric retrievals is not only relevant for the analysis of existing observational data (e.g., from the James Webb Space Telescope), but can also be beneficial during the planning stages of future exoplanet science missions, such as LIFE (Quanz et al., 2022) or the recently announced Habitable Worlds Observatory (Clery, 2023), which require the simulation of thousands of retrievals.

Finally, looking to the future, it seems plausible that approaches that replace the entire Bayesian inference routine with machine learning may gain further traction in the context of atmospheric retrievals (see, e.g., Marquez-Neila et al., 2018; Soboczenski et al., 2018; Cobb et al., 2019; Ardevol Martinez et al., 2022; Yip et al., 2022, or Vasist et al., 2023). In this case, reducing the number of retrieval parameters may no longer be a major concern-once trained, the inference time of these models is relatively independent of the number of parameters. However, even in this case, our approach will still be useful to provide a parameterized description of the posterior over physically consistent PT profiles.

### The role of the training dataset

The practical value of our method stands and falls with the dataset that is used to train it. This should come as no surprise, but there are a few aspects to this that we will discuss in more detail here.

First, it is important to understand that the training dataset defines a distribution over the functions that our model learns to provide parameterized access to. When using our model in an atmospheric retrieval, this distribution then becomes the prior over the set of PT profiles to be considered. Compared to, for example, low-order polynomials, this is a much narrower prior: While polynomials can fit virtually anything (this is the idea of a Taylor series), our model will only produce outputs that resemble the PT profiles that it encountered during training. If this is a good fit for a given target profile, we can expect to significantly outperform the broader prior of polynomials. Conversely, if our prior does not cover the target profile, we cannot expect good results. However, this is not a limitation of our method in particular, but is true for all approaches that use a parameterization for the PT profile instead of, for example, a level-by-level radiative transfer approach. As Line et al. (2013) explain, using "[a] parameterization does force the retrieved atmospheric temperature structure to conform only to the profile shapes and physical approximations allowed by that parameterization." Thus, as always in Bayesian inference, the decision to use our method as a prior for the PT profile (rather than, e.g., low-order polynomials) represents an assumption about the space of possible explanations for the data (e.g., the observed spectrum) that one is willing to accept, and it is up to the scientist performing an atmospheric retrieval to decide whether that assumption is justified.

Finally, it is important to keep in mind that our model will learn to replicate any anomalies and imbalances present in the training data. For example, if atmospheres without inversions are overrepresented in the training data (compared to atmospheres with inversions), our learned function distribution will also place a higher probability on PT profiles without inversions, which will obviously affect the results of a retrieval.

\begin{table}
\begin{tabular}{l l l c c c} \hline \hline Parameter & Description & Prior & True value & Retrieved (polynomial) & Retrieved (our method) \\ \hline \(a_{4}\) & PT profile parameter (polynomial) & Uniform(0, 10) & — & 1.33 \({}^{+0.39}_{-0.67}\) & — \\ \(a_{3}\) & PT profile parameter (polynomial) & Uniform(0, 100) & — & 30.49 \({}^{+0.83}_{-2.73}\) & — \\ \(a_{2}\) & PT profile parameter (polynomial) & Uniform(0, 300) & — & 106.78 \({}^{+15.09}_{-87.47}\) & — \\ \(a_{1}\) & PT profile parameter (polynomial) & Uniform(0, 500) & — & 204.71 \({}^{+118.24}_{-118.24}\) & — \\ \(a_{0}\) & PT profile parameter (polynomial) & Uniform(0, 600) & — & 384.86 \({}^{+178.00}_{-500.09}\) & — \\ \(z_{1}\) & PT profile parameter (our method) & Uniform(\(-4,4\)) & — & — & \(-2.56\) \({}^{+1.69}_{-1.34}\) \\ \(z_{2}\) & PT profile parameter (our method) & Uniform(\(-4,4\)) & — & — & \(-0.04\) \({}^{+1.19}_{-1.55}\) \\ \(\log_{10}(P_{0})\) & Surface pressure (in bar) & Uniform(\(-2,2\)) & 0.013 & \(-0.48\) \({}^{+0.60}_{-0.61}\) & \(-0.05\) \({}^{+0.14}_{-0.16}\) \\ \(R_{\mu\mu}\) & Planet radius (in \(R_{\mu}\)) & Normal(\(1,0.2\)) & 1.00 & 1.00 \({}^{+0.03}_{-0.57}\) & 1.00 \({}^{+0.03}_{-0.65}\) \\ \(\log_{10}(M_{\mu\mu})\) & Planet mass (in \(M_{\mu\mu}\)) & Normal(\(0,0.4\)) & 0.00 & 0.02 \({}^{+0.87}_{-0.57}\) & 0.03 \({}^{+0.65}_{-0.65}\) \\ \(\log_{10}(X_{\mathrm{N_{2}}})\) & N\({}_{2}\) mass fraction & Uniform(\(-2,0\)) & \(-0.11\) & \(-1.04\) \({}^{+0.89}_{-0.89}\) & \(-1.05\) \({}^{+0.93}_{-0.83}\) \\ \(\log_{10}(X_{\mathrm{O_{2}}})\) & O\({}_{2}\) mass fraction & Uniform(\(-2,0\)) & \(-0.70\) & \(-1.10\) \({}^{+0.86}_{-0.88}\) & \(-1.11\) \({}^{+0.99}_{-0.94}\) \\ \(\log_{10}(X_{\mathrm{CO_{2}}})\) & CO\({}_{2}\) mass fraction & Uniform(\(-10,0\)) & \(-3.40\) & \(-2.72\) \({}^{+1.33}_{-1.25}\) & \(-3.22\) \({}^{+0.80}_{-0.75}\) \\ \(\log_{10}(X_{\mathrm{CH_{4}}})\) & CH\({}_{4}\) mass fraction & Uniform(\(-10,0\)) & \(-5.77\) & \(-5.06\) \({}^{+1.56}_{-1.27}\) & \(-5.61\) \({}^{+0.84}_{-0.84}\) \\ \(\log_{10}(X_{\mathrm{H_{2}O}})\) & H\({}_{2}\)O mass fraction & Uniform(\(-10,0\)) & \(-3.00\) & \(-2.29\) \({}^{+1.36}_{-1.32}\) & \(-2.82\) \({}^{+0.77}_{-0.73}\) \\ \(\log_{10}(X_{\mathrm{O_{2}}})\) & O\({}_{3}\) mass fraction & Uniform(\(-10,0\)) & \(-6.52\) & \(-6.00\) \({}^{+1.07}_{-1.04}\) & \(-6.39\) \({}^{+0.72}_{-0.69}\) \\ \(\log_{10}(X_{\mathrm{CO}})\) & CO mass fraction & Uniform(\(-10,0\)) & \(-6.90\) & \(-7.47\) \({}^{+2.93}_{-2.40}\) & \(-7.76\) \({}^{+2.45}_{-2.13}\) \\ \hline \hline \end{tabular}
\end{table}
Table 1: Overview of all parameters in our simulated atmospheric retrieval, including the respective priors, true values, and retrieved values for both PT profile parameterization schemes. For the latter, we report the median as well as the 2.5% and 97.5% percentiles.
### Choosing a prior distribution for \(\mathbf{z}\)

The last point in the previous section is closely related to the question of which prior one should choose for \(\mathbf{z}\) during a retrieval. The natural choice, of course, is to assume a \(d\)-dimensional standard Gaussian prior for \(\mathbf{z}\) as that is the distribution that we encourage during training by using the \(\mathcal{L}_{\text{MMD}}\) loss. However, there are a few additional considerations to keep in mind here. First, using a Gaussian prior for \(\mathbf{z}\) means that the induced distribution over PT profiles, given by \(D(\cdot\,|\,\mathbf{z})\), will (approximately) match the distribution in the training dataset. This is a good choice if the training dataset itself constitutes a good prior for the PT profiles one wants to consider during a retrieval. On the other hand, consider now the case mentioned above where the training dataset contains all the expected shapes, but is highly biased towards one of the modes, not because this reflects the (assumed) true distribution of atmospheres, but as an artifact of the dataset generation process. In this case, it may make sense to choose a non-Gaussian prior (e.g., a uniform distribution) for \(\mathbf{z}\) to compensate for this imbalance. When doing so, it is important to ensure that \(\mathbf{z}\) stays in the value range that the decoder \(D\) encountered during training (i.e., \(\|\mathbf{z}\|<\tau\)). Another (and perhaps more principled) way to solve this problem is to resample the training dataset to balance the different modes (e.g., "inversion" and "no inversion"). For our experiment in Section 4.4, we have validated that such a resampling-based approach does indeed work, but have chosen not to include all the details of this additional experiment so as not to distract too much from the main message of our paper.

Finally, there is the fact that our trained model does not only return physically consistent PT profiles, but also allows smooth interpolation (or, to some extent, extrapolation) between them. In cases where this is not desired, for example because one wants to be sure that only PT profiles very close to the training dataset are produced, one can choose a data-driven prior distribution for \(\mathbf{z}\). A simple way to achieve this is to train a normalizing flow (see, e.g., Kobyzev et al.2021 or Papamakarios et al.2021 for recent introductions) to match the distribution of \(\mathbf{z}\)'s obtained from the training dataset. If we then use the normalizing flow to sample \(\mathbf{z}\), we will almost perfectly reproduce the distribution of the training data set. However, this can also exacerbate some of the problems we discussed earlier, for example, if the training data is unbalanced.

### Directions for future research

We see several directions for future work. For example, to account for the fact that not all parts of the atmosphere contribute equally to the observed spectrum, one possible approach could be to use a modified reconstruction loss weighted by the respective contribution function for each point in the profile. In the case of emission spectra, this could focus the expressiveness of the model on regions of higher pressure, which have a greater influence on the spectrum, rather than regions at high altitudes, where the thermal structure of the atmosphere is difficult to constrain from the available observational data. For transmission spectra, which probe the upper layers of the atmosphere, one would do the opposite and give more weight to high altitudes.

Another promising direction could be to extend our method to parameterize not only temperature as a function of pressure, but also chemical abundance profiles (e.g., the concentration of CO\({}_{2}\) in each layer), to make them accessible during a retrieval.

Finally, given a sufficiently large grid of training data, our approach can also be used to parameterize the thermal structure of 3-dimensional atmospheres, where temperature depends not only on pressure but also on the longitude and latitude.

## 6 Summary and conclusion

In this work, we have introduced a new approach to parameterizing PT profiles in the context of atmospheric retrievals of exoplanets. Unlike existing approaches for this problem, we do not make any a priori assumptions about the family of functions that describe the relation between \(P\) and \(T\), but instead learn a distribution over such functions from data. In our framework, PT profiles are represented by low-dimensional real vectors that can be used to condition a particular neural network, which, once conditioned, _is_ the corresponding PT profile (i.e., a single-argument, single-output function that maps pressure onto temperature values).

Figure 9: Results for our simulated atmospheric retrieval of an Earth-like planet using petitRADTRANS, LIFEsim, and PyMultiNest. _Left:_ The retrieved PT profiles for the polynomial baseline and our model, together with the emission contribution function. _Right:_ The relative fitting error of the spectrum, computed as (true spectrum – simulated spectrum) / true spectrum.
We experimentally demonstrated that our proposed method works for both terrestrial planets and hot Jupiters, and when compared with existing methods for parameterizing PT profiles, we found that our methods give better approximations on average (i.e., smaller fitting errors) while requiring a smaller number of parameters. Furthermore, we showed that our learned latent spaces are still interpretable to some extent: For example, the latent representations are often directly correlated with other properties of the atmospheres, meaning that physically similar profiles are grouped together in the latent space. Finally, we have demonstrated by example that our method can be easily integrated into an existing atmospheric retrieval framework, where it resulted in a significant speedup in terms of inference time, and produced a tighter and more accurate posterior for the PT profile than the baseline.

Given access to a sufficient training dataset, we believe that our method has the potential to become a valuable tool for exoplanet characterization: Not only could it allow the use of physically consistent PT profiles during a retrieval, but it can also reduce the number of parameters required to model the thermal structure of an atmosphere. This can either reduce the overall computational cost or free up resources for other retrieval parameters of interest, allowing for more efficient and accurate atmospheric retrievals, which in turn could improve our understanding of exoplanetary habitability and the potential for extraterrestrial life.

###### Acknowledgements.

 We thank Sandra T. Bastelberger and Nicholas Wogan for pointing out the potential problems with the PyATMOS simulations and taking the time to discuss the issue with us. Furthermore, we thank Markus Bonse, Felix Dannert, Maximilian Dax, Emily Garvin, Jean Hayor and Vincent Simper for their helpful comments on the manuscript. Part of this work has been carried out within the framework of the National Centre of Competence Research Planets supported by the Swiss National Science Foundation under grants 51NF40-182901. 851940. 202606. SPQ and EA acknowledge the financial support of the SNSF. TDG acknowledges funding through the Max Planck ETH Center for Learning Systems. BSK acknowledges funding through the European Space Agency's Open Space Innovation Platform (ESA OSFP) program.
