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

Technische Samenvatting: Benchmarking van Onzekerheidskwantificatie op AION-1

Probleemstelling
Fundamentele modellen voor astronomische surveys, zoals AION-1, bieden krachtige geleerde representaties voor downstream regressietaken zoals de schatting van galactische eigenschappen. Echter, puntvoorspellingen alleen zijn onvoldoende voor wetenschappelijke inferentie, die betrouwbare Onzekerheidskwantificatie (Uncertainty Quantification, UQ) vereist om fouten te propageren, out-of-distribution inputs te identificeren en robuuste populatiebrede inferenties te waarborgen. Hoewel UQ cruciaal is, blijft het slecht begrepen voor grote, black-box fundamentele modellen met hoogdimensionele embeddings. Bestaande methoden missen vaak formele dekkinggaranties of slagen er niet in aan te passen aan lokale voorspellingsproblemen in complexe, heterogene astrofysische data.

Methodologie
De auteurs voeren een directe vergelijking uit van zeven UQ-methoden toegepast op het bevroren AION-1 fundamentele model (specifiek de AION-B variant met 300 miljoen parameters). De studie richt zich op de regressie van vijf fysieke eigenschappen van Legacy Survey fotometrie/imaging en DESI spectra: roodverschuiving ($z$), stellaire massa ($\log M_\star$), stellaire populatie-leeftijd ($t_{age}$), gasfase-metalliciteit ($\log Z$) en specifieke stervormingssnelheid ($\log sSFR$). Labels zijn afgeleid van PROVABGS SED-fitting posteriors.

De experimentele opzet omvat:

1. Basisarchitectuur: Een drie-laags MLP brengt de 768-dimensionale AION-1 embeddings in kaart naar scalaire voorspellingen. De encoder-gewichten zijn bevroren; alleen de regressie-head wordt getraind.
2. Geëvalueerde UQ-methoden:
   • Conforme methoden (distributievrije garanties):
     • VanillaSplit: Standaard split conformal prediction met globale kwantielen (constante breedte intervallen).
     • MADSplit: Een MAD-genormaliseerde variant die de intervalbreedte aanpast aan lokale moeilijkheidsgraad.
     • CQR (Conformalized Quantile Regression): Traint quantile networks gevolgd door een conforme correctie.
     • LVD (Locally Valid and Discriminative): Leert een kernel in een gereduceerde feature-ruimte om conditionele non-conformity kwantielen te schatten. Getest in twee ruimtes: op de ruwe AION-1 embeddings (LCAion) en op de 256-dimensionale MLP hidden activaties (LCNet).
     • LVD-MAD: LVD uitgebreid met een geleerde residual-magnitude predictor voor schaal-adaptieve intervallen (LCMADAion en LCMADNet).
   • Niet-conforme baselines:
     • Deep Ensembles (DE): Vijf onafhankelijk geïnitialiseerde MLPs.
     • MC Dropout: Stochastische forward passes met dropout ingeschakeld.

De evaluatie meet marginale dekking, tail coverage (10% boven/onder), worst-predicted bin coverage (10% moeilijkste punten) en efficiëntie (gemiddelde voorspellingsinterval breedte).

Belangrijkste Resultaten

• Marginale Dekking: Alle conforme methoden (LVD varianten, VanillaSplit, MADSplit, CQR) bereiken een marginale dekking binnen $\sim$1 procentpunt van de nominale 90% doelstelling voor alle eigenschappen. In tegenstelling hiertoe falen niet-conforme baselines in kalibratie: Deep Ensembles onderschat systematisch (12–38%), terwijl MC Dropout sterke overdekking vertoont (57–99%) met excessief brede intervallen.
• Tail en Worst-Predicted Performance:
  • In de worst-predicted bin bereikt CQR de hoogste dekking (44–66%), gevolgd door LVD-MAD varianten (22–42%).
  • VanillaSplit stort in tot $\sim$0% dekking in de moeilijkste regio's.
  • LVD varianten die opereren op de ruwe AION-1 embeddings (LCAion, LCMADAion) presteren over het algemeen beter dan die gebruikmaken van gecomprimeerde MLP activaties, wat suggereert dat de 768-dimensionale ruimte structuur behoudt die gunstig is voor lokale kalibratie.
  • De MAD varianten (LCMADNet, LCMADAion) presteren consistent beter dan hun niet-MAD tegenhangers, wat het belang van schaal-adaptieve normalisatie voor heteroscedasticiteit benadrukt.
• Efficiëntie: Onder de goed gekalibreerde methoden produceren LVD varianten de nauwste voorspellingsintervallen terwijl ze de dekking behouden. CQR produceert bredere intervallen om marginale validiteit te garanderen.

Belangrijkste Bijdragen

1. Systematische Benchmarking: Het artikel biedt de eerste directe vergelijking van zeven UQ-methoden op bevroren AION-1 embeddings voor vijf verschillende galactische eigenschap-regressietaken.
2. Demonstratie van Lokale Validiteit: De studie demonstreert dat hoewel conforme methoden marginale validiteit garanderen, het Locally Valid and Discriminative (LVD) framework — met name wanneer het opereert op fundamentele model embeddings — eindige-steekproef lokale validiteit biedt. Dit maakt het mogelijk dat voorspellingsintervallen zich aanpassen aan de specifieke moeilijkheidsgraad van de voorspelling van elke sterrenstelsel, in plaats van enkel te vertrouwen op globale gemiddelden.
3. Analyse van de Embedding Ruimte: De resultaten wijzen erop dat het gebruik van de ruwe hoogdimensionele embeddings (768-dim) in plaats van gecomprimeerde intermediaire features (256-dim) de robuustheid van lokale kalibratie verbetert, met name in de tail regio's.

Significantie en Claims
De auteurs concluderen dat conforme predictie, en specif specifiek het LVD framework, de voorkeursbenadering voor UQ moet zijn voor onzekerheidsbewuste inferentie op fundamentele model embeddings in de astrofysica. In tegenstelling tot niet-conforme baselines die geen formele garanties bieden, en standaard conforme methoden die kunnen lijden onder conditionele miskalibratie, biedt LVD een principieel evenwicht tussen distributievrije garanties en lokale adaptiviteit. Dit is cruciaal voor wetenschappelijke inferentie waarbij de betrouwbaarheid van individuele voorspellingen even belangrijk is als populatiebrede statistieken. Het artikel merkt op dat hoewel AION-B werd gebruikt, het voordeel van LVD op embeddings verdere investigatie rechtvaardigt bij grotere model schalen (AION-L en AION-XL).