Robust Conditional Conformal Prediction via Branched Normalizing Flow

Technical Summary: Robust Conditional Conformal Prediction via Branched Normalizing Flow

1. Problem Statement

Conformal Prediction (CP) provides prediction sets with marginal coverage guarantees under the assumption that calibration and test data are independent and identically distributed (i.i.d.). However, in the presence of distribution shifts (where the test distribution $Q_{XY}$ differs from the calibration distribution $P_{XY}$), existing methods typically align only the marginal conformal score distributions. While this preserves marginal coverage, it fails to control the conditional coverage error at individual test inputs. Consequently, CP sets may become unreliable in specific regions where the conditional score distributions are mismatched, leading to either under-coverage or excessive over-coverage for specific inputs. The core challenge is to ensure robust conditional coverage guarantees ($Pr(Y \in C(X) | X=x) \ge 1-\alpha$) under joint distribution shifts without requiring knowledge of the true test labels during inference.

2. Methodology

The proposed solution introduces a theoretical framework linking distribution shift to conditional coverage degradation and a novel architecture to mitigate it.

Theoretical Framework:

• Conditional Coverage Gap (CCG): The authors define CCG to measure the discrepancy between the conditional coverage of a test input under $Q_{XY}$ versus the guarantee under $P_{XY}$.
• Integrated Coverage Gap (ICG): Defined as the expectation of CCG over the test feature distribution $Q_X$, serving as a holistic robustness metric.
• Wasserstein Bound: The paper derives an upper bound for ICG using the Wasserstein distance ($W$) between the joint calibration and test distributions ($P_{XY}$ and $Q_{XY}$). Specifically, $ICG(P, Q) \le \sqrt{2\kappa L} (\eta \cdot W(P_{XY}, Q_{XY}) + 1/4)$. This bound implies that minimizing the Wasserstein distance between the distributions can theoretically promote robust conditional coverage.

Algorithm: Branched Normalizing Flow (BNF)
To operationalize the Wasserstein bound, the authors propose Branched Normalizing Flow (BNF), an invertible mapping $f_\theta(x, y) = (f_{\theta_X}(x), f_{\theta_Y}(y))$ designed to transform the test distribution $Q_{XY}$ to match the calibration distribution $P_{XY}$.

• Architecture: BNF utilizes two independent branches. The $X$-branch ($f_{\theta_X}$) normalizes test inputs, while the $Y$-branch ($f_{\theta_Y}$) transforms the labels. Crucially, the branches do not share parameters, ensuring that the transformation of $X$ does not depend on the unobserved test label $Y$.
• Inference Process:
  1. Transform a test input $X_{n+1}$ to a normalized input $\bar{X}_{n+1} = f_{\theta_X}(X_{n+1})$.
  2. Construct a conditional prediction set $CC(\bar{X}_{n+1})$ using the calibration distribution $P_{XY}$ (e.g., via Conformalized Quantile Regression).
  3. Inverse transform the set back to the test space using $f_{\theta_Y}^{-1}$ to obtain the final robust prediction set $C_{BNF}(X_{n+1})$.
• Augmented BNF: To handle complex distributions where univariate monotonic transformations (required for simple inverse mapping) are insufficient, the authors introduce Augmented BNF. This variant adds a stochastic noise component $\epsilon \sim \mathcal{N}(0,1)$ to the $Y$-branch transformation ($f_{\theta_Y}^{aug}(y; \epsilon)$). This increases expressiveness but requires solving a root-finding problem (or grid search) during inference to construct the prediction set, as the transformation is no longer strictly monotonic.
• Multi-Source Domain Generalization (MSDG): In scenarios where the test distribution is a mixture of $K$ source distributions, the method minimizes the average Wasserstein distance between the calibration set and the transformed source sets, assuming uniform mixing weights.

3. Key Contributions

1. Quantification of Robustness: The definition of Conditional Coverage Gap (CCG) and Integrated Coverage Gap (ICG) provides a theoretical metric to evaluate conditional coverage robustness under distribution shift.
2. Theoretical Bound: The derivation of an upper bound for ICG in terms of the Wasserstein distance between joint distributions, establishing a theoretical link between transport minimization and conditional coverage preservation.
3. Branched Normalizing Flow (BNF): The introduction of a two-branch invertible architecture that enables the transformation of test inputs to a calibration-like distribution without requiring test labels, thereby preserving conditional guarantees upon inverse transformation.
4. Augmented BNF: An extension using stochastic augmentation to improve fitting capability for complex distributions, coupled with a method to recover conditional guarantees despite the loss of monotonicity.

4. Experimental Results

The method was evaluated on nine datasets covering synthetic shifts and real-world challenges (sales prediction, traffic forecasting, medical decision-making, and epidemic modeling) under Multi-Source Domain Generalization (MSDG) settings.

• Metrics: Performance was measured using Marginal Coverage, Worst-Slice Coverage (WSC), and the proposed Worst-Slice Coverage Gap (WSCG), which penalizes both under- and over-coverage.
• Findings:
  • Robustness: Augmented BNF combined with Conformalized Quantile Regression (CQR) consistently achieved the lowest WSCG across all datasets and confidence levels, outperforming baselines including Split CP, Importance-Weighted CP, Worst-Case CP, and Wasserstein-Regularized CP.
  • Coverage: The method maintained marginal coverage and WSC close to the target confidence level ($1-\alpha = 0.9$) while significantly reducing the deviation from the target in conditional slices.
  • Efficiency: While standard Augmented BNF produced larger prediction sets to account for source uncertainty, an "Augment-Conditioned" variant (which conditions on the known source during inference) significantly reduced set size while maintaining coverage.
  • Generalization: The method demonstrated robustness across varying sample sizes, numbers of source domains ($K$), and different confidence levels.

5. Significance and Claims

The paper claims that its primary significance lies in addressing the conditional coverage degradation of Conformal Prediction under distribution shifts, a problem largely unaddressed by existing marginal-focused methods. By bounding the conditional invalidity via Wasserstein distance and implementing this via an invertible branched flow, the work provides a theoretically grounded and empirically effective mechanism to ensure that uncertainty quantification remains reliable for individual test inputs, even when data distributions shift. The authors position this as a step toward safer AI applications in critical domains (e.g., healthcare, traffic, epidemic modeling) where distribution shifts are inevitable and conditional reliability is paramount.

The paper acknowledges limitations, including the computational cost of root-finding in Augmented BNF, the stochastic nature of the prediction sets, and architectural constraints requiring multivariate inputs. It also notes that while calibration data is reused for training the flow (a potential tuning bias), ablation studies suggest this impact is minimal in practice.