Increasing the Precision of Surrogate Models for Weak Lensing Mass Maps with Flow Matching

Technische Samenvatting: Het Verbeteren van de Precisie van Surrogaatmodellen voor Zwakke Lensing Massakaarten met Flow Matching

Probleemstelling
Zwakke gravitationele lensing biedt een directe sonde voor de kosmische grootte-schaalstructuur, waarbij niet-Gaussische informatie wordt gecodeerd in convergentie (massa)kaarten die ontoegankelijk zijn voor traditionele tweepuntsstatistieken. Nauwkeurige forward modeling van deze kaarten is essentieel voor kosmologische inferentie, waarbij doorgaans wordt vertrouwd op rekenkundig kostbare $N$-lichaamssimulaties. Hoewel diep leernetwerk-emulators, met name Generative Adversarial Networks (GAN's), zijn voorgesteld om dit proces te versnellen, lijden bestaande methoden aan aanzienlijke beperkingen. Vorige op GAN's gebaseerde surrogaten slagen er vaak niet in om samenvattende statistieken met voldoende precisie te benaderen (typisch fouten van 5–10%), produceren over-gesmoorde kaarten, worstelen met het vastleggen van de juiste verdeling van gegenereerde kaartsets (het niet reproduceren van de volledige probabilistische verdeling, maar slechts de "modus"), en vertonen trainingsinstabiliteit. Bovendien blijft conditionele generatie over continue kosmologische parameters (zoals materie-dichtheid $\Omega_m$ en clustering-amplitude $\sigma_8$) uitdagend vanwege heterogene leermoeilijkheden over de labelruimte.

Methodologie
De auteurs stellen een Residual Label-Conditional Flow Matching (RLCFM)-raamwerk voor om hoogtrouwheidsconvergentiekaarten te genereren die zijn geconditioneerd op kosmologische parameters. De aanpak bouwt voort op Continuous Normalizing Flows (CNF) en Flow Matching (FM), die een tijdsafhankelijk vectorveld leren om een a-prioriverdeling naar een doelpopulatieverdeling te transporteren.

Belangrijke methodologische componenten zijn:

1. Residual Label-Conditional Flow Matching: In plaats van een standaard label-conditional flow te trainen, introduceren de auteurs een herparameterisatie in een residuele ruimte. Ze trekken het labelspecifieke gemiddelde ($\mu_0(y)$) af van de doelkaarten voordat ze trainen. Dit verwijdert triviale label-afhankelijke gemiddelde verschuivingen, waardoor het netwerk zich kan richten op het leren van de overdraagbare residuale structuur. Dit ontwerp beoogt interpolatie over kosmologische parameters te stabiliseren.
2. Label-specifieke Gaussische A-prioriverdeling: Het model maakt gebruik van een datagedreven, label-specifieke Gaussische a-prioriverdeling $p(x|y)$ in plaats van één gedeelde a-prioriverdeling. Het gemiddelde en de variantie van de a-prioriverdeling worden afgeleid uit de statistieken van de conditionele verdeling $q(x|y)$ bij $t=0$, waarbij deze specifieke a-prioriverdeling wordt getransporteerd naar de doelpopulatieverdeling bij $t=1$.
3. Trainingsdoel: Het netwerk wordt getraind om de Flow Matching-loss te minimaliseren, door het tijdsafhankelijke vectorveld te regresseren dat de label-specifieke a-prioriverdeling transporteert naar de doelconvergentiekaarten. De training wordt uitgevoerd op een dataset van 57 kosmologische modellen ($\Omega_m, \sigma_8$) binnen een plat $\Lambda$CDM-raamwerk, bestaande uit 684.000 kaarten van $128 \times 128$ pixels.
4. Data-voorverwerking: Convergentiekaarten worden getransformeerd via een logaritme (benaderend een lognormaal veld) en gestandaardiseerd. Het model leert het residuale veld te genereren $r = \tilde{x} - \mu_0(y)$, dat bij inferentie wordt gereconstrueerd door het labelspecifieke gemiddelde weer toe te voegen.

Belangrijkste Bijdragen

• Nieuwe Architectuur: Introductie van de RLCFM-architectuur, die labelspecifieke a-prioriverdelingen combineert met training in residuale ruimte om de instabiliteit en interpolatie-uitdagingen van conditionele generatieve modellering in de kosmologie aan te pakken.
• Verbeterde Fidelity: De methode bereikt aanzienlijk hogere precisie dan eerdere GAN-benchmarks. Het reduceert typische fouten in basisstatistieken (pixel-/piekverdelingen, powerspectrum) van $<10\%$ naar $<1\%$, en in hogere-orde statistieken (bispectrum, Minkowski-functionalen) van $<20\%$ naar $<5\%$.
• Nauwkeurigheid op Verdelingsniveau: In tegenstelling tot GAN's die vaak de "modus" van de data benaderen, reproduceert het RLCFM-model succesvol de volledige verdeling van samenvattende statistieken gegenereerd uit willekeurige beginvoorwaarden, en komt overeen met de variabiliteit van $N$-lichaamssimulaties.
• Robuuste Interpolatie: Het model vertoont stabiele prestaties over de ruimte van kosmologische parameters, inclusief bij achtergehouden testkosmologieën en randen van het rooster, waarbij fouten aan de randen slechts licht toenemen in vergelijking met het centrum.

Resultaten
De auteurs evalueerden het RLCFM-model tegen $N$-lichaamssimulaties met behulp van een uitgebreide suite van metrieken:

• Visuele Kwaliteit: Gegenereerde kaarten zijn visueel niet te onderscheiden van $N$-lichaamreferenties, met consistente morfologische evolutie naarmate $\Omega_m$ en $\sigma_8$ variëren.
• Samenvattende Statistieken:
  • Pixel- en Piekgewichten: Genormaliseerde Wasserstein-1-afstanden zijn typisch $<0,01$.
  • Powerspectrum ($C_\ell$): Fractieve verschillen zijn typisch $<1\%$.
  • Bispectrum ($B_\ell$): Fractieve verschillen zijn typisch $<5\%$.
  • Minkowski-functionalen ($V_0, V_1, V_2$): Fractieve verschillen zijn over het algemeen onder de 3%.
  • Correlatiematrixen: Het Frobenius-normcontrast tussen RLCFM- en $N$-lichaamcorrelatiematrixen is typisch $<3\%$.
  • Structurele Gelijkenis (MS-SSIM): De significantie van het verschil in structurele gelijkenisscores is laag, wat aangeeft dat het model niet lijdt aan modusinstorting en multi-schaal structurele realisme behoudt.
• Efficiëntie: Het model werd getraind in ongeveer 7 uur (65.000 stappen) op een enkele NVIDIA A100 GPU, waarbij aanzienlijk minder epochs nodig waren om stabiele prestaties te bereiken in vergelijking met GAN-baselines.

Betekenis en Beweringen
Het artikel beweert dat conditionele flow matching, wanneer verrijkt met een formulering in residuale ruimte en een datagedreven a-prioriverdeling, een krachtig en efficiënt raamwerk biedt voor het genereren van kosmologische massa-kaarten. De primaire betekenis ligt in het praktische bruikbaarheid van de emulator: het vangt het kosmologische signaal met hoge fideliteit op zowel laag-orde als hoog-orde statistieken, terwijl het meerdere vormen van data-analyse ondersteunt. De auteurs stellen dat dit werk het vakgebied dichter bij een praktische massa-kaartemulator brengt die forward modeling kan ondersteunen voor aankomende wide-field surveys (zoals Euclid, LSST) zonder de rekenkosten van het uitvoeren van volledige $N$-lichaamssimulaties voor elk punt van de kosmologische parameter. De methode wordt gepresenteerd als een stap naar volledig generatieve, op simulaties gebaseerde kosmologische inferentiepijplijnen.