Optimal transport-based conformal prediction

We thank the reviewer for the careful reading and evaluation. We agree that our numerical experiments primarily compare against simple baselines. Our goal, however, is not to claim superiority over all methods but rather to motivate the use of transport-based quantiles in fundamental settings. We acknowledge that this may not have been sufficiently clear and we clarify this below. Thanks to your feedback, we also report additional metrics and results to provide a more comprehensive evaluation.

(computational costs and runtimes). We thank your for this comment which allows us to precise key properties of our methodology. Indeed, OT-CP+ requires solving multiple OT problems, which might be costly, especially for large-scale datasets. However, for the numerical experiments carried out in Fig.4, the computational time remains reasonable. An additional figure supporting this claim is now included: https://ibb.co/YTtv36ZM

(Comparison to existing related work). Our revised version now includes more related works as well as a discussion of limitations. Our main objective is to motivate the possibility of replacing univariate quantiles in the CP recipe by multivariate ones. OTCP should only be seen as a multivariate sorting of scores, which stands in contrast with the design of scores in applied settings. Our experiments suggest benefits of MK quantiles in terms of flexibility, but we do not claim to outperform all existing methods in multivariate regression. Rather, future work might profitably combine OTCP with more advanced multiple scores, different from componentwise residuals. Further numerical comparison is beyond the scope of our proposal and may perturbate the key message.

(Fig. 2). First, we stress that the ellipsoidal approaches do exploit correlations. Second, we intentionally only consider methods that can be interpreted as center-outward quantile regions, just as OT-CP, to ensure fair comparison. The mentioned references [1,2] differ in that they assume positive dependence, more related to a left-to-right ordering than a center-outward one.

(Fig 3b). We reported the empirical coverage averaged across subsets, which is drastically different from strategies that divide calibration examples into fixed groups (e.g., Mondrian CP). We also do not claim to achieve exact input-conditional coverage. We acknowledge the potential confusion and have revised our claim in Sec 3.2.

(Runtime cost of OT-CP+). Indeed, computational times is a limitation of OT-CP+, that is more costly than OT-CP. This cost depends on the parameter k. However, for experiments carried out in our paper, the price to pay seems relatively cheap, as illustrated in the new figure https://ibb.co/hpyCyFN comparing the time between OT-CP and OT-CP+ in the setting of Fig. 4.

(Type of coverage in Fig 4). The considered baseline shares our purpose of improving adaptivity, which is a common desiderata in CP. We agree that it does not guarantee asymptotic conditional coverage. In contrast, OT-CP+ satisfies this property at the price of larger set sizes on average when compared with local ellipsoids https://ibb.co/ymW5Mym6 A potential reason is that having more balanced conditional coverage requires including more points. We added a new figure tracking the marginal coverage of OT-CP+, to highlight that it does not tend to overcover https://ibb.co/1YNdSDny

(Recent CP methods for conditional coverage). As argued previously, our purpose is not to outperform all existing methods in conformal regression. This point has been clarified in the main body to ensure clarity when expressing our purposes.

(Discussion on related works). We discussed related work in "Introduction" and other paragraphs (CP methods for multi-output regression, in Sec. 3; CP methods for classification, in Sec. 4). We will integrate all these bibliographical elements into a dedicated subsection of the introduction.

Other Strengths And Weaknesses: We appreciate your comments on the writing, structure and illustrative figures. We have added more details on how to solve OT in practice.

(Motivation for Example 2) We believe that the results in Fig. 4 demonstrate OT-CP's ability to adapt to label confusions. Therein, classes 0 and 1 tend to be confused by the QDA classifier. In such cases, OT-CP achieves better trade-off between coverage and efficiency/informativeness. When the distribution is made easier (classes 0 and 1 become more distinct for QDA) all methods yield similar efficiency and informativeness: https://ibb.co/mVwMbsF6. This supports that OT-CP can effectively adapt to classification patterns where certain labels are prone to confusion.

