Multivariate Conformal Prediction using Optimal Transport

Many thanks for your detailed review and for the many kind suggestions to improve our work.

In conformal prediction, conditional coverage is more important than marginal validity ... but I believe only marginal validity alone may not be sufficient for publication in a top conference like ICML.

Many thanks for this very constructive point.

First, one should aknowledge that in vector-valued setting marginal validity is not straightforward, in contrast to $d=1$ where the conformity functions can be ranked. These canonical order no longer exists in multiple dimensions and unfortunately, marginal validity estimation can be lost when these multivariate quantile are not accurately estimated. These approximation biases can cause failure of marginal coverage even in recent works, for example Semi-parametric Conformal Prediction (J. W. Park etal, 2024)[Section 7: Limitations] which do not preserve marginal validity. In our setting, we additionally leverage entropic map whose approximation errors can also impact the coverage guarantee. This is why our conformalization step (Remark 3.6) is crucial. We will clarify.

We agree with the reviewer that conditional coverage is an important point to consider.

We clarified this on two fronts: metrics and by extending OTCP.

Results are visible in (anonymized link) https://shorturl.at/alFNM

• We leverage the implementation of [Dheur et al 25] and report the Worst Slab Coverage (WSC) and CEC-X [Appendix F.6] computed by their pipeline.

• In general, pointwise conditional coverage is impossible to achieve without stronger assumptions on the ground-truth distribution see https://arxiv.org/abs/1203.5422 or https://arxiv.org/abs/1903.04684. We have the same issue here. However, following your comment and that of Reviewer 8VoF, we also propose a simple adaptation of OTCP to approximate conditional coverage by partitioning the features space into regions $\mathcal{X} = \cup_{k=1}^{K} A_k$ and computing a transport map $T_{A_k}$ for every region. Our proof technique directly applies conditional on $A_k$​, and under exchangeability we have $\mathbb{P}(Y_{n+1} \in \mathcal{R}\_{\alpha}(X_{n+1}) \mid X_{n+1} \in A_k) \geq 1-\alpha$, for every $k \in [K]$. The partition $(A_k)$ are obtained by running $K$-means on the training set. These two baselines are OTCP-CLS (5) and OTCP-CLS (10), with $k=5$ and $k=10$, respectively.

The proposed score (16) seems to rely on a pre-selected score function ( S ) for scalar responses. The choice of ( S ) likely has a significant impact on the final results, and I encourage the authors to provide further discussion on this.

We agree, but this is the case for any other vector-valued conformal method. Even in one dimension, the choice of score function depends on the specific task and problem at hand. In this benchmark, we focused on simple residuals for OTCP to showcase more clearly the benefit of remapping these residues to their quantiles. Our framework covers arbitrary user-specified vector-valued score function.

The numerical results do not demonstrate a clear advantage of the proposed method. The region size in Figure 1 suggests that the proposed score yields comparable results to existing methods on most datasets.

