Joint inference of weak lensing convergence map and cosmology with diffusion models

Sintesi Tecnica: Inferenza Congiunta di Mappe di Convergenza di Lente Debole e Cosmologia con Modelli di Diffusione

Definizione del Problema
Le attuali analisi di lensing gravitazionale debole si affidano tipicamente a statistiche riassuntive (ad esempio, funzioni di correlazione a due punti) che scartano l'informazione non gaussiana presente nella distribuzione della materia, in particolare sulle scale piccole e non lineari. Sebbene gli approcci a campo completo offrano un modo rigoroso per catturare tutta l'informazione statistica, i metodi esistenti affrontano una dicotomia:

1. Inferenza Marginale: I metodi basati su reti neurali inferiscono direttamente i parametri cosmologici ($\theta$) dai dati, ma marginalizzano sul campo di convergenza ($\kappa$) sottostante, fallendo nel ricostruire la mappa di massa.
2. Ricostruzione del Campo: Metodi come la mappatura della massa ricostruiscono $\kappa$ a partire dallo shear rumoroso ($\gamma_{obs}$), ma trattano i parametri cosmologici come fissi, anziché inferirli congiuntamente.
3. Inferenza Esplicita (MCMC): Approcci come BORG eseguono l'inferenza congiunta ($p(\theta, \kappa | \gamma_{obs})$) utilizzando il Markov Chain Monte Carlo (MCMC). Tuttavia, questi richiedono un modello forward differenziabile che colleghi le condizioni iniziali alle osservabili, limitando la complessità della fisica che può essere impiegata, e soffrono di elevati costi computazionali in spazi ad alta dimensionalità.

Gli autori mirano a sviluppare un metodo per l'inferenza congiunta che ricostruisca simultaneamente il campo di convergenza e inferisca i parametri cosmologici senza richiedere un modello forward differenziabile o un campionamento MCMC iterativo durante l'inferenza.

Metodologia: JADE
Gli autori introducono JADE (Joint Architecture for fielD and cosmological parameter Estimation), un framework di inferenza implicita basato su modelli di diffusione condizionati.

• Inferenza Implicita tramite Modelli di Diffusione: Invece di apprendere una funzione di verosimiglianza, JADE apprende la distribuzione posteriore congiunta $p(\theta, \kappa | \gamma_{obs})$ direttamente dalle simulazioni. Il metodo utilizza il Flow Matching, un framework generativo in cui una rete neurale apprende un campo di velocità per trasportare una distribuzione di base (rumore gaussiano) verso la distribuzione target. Il modello è addestrato per predire il campione "pulito" ($x_1 = [\theta, \kappa]$) da un interpolante rumoroso ($x_t$) dato il campo di shear osservato.
• Architettura Transformer: Il nucleo di JADE è un Vision Transformer in spazio di pixel adattato per dati multimodali.
  • Sequenza di Token Unificata: Il campo di convergenza $\kappa$ viene decomposto in patch, e i parametri cosmologici $\theta$ vengono proiettati in embedding della stessa dimensione. Questi vengono concatenati in una singola sequenza di token.
  • Condizionamento: Anche il campo di shear rumoroso osservato $\gamma_{obs}$ viene patchificato e proiettato, quindi concatenato alla sequenza. Il meccanismo di attenzione del transformer permette ai token target (campo e parametri) di prestare attenzione all'input di condizionamento (shear) in tutte le posizioni spaziali.
  • Output: La rete produce stime denoise sia per la mappa di convergenza che per i parametri cosmologici simultaneamente.
• Inferenza Ammortizzata: Una volta addestrato, il modello può generare campioni posteriori per qualsiasi nuova osservazione integrando un'equazione differenziale ordinaria (ODE) condizionata sulla specifica $\gamma_{obs}$. Questo processo è "ammortizzato", il che significa che non è richiesto alcun riaddestramento o ottimizzazione per singola osservazione (come nel caso di MCMC).

Configurazione della Simulazione
Il metodo è stato validato su dati di lensing debole simulati:

• Cosmologia: Un modello $w$CDM con sei parametri ($\Omega_c, \Omega_b, \sigma_8, h_0, n_s, w_0$).
• Generazione dei Dati: Le mappe di convergenza sono state generate utilizzando un modello log-normale 2D (un modello fenomenologico per la crescita non lineare) piuttosto che simulazioni N-body complete, per consentire la generazione rapida del grande dataset di addestramento (100.000 campioni).
• Osservabili: Campi di shear simulati con profondità LSST Year 10 (Y10) con rumore di forma (shape noise), coprendo una patch di $5 \times 5$ deg$^2$ con $128 \times 128$ pixel.

