Wasserstein-Regularized Conformal Prediction under General Distribution Shift
Authors: Rui Xu
Chao Chen
Yue Sun
Parvathinathan Venkitasubramaniam
Sihong Xie
Venue: ICLR 2025 Poster (2025)
Keywords: conformal prediction
distribution shift
non-exchangeability
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: For conformal prediction under distribution shift, we bound of coverage gap by Wasserstein distance between calibration and test conformal scores, which is further decomposed and minimized based on impacts from covariate and concept shifts.

Abstract: Conformal prediction yields a prediction set with guaranteed $1-\alpha$ coverage of the true target under the i.i.d. assumption, 
which can fail and lead to a gap between $1-\alpha$ and the actual coverage. Prior studies bound the gap using total variation distance, which cannot identify the gap changes under distribution shift at different $\alpha$, thus serving as a weak indicator of prediction set validity. Besides, existing methods are mostly limited to covariate shifts, while general joint distribution shifts are more common in practice but less researched. In response, we first propose a Wasserstein distance-based upper bound of the coverage gap and analyze the bound using probability measure pushforwards between the shifted joint data and conformal score distributions, enabling a separation of the effect of covariate and concept shifts over the coverage gap. We exploit the separation to design algorithms based on importance weighting and regularized representation learning (WR-CP) to reduce the Wasserstein bound with a finite-sample error bound. WR-CP achieves a controllable balance between conformal prediction accuracy and efficiency. Experiments on six datasets prove that WR-CP can reduce coverage gaps to 3.2% across different confidence levels and outputs prediction sets 37% smaller than the worst-case approach on average.

Published as a conference paper at ICLR 2025
WASSERSTEIN-ABGGULARZZZD CONFORMAL
PREDICTION UNDER GENERAL DISTRIBUTION SHIFT
Rui Xu, Sihong Xie
The Hong Kong University of Science and Technology (Guangzhou) rxu2330 connect. hkust -22. edu. en, sihongxielhkuust-gz... -ew.ettttgg.eduue...
Chao Chen
Yue Sun, Parvathinathan Venkitasubrammnniam Lehigh University yus Lew.geeeeehhhh.e. pav309elehigh.edi
Harbin Institute of Technology cha01nbex0 gmail. com
ABSTRACT
Conformal prediction yields a prediction set with guaranteed a coverage of the target under the i.i.d. assumption. which may not hold and lead gap between - a and the actual coverage. Prior studies bound the gap using total variation distance, which cannot identify the gap changes under distribution shift a given a. Besides, existing methods are mostly limited to covariate shift, while general joint distribution shifts are more common in practice but less researched. In response, we first propose a Wasserstein distance-based upper bound of the cov- erage gap and analyze the bound using probability measure pushforwards between the shifted joint data and conformal score distributions, enabling separation of the effect of covariate and concept shifts over the coverage gap, We exploit the separation to design algorithm based importance weighting and regularized representation learning (WR-CP) to reduce the Wasserstein bound with a finite- sample error bound. WR-CP achieves a controllable balance between conformal prediction accuracy and efficiency. Experiments on six datasets prove that WR-CP reduce coverage gaps 3.2% across different confidence levels and outputs prediction sets 37% smaller than the worst-case approach on average.
INTRODUCTION
Because of data noise, unobservable factors, and knowledge gaps, stakeholders must also consider
prediction uncertainty in machine learning applications, especially in areas such as fintech (Ryu &
Ko 2020), healthcare (Feng et al. 2021), and autonomous driving (Seoni et al. 2023). Conformal
prediction (CP) addresses prediction uncertainty by generating set of possible targets instead of
single prediction (Vovk et al. 2005 Shater Vovk 2007 Angelopoulos & Bates 202T). We
focus on CP in regression tasks. With trained model h, CP calculates the difference (conformal
score) between the predicted and actual target via score function s(z, y) !h(x) y/ over some calibration instances. With the empirical 1 quantile of the conformal scores, the prediction set (ar) of a test input æ contains targets whose scores are smaller than r. If calibration and test data are independent and identically distributed (i.i.d)), the probability that the prediction set C(x) contains the true target of is close to a (i.c. the coverage guarantee).
Denote Pxy and QXy the calibration and test distributions, respectively, in space X y. We assume B/ ~ N (fp(x),ep) for (z. y) ~ Pxy and y|ı ~ N (fq(x), for (æc 8) ~ Qxy. In practice, be violated by a joint distribution shift such that Pxy + QXY. due to
the i.i.d. assumption covariate shift (Px * Qx), a concept shift (fr * fq). or both (Figure [[a) left) (Kouw & Loog 2018). With distribution shift, the coverage guarantee fails, leading gap between the probability that E C(x) and a. Formally, denoting Py and Qv the calibration and test conformal score distributions, respectively, the coverage gap is the difference between the cumulative density functions (CDFs) of Pv and Qy at quantile 7 (Figure [ left). Prior methods are concerned with
"Corresponding author; Sihong Xie, a hongxi e@thust--z....m.. edu.

---

Published as a conference paper at ICLR 2025
(a)
1. Initial Data Distribution Model
Importance Weighting for Covar late Shift
Wasserstein Minimization
Concept Shift
sample cample
weighted Pay somple Cay test sample
weighted Pay cal xample
Co best sampo regularioed f
iritial
Operations on Finito Samples
s(x.y) IM00 yl
Curuutatte diltutttee Prebabiiity
Cumaslative Probabiitt
Cumulative Probabiity
weighting
Wasserstein Distance Decomposition and Minimization
coverage gap
distance
W-disbance
W-distance
weighted P.
weighted P. Uy
1-s quantle
(b)
Prctabiity Serve
TV di-dizance
Choose the Right Metric to Bound Coverage Gap
W-disancee
TV-dirannee
Prior: Total variation distance
Proposed: Wasserstein distance
Figure 1: (a) Joint distribution shift can include both covariate shift (Px & Qx) and concept shift (fr + fo). Coverage gap (Eq. (3) is the absolute difference in cumulative probabilities calibration and test conformal scores at the 1 a quantile r. We address covariate-shin shitte--siit-iie shift--ndneeee Wasserstein distance by applying importance weighting (Tibshirani et al. 2019) to calibration samples, and further minimize
induued Wasserstein
distance to obtain accurate and efficient prediction sets; (b) Q and ? are two distinct test conformal score distributions. Wasserstein distance (Eq. (5)) integrates the vertical gap between two cumulative probability distributions overall all quantiles, and is sensitive to coverage gap changes at any quantile. Total variation distance fails to indicate coverage gap changes thoroughly as it agnostic about where two distributions diverge.
the worst-case shifts and passively expand prediction sets as much as possible to meet the coverage guarantee for any shifted test distribution, leading to excessively large and inefficient prediction sets (Gendler al. 2021, Cauchois et al., 2024 Zou & Liu 2024 Yan et al., 2024). Recent works
assume knowledge about the distribution shifts between test and calibration distribution (Barber et al.,
2023 Angelopoulos et al. 2022; Colombo 2024). The knowledge is further embedded as the total
variation (TV) distance between conformal score distributions Pv and Qy to bound and minimize the coverage gap. However, the TV distance ignores where two conformal score distributions differ. while the coverage gap is defined at specific and is location-depender indicative of coverage gap during model optimization (Figure (b) right).
making TV distance less
Opposing to TV distance, we adopt Wasserstein distance over the space of probability distributions of conformal score to upper bound the coverage gap under joint distribution shift. Such an upper bound integrates the vertical gap between the CDFs of two conformal score distributions Pv and Qv and measures the gap at any a (Figure Tb) left), indicating coverage gap at a given a for distribution discrepancy minimization and coverage guarantee (Section 3.1 Appendix BB. Targeting more effective algorithms specifically for covariate and concept shifts that constitute joint distribution shift, we further penetrate the complex landscape of joint distributional shift. We disentangle the complex dependencies between the Wasserstein upper bound and covariate and concept shifts using a novel pushforwards of probability measure, decomposing the bound into two Wasserstein terms so that the effects of covariate and concept shifts on the coverage gap are independent (Eq. 7 Theoretical analyses crystalize the link between CP coverage gap, the smoothness of the conformal residue and predictive model, and the amount of covariate and concept shifts (Section T... The decomposition allows representation learning using importance weighting (Tibshirani et al., 2019) that reduces the covariate-she shit--iiiitttteeddd term, and minimization of the concept-shiftindduce term (Figure aT) with finite samples with an empirical error bound (Section 4.2) We proved the effectiveness of the resulting algorithm, pessserstein-reulaaiize generalization where the test distribution is an unknown mixture of training distributions. On six datasets from applications including AI4S (Brooks & Marcolini 2014), smart transportation (Cui et al. 2019: Guo et al., 2019). and epidemic spread forecasting Deng et al.) 2020), experiments across various a values (0.1 to 0.9) demonstrate that coverage gaps are reduced to 3.2% and the prediction set sizes are 37% smaller than those generated by the worst-case approach on average. Besides, WR-CP allows a smooth balance between prediction coverage and efficiency (Figure 5).
conformal prediction (WR-CP), for multi- domain
2

---

Published as a conference paper at ICLR 2025
2 BACKGROUND AND RELATED WORKS
2.1 CONFORMAL. PREDICTION
Let X E EEX Rd and Y V C R denote the input and output random variable, respectively. A hypothesis h N' - y is a model trained to predict target Y from feature X. We observe n instances (X1, Vi)...., (Xn; Yn) from calibration distribution Pxy. Taking (I, as a realization of (X, Y), a score function s(x, y) X''x x + V C R quantifies how (x, y) conforms to the model h. For regression tasks, typically s(x, y) [h(x) y/. Split conformal prediction is widely-used, and defines calibration conformal scores V, =X,,, Y2) for i = =1, 1,...," (Papadopoulos et al. 2002).
Letting E be the (1 1)1/n quantile of Vi.... Vn the prediction set of input
is
c(X++1) = {0 s(Xn+1 0) I T. J/.
(1)
Consider the instance (Xx+1, 1)) following test distribution Qxy. If the test and calibration instances are i.i.d. (i.e. Pxy Qxy), the probability that the true target Y,+1 is included in
1(X,+1) is at least 1 1 a. If calibration conformal scores are almost surely distinct, we can also bound the probability from above by 1 1/(n + 1) (Angelopoulos & Bates 2021). The bounded
probability called coverage guanrantee:
Pr C(Xn+1)) E [1 a- 1/(n + 1)).
Vovk et al. (005) proved that the assumption of i.i.d instances can be relaxed to exchangeability of calibration and test instances. With exchangcability, prior CP methods proposed to improve the adaptiveness of prediction set to different test inputs (Romano ct al. 2019; 2020 Guan 2023
Amoukou & Brunel 2023 Han et al., 20.3) and maintain conditional coverage guarantee for sub-
populations of the test distribution (Gibbs et al.] 2023 Jung et al. 2022 Feldman et al. 2021
Cauchois et al. [022]] Foygel Barber al. 2021 Stutz et al. 2021 Einbinder et al. 2002b). However,
when the assumption is violated so that PxY QXy, coverage guarantee may not hold.
2.2 CONFORMAL PREDICTION UNDER DISTRIBUTION SHIPTS
Covariate shift (Px F Qx) ): Tibshirani et al. ratio between Px and Qx to satisfy the i.d.d assumption, so coverage is ensured under covariate shift. Concept shift (ʃp + fQ): [iibinder et al. (2022a): Sesia al. (2023) addressed CP under concept shift which is represented by label noise.
2019) adopted importance weighting by likelihood
Joint distribution shift (Pxy * Qxy) consists of covariate shift (Px * Qx) and/or concept
shift (fp * fo) (Kouw & Loog 2018). Barber et al. (2023), Angelopoulos et al. (2022), and
Angelopoulos & Bates (2021) upper- coverage gap via total variation distance, but TV distance cannot identify gap changes at fixed a. To reduce the gap, Gibbs & Candes (2021), Xu & Xie (202I). and Gibbs Candes (2024) focus on CP under dynamic shift (test distribution changes over time). Meanwhile, some works concentrate on static shift (test distribution unchanged). These works can be categorized into two pipelines. The first pipeline modifies vanilla CP upon a residual-driven model for robust coverage (Gendler et al. 2021; Cauchois et al.. 2024 Zou & Liu 2024). The second pipeline incorporates a conformal based loss during training to obtain robust and efficient prediction sets (Yan al. 2024). However, these works treat joint distribution shift as whole and adopt worst-case principle for prediction.
In this work, we explore CP under multi-source domain generalization, which focuses on developing a model that generalizes effectively to unseen test distributions by leveraging the data from multiple source distributions (Sagawa al. 2019 Krueger al.) (021). A related, yet distinct area is federated CP, which aims to train a model across decentralizer data sources to perform well on a known test distribution (typically uniformly weighted mixture of source distributions) without requiring centralization to ensure privacy. Regarding federated CP, FCP (Lu et al.) (202)) and FedCP- QQ (Humbert et al. (023) aim for a coverage guarantee when the test and calibration samples are exchangeable from the same mixture. When exchangeability does not hold, DP-FedCP (Plassier et al.) 2023) addresses scenarios where test samples are drawn from a single source distribution, assuming that only label shifts (Py + Qy) occur among the source distributions. Besides, CP with missing outcomes is studied by Liu et al. (2024) where the samples from the test distribution are accessible. The proposed WR-CP does not consider privacy but works on more generalized setup: the test samples are drawn from an unknown random mixture where both concept and covariate shifts can occur among the source domains.
3

