[Paper Review] Bayesian Deep Learning for Exoplanet Atmospheric Retrieval

Frank Soboczenski, Michael D. Himes+9 more|arXiv (Cornell University)|Nov 8, 2018

Geochemistry and Geologic Mapping21 references18 citations

TL;DR

This paper introduces INARA, a Bayesian deep learning framework for exoplanet atmospheric retrieval that uses a 1D convolutional neural network trained on 3 million synthetic spectra generated with the NASA Planetary Spectrum Generator. It achieves inference in seconds—orders of magnitude faster than traditional methods—while predicting 12 atmospheric molecules with uncertainty estimates via Monte Carlo dropout, enabling probabilistic retrieval for rocky exoplanets.

ABSTRACT

Over the past decade, the study of extrasolar planets has evolved rapidly from plain detection and identification to comprehensive categorization and characterization of exoplanet systems and their atmospheres. Atmospheric retrieval, the inverse modeling technique used to determine an exoplanetary atmosphere's temperature structure and composition from an observed spectrum, is both time-consuming and compute-intensive, requiring complex algorithms that compare thousands to millions of atmospheric models to the observational data to find the most probable values and associated uncertainties for each model parameter. For rocky, terrestrial planets, the retrieved atmospheric composition can give insight into the surface fluxes of gaseous species necessary to maintain the stability of that atmosphere, which may in turn provide insight into the geological and/or biological processes active on the planet. These atmospheres contain many molecules, some of them biosignatures, spectral fingerprints indicative of biological activity, which will become observable with the next generation of telescopes. Runtimes of traditional retrieval models scale with the number of model parameters, so as more molecular species are considered, runtimes can become prohibitively long. Recent advances in machine learning (ML) and computer vision offer new ways to reduce the time to perform a retrieval by orders of magnitude, given a sufficient data set to train with. Here we present an ML-based retrieval framework called Intelligent exoplaNet Atmospheric RetrievAl (INARA) that consists of a Bayesian deep learning model for retrieval and a data set of 3,000,000 synthetic rocky exoplanetary spectra generated using the NASA Planetary Spectrum Generator. Our work represents the first ML retrieval model for rocky, terrestrial exoplanets and the first synthetic data set of terrestrial spectra generated at this scale.

Motivation & Objective

• To develop a fast, scalable, and probabilistically principled method for retrieving atmospheric parameters of rocky exoplanets from observed spectra.
• To overcome the computational infeasibility of traditional Monte Carlo-based retrieval methods, which require hundreds of CPU hours for a single retrieval.
• To generate a large-scale, physically motivated synthetic data set of 3 million spectra for rocky exoplanets to train machine learning models.
• To integrate uncertainty quantification into deep learning retrieval via Monte Carlo dropout, emulating Bayesian posterior distributions.
• To enable retrieval of a larger number of atmospheric molecules (12) than previous ML methods (3–4), including key biosignatures like O2 and CH4.

Proposed method

• A supervised deep learning framework trains a 1D convolutional neural network (CNN) to map observed spectra to atmospheric parameters using a synthetic data set of 3 million parameter-spectrum pairs.
• The data set is generated using the NASA Planetary Spectrum Generator (PSG) with physically motivated priors on planetary and atmospheric parameters, including 12 molecules.
• The CNN architecture uses multiple Conv1D, ReLU, and MaxPool layers followed by fully connected layers, with 18 million trainable parameters.
• Monte Carlo dropout is applied during inference to estimate predictive uncertainty, approximating a Bayesian posterior distribution.
• The framework is implemented in PyTorch and deployed on Google Cloud using distributed computing with 2,000 high-end VMs and a custom pypsg package for PSG integration.
• Model training uses mean squared error loss with early stopping and a validation set of 10,000 samples to prevent overfitting.

Experimental results

Research questions

• RQ1 Can a deep learning model trained on synthetic spectra achieve fast and accurate retrieval of atmospheric parameters for rocky exoplanets?
• RQ2 Can Monte Carlo dropout in a deep neural network provide uncertainty estimates comparable to traditional Bayesian retrieval methods?
• RQ3 How does the performance of the proposed model scale with the number of atmospheric molecules retrieved?
• RQ4 Can a large-scale synthetic data set of 3 million spectra improve the generalization and robustness of ML-based atmospheric retrieval?
• RQ5 How does the inference speed of the ML model compare to traditional Monte Carlo sampling-based retrieval methods?

Key findings

• INARA achieves inference in seconds, representing a speedup of several orders of magnitude over traditional methods, which require hundreds of CPU hours.
• The model successfully retrieves 12 atmospheric molecules, including key biosignatures such as H2O, CO2, O2, N2, and CH4, outperforming prior ML methods that were limited to 3–4 molecules.
• Monte Carlo dropout provides predictive distributions that contain the true atmospheric parameters within the uncertainty bounds, as demonstrated by the red star and line falling within the predicted distributions in Figure 1.
• The best-performing model is a 1D CNN with the architecture: Conv1d(64)–tanh–MaxPool–Conv1d(64)–relu–MaxPool–Conv1d(128)–relu–MaxPool–Conv1d(256)–relu–FC(256)–relu–FC(12).
• The framework was trained on 110,000 parameter-spectrum pairs in 64 epochs, with early stopping and validation to prevent overfitting, and generalizes well to unseen spectra.
• The study presents the first large-scale synthetic data set of 3 million spectra for rocky exoplanets, enabling future research in ML-based atmospheric retrieval.