🔭 astrophysics

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

This paper benchmarks seven uncertainty quantification methods on the AION-1 astronomical foundation model for galaxy property regression, demonstrating that conformal prediction approaches—specifically the Locally Valid and Discriminative (LVD) framework—outperform non-conformal baselines by providing reliable, adaptive, and locally valid uncertainty intervals essential for scientific inference.

Original authors: Karla Tame-Narvaez, Aleksandra Ćiprijanović, Shubhendu Trivedi

Published 2026-06-09

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Karla Tame-Narvaez, Aleksandra Ćiprijanović, Shubhendu Trivedi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to guess the weight of a mysterious object just by looking at a blurry photo of it. In the world of astronomy, scientists do something similar: they look at images and light spectra of galaxies to guess their physical properties, like how old they are or how much mass they contain.

For a long time, computer programs (specifically "Foundation Models" like AION-1) have gotten very good at making these guesses. They can look at a galaxy and say, "This one is 5 billion years old." But here's the problem: a single number isn't enough for real science. If you tell a scientist, "This galaxy is 5 billion years old," but you don't tell them how sure you are, they can't trust the result. Maybe the photo was blurry, or maybe the galaxy is weird. They need to know the range of possibilities, like, "It's probably between 4 and 6 billion years, but I'm not 100% sure."

This paper is like a taste test for different "confidence meters." The authors took the powerful AION-1 computer model and tested seven different ways to add that "confidence meter" (called Uncertainty Quantification) to its predictions.

The Contestants: Seven Ways to Guess the "Range"

The authors compared seven different methods to see which one gives the most honest and useful "confidence intervals" (the range of likely answers).

The "Safe but Boring" Group (Conformal Methods):
- Think of these like a weather forecaster who guarantees a 90% chance of rain. They use strict mathematical rules to ensure that, over a long period, their predictions are right 90% of the time.
- VanillaSplit is the simplest: It gives everyone the same size umbrella, regardless of the weather.
- CQR (Conformalized Quantile Regression) is smarter: It adjusts the umbrella size based on how hard the prediction is, but it still mostly focuses on the "average" performance.
- The LVD Family (Locally Valid and Discriminative): This is the star of the show. Imagine a weather forecaster who doesn't just look at the city average, but looks specifically at your backyard. If your backyard is foggy and hard to predict, they give you a huge umbrella. If it's sunny and easy, they give you a tiny one. This method adapts to the specific difficulty of each galaxy.
The "Gut Feeling" Group (Non-Conformal Methods):
- These are like asking five different experts to guess the weight and averaging their answers (Deep Ensembles) or asking one expert to guess 100 times while slightly changing their mind each time (MC Dropout).
- The paper found these methods were unreliable. One group was way too confident (giving tiny, dangerous ranges), and the other was way too scared (giving huge, useless ranges). They failed to calibrate properly, meaning their "confidence" didn't match reality.

The Results: Who Won?

The authors tested these methods on five different galaxy properties (like redshift, mass, and age). Here is what they found:

The "Average" Winner: The standard "Safe" methods (like CQR) did a decent job. They hit the 90% target on average. If you looked at 100 galaxies, they were right about 90 of them.
The "Hard Mode" Winner: When the predictions got really hard (the "worst-predicted" galaxies), the standard methods started to fail. They gave intervals that were too small, missing the true answer.
The Overall Champion: The LVD method (specifically the version using the raw AION-1 data) was the clear winner.
- Why? It didn't just promise to be right on average. It promised to be right for each specific galaxy.
- The Analogy: If you are trying to guess the weight of a feather, a standard method might give you a range of "1 to 10 grams." But the LVD method realizes, "Hey, this is a feather, it's easy to guess," and gives you "0.1 to 0.2 grams." If you are guessing the weight of a cloud (which is hard), it says, "This is tricky," and gives you a huge range like "1 to 1,000 tons."
- It adapts its "confidence interval" to the local difficulty of the prediction, making it much more useful for scientists dealing with weird or complex data.

The Bottom Line

The paper concludes that for astronomy, where data is messy and complex, you shouldn't just trust a single number. You need a confidence interval that knows when it's struggling.

While the "average" methods are okay, the LVD framework is the best tool for the job. It acts like a smart assistant that knows exactly when to say, "I'm pretty sure about this," and when to say, "This is a really tough one, so here is a wide range of possibilities." This ensures that when astronomers use these AI models to make big discoveries, they aren't accidentally trusting a guess that the computer knew was shaky.

Technical Summary: Benchmarking Uncertainty Quantification on AION-1