Risultati Chiave

1. Ricostruzione Congiunta: Il modello ha generato con successo campioni posteriori sia di $\kappa$ che di $\theta$ condizionati dallo shear rumoroso. L'ispezione visiva e i test statistici hanno confermato che le mappe di convergenza ricostruite recuperano le strutture reali pur riflettendo correttamente l'incertezza residua sulle scale piccole.
2. Consistenza Statistica:
   • Spettri di Potenza: Gli spettri di potenza auto e cross dei campioni posteriori sono corrispondenti a quelli delle simulazioni indipendenti estratte alle cosmologie campionate, con un errore relativo di circa il 6,7% mediato su bin e scale.
   • PDF a un punto: Le distribuzioni recuperate hanno esibito la forma log-normale attesa, confermando la cattura delle statistiche non gaussiane.
   • Residui: La statistica chi-quadro ridotto ($\chi^2_{red} \approx 1.002$) ha confermato che i residui tra osservazioni e campioni posteriori erano dominati dal rumore di forma iniettato, come previsto.
3. Vincoli Cosmologici: I posteri marginali cosmologici ($p(\theta | \gamma_{obs})$) derivati da JADE hanno mostrato un eccellente accordo con le catene MCMC NUTS (No-U-Turn Sampler) tradizionali in termini di regioni di credibilità e direzioni di degenerazione.
4. Calibrazione:
   • Punteggio MIRA: Il posteriore congiunto ha raggiunto un punteggio di calibrazione MIRA (Mass In Random Areas) di $0.635 \pm 0.017$ (dove 0.667 è l'ottimale), indicando una buona calibrazione con una lieve eccessiva confidenza. Il posteriore cosmologico marginale ha ottenuto $0.659 \pm 0.0139$.
   • TARP: I test di Accuratezza con Punti Casuali (Tests of Accuracy with Random Points) hanno confermato che le regioni di credibilità hanno raggiunto la corretta copertura frequentista attraverso lo spazio dei parametri.
5. Efficienza: Generare un singolo campione posteriore congiunto richiede circa 0,2 secondi (256 passi ODE). Estrarre 1.000 campioni richiede circa 3 minuti, rispetto al costo computazionale significativamente più elevato dell'esecuzione di catene MCMC (che richiedevano circa 2,18 milioni di valutazioni del gradiente per il benchmark).

Significato e Rivendicazioni
L'articolo sostiene di aver stabilito l'inferenza implicita basata sulla diffusione come una via percorribile verso analisi cosmologiche a livello di campo completo. Il significato primario risiede nell'unificare la ricostruzione della mappa di massa e la stima dei parametri cosmologici all'interno di un unico framework non differenziabile.

Le principali rivendicazioni riguardanti i vantaggi del metodo includono:

• Nessun Modello Forward Differenziabile Richiesto: A differenza dell'inferenza a livello di campo basata su MCMC, JADE non richiede gradienti analitici del modello forward, permettendo teoricamente l'uso di simulatori arbitrariamente complessi (ad esempio, codice N-body completi o idrodinamici) una volta disponibili i dati di addestramento.
• Ammortizzazione: Il modello viene addestrato una sola volta e applicato istantaneamente a qualsiasi nuova osservazione, evitando la necessità di campionamento per singola osservazione.
• Inferenza Congiunta: È il primo metodo presentato che inferisce congiuntamente il posteriore completo $p(\theta, \kappa | \gamma_{obs})$ utilizzando un singolo modello di diffusione condizionato, invece di marginalizzare sul campo o inferire i parametri separatamente.

Limitazioni e Lavori Futuri
Gli autori osservano che la dimostrazione attuale si basa su simulazioni log-normali, che sono più veloci dei codici N-body ma meno complete dal punto di vista fisico. Scalare a simulatori realistici ed onerosi presenta sfide riguardanti la densità del set di addestramento nello spazio dei parametri cosmologici. Inoltre, come metodo ammortizzato, JADE deve essere riaddestrato se il modello dei dati (ad esempio, maschera del sondaggio, modello di rumore) cambia, a differenza dei metodi di inferenza esplicita che possono adattarsi al momento dell'inferenza. Gli autori propongono come lavoro futuro l'estensione di questa architettura all'inferenza congiunta delle condizioni iniziali e dei parametri cosmologici, e l'esplorazione di approcci di diffusione latente per gestire campi a risoluzione più elevata in modo più efficiente.