Robust Conditional Conformal Prediction via Branched Normalizing Flow

TL;DR

This paper introduces Branched Normalizing Flow, a novel method to improve the robustness of conformal prediction under distribution shifts by controlling conditional coverage errors.

Contribution

The paper proposes BNF, a two-branch architecture that normalizes test inputs to calibration distribution, enhancing conditional coverage guarantees under distribution shift.

Findings

BNF improves conditional coverage robustness across nine datasets.

Theoretical bounds relate conditional invalidity to Wasserstein distance.

Empirical results show BNF outperforms existing methods under distribution shift.

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.