Geometry-Aware Uncertainty Quantification via Conformal Prediction on Manifolds

Technical Summary: Geometry-Aware Uncertainty Quantification via Conformal Prediction on Manifolds

Problem Statement
Conformal prediction (CP) provides finite-sample, distribution-free coverage guarantees for regression tasks. However, standard CP constructions are designed for Euclidean output spaces, utilizing Euclidean residuals and coordinate-based regions. When the response variable $Y$ lies on a Riemannian manifold (e.g., the unit sphere $S^2$ for geomagnetic directions, the torus for protein angles, or rotation groups for spacecraft orientations), these standard approaches fail to respect the intrinsic geometry of the response space.

Two primary issues arise in this setting:

1. Geometric Mismatch: Standard residuals often rely on coordinate-based distances (e.g., latitude/longitude). On curved spaces, these distances depend on the chosen chart rather than the intrinsic discrepancy between responses. For instance, a fixed-width coordinate rectangle on a sphere has an area that varies with latitude (shrinking near the poles), whereas a geodesic cap of fixed intrinsic radius maintains a constant area regardless of location.
2. Heteroscedasticity: Prediction difficulty is rarely constant across the input space. In tasks like geomagnetic forecasting, secular variation rates differ significantly by region. A constant-width prediction set tends to over-cover "easy" regions and under-cover "hard" regions. While Euclidean regression addresses this by normalizing residuals with local difficulty estimates, existing manifold-valued CP methods have not explicitly adapted this concept to intrinsic geometry.

Methodology
The authors propose Adaptive Geodesic Conformal Prediction, a framework that integrates intrinsic geometry with local adaptivity within the split-conformal prediction setting.

• Setup: The data is split into a training set ($D_{tr}$), a calibration set ($D_{cal}$), and a test set. A base point predictor $\hat{y}: \mathcal{X} \to \mathcal{M}$ and a difficulty estimator $\hat{\sigma}: \mathcal{X} \to \mathbb{R}^+$ are trained using only $D_{tr}$.
• Nonconformity Scores: The method compares three scoring strategies:
  1. Naive Coordinate: Uses $L_\infty$ distance in spherical coordinates, producing axis-aligned rectangles.
  2. Standard Geodesic: Uses the unnormalized geodesic distance $s_{geo}(x, y) = d_{geo}(\hat{y}(x), y)$, producing constant-radius geodesic balls (spherical caps).
  3. Adaptive Geodesic (Proposed): Defines the score as $s_{adapt}(x, y) = d_{geo}(\hat{y}(x), y) / \hat{\sigma}(x)$.
• Difficulty Estimation: The estimator $\hat{\sigma}(x)$ is trained via 5-fold cross-validation on $D_{tr}$ to predict the geodesic residuals of the base model. It is trained without access to calibration labels to ensure valid conformal guarantees.
• Prediction Region: For a test point $x$, the prediction set is a geodesic ball centered at $\hat{y}(x)$ with a radius determined by the conformal quantile $\hat{q}_\alpha$ scaled by the local difficulty: $C(x) = B_{geo}(\hat{y}(x), \hat{q}_\alpha \cdot \hat{\sigma}(x))$. This allows the region size to expand in difficult areas and contract in easy ones while maintaining the shape of a geodesic cap.

Key Contributions

1. Geometry-Aware Regions: The method constructs prediction sets using geodesic distances, ensuring that regions (e.g., spherical caps) have position-independent area, thereby respecting the manifold's intrinsic geometry.
2. Locally Adaptive Calibration: By scaling residuals with a covariate-dependent difficulty estimate trained solely on the proper training data, the method adapts region sizes to local prediction difficulty without using calibration labels for model fitting.
3. Empirical Validation: The approach is validated on synthetic sphere data and a real-world geomagnetic forecasting task, demonstrating improvements in conditional coverage uniformity and worst-case coverage while preserving marginal coverage guarantees.

Experimental Results
The paper evaluates the method on two benchmarks: a synthetic sphere with heteroscedastic noise and the IGRF-14 geomagnetic field forecasting task.

• Synthetic Sphere ($S^2$):
  • All methods achieved valid marginal coverage near the nominal 0.90 level.
  • The Adaptive Geodesic method reduced the standard deviation of conditional coverage by 19% compared to the Standard Geodesic method.
  • Worst-bin coverage improved from 0.814 (Standard) to 0.839 (Adaptive).
  • The Naive Coordinate method required approximately 26% more area than the geodesic methods to achieve similar coverage, highlighting the cost of coordinate distortion.
• IGRF-14 Geomagnetic Forecasting:
  • The task involved predicting geomagnetic field vectors on $S^2$ across 352 locations with naturally heteroscedastic secular variation (ranging 274-fold in difficulty).
  • The difficulty estimator showed a Pearson correlation of $r=0.516$ with actual residuals.
  • The Adaptive Geodesic method reduced the conditional coverage standard deviation by 71% (0.031 vs. 0.107 for Standard Geodesic).
  • Worst-bin coverage increased significantly from 0.689 (Standard) to 0.855 (Adaptive).
  • The Standard Geodesic method was found to under-cover equatorial regions where field evolution is fastest, whereas the adaptive method maintained coverage closer to the nominal target across latitudes.

Significance and Claims
The paper claims that the proposed method offers a practical solution for uncertainty quantification in manifold-valued regression. The significance lies in the combination of two complementary choices:

1. Adaptivity: Scaling by $\hat{\sigma}(x)$ makes conformity scores comparable across regions of varying difficulty.
2. Geometry: Using geodesic caps ensures that the prediction set area is consistent with the manifold's metric, avoiding the distortions inherent in coordinate-based approaches.

The authors note that the effectiveness of the adaptive method correlates with the quality of the difficulty estimator. In the IGRF-14 experiment, a correlation of $r=0.516$ yielded substantial improvements. They suggest that if the correlation between the estimated difficulty and actual error is weak (e.g., below 0.15), the standard geodesic method may be preferable. The work does not claim to solve all manifold regression problems but provides a robust framework for cases where the response space is curved and prediction difficulty varies spatially. Future directions mentioned include extending the framework to non-exchangeable data and exploring anisotropic regions for direction-dependent uncertainty.