Neural Optimal Transport Meets Multivariate Conformal Prediction

TL;DR

This work tackles conditional multivariate distribution estimation by marrying neural optimal transport with vector quantile regression, enabling geometry-aware conditional quantiles via a convex-potential map learned with ICNNs. The learned CVQF yields multivariate ranks that feed into conformal prediction, producing distribution-free predictive regions that adapt to the conditional geometry of Y|X. The authors introduce amortized optimization to accelerate training and inference, and explore entropic and non-entropic variants, achieving tighter, more informative CP regions than coordinatewise or density-only approaches. Empirical results on synthetic and real multi-target datasets demonstrate improved coverage-utility trade-offs and scalability to high-dimensional targets. Overall, the framework offers a principled, scalable path to multivariate uncertainty quantification with robust finite-sample guarantees and geometry-aware predictive sets.

Abstract

We propose a framework for conditional vector quantile regression (CVQR) that combines neural optimal transport with amortized optimization, and apply it to multivariate conformal prediction. Classical quantile regression does not extend naturally to multivariate responses, while existing approaches often ignore the geometry of joint distributions. Our method parametrizes the conditional vector quantile function as the gradient of a convex potential implemented by an input-convex neural network, ensuring monotonicity and uniform ranks. To reduce the cost of solving high-dimensional variational problems, we introduced amortized optimization of the dual potentials, yielding efficient training and faster inference. We then exploit the induced multivariate ranks for conformal prediction, constructing distribution-free predictive regions with finite-sample validity. Unlike coordinatewise methods, our approach adapts to the geometry of the conditional distribution, producing tighter and more informative regions. Experiments on benchmark datasets show improved coverage-efficiency trade-offs compared to baselines, highlighting the benefits of integrating neural optimal transport with conformal prediction.

Figures (11)

• Figure 1: Example of points sampled from reference distribution of all the datasets we study and points sampled from approximation constructed by C-NQR$_U$ method.
• Figure 2: S-W2 on Neal's funnel distribution. We scale the dimension of a funnel from 2 to 16.
• Figure 3: Log-volume of the prediction sets, normalized by $d_y$, of the resulting prediction sets for different methods. Results averaged over 10 independent data splits. Nominal miscoverage level $\alpha=0.1$
• Figure 4: Worst-slab coverage for different methods. Results averaged over 10 independent data splits. Nominal miscoverage level $\alpha=0.1$
• Figure 5: Worst slab coverage at different nominal miscoverage $\alpha$ levels for conformal prediction methods, achieved on large datasets.

Theorems & Definitions (9)

• Theorem 1: carlier2016vector, Theorems 2.1 & 2.2
• Theorem 2: carlier2016vector, Theorem 2.3
• Remark 1
• Theorem 3: Volume–optimality of pullback balls under radiality
• Remark 2
• Theorem 4: Volume–optimality of pullback balls under radiality
• proof
• Remark 3: Examples satisfying assumptions of \ref{['thm:radial_volume_optimality']}
• Remark 4