(Solving the OT problem in practice) We used POT library [R2], which solves OT with the simplex algorithm. This solver is written in C++/Cython and is thus generally faster than pure Python implementations. We have added these details in Appendix. [R2] Flamary et al. “POT: Python Optimal Transport.” JMLR 2021

We thank the reviewer for the positive evaluation and address their comments with additional discussion and empirical results to illustrate the relevance and computational properties of OT-CP.

(Region size in regression). In Fig. 2, we reported the volumes to demonstrate that OT-CP achieves the desired coverage while producing smaller prediction sets than ELL, RECT. In Fig. 4, we initially did not include volumes since ELL fails to attain the desired conditional coverage, unlike OT-CP+. Based on your suggestion, we now monitor the volumes: https://ibb.co/ymW5Mym6 Overall, our results suggest that OT-CP(+) achieve coverage while ensuring efficiency, whereas other CP methods either fail to attain coverage or do so by producing unnecessarily large prediction sets. ELL generates smaller regions than OT-CP+, which may explain why it fails to achieve the desired coverage. We will add these results in our paper.

(OT-CP in Figures 8, 9, and 10). OT-CP provides more balanced coverage across classes, despite not being specifically designed for this task, while non-adaptive scores IP/MS can have low coverage (e.g., label 6 in Fig. 8). This benefit comes with better efficiency/informativeness than APS. This justifies our claim that OT-CP strikes "a favorable balance across all the considered metrics". Our revised version includes results averaged over labels: https://ibb.co/1GV3bsdq , https://ibb.co/nqTH1VCq.

(Average size of the prediction sets and other metrics with OT-CP+). Additional metrics complementing Fig. 4 will be included: https://ibb.co/ymW5Mym6 , https://ibb.co/YTtv36ZM Volumes of OT-CP+ are often larger than those from the ellipsoidal approach, but the smallest average set size is not necessarily the best as argued in Angelopoulos and Bates (2023). The motivation for OT-CP+ is precisely to enhance adaptivity with respect to the input.

([1] could be mentioned). This relevant reference was already included in the initial submission ("Conclusion and Perspectives").

(OT-CP+ and marginal coverage). We agree that the core guarantee of CP is finite-sample coverage. However, achieving adaptivity (i.e., obtaining prediction intervals with length adapted to the uncertainty relative to the considered test point), which is also a desirable property, is unattainable in finite samples. To introduce adaptivity, a sound strategy is to leverage universal consistent estimators, such as $k$NN. While this improves flexibility, it offers only asymptotic coverage, providing a different trade-off.

(OT-CP+ vs. [1], [2], and [3]). A general comparison with [1-3] is an interesting direction for future work, but we explain below why it is beyond our current scope.

(Other relevant baselines). [1,2] were already cited and discussed in our initial submission. We now cite [3] and note that this work proposes a CP method tailored for time series forecasting. Comparing our approach to theirs would thus require adapting OT-CP to handle time series, which is challenging since extensions of MK quantiles to such data have never been studied.

(Multivariate prediction regions of [1], [2], and [3]). Thank you for the question, which helps us refine our conclusions:

• We do not claim OT-CP outperforms all CP methods across every metric, but rather present it as a general methodology, valid in regression settings, improving simple baselines with basic multivariate scores via optimal transport.
• The suggested references are certainly relevant but appear to be more complementary than concurrent: [1] does not apply to black-box models, unlike OT-CP+ ; [2,3] propose CP strategies for handling complex univariate scores or time series.
• Combining OT-CP and [1-3] can foster interesting future work. For instance, one can replace the ball-shaped regions with conformalized radii from [3] by our MK quantile regions for greater flexibility.

(Available code). We provided the code in the supplementary material.

(Assumption 3.1). It implies $p(\cdot | x)$ is bounded away from $0$ and $+\infty$ on any compact subset of its convex support, a common assumption for transport quantiles (e.g., del Barrio et al., 2024).

("isn’t there a high likelihood that the ranking induced by this mapping would still assign very similar ranks to these two vectors?"). Our mapping does not necessarily lead to similar rankings for these vectors, precisely because we use optimal transport rather than simple normalization. Therefore, MK ranks are $d$-dimensional vectors that integrate the overall geometry of the underlying distributions and capture directional differences (and not just the magnitude).

