Conformal Inference under High-Dimensional Covariate Shifts via Likelihood-Ratio Regularization

Introduction

Conformal prediction (CP) has emerged as a powerful framework for uncertainty quantification. By generating "prediction sets" for black-box machine learning models, CP guarantees that the true label is included with a user-specified probability (e.g., 90%). However, standard CP relies heavily on the assumption of exchangeability—meaning the training and test data must originate from the exact same distribution.

In real-world deployments, models frequently encounter covariate shift, where the distribution of input features changes between training and testing (e.g., a model trained on sunny driving data but tested in the rain). Under these conditions, standard CP loses its coverage guarantees. This paper introduces a novel algorithm, Likelihood-Ratio Regularized Quantile Regression (LR-QR), which maintains valid uncertainty guarantees under covariate shift without explicitly estimating notoriously difficult high-dimensional likelihood ratios.

The Challenge of Covariate Shift

Existing methods attempt to correct for covariate shift by reweighting the data using the likelihood ratio—the probability density of seeing a data point in the test set versus the training set.

The fatal flaw in this approach is that estimating this density ratio is incredibly difficult and statistically inefficient for high-dimensional data, such as images or text embeddings. If the ratio estimate is inaccurate, the resulting conformal prediction sets become unreliable, either severely undercovering (failing to include the true label) or overcovering (producing uselessly large prediction sets). The authors ask a fundamental question: Can we achieve valid conformal coverage under covariate shift without actually estimating this high-dimensional likelihood ratio?

Methodology: LR-QR

The authors propose LR-QR to bypass explicit density estimation entirely. Instead of finding a conformal threshold by plugging a likelihood ratio estimate into a reweighting scheme, LR-QR frames the threshold discovery as a machine learning optimization problem:

1. The Base Loss: The algorithm utilizes the pinball loss, a standard loss function in quantile regression used to find specific percentiles of a distribution.
2. The Novel Regularizer: The authors introduce a custom, data-dependent regularization term. Geometrically, this regularizer forces the learned threshold function to align with the true likelihood ratio implicitly.
3. The Optimization Trick: The mathematical elegance of this regularizer lies in its simplicity; it can be computed using simple 1D sample averages over unlabeled source and target data.

By jointly minimizing the pinball loss (using labeled training data) and this new regularizer (using unlabeled training and test data), LR-QR finds a threshold that naturally adapts to the covariate shift. The regularization strength ($\lambda$) acts as a dial: a small $\lambda$ relies heavily on the labeled source data, while a larger $\lambda$ forces the model to incorporate the unlabeled target data to correct for the shift.

Experimental Validation & Key Results

The authors evaluate LR-QR across diverse, high-dimensional modalities, including tabular data (Communities and Crime), biological images (RxRx1 from WILDS), and Large Language Models (MMLU benchmark).

Expand

Source, p. 8

Across all datasets, LR-QR demonstrated superior performance:

• Robust Coverage in High Dimensions: LR-QR tracked the target coverage reliably. For instance, on the Communities and Crime dataset (where shift is induced by splitting data based on demographic subgroups), LR-QR consistently hit the nominal coverage target. Methods relying on explicit likelihood ratio estimation (like Weighted-CP or Doubly Robust CP) often failed because their ratio estimates degraded in high dimensions.
• Efficiency: On the MMLU LLM benchmark (targeting 90% coverage), LR-QR achieved 89.6% coverage while maintaining the smallest average prediction set size (3.38 options) among all successful methods.
• Resilience to Image Shifts: On the RxRx1 dataset, standard split conformal prediction fluctuated wildly due to batch effects across different laboratory experiments. LR-QR consistently stayed near the 90% nominal coverage across 14 distinct test environments without artificially inflating prediction set sizes.

Theoretical Insights and Future Work

A significant theoretical contribution of this work is proving that valid coverage only requires the threshold function to satisfy specific first-order optimality conditions, which can be achieved via regularization rather than direct density estimation. The authors support this with a novel analysis based on algorithmic stability bounds from learning theory, proving that the empirical LR-QR method converges to the true target coverage up to a small, controllable error.

Limitations: While the authors provide strong general guarantees that coverage error vanishes as data increases, they note that they do not currently have theoretical results that tightly control the "slack" (exact error bounds) for specific, localized scenarios based on the exact structure of the likelihood ratio.

Future work may involve deriving tighter, scenario-specific bounds or extending this regularization technique to other types of distribution shifts, such as label shift.

@inproceedings{joshi2025conformal,
  author = {Sunay Joshi and Shayan Kiyani and George J. Pappas and Edgar Dobriban and Hamed Hassani},
  title = {{Conformal Inference under High-Dimensional Covariate Shifts via Likelihood-Ratio Regularization}},
  booktitle = {Advances in Neural Information Processing Systems},
  year = {2025},
  url = {https://openreview.net/forum?id=2bgwni6Ber},
}