Robust Conditional Conformal Prediction via Branched Normalizing Flow

2026-05-03

AI Summary

This work addresses the challenge that conventional conformal prediction methods fail to guarantee reliable conditional coverage for individual test samples under distribution shift. To mitigate this limitation, the authors propose a Branch Normalizing Flow (BNF) architecture that employs a two-branch invertible mapping to align test samples with the calibration distribution and then inversely transforms them to construct prediction sets, thereby preserving conditional coverage guarantees even in the presence of distributional shifts. Innovatively linking Wasserstein distance to conditional coverage failure, the BNF is grounded in optimal transport theory, substantially enhancing coverage robustness. Extensive experiments across nine datasets and multiple confidence levels demonstrate that the proposed method consistently and significantly outperforms existing approaches.

Abstract

Conformal prediction (CP) constructs prediction sets with marginal coverage guarantees under the assumption that the calibration and test distributions are identical. However, under distribution shift, existing approaches primarily align marginal conformal score distributions, which is sufficient to preserve marginal coverage but does not control the conditional coverage error at individual test inputs. As a consequence, CP can remain unreliable in regions where the conditional score distributions are mismatched. In this work, we bound the conditional invalidity of CP under distribution shift in terms of the Wasserstein distance between the calibration and test distributions. This result highlights the role of invertible transport in mitigating conditional coverage degradation. Motivated by this insight, we introduce Branched Normalizing Flow (BNF), a two-branch architecture that normalizes a test input to the calibration distribution and transforms the prediction set of the normalized input back to the test distribution while preserving conditional guarantees. Empirically, BNF consistently improves conditional coverage robustness on nine datasets across a wide range of confidence levels.

Problem

• Conformal Prediction
• Distribution Shift
• Conditional Coverage
• Wasserstein Distance

Innovation

• Branched Normalizing Flow
• Conditional Conformal Prediction
• Distribution Shift
• Wasserstein Distance
• Invertible Transport