A Probabilistic Autoencoder for Galaxy SED Reconstruction and Redshift Estimation: Application to Mock SPHEREx Spectrophotometry

1. Il Problema

Le moderne indagini astronomiche, in particolare la missione SPHEREx della NASA, stanno generando volumi senza precedenti di fotometria multibanda e spettroscopia a bassa risoluzione (0.75–5 µm). Questo crea una sfida critica per la stima delle redshift fotometriche (photo-z):

• I metodi tradizionali basati sul fitting di template (TF) utilizzano librerie discrete di spettri. Sebbene efficienti, possono fallire nel catturare la piena variabilità delle SED (Distribuzioni di Energia Spettrale) delle galassie e tendono a produrre stime di redshift eccessivamente confidenti ma errate quando i parametri di disturbo (polvere, tipo spettrale) sono correlati con il redshift o scarsamente vincolati dai dati a basso rapporto segnale/rumore (SNR).
• I metodi guidati dai dati (data-driven) spesso mancano di un quadro probabilistico robusto e di priors ben definiti, rendendo difficile la calibrazione delle incertezze.
• È necessario un metodo che modelli le SED in modo continuo, permetta inferenze bayesiane congiunte su redshift e parametri intrinseci, e sia scalabile per centinaia di milioni di sorgenti.

2. Metodologia: PAE e PAESpec

Gli autori propongono un framework basato su un Autoencoder Probabilistico (PAE), implementato nel codice pubblico PAESpec. L'approccio si articola in diverse fasi:

• Modellazione della SED nel frame di riposo: Un autoencoder (AE) con architettura CNN (Convolutional Neural Network) comprime spettri ad alta risoluzione nel frame di riposo in un vettore latente a bassa dimensionalità ($\zeta$). A differenza dei Variational Autoencoders (VAE), non impone un prior gaussiano durante l'addestramento iniziale, permettendo al modello di catturare la varianza spiegata dai dati senza regolarizzazione eccessiva.
• Flussi Normalizzanti (Normalizing Flows - NF): Una volta appresa la rappresentazione latente, un flusso normalizzante (basato su Inverse Autoregressive Flow - IAF) mappa la distribuzione complessa dei vettori latenti $\zeta$ su una distribuzione base semplice (gaussiana standard $u$). Questo permette un campionamento efficiente e definisce un prior esplicito sullo spazio latente.
• Modellazione Forward e Inferenza Bayesiana: Il modello combina la SED ricostruita con operazioni fisiche di spostamento verso il rosso (redshifting) e convoluzione con i filtri di SPHEREx. L'inferenza avviene campionando la distribuzione a posteriori congiunta $p(u, z | f)$ (dove $z$ è il redshift e $f$ i dati osservati) utilizzando il Microcanonical Langevin Monte Carlo (MCLMC), un campionatore basato su gradienti ottimizzato per GPU.
• Alternativa SBI: Per scenari che richiedono una velocità estrema, gli autori sviluppano anche un approccio di Inferenza Basata su Simulazione (SBI) utilizzando un encoder Transformer e un flusso normalizzante condizionale. Questo metodo apprende direttamente la posterior $p(z|x)$ dai dati simulati rumorosi.

3. Contributi Chiave

• Framework PAE per Photo-z: Introdurre un metodo che sostituisce il set discreto di template con una varietà continua di SED appresa, permettendo una migliore gestione delle incertezze e delle degenerazioni.
• Diagnostica degli Outlier: Sviluppo di un criterio semplice (il rapporto tra le incertezze $\sigma_{PAE}/\sigma_{TF}$) per identificare sorgenti in cui il metodo TF fallisce. Il PAE produce distribuzioni posteriori più ampie e realistiche in questi casi, segnalando la mancanza di potere vincolante nei dati, mentre il TF produce stime "troppo sicure" ma errate.
• Analisi delle Superfici di Verosimiglianza: Dimostrazione attraverso il confronto delle profile likelihoods che le discrepanze tra PAE e TF derivano dalla struttura del modello: il TF, su una griglia discreta, può trovare minimi locali artificiali, mentre il PAE, marginalizzando su un manifold continuo, riflette più fedelmente l'incertezza intrinseca dei dati a basso SNR.
• Scalabilità e Performance: Implementazione efficiente su GPU (PAESpec) che raggiunge 300 campioni al secondo per sorgente, e un metodo SBI alternativo che offre un throughput **200 volte più veloce** rispetto al PAE, mantenendo prestazioni simili.

4. Risultati Principali

Il metodo è stato testato su dati mock di SPHEREx (102 bande) derivati dal catalogo COSMOS-2020:

• Recupero della SED: Il PAE ricostruisce le SED con errori frazionali dell'1-3%, con una buona capacità di recupero anche in presenza di dati mancanti (fino al 75% delle bande).
• Stima del Redshift:
  • Rispetto al TF, il PAE riduce significativamente la frazione di outlier (errori > 3$\sigma$), specialmente per sorgenti a basso SNR.
  • Per sorgenti ad alta precisione ($\sigma_z/(1+z) < 0.01$), le prestazioni sono comparabili (NMAD ~0.0057 per PAE vs 0.0055 per TF).
  • Per campioni più ampi e meno precisi, il PAE recupera fino al 40% in più di sorgenti rispetto al TF, mantenendo un bias basso e una calibrazione delle incertezze superiore.
• Impatto dei Prior: L'uso di un prior informativo sulla distribuzione del redshift ($p(z)$) e del prior del flusso normalizzante ($p(u)$) è cruciale per le sorgenti deboli, riducendo gli outlier e migliorando la convergenza delle catene di Markov (statistica di Gelman-Rubin).
• Confronto SBI vs PAE: Il metodo SBI basato su Transformer raggiunge prestazioni quasi identiche al PAE (NMAD e bias simili) ma con un throughput di inferenza di circa $9 \times 10^6$ galassie per nodo-ora, rendendolo ideale per l'elaborazione di dataset massivi.

5. Significato e Implicazioni

Questo lavoro rappresenta un passo avanti significativo per l'analisi dei dati di SPHEREx e di future indagini cosmologiche:

• Calibrazione delle Incertezze: Fornisce un metodo per ottenere stime di redshift non solo accurate, ma anche ben calibrate probabilisticamente, essenziale per studi di struttura su larga scala (LSS) e lensing gravitazionale.
• Gestione dei Fallimenti del Modello: La capacità di identificare e "pulire" i campioni TF esistenti utilizzando il PAE come diagnostica permette di migliorare la qualità dei cataloghi esistenti senza dover scartare l'intero set di dati.
• Flessibilità: L'approccio bayesiano esplicito permette di incorporare facilmente nuovi dati (es. fotometria da altre survey) e di adattare i prior in base alle esigenze scientifiche specifiche.
• Scalabilità: La dimostrazione di un metodo SBI ultra-veloce apre la strada all'applicazione di tecniche di inferenza avanzate su dataset di miliardi di galassie, superando i colli di bottiglia computazionali dei metodi tradizionali.

In sintesi, il paper dimostra che l'uso di autoencoder probabilistici combinati con flussi normalizzanti offre un equilibrio superiore tra flessibilità del modello, accuratezza fisica e gestione delle incertezze rispetto ai metodi di fitting di template tradizionali, fornendo al contempo strumenti pratici per la validazione e la pulizia dei dati.