Strong Gravitational Lensing Posterior Sampling in Pixel-Space Using Diffusion Models and Recurrent Inference Machines

Sintesi Tecnica: Campionamento del Posteriore di Lente Gravitazionale Forte nello Spazio dei Pixel utilizzando Modelli di Diffusione e Macchine di Inferenza Ricorrenti

Definizione del Problema
Modellare lenti gravitazionali forti galassia-galassia per inferire la distribuzione di luminosità di una galassia sorgente e la distribuzione di massa di una galassia in primo piano è un problema inverso computazionalmente oneroso. Questa sfida è esacerbata nel regime di osservazioni ad alta risoluzione e alto rapporto segnale-rumore (SNR), dove le assunzioni parametriche a bassa dimensionalità (ad esempio, semplici leggi di potenza ellittiche) sono insufficienti per catturare la complessità dei dati. Per modellare le osservazioni fino al livello del rumore, sono necessarie rappresentazioni pixelate ad alta dimensionalità sia della sorgente che della massa in primo piano. Tuttavia, rappresentare la lente e la sorgente come immagini pixelate rende il modello diretto (forward model) non lineare rispetto alla massa in primo piano, creando un problema inverso ad alta dimensionalità e non lineare caratterizzato da distribuzioni posteriori complesse e multimodali e da severe degenerazioni dei parametri. Gli algoritmi di campionamento tradizionali e gli approcorsi standard di machine learning spesso faticano a campionare efficacemente tali distribuzioni.

Metodologia: Il Framework DiRIM
Gli autori propongono il Diffusion Recurrent Inference Machine (DiRIM), un metodo bayesiano progettato per generare campioni del posteriore congiunto della luminosità della galassia sorgente ($s$) e della distribuzione di massa in primo piano (rappresentata come una mappa di convergenza, $\kappa$) condizionati dalle osservazioni ($y$). Il framework integra due distinti paradigmi di machine learning:

1. Modellazione Generativa basata su Diffusione: Il metodo utilizza modelli di diffusione basati sullo score per campionare da distribuzioni di probabilità complesse. Il processo prevede un SDE (equazione differenziale stocastica) in avanti che aggiunge gradualmente rumore ai dati, e un SDE inversa che ricostruisce i dati dal rumore. Per campionare dal posteriore $p(x|y)$, l'SDE inversa richiede lo score del posteriore $\nabla_x \log p_t(x|y)$.
2. Macchine di Inferenza Ricorrenti (RIMs): Le RIM sono impiegate per risolvere il problema inverso apprendendo un algoritmo di inferenza iterativo guidato dai dati. Invece di produrre un singolo stimatore MAP (Maximum A Posteriori), il framework RIM è qui adattato per raffinare iterativamente la stima dello score del posteriore.

L'Architettura DiRIM:

• Input: La rete neurale $g_\theta$ (un'architettura U-Net) prende in input i parametri del modello rumorosi al tempo $t$ ($x_t$), l'osservazione $y$, il passo temporale $t$, i gradienti della verosimiglianza ($\nabla L$) e i residui dell'immagine lenti normalizzati ($\delta y / \sigma_y$).
• Raffinamento Iterativo: A differenza dei modelli di diffusione standard che eseguono un singolo passaggio di denoising per ogni passo temporale, DiRIM utilizza il framework RIM per aggiornare iterativamente la propria stima della sorgente e della mappa di convergenza denoisate ($\hat{x}_0$) all'interno di ogni passo temporale di diffusione.
• Stima dello Score: La rete è addestrata per minimizzare una perdita di score-matching di denoising condizionale. Lo $\hat{x}_0^{(M)}$ finale delle iterazioni RIM viene utilizzato per calcolare lo score del posteriore richiesto per l'SDE inversa:
  $$\nabla_{x_t} \log p_t(x_t|y) = \frac{\hat{x}_0^{(M)} - x_t}{\sigma^2(t)}$$
• Dati di Addestramento: Il modello è addestrato su osservazioni mock realistiche generate da simulazioni cosmologiche magneto-idrodinamiche (IllustrisTNG) e profili analitici. Le immagini sorgente sono estratte dal dataset SKIRT-TNG (processato tramite trasferimento radiativo della polvere), e le mappe di convergenza sono o smussate da aloni di simulazione o renderizzate da profili analitici (Elliptical Power Law + multipoli + subaloni NFW).

Contributi Chiave

• Primo Campionamento del Posteriore Congiunto nello Spazio dei Pixel: Questo lavoro presenta il primo metodo capace di campionare il posteriore congiunto di sorgente e massa in primo piano come immagini pixelate per il lensing gravitazionale forte, andando oltre le stime puntuali o i fit parametrici a bassa dimensionalità.
• Integrazione di RIM e Diffusione: Il documento introduce una nuova architettura in cui una RIM è incorporata in un modello di diffusione condizionale per raffinare ricorsivamente gli score del posteriore, affrontando la non linearità e l'alta dimensionalità del problema del lensing.
• Modellazione al Livello del Rumore: Il metodo dimostra la capacità di modellare le osservazioni fino al livello del rumore, una capacità che spesso si perde quando si utilizzano modelli parametrici semplificati.

Risultati
Gli autori hanno addestrato e valutato DiRIM su due dataset: uno utilizzando mappe di convergenza da simulazioni cosmologiche e uno utilizzando profili analitici.

• Calibrazione: I test di Accuratezza con Punti Casuali (TARP) indicano che i campioni del posteriore sono quasi perfettamente calibrati, suggerendo che lo stimatore è non distorto.
• Goodness of Fit: La distribuzione dei valori $\chi^2$ per le immagini ricostruite corrisponde strettamente alla distribuzione teorica $\chi^2$, confermando che i modelli catturano i dati fino al livello del rumore.
• Gestione della Complessità: Il metodo modella con successo sistemi con distribuzioni di massa complesse, inclusi molteplici deflettori principali e subaloni. Nei casi in cui i modelli parametrici analitici non riuscivano a far aderire i dati al livello del rumore a causa della mancanza di flessibilità, DiRIM ha avuto successo.
• Rilevamento di Subaloni: Applicato a mappe analitiche contenenti subaloni NFW, il metodo ha rilevato i subaloni in oltre il 98% dei campioni del posteriore, raggiungendo una significatività di rilevamento $>3\sigma$.
• Campionamento delle Degenerazioni: Il metodo campiona efficacementamente le direzioni degenerate nello spazio dei parametri, come la correlazione tra la dimensione della sorgente e la pendenza radiale della mappa di convergenza.
• Robustezza (OOD): Il modello ha dimostrato robustezza verso i dati fuori distribuzione (OOD), inferendo con successo la morfologia delle sorgenti e delle mappe di convergenza che differivano significativamente dalla distribuzione di addestramento (ad esempio, galassie a forma di "7" generate da AI e mappe di convergenza sommate).
• Ablazione: Rimuovere la componente RIM (impostando le iterazioni $M=1$) ha causato un'inferenza distorta e residui con strutture significative, confermando che il raffinamento iterativo degli score è essenziale per il successo del metodo.

Significato e Rivendicazioni
Il documento afferma che DiRIM costituisce una "prima prova di concetto" riguardo alla capacità dei modelli generativi basati sulla diffusione combinati con le macchine di inferenza ricorrenti di risolvere problemi inversi ad alta dimensionalità e non lineari in astrofisica. Gli autori sottolineano che il loro metodo sfrutta rappresentazioni flessibili e ad alta dimensionalità per estrarre informazioni dai dati che sono inaccessibili alla modellazione parametrica tradizionale. Posizionano questo lavoro come un passo necessario verso la realizzazione del potenziale scientifico delle imminenti survey su larga scala (come LSST ed Euclid), che prevedono la scoperta di oltre 100.000 lenti gravitazionali forti. Gli autori notano con modestia che, sebbene il metodo sia efficace, rimangono piccole distorsioni nei sistemi complessi (ad esempio, molteplici deflettori, e il lavoro futuro affronterà questi aspetti ed esplorerà formulazioni nello spazio latente per ridurre i costi computazionali.