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

Technische Samenvatting: Gezamenlijke Inferentie van Zwakke Lenswerking Convergentiekaarten en Kosmologie met Diffusiemodellen

Probleemstelling
Huidige analyses van zwakke gravitationele lenswerking vertrouwen doorgaans op samenvattende statistieken (bijv. twee-punts correlatiefuncties), die niet-Gaussische informatie in de materieverdeling verliezen, met name op kleine, niet-lineaire schalen. Hoewel full-field benaderingen een principiële manier bieden om alle statistische informatie te vangen, kampen bestaande methoden met een dichotomie:

1. Marginale Inferentie: Neurale netwerkgebaseerde methoden leiden kosmologische parameters ($\theta$) direct af uit data door te marginaliseren over het onderliggende convergentieveld ($\kappa$), waardoor de reconstructie van de massakaart wordt verzuimd.
2. Veldreconstructie: Methoden zoals massakaart-reconstructie reconstrueren $\kappa$ uit ruisende shear ($\gamma_{obs}$), maar behandelen kosmologische parameters als vaststaand, in plaats van ze gezamenlijk af te leiden.
3. Expliciete Inferentie (MCMC): Benaderingen zoals BORG voeren gezamenlijke inferentie uit ($p(\theta, \kappa | \gamma_{obs})$) met behulp van Markov Chain Monte Carlo (MCMC). Deze vereisen echter een differentieerbaar voorwaarts model dat de begincondities koppelt aan de observeerbare grootheden, wat de complexiteit van de gebruikte fysica beperkt, en lijden onder hoge computationele kosten in hoog-dimensionale ruimtes.

De auteurs beogen een methode te ontwikkelen voor gezamenlijke inferentie die tegelijkertijd het convergentieveld reconstrueert en de kosmologische parameters afleidt, zonder een differentieerbaar voorwaarts model of iteratieve MCMC-sampling tijdens de inferentie te vereisen.

Methodologie: JADE
De auteurs introduceren JADE (Joint Architecture for fielD and cosmological parameter Estimation), een impliciet inferentiekader gebaseerd op conditionele diffusiemodellen.

• Impliciete Inferentie via Diffusiemodellen: In plaats van een likelihood-functie te leren, leert JADE de gezamenlijke posterieure distributie $p(\theta, \kappa | \gamma_{obs})$ direct van simulaties. De methode maakt gebruik van Flow Matching, een generatief framework waarbij een neuraal netwerk een snelheidveld leert om een basisdistributie (Gaussische ruis) naar de doeldistributie te transporteren. Het model is getraind om de "schone" sample ($x_1 = [\theta, \kappa]$) te voorspellen vanuit een ruisende interpolant ($x_t$) gegeven de geobserveerde shear-veld.
• Transformer Architectuur: De kern van JADE is een pixel-ruimte Vision Transformer aangepast voor multimodale data.
  • Verenigde Token-sequentie: Het convergentieveld $\kappa$ wordt gedecomponeerd in patches, en de kosmologische parameters $\theta$ worden geprojecteerd in embeddings van dezelfde dimensie. Deze worden samengevoegd tot een enkele sequentie van tokens.
  • Conditionering: Het geobserveerde ruisende shear-veld $\gamma_{obs}$ wordt eveneens gepatchificeerd en geprojecteerd, en vervolgens samengevoegd met de sequentie. Het aandachtmechanisme (attention mechanism) van de transformer stelt de doel-tokens (veld en parameters) in staat om bij alle ruimtelijke posities aandacht te besteden aan de conditionerende input (shear).
  • Output: Het netwerk geeft gelijktijdig gedempte (denoised) schattingen voor zowel de convergentiekaart als de kosmologische parameters.
• Geamortiseerde Inferentie: Eenmaal getraind kan het model posterieure samples genereren voor elke nieuwe observatie door een gewone differentiaalvergelijking (ODE) te integreren die geconditioneerd is op de specifieke $\gamma_{obs}$. Dit proces is "geamortiseerd", wat betekent dat er geen hertraining of per-observatie optimalisatie (zoals MCMC) nodig is.

Simulatieopzet
De methode werd gevalideerd op gesimuleerde zwakke lenswerking-data:

• Kosmologie: Een $w$CDM-model met zes parameters ($\Omega_c, \Omega_b, \sigma_8, h_0, n_s, w_0$).
• Data Generatie: Convergentiekaarten werden gegenereerd met behulp van een 2D log-normaal model (een fenomenologisch model voor niet-lineaire groei) in plaats van volledige N-body simulaties, om de snelle generatie van de grote trainingsdataset (100.000 samples) mogelijk te maken.
• Observeerbare grootheden: Gesimuleerde LSST Jaar 10 (Y10) diepte shear-velden met vormruis (shape noise), die een $5 \times 5$ deg$^2$ patch beslaan met $128 \times 128$ pixels.

