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

Sintesi Tecnica: Benchmarking della Quantificazione dell'Incertezza su AION-1

Problema
I modelli di fondazione per i rilievi astronomici, come AION-1, forniscono rappresentazioni apprese potenti per compiti di regressione a valle, come la stima delle proprietà galattiche. Tuttavia, le sole predizioni puntuali sono insufficienti per l'inferenza scientifica, che richiede una Quantificazione dell'Incertezza (UQ) affidabile per propagare gli errori, identificare input fuori distribuzione e garantire inferenze robuste a livello di popolazione. Sebbene l'UQ sia critica, rimane scarsamente compresa per i grandi modelli di fondazione "black-box" con embedding ad alta dimensionalità. I metodi esistenti spesso mancano di garanzie formali di copertura o non riescono ad adattarsi alle difficoltà locali di predizione in dati astrofisici complessi ed eterogenei.

Metodologia
Gli autori conducono un confronto diretto di sette metodi di UQ applicati al modello di fondazione congelato AION-1 (nello specifico la variante AION-B con 300 milioni di parametri). Lo studio si concentra sulla regressione di cinque proprietà fisiche dalla fotometria/imaging del Legacy Survey e dagli spettri DESI: redshift ($z$), massa stellare ($\log M_\star$), età della popolazione stellare ($t_{age}$), metallicità della fase gassosa ($\log Z$) e tasso specifico di formazione stellare ($\log sSFR$). Le etichette sono derivate dai posteri dell'SED-fitting di PROVABGS.

La configurazione sperimentale prevede:

1. Architettura Base: Un MLP a tre strati mappa gli embedding a 768 dimensioni di AION-1 in predizioni scalari. I pesi dell'encoder sono congelati; viene addestrata solo la testa di regressione.
2. Metodi UQ Valutati:
   • Metodi Conformal (Garanzie indipendenti dalla distribuzione):
     • VanillaSplit: Split conformal standard con quantili globali (intervalli a larghezza costante).
     • MADSplit: Variante normalizzata MAD che adatta la larghezza dell'intervallo alla difficoltà locale.
     • CQR (Conformalized Quantile Regression): Addestra reti quantili seguite da correzione conforme.
     • LVD (Locally Valid and Discriminative): Apprende un kernel in uno spazio di caratteristiche ridotto per stimare i quantili di non-conformità condizionale. Testato in due spazi: sugli embedding grezzi di AION-1 (LCAion) e sulle attivazioni nascoste dell'MLP a 256 dimensioni (LCNet).
     • LVD-MAD: LVD esteso con un predittore della magnitudo dei residui appreso per intervalli adattivi alla scala (LCMADAion e LCMADNet).
   • Baseline Non-Conformi:
     • Deep Ensembles (DE): Cinque MLP inizializzati indipendentemente.
     • MC Dropout: Passaggi forward stocastici con dropout abilitato.

La valutazione misura la copertura marginale, la copertura della coda (10% superiore/inferiore), la copertura del bin peggio predetto (10% dei punti più difficili) ed l'efficienza (larghezza media dell'intervallo di predizione).

Risultati Chiave

• Copertura Marginale: Tutti i metodi conformi (varianti LVD, VanillaSplit, MADSplit, CQR) raggiungono una copertura marginale entro circa l'1 punto percentuale dal target nominale del 90% per tutte le proprietà. Al contrario, le baseline non conformi non sono calibrate: i Deep Ensembles sottostimano sistematicamente la copertura (12–38%), mentre MC Dropout mostra una forte sovracopertura (57–99%) con intervalli eccessivamente ampi.
• Performance della Coda e del Peggior Predetto:
  • Nel bin peggior predetto, CQR raggiunge la copertura più alta (44–66%), seguito dalle varianti LVD-MAD (22–42%).
  • VanillaSplit crolla a circa lo 0% di copertura nelle regioni più difficili.
  • Le varianti LVD che operano sugli embedding grezzi di AION-1 (LCAion, LCMADAion) generalmente superano quelle che utilizzano attivazioni MLP compresse, suggerendo che lo spazio a 768 dimensioni conserva una struttura utile per la calibrazione locale.
  • Le varianti MAD (LCMADNet, LCMADAion) superano costantemente le loro controparti non-MAD, evidenziando l'importanza della normalizzazione adattiva alla scala per l'eteroschedasticità.
• Efficienza: Tra i metodi ben calibrati, le varianti LVD producono gli intervalli di predizione più stretti mantenendo la copertura. CQR produce intervalli più ampi per garantire la validità marginale.

Contributi Chiave

1. Benchmarking Sistematico: Il documento fornisce il primo confronto diretto di sette metodi di UQ sugli embedding congelati di AION-1 per cinque distinti compiti di regressione delle proprietà galattiche.
2. Dimostrazione della Validità Locale: Lo studio dimostra che, sebbene i metodi conformi garantiscano la validità marginale, il framework Locally Valid and Discriminative (LVD) — in particolare quando opera sugli embedding del modello di fondazione — fornisce una validità locale a campioni finiti. Ciò consente agli intervalli di predizione di adattarsi alla difficoltà specifica della predizione di ogni galassia anziché fare affidamento esclusivamente su medie globali.
3. Analisi dello Spazio degli Embedding: I risultati indicano che l'utilizzo degli embedding grezzi ad alta dimensionalità (768-dim) invece delle caratteristiche intermedie compresse (256-dim) migliora la robustezza della calibrazione locale, particolarmente nelle regioni di coda.

Significatività e Asserzioni
Gli autori concludono che la predizione conforme, e in particolare il framework LVD, dovrebbe essere l'approccio UQ preferito per l'inferenza consapevole dell'incertezza sugli embedding dei modelli di fondazione in astrofisica. A differenza delle baseline non conformi che mancano di garanzie formali, e dei metodi conformi standard che possono soffrire di miscalibrazione condizionale, LVD offre un equilibrio principiante tra garanzie indipendenti dalla distribuzione e adattabilità locale. Ciò è cruciale per l'inferenza scientifica dove l'affidabilità della singola predizione è importante quanto la statistica della popolazione. Il documento nota che, sebbene sia stato utilizzato AION-B, il vantaggio di LVD sugli embedding merita ulteriori investigazioni a scale di modello più grandi (AION-L e AION-XL).