---

Published as a conference paper at ICLR 2025
3 METHOD
3.1
SYPER-BOUNDING COVERAGE GAP BY WASSERSTEIN DISTANCE
As shown in Figure [b), Wasserstein distance can effectively indicate changes in coverage gap across different values of a. We formally upper-bound coverage gap via Wasserstein distance. Let V E V S IR be the random variable of conformal score. Py and Qv are calibration and test conformal score distributions, respectively. The guarantee in Eq. 2 indicates that Pr (s(Xn+11 Y S T) E [1 - a, a + 1/(n + 1)). FPV and Foy are CDFs of Pv and Qv. respectively. Under the i.i.d. assumption, Pv Qv, and thus For (T) Fpy (r) E a,, 1 + 1/ + 1)). However, the assumption can violated by joint distribution shift, which may results Pv * Qv. In this case, Fp, (T) is still bounded, but For (T) # FP, (T). Inadequate coverage renders prediction sets unreliable, while excessive coverage leads to large prediction sets, reducing prediction efficiency, and we define coverage gap the absolute difference"
Coverage gap
(F) Fqv (T)
(3)
Definition 1 (Kolmogorov Distance). (Gaunt & Li 2023) F2 and F, are the CDFs of probability measures y and on R. respectively. Kolmogorov distance between μ and v is given by ((,,v)
suPzeR [x)( (x)).
With Definition 1 as T E V C R. Eq. is bounded by K K(Pv, Qv):
Coverage gap = (T) - Fav (T) S supvEv FPy (v) - K K(P, (4)
Definition 2 (p-Wasserstein Distance). (Panaretos & Zemel 2019) Given two probability measures and v on metric space (X, cx), where X is set and ex is a metric on X, the Wasserstein distance
of order 2 1 between and v is
Wp(uvv) inf Urex cx(21, dy((1, 1/8 gET(pM)
(5)
where [(s, v) is the set of joint probability measures on X' X with marginals ((AxX) = =m(A) and 7xx x B) v(B) for all measurable sets A, X. Proposition 1. (Ross 2011) If a probability measure p. in space R has Lebesgue density bounded by L. then for any probability measure v, K(p,v) v) S (OLWi(,,))
In this work, let W denote Wasserstein distance with p p= 1. Applying Eq. (4) and Proposition I with I. as the Lebesgue density bound of Pv, we can develop an upper bound by
Coverage gap S 22LW [v)...
(6)
3.2 WASSERSTEIN DISTANCE DECOMPOSITION AND MINIMIZATION
In Eq. (6). we show that the Wasserstein distance W (Pv, Qv) can effectively bound the coverage gap caused by a joint distribution shift. However, it is still not clear how the two components of joint distribution shift, namely, covariate shift in space and concept shift in space y lead to W (Pv, Qv) in space V. Besides, we want the quantified contributions amenable to optimization techniques to reduce (PPv, Qv). To the best of our knowledge, there is no prior work that suits this need. Therefore, propose to upper-bout bound W(Pv, Qv) with two discrepancy terms due covariate and concept shifts, and corresponding optimization methods to reduce W W(Pr (Pv. Qv) via minimizing the two terms.
Definition 3 (Pushforward Measure). If X' and y are separate measurable spaces, is a prbability measure on X, and f X + y a measureable function, define the pushforward fal of μ through such that (+)( ( for all measurable set A [ y.
With Definition 3 we have Py = SP# Px and Qr = JQ# Besides, we define sp(x) = s(x, ff(2))
(xx) - fr(x)) for x Px, and sQ(x) = s(x, fx(x)) = [h(x) fo(x) for
I Qx, leading to pushforwards of the conformal score Py = s Px and Qv = SQ# Qx
In this study, we assume that n is sufficiently large for Fry (7) to be approximated as allowing us to view Eq. (3) the difference between Fov (T) and
4

---

Published as a conference paper at ICLR 2025
To upper-bound W (Pv, Qv) about conformal scores according to covariate (concept, resp.) shifts in the X' (), resp.) space, we introduce a pushforward Qv.sr s =p#Qx on V. Since Pv and QV.er pushforward measures by the same function sp from Px and Qx, respectively, ((vv, Qv.nr) is a measure of covariate shift (Px * Qx). Also, as Qv.sp and Qy are pushforward measures from the same source Qx by Sp and sq. respectively, spp,.sp Qv) can indicate the extent of concept shift (fp a fq. and thus sp a sq). The relationships among the pushforward measures are shown in Figure As Panaretos & Zemel (2019) states, the triangle inequality holds that
W(Pv, Qv) & ((v,, Qv.sp) + (Qv,spp Qv).
With Eq. 77) bounding ((vv. Qv), we design an approach to minimize the upper bound. First, we adopt importance weighting, which weights calibration conformal scores with the likelihood ratio dQx (x)/ (z). Tibshirani et al. 2019) prove that impor- tance weighting can preserve the coverage guarantee when only a covariate shift occurs. However, existing works do not include the weighting technique when dealing with a joint distribution shift. We prove that importance weighting can minimize covariate- shift-induced Wasserstein distance, W W(Pv, Qv.sp), even if concept shift coincides. Given any measurable set A S X. B = {sp(x) I E
V
Sp
Px
fr
Py
Pr
covariate
Sp
2v.5p
concept shift
Qx
Qr
Note: (ot):/)()>-/b)>- (((+0) (e)))/ - fx0)
Figure 2: Pushforward measures.
A} CV. With Definition 3
Pv(B) = / dPv(v) = d(sp)Px)(o)= = //100/////0 dPx(x) weighting
- = dQx(1) dPx(x)= / (Qx(())= = (((s=*xx)())) 1/100000000 dQv.sp (v) = Qv. (B). (x)
Since importance weighting can transform Pv to Qv..p. W(Pv, Qv.ap) is minimized, and the remaining term in the upper bound in Eq. (7) is the concept-shift-ineee indift-induued component WQQ,, Qv)- Next, we further minimize it during training, as illustrated in Figure [ The reasoning behind distinguishing between covariate and concept shifts elaborated in Appendix
4 THEORY
4.1 UPPER-BOUUDDE WASSERSTEIN DISTANCE BY COVARIATE AND CONCEPT SHIFTS
Although Eq. 7 upper bounds for W(Pv, Qv) by W(v,, Qv.sp) and WQ Qv), it remains unclear how these shifts lead to these terms. Covariate shift can be more accurately quantified by (Px. Qx). Also, with Qy
Qp W(Qy.fr. Qy) is more direct way to measure concept shift by comparing fo and fo based on Qx. Therefore, we further upper-boung the two terms on the right-hand side of Fq. 7 using W(Px, Qx) and W(Q..ee Qr). We extend theorem in Aolaritei et al.] (2022), which pushes two probability measures with the same function, while Theorem considers pushing with different functions.
Theorem 1. For probability measures I and v on metric space (X, letting f.g xxx: X + y be measurable functions, jij and Vg on metric space (D,cy) are pushforwards of μ and v under functions f and respectively. The Wasserstein distance between μ and Vy holds the equivalence:
/
v2) (v1.))))) inf geT(U,)) /
W(ur,n)
inf vET(ay.s) yxy
Xxx
As cv = cy(x1, = =|x1 322. Theorem leads to the following
inf / (xr(xx.xx) Ixxx (sp(xr) sQ(x2)) dx),,,
W(Qv.ops
(8)
inf / W(Qy....Qr) 7Er(Qx(xx) xxxx (ff(21) f2((2)) (2(
(9)
5

---

Published as a conference paper at ICLR 2025
Let be the optimal transport plan of ((Qrssn:Q With 7 = max (p(21)=ag(02)
T1 [/((e))-1((=2]1 we have
W(Qv.up: Qv) S JxXX - (0(2)) - (21,,
(10)
$ 5xxxx nlfr(x11) fo((z2) dy" x2) = nw (Qr.fo Qy)
In Eq. (O)), the first inequality holds as "- may not be the optimal transport plan of W W(Qv.sr Qv). and the second inequality follows the definition of nJ. Appendix Dishows a geometric intuition of rj-
Theorem 2. For probability measures and on metric space (X, cx) with measurable function f + Y, Ps and v on metric space (Y, ey) are the pushforward of y and v through function S. respectively. If f has Lipschitz continuity constant K, ie., (o(f(x).f(aa)) S K, M, 22 E X,
cx(21122)
Vy) (WW
(11)
As X S (sp(a)))spar)) K, 12 X. With Theorem 121-22l2 Qx through 8p, we have
RR, (x) = || 22ll2. is the Lipschitz constant of sp X + V such that
1 as Pv and Qv are pushforwards of Px and
W(vv, Qv,sp) S <xW(Px,QN)
(12)
Plugging Eq. 10 and Eq. 12
into Eq. (7), W(Pv, Qv) s kW Qx) n Qy) we can further bound the coverage gap using the magnitudes of
Therefore, by utilizing Eq. covariate and concept shifts:
(6)
Coverage gap S V 222 (xW Qx) + m/ (Qx, Qy)).
(13)
highlights how covariate and concept shifts impact the coverage gap. While the values of W (Px, Qx) and W (Qy.ff: Qy) are inherent properties of given data and cannot be altered, the parameters and 17 are linked to the model h, allowing minimizing and via optimizing h.
Equation
13
4.2 EMPIRICAL UPPER BOUND OF COVERAGE GAP
In practice, Py and Qv are rarely available. Sometimes we may have access to their empirical distributions via the score function s, where Py is derived from n calibration samples and Qy is obtained from test samples. Having the Wasserstein distance between the two empirical distributions W W(Pv, Qv), we derive the error bound between the empirical form and W (Pv,Qv) by asymptotic estimation.
Definition 4 (Upper Wasserstein Dimension). DDulley 1969) Given a set A S X, the covering number, denoted (e(AA), is the minimun such that b closed balls, B1...,,B, of di- ameter achieve A C UisiciB, For distribution μ in X, the (c, dimienssoo is de( () - log (Nc(A) ((AA) 2 /)o//oge. The upper Wassersteion dimension with
is
dw (1)) = inf E oo): lim d. (1,, -= S sp)
(14)
With the definition of upper Wasserstein dimension, Weed & Bach (2019) conducted how an empirical distribution converges to its population by the Wasserstein distance between them.
Proposition 2. WWed & Bach 2019) Given a probability measure o >> dw (u). If Pa is an empirical measure corresponding to i.i.d. samples from μ,, 3λ E R such that E(W (p, An)) s An-1/o, Furthermore, for > O, Pr(W (n, Pm)) 2 (H, P.m) + t) S c-2ns"
Theorem 3. Given two probability measures and v, on >> dw dw(H) and o, >> (w(v). Pn and im are empirical measures corresponding to n i.i.d. samples from I and m iidd. samples from v,
respectively. For tu to > 0, 3r A E R with probability at least (1 -- -2nt, (1 2mt that
W(u,)) & W(in, Pm) + Aum -1/a, + >um 1//0 -1/ov + tp +t..
(15)
Applying Theorem 3 to Eq. (6). we derive empirical upper bound of coverage gap. Specifically, Py has Lebesgue density bounded by L, for tp, to > 0, op > dw (Pv), and oQ > dw (Qv).
EXp, XQ E R with probability at least (1 - e-ntpp" (1 e- 2m2oq" that
Coverage gap S V2 (W(Pr, Qv) + Apn Apm-111p 1/ar + >qm=-1/a 1/00 +tp +tp+to) tt).
(16)

---

Published as a conference paper at ICLR 2025
5 APPLICATION TO MULTI-SOURCE CONFORMAL PREDICTION
In this work, we consider the test distribution to be an unknown random mixture of multiple training distributions, referred to as multi-source domain generalization (Sagawa et al. 2019). As highlighted by |Cauchois et al. (2024), achieving 1-a coverage for each of the training distributions ensures that the coverage on test data remains 1-a the test distribution is any mixture of the training distributions. We apply the methodology outlined in Section 3 to this scenario, namely multi-source
conformal prediction. Given training distributions XY for I,.,,, k, we require Qxy follows
QX E E We W; XY 01 we 2
Wi
(17)
In other words, Qxxy is an unknown random mixture of D Xy io I,... k. Next, we introduce a surrogate of W (Qv Qv), allowing the minimization of W(Qv,sn, Qv) even when the test
distribution Qx is unknown in practice. With the score function s(x, y) and D = =s D Dxy
Qv s QxY
8# En w. D.
k E
W48# Do E.
w Wr
(18)
By marginalizing out Y in Eq. (17 we obtain Qx = Ei= D Similar to 18 with score
Eq.
function s2(z) and ( V,ar = SP D Qv. = SP#0x = Ei-1 w. ( V.er*
Theorem 4. In space X C RR is a mixture distribution of multiple distributions u I,....k. such that v = -Ein wayl) with w1, W2 2 0, Σx=
1= W; = 1. For any distribution μ on X, Wasserstein distance has the inequality that W (J,, v) Σ (w,W(u,,
By Theorem 4 W(Qv,ep: Qv) S Et=, u. (. D(") E w/ E. Ei=1 WD
Visp' D
The inequality offers surrogate of (QQ.ssp: Qv). Even if QXY is unknown, with uniformly distributed weights, we minimize the expectation of the surrogate with w; = 1/k for i = 1.... k:
min Ein W(D.... D."). Besides reducing the coverage gap, we also want smaller prediction errors, SO we include empirical risk minimization (ERM) (Vapnik 1991) during training. Hence, with a loss function and a parameterized model hg, we merge the constant 1/k with a hyperparameter B, and introduce the objective function
E (2 [ho(z)) + 3 Σ * W(D Vap DVD)
min
(19)
We design Wasserstein-regullaiize Conformal Prediction (WR-CP) to optimize he by Eq. 199 with c(i) finite samples and generate prediction sets with small coverage gaps. SXY is the sample set drawn from D for i = 1, 1...., k, and Sky is the sample set drawn from Pxv. SKy is a test set containing samples from an unknown distribution QXY. Algorithm [ shows the implementation of WR-CP. Kernel density estimation (KDE) is applied to obtain Px. D and Qx for the calculation of likelihood ratios, whereas D06) and (:()) are estimated as discontinuous, point-wise distributions ensure differentiability during training. We show the details of distribution estimation in Appendix E As Algorithm [ indicates, in the prediction phase, WR-CP follows the inference procedure of importance-weighter conformal prediction (IW-CP) proposed by Tibshirani et al. (2019). When 3 = O, Eq. (19) returns to empirical risk minimization, and thus WR-CP becomes IW-CP.
6 EXPERIMENTS
6.1 DATASETS AND MODELS
Experiments were conducted on six datasets: (a) the airfoil self-noise dataset (Brooks & Marcolini 2014); (b) Seattle-loop (Cui et al. 2019), PeMSD4, PeMSD8 (Guo et al., 2019) for traffic speed prediction; (c) Japan-Prete Prefecturs,, and U.S.-States (Deng et al. 2020) for epidemic spread forecasting. k = 3 for the airfoil self-noise dataset, and & = 10 for the other five datasets. We conducted s(i) 10 sampling trials for each dataset. Within each trails, we sampled Sxy from each subset for
7

---

Published as a conference paper at ICLR 2025
Algorithm 1 Wasserstein-regularizze Conformal Prediction (WR-CP)
Require:
training set S from distribution x for = 1,... k; calibration set S from Pxy: N training epochs; model he: score function s(x, y) = he (x) y loss function I: balancing hyperparameter B.
Training Phase:
I: Ohtain Px and D for 1,.... by kernel density estimation;
for =
do
3: SP = [s(2, y) (r, E &
4:
for = 1 to k do
5:
Obtain D from S {a v) (z, v) E S by point-wise distribution estimation;
dDx (x) where the feature that (æ,y) v) &Xx, - v)
6:
Weight all SP with normalized
ax)
7: Obtain D V.P from the weighted sp by point-wise distribution estimation;
8: end for
9: Optimize he by mino E E (n.iessf [1hh(()), + W D
10: end for
Prediction Phase:
11: Obtain Qx by kernel density estimation;
12: of {=( v) (x.y) v) rx
Weight all e SE with normalized d(x(=) where æ the feature that (x, E Sky, (#)
14: ==1 quantile of the weighted Sp;
15: for (I. y) E 2 do
16: C(z) = (0 a(x, 9) ST E
17: end for
L, k. Given that calibration and training data are commonly assumed follow the same distribution in CP, we sampled Sky from the union of the k subsets. Additionally, we generated 10k: test sets for each dataset in every trial. A multi-layer perceptron (MLP) with an architecture of (input dimension, 64, 64, 1) was utilized in all experimental setups to maintain comparison fairness. The detailed information about datasets and sampling procedure shown in Appendix F. our work is released on https: / thub. com rxu01 12 /WRRCP.
The code of
6.2 CORRELATION BETWEEN WASSERSTEIN DISTANCE AND COVERAGE GAP
We demonstrated Wasserstein distance indicate coverage gap changes across from 0.1 to 0.9 comprehensively, as illustrated in Figure [b). Specifically, for each dataset, he was optimized by empirical risk minimization. Then, applied vanilla conformal prediction and calculated the average value of coverage gaps for a values from 0.1 to 0.9. Meanwhile, we also computed the Wasserstein distances between the calibration and each conformal score distributions. Our findings highlighted a strong positive monotonic relationship between Wasserstein distance and the average value, indicating its sensitivity to coverage gap changes across different a.
each test set
Baselines. Three baseline distance measures were selected. First of all, total variation (TV) distance was chosen as Barber et al. (2023) aimed to use bound coverage gap. Besides, Kullback-Lein (KL)-divergence and expectation differenec (AEE) were selected as they are widely applied in domain adaptation researches (Nguyen et al. 2021 Magliacane et al., 2018).
Metric. We applied Spearman's 's coefficient, -1 -1 -115,, S 1 to quantify the monotonic relationship between distance measures and the average coverage gap. The absolute value of the coefficient represents the strength of the correlation. Its sign indicates if a correlation is positive negative. A higher positive rs means a stronger positive monotonic relation. We show the detailed definition of Spearman's 's coefficient in Appendix F2
Result. Table I presents Spearman's coefficients between distance measures and the average coverage gap across the six datasets, with the standard deviations shown in parentheses. The highest coefficient is bold and the second- highestt coefficient is underlined. The result shows that the Wasserstein distance consistently exhibits a high coefficient, suggesting that Wasserstein distance is an effective indicator of the average coverage gap, and establishing it as a suitable optimization metric for maintaining coverage guarantees across various values.
8

---

Published as a conference paper at ICLR 2025
Table 1: Spearman's 's coefficients between distance measures and the average coverage gap
Dataset
Airfoil
PeMSD4
PeMSD8
Seattle
U.S.
Japan
W
0.59 (0.24)
0.84 (0.03)
0.90 (0.03)
0.84 (0.05)
0.77 (0.06)
0.57 (0.05)
TV
0.45 (0.16)
(0.03)
0.86 (0.06)
0.75 (0.09)
0.67 (0.10)
0.37 (0.06)
KL
0.40 (0.21)
0.49 (0.17)
0.51 (0.09)
0,45 (0.17)
0.60 (0.11)
0.53 (0.05)
AE
0.55 (0.19)
0.78 (0.05)
0.85 (0.04)
0.71 (0.06)
0.68 (0.08)
0.37 (0.09)
6.3 EVALUATION OF WR-CP IN WASSERSTEIN DISTANCE MINIMIZATION
We proved that WR-CP, utilizing importance weighting, can effectively minimize the Wasserstein distances resulting from both concept shift and covariate shift.
Baselines. Besides WR-CP, we also conducted vanilla CP and IW-CP on all sampled datasets.
Metric. These approaches were compared based the Wasserstein distance between calibration and test conformal scores. To place greater emphasis on the vertical coverage gap between conformal score CDFs, the distances were normalized to mitigate the impact of varying score scales across datasets, enabling more meaningful comparisons.
Result. Figure 3 shows that WR-CP consistently reduces Wasserstein distance. The extent of these reductions is dependent on the value of B. However, despite the ability to address covariate-shift induced Wasserstein distance, importance weighting may not always lead a reduction, as seen the case of the Seattle-lopp dataset. Further explanation of the phenomenon is provided in Appendix C
Airfoil
NAMSD4
NiMED8
Seattle
US
Japan
Normaltzed
1.50
5.10
0.62
0.60
0.27
0.50
Vanilla CP
CP
WR.CP
Figure 3: Comparison of vanilla CP, IW-CP, and WR-CP based on normalized Wasserstein distance between calibration and test conformal scores: IW-CP can only address the distance caused by covariate shift, while WR-CP reduces the distance from concept shift. The 33 values for the WR-CP method are 9, 11, 10, 13, and 13, respectively.
6.4 ROBUST AND EFFICIENT PREDICTION SETS BY WR-CP
We experimentally demonstrated that, compared with prior works, WR-CP is capable of reducing coverage gap without significantly sacrificing prediction efficiency.
Baselines. Besides vanilla CP and IW-CP (Tibshirani et al. 2019), conformalized quantile regression (CQR) (Romano al. 2019) was chosen as a representative method for adaptive CP. We also included the worst-case conformal prediction (WC-CP), which is an implement of the worst-case approach proposed by Gendler et al. [202T] Cauchois al, (2024) Zou Liu (2024) in the convex hull setup.
Metric. We compared the coverage gaps and sizes of prediction sets generated by WR-CP and baselines as a ranges from 0.1 0.9 across all sampled datasets. Prediction sets are better when actual coverages are more concentrated around a and have smaller sizes.
Result. With = 0.2, Figure Saare4coom confirms that WR-CP consistently exhibits the most concentrated coverages around 1 compared to vanilla CP. IW-CP, and CQR across datasets. While WC-CP maintains coverage guarantees under joint distribution shift, leads inefficient predictions. In contrast, WR-CP mitigates this inefficiency through smaller set sizes. We show the results with other a values in Appendix F3 have smaller prediction sets than WR-CP. Since WR-CP is trained with the additional Wasserstein regularization term in Eq. (19), the trade-off inevitably causes an increase in prediction errors, which proportional to conformal scores. Consequently, methods based empirical risk minimization,
important to observe that vanilla CP and and IW-CP always

---

Published as a conference paper at ICLR 2025
Airfoil
PeMSD4
PoMSD8
Scattle
VS
Japan
2.0
0.9
0.8
9.8
0.5
6.6
0.6-
07-
0.8
2.05-
3,00
4.00
0.50
Vanilla CP
IW-CP COR
WC-CP
WR-CP
1-
Figure 4: Coverages and prediction set sizes of WR-CP and baselines with & 0.2: WR-CP makes coverages test data concentrated around the level compared to vanilla CP. IW-CP, and CQR. While WC-CP ensures coverage guarantecs, leads to inefficient predictions due large set sizes, whereas WR-CP mitigates this inefficiency. The 3 values
WR-CP method are 4.5, 9, 6, 8, and 20, respectively.
Airfoil
PeM4SD4
PeMSD8
Seattle
US
Japan
Coverage 030
0.10
3.66
0.7-
0.00
0.02
3:04
0.09
0.04
0.04
0.3
Pred.
Pred.
Size
Pred. Size
IW-OP
WR-CP
Figure 5: Pareto fronts of coverage gap and prediction set size obtained from WR-CP with varying 8:: WR-CP effectively balances conformal prediction accuracy and efficiency, providing flexible and customizable solution. When 3 0, WR-CP returns IW-CP.
like vanilla CP and IW-CP, tend to yield smaller prediction sets compared to WR-CP due to their lower conformal scores. We further discuss the trade-off in WR-CP Subsection 6.5 Lastly, we can see IW-CP have worse coverages than vanilla CP on Seattle-loop dataset, reflecting the fact that importance weighting enlarges Wasserstein distance on that dataset in Figure
6.5 ABLATION STUDY
As outlined in Eq. 19,, WR-CP is regulated by a hyperparameter which governs the trade-off between coverage gap and prediction set size. is essential to investigate the performance of WR-CP under different 3 values, which are listed Appendix F.4 To achieve this, we conducted WR-CP on all sampled datasets with varying 3 values. At each 8 value, we calculated the average coverage gap and set size over a from to 0.9. Finally, we obtained a Pareto front for each dataset in Figure In particular, when A O, WR-CP reverts back to IW-CP, so we emphasize the outcomes in this scenario as boundary solutions derived from IW-CP. The results indicate that WR-CP allows users to customize the approach based on their preferences for conformal prediction accuracy and efficiency. We further explore whether WR-CP can achieve efficient prediction with a coverage guarantee in Appendix G
The limitations of our study presented in Appendix H
CONCLUSION
In this work, point out that the coverage gap of conformal prediction under joint distribution shift relies on the distance between the CDFs of calibration and test conformal score distributions. Based on this observation, we propose an upper bound of coverage gap utilizing Wasserstein distance, offering better identifiability of gap changes at different a. We conduct a detailed analysis of the bound by utilizing probability measure pushforwards from the shifted joint data distribution to conformal score distributions. This approach allows us to explore the separation of the impact of covariate and concept shifts on the coverage gap. Based on the separation, we design musserrstei-regglaare prediction (WR-CP) via importance weighting and regularized representation learning, which can obtain accurate and efficient prediction sets with controllable balance. The performance WR-CP experimentally analyzed with diverse baselines and datasets.
conformal
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENT
Sihong Xie was supported by the Department of Science and Technology of Guangdong Province (Grant No. 2023CX10X079), the National Key R&D Program of China (Grant No. 2023YFF0225011,, the Baanghhou-HKUSTCGE Joint Funding Program (Grant No. 2023/0300008), and Education Bureau Guangzhou Municipality.
REFERENCES
Salim I Amoukou and Nicolas JB Brunel. Adaptive conformal prediction by reweighting nonconfor- mity score. arXiv preprint arXiv:2303. 12695, 2023.
Anastasios N Angelopoulos and Stephen Bates. A gentle introduction to conformal prediction and distribution-free uncertainty quantification. arXiv preprint arXiv:2107.07511 2021.
Anastasios N Angelopoulos. Stephen Bates, Adam Fisch, Lihua Lei, and Tal Schuster. Conformal risk control. arXiv preprint wwxiv:2208.028//1 2022.
Liviu Aolaritei, Nicolas Lanzetti, Hongruyu Chen, and Florian Dörfler. Distributional uncertainty propagation via optimal transport. arXiv preprint N0i::2205.00344, 2022.
Rina Foygel Barber, Emmanuel Candes, Aaditya Ramdas, and Ryan J Tibshirani. Conformal prediction beyond exchangeability. The Annals of Statistics, (2):81 1-855, 2023.
Alberto Bernacchia and Simone Pigolotti. Self-consistent method for density estimation. Journal of the Royal Statistical Society Series B: Statistical Methodology, 73(3):407-422, 2011.
Pope D. Brooks, Thomas and Michael Marcolini. Airfoil Self-Noise. UCI Machine Learning Repository, 2014. DOI:
Maxime Cauchois, Suyash Gupta, and John Duchi. Knowing what you know: valid and validated confidence sets in multiclass and multilabel prediction. Journal of machine learning research, 22 (81):1-42,20) 2021.
Maxime Cauchois, Suyash Gupta, Alnur Ali, and John C Duchi. Robust validation: Confident predictions even when distributions shift. Journal of the American Statistical Association, pp. 1-66, 2024.
Nicolo Colombo. Normalizing flows for conformal regression. arXiv preprint arXiv:2406.03346, 2024.
Zhiyong Cui, Kristian Henrickson, Ruimin Ke, and Yinhai Wang. Traffic graph convolutional recurrent neural network: A deep learning framework for network-scale traffic learning and forecasting. IEEE Transactions on Intelligent Transportation Systems, 2019.
Songgaojun Deng, Shusen Wang, Huzefa Rangwala, Lijing Wang, and Yue Ning. Cola-gnn: Cross- location attention based graph neural networks for long-term ili prediction. In Proceedings of the 29th ACM international conference on information & knowledge management, pp. 245 2544, 2020.
Richard Mansfield Dudley. The speed of mean glivenko-cantell convergence. Mathematical Statistics, 40(1):40-50, 1969.
The Annals of
Bat-Sheva Einbinder, Stephen Bates, Anastasios N Angelopoulos, Asaf Gendler, and Yaniv Romano. Conformal prediction robust to label noise. arXiv preprint arXiv:2209. 14295, 2, 2022a.
Bat-Sheva Einbinder, Yaniv Romano, Matteo Sesia, and Yanfei Zhou. Training uncertainty-aware classifiers with conformalized deep learning. Advances in Neural Information Processing Systems, 15:22380-223349 2022b.
Shai Feldman, Stephen Bates, and Yaniv Romano. Improving conditional coverage via orthogonal quantile regression. Advances in neural information processing systems, 34:2060-2071, 2021.
11

---

Published as a conference paper at ICLR 2025
Di Feng, Ali Harakeh, Steven L Waslander, and Klaus Dietmayer. A review and comparative study on probabilistic object detection in autonomous driving. IEEE Transactions on Intelligent Transportation Systems, 23(8)99911-9980 2021.
Rina Foygel Barber, Emmanuel Candes, Aaditya Ramdas, and Ryan J Tibshirani. The limits of distribution-free conditional predictive inference. Information and Inference: A Journal of the IMA, 10(2):455-433 482, 2021.
Robert E Gaunt and Siqi Li. Bounding kolmogorov distances through wasserstein and related integral probability metrics. Journal of Mathematical Analysis and Applications, 522(1): 126985, 2023.
Asaf Gendler, Tsui-Wei Weng, Luca Daniel, and Yaniv Romano. Adversarially robust conformal prediction. In International Conference on Learning Representations, 2021.
Isaac Gibbs and Emmanuel Candes. Adaptive conformal inference under distribution shift. Advances in Neural Information Processing Systems, 14:1666-1 2021.
Isaac Gibbs and Emmanuel J Candès. Conformal inference for online prediction with arbitrary distribution shifts. Journal of Machine Learning Research, 25(162):1-36, 2024.
Isaac Gibbs, John J Cherian, and Emmanuel Candès. Conformal prediction with conditional
guarantees. arXiv preprint arXiv:2305. /2616, 2023.
Leying Guan. Localized conformal prediction: A generalized inference framework for conformal prediction. Biometrika, 110(1)333-50, 2023.
Shengnan Guo, Youfang Lin, Ning Feng. Chao Song, and Huaiyu Wan. Attention based spatial- temporal graph convolutional networks for traffic flow forecasting. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 922-929, 2019.
Xing Han, Ziyang Tang, Joydeep Ghosh, and Qiang Liu. Split localized conformal prediction, 2023. abs wwwpaab......... 13092]
URL http:/ttss/s/xii..o..a
Pierre Humbert, Batiste Bars, Aurélien Bellet, and Sylvain Arlot. One-shot federated conformal prediction. In International Conference on Machine Learning, pp. 14153-14177. PMLR, 2023.
Christopher Jung. Georgy Noarov, Ramya Ramalingam, and Aaron Roth. Batch multivalid conformal prediction. arXiv preprint arXiv:2209. /5/45, 2022.
Wouter M Kouw and Marco Loog. An introduction domain adaptation and transfer learning. arXiv preprint 1rXiv:1812.11880 2018.
David Krueger, Ethan Caballero, -eeen--errik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-disttibution generalization via risk extrapola- tion (rex). In International conference on machine learning, pp. 5815-5826. PMLR, 2021.
Yi Liu, Alexander Levis, Sharon-Lise Normand, and Larry Han. Multi-source conformal inference under distribution shift. arXiv preprint arXiv:2405.09331, 2024.
Charles Lu, Yaodong Yu, Sai Praneeth Karimireddy, Michael Jordan, and Ramesh Raskar. Federated conformal predictors for distributed uncertainty quantification. In International Conference on Machine Learning, pp. 22942-22966 PMLR, 2023.
Sara Magliacane. Thijs Van Ommen. Tom Claassen, Stephan Bongers, Philip Versteeg, and Joris M Mooij. Domain adaptation by using causal inference to predict invariant conditional distributions. Advances in neural information processing systems, 31, 2018.
A Tuan Nguyen, Toan Tran, Yarin Gal, Philip HS Tor, and Atlim Güneş Baydin. Kl guided domain adaptation. arXiv preprint arXiv:2106.0772 077800 2021.
Travis A O' Brien, Karthik Kashinath, Nicholas R Cavanaugh, William D Collins, and John P O' Brien. A fast and objective multidimensional kernel density estimation method: fastkde. Computational Statistics & Data Analysis, 101:148- 160, 2016.
12

---

Published as a conference paper at ICLR 2025
Victor M Panaretos and Yoav Zemel. Statistical aspects of wasserstein distances. Annual review of statistics and its application, 6(1):405-43 2019.
Harris Papadopoulos, Kostas Proedrou, Volodya Vovk, and Alex Gammerman. Inductive confidence machines for regression. In Machine learning: ECML 2002: 13th European conference machine learning Helsinki, Finland, August 19-23, 2002 proceedings 13, pp. 345-356. Springer, 2002.
Fabian Pedregosa, Gaäl Varoquaux, Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, et al. Scikit-learn: Machine learning in python. the Journal of machine Learning research, 12:2825-2830, 2011.
Vincent Plassier, Mehdi Makni, Aleksandr Rubashevskii, Eric Moulines, and Maxim Panov. Con- formal prediction for federated uncertainty quantification under label shift. In International Conference on Machine Learning, pp. 27907-27947. PMIR, 2023.
Yaniv Romano, Evan Patterson, and Emmanuel Candes. Conformalized quantile regression. Advances in neural information processing systems, 32, 2019.
Yaniv Romano, Matteo Sesia, and Emmanuel Candes. Classification with valid and adaptive coverage. Advances in Neural Information Processing Systems, 33:3581-3591, 2020.
Nathan Ross. Fundamentals of stein's method. 2011.
Hyun-Sun Ryu and Kwang Sun Ko. Sustainable development of fintech: Focused uncertainty and perceived quality issues. Sustainability, 12(18):7669, 2020.
Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-casc generalization. arXiv preprint 191.v:1911.00000 108331, 2019.
Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-63):94, 2015.
Silvia Seoni, Vicnesh Jahmunah, Massimo Salvi, Prabal Datta Barua, Filippo Molinari, and U Rajen- dra Acharya. Application of uncertainty quantification artificial intelligence healthcare: review of last decade (2013-2023). Computers in Biology and Medicine, pp. 107441, 2023.
Matteo Sesia, YX Wang, and Xin Tong. Adaptive conformal classification with noisy labels. arXiv preprint arXiv:2309 200002200 2023.
Glenn Shafer and Vladimir Vovk. A tutorial conformal prediction, 2007. URL https:
arxiv. org. abs/070s 0700.
David Stutz, Ali Taylan Cemgil, Arnaud Doucet, et al. Learning optimal conformal classifiers. arXiv preprint ww.xv:22100091990 2021.
Ryan J Tibshirani, Rina Foygel Barber, Emmanuel Candes, and Aaditya Ramdas. Conformal prediction under covariate shift. Advances in neural information processing systems, 32, 2019.
Vladimir Vapnik. Principles of risk minimization for learning theory. Advances neural information processing systems, 4, 1991.
Vladimir Vovk, Alexander Gammerman, and Glenn Shafer. Algorithmic learning in random world, volume 29, Springer, 2005.
Jonathan Weed and Francis Bach. Sharp asymptotic and finite-sample rates of convergence of empirical measures in wasserstein distance. 2019.
Chen Xu and Yao Xie. Conformal prediction interval for dynamic time-series. In International Conference on Machine Learning, pp. 11569, PMLR, 2021.
Ge Yan, Yaniy Romano, and Tsui-Wei Weng. Provably robust conformal prediction with improved efficiency. arXiv preprint 24ti:22004.16511, 2024.
Xin Zou and Weiwei Liu. Coverage-gunee guaranteed prediction sets for out-of-distribution data. In Proceedings of the AAAI Conference Artificial Intelligence, volume 38, pp. 17263-17270 17270, 2024.
13

---

Published as a conference paper at ICLR 2025
A PROOFS OF THEOREMS
A.1 PROOF OF THEOREM !
Proof. We define g by g(mn, x2) (((11)) (n,, U2)- Let Idx be the identity mapping function on X, and let ny be the mapping function to the i-th marginal. The proof follows Proposition 3 in the work by Aolaritei al. (2022).
First, we prove the inclusion that (ʃ 9) #) v) C T( Dw.). Consider E TTu,,)) v), so it is equivalent to prove that (J 9) y (U, which means the marginals of (ʃ # f# and IgV. For any continuous and bounded function Φ y + R, we have
IyL ((n)) d((S x g)47)((n, = Sx d7
(20)
J o df(())) / o(v)) d(f
so we obtain I ((s x 9) #7) - fy and similarly derive T2 (( x g) ()) - UwV.
Secondly, need prove =. ggv) C (J g), With 7
pr, ((
seek 7 E T(a.e) such that (ʃ x # 7 To do so, let 712 = (Idx x f) E [1,, f (1), 723 : > [f g2v), and 731 == (g x Idxx #v L(ggv, As T2# 1 723 f+ll. and 734 = T2# 723 = I#v. Santambrogio (0155) ensures joint probability measure 7 on
X' x X X' satisfying ( T2) # 712- T2 # 723, and (Tg TT) = 734 We
demonstrate that 7 (πı x # is the probability measure we are seeking. For this, we prove y v) with any continuous and bounded function Ø R by
J x2)
JXXXxUXX
((x1) d5 (21, V/, 92,
(21)
- ((x1) d72 v1) / 6(z1) du(x1).
xxx
Jx
Eq. 21 indicates T1#Y μ.. Similarly, can derive 72#7 As result, prove (f x 9 #7 = with any continuous and bounded function Ø y x y + by
S (2)1 x 7)
xxxx +(f(z)). dr
I ((s(1)) (((m)). (22)) (
(22)
)(1 v2) dz V1, x2)
Jy yxy 0(v1, v2) d72a(11.22) / Jxxy 2(v1, d1/vvi (21,
As (f #)# u) = T(S#U, u) and T( Un C (S #(,,,, we obtain (ʃ x #)# (x) =T
V). Finally, we prove Theorem by
" =W(#H,, Ixv) inf T'E((=Hua)) cy(11, U2) 82)
I
inf "'((fx0) ((mu) JyX cy(21, dr (v/ U2)
(23)
inf inf I cylun. U2) d(lf 4) y2)
7El(uu)) pxx I
in inf nEN(uu) Jxxy
cy(s(m)), d7()))
14

---

Published as a conference paper at ICLR 2025
A.2 PROOF OF THEOREM
Proof. Let 7 E Uf) be the pushforward of E D) via function f x J. We can apply
Theorem W(yg, u5) and obtain
(() = inf
(24)
TEI(u.M) xxx
the optimal transport plan for W (H, v) is y", and k. bounds the Lipschitz continuity of f, we have
W(yg & f xxx
((((1),,((xz) dr (21,
(25)
$ Srxx (cx(x1, dy (21, = kWW,u
In Eq. (25). the first inequality holds because 7 may not be the optimal transport plan for W(y vy). and the second inequality holds due to the definition of s.
A.3 PROOF OF THEOREM
Proof. As Wasserstein distance satisfies triangle inequality, (u, v) and W (fm, Dm) follow
W(,,v) < W (fn,.)) + W ((in.)) v) S W W(ins 14) + W (fn, 0m) + ((D,,
(26)
1/0p Given IE[ (m)) S 2ym 11/op and E[ (u, (m S -1/00 from Proposition with probabil- ities at least 1 - é 2nz and 1 - e -2mt, respectively, we have
W(n, Am) S Ann 1/a, 1/a, + t W(v, Dm)) I Aum >um- 1//0 tv.
(27)
It is reasonable to assume the two events in Eq. (27) are independent, so we can apply them to Eq. 26 and thus obtain Eq. 15) with probability at least (1 -2m4,² e 2ma,") )
A.4 PROOF OF THEOREM 4
Proof. We denote Fp. Fu, and Fute the corresponding CDFs of μu, v, and v10 for = 1,.... k.
When two distributions tions equals the area between their CDFs. Therefore, the I-Wasserstein distance between and is given by
the real number set R with Euclidean distance, W of the two distribu-
(x)) W(n,v)= / -F())d))
(28)
Since = EEEe w, we have Fx(x) Eie Wi (x))) As v, vl) and μ are definded on X C TR we can derive
w(, = / (Fp(z) - Fo(x)) do dx= / Fx(x) - Σ
Wi Foc dx
==1
will WxFF (x)) Σ w. Foi (2)
dx
Σ wi (Fu(2)
Futo (x)) dx
xx i=1
(29)
Wi [F((x)) - Faro (x) dr =
Σ
wi
[F)()) - Poco (x)
1=1
Σ
wi W (1,
1=1
15

---

Published as a conference paper at ICLR 2025
B COMPARISON BETWEEN TOTAL VARIATION AND WASSERSTEIN DISTANCE
The total variation (TV) distance between two univariate distributions is defined as half of the absolute area between their probability density functions (PDFs), For instance, given two distributions μ and with PDFs Py and Pr, respectively, on space R2., the TV distance is given by
TV v)
(p(()) - Pv(x)) dx.
(30)
2 R20
In contrast, we expand W (j, v) according to Eq. (8 by
W( = / [F((F) - FF(()) dr=
1 Pu(t) dt - pa(t) (o(t
3200
B20
(31)
R20 !' Pu(t) pu(() dt dr.
The inner integration between 0 and æ indicates Wasserstein distance cares where two distributions and differ, whereas the total variation distance Eq. (30) does not take this into consideration.
We would like to introduce toy example to illustrate further why total variation distance can not consistently capture the closeness between cumulative distribution functions (CDFs). Consider
three conformal score distributions Pv, Q Q on space Ryo with their PDFs:
(2)
PPv (v) = 1, v E [0, 1]:
1 if E 0.9)
Pag (v) = 2
if v (0.9,
[2 if [0,0.04].
Pq (v) (o)
[ if (0.04, 0.96).
Therefore, we calculate (V 2 0.05 and TV Q() = 0.04, while W(Pv, = 0.0025 and W(Pv, Q( 0.0384. In this example, a reduction in total variation distance results in a larger Wasserstein distance between two CDFs. Intuitively, TVD only measures the overall difference between two distributions without accounting for the specific locations where they diverge. In contrast, the Wasserstein distance will be high when divergence occurs early (i.e., at small quantile), especially if the discrepancy persists until the "lagging" CDF catches up. We visualize the example in Figure
Py and 0
Py and of
2.5
Py-
1.0-
2.0
0.8
1.5
0.6
1.0
0.4
P
0.2
0.0
0.0
PDF: P. and o!"1
CDF: P. and 0f
2.0
oZ
1.5
1.0
0.5
0
0.0
0.7
0.4
0.5
0.8
1.0
0.0
0.2
0.A
0.6
1.0
Figure 6: Comparison between total variation distance and Wasserstein distance: reduction in the total variation distance does necessarily result in CDFs becoming closer.
16

---

Published as a conference paper at ICLR 2025
C RATIONALE FOR DIFFERENTIATING COVARIATE AND CONCEPT SHIFTS
There are two key reasons to differentiate between covariate and concept shifts. First, making this distinction enables the application of importance weighting. Minimizing the Wasserstein regulariza- tion term inevitably increases prediction residuals. By applying importance weighting, we expect to reduce the distance, mitigating the adverse effects regularization on optimizing the regression loss function in Eq. [19). Figure 3 shows this expectation is met five out of the six datasets. This occurs because, in most cases, covariate shifts exacerbate the distance caused by concept shifts (fp * fo). Consequently, importance weighting effectively reduces this distance, as illustrated in Figure aa and evidenced by the results for the airfoil self-noise, PeMSD4, PeMSD8, U.S.-States and Japan-Prefecture, datasets in Figure 3 However, there are instances where covariate shifts can alleviate the Wasserstein distance induced by concept shifts. In such cases, applying importance weighting may incraase the distance, as demonstrated in the results for the Seattle-looy dataset in Figure This phenomenon is further illustrated in Figure bb).
(b)
Cunddatte Frababilty
Ce sampl givn mosd
DumTuaat Probasttt
fg
Pr
Art dlringg
Pry P. sarple Okiss rodel
Curubrtvee Potsecityy
wegtted model
Cumtirine Pro
fe
Figure 7: Efffect of importance weighting on Wasserstein distance: (a) Scenario where importance weighting reduces Wasserstein distance; (b) Scenario where importance weighting enlarges Wasserstein distance.
Secondly, in multi-source CP, different training distributions D can suffer from different degrees of covariate and concept shifts. Importance weighting allows the regularized loss in Eq. (19) to minimize the distance between training conformal score distribution Dif) and its correspondingly weighted calibration conformal score distribution V.sp. so the model can be more targeted on those whose remaining Wasserstein distances are large. Also, since various non-exchangeable test distributions will weight calibration conformal score distribution differently in the inference phase, prediction set sizes can be adaptive to different test distributions. In contrast, without importance
weighting, the model can only regularize Σ W(Pv, D( and use the same quantile of Pv to generate prediction sets for samples from all test distributions, resulting in the same prediction set size and lack of adaptiveness.
To further demonstrate the two reasons we mentioned above, we modify Wasserstein-regullrizzii
based on unweighted calibration conformal scores (i.e. Dt W (Pv, D())) )) during training. Also, the weighting operation in the prediction phase in Algorithm ii removed accordingly. This method denoted WR-CP(uw). We performed WR-CP(uw) on the sampled data from the 10 trials of each dataset at a 0.2 and compared its results with those of WR-CP.
The comparison is depicted in Figure Although the average coverage gaps between WR-CP and WR-CP(uw) are quite similar, at 3.1% and 2.3% respectively, the average prediction set size for WR-CP is 28.0% smaller than that of WR-CP(uw). This observation proves our first reason that importance weighting effectively reduces the Wasserstein distance between calibration and test conformal scores. By doing so, mitigates the side effect of optimizing the regularized objective function in Eq. 19), which increases prediction residuals. Since larger residuals result in larger prediction sets, reducing residuals directly helps minimize prediction set size. Additionally, the standard deviations of the prediction set sizes observed in WR-CP(uw) are typically smaller than those found in WR-CP. This proves the second reason that removing importance weighting will make prediction sets less adaptive different test distributions.
17

---

Published as a conference paper at ICLR 2025
Airfoi
PeMSD4
PcMS08
Scattle
VS
J3pan
1.5
1.0
1.0
1.0
13-
3.5
2.8-
9.4
0.4
0.4
1.02
0.50
0:25
WR.CP
WR-CPIUWI
1-e
Figure 8: Comparison between WR-CP and WR-CP(uw) at a 0.2, Both methods were implemented using the same values of 4.5, 9, 9, 6, 8, and 20 across the datasets.
D GEOMETRIC INTUITION OF r]
To provide geometric intuition of 7, we expand the definition r as
(sp(01) sQ(22)
n #1
(fp(z1) fz(22)]
((11,//((1)) (12, (Q(22))
(32)
#1,8
[fr(z1) - fq (1z))
[1((a)) - fp(x1)] 11(x2) fQ(
21
- fQ(222)
We first simplify the definition by assuming ø1 x2, the denominator is the absolute difference between two ground-truth mapping functions fr and fo at æ1, and the numerator the absolute difference of the residuals of fr and fo with given model h T1. is the largest ratio between the two absolute differences. A small r) means even if fp and fo differ significantly. results in similar prediction residuals on fP and fq. When #ı A x2, 7) the largest ratio of the absolute differences at two positions, æi and 22, so a small v means that / can lead to similar residuals when f( and fx(2)) differ. The expanded definition above includes both ø1 x2 and 21 22 conditions and Figure (a) and (b) present the two conditions, respectively. Intuitively, the residual difference caused by concept shift will constrained by r).
(a)
(b)
(x)()) - f((r)
11-04)-46001
Sola,
9g(xr)
sp(Hi)
[p((t))
X1 X2
X2
33
Figure 9: Geometric intuition 7 when (a) æ1 32 and (b) œ1 o Jzt Intuitively, the residual difference caused by concept shift will he constrained by
18

---

Published as a conference paper at ICLR 2025
E DISTRIBUTION ESTIMATION
E. KERNEL DENSITY ESTIMATION
Px and D for i L,..., are obtained by kernel density estimation (KDE), and based on the estimated distributions we calculate the likelihood ratio.
In our experiments, we applied the Gaussian kernel, which is a positive function of z E X' C R given by
lidl
K (xxbb)
26
(33)
(V2rb)
where ||-|| is Euclidean distance and is bandwidth. Given this kernel form, the estimated probability density, denoted by P, position za within group of points "1,",,, is
a K)
K( Iz, 6).
(34)
To find the optimized bandwidth value of Px and D for i 1, ..... & on each dataset, we applied the grid search method with bandwidth pool using scikit-learn package (Pedregosa et al. 2011). With the approximated marginal distribution densities, we can calculate the likelihood ratio to implement the weighting technique proposed by Tibshirani et al. (2019).
E.2 POINT-WISE DISTRIBUTION ESTIMATION
D and DVrr for = 1...., h are estimated as discontinuous, point-wise distributions to ensure differentiability during training. Specifically, as D and (() Visp are conformal score distributions on real number set R. W(D D( is equal to area between their CDFs, Eq. (28) shows. Hence, focus is estimating the CDFs of D and V. for 1,.,,.
For the details of point- woie distribution estimation, consider we have 21, ... drawn from probability measure y in space X' C TR so the approximated CDF of js is given by
1 a E 8z, 1sse
Fa(x)
(35)
where 1 is the indicator function and 62, represents the point mass at v, (i.e., the distribution placing all mass at the value x,). other words, Eq. (35) counts the partition samples that are smaller than This point-wise estimation ensures that the Wasserstein- I distance between the estimated distributions is differentiable.
F SUPPLEMENTARY EXPERIMENTAL INSIGHTS
F.1 DATASETS
The airfoil self-noise dataset from the UCI Machine Learning Repository (Brooks & Marcolini 2014 was intentionally modified to introduce covariate shift and concept shift among them. It includes 1503 instances. The target variable is the scaled sound pressure level of NASA airfoils, and there 5 features: log frequency, angle of attack, chord length, free-stream velocity, and log displacement thickness of the suction side. To introduce covariate shift, we divided the original dataset into three subsets based on the 33% and 66% quantiles of the first dimension feature, log frequency, and partially shuffled them. Therefore, & for this dataset. We further introduced concept shifts among the three subsets by modifying target values. With following normal distribution N (0, 10), for y in the first set, y+ = y/ 1000 * E; for y in the second set, y+ = y/ę: for y in the third set, y+ = ç. With the modified data, conducted sampling trials to generate 10 randomly sampled datasets.
The Seattle-lopp dataset Cui et al. 2209 as well as the PeMSD4 and PeMSEDS datasets Guo et al. (2019), consist of sensor-observed traffic volume and speed data gathered in Seattle, San Francisco, and San Bernardino, respectively. The data was collected at 5-minute intervals. Our goal for each dataset is to forecast the traffic speed of specifie interested local road segment in the next time step
19

---

Published as a conference paper at ICLR 2025
by utilizing the current traffic speed and volume data from both the local segment and its neighboring segments. Before sampling, we selected 10 segments of interest for each dataset randomly, setting k 10 for them. There natural joint distribution shifts present among these segments because of the varying local traffic patterns.
The U.S.-States and Japan- Prefettreeesturess datasets Deng al. (2020) contain data the number of patients infected with influenza-lid illness (ILI)) reported by the U.S. Department of Health and Human Services, Center for Disease Control and Prevention (CDC), and the Japan Infectious Diseases Weekly Report, respectively. The data in each dataset is structured based on the collection region. Our objective is utilize the regional predictive features, including population, the increase in the number of infected patients observed in the current week, and the annual cumulative total of infections, to forecast the rise in infections for the following week the corresponding region. We also randomly selected 10 regions for both datasets, so k = 10. Due to the diverse regional epidemiological conditions, there are inherent joint distribution shifts among these regions.
For each dataset, we began by sampling X from each subset i, for i = L..., k. without replacement. After this step, we allocated the remaining elements within each subset for calibration and testing purposes. The parts intended for calibration across all subsets were then unified to form Sky- Lastly, to create diverse testing scenarios, we generated multiple test sets by randomly mixing the parts designated for testing from each subset with replacement. For each dataset, conducted the sampling trial for 10 times, and calculated the mean and standard deviation of the results from these trials, as shown in Figure 3 Figure 4 and Figure 5 For efficiency, all CP methods were conducted as split conformal prediction.
We introduce toy example to further illustrate that exchangeability does not hold. Consider we have two training distributions:
D = N [0.0],
(10.0,
0.7 0.7 1 ),7)
D =N LVL..11 -b.. 1 -0.6])
calibration distribution is mixture of these two training distributions with known weights, such as a uniformly weighted mixture (w1 = w2 = 0.5). A test distribution is a mixture of D and D (22) with unknown random weights. To visualize the non-exchangeabliity in Figure I0 assume the unknown test distribution has weights of 0.2 for D ( and 0.8 for ( XY
Training distribution oi!
Training distribution D
Calibration distribution Per
Test distribution Our
X
Figure 10: Calibration and test samples not exchangeable they from different distributions.
F.2 SPEARMAN'S COEFFICIENT
We first provide the definition of Pearson coefficient.
Definition (Pearson coefficient). With n pairs of samples, (zt, y:) for 1....,,, n. of two random variables X and Y, Pearson coefficient, Yp. is calculated as the covariance of the sumples divided by the product of their standard deviations. Formally, it is given by
Ei Z) D)
Tp
(36)
VIL VE (v: (M--T7 -)
where T and y are the means of the samples of X and Y, respectively.
20

---

Published as a conference paper at ICLR 2025
Based Pearson coefficient, the definition of Spearman's coefficient is given as follows.
Definition 6 (Spearman's 's coefficient). With n pairs of samples, (24, yi) for i = 1,.... of two random variables X and Y, letting r(.) the rank function (i.e., r(z1) ) = indicates that In is the third largest sumple umong 21, 21,...;a Spearman's 's coefficient, Ts is defined as the Pearson coefficient between the ranked samples:
E (r(x))))))))))))))))))()))))(( r(T)) (r(ye) r(T))
7s
(37)
V E (r(x)) V Ei (r(w)) r(v))²
where x and y are the means of the samples of X and Y, respectively.
We calculated Spearman's s coefficient between each distance measure and the largest coverage gap in Section 6 to confirm that Wasserstein distance holds the strongest positive correlation compared with other distance measures.
F.3 ADDITIONAL EXPERIMENT RESULTS OF SUBSECTION 6.4
In addition to the results shown Figure we present further experimental findings from Subsec- tion 6.4 with values of 0, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, and 0.9 on Figure 11 12
13
14
15
T6
17
and 18 respectively. Clearly, WR-CP demonstrates the ability to generate more tightly concentrated compared to vanilla CP and IW-CP. Additionally, it yields smaller prediction
coverages near set sizes than the state-of-the-an method WC-CP. These figures also reveal a trend where as the a value increases, WR-CP requires smaller 3 to achieve acceptable coverages around - a, the prediction set sizes produced by WR-CP are closer with those of vanilla CP and IW-CP, as evidenced by the results the PeMSD4 in Figure [1 and Figure 18 This phenomenon could be attributed to the trade-off between conformal prediction accuracy and efficiency under joint distribution shift. The Wasserstein regularization term Eq. (19) tends prioritize aligning smaller conformal scores initially, as it reduces the Wasserstein penalty with a lesser increase in the empirical risk minimization term. Hence, as the hyperparameter I increases, the model gradually aligns larger conformal scores from two different distributions, which will adversely impact the risk-driven term more. When considering a higher a value, the focus is on ensuring that the coverages on test data are close to the smaller indicating the importance of aligning small conformal scores. Consequently, high B value is not necessary this case, leading to smaller prediction set sizes being achieved.
F.4 EXPERIMENT SETUPS IN ABLATION STUDY
To visualize a comprehensive and evenly -eett--dittibuted set of optimal solutions Pareto fronts, utilized WR-CP with varying values of 3 to produce the results depicted in Figure 5 As mentioned in Section 5 it is worth noting that when 3 0, WR-CP reverts IW-CP. The selected / values for the results of Figure 5 are shown in Table
Table 2: B values of WR-CP in ablation study
Dataset
B values
Airfoil
1.5,, 2,5, 3.5, 4.5, 6, 8, 9, 13, 20.
PeMSD4
1, 1.5,, 2,5, 9,9, 11, 15, 20.
PeMSD8
7,9, 17.
Scattle
1, 3,5,4, 5,5, 6, 7, 8, 10, 13, 15, 20.
U.S.
1.5, 2,5,,, 8, 13.
Japan
1,2, 6, 8, 10, 13, 20.
21

---

Published as a conference paper at ICLR 2025
Airfail
PrMSD4
NrMSD8
Seattlee
1.24
2.0
1.0
Ooverage
0.9
5.4
0.8
0.5
2.5
15,00
6:00
10.00
4.00
2.53
Vanilla
COR
WC-EP
WR-EP
1-a0
Figure 11: Coverages and set sizes WR-CP and baselines with a 0.1: The values for WR-CP method are 9, 11, 9, 8, 13, and 20, respectively.
Airfoil
PeMSD4
PeNSD8
Scattle
1.4
1.0
2.8
0.8-
0.5
0.5
2.6-
Sice
0.40
5.00
1.00
0.50
Vanill
IW-CP
COR
WC-CP
WR-CP
1-
Figure 12: Coverages and set sizes WR-CP and baselines with a 0.3: The values for WR-CP method are 3, S, 5, 5, 8, and 13, respectively.
Airfoil
PeMSD4
PeNSD8
Seattle
1.0
Coveraee
0.8
0.75
0.0
0.8
0.7- 0.6
0.6
5.4
0.35
0.4
0.5
2,00
2.00
4.00
107-
CQR
Figure 13: Coverages and set sizes WR-CP baselines with a 0.4: The values for WR-CP method are 3, S, S, 8, and 13, respectively.
Airfoil
PeM504
PeMSD8
Scattle
1.00
Covergee
0.8
0.5
7.6
0.50
0.50
0.5
0.22
0.45
2,50
0.22
103
0.50
CQR
WR-CP
Figure 14: Coverages and set sizes WR-CP baselines with a 0.5: The values for WR-CP method are 3, S, 3. 5, 8, and 13, respectively.
22

---

Published as a conference paper at ICLR 2025
Airfail
PrMSD4
NrMSD8
Seattle
0.6
0.75
Coverage
6.6-
6.75-
0.5
0.30
0.3
6.4
6.35
6.2
Size
0.32
1:50
3,00
100
0.20
1.40
6,50
0.25
Vanilla
COR
WR-EP
Figure 15: Coverages and set sizes WR-CP and baselines with a 0.62 The values for WR-CP method are 3, 5, 3, 5, 8, and 13, respectively.
Airfoil
PeMSD4
PeNSD8
Scattle
1.00
0.6-
0.75
0.4
0.50
0.4
0.50
0.5-
0.22
0,2
0.3
0.25
0.22
0.28
0.30
3.50
0.25
Vanilla
IW-CP
COR
WC-CP
WR-EP
1--0
Figure 16: Coverages and set sizes WR-CP and baselines with a 0.7: The / values for WR-CP method are 2, 2, 2, 5, 8, and 10, respectively.
Airfoil
PeMSD4
PeNSD8
Seattle
Goveraeb
0.22
0.4
5.75-
1.0
2.5
0.3-
0.25
0.20
Prod. Ste
0.15
0.75
0.00
5.40
0.30
3,25
CQR
Figure 17: Coverages and set sizes WR-CP baselines with a method are 2, 2, 5, and 10, respectively.
The values for WR-CP
Airtoil
PeM504
PeMSD8
Scattle
0.3
6.4
0.7
0-0-
0.75
5.4
5.2
3.0
0.0
0.22
5.30
0.05
6,20
0.13
COR
WR-CP
Figure 18: Coverages and set sizes WR-CP baselines a 0.9: The values for WR-CP method are 2. 1, 1, 2, and respectively.
23

---

Published as a conference paper at ICLR 2025
G PREDICTION EFFICIENCY WITH COVERAGE GUARANTEE
Although Wasserstein-regullrizze loss in Eq. 199 offers controllable trade-off with significantly improved prediction efficiency and a mild coverage loss, is worth investigating if this efficiency can be achieved with coverage guarantee. In this section, we first derive a coverage lower bound of WR-CP via the multi-source setup in Appendix G. Then, we show that combination of WC-CP and the Wasserstein-rrggularized loss not achieve small prediction sets with ensured coverage in Appendix G.2
G.1 COVERAGE GUARANTEE FROM MULTI-SOURCE SETUP
Under the setup of multi-source conformal prediction, with as the - a quantile of the weighted calibration conformal score distribution Qv we can derive the coverage gap upper bound by
[FQv.pp (T) - -FQ
Σ F.F D
(T)
Σ F (OFFD((())
(38)
Wel F DV (T) D--ppy b()
sup
IFD Vip
(T) (p)
iE(......k
In other words, the coverage gap test distribution must be less equal to the largest gap among multiple training distributions. Denoting aD = su..e()......) |F Viap
(T) -F D (T)], we
have a coverage guarantee Pr(Yn+1 E xn+1) 2
ap.
The regularization term Σ W(D Vap D in Eq. 119 can minimize ap, and thus making 1- a - &D closer the desired a. It is important to highlight that ap is adaptive to variations in test distribution Qv, as evident from Eq. (38). This adaptivity ensures that the lower bound dynamically adjusts to different Qv. To evaluate the prediction efficiency guarantee, we set 0.1 and computed the corresponding Additionally, we calculated the coverage and prediction set size of WC-CP on each test distribution, using the corresponding guarantee at
WR-CP under this for various test distributions.
an for comparison.
Airfoil
PeMSD4
PeMSD8
Scottle
US
Japan
Covarage
$
+
0.75
2.35
0.75
0.75
0.52
3.25
6.25
0.25
0.66
4:20
10.00
5.00
7.00
2.00
0.22
WC-CP
WR.CP
...:
Lowest Guaranteed Coverage
Figure 19: Coverages and set sizes of WC-CP and WR-CP with coverage guarantee at 1 -
The experiment results are depicted in Figure 19 demonstrating improved prediction efficiency on the PeMSD4, PeMSD8, U.S.-States, and Prefecttrre datasets. However, the efficiency remains almost unchanged on the Seattle-loop dataset and even declines on the airfoil self-noise dataset. This phenomenon can attributed to the regularization mechanism. While WR-CP enhances prediction efficiency by leveraging the calibration distribution to generate prediction sets, regularization inevitably increases prediction residuals, leading to larger prediction sets. These two opposing effects can interact differently depending on the dataset characteristics. When the efficiency gains outweigh the drawbacks of regularization, we observe reduced prediction set size. Conversely, in datasets like the Seattle-loop and airfoil self-noise, the benefits of regularization are outweighed by the increased prediction residuals, resulting in unchanged or diminished efficiency. The averaged prediction set size reduction across the six datasets is 26.9%.
24

---

Published as a conference paper at ICLR 2025
G.2 POOR COMPATIBILITY BETWEEN WASSERSTEIN-
REGULARIZED DASSERSTEIN-REGGULAZZZED LOSS AND WC-CP
Since the WC-CP is conservative post-hoc uncertainty quantification method but the proposed regularized loss Eq. (19 is applied during training, one may consider applying WC-CP upon the model trained by the regularized loss to obtain guaranteed coverage. However, WC-CP and the model are not suitable for complementing each other. While regularization enhances the reliability of calibration distributions, the worst-case approach depends exclusively on the upper bound of - test conformal score quantile, rendering it unable to benefit from regularization. In contrast, the under this condition, the regularization inevitably increases the prediction residuals, which turn increases the upper bound of the test conformal score quantile. Experiment results in Figure 20 demonstrate the analysis, where WC-CP is the worst-case method based on a residual-driven model (same as the WC-CP method in Section 6.4). and Hybrid WC-WR represents applying WC-CP to model trained by Eq. (19..
WC-CP may result in larger prediction
Airtoil
PeM504
PeMS08
Seattle
Japan
1.2
T
1.0
1.0
1.0-
1.3-
T
3.0-
6.9
0.9
0.9
0.9-
5.9
2.8
0.8
0.6
0.5
0.8
0.8
Frod SetSioe
20.00
10.03
10.00
2.02
10.00
5.00
WC-CP
Hybrid WR-WC
1-9
Figure 20: Coverages and set sizes of WC-CP and Hybrid WC-WR with coverage guarantee a = =09..
H LIMITATIONS
H.1 SUSCEPTIBILITY TO DENSITY ESTIMATION ERRORS
Given that Wasserstein regularization relies on sepporaace--eighhee conformal scores, its perfor- mance is greatly influenced by the accuracy of the estimated likelihood ratio obtained through KDE. Inaccurate estimation can significantly impact the effectiveness of WR-CP. For instance, in Figure 4 WR-CP yields larger prediction set sizes with less concentrated coverages the airfoil self-noise dataset compared to other datasets. This can be attributed to the airfoil self-noise dataset having the highest feature dimension (5) and the smallest size of the sampled Sky (500). These challenges in KDE lead to suboptimal performance of WR-CP on the airfoil self-noise dataset when compared to its performance on others.
The main reason for KDE error is numerical instability, which can arise from several factors. A poor choice of kernel a critical contributor; for instance, kernels with sharp edges or discontinuities, such as rectangular or triangular kernels, can result in jagged density estimates and amplify errors near boundaries. Fat-tailed kernels, such as the Cauchy kernel, may assign excessive weight to distant data points, leading to inaccuracies in density estimates and numerical precision challenges. Additionally, the lack of feature normalization can exacerbate the effects of extreme values, skewing the density estimation process and reducing computational stability. Lastly, inappropriate bandwidth selection, either too small (overfitting) or too large (underfittin)), can disrupt the balance between bias and variance, further contributing to instability in the estimation.
In our work, we first adopted the Gaussian kernel, valued for its smoothness and numerical stability. To mitigate the influence of extreme values, applied feature normalization, ensuring more stable density estimation process. Additionally, we conducted a comprehensive grid search to fine-tune the bandwidth, achieving optimal balance between bias and variance for robust and accurate results. The bandwidth candidates were selected from a logarithmically spaced range between 10 2 and 100.5 consisting of 20 evenly distributed values logarithmic scale.
25

---

Published as a conference paper at ICLR 2025
H.2 COMPUTATIONAL CHALLENGES KDE
We applied a grid search approach to identify the optimal bandwidth for KDE, which ensures an effective balance between bias and variance in density estimation. However, this method often involves extensive computational effort, particularly when working with high-dimensional datasets, as it requires repeated calculations over a range of bandwidth values. To address this challenge. Bemnachia-Pigolott KDE (Bernacchia & Pigolottil 201T) introduces an innovative framework that combines Fourier-based filter with systematic approach for simultaneously determining both the kernel shape and bandwidth. This method not only reduces subjectivity in kernel selection but also offers a more efficient computational pathway. Building on this foundation, FastKDE (O Briei et al. 2016) adapts and extends the Bernacchia-Pigoloti approach for high-dimensional scenarios, incorporating optimizations that significantly improve computational speed and scalability. These advancements represent promising directions for mitigating the computational overhead in our own work, where similar strategies could be leveraged to streamline the bandwidth selection process and enhance the overall efficiency of KDE in complex datasets.
H.3 OTHER CHOICES OF THE CALIBRATION DISTRIBUTION
In the experiments conducted in Section 6 we specifically examine the scenario where the calibration
data follows mixture distribution of D X for I,... k with equal weights. However, this may not always be the case in real-world situations. Given that the calibration distribution plays
crucial role in determining the difficulty of minimizing Eq. (19) during training, it is valuable to investigate the performance of WR-CP with a calibration distribution different from a mixture of training distributions.
26


Reviews:
[{"review_id": "tUryMvj0ku", "reviewer": "Reviewer_C65Z", "summary": "This paper addresses conformal prediction under general distribution shift, including both covariate and concept shift. The authors propose using the Wasserstein distance to bound the coverage gap of conformal prediction methods in shifted distributions. Then, they decompose the Wasserstein distance between calibration and test distributions into two parts which relates to the covariate shift and concept shift separately. For scenarios where test samples are available, they use the Wassertein distance between empirical calibration and test distributions to approximate the true Wasserstein distance with concentration bound. Finally, they also consider the case without samples from the test distributions, in a multi-source conformal prediction setting. In this setting, the test distribution is an unknown mixture of training distributions. In this case, the authors introduce a Wasserstein regularizer in the training loss, which is the uniform average of the Wasserstein distance between the pushforward measure of the conformal score of the entire training data and the conformal score of each compoent over that component.", "questions": "(1) After decomposing the Wasserstein distance, the covariate shift component seems to be captured by the Lipschitz constant of the difference between the hypothesis and the training distribution. Is it correct to interpret that a smoother difference between the hypothesis and the training distribution would improve robustness against covariate shift, since intuitively, if this difference is smooth over x, then covariate shift will change the prediction very little? Also, this kappa doesn't depend on the test distribution, so it also means the covariate shift part is easier?\n\n(2) Is there any geometric intuition on the term eta, which captures the concept shift? \n\n(3) When the test distribution is an unknown mixture of the training distribution, the Wasserstein regularizer term seems to ask the conformal prediction has a small average wasseerstein distance over all components. Would it has negative effects on the results if some component with small weights in practice?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "(1) The paper rigorously explores conformal prediction under complex distribution shifts, going beyond covariate shift to address both covariate and concept shifts. \n\n(2) The decomposition of the Wasserstein distance to measure and address different types of shifts is novel and provides a clearer structure for bounding the coverage gap. \n\n(3) The optimization method proposed in this paper by minimizing the Wasserstein distance is very practical. The application of a Wasserstein regularizer in cases where test samples are unavailable (i.e., unknown mixture of training distributions) is a practical contribution, especially for real-world applications. \n\n(3) The paper provides theoretical guarantees and concentration bounds. The authors validate their method across multiple datasets, showing significant improvements in coverage gap reduction and prediction set size.", "weaknesses": "(1) The Wasserstein distance upper bound on the coverage gap is decomposed into two parts. The term relates to the concepts shift, eta depends on the test distribution and is a little hard to interpret for me.\n\n(2) In cases where some training components have smaller weights in the unknown mixture, the uniform averaging of the Wasserstein distance may potentially lead to suboptimal regularization effects.\n\n(3) minor thing: in Line 243, there is no plot (c) in figure 1.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HKUwSOQRo5", "reviewer": "Reviewer_GF9E", "summary": "This paper investigates the effect of distribution shift between calibration and test data on conformal prediction guarantees. Theoretically, they upper bound the total distribution shift by two terms that decouple the effect of covariate and concept shifts. They are able to do so by using Wasserstein distance based upper bound of coverage gap ( which they also introduce.) Furthermore, they propose methods by which they minimize each term in the upper bound. Empirically they show that their proposed algorithm (WR-CP) can reduce coverage gaps and output smaller prediction sets.", "questions": "1. With respect to point 1 in the \"Weaknesses\",  I am concerned about the practicality of the algorithm 1 ( WR-CP), and computational overhead. Further, could you elaborate on how alternate solutions in the existing literature can potentially solve the density-estimation problem - or be integrated within your pipeline ( if any according to your literature-review ). I believe this is a major over-head that can be addressed if the algorithm allows for alternate solutions. \n\n2. With respect to points 2a, 2b, could you elaborate where your contributions stand with respect to the CP pipeline ( as described ).  I believe the paper can benefit from an alternative narrative where the contributions are more on the conformal-based training procedure, and formal gaurantees are provided. Providing formal gaurantees during the trainning phase might be difficult, however applying a calibration procedure with the gaurantees to the \"changed\" model could mitigate this issue.\n\n3. with regard to point 2c above, I am concerned about \"the fairness\" in the comparison with the baselines in which coverage is gauranteed. \n\n4. I believe the related works should include conformal training based procedures where the model parameters are explicitly changed in order to optimize a conformal-based efficiency, in order to clearly place your contributions with respect to the literature. Further, the application to Multi-source conformal prediction calls for a more thorough literature review in this context. there are works that use a similar framework and address distribution shift in a multi-agent CP framework, or use a similar global distribution to introduce distributed/federated conformal procedures. None of those were addressed in this work.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. They address an important problem in CP, namely decoupling the effect of covariate shift and concept shift on the coverage gap.\n2. They upper bound the coverage gap by adopting the Wasserstein distance over the space of probability distribution of conformal scores. This allows to indicate changes across different values of mis-coverage rate $\\alpha$.\n3. The upper bound is then decomposed to explicitly depend on two terms where the effect of covariate/concept shift is decoupled. This is an important and novel contribution. The theory seems to extensively validate the results.", "weaknesses": "1. The minimization of the term corresponding to covariate shift, seems to just use previous contributions in the field (Tibshirani et. al 2019). Furthermore the optimization algorithm heavily depends on good density estimation, which is very difficult to do in practice.\n\n2. The insights from the newly derived Wasserstein-based upper bound on the coverage-gap are the basis of the optimization procedure introduced in algorithm 1 (WR-CP). To this end:\n\n2.a) There are no formal robustness/coverage guarantees on the final prediction sets. The way the problem is introduced, this becomes a major limitation.\n\n2.b) The problem is considered within the context of \"regression\" tasks. However, when exploring the literature including classification tasks, we can see a clear distinction between training phase (where the predictor model parameters are changed in order to optimize a conformal-based loss ( e.g (Stutz et. al 2022, Yan et. al 2024)) , and the calibration procedure where the model is treated as a black-box and formal coverage guarantees are the main point. The same pipeline seamlessly extends to the regression setting as well. I do not believe the authors have made a clear distinction where their contribution lies in this context. \n\n2.c) The experiments compare the proposed algorithm WR-CP to Worst-Case-CP ( WC-CP) to indicate a reduction in length-inefficiency of the prediction sets. I do not believe the experiments are a fair comparison, as the model parameters are \"changed\". The fair comparison would apply a fixed calibration procedure, with guarantees to both models ( one baseline trained without any insights from the Wasserstein upper bound, one trained according to WR-CP) in order to measure the gains in terms of coverage gap or the length-inefficiency. To indicate this point further, (https://lilywenglab.github.io/Provably-Robust-Conformal-Prediction/) applied \"Conformal Training\" (Stutz. et. al) where the model parameters are changed similar to WR-CP, on top of an algorithm that certifies robustness in case of adversarial noise, and reduce the inefficiency of the prediction sets. Then when applying a CP method with guarantees on top of the new model, we obtain smaller sets. This procedure seems sound as it does not break the main point of CP ( coverage guarantees) and yet reduces inefficiency. The current experimental set-up of this paper does not support their claim ( (-38%) reduction in the prediction set sizes than the worst-case approach on average), and are insufficient. Based on this, I don't believe the experiments fairly and constructively asses performance of WR-CP compared to the SOTA or the prior works. \n\n5. very very Minor typos ~ there is no fig 1(c) , line 153 \"adaption\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pFiANt6xzY", "reviewer": "Reviewer_HBTg", "summary": "Conformal Prediction (CP) sets may not have the guaranteed lower-bounded coverage if calibration and test samples are not exchangeable. The authors derive bounds for and optimize the gap between nominal and empirical coverage in such a distribution-shift setup. Unlike existing work, the bounds and the optimization loss are given in terms of the Wasserstein Distance.", "questions": "- If the goal is to optimize the size of the prediction intervals, why is it important to distinguish between covariate and concept shift? And why should one avoid the worst-case approach of an analogous TVD minimization? \n- Why does the TVD fail *to indicate coverage gap changes thoroughly\" and  *is agnostic about where two distributions diverge*? The WD integrates over local divergences and does not explicitly depend on $\\alpha$. Can one prove that the TVD fails to summarize the distance between the quantiles of two distributions?\n- Are the test and calibration samples of Section 5 non-exchangeable?\n- Is the loss in Eq.19 tractable in practice? Doesn't the KDE cause numerical instabilities for bad kernel choices?  \n- Have you compared the proposed optimization with a naive CP baseline trained on all available original and shifted samples?\n- Do the proofs of Theorems 1 and 2 depend on the specific conformity measure?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- Improving CP under distribution shifts is timely and challenging. The provided bounds and algorithm may increase CP applicability to new kinds of data.\n- Optimizing the uncertainty estimation of a CP algorithm is an increasingly popular topic. It is interesting to focus on reducing a model's uncertainty instead of improving the accuracy of its point predictions.\n- Looking at CP-aware optimization as the joint optimization of a standard objective and a CP-aware regularization term.\n- The proposed penalty term does not seem to involve the (non-differentiable) sample quantiles of the conformity score distribution.", "weaknesses": "- Using WD to optimize the prediction intervals is justified by saying that the Total Variation Distance (TVD) *ignores where two conformal score distributions differ* by focusing on the worst-case discrepancy and is location independent. The authors should clarify the difference between the integral and the infimum in the WD definition, Eq.5, and the sup appearing in the definition of TVD.\n- The authors may justify better their idea of separating covariate and concept shift. In practice, such separation requires making strong assumptions on the data-generating distribution. The advantages are less clear.\n- The optimization scheme assumes the availability of samples from the base and the shifted distributions. This is not the case in the standard distribution-shift setup. Access to the test samples at training time would allow a more straightforward approach where the CP algorithm is trained on all samples and hence not affected by distribution shift.\n- The authors should discuss the exchangeability of the samples in the proposed application of Section 5.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)