Non-exchangeable Conformal Prediction with Optimal Transport: Tackling Distribution Shift with Unlabeled Data

Thank you for taking the time to review our work. We sincerely appreciate the encouraging feedback and the thoughtful comments, which we address in detail below.

Uncertainty representation vs quantification

That is a good point. We agree that conformal prediction is more accurately described as an uncertainty representation technique, as it conveys uncertainty through the size of the prediction set rather than assigning a numerical value. We will add a brief footnote to clarify this distinction in the revised version.

Ambiguous ground-truth

Thanks for highlighting another interesting scenario, where our methods are applicable. We will also include a discussion on ambiguous ground-truth in the paper.

Our approach is designed to operate directly on (non)conformity scores. This is convenient because these scores are typically unidimensional, which greatly facilitates computation. More importantly, it allows us to handle arbitrary distribution shifts. This advantage also extends to cases with ambiguous ground truth. For instance, we can directly apply our method to the unidimensional score proposed by Stutz et al., which is computed by averaging the conformity scores of each class, weighted by their plausibilities.

Formally, assuming $K$ possible classes and plausibilities $\lambda$ in the simplex $\Delta^{K-1}$, our method also applies to the score function of the form $s'(x, \lambda) := \sum_{K} \lambda_k s(x, y_k)$. This enables us to fit a new set of weights over the calibration data to obtain a threshold that adapts under distribution shift. Once this threshold is established, one can proceed with the construction of predictive sets as in (Stutz et al. 2023), or following the approach of Caprio et al., as you referenced.

The main challenge in applying our methods to ambiguous ground truth lies not in the theoretical construction, but in the design of the auxiliary distributions. This is because, rather than working with a discrete set of $K$ classes, we must consider the full simplex $\Delta^{K-1}$. Although the min-max construction remains straightforward to derive, it may be suboptimal, as the resulting auxiliary distributions can be "too far apart" to be effective. Identifying more effective strategies for constructing auxiliary distributions in this setting is a promising direction for future work. We thank the reviewer again for pointing us in this direction.

References

Stutz et al. "Conformal prediction under ambiguous ground truth". TMLR, 2023

Caprio et al. "Conformalized Credal Regions for Classification with Ambiguous Ground Truth." TMLR, 2024.

We thank the reviewer for their detailed feedback and constructive criticism. We address each concern systematically below.

1. Motivation for Total Coverage Gap

The total coverage gap serves two key purposes:

1. General Optimization Objective: The total coverage gap captures the overall effect of distribution shift, being zero when the distributions are identical. Similarly, the likelihood ratios of Tibshirani et al. (2019) do not target a specific $\alpha$, yet learning them remains both reasonable and effective for mitigating covariate shift effects. Both methods aim to learn a set of of weights only once using some unlabeled test data, such that applying standard split CP (with any $\alpha$) to the weighted calibration data, approximates calibration on the test distribution.

2. Theoretical Insight: The total coverage gap also establishes a novel connection between conformal prediction and optimal transport, specifically through the 1-Wasserstein distance. This provides a broader theoretical perspective, even though it is not directly tied to a particular $\alpha$

That said, we agree that some values of $\alpha$ are more relevant in practice. For this reason, we derive a bound for specific $\alpha$ values (Theorem A.6 in the appendix) and consider the coverage gap over a given range of $\alpha$ values (see discussion with reviewer XezP). However, in our experiments, these more targeted objectives have so far been less effective for optimization. We believe improving them is a promising direction for future work.

2. Lack of Theoretical Coverage Guarantees

We agree that statistical guarantees are central to conformal prediction. However, our work specifically addresses the setting of arbitrary distribution shifts, where standard conformal guarantees no longer hold. That is, our method does not invalidate any existing guarantees but rather operates in a regime where those guarantees are already compromised.

Moreover, unlike prior approaches, we make no assumptions about the nature of the distribution shift. For example, Tibshirani et al. (2019) rely on specific shift types, and Barber et al. (2023) require prior knowledge of the shift mechanism. In the general setting we consider, deriving formal coverage guarantees is extremely challenging. Nevertheless, we show that it is still possible to bound the coverage gap. Our key contribution is to use these bounds, even if they are not always tight, as a principled and effective objective for optimization. This allows us to approach the desired coverage level using only unlabeled samples from $Q$.

3. A Mismatch Between the Theory and the Algorithm

