Beyond Uncertainty Sets: Leveraging Optimal Transport to Extend Conformal Predictive Distributions to Multivariate Settings

Abstract

Conformal prediction (CP) constructs uncertainty sets for model outputs with finite-sample coverage guarantees. Yet ranking scores is straightforward only when they are scalar-valued, limiting CP to real-valued scores or ad-hoc one-dimensional reductions. Vector-valued scores arise naturally in multi-output regression and model aggregation, where each predictor in an ensemble provides its own score. Optimal transport (OT) defines vector ranks and center-outward multivariate quantile regions, though generally with asymptotic coverage guarantees. Applying a fixed transport map learned from calibration data to a new point introduces an uncontrolled approximation error. We restore finite-sample, distribution-free coverage by conformalizing vector-valued OT quantile regions. Each candidate's rank is defined by transporting the calibration scores augmented with that candidate's score, preserving the symmetry needed for validity. This appears to require a continuum of OT problems. However, we prove that the optimal assignment is piecewise constant across a fixed polyhedral partition of score space. This lets us characterize the entire prediction set in $O(n^3)$ time, matching the cost of a single assignment solve. It also addresses a limitation of prediction sets: they indicate which outcomes are plausible, but not their relative likelihood. In one dimension, conformal predictive distributions (CPDs) fill this gap by producing a predictive distribution with finite-sample calibration. Extending CPDs beyond one dimension remained an open problem. We construct, to our knowledge, the first multivariate CPDs with finite-sample calibration: a center-outward predictive distribution whose derived uncertainty regions have conformal coverage. We present both conservative and exact randomized versions; the latter generalizes the classical Dempster-Hill procedure.