Domain-Shift-Aware Conformal Prediction

• Domain-Shift-Aware Conformal Prediction is a framework that adjusts prediction sets using importance weighting to counteract covariate, label, and group shifts.
• It leverages adaptive quantile estimation and density ratio techniques to restore finite-sample coverage even under non-exchangeable training-test conditions.
• Empirical validations in drug discovery, computer vision, and federated learning demonstrate DS-CP's ability to deliver robust, efficient prediction intervals in diverse shifts.

Domain-Shift-Aware Conformal Prediction (DS-CP) is a family of methodologies in conformal inference that systematically adapts prediction set construction to account for domain or distributional shift between calibration (training) data and deployment (test) data. Central to DS-CP is the use of importance weighting, reweighting, or adaptive quantile estimation—addressing violations of the exchangeability assumption required for standard conformal prediction. This paradigm not only enables robust uncertainty quantification when the test and training distributions differ but also offers theoretical and empirical tools for applications encountered in high-dimensional, dynamic, or structured-data regimes, such as covariate shift, label shift, group shift, and geometric transformations.

1. Weighted and Likelihood Ratio-Based Conformal Prediction

The foundational principle of DS-CP is the weighted extension of conformal prediction to accommodate covariate shift, as introduced in "Conformal Prediction Under Covariate Shift" (Tibshirani et al., 2019). In this setting, if the training data are drawn from $P_X$, but test data are drawn from %%%%1%%%% with $P(Y|X)$ fixed, each calibration example is assigned a weight proportional to the likelihood ratio $w(x) = \frac{d\widetilde{P}_X}{dP_X}(x)$. For a test point $x$, the prediction set is determined by forming a weighted empirical measure over nonconformity scores:

$$p_k^{w}(x) = \frac{w(X_k)}{\sum_{j=1}^n w(X_j) + w(x)}, \quad p_{n+1}^{w}(x) = \frac{w(x)}{\sum_{j=1}^n w(X_j) + w(x)}$$

$$\text{Prediction set: } \widehat{C}_n(x) = \big\{y \in \mathbb{R}: V_{n+1}^{(x, y)} \leq \text{Quantile}(1-\alpha; \sum_{k=1}^n p_k^w(x) \delta_{V_k^{(x,y)}} + p_{n+1}^w(x) \delta_\infty) \big\}$$

This approach covers covariate shift by simulating exchangeability between weighted calibration scores and the new test environment. If the ratio is unknown, it can be estimated using density ratio estimation, often via a domain classifier distinguishing training and test covariates (Tibshirani et al., 2019, Laghuvarapu et al., 2023).

The methodology extends to settings such as latent variable and missing data problems, graphical models with structured shifts, and more generally any scenario with "weighted exchangeability," where the data-generating process allows reweighting to restore calibration validity.

2. Specialized DS-CP Frameworks for Diverse Shift Types

Recent work generalizes DS-CP to a variety of distributional shift scenarios:

• Label Shift: For shifts where $p(x|y)$ remains fixed but $p(y)$ changes (label shift), efficient algorithms compute importance weights by inverting the classifier confusion matrix, propagating Clopper–Pearson confidence intervals via Gaussian elimination, and constructing PAC-valid prediction sets under uncertainty (Si et al., 2023).
• Group-Weighted and Stratified Shift: If samples can be partitioned into a finite number of groups and the shift is driven by group proportion changes (e.g., medical multi-site studies), Group-Weighted Conformal Prediction (GWCP) forms empirical distributions and quantile calculations at the group level, yielding explicit and often tighter coverage guarantees than generic WCP:

$$\text{Threshold: } \widehat{q} = \text{Quantile}_{1-\alpha}\left(\sum_{k=1}^K q_k \widehat{P}_\text{score}^{(k)} \right)$$

$$\text{Coverage guarantee: } P\{ Y_{n+1} \in \widehat{C}_n(X_{n+1}) \} \geq 1 - \alpha - \max_{k} \frac{q_k}{n_k}$$

(Bhattacharyya et al., 30 Jan 2024)

• Data Heterogeneity and Federated Learning: In federated or multi-source environments, DS-CP can be instantiated by assigning importance weights using density ratios between each agent's local distribution and a global or target distribution, often via GMM-based estimation and federated quantile computation (Plassier et al., 2023).
• Geometric Shift: For geometric data shifts (rotations, flips), canonicalization networks predict and unwind group transformations. Conformal prediction is performed on canonicalized data or with group-probability-informed weights, maintaining validity even under complex shift patterns (Linden et al., 19 Jun 2025).
• High-Dimensional Covariate Shift: When direct density ratio estimation is infeasible, one can solve a likelihood-ratio-regularized quantile regression problem—projecting the likelihood ratio onto a manageable function space and enforcing coverage via penalized pinball loss optimization (Joshi et al., 18 Feb 2025).

3. Algorithmic Realizations and Practical Implementation

Implementation of DS-CP methods requires several key components:

• Weight Estimation:
  • Density/Class Probability Ratio: Via domain classifiers or energy-based models, estimating $w(x) = \frac{\widetilde{p}_X(x)}{p_X(x)}$ with logistic regression, random forests, or KDE in learned feature space (e.g., energy-based "CoDrug" (Laghuvarapu et al., 2023)).
  • Group/Domain Labels: When group structure is known, empirical counts suffice (Bhattacharyya et al., 30 Jan 2024).
  • Federated Statistics: In federated learning, GMM parameters or sufficient statistics can be shared to secure privacy and computational efficiency (Plassier et al., 2023).
