Conformal Prediction for Molecular Properties under Label Shift

Hyeonsu Lee; Juyeon Kim; Erkhembayar Jadamba; Seungjin Choi; Hyunjin Shin

cs.LGAug 18, 2026

Conformal Prediction for Molecular Properties under Label Shift

Authors: Hyeonsu Lee, Juyeon Kim, Erkhembayar Jadamba, Seungjin Choi, Hyunjin Shin

Organizations: Mogam Institute for Biomedical Research · Intellicode

Abstract

Drug discovery and development underpins healthcare but remains costly and failure-prone. A critical bottleneck lies in predicting molecular properties such as solubility, potency, and toxicity, which directly determine whether a candidate can advance from preclinical to clinical trials. Artificial Intelligence (AI) has accelerated this process, yet its reliability is often undermined by distribution shift, as experimental conditions frequently diverge from training data. In addition, conventional point predictions provide only single-value estimates, offering limited guidance for high-stakes experimental design. We address these challenges with a conformal prediction framework tailored to label shift. By weighting conformal scores using marginal label probability ratios, our method produces statistically rigorous prediction intervals without retraining. This enables robust uncertainty quantification even when property distributions drift, directly tackling one of the most pervasive obstacles to applying AI in real-world drug development. By moving beyond accuracy alone to provide actionable confidence measures, our approach enhances the trustworthiness of AI-driven predictions. This further aligns predictive modeling with regulatory demands for transparency and uncertainty reporting and ultimately supports more reliable decision-making in billion-dollar development pipelines.

Figures & tables

Figure 1: Schematic diagram of conformal prediction for molecular properties under label shift.

Figure 2: The KDE distributions for 1000 repeated experiments using standard split conformal prediction and the proposed methods are shown. The triangular markers on the x-axis correspond to the mean value of each distribution. The gray and red lines indicate the coverage distributions obtained by applying standard split conformal prediction to the test data without and with label shift (non-uniformly subsampled), respectively.

Figure 3: The KDE distributions for 1000 repeated experiments using standard split conformal prediction and the proposed methods are shown. The triangular markers on the x-axis correspond to the mean value of each distribution. Coverage distributions are shown for the test data under label shift: the red line corresponds to standard split conformal prediction, while the purple, green, and light blue lines correspond to conformal prediction with weights calculated via BBSE, RLLS, and MLE (BCTS), respectively.

Figure 4: The KDE distributions for 1000 repeated experiments using standard split conformal prediction and the proposed methods are shown. The triangular markers on the x-axis correspond to the mean value of each distribution. Interval length distributions are shown for the test data under label shift: the red line corresponds to standard split conformal prediction, while the purple, green, and light blue lines correspond to conformal prediction with weights calculated via BBSE, RLLS, and MLE (BCTS), respectively.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	BBSE	RLLS	MLE
Optimizer	AdamW [ 34 ]	AdamW [ 34 ]	AdamW [ 34 ]
Adam betas	(0.9, 0.99)	(0.9, 0.99)	(0.9, 0.99)
Learning rate	5e-5	5e-5	5e-5
Weight decay	0.1	0.1	0.1
Warmup steps	100	100	100
Error rate $\alpha$	0.1	0.1	0.1

Appendix

Table 1: Hyperparameters for training BBSE, RLLS, and MLE

Bins	BBSE		RLLS		MLE (BCTS)
Bins	Coverage	Length	Coverage	Length	Coverage	Length
5	0.865 $\pm$ 0.033	4.346 $\pm$ 0.316	0.886 $\pm$ 0.034	4.804 $\pm$ 0.590	0.899 $\pm$ 0.034	5.163 $\pm$ 0.832
15	0.845 $\pm$ 0.032	3.993 $\pm$ 0.236	0.877 $\pm$ 0.032	4.577 $\pm$ 0.42	0.895 $\pm$ 0.034	5.031 $\pm$ 0.620
50	0.804 $\pm$ 0.033	3.414 $\pm$ 0.149	0.859 $\pm$ 0.032	4.204 $\pm$ 0.297	0.892 $\pm$ 0.034	4.947 $\pm$ 0.581
100	0.782 $\pm$ 0.033	3.165 $\pm$ 0.123	0.843 $\pm$ 0.032	3.934 $\pm$ 0.222	0.858 $\pm$ 0.031	4.192 $\pm$ 0.242

