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

Eugene Ndiaye

stat.MLNov 19, 2025

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

Authors: Eugene Ndiaye

Organizations: Apple

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.

Figures & tables

Figure 1 : Choosing and pulling back the reference set. Left: for a known CDF,

B_{\alpha}

is chosen by length and pulls back to a quantile interval. Center: under the discrete law

\nu_{n+1}

B_{\alpha,n+1}

is chosen by counting labels; here five observations are drawn from

N(0,9)

and

\alpha=0.4

, so four of the six labels are retained. Right: randomization fills the reference intervals and restores

\nu

; the bold orange curve shows the realization

\tau_{0}=0.85

. The horizontal brackets are the corresponding pullbacks on the

z

-scale.

Figure 2 : An equal-mass reference partition with

n+1=11=1+2\cdot 5

. The central remainder and the two radial levels crossed with five directions form eleven cells of equal

\nu

-mass. Their Euclidean areas differ because

\nu

is uniform in radius.

Figure 3 : Two samples of $n=640$ observations: a correlated Gaussian (top) and a skewed distribution (bottom). One further draw is the candidate in the augmented assignment. Left: observations coloured by the radius of their assigned target. Centre: source cells and the regions $\Omega_{r}$ of coverage $25.9$ , $53.2$ , $80.5$ and $96.1\%$ ; the $80.5\%$ region is filled. Right: the map $T$ . The targets form the grid of Section 4.1 with $n_{R}=n_{S}=25$ and $16$ copies of the origin.

Figure 4 : Fitting time against

n

for draws from the correlated Gaussian law of Figure 3 , using barycenters of equal-mass reference cells. One assignment of size

n+1

with the candidate known (reference);

n+1

assignments of size

n

, one per omitted label (naive); one assignment followed by the shortest-path search, which gives every

\ell_{j}

(proposed). For

n\geq 100

the proposed fit takes at most twice as long as one assignment; the naive method takes about

n

times longer at the tested sizes. The reference and naive fits use the SciPy assignment solver ( Crouse, 2016 ; Virtanen et al., 2020 ) ; the proposed fit uses Numba ( Lam et al., 2015 ) . Times are from one thread of an Apple M5 processor, with five instances for

n\leq 500

, three up to

n=10^{4}

, and one at

n=2\cdot 10^{4}

Panel of Figure 5 — Figure 5 : Nested unions for $n=99$ draws from $N_{2}(0,\tfrac{1}{4}I_{2})$ and $100$ targets at radii $k/10$ , $k=1,\ldots,10$ , along ten directions drawn uniformly on the circle. Retained cells are filled with the colour of their target’s radius, as in Figure 3 ; white cells are excluded. At $r=0.999$ the outer shell is excluded; at $r=1$ it is included. The common plotting window clips unbounded cells. For the source cells of our construction, Proposition 5.9 gives a sufficient condition for boundedness.

Figure 6 : Coverage of the prediction set against the nominal level, for Gaussian and skewed noise, with

n=99

(top) and

n=999

(bottom) calibration pairs. Each point averages

500

repetitions with

2000

fresh test pairs each. Targets are barycenters of equal-mass reference cells. Hard regions follow the staircase

|J_{r}|/(n+1)

and cover more than asked; randomized-region averages lie near

1-\alpha

. The steps shrink as

n

grows.

Figure 7 : The transform $G_{x}$ (left) and the field $H_{x}=A^{\top}G_{x}$ (right) for an affine score. Grey lines are level sets of $\Phi\circ S_{x}$ . $H_{x}$ is its gradient and crosses them outward; $G_{x}$ turns along them and is not monotone. Colour gives $\|G_{x}(y)\|$ .

Explore similar work

Multi-Variable Conformal Prediction: Optimizing Prediction Sets without Data Splitting
May 12, 2026Laura Lützow, Simone Garatti, Marco C. Campi +2Online Conformal Prediction Cross Validation
TU Munich & MCML, Germany · Politecnico di Milano, Italy · University of Brescia, Italy +1
A Unified Perspective on Conformal Prediction and Wasserstein Distributionally Robust Optimization for Uncertainty Quantification
Aug 30, 2026Kehan Long, Yiqi Zhao, Pol Mestres +3Distributionally-Robust Optimization Uncertainty Quantification
Contextual Robotics Institute, University of California San Diego, La Jolla, CA, USA · Thomas Lord Department of Computer Science, University of Southern California, Los Angeles, CA, USA · Department of Mechanical and Civil Engineering, California Institute of Technology, Pasadena, CA, USA +1
Multi-Distribution Robust Conformal Prediction
Jan 6, 2026Yuqi Yang, Ying JinOnline Conformal Prediction Nonconformity Scores
Mathematics Department, The Hong Kong University of Science and Technology, Hong Kong, China · Department of Statistics and Data Science, University of Pennsylvania, Philadelphia, PA, USA.

May 12, 2026, stat.MLJ/K move · Enter open · S save

Multi-Variable Conformal Prediction: Optimizing Prediction Sets without Data Splitting

Laura Lützow, Simone Garatti, Marco C. Campi, Lars Lindemann+1

TU Munich & MCML, Germany · Politecnico di Milano, Italy · University of Brescia, Italy+1

Online Conformal Prediction Cross Validation

Open paper AskPDF arXiv

Conformal prediction constructs prediction sets with finite-sample coverage guarantees, but its calibration stage is structurally constrained to a scalar score function and a single threshold variable - forcing shapes of prediction sets to be fixed before calibration, typically through data splitting. We introduce multi-variable conformal prediction (MCP), a framework that extends conformal prediction to vector-valued score functions with multiple simultaneous calibration variables. Building on scenario theory as a principled framework for certifying data-driven decisions, MCP unifies prediction set design and calibration into a single optimization problem, eliminating data splitting without sacrificing coverage guarantees. We propose two computationally efficient variants: RemMCP, grounded in constrained optimization with constraint removal, which admits a clean generalization of split conformal prediction; and RelMCP, based on iterative optimization with constraint relaxation, which supports non-convex score functions at the cost of possibly greater conservatism. Through numerical experiments on ellipsoidal and multi-modal prediction sets, we demonstrate that RemMCP and RelMCP consistently meet the target coverage with prediction set sizes smaller than or comparable to those of baselines with data split, while considerably reducing variance across calibration runs - a direct consequence of using all available data for shape optimization and calibration simultaneously.

Abstract

Explore similar work

Multi-Variable Conformal Prediction: Optimizing Prediction Sets without Data Splitting

A Unified Perspective on Conformal Prediction and Wasserstein Distributionally Robust Optimization for Uncertainty Quantification

Multi-Distribution Robust Conformal Prediction