(Choice of the reference distribution). This is a relevant question, as choosing a reference distribution is an open research problem (even beyond CP) and is largely guided by heuristics. We chose the $L_1$-norm ball because the $L_1$-norm of our score yields IP, a commonly used univariate score in CP for classification: see below eq.(11).

Thank you for the positive and detailed feedback.

(OT-CP+ and the use of k-nearest neighbor). We agree that OT-CP+ is not the only way to make OT-CP adaptive: our proposed methodology aims to demonstrate that conditional coverage can be achieved with only a slight modification of the generic OT-CP framework. In addition to being easy to implent, the added k-NN step allows us to leverage established results on the consistency of quantiles (del Barrio et al., 2024), which serve as the foundation for our Theorem 3.2. Therefore, OT-CP+ should be seen as evidence of OT-CP's inherent flexibility, showing its ability to incorporate refinements to achieve specific properties, such as adaptivity.

(OT-CP or OT-CP+). The results for base OT-CP and OT-CP+ are resp. presented in Sections 3.1 and 3.2: Figure 2 for OT-CP, and Figures 3 and 4 for OT-CP+. The figure titles now explicitly indicate the used method.

(Numerical results per class). We will add results marginalized over $Y$ to improve readability: https://ibb.co/1GV3bsdq ; https://ibb.co/nqTH1VCq ; https://ibb.co/p6jGSmXj See also our answer "OT-CP in Figures 8, 9, and 10" for reviewer bVvm.

(Real datasets). In addition to the reference, we included a table in Appendix B.1 summarizing the main statistics for each dataset. Since these are directly sourced from the literature rather than created in our work, we felt that no further details were necessary. That said, we are open to add any relevant information if needed.

(Reference not discussed). Thank you for the suggestion. We will expand Remark 2.3 by referencing Sections 3 and 4 of Peyré & Cuturi (2019) for an overview of solvers for OT.

(Computational comparison between baselines and the OT-CP). The calibration time corresponding to Fig. 4 (OT-CP+ compared with ELL) will be added in the appendix, https://ibb.co/YTtv36ZM

(Computational complexity). Including the k-NN increases the required complexity, as illustrated in https://ibb.co/hpyCyFN. For all the experiments carried out in this paper (up to $n=10\,000$ points), the computational time of OT-CP is fast thanks to the Python Optimal Transport library (which solves OT with a simplex implemented in C++). OT-CP+ is more demanding than OT-CP, as it solves one OT problem per new test point. The parameter $k$ drives the computational time, and our Theorem 3.2 only requires that it grows slower than $n$ (see next answers).

(Choice of k in k-NN). As the reviewer correctly pointed out, $k$ should be chosen carefully to ensure asymptotic conditional coverage. More precisely, our Theorem 3.2 shows that k should grow to $+\infty$ as $n \to +\infty$ at a slower rate than $n$ (this point is further clarified below, in "Questions for Authors"). In our experiments, setting $k = n/10$ provides a tradeoff between adaptive results and fast computational complexity. Thank you for your comment: we agree that this is an important aspect to highlight that has been added in our paper.

(OT-CP and full CP). OT-CP can be adapted to full CP and would consist in solving OT on the entire training dataset and every candidate $(X_{test}, y)$. Similar to full CP methods, this would improve statistical efficiency and reduce variability by avoiding data splitting, at a price of increased computational cost.

(Asymptotic conditional coverage). Indeed, Theorem 3.2 applies to the quantile region computed with the $k$NN step, under some mild assumptions on $k$ as a function of $n$. We provide more explanations on this point below.

(Assumption on k and n). The assumption is $k \to +\infty$ as $n \to +\infty$ and $\frac{k}{n} \to +\infty$. This means $k$ grows to $+\infty$ as $n \to +\infty$, and at a slower rate than $n$ (for instance, $k = \log(n)$). This assumption is taken from del Barrio et al. (2024) and ensures the consistency of the sequence of weight functions in their Theorem 3.3, which we use (through their ensuing Corollary 3.4) to prove our Theorem 3.2. Note that this is a classical assumption to ensure the universal consistency of $k$NN estimators (e.g., Corollary 19.1 in [R1]).

