Beyond Point Estimates: Benchmarking Uncertainty Quantification Methods on the AION-1 Astronomical Foundation Model

2026-06-05

AI Summary

This work addresses the lack of reliable uncertainty quantification in astronomical foundation models for galaxy property regression, which hinders robust scientific inference. Building upon frozen AION-1 embeddings, the study systematically evaluates seven uncertainty estimation methods across five tasks—including redshift and stellar mass—using Legacy Survey and DESI data. It introduces, for the first time in this domain, the locally valid and discriminative (LVD) conformal prediction framework, integrating techniques such as conformalized quantile regression (CQR), deep ensembles, and Monte Carlo Dropout. The LVD approach achieves marginal coverage close to the nominal 90% while significantly outperforming non-conformal baselines. Crucially, it provides per-galaxy uncertainty estimates that adaptively reflect individual prediction difficulty, ensuring both local validity and practical reliability for downstream scientific analysis.

Abstract

Foundation models for astronomical surveys offer powerful learned representations that can be transferred to downstream regression tasks such as galaxy property estimation. However, point predictions alone are insufficient for scientific inference; reliable uncertainty quantification (UQ) is essential. We compare seven UQ methods on galaxy property regression using frozen AION-1 foundation-model embeddings, predicting redshift, stellar mass, stellar-population age, gas-phase metallicity, and specific star-formation rate, from Legacy Survey photometry/imaging and DESI spectra, with PROVABGS-derived labels. Distribution-free conformal methods achieve marginal coverage within $\sim$1\,pp of the nominal 90\% across all properties, while non-conformal baselines (Deep Ensembles, MC~Dropout) fail to calibrate reliably. Among conformal approaches, Conformalized Quantile Regression (CQR) delivers the best coverage in the bin with the poorest model predictions. More importantly, only the Locally Valid and Discriminative (LVD) framework -- particularly when operating on AION-1 embeddings -- also provides finite-sample \emph{local validity}, producing intervals that adapt to each galaxy's local prediction difficulty rather than relying on marginal guarantees alone. These results establish conformal prediction, and LVD in particular, as the preferred UQ framework for uncertainty-aware inference on foundation-model embeddings in astrophysics.

Problem

• uncertainty quantification
• foundation models
• astronomical surveys
• galaxy property estimation
• scientific inference

Innovation

• conformal prediction
• uncertainty quantification
• foundation models
• local validity
• astronomical inference