Generative modelling for mass-mapping with fast uncertainty quantification

TL;DR

Mass-mapping from weak lensing is an ill-posed problem requiring both high-fidelity reconstructions and reliable uncertainty quantification. MMGAN introduces a regularised conditional Wasserstein GAN to generate approximate posterior samples of the convergence field $\kappa$ from shear data $\gamma$, enabling rapid uncertainty maps for downstream cosmology. Trained on 10,000 mock COSMOS-like maps from $\kappa$TNG and applied to real COSMOS data, MMGAN achieves reconstruction quality competitive with state-of-the-art methods while producing samples in under a second per posterior draw. The approach yields pixel-wise uncertainties that correlate with reconstruction errors and is designed for integration into large cosmology pipelines, with code and data publicly available.

Abstract

Understanding the nature of dark matter in the Universe is an important goal of modern cosmology. A key method for probing this distribution is via weak gravitational lensing mass-mapping - a challenging ill-posed inverse problem where one infers the convergence field from observed shear measurements. Upcoming stage IV surveys, such as those made by the Vera C. Rubin Observatory and Euclid satellite, will provide a greater quantity and precision of data for lensing analyses, necessitating high-fidelity mass-mapping methods that are computationally efficient and that also provide uncertainties for integration into downstream cosmological analyses. In this work we introduce MMGAN, a novel mass-mapping method based on a regularised conditional generative adversarial network (GAN) framework, which generates approximate posterior samples of the convergence field given shear data. We adopt Wasserstein GANs to improve training stability and apply regularisation techniques to overcome mode collapse, issues that otherwise are particularly acute for conditional GANs. We train and validate our model on a mock COSMOS-style dataset before applying it to true COSMOS survey data. Our approach significantly outperforms the Kaiser-Squires technique and achieves similar reconstruction fidelity as alternative state-of-the-art deep learning approaches. Notably, while alternative approaches for generating samples from a learned posterior are slow (e.g. requiring $\sim$10 GPU minutes per posterior sample), MMGAN can produce a high-quality convergence sample in less than a second.

Figures (8)

• Figure 1: Illustration of the architecture of the MMGAN generator. The shear map, comprising real and imaginary components, is used to produce a pseudo-reconstruction, which is similarly decomposed into its real and imaginary parts. These components are concatenated with a two-channel random latent vector $z$, and subsequently passed through our U-Net generator, which outputs a single sample of the convergence from the learned posterior distribution. The numbers below each block indicate the number of channels in each layer. Additionally, the colour of the blocks indicate the series of operations applied, as dictated by the legend. The residual block has been illustrated in greater detail below the main generator architecture. It consists of a $3\times3$ convolution, followed by batch normalisation, then parametric ReLU. This is done twice, and then the output of this is added to a $1\times1$ convolution of the original input.
• Figure 2: PSNR and Pearson correlation coefficient values of MMGAN reconstruction dependant on the number of approximate posterior samples used to create that reconstruction, which is given by the mean of the approximate posterior samples. The curve flattens out for both metrics, indicating there is little need to consider $N > 32$.
• Figure 3: A reconstructed convergence map for one of the mock COSMOS maps. Our reconstruction is the average over 32 approximate posterior samples. On the bottom row is the pixel-wise absolute error between the reconstruction and the ground truth, and the standard deviation between the 32 samples used to build the reconstruction. The white contour indicates the outer border of the mask applied to the data. We achieve superior visual quality as compared to the Kaiser-Squires reconstruction, with no peak suppression. Additionally, we see visual correlation between the absolute error and the standard deviation map.
• Figure 4: A selection of generated approximate posterior samples for a given shear map, in comparison with the ground truth. We have zoomed in on a region of the samples, to better show the variation within different samples.
• Figure 5: Demonstration of how the $N$-sample reconstruction varies for $N \in \{1,4,8,32\}$, $N=1$ being a single posterior sample, for the zoomed-in region shown in the red box. The figure also shows the Kaiser-Squires map for the same region. As can be seen, the reconstruction becomes smoother as $N$ increases, however, the prominent features remain. An individual sample has a far higher level of detail, comparable with the true map, however, it can be seen that features differ slightly to the truth, indicating why it is necessary to average over a number of samples. Despite some loss of the smallest-scale structure for $N=32$, there is less peak suppression than the Kaiser-Squires reconstruction.