📊 statistics

A plug-and-play approach with fast uncertainty quantification for weak lensing mass mapping

The paper introduces PnPMass, a flexible and efficient plug-and-play framework for weak-lensing mass mapping that combines a single trained deep-learning model with gradient descent for fast reconstruction and conformal prediction for calibrated uncertainty quantification, making it well-suited for upcoming large-scale cosmological surveys.

Original authors: Hubert Leterme, Andreas Tersenov, Jalal Fadili, Jean-Luc Starck

Published 2026-07-16

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Hubert Leterme, Andreas Tersenov, Jalal Fadili, Jean-Luc Starck

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine the universe as a giant, invisible web. Most of this web isn't made of the stars and planets we can see; it's made of "dark matter," a mysterious substance that we can't touch or see directly. We only know it's there because it acts like a cosmic lens, bending the light from distant galaxies as it travels toward us. When astronomers look at these bent shapes, they are trying to solve a giant, blurry puzzle: "Where is all this invisible mass hiding?" This process is called "mass mapping." It's like trying to figure out the shape of a hidden object by looking at how it distorts the reflection in a funhouse mirror. The problem is, the mirrors are dirty, the reflections are fuzzy, and the data is so massive that old ways of solving the puzzle are either too slow or require rebuilding the whole puzzle-solving machine every time the sky changes.

Enter a new approach called PnPMass, a clever method designed to map this invisible dark matter quickly and accurately, while also telling us exactly how much we can trust the result. Think of it as a smart, adaptable detective that doesn't need to be retrained every time it looks at a new crime scene. Instead of just guessing the location of the dark matter, PnPMass also puts up "confidence fences" around its guesses, showing us where the answer is rock-solid and where it's a bit shaky. This is crucial because if we want to understand how the universe began and how it will end, we need to know not just what the map looks like, but how sure we are about every single pixel on that map.

The Detective's New Toolkit

For a long time, scientists had two main ways to solve this dark matter puzzle, and both had a catch. The first method was like a super-fast, super-smart robot (called DeepMass) that could instantly guess the map. But this robot was rigid; if the "noise" on the mirror changed even a little—like if the telescope looked at a different part of the sky with different star density—the robot would get confused. You'd have to stop, retrain the robot from scratch, and start over. The second method was like a very thorough, but incredibly slow, detective (called DeepPosterior) who would run thousands of simulations to find the answer. It was flexible and didn't need retraining, but it took so long to run that it wasn't practical for the huge amounts of data coming from new telescopes like Euclid and Rubin.

The authors of this paper introduced PnPMass, a "plug-and-play" approach that tries to get the best of both worlds. Imagine a detective who has a single, highly trained "denoising" tool. This tool is taught to clean up images that are covered in static, like a TV screen with snow on it. The clever trick here is that the authors taught this tool using a specific kind of "white noise" (static) that is random and uniform. Then, they built a mathematical "cleaning loop" around it.

Here is how the loop works: The detective takes a blurry, noisy map of the sky. First, they take a step to make the map look more like the actual data they observed (a "forward step"). Then, they hand the result to their trained denoising tool to clean it up (a "backward step"). They repeat this process, back and forth, just a few times. Because the tool was trained on a universal type of noise, it doesn't matter if the actual sky data has a weird noise pattern; the math inside the loop adjusts for it automatically. This means the detective doesn't need to be retrained for every new sky region. It's like having a single pair of glasses that can automatically adjust their focus whether you are looking at a rainy window or a dusty one.

The "Confidence Fence"

One of the biggest hurdles in science is knowing how wrong you might be. If a method says "there is a dark matter clump here," how sure is it? Previous methods either didn't tell you, or they had to run thousands of slow simulations to guess the uncertainty.

PnPMass introduces a fast way to build "confidence fences" without the slow simulations. The authors used a technique called conformal prediction. Think of this like a quality control check. Before the detective goes out to solve the real mystery, they test their method on a bunch of fake, known puzzles (simulations). They see how often their "confidence fences" were too small or too big. Then, they use that data to mathematically adjust the size of the fences for the real job.

The result is a map where every single pixel has an error bar (a fence) that is statistically guaranteed to be correct a certain percentage of the time. In their tests, PnPMass produced tighter, more precise fences than the other methods, meaning it could pinpoint the dark matter with more confidence than the slow, thorough detective, and without the rigid limitations of the fast robot.

What the Tests Showed

The authors tested PnPMass using a massive dataset of simulated universes (called κTNG). They ran the algorithm for just a few iterations—often fewer than ten—and it converged to a very accurate answer.

Speed and Flexibility: Unlike the rigid robot, PnPMass didn't need to be retrained for different noise levels. It handled the messy, real-world data conditions effortlessly.
Accuracy: It performed almost as well as the best existing deep-learning methods, which usually require specific retraining for every new observation.
Uncertainty: The "confidence fences" it built were tighter than those from other methods. This suggests that PnPMass is better at distinguishing between real dark matter peaks and just random noise.

Interestingly, the authors found a trade-off. The method that was slightly less accurate at finding the exact shape of the dark matter peaks (PnPMass) actually produced smaller error bars than the method that was slightly more accurate at finding the peaks (DeepMass). They suspect this is because the more accurate method sometimes "hallucinates" peaks that aren't there, which makes it less confident overall. PnPMass, by being more conservative, avoids these false alarms, resulting in a more reliable map with tighter, trustworthy error bars.

The Bottom Line

