🤖 machine learning

Conformal Prediction for Molecular Properties under Label Shift

This paper introduces a conformal prediction framework tailored to label shift that generates statistically rigorous prediction intervals for molecular properties by weighting scores with marginal label probability ratios, thereby enhancing the reliability and regulatory compliance of AI-driven drug discovery without requiring model retraining.

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

Published 2026-08-19

📖 5 min read🧠 Deep dive

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The journey of turning a chemical compound into a life-saving medicine is a marathon of uncertainty. It is a process that can take more than a decade and cost billions of dollars, often ending in failure because a molecule that looked perfect in a computer model behaves unpredictably in a living body. To speed this up, scientists increasingly rely on artificial intelligence to predict how a new molecule will act, such as whether it will dissolve in water or poison a cell. However, these digital predictions are only as good as the data they were trained on. When a model encounters a new type of molecule that differs significantly from the examples it studied, its confidence can become a liability, offering a single, precise number that is confidently wrong. To make these tools safe for real-world use, researchers need a way to measure not just the answer, but the reliability of that answer, especially when the rules of the game change.

In a recent study, researchers at the Mogam Institute for Biomedical Research and Intellicode tackled this specific problem of shifting rules, known in technical terms as label shift. This occurs when the distribution of the properties a model is trying to predict changes between the training phase and the real-world application. For instance, a model might be trained on a wide variety of common molecules but then asked to find rare compounds with exceptionally high potency. In this scenario, the standard way of measuring uncertainty breaks down, often leading to prediction intervals that are too narrow to catch the true value. The team developed a new framework that adjusts the confidence intervals of these AI predictions without needing to retrain the expensive underlying models. By using a statistical technique called conformal prediction, which constructs a range of likely values rather than a single point estimate, they created a system that remains reliable even when the data distribution drifts.

The researchers tested their method using a large dataset of nearly ten thousand compounds with known solubility levels, a critical factor in drug design. They employed a sophisticated language model, trained on hundreds of millions of chemical structures, to make its initial predictions. To simulate the real-world challenge of label shift, they artificially altered the test data so that the distribution of solubility values differed from the training set. When they applied the standard, unadjusted method to this shifted data, the system failed to capture the true values within its predicted ranges as often as it should have, leaving the results dangerously overconfident. The traditional approach, which assumes the test data looks like the training data, simply could not keep up with the change.

To fix this, the team introduced a weighting system that acts as a corrective lens for the predictions. They first divided the available data into three distinct groups: one to train the model, one to estimate how the data distribution had shifted, and a third to calibrate the final prediction ranges. Using the second group, they calculated the ratio between the expected frequency of different solubility levels in the new environment versus the old one. They then applied these ratios as weights to the errors the model made during calibration. This process effectively told the system to pay more attention to the types of errors that were becoming more common in the new data and less attention to those that were fading away. By doing so, they could construct prediction intervals that maintained their statistical validity, ensuring that the true value fell within the predicted range at the desired confidence level, regardless of how the data had shifted.

The results of their simulations, repeated a thousand times to ensure robustness, showed a clear improvement. While the standard method produced intervals that frequently missed the true values under label shift, the new weighted approach consistently restored the coverage to the target level. The researchers found that the method worked best when they used a specific statistical technique called maximum likelihood estimation to calculate the necessary weights, though other methods also showed promise. Interestingly, the most accurate method for recovering coverage tended to produce slightly wider intervals than the others, a trade-off that reflects a more honest assessment of uncertainty. The study demonstrated that it is possible to adapt existing, powerful AI models to new, shifting chemical landscapes without the prohibitive cost of retraining them from scratch.

This work offers a practical path forward for integrating artificial intelligence into the high-stakes environment of drug discovery. By providing a way to generate statistically rigorous confidence measures that adapt to changing conditions, the framework helps bridge the gap between theoretical model performance and regulatory requirements for transparency. It ensures that when scientists look at a prediction for a novel molecule, they are not just seeing a number, but a realistic assessment of its potential, complete with a clear understanding of the risks involved. This shift from blind confidence to measured reliability is a crucial step toward making AI a trusted partner in the billion-dollar endeavor of bringing new medicines to patients.

Technical Summary: Conformal Prediction for Molecular Properties under Label Shift

Problem Statement
The application of Artificial Intelligence (AI) in drug discovery faces a critical reliability bottleneck: distribution shift. While AI models accelerate the prediction of molecular properties (e.g., solubility, potency, toxicity), their performance often degrades when test data diverges from training distributions. Specifically, this work addresses label shift, a scenario where the marginal distribution of the target property ( $P(y)$ ) changes between the source (training) and target domains, while the conditional distribution of features given the label ( $P(x|y)$ ) remains invariant. This is common in drug discovery when research priorities shift toward discovering molecules with property values (e.g., extreme potency or low toxicity) that were underrepresented in the original training data.

