# Robust Conditional Conformal Prediction via Branched Normalizing Flow

arXiv: https://arxiv.org/abs/2605.01868 · HTML: https://mlsanity.com/paper/2605.01868/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 45 of its 80 references that are in the map.
- First posted on arXiv 3 May 2026.

## Builds on (most cited first, 30 of 45)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction](https://mlsanity.com/paper/1802.03426/) (2018)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Normalizing Flows for Probabilistic Modeling and Inference](https://mlsanity.com/paper/1912.02762/) (2019)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [Distribution-Free Predictive Inference For Regression](https://mlsanity.com/paper/1604.04173/) (2016)
- [Improving Variational Inference with Inverse Autoregressive Flow](https://mlsanity.com/paper/1606.04934/) (2016)
- [Conformalized Quantile Regression](https://mlsanity.com/paper/1905.03222/) (2019)
- [Adaptive Graph Convolutional Recurrent Network for Traffic Forecasting](https://mlsanity.com/paper/2007.02842/) (2020)
- [Conformal Prediction Under Covariate Shift](https://mlsanity.com/paper/1904.06019/) (2019)
- [Classification with Valid and Adaptive Coverage](https://mlsanity.com/paper/2006.02544/) (2020)
- [Uncertainty Sets for Image Classifiers using Conformal Prediction](https://mlsanity.com/paper/2009.14193/) (2020)
- [Conformal prediction beyond exchangeability](https://mlsanity.com/paper/2202.13415/) (2022)
- [Statistical Aspects of Wasserstein Distances](https://mlsanity.com/paper/1806.05500/) (2018)
- [Why and When Can Deep -- but Not Shallow -- Networks Avoid the Curse of Dimensionality: a Review](https://mlsanity.com/paper/1611.00740/) (2016)
- [Conformal Prediction With Conditional Guarantees](https://mlsanity.com/paper/2305.12616/) (2023)
- [Traffic Graph Convolutional Recurrent Neural Network: A Deep Learning Framework for Network-Scale Traffic Learning and Forecasting](https://mlsanity.com/paper/1802.07007/) (2018)
- [Localized Conformal Prediction: A Generalized Inference Framework for Conformal Prediction](https://mlsanity.com/paper/2106.08460/) (2021)
- [Learning Optimal Conformal Classifiers](https://mlsanity.com/paper/2110.09192/) (2021)
- [Robust Validation: Confident Predictions Even When Distributions Shift](https://mlsanity.com/paper/2008.04267/) (2020)
- [A comparison of some conformal quantile regression methods](https://mlsanity.com/paper/1909.05433/) (2019)
- [Domain Adaptation for Time Series Forecasting via Attention Sharing](https://mlsanity.com/paper/2102.06828/) (2021)
- [Augmented Normalizing Flows: Bridging the Gap Between Generative Flows and Latent Variable Models](https://mlsanity.com/paper/2002.07101/) (2020)
- [Regression Conformal Prediction with Nearest Neighbours](https://mlsanity.com/paper/1401.3880/) (2014)
- [Adaptive, Distribution-Free Prediction Intervals for Deep Networks](https://mlsanity.com/paper/1905.10634/) (2019)
- [normflows: A PyTorch Package for Normalizing Flows](https://mlsanity.com/paper/2302.12014/) (2023)
- [Selection and Aggregation of Conformal Prediction Sets](https://mlsanity.com/paper/2104.13871/) (2021)

## Abstract (arXiv)

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.