Fast Sampling of Cosmological Initial Conditions with Gaussian Neural Posterior Estimation

Sintesi Tecnica: Campionamento Rapido delle Condizioni Iniziali Cosmologiche con la Stima del Posteriore Gaussiano tramite Reti Neurali

Enunciato del Problema
Ricostruire il campo di densità della materia primordiale (condizioni iniziali cosmologiche, o IC) dalle osservazioni tardive della struttura su larga scala (LSS) è una sfida fondamentale in cosmologia. Sebbene l'Universo primordiale sia ben descritto da fluttuazioni gaussiane lineari, la rete cosmica attuale è altamente non lineare e non gaussiana. Invertire questa evoluzione è un problema mal posto: la mappatura dal tempo iniziale al tempo finale è non iniettiva (diversi IC possono generare lo stesso campo tardivo) e la formazione delle strutture sulle scale degli aloni è caotica, amplificando il rumore di misura.

Gli approcci probabilistici tradizionali, come il Bayesian Origin Reconstruction from Galaxies (BORG), utilizzano l'Hamiltonian Monte Carlo (HMC) per campionare la distribuzione posteriore. Tuttavia, questi metodi richiedono simulatori differenziabili (spesso basati su schemi di approssimazione) e sono computazionalmente proibitivi, richiedendo mesi di tempo CPU per convergere su spazi di parametri ad alta dimensionalità ($\gtrsim 10^6$ parametri). Al contrario, molti approcci di machine learning in questo dominio si affidano a stimatori puntuali (ad esempio, U-Net) che forniscono una singola ricostruzione "migliore" senza quantificare l'incertezza, o a modelli generativi (ad esempio, la diffusione basata sul punteggio/score-based diffusion) che rimangono computazionalmente costosi anche durante l'inferenza.

Metodologia
Gli autori propongono un framework di Inferenza Basata su Simulazione (SBI) denominato Gaussian Neural Posterior Estimation (NPE). L'innovazione centrale consiste nel modellare la distribuzione posteriore delle condizioni iniziali, $p(\mathbf{z}|\mathbf{x})$, come una distribuzione Gaussiana nel campo inferito $\mathbf{z}$, con una matrice di covarianza diagonale nello spazio di Fourier.

La metodologia procede come segue:

1. Formulazione Statistica:

   • Prior: Si assume che il prior $p(\mathbf{z})$ sia un campo casuale gaussiano definito dallo spettro di potenza della materia lineare ad alto redshift ($z=127$).
   • Likelihood: La verosimiglianza $p(\mathbf{x}|\mathbf{z})$ è modellata come una Gaussiana con una media dipendente dall'osservazione $\hat{\mathbf{z}}_\theta(\mathbf{x})$ e una matrice di precisione $\mathbf{Q}^L_\theta$. Fondamentalmente, $\mathbf{Q}^L_\theta$ è vincolata ad essere diagonale nello spazio di Fourier (e approssimativamente invariante per rotazione), riflettendo l'invarianza traslazionale della fisica.
   • Posteriore: La combinazione di un prior gaussiano e di una verosimiglianza gaussiana produce un posteriore gaussiano. La matrice di precisione del posteriore è la somma delle matrici di precisione del prior e della verosimiglianza ($\mathbf{Q}_\theta = \mathbf{Q}_P + \mathbf{Q}^L_\theta$).

2. Architettura Neurale:

   • Stimatore MAP: La media del posteriore, $\hat{\boldsymbol{\mu}}_\theta(\mathbf{x})$, è modellata utilizzando un'architettura U-Net 3D. Per migliorare le prestazioni, gli autori introducono uno strato di attivazione personalizzato che separa la ricostruzione in due scale: un fattore di scala addestinabile $\alpha_\theta(\mathbf{k})$ gestisce l'evoluzione lineare su larga scala, mentre la U-Net si concentra sulla ricostruzione dei residui non lineari su piccola scala.
   • Matrice di Precisione: Sono appresi gli elementi diagonali della matrice di precisione della verosimiglianza, $D^L_\theta(k)$. Gli autori hanno testato due parametrizzazioni: (1) trattare i valori come parametri liberi per ogni modo di Fourier, e (2) imporre la simmetria rotazionale tramite una rete neurale densa che prende come input il numero d'onda $k$. L'approccio con la matrice "libera" ha mostrato una stabilità leggermente superiore.

