Joint inference of weak lensing convergence map and cosmology with diffusion models

Technical Summary: Joint Inference of Weak Lensing Convergence Maps and Cosmology with Diffusion Models

Problem Statement
Current weak gravitational lensing analyses typically rely on summary statistics (e.g., two-point correlation functions) which discard non-Gaussian information present in the matter distribution, particularly on small, non-linear scales. While full-field approaches offer a principled way to capture all statistical information, existing methods face a dichotomy:

1. Marginal Inference: Neural network-based methods infer cosmological parameters ($\theta$) directly from data but marginalize over the underlying convergence field ($\kappa$), failing to reconstruct the mass map.
2. Field Reconstruction: Methods like mass mapping reconstruct $\kappa$ from noisy shear ($\gamma_{obs}$) but treat cosmological parameters as fixed, rather than jointly inferring them.
3. Explicit Inference (MCMC): Approaches like BORG perform joint inference ($p(\theta, \kappa | \gamma_{obs})$) using Markov Chain Monte Carlo (MCMC). However, these require a differentiable forward model linking initial conditions to observables, limiting the complexity of physics that can be employed, and suffer from high computational costs in high-dimensional spaces.

The authors aim to develop a method for joint inference that simultaneously reconstructs the convergence field and infers cosmological parameters without requiring a differentiable forward model or iterative MCMC sampling at inference time.

Methodology: JADE
The authors introduce JADE (Joint Architecture for fielD and cosmological parameter Estimation), an implicit inference framework based on conditional diffusion models.

• Implicit Inference via Diffusion Models: Instead of learning a likelihood function, JADE learns the joint posterior distribution $p(\theta, \kappa | \gamma_{obs})$ directly from simulations. The method utilizes Flow Matching, a generative framework where a neural network learns a velocity field to transport a base distribution (Gaussian noise) to the target distribution. The model is trained to predict the "clean" sample ($x_1 = [\theta, \kappa]$) from a noisy interpolant ($x_t$) given the observed shear field.
• Transformer Architecture: The core of JADE is a pixel-space Vision Transformer adapted for multimodal data.
  • Unified Token Sequence: The convergence field $\kappa$ is decomposed into patches, and the cosmological parameters $\theta$ are projected into embeddings of the same dimension. These are concatenated into a single sequence of tokens.
  • Conditioning: The observed noisy shear field $\gamma_{obs}$ is also patchified and projected, then concatenated to the sequence. The transformer's attention mechanism allows the target tokens (field and parameters) to attend to the conditioning input (shear) at all spatial positions.
  • Output: The network outputs denoised estimates for both the convergence map and the cosmological parameters simultaneously.
• Amortized Inference: Once trained, the model can generate posterior samples for any new observation by integrating an Ordinary Differential Equation (ODE) conditioned on the specific $\gamma_{obs}$. This process is "amortized," meaning no retraining or per-observation optimization (like MCMC) is required.

Simulation Setup
The method was validated on simulated weak lensing data:

• Cosmology: A $w$CDM model with six parameters ($\Omega_c, \Omega_b, \sigma_8, h_0, n_s, w_0$).
• Data Generation: Convergence maps were generated using a 2D log-normal model (phenomenological model for non-linear growth) rather than full N-body simulations, to allow for rapid generation of the large training dataset (100,000 samples).
• Observables: Simulated LSST Year 10 (Y10) depth shear fields with shape noise, covering a $5 \times 5$ deg$^2$ patch with $128 \times 128$ pixels.

Key Results

1. Joint Reconstruction: The model successfully generated posterior samples of both $\kappa$ and $\theta$ conditioned on noisy shear. Visual inspection and statistical tests confirmed that the reconstructed convergence maps recover ground-truth structures while correctly reflecting residual uncertainty on small scales.
2. Statistical Consistency:
   • Power Spectra: The auto- and cross-power spectra of the posterior samples matched those of independent simulations drawn at the sampled cosmologies, with a relative error of ~6.7% averaged over bins and scales.
   • One-Point PDFs: The recovered distributions exhibited the expected log-normal shape, confirming the capture of non-Gaussian statistics.
   • Residuals: The reduced chi-squared statistic ($\chi^2_{red} \approx 1.002$) confirmed that residuals between observations and posterior samples were dominated by injected shape noise, as expected.
3. Cosmological Constraints: The marginal cosmological posteriors ($p(\theta | \gamma_{obs})$) derived from JADE showed excellent agreement with traditional NUTS (No-U-Turn Sampler) MCMC chains in terms of credible regions and degeneracy directions.
4. Calibration:
   • MIRA Score: The joint posterior achieved a MIRA (Mass In Random Areas) calibration score of $0.635 \pm 0.017$ (where 0.667 is optimal), indicating good calibration with mild overconfidence. The marginal cosmological posterior scored $0.659 \pm 0.0139$.
   • TARP: Tests of Accuracy with Random Points confirmed that the credible regions achieved correct frequentist coverage across the parameter space.
5. Efficiency: Generating a single joint posterior sample takes approximately 0.2 seconds (256 ODE steps). Drawing 1,000 samples takes ~3 minutes, compared to the significantly higher computational cost of running MCMC chains (which required ~2.18 million gradient evaluations for the benchmark).

Significance and Claims
The paper claims to establish diffusion-based implicit inference as a viable route toward full field-level cosmological analyses. The primary significance lies in unifying mass map reconstruction and cosmological parameter estimation within a single, non-differentiable framework.

Key claims regarding the method's advantages include:

• No Differentiable Forward Model Required: Unlike MCMC-based field-level inference, JADE does not require analytical gradients of the forward model, theoretically allowing the use of arbitrarily complex simulators (e.g., full N-body or hydrodynamical codes) once training data is available.
• Amortization: The model is trained once and applied to any new observation instantly, avoiding the need for per-observation sampling.
• Joint Inference: It is the first method presented to jointly infer the full posterior $p(\theta, \kappa | \gamma_{obs})$ using a single conditional diffusion model, rather than marginalizing over the field or inferring parameters separately.

Limitations and Future Work
The authors note that the current demonstration relies on log-normal simulations, which are faster than N-body codes but less physically complete. Scaling to realistic, expensive simulators presents challenges regarding the density of the training set in cosmological parameter space. Additionally, as an amortized method, JADE must be retrained if the data model (e.g., survey mask, noise model) changes, unlike explicit inference methods that can adapt at inference time. The authors propose future work to extend this architecture to the joint inference of initial conditions and cosmological parameters, and to explore latent diffusion approaches to handle higher-resolution fields more efficiently.