Exoformer: Accelerating Bayesian atmospheric retrievals with transformer neural networks

Introduction

The deployment of the James Webb Space Telescope (JWST) and the forthcoming Ariel mission marks a new chapter in exoplanetary science, providing unprecedented spectroscopic observations of planetary atmospheres. To fully characterize the spectral features uncovered by these missions, we must employ increasingly complex atmospheric models that account for different phenomena, such as convection and chemical disequilibrium. However, the traditional approach to analyzing these spectra involves Bayesian retrieval frameworks, which can be extremely computationally expensive and time-consuming. Because these tools often require millions of atmospheric forward models to achieve convergence, they create a significant bottleneck that could slow down the rapid analysis of the thousands of spectra expected from next-generation telescopes. In this work, we address this challenge by introducing Exoformer, a novel transformer-based neural network designed to rapidly generate informative prior distributions for the atmospheric transmission spectra of hot Jupiters. By integrating deep learning with classical Bayesian methods, we aim to significantly enhance the efficiency of atmospheric retrieval tools without compromising their scientific accuracy.

The Exoformer Architecture

To overcome the limitations of previous architectures like convolutional neural networks or long-short term memory networks, we adopted the transformer architecture (Vaswani et al., 2017), which is uniquely capable of capturing long-range dependencies and correlations within sequential data. The core of our model is the multi-head self-attention mechanism, which allows the network to relate every element of a spectral sequence to all other elements.

Our architecture uses a series of transformer encoder blocks, each consisting of multi-head self-attention, residual skip connections, and layer normalization to ensure numerical stability, followed by a feed-forward neural network. Finally, we employ a prediction layer that uses average pooling and a two-layer multilayer perceptron to output the values of six target atmospheric parameters. To estimate parameter uncertainties, we implemented the Monte Carlo (MC) dropout technique (Gal & Ghahramani, 2016), which allows us to sample the outputs multiple times during the inference phase and retrieve approximated posterior distributions.

 

Training and Methodology

We trained Exoformer using a dataset of 10^7 atmospheric transmission spectra of hot Jupiters, which were generated using the TauREx analytical forward model (Zingales & Waldmann, 2018). These spectra were parameterized by seven free parameters, though we focused our regression task only on six of them: the abundances of H2O, CH4, CO, and CO2, the isothermal temperature, and the planetary radius. To ensure numerical stability and prevent the model from being biased by varying transit depth magnitudes, we applied a normalization scheme where each spectrum was divided into 14 wavelength bands, with the spectral points in each interval normalized between 0 and 1. We also applied min-max scaling to the atmospheric parameters themselves to facilitate the training process.

The training process was conducted using the AdamW optimizer and mean squared error (MSE) as our training loss function. We partitioned our dataset into training (90%), testing (9%), and validation (1%) subsets to reliably assess model performance and generalization while keeping computational demands low.

 

Performance on Simulated and Real Observations

We evaluated the performance of Exoformer through both simulated observations and the analysis of real JWST data. In our first test, we used the TauREx 3 model and the Pandexo tool to simulate a NIRSpec PRISM observation of a hot Jupiter, incorporating realistic Gaussian-distributed noise. We found that Exoformer’s predictions were consistent with the ground truth values within the 1-sigma error bars, even when spectral data from certain wavelengths were missing, demonstrating the model's ability to leverage correlations captured within its embeddings.

To demonstrate the robustness of our tool in real-world scenarios, we applied Exoformer to JWST transmission spectra of the hot Jupiters WASP-39b and WASP-17b. Despite these real observations containing atmospheric phenomena not present in our training set — such as traces of SO2 and H2S molecules or complex cloud effects — the posterior distributions recovered by Exoformer remain statistically compatible with those obtained through the classical Bayesian TauREx framework within 1-sigma. This comparison highlights that while Exoformer exhibits slightly lower accuracy than the Bayesian retrieval, it remains a reliable tool for characterizing real exoplanetary atmospheres.

 

Accelerating Bayesian Retrievals

The most significant application of our work is the development of a hybrid retrieval approach. Bayesian algorithms, such as those using nested sampling, benefit from informative prior distributions, which constrain the search to high-probability regions of the parameter space. By using Exoformer to generate informative priors, we can significantly reduce the time of the retrieval process.

Our results demonstrate a significant improvement in computational efficiency. For WASP-17b (Figs 1 and 2), we observed a speedup of approximately three times. For the WASP-39b retrieval (Figs 3 and 4), our hybrid method achieved a speedup of nearly eight times compared to the classical approach using uniform priors. Crucially, we confirmed the statistical consistency of this hybrid approach by comparing the log-Bayesian evidence between the uniform and informative prior models. In both cases, the absolute difference in the Bayes factor was less than five, indicating no strong preference for either model (Kass & Raftery 1995; Trotta2007) and ensuring that our method does not compromise the integrity of the scientific results.

 

Conclusion

Exoformer can bridge the gap between the speed of deep learning and the rigour of Bayesian statistics. By providing informative priors, we enable Bayesian tools to focus their computational resources on the most relevant areas of the parameter space. While our current model is optimized for hot Jupiters, the fixed architecture of Exoformer allows for efficient retraining on new datasets, making it a scalable solution for exploring a wider range of exoplanet populations.

We acknowledge that further improvements, such as incorporating more sophisticated physical models and developing finer wavelength grids to match the high resolution of JWST and Ariel, will be necessary to increase accuracy. This paves the way for the rapid and efficient characterization of the vast amount of spectroscopic data that will be provided by the next generation of space-based observatories.

Fig 1: WASP-17 model fit

Fig 2: WASP-17 Corner Plot