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

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

Problem Statement
Weak gravitational lensing provides a direct probe of the cosmic large-scale structure, encoding non-Gaussian information in convergence (mass) maps that is inaccessible to traditional two-point statistics. Accurate forward modeling of these maps is essential for cosmological inference, typically relying on computationally expensive $N$-body simulations. While deep learning emulators, particularly Generative Adversarial Networks (GANs), have been proposed to accelerate this process, existing methods suffer from significant limitations. Previous GAN-based surrogates often fail to match summary statistics with sufficient precision (typically 5–10% errors), produce over-smoothed maps, struggle to capture the correct distribution of generated map sets (failing to reproduce the full probabilistic distribution rather than just the "mode"), and exhibit training instability. Furthermore, conditional generation across continuous cosmological parameters (e.g., matter density $\Omega_m$ and clustering amplitude $\sigma_8$) remains challenging due to heterogeneous learning difficulties across the label space.

Methodology
The authors propose a Residual Label-Conditional Flow Matching (RLCFM) framework to generate high-fidelity convergence maps conditioned on cosmological parameters. The approach builds upon Continuous Normalizing Flows (CNF) and Flow Matching (FM), which learn a time-dependent vector field to transport a prior distribution to a target data distribution.

Key methodological components include:

1. Residual Label-Conditional Flow Matching: Instead of training a standard label-conditional flow, the authors introduce a reparameterization in a residual space. They subtract the label-specific mean ($\mu_0(y)$) from the target maps before training. This removes trivial label-dependent mean shifts, allowing the network to focus on learning the transferable residual structure. This design aims to stabilize interpolation across cosmological parameters.
2. Label-Specific Gaussian Prior: The model employs a data-driven, label-specific Gaussian prior $p(x|y)$ rather than a single shared prior. The prior's mean and variance are derived from the statistics of the conditional distribution $q(x|y)$ at $t=0$, transporting this specific prior to the target distribution at $t=1$.
3. Training Objective: The network is trained to minimize the Flow Matching loss, regressing the time-dependent vector field that transports the label-specific prior to the target convergence maps. The training is performed on a dataset of 57 cosmological models ($\Omega_m, \sigma_8$) within a flat $\Lambda$CDM framework, comprising 684,000 $128 \times 128$ pixel maps.
4. Data Preprocessing: Convergence maps are transformed via a logarithm (approximating a lognormal field) and standardized. The model learns to generate the residual field $r = \tilde{x} - \mu_0(y)$, which is then reconstructed at inference by adding the label-specific mean back.

Key Contributions

• Novel Architecture: Introduction of the RLCFM architecture, which combines label-specific priors with residual-space training to address the instability and interpolation challenges of conditional generative modeling in cosmology.
• Improved Fidelity: The method achieves significantly higher precision than previous GAN benchmarks. It reduces typical errors in basic statistics (pixel/peak distributions, power spectrum) from $<10\%$ to $<1\%$, and in higher-order statistics (bispectrum, Minkowski functionals) from $<20\%$ to $<5\%$.
• Distribution-Level Accuracy: Unlike GANs which often fit the "mode" of the data, the RLCFM model successfully reproduces the full distribution of summary statistics generated from random initial conditions, matching the variability of $N$-body simulations.
• Robust Interpolation: The model demonstrates stable performance across the cosmological parameter space, including at held-out test cosmologies and grid edges, with errors increasing only mildly at boundaries compared to the center.

Results
The authors evaluated the RLCFM model against $N$-body simulations using a comprehensive suite of metrics:

• Visual Quality: Generated maps are visually indistinguishable from $N$-body references, with consistent morphological evolution as $\Omega_m$ and $\sigma_8$ vary.
• Summary Statistics:
  • Pixel and Peak Distributions: Normalized Wasserstein-1 distances are typically $<0.01$.
  • Power Spectrum ($C_\ell$): Fractional differences are typically $<1\%$.
  • Bispectrum ($B_\ell$): Fractional differences are typically $<5\%$.
  • Minkowski Functionals ($V_0, V_1, V_2$): Fractional differences are generally below 3%.
  • Correlation Matrices: The Frobenius norm contrast between RLCFM and $N$-body correlation matrices is typically $<3\%$.
  • Structural Similarity (MS-SSIM): The significance of the difference in structural similarity scores is low, indicating the model does not suffer from mode collapse and preserves multi-scale structural realism.
• Efficiency: The model was trained in approximately 7 hours (65,000 steps) on a single NVIDIA A100 GPU, requiring significantly fewer epochs to reach stable performance compared to GAN baselines.

Significance and Claims
The paper claims that conditional flow matching, when augmented with a residual-space formulation and a data-driven prior, provides a powerful and efficient framework for generating cosmological mass maps. The primary significance lies in the practical usability of the emulator: it captures the cosmological signal with high fidelity across both low-order and high-order statistics while supporting multiple forms of data analysis. The authors assert that this work brings the field closer to a practical mass map emulator capable of supporting forward modeling for upcoming wide-field surveys (e.g., Euclid, LSST) without the computational cost of running full $N$-body simulations for every cosmological parameter point. The method is presented as a step toward fully generative, simulation-based cosmological inference pipelines.