We thank the reviewer for highlighting this important point and will clarify it in the paper. By optimizing the calibration weights, we are directly minimizing the total coverage gap, but indeed this introduces dependence between the weights and the calibration scores. However, this can be easily addressed by fitting a weight function $w: s(x, y) \rightarrow \mathbb{R}$ on a separate set of calibration samples from $P$, thereby restoring independence. This is similar to the setup in Tibshirani et al. (2019), where weights are functions of covariates only, i.e., $w: x \rightarrow \mathbb{R}$.

To address this concern, we conducted additional experiments comparing our original method with the weight function approach trained on held-out data. The optimization remains essentially same: we use a set of labeled samples from $P$ and a set of unlabeled samples from $Q$ to learn the weight function $w: s(x, y) \rightarrow \mathbb{R}$ via our objectives. Once this function is learned, we apply it to another set of samples from $P$, the calibration set proper, before applying weighted split CP.

The results for iWildCam and ImageNet-C with severity 3 are shown in Table 1 below. We observe almost no difference in performance between the two approaches. The only downside of learning the weight function is that one needs another set of samples from $P$ (we use 300 samples for calibration in both cases, and another 300 to learn the weight function). However, once more, this is the same scenario considered in Tibshirani et al. (2019).

Table 1: Comparison of Original Method vs. Weight Function Approach

4. Mismatch Between Theory and Experiments

This apparent mismatch is by design and reflects practical considerations. We optimize the total coverage gap (which does require specifying $\alpha$) but evaluate at specific coverage levels (90%) because that aligns with how conformal prediction is typically used in practice, as you rightly pointed out.

To provide a more complete picture, we also report the total coverage gap itself in Tables 2, 3, and 4 in the appendix. We are also happy to include coverage results for additional $\alpha$ values if that would help strengthen the paper.

5. Using All Examples Simultaneously

There seems to be a misunderstanding regarding this point. Similarly to Tibshirani et al. (2019), we assume access to a small unlabeled dataset from the test distribution $Q$, denoted $\mathcal{D}_Q^{(1)}$, which we use to learn the weights (or weight function). Once the weights are learned, we have standard split conformal prediction but with a weighted calibration set, and it does not depend on how the test samples $\mathcal{D}_Q^{(2)}$ are observed. Our methods are applicable to sequential test examples to the same extent that standard split CP is and does not use all examples simultaneously. Note that we use disjoint subsets of $Q$ for weight learning ($\mathcal{D}_Q^{(1)}$) and evaluation ($\mathcal{D}_Q^{(2)}$).

In the case test samples come sequentially, we can always wait until we have enough samples before learning the weights. As shown, in Appendix B.3.1, we do not need many unlabeled test samples to learn the weights effectively. That means we use standard split CP for first $m$ samples, learn the weights, and then use these weights when applying split CP on the following samples. It should also be possible to refine the weights on the fly as new test samples arrive using the same objectives, but we leave that for future work.

6. On the Fairness of the Comparison with Other Baselines

1. Tibshirani et al. (2019): As mentioned previously, our method is directly comparable to the density ratio approach proposed by Tibshirani et al. (2019), and we follow a similar setup for a fair comparison.

2. Kasa et al. (2024): The method of Kasa et al. (2024) is designed for sequential test samples and does not use the unlabeled test samples in $\mathcal{D}_Q^{(1)}$. However, we adapted their method to the split CP setting as described in Appendix B.1.3 such that the comparison remains informative. In our experiments, we estimate the entropy required by their method using all test samples in $\mathcal{D}_Q^{(2)}$, which means, in our experiments, their method effectively uses the entire test set simultaneously, unlike ours.

3. Gibbs et al. (2023): The method of Gibbs et al. (2023) targets conditional coverage and also does not use the unlabeled test samples in $\mathcal{D}_Q^{(1)}$. As discussed in Appendix B.1.4, it offers stronger guarantees, their method comes with significant computational cost and relies on stronger assumptions. Still, we believe the comparison is informative, especially in highlighting the difficulty of prespecifying the class of distribution shifts, as required by their method, in large-scale classification tasks like those we consider.

6. Mathematical Formulation Issues

Thanks for pointing this out. We will improve the notation accordingly. Specifically, in *Equation (16)**, we will use the density function notation $p_{s \sharp P}(s_c)$ instead of the measure notation $s \sharp P(s_c)$, since $s\sharp P$ is a pushforward measure and should take measurable sets as input, not point values.

7. Overall Assessment and Significance

We believe our work merits higher scores based on several key contributions:

