Self-Supervised Conformal Prediction with Equivariant Bootstrapping for Image Uncertainty Quantification

Technical Summary: Self-Supervised Conformal Prediction with Equivariant Bootstrapping for Image Uncertainty Quantification

Problem Statement
Inverse imaging problems, defined as recovering an underlying signal $x^\star$ from noisy observations $y = Ax^\star + \epsilon$, are ubiquitous in scientific fields such as astronomy and medical diagnostics. These problems are often ill-posed, leading to multiple plausible solutions and significant reconstruction uncertainty. While accurate Uncertainty Quantification (UQ) is critical for trustworthy scientific analysis, existing methods face two primary limitations:

1. Lack of Accuracy Guarantees: Many heuristic UQ methods lack statistical rigor.
2. Dependence on Ground Truth: Statistically rigorous methods, such as Conformal Prediction, typically require exchangeable ground truth data for calibration. In many domains (e.g., cosmology), ground truth is unavailable, expensive to simulate, or prone to distribution shifts relative to real observations. Relying on simulated calibration data can introduce hidden biases into downstream scientific inferences.

Methodology
The authors propose a fully self-supervised UQ framework that combines Equivariant Bootstrapping with Risk-Controlling Prediction Sets (RCPS) calibrated via Stein's Unbiased Risk Estimator (SURE). The methodology proceeds as follows:

1. Equivariant Bootstrapping:

   • Instead of standard parametric bootstrapping (which suffers from bias when bootstrapping from a single point estimate), the method leverages symmetries in the imaging task.
   • It utilizes a set of invertible transformations $\mathcal{G}$ (e.g., rotations, flips, cyclic shifts) that leave the signal space invariant but do not commute with the forward operator $A$ or the reconstruction method.
   • Bootstrap observations are generated by applying random transformations $T_{g_i} \in \mathcal{G}$ to the estimator $\hat{x}(y)$, passing them through the noise model, and inverting the transformation. This generates a distribution of bootstrap samples $\{\tilde{x}_i\}$ that reflects reconstruction variability without requiring ground truth.

2. Self-Supervised Calibration (RCPS with SURE):

   • To convert heuristic bootstrap regions into statistically rigorous confidence sets with guaranteed coverage, the framework employs RCPS.
   • The Challenge: RCPS requires a calibration dataset to estimate the risk (error rate). Since ground truth is unavailable, the authors replace the true non-conformity score $s(x^\star, \hat{x})$ with an unbiased estimate derived from Stein's Unbiased Risk Estimator (SURE).
   • Non-Conformity Score: The score is defined in the measurement space as $s_{A^\top A}(x^\star, \hat{x}(y)) = \frac{1}{m}\|Ax^\star - A\hat{x}(y)\|_2^2$.
   • SURE Estimation: The SURE formula provides an unbiased estimate of this score using only the observation $y$ and the reconstruction operator $\hat{x}(y)$:
     $$\text{SURE}(y) = \frac{1}{m}\|y - A\hat{x}(y)\|_2^2 - \text{tr}(\Sigma) + \frac{2}{m}\text{tr}\left(\Sigma \frac{\partial}{\partial y} A\hat{x}(y)\right)$$
   • Calibration: A scaling factor $\lambda$ is determined via a root-finding algorithm to ensure the prediction set $C_{\lambda; \alpha}$ satisfies the risk constraint $P(R(C) \le \alpha) \ge 1 - \delta$, where the risk is estimated empirically using the SURE values.

Key Contributions

• Self-Supervised UQ Framework: The paper introduces a method that generates statistically rigorous, per-image uncertainty quantification without any ground truth calibration data, thereby eliminating distribution shift biases between calibration and observation sets.
• Integration of Equivariant Bootstrapping and RCPS: The work combines the method-agnostic uncertainty generation of equivariant bootstrapping with the statistical guarantees of RCPS, refined by SURE.
• Application to Weak Lensing Mass-Mapping: The framework is demonstrated on the cosmological inverse problem of reconstructing dark matter convergence fields from galaxy shear measurements. This application is particularly sensitive to calibration biases, as existing methods often rely on simulated data tied to specific cosmological models.
• Method Agnosticism: The UQ framework is independent of the specific reconstruction algorithm used (e.g., Kaiser-Squires, deep learning, wavelet sparsity), serving as a post-hoc calibration layer.

Experimental Results
The method was evaluated on a weak gravitational lensing mass-mapping task using 300$\times$300 pixel mock COSMOS survey convergence maps.

• Setup: Reconstructions were performed using the Kaiser-Squires method. Equivariant bootstrapping generated 100 samples per image using geometric transformations and frequency shelving filters.
• Calibration: RCPS was applied with a target error rate $\delta = 0.1$ using SURE for calibration.
• Performance: The conformalized prediction sets closely tracked the ideal coverage diagonal across various confidence levels.
• Comparison: The self-supervised conformalized approach significantly outperformed unconformalized equivariant bootstrapping and parametric bootstrapping. Parametric bootstrapping provided negligible coverage, failing to capture adequate reconstruction variability.

Significance and Claims
The authors claim that this framework provides a robust solution for UQ in inverse imaging where ground truth is inaccessible. By avoiding simulated calibration data, the method prevents the introduction of hidden biases that could compromise downstream scientific inference, such as the estimation of cosmological parameters. The approach offers a statistically rigorous alternative to existing methods that either lack accuracy guarantees or rely on potentially biased simulated datasets.

Limitations
The paper acknowledges several limitations:

• Exchangeability Assumption: The calibration step assumes access to multiple observations drawn from the same distribution. In real-world mass-mapping, a single large sky observation is often used, and generating exchangeable emulations from it may not perfectly satisfy this assumption.
• Nullspace Insensitivity: The SURE-based score is blind to errors in the nullspace of the forward operator $A$. While this is less critical for mass-mapping (where $A$ is nearly full-rank), it could lead to underestimated uncertainty in severely rank-deficient inverse problems.
• Granularity: The current framework quantifies uncertainty at the image level (MSE) rather than providing per-pixel uncertainty estimates, which are common in supervised approaches.
• Computational Cost: Generating numerous bootstrap samples requires efficient reconstruction algorithms.