• Score Aggregation and Quantile Computation:
  • Weighted empirical distributions are used for ranking nonconformity scores, with Monte Carlo or closed-form quantile evaluation as necessary.
• Test-Time Adaptation and Regularization:
  • To prevent degenerate behavior (e.g., over-conservativeness), regularization such as capping the test-point weight is applied (Lin et al., 7 Oct 2025). Other methods use entropy scaling, test-time adaptation (EACP) via entropy minimization to adjust uncertainty measures without requiring labeled test data (Kasa et al., 3 Jun 2024).
• Computational Considerations:
  • Split conformal methodologies and federated smoothing techniques enable scalable computation (Tibshirani et al., 2019, Plassier et al., 2023).
  • Tools exist for reducing sensitivity to weight outliers, such as effective sample size corrections (Tibshirani et al., 2019).

4. Coverage Guarantees and Theoretical Properties

DS-CP methodologies are accompanied by rigorous coverage statements:

• Weighted Coverage: The weighted empirical approach ensures that under correct specification (or sufficiently precise estimation) of weights, the marginal coverage in the test domain approaches $1-\alpha$.
• Explicit Gap Bounds: For group-weighted and high-dimensional methods, coverage errors are upper-bounded by easily computable terms, e.g., $\max_k q_k/n_k$ in GWCP (Bhattacharyya et al., 30 Jan 2024), or a stability-controlled gap in LR-QR (Joshi et al., 18 Feb 2025).
• Robustness Under Model Misspecification: Theoretical analyses provide both upper and lower bounds on possible miscoverage as a function of the total variation between calibration and test nonconformity distributions, and regularization controls over-conservativeness (Lin et al., 7 Oct 2025).
• Ball-based Certificates: For worst-case adversarial or domain shift bounded by a known radius, robust CP constructs guarantee $1-\alpha$ coverage within the specified perturbation set (Zargarbashi et al., 12 Jul 2024).

5. Empirical Performance and Applications

Empirical validation of DS-CP is extensive and spans domains:

• Drug Discovery: In small-molecule property prediction under scaffold/fingerprint split and de novo molecule generation, density-ratio-weighted CP (CoDrug) restores valid coverage, reducing gaps by up to 60% relative to unweighted CP (Laghuvarapu et al., 2023).
• LLMs: For LLMs on the MMLU benchmark, DS-CP reduces undercoverage associated with substantial domain (subject) shift, consistently providing more reliable uncertainty estimates than standard CP (Lin et al., 7 Oct 2025).
• Federated/Corrupted Vision: Across federated ImageNet, CIFAR datasets subject to blur or corruption-induced shift, federated DS-CP achieves target coverage with stability and privacy preservation (Plassier et al., 2023).
• Dynamic Environments: In settings with rapidly changing distribution (e.g., continual learning), ensemble DS-CP strategies with adaptive regret guarantees track the best-performing conformal predictor and maintain valid coverage (Hajihashemi et al., 6 Nov 2024).
• High-Dimension/Adversarial: DS-CP methods leveraging regularization or worst-case perturbation offer practical prediction set efficiency while sustaining formal guarantees in image, tabular, and text domains (Zargarbashi et al., 12 Jul 2024, Joshi et al., 18 Feb 2025).
• Physics and Causality: Structural causal modeling (PI-SCM) is exploited to minimize the Wasserstein distance between calibration and test conformal score distributions, stabilizing predictive coverage across heterogeneous physical domains, as shown in traffic and epidemic forecasting (Xu et al., 22 Mar 2024).

6. Limitations, Extensions, and Open Questions

While DS-CP has demonstrated robust performance, certain limitations persist:

• Density Ratio Estimation Accuracy: Performance is sensitive to the quality of estimated weights. In high-dimensional or complex-structured data, this remains a limiting factor (Joshi et al., 18 Feb 2025, Laghuvarapu et al., 2023).
• Choice of Regularization: The parameterization of regularizers (e.g., the test-point constant or entropy scaling) influences conservativeness versus informativeness, and optimal tuning is subject to further research (Lin et al., 7 Oct 2025, Kasa et al., 3 Jun 2024).
• Generalizing Beyond Covariate Shift: Extensions have been proposed for label shift (via importance interval computation (Si et al., 2023)), group shift, and geometric shift (using canonicalization and Mondrian/weighted CP (Linden et al., 19 Jun 2025)), but comprehensive treatment of multiple simultaneous shift types is ongoing.
• Open-Ended Generation and Other Modalities: Scaling DS-CP to open-vocabulary, multi-modal, or sequential prediction tasks (e.g., code/output generation in LLMs) presents additional challenges highlighted as future directions (Lin et al., 7 Oct 2025).

7. Methodological Summary Table

Shift Type	Weighting Mechanism	Coverage Guarantee
Covariate shift	Density/class ratio w(x)	$\geq 1-\alpha$ (with estimated ratio)
Label shift	Inverse confusion matrix	PAC coverage with propagated intervals
Group/stratified	Group frequencies q_k/p_k	$1-\alpha - \max_k q_k/n_k$
Geometric shift	Canonicalization network	Empirical recovery via group-informed weighting
High dimensions	LR-QR regularization	Stability-controlled gap from nominal
Dynamic/adaptive	Adaptive ensemble/expert	Strongly adaptive regret bounds

In summary, Domain-Shift-Aware Conformal Prediction defines a principled framework for uncertainty quantification under distribution shift, leveraging weighted empirical measures, explicit importance weighting, and adaptive quantile procedures to restore finite-sample coverage and maintain efficient, informative prediction sets across a spectrum of modern learning scenarios.