A plug-and-play approach with fast uncertainty quantification for weak lensing mass mapping
Paper introduces PnPMass, a plug-and-play framework for weak lensing mass mapping that reconciles reconstruction accuracy with practical deployment constraints of upcoming Stage-IV surveys. The key innovation is a carefully chosen data-fidelity operator that decouples denoiser training from observation-specific noise statistics, enabling a single trained model to handle varying survey conditions without retraining. Coupled with moment-network-based uncertainty quantification and conformal calibration, the method offers fast inference with coverage guarantees, addressing limitations of both end-to-end deep learning and costly MCMC sampling approaches.
The paper presents a methodologically sound and practically valuable solution to a critical bottleneck in weak lensing analysis. By deriving that the operator $\tilde{\mathbf{B}} := \tilde{\mathbf{A}}^{\top} \tilde{\mathbf{\Sigma}}^{-1/2}$ renders the training noise covariance approximately isotropic, the authors enable genuine plug-and-play capability across different noise realizations. The integration of fast uncertainty quantification via moment networks, calibrated with Conformalized Quantile Regression (CQR), provides non-asymptotic coverage guarantees without posterior sampling. While the theoretical convergence results rely on assumptions difficult to verify in practice, the empirical validation on $\kappa$TNG simulations demonstrates competitive accuracy with substantially improved computational efficiency.
The mathematical foundation for training the denoiser on white Gaussian noise independent of the observation covariance is rigorous and novel. The explicit derivation showing $\mathbf{\Sigma}_0 = \tau^2 \tilde{\mathbf{A}}^{\top}\tilde{\mathbf{A}} \approx \tau^2 \mathbf{I}$ justifies the algorithm's flexibility, while Proposition 1 provides concrete linear convergence guarantees under stated assumptions. The uncertainty quantification pipeline elegantly adapts moment networks to the PnP framework and applies CQR to achieve distribution-free calibration with marginal coverage guarantees. The empirical comparison with existing methods is fair and transparent, acknowledging the trade-offs between DeepMass's retraining requirements, DeepPosterior's sampling costs, and PnPMass's balance of accuracy and speed.
A critical gap exists between the theoretical requirement for a non-expansive denoiser (assumption H.2) and the actual implementation, where the authors explicitly state they ignore this constraint to avoid computational expense, leaving convergence in practice unverified. Additionally, the CQR calibration provides only marginal rather than conditional coverage, meaning the guarantee is marginalized over all possible outcomes rather than holding conditional on specific high-noise observations. The approach assumes perfect knowledge of the noise covariance matrix and tests only on a single cosmological model without robustness analysis to distributional shifts or model misspecification.
The evidence supports the core claim that PnPMass achieves reconstruction accuracy close to state-of-the-art deep learning methods while avoiding per-observation retraining. The experiments on $\kappa$TNG simulations compare correlation coefficients and power spectra against DeepMass, DeepPosterior, and classical methods under realistic noise and masking conditions. The comparison with DeepMass fairly reflects the practical limitation of end-to-end methods requiring retraining for new noise covariances, though the absolute performance depends on the retraining protocol. The calibration study validates that CQR achieves the target coverage rates empirically, supporting the reliability of the uncertainty estimates despite the Gaussian approximation in pre-estimation.
While the paper provides detailed algorithmic pseudocode (Algorithm 1), mathematical derivations for step-size bounds $\frac{1-\rho}{\lambda_{\min}} \leq \tau \leq \frac{1+\rho}{\lambda_{\max}}$, and architectural descriptions in the appendix, no code repository or software release is mentioned, posing a significant barrier to independent reproduction. The experiments use the public $\kappa$TNG dataset, but the training relies on specific hyperparameters including a noise standard deviation range of $[0, 0.2]$ and Wiener filtering for the residual variant, whose sensitivity to deviations is not fully characterized. The assumption that the noise covariance $\mathbf{\Sigma}$ is exactly known for constructing the forward operator may not hold in real survey conditions where galaxy number counts and ellipticity dispersions are themselves uncertain.
Upcoming stage-IV surveys such as Euclid and Rubin will deliver vast amounts of high-precision data, opening new opportunities to constrain cosmological models with unprecedented accuracy. A key step in this process is the reconstruction of the dark matter distribution from noisy weak lensing shear measurements. Current deep learning-based mass mapping methods achieve high reconstruction accuracy, but either require retraining a model for each new observed sky region (limiting practicality) or rely on slow MCMC sampling. Efficient exploitation of future survey data therefore calls for a new method that is accurate, flexible, and fast at inference. In addition, uncertainty quantification with coverage guarantees is essential for reliable cosmological parameter estimation. We introduce PnPMass, a plug-and-play approach for weak lensing mass mapping. The algorithm produces point estimates by alternating between a gradient descent step with a carefully chosen data fidelity term, and a denoising step implemented with a single deep learning model trained on simulated data corrupted by Gaussian white noise. We also propose a fast, sampling-free uncertainty quantification scheme based on moment networks, with calibrated error bars obtained through conformal prediction to ensure coverage guarantees. Finally, we benchmark PnPMass against both model-driven and data-driven mass mapping techniques. PnPMass achieves performance close to that of state-of-the-art deep-learning methods while offering fast inference (converging in just a few iterations) and requiring only a single training phase, independently of the noise covariance of the observations. It therefore combines flexibility, efficiency, and reconstruction accuracy, while delivering tighter error bars than existing approaches, making it well suited for upcoming weak lensing surveys.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.