Novel Theoretical Framework: We derive new bounds to the coverage via a new connection between conformal prediction and optimal transport theory, establishing new theoretical insights through the 1-Wasserstein distance. This represents a significant theoretical advance in understanding distribution shifts in conformal prediction.

Practical Algorithm: Our method handles arbitrary distribution shifts using only unlabeled test data, without requiring prior knowledge of the shift type or mechanism. Strong empirical results across diverse datasets (ImageNet-C, iWildCam, synthetic regression) and shift types demonstrate its practical relevance.

Methodological Innovation: The use of auxiliary distributions to derive bounds without labeled test data is a novel and flexible approach that could inspire future work in non-exchangeable conformal prediction.

Broad Impact: Our framework applies to both classification and regression tasks and offers a principled solution to a fundamental challenge in deploying conformal prediction under real-world conditions.

We thank the reviewer for their thoughtful and constructive feedback. We address each of the main concerns below.

Coverage Level Weighting and Restricted Total Coverage Gap

That is a very interesting point. Thanks for bringing that up.

One could extend the definition to accept a range $\alpha^{-}$ and $\alpha^{+}$ such that $\Delta_{P,Q}(\alpha^{-}, \alpha^{+}) := \int_{\alpha^{-}}^{\alpha^{+}} \frac{\Delta_{P,Q}(\alpha)}{\alpha^{+} - \alpha^{-}} d\alpha$.

The proof of Theorem 3.2 in appendix A.1 would follow in a similar fashion, except now we only consider score values in a restricted range. We would lose the connection to the 1-Wassertein distance, but for the weighted-CDF version of the bound we would get

$\Delta_{P,Q}(\alpha^{-}, \alpha^{+}) \leq \frac{1}{\alpha^{+} - \alpha^{-}} \int_\mathbb{R} s\sharp P(s_c) | F_{s\sharp P}(s_c) - F_{s\sharp Q}(s_c) | \mathbf{1} (s_c \in [F_{s\sharp P}^{-1}(\alpha^{-}), F_{s\sharp P}^{-1}(\alpha^{+})]) ds_c,$ where $\mathbf{1}$ is a indicator function.

We tried to optimize this weighted-CDF bound during the rebuttal phase, but it underperformed the total coverage gap version of the bound. In the table below, we see the improvement over the unweighted version is small when optimizing $\Delta_{P,Q}(0.0, 0.15)$. We see two possible reasons for this. First, optimizing the bound for a specific coverage gap range may be inherently more difficult; for example, the upper bound for a fixed $\alpha$ derived in Appendix A.3 also failed to yield better empirical performance. Second, the total coverage gap may already serve as a strong objective, leaving limited room for improvement through alternative formulations. Yet, these more targeted objectives are a very promising avenue for future work, and we will include this discussion in the paper.

Empirical Comparison: Total Coverage Gap vs. High Coverage Gap ($\alpha \leq 0.15$)

Empirical Robustness to Sandwiching Distribution Misspecification

That is an interesting point. As mentioned in the paper, $(s_{\sharp}\hat{Q}_f, s_{\sharp}\hat{Q}_U)$ might already not safisfy the stochastic dominance relationship in practice. However, our experiments show that coverage often improves even in such cases. This suggests that the sandwiching approach can still yield meaningful bounds, even when the sandwiching distributions are somewhat misspecified. More generally, one can expect improved coverage whenever these distributions better approximate the true test score distribution than the calibration scores do.

Performance Analysis of 1-Wasserstein vs. Weighted CDF Bounds

One might expect the weighted-CDF bound to outperform the 1-Wasserstein bound, given that it is provably tighter. However, in some experiments, particularly the regression task, the 1-Wasserstein bound performed better. As discussed in Section~B.3.2, we hypothesize that this is because the 1-Wasserstein bound is easier to optimize. It only requires identifying the maximum value of the score density, while the weighted-CDF bound depends on accurately estimating the entire density. It is not entirely clear why regression tasks are more sensitive to this difference. One possibility is that, in conformal regression-as-classification, the output space must be discretized into bins somewhat arbitrarily. This binning may introduce additional noise into the score function, especially under distribution shift, making density estimation more challenging.

Comparison with Likelihood Ratio Methods

Learning weights via the proposed procedure --- is it possible to visualize by how much more efficient/robust/etc the weights learned via this procedure are vs. the Tibshirani et al brute-force density ratio proposal?

