Source-linked AI summary
Conformal Prediction for Molecular Properties under Label Shift
Hyeonsu Lee, Juyeon Kim, Erkhembayar Jadamba, Seungjin Choi, Hyunjin Shin
TL;DR
Molecular property predictors can become overconfident and unreliable under label shift, limiting uncertainty guidance for drug discovery. This paper introduces weighted conformal prediction using target-to-source label probability ratios, and finds that it improves coverage without retraining, with maximum-likelihood corrections achieving the best recovery but wider intervals.
Problem
AI-based molecular property predictors can produce overconfident, incorrect point estimates and lose uncertainty reliability when test property distributions differ from training data.
Method
The framework weights conformal prediction scores by target-to-source marginal label probability ratios estimated from unlabeled target data, avoiding model retraining.
Results
Weighted conformal prediction consistently improves coverage over split conformal prediction under label shift, while MLE achieves the highest coverage but produces wider intervals.
Takeaways & Limitations
The framework provides statistically grounded uncertainty intervals for molecular property prediction under shifting property distributions without costly retraining.
Takeaways & Limitations
Splitting source data into training, weighting, and calibration sets can reduce effective training size and hurt performance in low-sample regimes, while standard intervals are suboptimal for heteroscedastic data.
Abstract
from arXiv · showhide
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.
1 Introduction
Drug discovery is costly and failure-prone, while distribution shift can make molecular-property predictions overconfident and unreliable. This work proposes conformal prediction with label-distribution-ratio weighting to generate rigorous, adaptive intervals without retraining.
- A single new drug can cost $314 million to $2.8 billion, take over a decade, and face failure rates exceeding 90%.
- Distribution shift undermines AI reliability when novel compounds occupy chemical spaces unseen during training, creating overconfident predictions in high-stakes decisions.
- Conformal prediction provides distribution-free, finite-sample-valid intervals that contain the true label with user-specified probability, such as 90%.
- Label shift changes the target-property distribution P(y) while preserving P(x | y), including when research prioritizes underrepresented potency or toxicity values.
- The proposed framework weights conformal scores using target-to-source label-distribution ratios estimated by BBSE, RLLS, or MLE.
- The method mitigates label-shift effects without costly retraining and produces statistically rigorous intervals that adapt to changing molecular-property distributions.
2 Methods
The framework constructs conformal prediction intervals under label shift by estimating target-to-source label probabilities and incorporating them into weighted calibration. It assumes invariant feature distributions conditional on labels and uses weighted quantiles to restore marginal coverage.
- Label shift assumption: The method assumes that p(x|y) remains constant across source and target domains while the marginal label distribution changes.This is the label shift assumption underlying the framework.
- Weight estimation: Continuous responses are discretized into K equally sized bins, used only to estimate importance weights from marginal pseudo-label probabilities.The bin range is determined by the minimum and maximum values of data used to calculate the marginal probability ratios.
- Weight estimation: Each bin’s importance weight is the ratio of its target pseudo-label probability to its source pseudo-label probability.Source probabilities come from held-out empirical frequencies, while target probabilities are estimated from unlabeled target data using BBSE, RLLS, or MLE.
- Weighted calibration: The source data are split into disjoint training, weight-estimation, and calibration sets to prevent leakage between model fitting, shift estimation, and interval calibration.The calibration set supplies nonconformity scores, while the weights set estimates label-shift importance weights.
- Weighted calibration: Weighted conformal prediction assigns label-bin weights to calibration residuals, computes their weighted quantile, and centers the resulting interval on the model’s point prediction.Under the label shift assumption, this weighted quantile corrects distributional shift and restores the marginal coverage guarantee.
3 Experimental Settings
The experiments use AqSolDB to evaluate molecular solubility prediction with a chemically pretrained BART-based language model. A 1,000-repetition split-conformal simulation assesses coverage under continuous label shift using source and target data partitions.
- Datasets: AqSolDB contains 9,982 compounds with experimentally measured logarithmic solubility for benchmarking molecular physicochemical-property prediction.The dataset provides solubility measurements in aqueous solutions and supports drug-design and delivery research.
- Model: The chemical-data LLM uses a BART architecture pretrained on approximately 200 million unlabeled SMILES from six chemical databases.The pretraining sources are ChEMBL, PubChem, ZINC, Enamine, Coconut, and DrugBank.
- Conformal simulation: The simulation applies split conformal prediction to continuous label shift, weighting scores by marginal probability ratios to target exchangeability and minimum coverage of 1 −α.The exchangeability requirement concerns nonconformity-score distributions between training and test datasets when label distributions differ.
- Experimental protocol: 1,000 repetitions divide data into source and target subsets at a 60%–40% ratio, with source data split equally into training, weighting, and calibration sets.The target subset is divided into no-shift and shifted test data to evaluate coverage under both conditions.
4 Experimental Results
Experiments show that standard split conformal prediction loses coverage under label shift, whereas marginal-probability-ratio weighting recovers more reliable coverage. Among BBSE, RLLS, and MLE with bias-corrected temperature scaling, MLE performs best and remains robust across bin counts.
- Split conformal prediction fails under label shift: Standard split conformal prediction exhibited substantially reduced coverage on label-shifted test data relative to non-shifted data.The shifted coverage distribution moved left, with average coverage falling below the nominal target.
- Ratios of marginal probabilities recover coverage loss from label shift: Weighted conformal prediction using marginal probability ratios produced statistically valid, more robust prediction intervals under continuous label shift.Ratios were estimated using BBSE, RLLS, and MLE with bias-corrected temperature scaling.
- Approaches for estimating marginal probability ratios: MLE achieved the highest coverage among the evaluated ratio-estimation methods, followed by RLLS and BBSE.The comparison is reported in Figure 3 for label-shifted test data.
- Approaches for estimating marginal probability ratios: MLE maintained robust coverage recovery across different numbers of bins.The reported improvement was attributed to bias-corrected calibration reducing systematic class bias and to MLE’s theoretical advantages.
5 Limitation
The approach requires separate training, weighting, and calibration splits, which can reduce effective training data and hurt low-sample performance. Standard conformal intervals are also suboptimal for heteroscedastic data, motivating adaptive methods as future work.
- Data requirements: Separate training, weighting, and calibration splits can reduce effective training size and hurt performance in low-sample regimes.Data augmentation methods may mitigate this limitation.
- Interval adaptivity: Standard conformal intervals are suboptimal for heteroscedastic data.Conformalized Quantile Regression could provide more adaptive intervals, but exploring this remains future work.
6 Summary
The paper presents a practical, statistically grounded framework for reliable molecular-property prediction intervals under label shift. It weights conformal prediction with target-label distribution estimates to restore coverage guarantees lost under distribution shift.
- Framework: The framework produces reliable prediction intervals for molecular property prediction under label shift.It is presented as both practical and statistically grounded.
- Label-shift correction: Weighting conformal prediction with target-label distribution estimates restores coverage guarantees lost under distribution shift.The estimates are obtained using BBSE, RLLS, and MLE.
- Evaluation: The method was tested on the AqSolDB dataset with a large-scale pretrained chemical language model.
A Detailed Experimental Settings
The experiments trained all models on NVIDIA A100 SXM4 40GB GPUs, while the approximately 7.5 GB training-memory requirement also permits lower-specification GPUs. Training hyperparameters for BBSE, RLLS, and MLE are provided in Table 1.
- Hardware and hyperparameters: All models were trained on NVIDIA A100 SXM4 40GB GPUs.The stated hardware configuration used NVIDIA A100 SXM4 40GB GPUs for training.
- Hardware and hyperparameters: Approximately 7.5 GB of training memory makes the model feasible to run on lower-specification GPUs.The passage states that the model requires approximately 7.5 GB of memory for training.
- Hardware and hyperparameters: Table 1 reports training hyperparameters for BBSE, RLLS, and MLE.The table is specifically titled as covering hyperparameters for training BBSE, RLLS, and MLE.
B Additional Results
Additional experiments examined how the number of bins affects coverage and interval-length distributions when applying WCP with BBSE, RLLS, and MLE. Results were summarized across 1000 repeated trials using average metrics and distribution plots.
- Coverage recovery performance based on the number of bins: The study varied bin numbers for BBSE, RLLS, and MLE, reporting mean and standard deviation of coverage and interval length over 1000 trials.These results are presented in Table 2.
- Coverage recovery performance based on the number of bins: Fewer bins produced high coverage in some trials, while the overall coverage distribution varied with the discretization used for WCP.The passage indicates that bin count changes the characteristics of the coverage distribution.
- Coverage recovery performance based on the number of bins: Figures 5–7 show KDE distributions of coverage and interval length for standard split conformal prediction and BBSE-, RLLS-, and MLE-based methods across bin numbers.Triangular x-axis markers indicate each distribution’s mean.