Shape-Adaptive Conditional Calibration for Conformal Prediction via Minimax Optimization

1. Problem Statement

The paper addresses a fundamental limitation in Conformal Prediction (CP): while standard CP guarantees marginal coverage (i.e., the prediction set covers the true label with probability $1-\alpha$ on average), it often fails to provide conditional coverage (i.e., coverage holds for specific subgroups or local regions of the covariate space).

• The Challenge: Achieving exact test-conditional coverage ($P(Y \in C(X) | X) \ge 1-\alpha$) is theoretically impossible in a distribution-free setting without resorting to trivial, infinitely large prediction sets.
• Existing Limitations: Prior methods (e.g., Conditional Calibration by Gibbs et al., 2025) approximate conditional coverage by calibrating a threshold function $f(Z)$ for a fixed score function $s(X, Y)$. This approach suffers from two main issues:
  1. Geometric Rigidity: It can only adjust the volume of the prediction set (by scaling the threshold) but cannot alter the shape (e.g., orientation, aspect ratio) to match local heteroscedasticity or correlations in multivariate outputs.
  2. Masked Conditioning: It struggles when the conditioning variable $Z$ (e.g., sensitive attributes like race or gender) is available during calibration but unobserved at test time. Existing methods typically require $Z \subseteq X$.

2. Methodology: Minimax Optimization Predictive Inference (MOPI)

The authors propose MOPI, a framework that reformulates conditional predictive inference as a minimax optimization problem.

Core Formulation

Instead of fixing the score function and optimizing a threshold, MOPI optimizes over a structured class of set-valued mappings $\mathcal{C}$ while simultaneously solving for a weight function $f$ in a class $\mathcal{F}$.

The objective is to minimize the maximum deviation from the target coverage:
$$\min_{C \in \mathcal{C}} \max_{f \in \mathcal{F}} \mathbb{E}\left[ f(Z) \left( \mathbb{1}\{Y \notin C(X)\} - \alpha \right) - f^2(Z) \right]$$

• The Set-Valued Class $\mathcal{C}$: Defined as $C(x; h) = \{y \in \mathcal{Y} : T(h(x), y) \le 0\}$.
  • $h(x)$ encodes the geometric structure (e.g., mean $\mu(x)$ and covariance $\Sigma(x)$ for ellipsoids, or scale parameters for boxes).
  • This allows the prediction set to adapt its shape (rotation, aspect ratio) to local data density, rather than just scaling a fixed shape.
• The Weight Class $\mathcal{F}$: A class of functions on the conditioning variable $Z$.
  • If $Z$ is finite (group-conditional), $\mathcal{F}$ is a parametric class.
  • If $Z$ is continuous (test-conditional), $\mathcal{F}$ is a Reproducing Kernel Hilbert Space (RKHS).
• The Minimax Logic: The inner maximization over $f$ identifies the "worst-case" subgroup where coverage is violated. The outer minimization adjusts the shape parameters $h$ to correct this violation.

Key Technical Innovations

1. Decoupling $Z$ and $X$: The framework separates the coverage enforcement (dependent on $Z$) from the predictive mapping (dependent on $X$). This allows the model to learn a threshold/shape function using $Z$ during calibration, even if $Z$ is masked at test time.
2. Mean Squared Coverage Error (MSCE) Connection: The authors prove that under mild conditions, minimizing the minimax objective is equivalent to minimizing the MSCE:
   $$\text{MSCE}(C) = \mathbb{E}\left[ (\mathbb{P}(Y \notin C(X) | Z) - \alpha)^2 \right]$$
   This provides a rigorous statistical justification for the method.
3. Smooth Optimization: Since the indicator function $\mathbb{1}\{Y \notin C(X)\}$ is non-differentiable, MOPI replaces it with a smooth surrogate (e.g., Sigmoid or Gaussian error function) to enable gradient-based optimization (e.g., Stochastic Gradient Descent).

3. Key Contributions

1. Shape Adaptivity: MOPI generalizes previous conditional calibration methods by allowing the prediction set's geometry to co-evolve with local data density. It can learn optimal covariance matrices for ellipsoids or scaling factors for boxes, capturing local correlations that fixed-score methods miss.
2. Handling Masked Sensitive Attributes: MOPI is the first conformal method capable of enforcing valid conditional coverage on sensitive attributes (e.g., race, gender) that are available during calibration but unobserved at test time. This is crucial for fairness-aware AI where sensitive attributes must be hidden at inference.
3. Theoretical Guarantees:
   • Oracle Inequalities: The paper establishes non-asymptotic bounds on the MSCE of the empirical MOPI predictor.
   • Convergence Rates: It derives optimal convergence rates (up to logarithmic factors) for both finite-dimensional (group-conditional) and infinite-dimensional (RKHS) settings.
   • Approximation Error: Theoretical analysis shows that the error depends on the correlation between the covariate $X$ and the conditioning variable $Z$. If $Z$ is a deterministic function of $X$, perfect conditional coverage is achievable.
4. Computational Efficiency: By deriving closed-form solutions for the inner maximization (for finite $Z$) and using kernel methods (for continuous $Z$), the problem is reduced to a tractable empirical risk minimization problem.

4. Experimental Results

The authors evaluated MOPI on synthetic and real-world datasets against baselines: Split Conformal Prediction (SCP), Conditional Calibration (CC), and Randomized Localized CP (RLCP).

• Multivariate Regression (Ellipsoids/Boxes):
  • In synthetic data with heteroscedastic noise and vector-valued labels, MOPI produced significantly smaller prediction sets (higher efficiency) while maintaining coverage closer to the nominal level (90%) compared to baselines.
  • Baselines often produced overly conservative (large) sets because they could not adjust the shape to match the local covariance structure.
• Equalized Coverage (Masked Attributes):
  • In settings where sensitive attributes ($Z$) were masked at test time, baselines (CC, RLCP) failed to maintain equalized coverage across groups.
  • MOPI successfully utilized $Z$ during calibration to learn a shape function that balanced coverage across groups, even though $Z$ was unavailable at test time.
• Real-World Datasets:
  • Households Dataset: MOPI achieved more uniform local coverage across income levels compared to baselines.
  • Communities and Crime: MOPI maintained balanced coverage across racial groups in the "masked" scenario, whereas baselines showed significant coverage gaps for specific racial minorities.
  • FashionMNIST: MOPI achieved the best conditional coverage metrics (cluster and class-conditional) while maintaining marginal validity.

5. Significance

This work represents a significant advancement in Uncertainty Quantification and Fair AI:

• Beyond Volume: It moves conformal prediction beyond simple volume scaling to geometric shape adaptation, which is critical for high-dimensional and correlated outputs.
• Fairness in Practice: It solves a practical bottleneck in algorithmic fairness: enforcing fairness constraints (equalized coverage) on sensitive attributes that cannot be used at inference time.
• Theoretical Rigor: It bridges the gap between minimax moment-restriction problems and conformal prediction, providing the first non-asymptotic oracle inequalities for shape-adaptive conditional coverage.

In summary, MOPI offers a principled, flexible, and theoretically sound framework for constructing prediction sets that are not only statistically valid on average but also reliable for specific subgroups and adaptable to complex local data structures.