Standard conformal prediction, a framework for generating statistically rigorous prediction intervals with finite-sample validity, relies on the assumption of exchangeability between training and test data. Under label shift, this assumption is violated, causing standard methods to produce overconfident yet unreliable predictions with coverage rates falling significantly below the nominal confidence level. Furthermore, conventional point predictions offer limited guidance for high-stakes experimental design, and existing solutions for distribution shift often require costly model retraining or focus primarily on covariate shift rather than label shift in continuous regression tasks.

Methodology
The authors propose a Weighted Conformal Prediction (WCP) framework tailored to handle label shift in molecular property regression without retraining the base model. The methodology proceeds through the following steps:

Data Partitioning: The source dataset ( $D_s$ ) is split into three disjoint subsets to prevent data leakage:
- Training Set ( $D_{train}$ ): Used to train the base prediction model ( $f$ ).
- Weights Set ( $D_{weight}$ ): Used to estimate the importance weights representing the ratio of target to source label distributions.
- Calibration Set ( $D_{cal}$ ): Used to compute nonconformity scores.
- The target dataset ( $D_t$ ) is split into non-shifted and shifted subsets for evaluation.
Discretization and Weight Estimation: Since the target property is continuous, the response variable $y$ is discretized into $K$ bins to create pseudo-labels. Importance weights ( $w$ ) are estimated as the ratio of the target marginal label probability to the source marginal label probability ( $w(\tilde{y}) = q(\tilde{y})/p(\tilde{y})$ ). The authors employ three distinct techniques to estimate the target distribution $q(\tilde{y})$ from unlabeled target data:
- Black Box Shift Estimation (BBSE)
- Regularized Learning under Label Shifts (RLLS)
- Maximum Likelihood Estimation (MLE) with bias-corrected temperature scaling (BCTS).
- For MLE, a probabilistic approach models predictions as Gaussians centered at the model output with standard deviation equal to the RMSE, calculating bin probabilities via Cumulative Distribution Function (CDF) differences.
Weighted Conformal Prediction:
- Nonconformity scores ( $s_i$ ) are computed for the calibration set, typically as absolute residuals ( $|y_i - f(x_i)|$ ).
- Each score is assigned an estimated importance weight ( $\hat{w}_i$ ) based on the bin of the true label.
- A weighted quantile ( $\hat{q}_w$ ) is calculated such that the weighted proportion of scores below the threshold meets the desired confidence level ( $1-\alpha$ ).
- The final prediction interval for a new test point $x_t$ is constructed as $C(x_t) = [f(x_t) - \hat{q}_w, f(x_t) + \hat{q}_w]$ .

Key Contributions

Label Shift Adaptation: The paper introduces a specific adaptation of conformal prediction for continuous regression tasks under label shift, a gap previously underexplored compared to covariate shift.
Model-Agnostic Framework: The method is compatible with various base models (demonstrated here with a Large Language Model) and does not require costly model retraining. It relies solely on re-weighting the calibration step.
Integration of Estimation Techniques: The framework systematically evaluates and integrates BBSE, RLLS, and MLE for estimating marginal probability ratios, demonstrating how these techniques can be applied to correct distributional shifts using only the outputs of predictive models.
Statistical Rigor: The approach restores the marginal coverage guarantee ( $P(y_{test} \in C(x_{test})) \geq 1-\alpha$ ) even when property distributions drift, providing statistically valid uncertainty quantification.

Experimental Results
The framework was evaluated on the AqSolDB dataset (9,982 compounds) using a chemical Large Language Model (LLM) based on a BART architecture, pretrained on ~200 million SMILES strings and fine-tuned for solubility prediction (logS).

Baseline Failure: Standard split conformal prediction (CP) exhibited a significant decline in coverage on label-shifted test data, with average coverage falling substantially below the nominal target.
Performance Recovery: The proposed Weighted Conformal Prediction (WCP) methods successfully recovered coverage. The coverage distributions for WCP shifted to the right relative to standard CP, indicating that the intervals more frequently contained the true labels.
Comparison of Estimators: Among the weight estimation techniques, MLE achieved the highest coverage, followed by RLLS and BBSE. MLE demonstrated robustness regarding the number of bins used for discretization.
Trade-offs: While MLE provided the best coverage recovery, it resulted in slightly wider prediction intervals compared to BBSE and RLLS.

Significance and Claims
The authors position this work as a critical step toward the regulatory compliance and real-world adoption of AI in drug discovery. By moving beyond single-value accuracy to provide actionable, statistically rigorous confidence measures, the method addresses the FDA's 2025 guidance requiring "appropriate confidence intervals" and "uncertainty estimates" for AI systems in regulatory submissions.

The paper claims that this framework enhances the trustworthiness of AI-driven predictions in billion-dollar development pipelines. It enables robust uncertainty quantification even when property distributions drift, directly tackling a pervasive obstacle in applying AI to novel chemical spaces. Ultimately, the work offers a generalizable, model-agnostic solution that ensures models remain trustworthy when navigating the uncertain frontiers of drug development, aligning predictive modeling with the demands for transparency and reliability in high-stakes scientific decision-making.

Technical Summary: Conformal Prediction for Molecular Properties under Label Shift

More like this