Source-linked AI summary
A Missing Latent, Not a Missing Simulator: Radius-Augmented Inference for Real JWST Retrieval
Angshuman Chakravertty, P V V Raj
TL;DR
Amortized SBI can retrieve exoplanet atmospheres accurately from synthetic JWST spectra yet degrade on real reduced spectra. MIRAGE adds planet radius to a flow-matching posterior and calibrates it against nested sampling; the same method transfers across three real targets and two instruments, including a WASP-39b fit with χ2/N reduced from 301 to 0.06.
Problem
Method
MIRAGE adds planet radius to a flow-matching posterior, domain-randomizes training spectra, and calibrates against an independent nested-sampling reference using importance sampling and an optimal-transport map.
Source: §1 · p. 2, §2.1 · p. 3Results
Takeaways & Limitations
A factor encoded by the simulator but omitted from inference can make a well-specified forward model look misspecified, a lesson the paper describes as simulator-agnostic.
Source: §5 · p. 6Calibration requires an independent nested-sampling run for each target, and coverage against that same reference is not independent evidence about the real spectrum.
Source: §5 · p. 5
Abstract
from arXiv · showhide
Amortized simulation-based inference (SBI), which is trained on radiative-transfer simulators, recovers exoplanet atmospheres accurately on synthetic James Webb Space Telescope (JWST) spectra but collapses when it comes to real reduced spectra. The flow posterior collapsed on real WASP-39b (importance-sampling effective sample size (ESS) = 1, best-fit \c{hi}2/N = 301), and such a failure is usually blamed on missing the forward-model physics, but ruling these levers out with nested sampling first (a temperature gradient, SO2 opacity, and a high-fidelity opacity set) leaves the fit unchanged, meaning the collapse is not from the simulator misspecification but instead from a missing latent, the planet radius. To fix this, we build MIRAGE, a radius-augmented flow-matching posterior calibrated against an independent nested-sampling reference with importance sampling and an optimal-transport map, which yields a physical and literature-consistent retrieval of real WASP-39b (with \c{hi}2/N from 301 to 0.06). This same method transfers unchanged across two instruments and three real JWST targets, including one spectrum self-reduced end-to-end from raw Mikulski Archive for Space Telescopes(MAST) data. The lesson is cross-domain, as a latent the simulator encodes but the inference omits can masquerade as misspecification.
1 Introduction
Amortized SBI works accurately on synthetic JWST spectra but can degrade on real reduced spectra, where simulator mismatch is a common concern. The paper diagnoses a missing radius latent and presents MIRAGE, which calibrates a radius-augmented posterior and transfers unchanged across real targets.
1 Introduction
- Amortized SBI and NPE are standard for scientific inverse problems, but real reduced JWST spectra can challenge retrievers that work accurately on synthetic spectra [Cranmer et al., 2020] [Vasist et al., 2023, Gebhard et al., 2025].When simulator and instrument differ, these estimators can concentrate confidently in the wrong place.Source: §1 · p. 1
- The real-data collapse is diagnosed as an identifiability failure from omitting planet radius, rather than forward-model misspecification.Without radius, the inference cannot absorb an approximately 3% transit-baseline offset and forces the atmosphere toward an overconfident, wrong solution.Source: §1 · p. 1
- MIRAGE adds radius to the inferred parameters, domain-randomizes training spectra, and calibrates an overdispersed flow posterior against an independent nested-sampling reference.The calibration uses an optimal-transport map and importance sampling; the method is applied unchanged across three targets spanning planet class, instrument, and data provenance.Source: §1 · p. 2
- χ2/N fell from 301 to 0.06 for real WASP-39b, with coverage rising from 2/7 to 7/7 after radius augmentation and calibration.MIRAGE also transfers unchanged across three real JWST targets and two instruments, including a spectrum self-reduced from raw MAST data.Source: §1 · p. 2
2 Method
MIRAGE trains a flow-matching posterior on simulated spectra with randomized correlated noise, includes planet radius among inferred parameters, and calibrates the flow against a nested-sampling reference. Its architecture separates offline simulation and training from per-target retrieval and calibration.
2.1 Forward model and training data
- Training uses TauREx spectra on each target’s JWST wavelength grid, with a seven-parameter vector containing radius, temperature, and five molecular abundances.The training set contains 28,000 parameter, noisy-spectrum, and covariance triples; correlated noise is randomized and supplied as a conditioning input.Source: §2.1 · pp. 2–3
- Correlated-noise domain randomization addresses the mismatch between idealized training noise and wavelength-correlated systematics in real JWST spectra.The covariance assigns variance to white noise and squared-exponential and Ornstein–Uhlenbeck components, with hyperparameters randomized for each training draw.Source: §2.1 · pp. 2–3
2.2 Flow-matching posterior with radius
- The flow-matching posterior infers radius alongside atmospheric parameters, allowing the network to absorb transit-baseline offsets instead of distorting the atmosphere.Adding radius is the only structural change from a standard flow-matching posterior.Source: §2.2 · p. 3
2.3 Calibration
- Calibration importance-reweights the overdispersed flow proposal and transports it onto an independent nested-sampling reference, with quality assessed by interval coverage.The calibrated samples are the transported flow samples; the architecture places simulation and amortized training offline and retrieval plus calibration at each real target.Source: Figure 1 · p. 2, §2.3 · pp. 3–4
3 Diagnosis: the gap is identifiability, not simulator fidelity
Independent nested-sampling tests left the real WASP-39b collapse unchanged when candidate forward-model physics was added. Freeing the omitted planet radius instead restored a warm, water-rich fit.
3 Diagnosis: the gap is identifiability, not simulator fidelity
- The real-data flow collapsed to a cold, flat atmosphere with best-fit χ2/N = 301 and importance-sampling ESS = 1, although published retrievals fit WASP-39b with the same forward-model class.The published retrieval comparison suggests the physics may not be the issue, motivating tests of what differed in the inference setup.Source: §3 · p. 4
- Adding a temperature gradient, SO2 opacity, and R=15000 ExoMolOP opacities left the cold-flat solution at χ2 near 2.5 in nested-sampling refits.The radius was held fixed except in the final variant; χ2/N was measured over 47 real-covered bins.Source: §3 · p. 4, Table 1 · p. 4
- Freeing radius with the same model and data drops χ2/N from 38.44 to 0.76 and returns a warm (∼600 K), water-rich atmosphere consistent with the literature.A roughly 3% mismatch between the assumed transit baseline and the real spectrum’s continuum could otherwise be absorbed only by distorting the atmosphere.Source: §3 · p. 4
4 Results
After optimal-transport calibration, the radius-augmented posterior produces a physical, literature-consistent WASP-39b retrieval and transfers unchanged across three real JWST targets. The evaluation also includes a self-reduced spectrum and a negative learned-translation comparison.
4 Results
- Best-fit χ2/N falls from 301 to 0.06 on real WASP-39b, with recovered radius and water abundance matching the nested-sampling anchor and published retrieval.The result is shown in Figure 2 and Table 2.Source: §4 · p. 4, Figure 2 · p. 5, Table 2 · p. 5
- The unchanged method covers nested-sampling reference values for 7/7, 7/7, and 6/7 inferred parameters across three targets spanning planet class, instrument, and data provenance.In every case, the recovered radius is near the literature value and the abundances are broad but physical.Source: §4 · p. 4
- On self-reduced WASP-96b data from raw MAST observations, the retrieval recovers the 1.4 µm water feature and is as strong as on published spectra.The spectrum was reduced end-to-end with NIRISS SOSS.Source: §4 · p. 4
- CycleGAN translation did not improve on structured randomization; the translated spectra had MMD ≈1.22 against the target spectra.The authors retain structured randomization as the production training procedure because the learned translator added complexity without benefit for this problem.Source: §4 · p. 5
5 Limitations and conclusion
MIRAGE’s current boundary includes a per-target nested-sampling requirement, limited transfer of covariance conditioning to real noise, and no observed benefit from CycleGAN translation. The paper concludes that omitted inference factors can make a well-specified simulator appear misspecified.
5 Limitations and conclusion
- Calibration still requires an independent nested-sampling run for each new target, and same-reference coverage is not independent evidence of recovering the true posterior.The optimal-transport map, rather than the raw overdispersed flow alone, aligns samples to the reference.Source: §5 · p. 5
- Covariance conditioning improved synthetic ABC coverage from 0.587 to 0.687 but did not transfer to the evaluated real target because its noise fell outside the training distribution.A higher-resolution training grid improved posterior centering, not sharpness.Source: §5 · p. 6
- CycleGAN translation did not improve upon structured randomization, while the conclusion frames omitted simulator-encoded factors as a source of apparent misspecification.The authors present MIRAGE as one added parameter plus calibration, not a new architecture.Source: §5 · p. 6
NeurIPS Paper Checklist
The checklist reports that the manuscript specifies the experimental setup and training details, but does not provide statistical-significance information or a concrete open-access path for code and data. It also flags that released artifacts are not identified with documentation or licensing information.
3. Theory assumptions and proofs; 14–16. Research ethics and declarations
- The checklist describes the work as an empirical study without theorem-level results requiring formal proofs or human-subject research.It also reports no described LLM use as part of the scientific method or evaluation.Source: NeurIPS Paper Checklist · pp. 7, 9
4. Experimental result reproducibility; 6. Experimental setting/details
- The checklist says the method, calibration, training-set size, and training details needed to understand or reproduce results are specified in Section 2 and Appendix A.Reported details include the forward model, priors, structured-noise ranges, objective, optimizer, and training schedule.Source: NeurIPS Paper Checklist · p. 8
5. Open access to data and code; 12. Licenses for existing assets
- The manuscript lacks an anonymized code/data access path and concrete reproduction instructions, and it does not identify released assets, their documentation, or their licenses.The checklist says these answers should be revisited after the anonymized artifact package is finalized.Source: NeurIPS Paper Checklist · pp. 8–9
7. Experiment statistical significance
- The manuscript reports point values but provides no error bars, confidence intervals, repeated-run variability, or statistical-significance procedure for its main comparisons.The checklist marks this item “No.”Source: NeurIPS Paper Checklist · p. 8
A Training details and priors
The appendix specifies uniform priors, model-training settings, compute, and the real spectra and nested-sampling reference used. It distinguishes hot-Saturn and cold sub-Neptune prior ranges while holding the five abundance priors fixed.
Priors
- Hot-Saturn targets use Rp ∼U(0.8, 1.6) RJ and T ∼U(300, 2500) K, while K2-18b uses Rp ∼U(0.10, 0.40) RJ and T ∼U(100, 600) K.All five log-abundances use log Xi ∼U(−9, −1).Source: A Training details and priors · p. 10
Network and training
- Training uses a four-layer transformer encoder, AdamW at learning rate 3 × 10−4, batch size 512, and up to 512 epochs on 28,000 training pairs.The set includes 4,000 held-out validation pairs; early stopping typically converged near epoch 250.Source: A Training details and priors · p. 10
Compute
- Models train on a single Apple Silicon machine, with the covariance arm on CPU and other arms on MPS; each has about 1M parameters.Inference is amortized at milliseconds per target.Source: A Training details and priors · p. 10
Real spectra and nested-sampling anchor
- The real-data set combines WASP-39b NIRSpec PRISM, self-reduced WASP-96b NIRISS SOSS, and K2-18b NIRISS spectra, with a TauREx/Nestle reference using 300 live points.WASP-39b has 52 bins, of which 47 real-covered bins are used for χ2.Source: A Training details and priors · p. 10