Problem Statement
Foundation models for astronomical surveys, such as AION-1, provide powerful learned representations for downstream regression tasks like galaxy property estimation. However, point predictions alone are insufficient for scientific inference, which requires reliable Uncertainty Quantification (UQ) to propagate errors, identify out-of-distribution inputs, and ensure robust population-level inferences. While UQ is critical, it remains poorly understood for large, black-box foundation models with high-dimensional embeddings. Existing methods often lack formal coverage guarantees or fail to adapt to local prediction difficulties in complex, heterogeneous astrophysical data.

Methodology
The authors conduct a direct comparison of seven UQ methods applied to the frozen AION-1 foundation model (specifically the AION-B variant with 300M parameters). The study focuses on regressing five physical properties from Legacy Survey photometry/imaging and DESI spectra: redshift ( $z$ ), stellar mass ( $\log M_\star$ ), stellar-population age ( $t_{age}$ ), gas-phase metallicity ( $\log Z$ ), and specific star-formation rate ( $\log sSFR$ ). Labels are derived from PROVABGS SED-fitting posteriors.

The experimental setup involves:

Base Architecture: A three-layer MLP maps the 768-dimensional AION-1 embeddings to scalar predictions. The encoder weights are frozen; only the regression head is trained.
UQ Methods Evaluated:
- Conformal Methods (Distribution-free guarantees):
  - VanillaSplit: Standard split conformal prediction with global quantiles (constant-width intervals).
  - MADSplit: MAD-normalized variant adapting interval width to local difficulty.
  - CQR (Conformalized Quantile Regression): Trains quantile networks followed by conformal correction.
  - LVD (Locally Valid and Discriminative): Learns a kernel in a reduced feature space to estimate conditional nonconformity quantiles. Tested in two spaces: on raw AION-1 embeddings (LCAion) and on 256-dimensional MLP hidden activations (LCNet).
  - LVD-MAD: LVD extended with a learned residual-magnitude predictor for scale-adaptive intervals (LCMADAion and LCMADNet).
- Non-Conformal Baselines:
  - Deep Ensembles (DE): Five independently initialized MLPs.
  - MC Dropout: Stochastic forward passes with dropout enabled.

The evaluation measures marginal coverage, tail coverage (10% upper/lower), worst-predicted bin coverage (10% hardest points), and efficiency (mean prediction interval width).

Key Results

Marginal Coverage: All conformal methods (LVD variants, VanillaSplit, MADSplit, CQR) achieve marginal coverage within $\sim$ 1 percentage point of the nominal 90% target across all properties. In contrast, non-conformal baselines fail to calibrate: Deep Ensembles systematically undercover (12–38%), while MC Dropout exhibits strong overcoverage (57–99%) with excessively wide intervals.
Tail and Worst-Predicted Performance:
- In the worst-predicted bin, CQR achieves the highest coverage (44–66%), followed by LVD-MAD variants (22–42%).
- VanillaSplit collapses to $\sim$ 0% coverage in the hardest regions.
- LVD variants operating on the raw AION-1 embeddings (LCAion, LCMADAion) generally outperform those using compressed MLP activations, suggesting the 768-dimensional space retains structure beneficial for local calibration.
- The MAD variants (LCMADNet, LCMADAion) consistently outperform their non-MAD counterparts, highlighting the importance of scale-adaptive normalization for heteroscedasticity.
Efficiency: Among well-calibrated methods, LVD variants produce the tightest prediction intervals while maintaining coverage. CQR produces wider intervals to guarantee marginal validity.

Key Contributions

Systematic Benchmarking: The paper provides the first direct comparison of seven UQ methods on frozen AION-1 embeddings for five distinct galaxy property regression tasks.
Local Validity Demonstration: The study demonstrates that while conformal methods guarantee marginal validity, the Locally Valid and Discriminative (LVD) framework—particularly when operating on foundation model embeddings—provides finite-sample local validity. This allows prediction intervals to adapt to the specific difficulty of each galaxy's prediction rather than relying solely on global averages.
Embedding Space Analysis: The results indicate that utilizing the raw high-dimensional embeddings (768-dim) rather than compressed intermediate features (256-dim) enhances the robustness of local calibration, particularly in tail regions.

Significance and Claims
The authors conclude that conformal prediction, and specifically the LVD framework, should be the preferred UQ approach for uncertainty-aware inference on foundation-model embeddings in astrophysics. Unlike non-conformal baselines which lack formal guarantees, and standard conformal methods which may suffer from conditional miscalibration, LVD offers a principled balance of distribution-free guarantees and local adaptivity. This is crucial for scientific inference where individual prediction reliability is as important as population-level statistics. The paper notes that while AION-B was used, the advantage of LVD on embeddings warrants further investigation at larger model scales (AION-L and AION-XL).

The Contestants: Seven Ways to Guess the "Range"

The Results: Who Won?

The Bottom Line

Technical Summary: Benchmarking Uncertainty Quantification on AION-1

More like this