Belangrijkste Resultaten

1. Gezamenlijke Reconstructie: Het model genereerde succesvol posterieure samples van zowel $\kappa$ als $\theta$, geconditioneerd op ruisende shear. Visuele inspectie en statistische tests bevestigden dat de gereconstrueerde convergentiekaarten de grondwaarheid-structuren herstellen, terwijl ze correct de resterende onzekerheid op kleine schalen reflecteren.
2. Statistische Consistentie:
   • Power Spectra: De auto- en cross-power spectra van de posterieure samples kwamen overeen met die van onafhankelijke simulaties getrokken bij de gesamplede kosmologieën, met een relatieve fout van ~6,7% gemiddeld over bins en schalen.
   • Eén-punt PDF's: De gereconstrueerde distributies vertoonden de verwachte log-normale vorm, wat bevestigt dat de niet-Gaussische statistiek is gevangen.
   • Residualen: De gereduceerde chi-kwadraat statistiek ($\chi^2_{red} \approx 1,002$) bevestigde dat de residualen tussen observaties en posterieure samples werden gedomineerd door geïnjecteerde vormruis, zoals verwacht.
3. Kosmologische Beperkingen: De marginale kosmologische posteriors ($p(\theta | \gamma_{obs})$) afgeleid van JADE vertoonden uitstekende overeenstemming met traditionele NUTS (No-U-Turn Sampler) MCMC-ketens wat betreft geloofwaardige regio's en degeneratierichtingen.
4. Kalibratie:
   • MIRA Score: De gezamenlijke posterieure bereikte een MIRA (Mass In Random Areas) kalibratiescore van $0,635 \pm 0,017$ (waarbij 0,667 optimaal is), wat wijst op goede kalibratie met milde overmoedigheid. De marginale kosmologische posterieure scoorde $0,659 \pm 0,0139$.
   • TARP: Tests of Accuracy with Random Points bevestigden dat de geloofwaardige regio's de correcte frequentistische dekking bereikten over de parameterruimte.
5. Efficiëntie: Het genereren van een enkele gezamenlijke posterieure sample duurt ongeveer 0,2 seconden (256 ODE-stappen). Het trekken van 1.000 samples duurt ~3 minuten, vergeleken met de aanzienlijk hogere computationele kosten van het draaien van MCMC-ketens (die ~2,18 miljoen gradiënt-evaluaties vereisten voor de benchmark).

Betekenis en Claims
Het artikel beweert vast te stellen dat diffusie-gebaseerde impliciete inferentie een levensvatbare route is naar volledige field-level kosmologische analyses. De primaire betekenis ligt in het verenigen van massakaart-reconstructie en kosmologische parameter-schatting binnen één enkel, niet-differentieerbaar framework.

Belangrijke claims met betrekking tot de voordelen van de methode zijn:

• Geen Differentieerbaar Voorwaarts Model Vereist: In tegen tegenoverkomend met MCMC-gebaseerde field-level inferentie, vereist JADE geen analytische gradiënten van het voorwaartse model, wat theoretisch het gebruik van willekeurig complexe simulators (bijv. volledige N-body of hydrodynamische codes) mogelijk maakt zodra de trainingsdata beschikbaar is.
• Amortisatie: Het model wordt één keer getraind en kan direct op elke nieuwe observatie worden toegepast, waardoor per-observatie sampling wordt vermeden.
• Gezamenlijke Inferentie: Het is de eerste methode die wordt gepresenteerd die de volledige gezamenlijke posterieure $p(\theta, \kappa | \gamma_{obs})$ afleidt met behulp van een enkele conditionele diffusie-model, in plaats van het veld te marginaliseren of parameters apart af te leiden.

Beperkingen en Toekomstig Werk
De auteurs merken op dat de huidige demonstratie steunt op log-normale simulaties, die sneller zijn dan N-body codes maar minder fysisch compleet. Het schalen naar realistische, dure simulators brengt uitdagingen met zich mee wat betreft de dichtheid van de trainingsset in de kosmologische parameterruimte. Daarnaast, als een geamortiseerde methode moet JADE opnieuw worden getraind als het datamodel (bijv. survey mask, ruismodel) verandert, in tegenstelling tot expliciete inferentiemethoden die zich op het moment van inferentie kunnen aanpassen. De auteurs stellen voor om deze architectuur in de toekomst uit te breiden naar de gezamenlijke inferentie van begincondities en kosmologische parameters, en om latente diffusie-benaderingen te verkennen om hogere-resolutie velden efficiënter te verwerken.