3. Addestramento e Dati:

   • Il modello è addestrato su 2.000 coppie di campi di densità della materia oscura iniziale ($z=127$) e finale ($z=0$) provenienti dalla suite di simulazioni N-body Quijote.
   • Il modello forward utilizza simulazioni N-body complete e non differenziabili (Gadget-III), evitando la necessità di approssimazioni differenziabili richieste dai metodi HMC.
   • L'addestramento minimizza il log-posteriore negativo (Eq. 7) su 30 epoche su una singola GPU Nvidia A100, richiedendo circa 1,5 ore.

Contributi Chiave

• Campionamento Efficiente: Assumendo un posteriore gaussiano con una matrice di precisione diagonale nello spazio di Fourier, gli autori consentono la generazione di migliaia di campioni posteriori in pochi secondi (circa 3 secondi per 1.000 campioni su una singola GPU), ordini di grandezza più velocemente rispetto ai metodi HMC o basati sulla diffusione esistenti.
• Simulatori Non Differenziabili: Il metodo non richiede che il modello forward sia differenziabile, permettendo l'uso di simulazioni N-body complete e all'avanguardia senza schemi di approssimazione.
• Fit Analitico per il Campionamento Rapido: Gli autori dimostrano che i valori diagonali della matrice di precisione appresa $D_\theta(k)$ possono essere adattati accuratamente con una semplice funzione analitica (una Gaussiana più una costante) fino a $k \approx 0.2 \, h/\text{Mpc}$. Ciò consente a qualsiasi stimatore puntuale (MAP) addestrato di essere convertito in un campionatore posteriore veloce senza dover riaddestrare l'intera matrice di precisione.
• Quantificazione dell'Incertezza: A differenza degli standard stimatori puntuali, questo metodo fornisce una distribuzione posteriore completa, consentendo una rigorosa quantificazione dell'incertezza per la ricostruzione del campo.

Risultati
Gli autori validano il metodo utilizzando una simulazione fiduciale non presente nel set di addestramento:

• Consistenza Visiva e Statistica: I campioni generati somigliano visivamente alla verità fondamentale (ground truth). Le statistiche di riepilogo (spettro di potenza, funzione di trasferimento, cross-correlazione) mostrano un'alta concordanza con le condizioni iniziali reali. Lo spettro di potenza corrisponde alla verità entro l'1–2% fino alla scala di Nyquist ($k_{Nyq} \approx 0.4 \, h/\text{Mpc}$).
• Cross-Correlazione: La cross-correlazione tra i campioni generati e la verità rimane elevata ($\gtrsim 70\%$) fino a $k \approx 0.3 \, h/\text{Mpc}$.
• Consistenza Bayesiana: Un test di consistenza dipendente dalla scala conferma che le differenze normalizzate tra i modi reali e la media del posteriore seguono una distribuzione normale standard ($N(0,1)$) attraverso vari bin di $k$. Ciò verifica che le incertezze inferite sono ben calibrate.
• Robustezza: Il metodo mantiene le prestazioni anche quando viene aggiunto del rumore gaussiano al campo finale osservato.

Significato e Rivendicazioni
L'articolo sostiene che questo approccio fornisce una via efficiente in termini di simulazione per l'SBI sequenziale in cosmologia. Riducendo la complessità del posteriore alla forma gaussiana con una matrice di precisione diagonale nello spazio di Fourier, gli autori ottengono un equilibrio tra accuratezza fisica e velocità computazionale che era precedentemente irraggiungibile per la ricostruzione di campi ad alta dimensionalità.

Gli autori notano con modestia che l'approssimazione gaussiana con diagonale di Fourier è sufficiente per l'intervallo di scala selezionato ($k \lesssim 0y.4 \, h/\text{Mpc}$), ma potrebbe richiedere matrici di precisione più sofisticate (ad esempio, basate su convoluzioni o wavelet) per scale più piccole o quando gli effetti osservativi rompono l'invarianza traslazionale. Essi sottolineano che, sebbene l'attuale lavoro assuma parametri cosmologici fissi e osservazioni idealizzate della materia oscura, il framework è progettato per essere esteso a cosmologie variabili e survey galattiche realistiche in lavori futuri. La principale significatività risiede nel rendere possibile la generazione di migliaia di realizzazioni vincolate in pochi secondi, aprendo la strada ad analisi statistiche robuste della LSS che erano precedentemente computazionalmente impraticabili.