Multivariate Conformal Prediction using Optimal Transport

Michal Klein, Louis Béthune, Eugene Ndiaye, marco cuturi

Abstract: Conformal prediction (CP) provides distribution-free uncertainty quantification by constructing prediction sets whose validity relies on ranking conformity scores. Because ranking requires an ordering, most CP methods use univariate scores; extending them to multivariate settings, where no canonical order for vectors exists, remains challenging. We build on the theory of Monge--Kantorovich quantiles and ranks to propose a geometry-aware scalarization of vector-valued scores: we transport multivariate conformity scores to the spherical uniform distribution on the unit ball via an entropic optimal transport (OT) map and use the transported radius as a scalar score. Standard split conformal calibration then applies directly, preserving finite-sample marginal coverage. The resulting method, OTCP, produces prediction regions that adapt to the empirical geometry of the score distribution, going beyond the ellipsoidal sets imposed by norm-based scalarizations. Across a benchmark of 24 multivariate regression datasets, OTCP improves efficiency and conditional-coverage metrics mainly in low output dimensions ($d \leq 4$), while we also study the computational and statistical trade-offs involved in estimating entropic OT maps.

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

Changes Since Last Submission:

Abstract: Rewritten to lead with CP definition, then introduce Monge-Kantorovich quantiles and OT scalarization to the unit ball.

Introduction:
- Added references to existing multivariate CP approaches (norms, max-aggregation, copulas, manifolds)
- Added forward reference to the three-way data split and pipeline.
- Added schematic figure (Figure 1) with descriptive caption.
- Updated concurrent work pointer to Section 3.1.
- New figure: Created figures/otcp_schematic.pdf, showing the OTCP pipeline (score transport to reference ball + CDF-based calibration), generated programatically, as per reviewer No3E explicit request.

Section 3.1 (OT Scalarization):
- Merged Section 6 (Concurrent Work) into the comparison paragraph, per Reviewer bWf2's explicit request.
- Added data budget trade-off note for the hold-out/calibration split.

Section 4 (Experiments):
- Clarified that $\varepsilon=0.1$ is relative to cost matrix std (OTT-JAX rescaling).
- Fixed broken appendix reference ("Ablations provided in Appendix" → \Cref{fig:sup-region}).

Conclusion:
- Added discussion that $d\approx6$ threshold likely depends on benchmark sample sizes.
- Added synthetic non-elliptical experiments as future work direction.

Code: https://ott-jax.readthedocs.io/_autosummary/ott.tools.conformal.OTCP.html

Assigned Action Editor: ~Matthew_J._Holland1

Submission Number: 6055