(Classification with K>=3). Our method can also be applied to binary classification ($K = 2$) but our score (11) would have an intrinsic dimensionality of 1, since the vector of estimated class probabilities belongs to the simplex. Because this setting is not the most relevant for us, we chose to focus on $K \geq 3$.

(One-vs-all strategy). OT-CP can be applied to one-vs-all classification by simply selecting an appropriate score function (i.e., based on K sigmoids). Our theoretical results remain valid in this setting. This is an interesting research direction that could further broaden the applicability of OT-CP.

(ELL). ELL refers to the ellipsoidal approach with adaptive covariance estimation from Messoudi et al. (2022): see l.292-297. We have made this point clearer in our revised paper.

[R1] Biau, Devroye (2015). Lectures on the nearest neighbor method.

General comment

We would like to thank all the reviewers for their time and feedback. We have revised the paper accordingly and provide detailed responses below.

1. We emphasize that integrating multivariate quantiles into the conformal prediction framework while ensuring theoretical coverage guarantees is non-trivial. The property that ranks follow a uniform distribution, under exchangeability, is crucial. In the literature of multivariate quantiles, optimal transport tools extend this critical property (see, e.g., Hallin et al. 2021). However, the stability arguments invoked in standard proofs of the quantile lemma (see, e.g., Tibshirani et al. 2019) do not directly apply here, as they rely on unresolved theoretical questions in optimal transport. To better highlight this subtle point, we clarified Step 2 of our OT-CP methodology and explicitly demonstrate how our approach provides theoretical coverage guarantees.
2. We have added the suggested references to improve the discussion on related work.
3. We systematically computed the volumes of the prediction regions, precised the running times of the method, and conducted new numerical experiments to better support our claims.

We now provide detailed reponses below. While we have carefully considered all reviewer comments, we are sometimes unable to provide an exhaustive answer for every point raised due to the character limit.

[Tibshirani et al.,2019] Tibshirani, Foygel Barber, Candès, Ramdas, Conformal prediction under covariate shift, Neurips 2019

[Hallin et al.,2021] Hallin, Del Barrio, Cuesta-Albertos, Carlos Matràn, Distribution and quantile functions, ranks and signs in dimension d: A measure transportation approach, Annals of Statistics, 2021

Answer to Reviewer RqtM

We thank the reviewer for their comments. We appreciate the opportunity to clarify our contribution.

"However, since the concepts of multivariate quantiles, ranks and confidence sets are already defined in [1] and derivative work, it is not clear to me what the main contribution is."

While optimal transport-based quantiles are not new, the novelty of our work lies in integrating them into an effective and flexible conformal prediction (CP) framework designed for multivariate scores. This aligns with the CP literature, where quantiles typically serve as building blocks for uncertainty quantification strategies, rather than primary contributions. More precisely, unlike [1], we address several nontrivial challenges specific to Monge-Kantorovich quantiles for CP, including establishing both nonasymptotic and asymptotic coverage guarantees, selecting reference rank vectors $\{U_i\}_{i=1}^n$ that are appropriate for multivariate scores, and comparing performance with existing CP methods.

"The paper follows the work of [1] and its derivative papers but it is not clear to me what the main novelty is."

Our paper goes beyond prior work on transport quantiles, including [1], by integrating them in conformal prediction, which had not been explored before. By combining these two lines of work, we make the contributions outlined above but also offer a novel perspective that could be of interest to both the machine learning and optimal transport communities.

"May be worth noting/including in the benchmarks the paper [2], that also applies the multivariate quantile function defined in [1] to regression tasks."

We thank the reviewer for the additional reference [2], which we have added to the introduction. Nevertheless, we emphasize that this paper leverages the Monge-Kantorovich quantiles from [1] to address quantile regression, and not conformal prediction.