Multivariate Conformal Prediction using Optimal Transport

TMLR Paper6055 Authors

Abstract: Conformal prediction (CP) quantifies the uncertainty of machine learning models by constructing sets of plausible outputs. These sets are constructed by leveraging a so-called conformity score, a quantity computed using the input point of interest, a prediction model, and past observations. CP sets are then obtained by evaluating the conformity score of all possible outputs, and selecting them according to the rank of their scores. Due to this ranking step, most CP approaches rely on a score functions that are univariate. The challenge in extending these scores to multivariate spaces lies in the fact that no canonical order for vectors exists. To address this, we leverage a natural extension of multivariate score ranking based on optimal transport (OT). Our method, OTCP, offers a principled framework for constructing conformal prediction sets in multidimensional settings, preserving distribution-free coverage guarantees with finite data samples. We demonstrate tangible gains in a benchmark dataset of multivariate regression problems and address computational \& statistical trade-offs that arise when estimating conformity scores through OT maps.

Submission Length: Long submission (more than 12 pages of main content)

Changes Since Last Submission:

This revision clarifies both the paper’s positioning and its technical narrative.

First, we rewrote the methodological framing around a general two-step view of multivariate conformal prediction: construct a vector-valued discrepancy $S(x,y)\in\mathbb{R}^d$, then scalarize it via a map $\phi:\mathbb{R}^d\to\mathbb{R}$. This now provides a clearer common perspective on M-CP, Merge-CP, and OTCP.

Second, we substantially clarified the validity story. The manuscript now makes explicit that OT contributes a geometry-aware scalarization of multivariate scores, while finite-sample marginal validity follows from standard split conformal inference once that scalarization is learned on data independent of calibration. We also revised the localization section to distinguish more carefully between: (i) global marginal validity, (ii) exact finite-sample cell-conditional validity in the hard-partition setting, and (iii) soft localization as a heuristic for approximate object-conditional behavior.

Third, we strengthened the conceptual motivation for OT and spherical ranks. The revised paper now explains more clearly why fixed norm-based scalarizations induce ellipsoidal level sets, whereas OTCP can adapt to non-elliptical empirical score geometry through the transport map and the preimage of spherical shells. We also expanded the discussion of the spherical reference, rank/sign interpretation, and alternative reference choices.

Fourth, we expanded and reorganized related work in multivariate conformal prediction, including ellipsoidal / Mahalanobis approaches, manifold / Jacobian-based score engineering, copula-based methods, and a more explicit comparison to concurrent OT-based work.

Finally, we improved presentation throughout: notation was cleaned up, citation style and terminology were made more consistent, figure captions were rewritten to be more interpretive, and the discussion of OT hyperparameters $m$ and $\varepsilon$ was expanded, including clarification of the “debiasing” remark.

Assigned Action Editor: ~Matthew_J._Holland1

Submission Number: 6055