$\texttt{Exoformer}$: Accelerating Bayesian atmospheric retrievals with transformer neural networks

1. Problem Statement

The advent of high-precision space telescopes like JWST and the upcoming Ariel mission has revolutionized exoplanetary spectroscopy, providing high-quality data that reveals complex atmospheric features. However, analyzing this data presents a significant computational bottleneck:

• Computational Cost: Classical Bayesian atmospheric retrieval methods (e.g., TauREx, NEMESIS) rely on nested sampling algorithms that require millions of forward model evaluations to converge on posterior distributions.
• Scalability: As models become more complex (incorporating multidimensional structures, chemical disequilibrium, and convection), the parameter space expands, making convergence slow and computationally prohibitive.
• Uniform Priors: Traditional retrievals often use uninformative (uniform) priors, forcing the sampler to explore vast regions of parameter space with low probability, further increasing computation time.

The authors propose a solution to accelerate these retrievals without sacrificing accuracy by leveraging deep learning to generate informative priors.

2. Methodology

The paper introduces Exoformer, a novel tool based on Transformer neural networks designed to predict atmospheric parameters from transmission spectra and generate informative prior distributions.

A. Architecture (Exoformer)

• Core Mechanism: The model utilizes a Transformer encoder architecture, specifically chosen for its self-attention mechanism. Unlike CNNs or LSTMs, Transformers excel at capturing long-range dependencies and correlations within sequential data (spectral wavelengths).
• Input Processing:
  • Input: Transmission spectra with 515 spectral points (0.3–50 µm).
  • Linear Embedding: Spectral points are transformed into 128-dimensional vectors via a Multi-Layer Perceptron (MLP).
  • Positional Encoding: Sine/cosine functions are added to preserve wavelength ordering information, which is critical for identifying absorption features.
• Encoder Blocks: The model uses 5 transformer blocks, each containing 8 multi-head self-attention layers and a feed-forward network (1024 hidden units).
• Output Layer: Average pooling reduces the sequence dimension, followed by a 2-layer MLP to output six scalar atmospheric parameters:
  1. Log abundances of H₂O, CH₄, CO, and CO₂.
  2. Isothermal temperature ($T_{iso}$).
  3. Planet radius ($R_p$).
  • Note: Planetary mass ($M_p$) was excluded from the direct prediction due to degeneracy with other parameters; it is fixed in subsequent retrievals.

B. Training and Uncertainty Estimation

• Dataset: Trained on 107,000 synthetic hot Jupiter spectra generated by TauREx 3, covering a wide range of physical parameters.
• Preprocessing: Spectra were normalized into 14 wavelength bands to remove transit depth magnitude biases. Parameters were min-max scaled.
• Uncertainty Quantification: To provide probability distributions rather than single point estimates, the authors employed Monte Carlo (MC) Dropout. By keeping dropout active during inference and running multiple forward passes ($N_{step}$), the model approximates the posterior distribution of the parameters, capturing both model and data uncertainty.

C. Hybrid Retrieval Strategy

The core innovation is a hybrid approach:

1. Exoformer Inference: The neural network rapidly analyzes an observed spectrum to generate informative prior distributions (Gaussian-like) for the atmospheric parameters.
2. Bayesian Retrieval: These informative priors replace standard uniform priors in the TauREx nested sampling framework.
3. Result: The nested sampler is constrained to high-probability regions, significantly reducing the volume of parameter space it must explore.

3. Key Contributions

• Novel Architecture: First application of Transformer networks (specifically self-attention) to exoplanetary atmospheric retrieval, demonstrating their suitability for capturing complex spectral correlations.
• Informative Priors Generation: A robust method to convert deep learning predictions into statistically valid priors for Bayesian inference, bridging the gap between data-driven speed and physics-based rigor.
• Uncertainty Propagation: Implementation of MC Dropout to provide full posterior distributions from the neural network, ensuring the generated priors reflect prediction confidence.
• Open Source: The code for Exoformer and the TauREx plugin (ExoformerPriors) is made publicly available.

4. Results

The method was validated using both simulated JWST data and real observations of WASP-39b and WASP-17b.

A. Performance on Simulated Data

• Exoformer successfully retrieved parameters consistent with ground truth values within $1\sigma$ error bars.
• The distributions were Gaussian-like, as expected from MC Dropout theory.

B. Comparison with Classical Retrievals (Real Data)

• Accuracy: The posterior distributions obtained using Exoformer-derived priors were consistent with those from standard TauREx retrievals using uniform priors (within $1\sigma$ for most parameters).
• Speedup: The most significant result was the reduction in computational time:
  • WASP-39b: Reduced from 498 hours (uniform priors) to 64 hours (informative priors) $\rightarrow$ ~7.8x speedup.
  • WASP-17b: Reduced from 86 hours to 27 hours $\rightarrow$ ~3.2x speedup.
  • Exoformer inference itself took only ~2 minutes on an NVIDIA A2 GPU.
• Statistical Consistency: Log-Bayes factors ($|\Delta \log Z|$) were calculated to compare the evidence of the two approaches.
  • For WASP-39b, $|\Delta \log Z| = 1.16$ (weak preference).
  • For WASP-17b, $|\Delta \log Z| = 4.87$ (moderate preference for uniform priors, likely due to broader priors yielding higher evidence when models fit poorly).
  • Crucially, no strong preference was found for either model, confirming that the hybrid approach does not bias the physical results.

C. Limitations Observed

• The models (both Exoformer and TauREx) struggled to fit specific spectral features (e.g., SO₂ in WASP-39b, strong H₂O bands in WASP-17b) because the training dataset and forward model lacked these specific chemical species and physical processes (e.g., photochemistry, 3D effects).
• The current wavelength grid (515 points) is optimized for HST and may lose resolution for high-resolution JWST data.

5. Significance

• Scalability for Future Missions: This hybrid approach is essential for handling the anticipated flood of data from JWST and Ariel. It enables the analysis of thousands of spectra that would otherwise be computationally intractable.
• Efficiency without Compromise: It demonstrates that deep learning can be used to accelerate rigorous Bayesian inference rather than replacing it, preserving the statistical robustness of physical models.
• Pathway to Complex Models: By reducing the computational cost of retrieval, this method paves the way for incorporating more realistic, complex atmospheric models (e.g., 3D dynamics, chemical disequilibrium) that are currently too expensive to run with standard uniform priors.
• Adaptability: The data-driven nature of the approach allows the model to be retrained on new datasets (e.g., for smaller exoplanets or different chemical regimes) without changing the underlying architecture.

In conclusion, Exoformer represents a paradigm shift in exoplanet characterization, successfully merging the speed of deep learning with the rigor of Bayesian statistics to overcome the computational bottlenecks of the next generation of space astronomy.