Appendix

Table 2: Comparison of average coverage and interval length across different bin numbers for BBSE, RLLS, and MLE methods.

Figure 5: The KDE coverage distribution (top) and interval length distribution (bottom) for 1000 repeated experiments using standard split conformal prediction and the proposed methods are shown. Each color represents the number of bins used to calculate marginal probability ratios via BBSE. The triangular markers on the x-axis indicate the mean value of each distribution.

Figure 6: The KDE coverage distribution (top) and interval length distribution (bottom) for 1000 repeated experiments using standard split conformal prediction and the proposed methods are shown. Each color represents the number of bins used to calculate marginal probability ratios via RLLS. The triangular markers on the x-axis indicate the mean value of each distribution.

Figure 7: The KDE coverage distribution (top) and interval length distribution (bottom) for 1000 repeated experiments using standard split conformal prediction and the proposed methods are shown. Each color represents the number of bins used to calculate marginal probability ratios via MLE. The triangular markers on the x-axis indicate the mean value of each distribution.

Explore similar work

Split Conformal Prediction with Label-Shift-Adjusted Bayesian Scores
Sep 11, 2026Hyeonsu Lee, Juyeon Kim, Erkhembayar Jadamba +2Label Shift Conformal Prediction
MOGAM Institute for Biomedical Research, Korea · CROID Research and aSSIST University, Korea
KMM-CP: Practical Conformal Prediction under Covariate Shift via Selective Kernel Mean Matching
Mar 27, 2026Siddhartha Laghuvarapu, Rohan Deb, Jimeng SunWeighted Conformal Prediction Covariate Shift
Siebel School of Computing and Data Science, University of Illinois Urbana-Champaign, IL, USA.
Audited Conformal Prediction for Classification under Unknown Distribution Shift
Jun 12, 2026Yanfei Zhou, Rizal Fathony, Nam H. Nguyen +1Distribution Shift Robustness Uncertainty Quantification
Department of Data Sciences and Operations University of Southern California Los Angeles, CA 90089 · AI Foundations, Capital One New York, NY 10017 · Department of Data Sciences and Operations Thomas Lord Department of Computer Science University of Southern California Los Angeles, CA 90089

Sep 11, 2026J/K move · Enter open · S save

Split Conformal Prediction with Label-Shift-Adjusted Bayesian Scores

Hyeonsu Lee, Juyeon Kim, Erkhembayar Jadamba, Seungjin Choi+1

MOGAM Institute for Biomedical Research, Korea · CROID Research and aSSIST University, Korea

Label Shift Conformal Prediction Regression Prediction Intervals

Open paper AskPDF arXiv

Conformal prediction provides distribution-free uncertainty quantification under exchangeability. However, this assumption is violated by label shift, where the marginal distribution of labels changes while the conditional distribution of inputs given labels remains stable. Under such shifts, standard conformal procedures no longer maintain their intended coverage behavior. Existing approaches address this via importance weighting. They pair the reweighting with residual-based nonconformity scores that ignore predictive uncertainty. The resulting intervals have uniform width. Bayesian conformal methods produce adaptive intervals by leveraging predictive distributions. They evaluate conformity under the source predictive, which is misaligned with the target domain under label shift. We propose the Label-Shift-Adjusted Bayesian Score (LSA score), a nonconformity score derived from a posterior predictive tilting identity. This identity shows that the target predictive is an importance-weighted transformation of the source predictive. We use it to derive a direct correction to the Bayesian score. We evaluate the method on molecular property prediction under controlled label shift. The LSA score consistently yields shorter intervals than residual-based and source-based Bayesian scores. Coverage in the target domain remains comparable. Under stronger shift, all methods incur some coverage loss due to pseudo-label-based density-ratio estimation. The LSA score is defined for any source predictive with a tractable log-density. We instantiate it with Bayesian Ridge Regression, where the correction admits a closed form.

Abstract

Explore similar work

Split Conformal Prediction with Label-Shift-Adjusted Bayesian Scores

KMM-CP: Practical Conformal Prediction under Covariate Shift via Selective Kernel Mean Matching

Audited Conformal Prediction for Classification under Unknown Distribution Shift