We do not claim dominance indeed, as the breadth of datasets targetted in this benchmark (variety in size, dimension) can be overwhelming. But we do see encouraging results overall (supporting the use of OTCP), and an interest for visualization (e.g. in the taxi dataset, see last picture of https://shorturl.at/alFNM)

Many thanks for your thoughtful feedback, precise description of shortcomings, and suggestion to add references and baselines.

Presentation is rushed, which thwarts clarity

We apologize for this. We decided to submit based on the publication of a concurrent submission with a similar idea in https://arxiv.org/abs/2501.18991. We have fixed many typos.

[1] Predates Chernozhukov [...] [2]

We agree. Although they were published almost concurrently https://arxiv.org/abs/1406.4643, https://arxiv.org/abs/1412.8434, that important reference was missing. We added it.

existing methods that use ideas of optimal transport for multivariate quantile regression

Thanks for this remark. We wholeheartedly agree that a more detailed discussion on VQR and DQR was missing.

However, we argue that while OTCP and VQR may seem similar, they are very different:

VQR methods solve a far more challenging problem than OTCP: they model the conditional quantile map via OT whithin a structural learning problem. While this should in principle allow VQR to have a very fine-grained view, and potentially outperform OTCP, we argue that the reality is messy for $d\geq 2$ and that aiming for the most ambitious goal (full view of conditional quantiles) does not necessarily translate into practical gains when used for the simpler goal of uncertainty quantification.

This difference appears clearly in VQR's computations. While OTCP uses a standard call to Sinkhorn, VQR requires solving a mean-independence constrained OT problem, https://arxiv.org/pdf/2205.14977 [Eq. 4], for which there is no efficient solver other than SGD ("solving VQR via Sinkhorn is very slow in practice"). VQR also requires an extra post-processing step (VMR, Section 5). These practical hurdles, with no consistency guarantees, stand in contrast to our understanding of the Sinkhorn entropic map ([Pooladian/Niles-Weed 21]). Finally, VQR does not inherently provide coverage guarantees unless explicitly combined with conformal techniques with scalar-valued score (e.g. Feldman et al 2023).

Practically, the public implementation of (NL)VQR hardcodes the number of target points, so that it grows exponentially in dimension. Please note that https://arxiv.org/pdf/2205.14977 only consider datasets of dimension $d\leq 2$.

An important value of our submission, and of the code we have released, is to target upfront scalability issues, which is why we consider higher dimensions (in double digits) in our experiments.

With all these caveats, we have incorporated the following baselines in our benchmark:

• VQR, and its conformalized counterpart VQR-CP
• Nonlinear VQR NL-VQR, and its conformalized counterpart NL-VQR-CP
• ST-DQR-CP [Feldman et al. 23]

Please look at (anonymized link) https://shorturl.at/alFNM where we added conditional coverage metrics along with a localized OTCP.

Note: the intervals $[0,1/T,.., 1]$ coded in VQR was set so that $T^d \approx 8000$, i.e. same # of target points as that used for OTCP.

As can be seen, OTCP can perform better than VQR methods. These results should not be construed as a criticism of VQR. VQR aims for a more challenging problem, but may run out of steam in higher dimensions for this specialized task.

Propositions 3.4 and 3.5, which are a significant part of the contribution, are presented without proofs.

Thanks for kindly pointing this out. The proof of Prop. 3.3 requires applying Lemma 3.3 to $Z_i​=S(X_i​,Y_i​$). The proof of Prop. 3.5 is a direct application to the specific case of discrete Spherical uniform. We will clarify.

I was not convinced by the coverage plots [...] expected coverage falls below the required level.

The coverage can fall below the target level when the exchangeability assumption is not exactly satisfied; Some datasets are very small. We did not attempt to account for robustness.

“Dimensionality of the ball used for optimal transport?

To ensure the existence of the quantile transport map (inverse of the CDF), the dimensionality of the ball must be matched with that of $d$, the multivariate score $S$. There is no other alternative in the literature.

conformal prediction does not apply to the "continuous case", could the authors clarify this claim?

We will remove this. We meant that the mechanics of (continuous) Monge map estimation may collide, at first sight, with the empirical CDF approach that is crucial to CP.

"is the message here that CP works for any function, even those that poorly approximate the map?”**

A poor map estimation will compromise results in practice, but since we reconformalize the norms of the transports scores, our coverage guarantee trivially holds, just as the coverage does not depend on accuracy of base prediction model. Previous approaches for providing (marginal) coverage for vector-valued map failed see (J. W. Park etal, 2024)[Section 7: Limitations]. See reply to Reviewer QqAt

Many thanks for your detailed and insightful review.

conditional coverage are relevant metrics

We agree and now include conditional coverage metrics. We also implement a localized variant of OTCP where the data is partitioned in the feature space using $k$-means ($k=5$ or $10$) and a separate transport map is learned in each region. This offers a simple way to capture conditional heterogeneity.

Additional experiments with conditional coverage metrics as suggested can be found in https://shorturl.at/alFNM

As expected, the localized version can indeed improve the worst-case conditional coverage but with additional computations to fit transport map on every partitions.

The computation time can be significantly higher for only a slight reduction in the size of the prediction set compared to other methods. […]For Figures 2 and 8, it is not evident that the size of the predicted regions decreases as m increases. Moreover, Merge-CP appears to perform better on certain datasets

This is a fair observation, as the computation of a transport map is more costly. However, the gain in region size can be substantial. We do not claim that OTCP outperforms all methods in all datasets. From our observations, higher dimensionality can degrade the quality of multivariate quantiles L.260. We can still guarantee that coverage is provably maintained but the size of the region can indeed be inflated depending on the approximation errors.

"The paper focuses exclusively on multi-output conformal methods that do not require estimating the joint distribution”

This is a valid comment. If one has access to the an estimate of the joint distribution, it should be interesting to see how to wrap it with the OTCP framework. In this paper, we focused on clarifying the situation in a model-agnostic and arbitrary score. If one has access to a smooth joint or conditional distribution $P_{Y \mid X}$ One natural strategy is to follow PCP style for examples and consider vector-valued scores on samples obtained from the generative model.

Additionally, one could also extract a conditional transport map T_x \\# P_{Y \mid X=x} = \mathbb{U} to a reference distribution. Our framework can still operate as a wrapper around generative models by being compatible with any estimated (conditional) transport map on the output. Indeed if a variable $Z$ can be transported by $T$, then an invertible function $f$ of $Z$ also induces a transport map $T_f = f\circ T \circ f^{-1}$. As such, a transport map on $Y$ also induces a natural transport map by composition on the conformity score $s(x,y)$ simply by applying this to $f= s(x, \cdot)$ for each $x$. As such, we can easily incorporate it in OTCP pipeline to provide valid coverage. We will add more clarification on this in the revised version.

While the vector-valued conformity score is conditioned on ( x ), the optimal transport map itself is not, raising concerns about the method’s ability to fully capture conditional uncertainty.

This is a good remark and corresponds to exactly what we implemented as localized OTCP with a $k$-means strategy. Basically each point will have its own transportation map and we can think of the OT merging score as $S_{\mathrm{OTCP}}^{A_k}(x, y) = \|T_{A_k} \circ S(x, y)\|$ where the collection of $(A_k)_{k\in[K]}$ is a clustering partition of the feature space. This matches the suggested approach leveraging conditional transport map in Hallin et al. (2021)

“How do the authors generate $n_S$ vectors uniformly on the unit sphere in dimension $d$?”

We formally describe it in the paper L.324 "Sampling on the sphere". We use a quasi-Monte Carlo method to generate points evenly spread on a sphere. Starting from well-distributed points in $[0,1]^d$, we transform them using the inverse normal distribution to get Gaussian-like vectors, then normalize them to lie on the unit sphere. This gives us a lower-discrepancy sampling of directions than random sampling.

Illustrations of the prediction regions would have been helpful in better understanding what OT-CP does, and providing concrete examples could have better motivated the proposed method.

We provided the taxi demand prediction task and included a new visualization with the localized transport maps. Thanks for the suggestions for improving our figures.

“The lack of publicly available code limits reproducibility.”

We have already released an OTCP module in a major optimal transport toolbox, that can be used to directly recover our experimental results. To preserve anonymity, we cannot include the link now but it will be included in the camera ready version of our paper.