This paper doesn't claim to have solved the mystery of dark matter forever. Instead, it offers a new, highly efficient tool for the upcoming era of massive sky surveys. By combining a flexible "plug-and-play" algorithm with a fast, mathematically guaranteed way to measure uncertainty, PnPMass allows astronomers to process the flood of data from future telescopes quickly and with a clear understanding of how much they can trust the results. It's a step toward turning the blurry, noisy pictures of our universe into a sharp, reliable map of the invisible web that holds it all together.

Technical Summary: PnPMass for Weak-Lensing Mass Mapping

Problem Statement
Upcoming stage-IV cosmological surveys (e.g., Euclid, Rubin) will generate vast amounts of high-precision weak-lensing shear data. A critical step in analyzing this data is reconstructing the dark matter distribution (convergence map, $\kappa$ ) from noisy shear measurements ( $\gamma$ ). While deep learning methods like DeepMass offer high accuracy and fast inference, they lack flexibility: they require retraining for every new observation because their performance is sensitive to varying noise covariances and masks. Conversely, flexible methods like DeepPosterior (which uses Markov chain Monte Carlo sampling) do not require retraining but suffer from computationally expensive inference. Furthermore, existing methods often lack rigorous uncertainty quantification (UQ) with coverage guarantees, which is essential for reliable cosmological parameter estimation.

Methodology: PnPMass
The authors introduce PnPMass, a plug-and-play (PnP) framework designed to balance accuracy, flexibility, and speed. The approach integrates a deep learning prior into a physical inverse problem solver.

Plug-and-Play Framework:
- The mass-mapping problem is formulated as a real-valued inverse problem: $\tilde{\gamma} = \tilde{A}\kappa + \tilde{n}$ , where $\tilde{A}$ is a linear operator derived from the Fourier transform and the noise covariance $\tilde{\Sigma}$ .
- The algorithm alternates between a forward step (gradient descent with a data fidelity term) and a backward step (denoising using a deep neural network $D_{\hat{\theta}}$ ).
- Key Innovation in Training: Unlike standard PnP approaches that might require specific noise training, the authors prove that by choosing a specific forward operator $\tilde{B} = \tilde{A}^\top \tilde{\Sigma}^{-1/2}$ , the residual noise seen by the denoiser becomes approximately white Gaussian noise with a standard deviation equal to the step size $\tau$ .
- Noise-Level Awareness: To ensure flexibility across different survey conditions, the denoiser is trained on data corrupted by white Gaussian noise with standard deviations randomly sampled from a wide interval $[\tau_{min}, \tau_{max}]$ . The noise level $\sigma$ is provided as an additional input channel to the network during inference. This allows a single trained model to handle varying noise covariances without retraining.
Uncertainty Quantification (UQ):
- The paper proposes a sampling-free UQ scheme. Instead of MCMC, it utilizes moment networks (adapted from Jeffrey & Wandelt, 2020) to estimate the posterior variance.
- An order-2 moment network $G_{\hat{\omega}}$ is trained to predict the squared error of the point estimate.
- Conformal Prediction: To ensure statistical rigor, the authors apply Conformalized Quantile Regression (CQR). This post-processing step calibrates the raw error bars derived from the moment network, providing per-pixel error intervals with non-asymptotic, marginal coverage guarantees (i.e., the true value falls within the interval with a specified probability, regardless of the underlying data distribution).
Residual Variant:
- A variant of PnPMass incorporates prior knowledge that convergence maps are a superposition of a Gaussian field and a non-Gaussian residual. It first estimates the Gaussian component via Wiener filtering and then applies PnPMass to reconstruct the non-Gaussian residual, aiming for higher accuracy in peak regions.

Key Contributions

Flexibility without Retraining: PnPMass requires only a single training phase, regardless of the observation's noise covariance or mask, unlike end-to-end methods like DeepMass.
Fast Inference: The method converges in a few iterations (typically 8), avoiding the computational cost of sampling-based methods like DeepPosterior.
Rigorous UQ: It provides a fast, sampling-free uncertainty estimation scheme with provable coverage guarantees via conformal prediction, a feature often missing or computationally prohibitive in deep learning mass mapping.
Theoretical Insight: The paper establishes that training a denoiser on white Gaussian noise is sufficient for solving inverse problems with arbitrary noise covariances, provided the forward operator is correctly whitened.

Results
Experiments were conducted on the $\kappa$ TNG simulated dataset (mock convergence maps based on $\Lambda$ CDM cosmology).

Reconstruction Accuracy: PnPMass achieves reconstruction accuracy (measured by normalized RMSE) comparable to the state-of-the-art DeepMass method and superior to classical methods (Wiener, MCALens). While DeepMass performs slightly better (likely due to fine-tuning to specific noise conditions), PnPMass offers a superior trade-off between accuracy and flexibility.
Uncertainty Quantification: After calibration, PnPMass produces tighter error bars than both model-driven methods (Wiener, MCALens) and the data-driven DeepMass method, while maintaining the target miscoverage rate (e.g., ~4.55% for 2 $\sigma$ confidence).
Convergence: The algorithm converges rapidly, with reconstruction accuracy stabilizing after a few iterations.

Significance and Claims
The paper claims that PnPMass is well-suited for upcoming weak-lensing surveys because it combines the flexibility of Bayesian methods (no retraining for new masks/noise) with the efficiency of deep learning (fast inference). By delivering accurate reconstructions with statistically calibrated uncertainty estimates, it addresses the critical need for reliable cosmological parameter estimation in the era of big data astronomy. The authors position PnPMass as a practical solution that overcomes the limitations of current deep learning approaches (lack of flexibility) and sampling-based approaches (computational cost), making it a viable candidate for processing data from Euclid and Rubin.

The Detective's New Toolkit

The "Confidence Fence"

What the Tests Showed

The Bottom Line

Technical Summary: PnPMass for Weak-Lensing Mass Mapping

More like this