Could you clarify what you mean by robustness in this context and what sort of visualization you have in mind? In our experiments, we found that our method consistently outperforms the likelihood ratio approach proposed by Tibshirani et al. across a variety of settings. We believe this is largely because learning accurate likelihood ratios is inherently difficult, especially in high-dimensional spaces. In contrast, our method focuses on unidimensional scores, which are much easier to estimate and optimize.

Additionally, our approach is not limited to covariate shift scenarios, which makes it more broadly applicable and potentially more robust in diverse real-world conditions.

---

We thank the reviewer again for the thoughtful and constructive feedback. These questions helped us clarify both theoretical and practical aspects of our method. We hope our responses address the main concerns, but we are happy to discuss further if anything remains unclear.

Thanks for your review and for the encouraging comments regarding the novelty of our approach and the thoroughness of our experiments. We address your questions below.

Can you provide some qualitative examples of the UQ provided by your algorithm vs baselines?

We agree that qualitative examples are helpful for understanding the behavior of uncertainty quantification methods, and we plan to include them in the final version of the paper. Since we cannot share images during the rebuttal phase, we refer to Angelopoulos et al. (2021), who also apply conformal prediction to ImageNet and provide relevant visual examples.

To help clarify our experimental setting for readers unfamiliar with ImageNet-C: it applies 15 corruption types (e.g., Gaussian noise, motion blur, frost, fog) at 5 severity levels to create realistic distribution shifts that models encounter in deployment. In our setting, we observe that prediction sets increase in size with the degree of distribution shift, as expected. For instance, sets are larger on ImageNet-C with severity five than with severity one, as shown in Table 1.

Regarding other baselines, our primary focus is on improving conformal prediction under distribution shift, so comparisons to non-conformal baselines are not the central aim of this work. That said, we’d be happy to consider including such comparisons if they could help clarify or strengthen the paper. Are there particular baselines you think would be especially valuable to include, or aspects of the comparison that you feel would add meaningful insight?

How long does it take (and on what machine) on your datasets?

To be precise, the computational complexity of evaluating the KDE is given by $\mathcal O(n \cdot k)$, where $n$ is the number of datapoints where the KDE is defined ($n$ samples from $P$) and $k$ is the number of points where the KDE is evaluated. While $n$ is fixed for a given sample from $P$, the value of $k$ can be tuned according to the compute budget. It directly affects the accuracy of downstream computations: a higher $k$ improves the precision of the approximate integral in the case of the weighted-CDF bound, and enhances the estimation of $\max s_\sharp P$ in the context of the 1-Wasserstein bound.

In practice, fitting the weights is relatively fast. On a standard commercial GPU (RTX 2080 Ti), the process takes under a minute for 300 evaluation points and less than 10 minutes for 10,000 points. These runtimes were sufficient for our use case, but there is still room for optimization. Several techniques exist to accelerate KDE evaluation, including the fast Gauss transform and tree-based approximations, which could be incorporated to reduce computational overhead. See, for example, Yang, Duraiswami, and Gumerov (2003) and the references therein.

Would increasing number of calibration samples beyond 300 point, for better calibration quality?

Yes, we have observed that larger numbers of calibration samples are correlated with better coverage. See Tables 2 and 3 in Appendix B for an analysis of how the number of calibration samples from $P$ and unlabeled samples from $Q$ affects coverage.

This trend is expected. More calibration data typically improves the reliability of coverage estimates, and in our setting, also facilitates learning as it increases the number of points we can use to approximate $Q$ via the weighted empirical distribution over calibration scores. More unlabeled samples from $Q$ are also beneficial, since it increases the amount of information we have from the test distribution. However, we chose to focus our experiments on smaller calibration sets to demonstrate that our method remains effective even in more challenging and realistic scenarios where both labeled calibration data and unlabeled test samples from $Q$ are scarce.

Would any approximation method for density estimation or computing the 1-Wasserstein distance work in place of full fitting?

Yes, alternative methods are applicable, provided the resulting bound remains differentiable with respect to the learnable weights. For density estimation, other approaches can be used, including the efficient KDE variants mentioned above. Similarly, various methods exist for approximating the 1-Wasserstein distance. However, since the scores are typically one-dimensional, the 1-Wasserstein distance can be computed exactly in closed form, leaving little room for improvement through approximation.

References

Angelopoulos et al. "Uncertainty Sets for Image Classifiers using Conformal Prediction." ICLR, 2021.

Yang, Duraiswami, and Gumerov. "Improved fast gauss transform and efficient kernel density estimation." ICCV, 2003.