Flow matching for atmospheric retrieval of exoplanets: Where reliability meets adaptive noise levels

¹ Max Planck Institute for Intelligent Systems, Max-Planck-Ring 4, 72076 Tübingen, Germany
² ETH Zurich, Institute for Particle Physics & Astrophysics, Wolfgang-Pauli-Strasse 27, 8093 Zurich, Switzerland
³ ELLIS Institute Tübingen, Maria-von-Linden-Straße 2, 72076 Tübingen, Germany
⁴ Max Planck Institute for Gravitational Physics (Albert Einstein Institute), Am Mühlenberg 1, 14476 Potsdam, Germany
⁵ ETH Zurich, Department of Earth and Planetary Sciences, Sonneggstrasse 5, 8092 Zurich, Switzerland
⁶ ETH Zurich, Department of Computer Science, Universitätsstrasse 6, 8092 Zurich, Switzerland

^★ Corresponding author; tgebhard@tue.mpg.de

Received: 11 August 2024
Accepted: 16 October 2024

Abstract

Context. Inferring atmospheric properties of exoplanets from observed spectra is key to understanding their formation, evolution, and habitability. Since traditional Bayesian approaches to atmospheric retrieval (e.g., nested sampling) are computationally expensive, a growing number of machine learning (ML) methods such as neural posterior estimation (NPE) have been proposed.

Aims. We seek to make ML-based atmospheric retrieval (1) more reliable and accurate with verified results, and (2) more flexible with respect to the underlying neural networks and the choice of the assumed noise models.

Methods. First, we adopted flow matching posterior estimation (FMPE) as a new ML approach to atmospheric retrieval. FMPE maintains many advantages of NPE, but provides greater architectural flexibility and scalability. Second, we used importance sampling (IS) to verify and correct ML results, and to compute an estimate of the Bayesian evidence. Third, we conditioned our ML models on the assumed noise level of a spectrum (i.e., error bars), and thus made them adaptable to different noise models.

Results. Both our noise-level-conditional FMPE and NPE models perform on a par with nested sampling across a range of noise levels when tested on simulated data. FMPE trains about three times faster than NPE and yields higher IS efficiencies. IS successfully corrects inaccurate ML results, identifies model failures via low efficiencies, and provides accurate estimates of the Bayesian evidence.

Conclusions. FMPE is a powerful alternative to NPE for fast, amortized, and parallelizable atmospheric retrieval. IS can verify results, helping to build confidence in ML-based approaches, while also facilitating model comparison via the evidence ratio. Noise level conditioning allows design studies for future instruments to be scaled up; for example, in terms of the range of signal-to-noise ratios.