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

Technische Zusammenfassung: Erhöhung der Präzision von Surrogatmodellen für Weak-Lensing-Massenkarten mittels Flow Matching

Problemstellung
Die schwache Gravitationslinsung bietet einen direkten Zugang zur großräumigen Struktur des Universums und kodiert nicht-gaußsche Informationen in Konvergenz-(Massen-)Karten, die für traditionelle Zwei-Punkt-Statistiken unzugänglich sind. Eine genaue Vorwärtsmodellierung dieser Karten ist für kosmologische Schlussfolgerungen unerlässlich und stützt sich typischerweise auf rechenintensive $N$-Körper-Simulationen. Obwohl Deep-Learning-Emulatoren, insbesondere Generative Adversarial Networks (GANs), vorgeschlagen wurden, um diesen Prozess zu beschleunigen, leiden bestehende Methoden unter erheblichen Einschränkungen. Frühere GAN-basierte Surrogate erfüllen Summary-Statistiken oft nicht mit ausreichender Präzision (typischerweise Fehler von 5–10 %), erzeugen übermäßig geglättete Karten, haben Schwierigkeiten, die korrekte Verteilung generierter Kartensätze zu erfassen (sie reproduzieren nicht die gesamte probabilistische Verteilung, sondern nur den „Modus"), und zeigen Trainingsinstabilität. Darüber hinaus bleibt die bedingte Generierung über kontinuierliche kosmologische Parameter hinweg (z. B. Materiedichte $\Omega_m$ und Clustering-Amplitude $\sigma_8$) aufgrund heterogener Lernschwierigkeiten im Label-Raum eine Herausforderung.

Methodik
Die Autoren schlagen ein Residual Label-Conditional Flow Matching (RLCFM)-Rahmenwerk vor, um hochfidele Konvergenzkarten zu generieren, die an kosmologische Parameter konditioniert sind. Der Ansatz baut auf Continuous Normalizing Flows (CNF) und Flow Matching (FM) auf, die ein zeitabhängiges Vektorfeld erlernen, um eine Prior-Verteilung in eine Ziel-Datenverteilung zu transportieren.

Zu den wesentlichen methodischen Komponenten gehören:

1. Residual Label-Conditional Flow Matching: Anstatt ein Standard-Label-Conditional-Flow zu trainieren, führen die Autoren eine Reparametrisierung in einem Residualraum ein. Sie subtrahieren den label-spezifischen Mittelwert ($\mu_0(y)$) von den Zielkarten vor dem Training. Dies entfernt triviale label-abhängige Mittelwertverschiebungen und ermöglicht es dem Netzwerk, sich auf das Erlernen der übertragbaren residualen Struktur zu konzentrieren. Dieses Design zielt darauf ab, die Interpolation über kosmologische Parameter hinweg zu stabilisieren.
2. Label-spezifische Gaußsche Prior: Das Modell verwendet eine datengesteuerte, label-spezifische Gaußsche Prior $p(x|y)$ anstelle einer einzigen gemeinsamen Prior. Der Mittelwert und die Varianz der Prior werden aus den Statistiken der bedingten Verteilung $q(x|y)$ bei $t=0$ abgeleitet und transportieren diese spezifische Prior zur Zielverteilung bei $t=1$.
3. Trainingsziel: Das Netzwerk wird trainiert, um den Flow-Matching-Verlust zu minimieren, indem es das zeitabhängige Vektorfeld regressiert, das die label-spezifische Prior in die Ziel-Konvergenzkarten transportiert. Das Training erfolgt auf einem Datensatz von 57 kosmologischen Modellen ($\Omega_m, \sigma_8$) innerhalb eines flachen $\Lambda$CDM-Rahmenwerks, der 684.000 $128 \times 128$ Pixel-Karten umfasst.
4. Datenvorverarbeitung: Konvergenzkarten werden mittels eines Logarithmus transformiert (Annäherung an ein lognormales Feld) und standardisiert. Das Modell lernt, das Residualfeld $r = \tilde{x} - \mu_0(y)$ zu generieren, das bei der Inferenz durch Hinzufügen des label-spezifischen Mittelwerts wieder rekonstruiert wird.

Hauptbeiträge

• Neuartige Architektur: Einführung der RLCFM-Architektur, die label-spezifische Priors mit Training im Residualraum kombiniert, um die Instabilität und Interpolationsherausforderungen der bedingten generativen Modellierung in der Kosmologie zu adressieren.
• Verbesserte Fidelity: Die Methode erreicht eine deutlich höhere Präzision als frühere GAN-Benchmarks. Sie reduziert typische Fehler in Basisstatistiken (Pixel-/Peak-Verteilungen, Leistungsspektrum) von $<10\%$ auf $<1\%$ und in höherordentlichen Statistiken (Bispektrum, Minkowski-Funktionale) von $<20\%$ auf $<5\%$.
• Genauigkeit auf Verteilungsebene: Im Gegensatz zu GANs, die oft nur den „Modus" der Daten anpassen, reproduziert das RLCFM-Modell erfolgreich die gesamte Verteilung von Summary-Statistiken, die aus zufälligen Anfangsbedingungen generiert wurden, und passt dabei die Variabilität von $N$-Körper-Simulationen an.
• Robuste Interpolation: Das Modell zeigt eine stabile Leistung über den Raum der kosmologischen Parameter hinweg, einschließlich bei zurückgehaltenen Test-Kosmologien und an den Rändern des Gitters, wobei die Fehler an den Grenzen im Vergleich zum Zentrum nur leicht ansteigen.

Ergebnisse
Die Autoren bewerteten das RLCFM-Modell gegenüber $N$-Körper-Simulationen unter Verwendung einer umfassenden Suite von Metriken:

• Visuelle Qualität: Generierte Karten sind visuell von $N$-Körper-Referenzen nicht zu unterscheiden, mit einer konsistenten morphologischen Evolution bei Variation von $\Omega_m$ und $\sigma_8$.
• Summary-Statistiken:
  • Pixel- und Peak-Verteilungen: Normalisierte Wasserstein-1-Abstände liegen typischerweise bei $<0,01$.
  • Leistungsspektrum ($C_\ell$): Fraktionale Unterschiede liegen typischerweise bei $<1\%$.
  • Bispektrum ($B_\ell$): Fraktionale Unterschiede liegen typischerweise bei $<5\%$.
  • Minkowski-Funktionale ($V_0, V_1, V_2$): Fraktionale Unterschiede liegen allgemein unter 3 %.
  • Korrelationsmatrizen: Der Frobenius-Norm-Kontrast zwischen RLCFM- und $N$-Körper-Korrelationsmatrizen liegt typischerweise bei $<3\%$.
  • Strukturelle Ähnlichkeit (MS-SSIM): Die Signifikanz des Unterschieds in den strukturellen Ähnlichkeitswerten ist gering, was darauf hindeutet, dass das Modell nicht unter Mode-Collapse leidet und die mehrskalige strukturelle Realität bewahrt.
• Effizienz: Das Modell wurde in etwa 7 Stunden (65.000 Schritte) auf einer einzelnen NVIDIA A100 GPU trainiert und benötigte im Vergleich zu GAN-Baselines signifikant weniger Epochen, um eine stabile Leistung zu erreichen.

Bedeutung und Behauptungen
Die Arbeit behauptet, dass bedingtes Flow Matching, ergänzt durch eine Formulierung im Residualraum und eine datengesteuerte Prior, ein leistungsfähiges und effizientes Rahmenwerk zur Generierung kosmologischer Massenkarten bietet. Die primäre Bedeutung liegt in der praktischen Nutzbarkeit des Emulators: Er erfasst das kosmologische Signal mit hoher Fidelity über sowohl niedrigordentliche als auch höherordentliche Statistiken hinweg und unterstützt gleichzeitig verschiedene Formen der Datenanalyse. Die Autoren betonen, dass diese Arbeit das Feld einem praktischen Massenkarten-Emulator näherbringt, der die Vorwärtsmodellierung für kommende Weitfeld-Übersichtssurveys (z. B. Euclid, LSST) unterstützen kann, ohne die Rechenkosten vollständiger $N$-Körper-Simulationen für jeden kosmologischen Parameterpunkt zu verursachen. Die Methode wird als Schritt hin zu vollständig generativen, simulationsbasierten kosmologischen Inferenzpipelines präsentiert.