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

May 12, 2026stat.ML

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

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.

Laura Lützow, Simone Garatti, Marco C. Campi +2

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

Online Conformal Prediction Cross Validation

Ask PDF

Aug 30, 2026math.OC

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

Uncertainty quantification from finite data is central to machine learning, optimization, and automation systems, where decisions must remain reliable under limited samples and test-time distribution shift. Conformal prediction (CP) and distributionally robust optimization (DRO) offer two complementary approaches: CP constructs data-dependent prediction sets with distribution-free finite-sample validity under exchangeability, while DRO optimizes worst-case performance over an ambiguity set around an empirical distribution. We develop a unified probabilistic perspective on CP and DRO by viewing both as ways to turn finite calibration data into a data-dependent quantile estimator that a test score falls below with high probability. From this perspective, CP and DRO correct the empirical quantile along two coordinates of the same family of estimators: CP inflates the quantile level, whereas DRO shifts the quantile value through an ambiguity radius. Both methods provide the same calibration-conditional guarantee for the true distribution, requiring the target coverage to hold with high probability over the calibration sample. Their constructions differ, however: CP uses a closed-form, distribution-free level correction, while DRO uses a value-space correction whose certified radius depends on properties of the unknown distribution and additionally guarantees coverage uniformly over the ambiguity set. This distinction emerges in the tails of the score distribution. Because CP relies on sparse upper-tail order statistics of the calibration samples, its level inflation barely moves the estimator when those samples are dense near the target quantile but overshoots when they are sparse, whereas a well-chosen DRO radius corrects in value space and may avoid this overshoot.

Kehan Long, Yiqi Zhao, Pol Mestres +3

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

Distributionally-Robust Optimization Uncertainty Quantification Robust Optimization Perspective

Ask PDF

Jan 6, 2026cs.LG

Multi-Distribution Robust Conformal Prediction

In many fairness and distribution robustness problems, one has access to labeled data from multiple source distributions yet the test data may come from an arbitrary member or a mixture of them. We study the problem of constructing a conformal prediction set that is uniformly valid across multiple, heterogeneous distributions, in the sense that no matter which distribution the test point is from, the coverage of the prediction set is guaranteed to exceed a pre-specified level. We first propose a max-p aggregation scheme that delivers finite-sample, multi-distribution coverage given any conformity scores associated with each distribution. Upon studying several efficiency optimization programs subject to uniform coverage, we prove the optimality and tightness of our aggregation scheme, and propose a general algorithm to learn conformity scores that lead to efficient prediction sets after the aggregation under standard conditions. We discuss how our framework relates to group-wise distributionally robust optimization, sub-population shift, fairness, and multi-source learning. In synthetic and real-data experiments, our method delivers valid worst-case coverage across multiple distributions while greatly reducing the set size compared with naively applying max-p aggregation to single-source conformity scores, and can be comparable in size to single-source prediction sets with popular, standard conformity scores.

Yuqi Yang, Ying Jin

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.

Online Conformal Prediction Nonconformity Scores Distributionally-Robust Optimization Long-Run Coverage Guarantees

Ask PDF