Title: CoDrug: Conformal Drug Property Prediction with Density Estimation under Covariate Shift
Venue: NeurIPS 2023
Accepted: yes
Keywords: drug discovery

molecule property prediction

conformal prediction
Abstract: In drug discovery, it is vital to confirm the predictions of pharmaceutical properties from computational models using costly wet-lab experiments. Hence, obtaining reliable uncertainty estimates is crucial for prioritizing drug molecules for subsequent experimental validation. Conformal Prediction (CP) is a promising tool for creating such prediction sets for molecular properties with a coverage guarantee. However, the exchangeability assumption of CP is often challenged with covariate shift in drug discovery tasks: Most datasets contain limited labeled data, which may not be representative of the vast chemical space from which molecules are drawn. To address this limitation, we propose a method called CoDrug that employs an energy-based model leveraging both training data and unlabelled data, and  Kernel Density Estimation (KDE) to assess the densities of a molecule set. The estimated densities are then used to weigh the molecule samples while building prediction sets and rectifying for distribution shift. In extensive experiments involving realistic distribution drifts in various small-molecule drug discovery tasks,  we demonstrate the ability of CoDrug to provide valid prediction sets and its utility in addressing the distribution shift arising from de novo drug design models. On average, using CoDrug can reduce the coverage gap by over 35% when compared to conformal prediction sets not adjusted for covariate shift.

CoDrug : Conformal Drug Property Prediction with Density Estimation under Covariate Shift

Abstract

In drug discovery, it is vital to confirm the predictions of pharmaceutical properties from computational models using costly wet-lab experiments. Hence, obtaining reliable uncertainty estimates is crucial for prioritizing drug molecules for subsequent experimental validation. Conformal Prediction (CP) is a promising tool for creating such prediction sets for molecular properties with a coverage guarantee. However, the exchangeability assumption of CP is often challenged with covariate shift in drug discovery tasks: Most datasets contain limited labeled data, which may not be representative of the vast chemical space from which molecules are drawn. To address this limitation, we propose a method called CoDrug that employs an energy-based model leveraging both training data and unlabelled data, and Kernel Density Estimation (KDE) to assess the densities of a molecule set. The estimated densities are then used to weigh the molecule samples while building prediction sets and rectifying for distribution shift. In extensive experiments involving realistic distribution drifts in various small-molecule drug discovery tasks, we demonstrate the ability of CoDrug to provide valid prediction sets and its utility in addressing the distribution shift arising from de novo drug design models. On average, using CoDrug can reduce the coverage gap by over 35% when compared to conformal prediction sets not adjusted for covariate shift.

1 Introduction

Drug discovery is a challenging and complex task, with a high failure rate and limited understanding of the chemical and biological processes involved. These contribute to making drug discovery an extremely costly and time-consuming endeavor. Recently, advances in deep learning have aimed to reduce the cost of drug discovery by proposing AI methods for developing accurate property prediction models and De Novo drug design models :

Property prediction models aim to aid the laborious and expensive stages of drug discovery by building accurate supervised learning models that take in a drug representation as input and output a target property [1, 2].

De novo drug design models, on the other hand, aim to discover new drug molecules that satisfy a set of pharmaceutical properties [3, 4, 5, 6].

With the high cost and significance of drug discovery, it is essential to have accurate and reliable uncertainty estimates in supervised learning models for property prediction. By providing set-valued or interval-valued estimates instead of solely relying on point estimates, uncertainty estimation enables more informed decision-making and reduces the risk of failures, making the drug discovery process more efficient. Conformal prediction (CP), pioneered by [7], offers a solution to uncertainty quantification for complex models like neural networks, by constructing provably valid prediction sets 1 in supervised learning models. Its application to drug property prediction has also been explored for various drug discovery tasks [8, 9, 10].

A crucial assumption in the CP framework is that the test samples are exchangeable with the holdout set used to calibrate the algorithm. In drug discovery, there is often a limited amount of available training or validation data. Furthermore, de novo drug design models or screening datasets sample molecules from a large chemical space, making the exchangeability assumption invalid.

In this paper, we deal with a situation where the training data originates from a distribution, P ( X ) , while the test data comes from a different distribution, P test ( X ) . However, in both cases, the molecular properties, determined by the conditional distribution P ( Y | X ) , remain the same as they are governed by nature and unaffected by shifts in the input distribution, assuming testing parameters are stable. This is referred to as covariate shift . Although recent research in Conformal Prediction (CP) [11] suggests a method for correcting covariate shift, accurately estimating the precise level of covariate shift remains a practical challenge.

This paper proposes a novel and practical method for Co nformal Drug property prediction, dubbed as CoDrug , to improve coverage in the conformal prediction framework under covariate shift. We address the problem of non-exchangeability by quantifying the underlying covariate shift at test time and leverage recent advances in conformal prediction to obtain prediction sets. Further, we demonstrate applying CoDrug to obtain valid uncertainty estimates w.r.t. a target property on molecules sampled from de novo drug design models. We summarize our main contributions below:

We propose a novel approach to create prediction sets for drug property prediction, dubbed as CoDrug . Using kernel density estimates (KDE) and recent advances in CP, CoDrug corrects for covariate shift at test time and creates prediction sets whose coverage rate is closer to the target.

We show that the kernel density estimates are consistent, which means that asymptotically, the covariate shift is precisely adjusted for, and the coverage guarantee is recovered.

We demonstrate the loss of coverage in property prediction tasks induced by two forms of distribution shifts - molecular scaffold splitting and molecular fingerprint splitting. Our experiments show that CoDrug effectively reduces the gap between actual and target coverage for prediction sets, with an average enhancement of up to 35% compared to the conformal prediction method without covariate shift adjustment. Additionally, in our experiments on molecules generated by de novo drug design models, we observe a 60% reduction in the coverage gap on average.

2 Related Work

Recently, deep learning techniques have been extensively studied for their potential in drug discovery, specifically in developing accurate predictive and generative models. This led to various architectures for predicting drug properties from SMILES/SELFIES strings [12], molecular graph representations[2, 1] and self-supervised learning [13]. Another area of research focused on building generative models to discover novel molecules using variational autoencoder [5, 6] and reinforcement learning [3, 4, 14].

Furthermore, several methods have been proposed for addressing uncertainty quantification in molecule property prediction, utilizing various Bayesian techniques [15]. Recently, conformal prediction methods have gained increasing attention for drug property prediction [8, 9, 10, 16, 17]. However, these studies primarily focus on generating efficient conformal predictors, without taking into account distribution shifts. Although several benchmarking datasets [18, 19] and methods [20] have been developed for drug property prediction under distribution shift, the problem of uncertainty quantification under distribution shift is still open.

Recent advancement in conformal prediction recovers the coverage guarantee for conformal prediction under known covariate shift [11]. [21] built upon [11] and proposed the Feedback Covariate Shift

(FCS) method for the task of protein design. In practice, one cannot know the exact densities to measure the covariate shift. Like [21], we also leverage [11], but a key difference is that the training density is well-defined in [21] but unknown in ours, requiring us to estimate it. Additionally, our focus diverges from [21] as we concentrate on molecule property prediction rather than protein design.

3 Preliminaries

Reliable estimation of drug properties is crucial for identifying potential drug candidates. Many essential drug properties, such as toxicity, efficacy, drug-drug interactions etc. are formulated as classification problems. Consider a classification task, with each data point Z = ( X,Y ) ∈ R d × [ K ] ( [ K ] = { 0 , 1 , 2 , ..., K -1 } ). For instance, in Fig. 1(a), we seek to construct prediction sets for the problem of solubility classification. (Note that in practice, most drug discovery tasks are formulated as binary classification problems, with K = 2 , but we present the general form of the methodologies.) While building an accurate base classifier ( f ) is important, we usually would like more than a point estimate of the solubility of the molecule, but also some 'confidence level'. This could be encoded in the form of a prediction set denoted as ˆ C ( X ) ⊆ [ K ] .

The main goal we seek in such prediction sets is valid coverage: Given a target (e.g. 90%), we would like to construct a set-valued prediction (Fig. 1(a)) such that, if a molecule is water soluble, this prediction set will include the label 'water soluble' with at least 90% probability. Formally, given 1 -α ∈ (0 , 1) , and a new test molecule ( X N +1 , Y N +1 ) , we would like ˆ C to be 1 -α valid:

Conformal Prediction(CP) framework enables us to achieve such validity in Eq. (1). We will expand the details in Section 4.2. Remarkably, the only requirement of CP is a hold-out calibration set where the base classifier f is not trained on 2 .

One critical assumption for typical CP methods is that the test and calibration data are i.i.d (or exchangeable) which is rarely realistic in drug discovery tasks. On the other hand, although the distribution of molecules X changes from calibration to test time, the conditional distribution Y | X is unlikely to change as the molecular properties are determined by nature and remain the same under similar experimental conditions. Formally, if we denote our calibration set as { ( X i , Y i ) } N i =1 and the test point as ( X N +1 , Y N +1 ) , we have:

It is important to note that the test distribution P test maintains the same conditional distribution P cal Y | X as the calibration distribution, a phenomenon known as covariate shift . This shift is prevalent in de novo drug design models, which require navigating a vast chemical space to pinpoint optimal molecules for a specific goal. However, in many drug discovery tasks, the datasets typically contain only a few thousand data points, representing a limited chemical space. Thus, when models trained on these smaller datasets are used on molecules drawn from the broader molecular space, they inevitably encounter covariate shift. Next we will lay out the exact details of constructing prediction sets with the presence of covariate shifts for supporting drug discovery applications.

4 CoDrug Method

4.1 Overview

In the subsequent subsections, we describe the three primary components of CoDrug . In Section 4.2, we first provide a brief overview of inductive conformal prediction, presenting a method for constructing valid prediction sets in scenarios both without and with distribution shifts, presuming oracle access to the unknown distributions P test X and P cal X . Next, in Section 4.3, we present the details of the training aspects of the base energy-based classifier, emphasizing additional regularization using unlabeled data to enhance its capability to model varying molecule distributions. Finally, in Section 4.4, we employ kernel density estimation (KDE) on the embeddings or logits of the energy model trained in Section 4.3 to estimate the unknown distributions P test X and P cal X , and rectify covariate shift using

(a)ConformalPredictionFramework

(b)CalibrationwithoutCovariateshift

(c)Trainingan Energybased model

X1

C(X1)

{Benzene Soluble)

v(x,y)=1-f,(x,)

Sortedv(x,y)

Training Set

(Water

0.23

0.11

Soluble)

0.23

X2

C(X2)

{Benzene soluble,

v(x,y)=1-f(x,)

0.35

0.43

(Benzene

Soluble)

Water soluble}

0.49

0.49

0.61

(α=0.3)

Unlabelled Set

0.64

70thpercentile

(anjosu)

X3

C(X3)

{lnsoluble}

v(xy)=1-f(x)

0.71

0.67

EBM

XD

0.11

0.82

min B(x,y)~Dn [-log(Fy(x)] +入 · Lenergy

Solubility Classification

6

(d)Density Estimation

111111

Pcar(x)

(e)Calibrationundercovariateshift

Calibration Set

0.35

Calibration set Densit

Sortedv(x,y)

0.30

0.25

Pcal(X)

Ptest (Xx)

W20=0.09

w=0.14

w=0.03

0.11

0.23

0.35

TestSet

Estimation

Density

Kernel

Ptes(x)

0.15

0.10

0.05°

0.00

QOOCiD:

w,=0.11

w=0.21

w=0.02

w=0.01

W=0.04

wg=0.07

0.64

0.43

0.49

0.61

0.67

0.71

70thpercentile

（α=0.3)

Ptest (x)>Ptes（x)

w=0.28

0.82

Trained EBM

Feature Extraction

KDE

Pca(x)<P(X)

W,>W0

W;= Ptest (X) / Pcal(X)

4.2 Conformal Prediction Set

In this section, we will explain how to use conformal prediction to construct valid prediction sets. We will start with the case without covariate shift, and then explain how to correct for covariate shift.

4.2.1 Conformal Prediction without Covariate Shift

Conformal prediction, pioneered by [7], is a powerful framework to construct prediction sets with the guarantee in Eq. (1). We assume that a base classifier f is trained on a training set D train , and we have a hold-out calibration set D cal . To simplify notation, we will denote the calibration set as { Z i } N i =1 and the test point of interest as Z N +1 . We will also abuse the notation to use D to denote both the empirical calibration/test set as well as the underlying distribution. Note that we ignored the training samples because they are no longer used after the classifier f is trained.

We first introduce some useful definitions: (empirical) CDF and quantile function.

The cumulative distribution function (CDF) F of a set of values { v i } N i =1 is defined as:

The quantile function with respect to a CDF F is:

Given a target coverage level 1 -α ∈ (0 , 1) , the Mondrian inductive conformal prediction set (Mondrian ICP) is given by:

where, p f y ( x ) corresponds to the softmax output of class y from model f . Here, { v i } N i =1 are defined by v ( x i , y i ) = 1 -p f y i ( x i ) , which are called 'nonconformity scores' [22] and measure how 'anomalous'

a point z = ( x, y ) is with respect to other points from this distribution. Intuitively, we assign to each molecule a score using the same rule using f , which is trained on a separate data split D train . Now, we choose a threshold t that is larger than 1 -α (e.g., 90%) of the molecules. Because of our i.i.d. assumption, if we sample another molecule Z N +1 from the same distribution, we expect its score to be lower than this threshold with a probability of 1 -α . For eg, in Fig. 1(b), notice how the threshold t , is computed as the value of the Quantile function at α = 0 . 7 ( t = 0 . 64 in this case). We formally state the coverage guarantee without covariate shift , in the following theorem:

Remarks : Theorem 4.1 is a result of classical Mondrian inductive conformal prediction [7]. In fact, in the classification setting, instead of the i.i.d. assumption, one could make a slightly milder assumption that data are exchangeable within each class.

4.2.2 Conformal Prediction with Covariate Shift

While Theorem 4.1 provides a nice first step, the i.i.d. assumption poses a significant limitation in drug discovery. As mentioned in Eq. (3), the distributions of X on D test and D cal can differ. In Eq. (7), we used the empirical CDF F (Eq. (4)) to choose the threshold t . This is because of our i.i.d. assumption: a particular molecule type appears with equal probability/density in both the calibration and test sets. This is no longer the case with covariate shift, which means our F needs to account for such difference in P X .

Formally, recall that P cal X and P test X represent the density of the molecule X for the calibration and test sets. We will assign a weight to each molecule x that is proportional to the density/likelihood ratio dP test X / dP cal X in the empirical CDF, leading to:

W = ∑ N +1 i =1 w ( x i ) is just a normalizing factor. The subscript x N +1 is used to highlight that our updated CDF now depends on the test molecule x N +1 through the weights. Here, w ( x ′ ) could be viewed as a likelihood ratio, and is crucial in adjusting for the covariate shift. For eg, in Fig. 1(e). Notice how the values of w ( x i ) depend on the densities P cal X and P test X . In the figure, the value of weighted Quantile at α = 0 . 3 or the threshold t is shifted from 0.64 to 0.71 to account for the shift. We formally state the modified theorem from [11] that recovers the coverage guarantee under covariate shift for Mondrian ICP:

Theorem 4.2. [11] Assume that ˜ P X is absolutely continuous with respect to P X . For any α ∈ (0 , 1) , let F w be defined as in Eq. (9) , and

However, in practice, both P test X and P cal X are unknown, rendering Theorem 4.2 impractical. In section 4.4, we will provide a viable way to estimate P test X , and recover the guarantee asymptotically (namely with large calibration and test sets). In the following section, we will delve into our training methodology, which harnesses unlabeled data to effectively model varying molecular distributions.

4.3 CoDrug Training Methodology

CoDrug handles distribution shift by proposing an energy-based model formulation [23]. The core idea behind an energy-based model is to construct a function E that maps an input x to a scalar value, known as energy. A collection of energy values can be transformed into a probability density function p ( x ) through the Gibbs distribution

Consider a discriminative neural network f used in a K class classification setting. f ( x ) maps an input x into K real-valued scalars, which are used to derive a conditional class-wise probability:

Then,

where, f y ( x ) refers to the y th logit of the classifier f ( x ) . In this setting, the energy function E ( x ) can be expressed in terms of the denominator of the softmax probabilities in Eq. (14).

Directly using embeddings from a model trained on labeled data may not yield reliable density estimates, as the model lacks knowledge of data outside the training distribution. To overcome this, we co-train the model with unlabeled molecule data. This aids the model f in effectively mapping molecules with distribution shifts to a distinct embedding space. We follow [24], and use an extra regularization term in the loss function to ensure energy separation between in-distribution and out-of-distribution data. The objective function is defined as follows:

where p f y ( x ) refers to the softmax outputs of the classification model f for class y , and D in is the in-distribution training data for which labels are available. The training objective is combined with an additional term L energy given by:

where D out refers to the subset of the unlabelled that is out-of-distribution (OOD). The objective of this term is to enforce a margin of separation between the training samples and the OOD data using the hyper-parameters m in , m out . Particularly, one term penalizes the model if the energy values for in-distribution data are higher than a certain value, while the other term penalizes if the OOD samples have an energy lower than a certain value. In the next section, we will explain how to use either the latent embedding of f or the logits to estimate the density and correct for covariate shift.

4.4 Density Estimation

As discussed earlier, we need to estimate P test X and P cal X to correct for the covariate shift. We resort to Kernel Density Estimation (KDE), a classical nonparametric method, to estimate the density of arbitrary distributions of molecules. For a set of data X 1 , . . . , X n i.i.d ∼ D , KDE is given by:

where K is a fixed non-negative kernel function, and h > 0 is a smoothing bandwidth . Such KDE estimates have nice asymptotic convergence properties, as formally stated in the following theorem:

Theorem 4.3. [25] Assume that the true density p is square-integrable and twice differentiable and that its second-order partial derivatives are bounded, continuous, and square-integrable. If K is spherically symmetric on R d , with a finite second moment, and we choose the bandwidth h such that

where P → means convergence in probability.

Note that commonly used kernels, such as Gaussian kernel, satisfy the requirements.

Since x here refers to molecular entities (e.g. SMILES strings), we cannot use a Gaussian kernel directly. Instead, we use embeddings or prediction logits produced by a trained model f as the input to the kernel. Under the assumption that KDE accurately reflects the true density of the underlying distribution, we could construct kernel density estimators for both the calibration set and test sets (remember that we do not have access to the test labels but have access to the input X ), and use

then as m →∞ ,

to replace the unknown w in Eq. (10), giving us the final prediction set:

where ˆ W = ∑ N +1 i =1 ˆ w ( x i ) is a normalizing factor. Here, ˆ p h test ( · ; D test ) is constructed using samples from the test data with an optimal bandwidth h test chosen on the test data via cross-validation, and ˆ p h cal ( · ; D cal ) is constructed similarly but on the calibration data. It is clear that, as the number of samples from D cal and D test increases, ˆ w converges to w in Eq. (10), and ˆ W CoDrug recovers the coverage guarantee asymptotically. In practice, recovering asymptotic coverage on a finite amount of data is challenging. However, the coverage tends to approach the target value as we observe in our experiments. The overall procedure for density estimation is depicted in Fig. 1(d).

Algorithm 1 summarizes all the components in Section 4. In Section 5, we will verify the efficacy of CoDrug in property prediction tasks, and molecules sampled from de novo drug design models.

Algorithm 1 Procedure for Property Prediction

Training :

Split the dataset into training set D and calibration set D = { z } N .

Train a neural net classifier f on D train

Compute the KDE ˆ p h cal ( · ; D cal ) for all points in D cal

Test Time , for a test set D test

:

Compute KDE ˆ p h test ( · ; D test ) for all points in D cal using Eq. (18). For any x N +1 ∈ D test , compute ˆ w ( x ) and ˆ w ( x i ) for x i ∈ D cal . ˆ CoDrug

Construct the prediction set C ( x N +1 ) using Eq. (22).

5 Experiments

In this section, we put our proposed method, CoDrug , to the test on various drug discovery tasks. Section 5.1 describes the datasets used and key implementation details. Section 5.2.1 empirically demonstrates the loss of validity in conformal prediction sets on different drug discovery datasets. Section 5.2.2 shows how the setup improves the validity of the conformal prediction sets. Section 5.3 confirms the utility of CoDrug in de novo drug design. We include additional details on implementation, datasets, and hyperparameters in the appendix.

5.1 Data and Implementation Details

Splitting Strategies: To demonstrate the effectiveness of CoDrug under covariate shift, we use two different strategies when creating calibration/test splits. In both strategies, we try to create calibration and test splits that are dissimilar to each other, which is a challenging but realistic setting in drug discovery. We used the DeepChem[26] library for splitting. In scaffold splitting , the dataset is grouped based on chemical scaffolds, representing core structures of molecules. The test set and train set consist of different scaffolds. In fingerprint splitting , the dataset is partitioned based on Tanimoto similarity of molecular fingerprints [27]. Molecules with the highest dissimilarity in terms of Tanimoto similarity are included in the test set.

Datasets: Weuse four binary classification datasets for toxicity prediction (AMES, Tox21, ClinTox) and activity prediction (HIV activity), obtained from TDC [28]. To train the Energy based model, we obtained the unlabelled data from the ZINC-250k dataset [29], a subset of the ZINC that covers a large chemical space. For each dataset and split type, we removed the molecules that are similar to the training (and calibration) set from the unlabelled dataset.

Classification Model: The architecture of our classifier f is AttentiveFP [1], a graph neural network-based model. We chose AttentiveFP as it has state-of-the-art results in several drug property prediction tasks. It is trained using the objective function described in Eq. (16).

De Novo Drug Design Experiments: In Section 5.3, we perform experiments to construct conformal prediction sets on molecules sampled from de novo drug design models. As generative models, we use REINVENT[14] and GraphGA[30] (top-ranked methods in MolOpt [31] benchmark). The models are optimized to sample molecules w.r.t. three popularly used computational oracles QED (quantitative estimate of drug-likeness), JNK3 activity, and GSK3B activity. For building the conformal prediction sets, we chose logP as our target property, assigning values in the range of

train cal i i =1 by minimizing Eq. (16). using Eq. (18).

[1.0,4.0] a class of Y=1, and Y=0 otherwise (representing the drug-like range [32]). We obtain the computational oracles from TDC [33], and generative models from MolOpt package[31].

5.2 Property Prediction Results

5.2.1 Unweighted conformal prediction (baseline)

In this section, we demonstrate the unpredictable behavior of the unweighted CP method without proper correction under distribution shift. Table 1 shows the results of conformal prediction under various distribution shift conditions. 'Random' refers to the ideal/unrealistic scenario where the test and calibration samples are split randomly (aka. no distribution shift). 'Scaffold' and 'Fingerprint' denote scenarios in which there is a distribution shift between the test and training data outlined in the Methods section. In all scenarios, 15% training set is held out for calibration, and prediction sets are calculated using the algorithm described in Algorithm 1 without any correction.

From Table 1, we observe that the Random configuration demonstrates little loss in coverage and coverage decreases under distribution shifts (Scaffold and Fingerprint). But for fingerprint and scaffold split, unweighted CP failed to provide target coverage and exhibit unpredictable behavior. For instance, at α = 0 . 2 , under fingerprint split, Unweighted has a coverage of 0.34 against a target coverage of 0.8 for the AMES dataset, while achieving a very different coverage of 0.77 with scaffold split on the same dataset.

Table 1: Unweighted CP's (baseline) coverage under various distribution shifts (or absence thereof) should ideally align closely with the target 1 -α . However, in most datasets with fingerprint and scaffold splits-reflective of more realistic scenarios-the baseline method falls short. Often, substantial deviations in coverage confirm the unpredictability of unweighted CP when exchangeability ceases to apply. Here, values not significantly deviating from 1 -α at a p-value of 0.05 are highlighted in bold, indicating desirable performance.

5.2.2 Weighted conformal prediction improves coverage

In Table 2, we present the benefits of using weighted CP via CoDrug . The table depicts results from conformal prediction using 3 different schemes.

CoDrug (Energy): This variant of CoDrug uses weights computed from KDE on the prediction logits of the trained EBM, as described in Section 4.3.

CoDrug (Feature): This variant of CoDrug builds the KDE instead on the features extracted from the penultimate layer of the trained EBM.

Unweighted: Refers to the unweighted prediction conformal prediction (baseline).

In both weighting schemes of CoDrug , we use KDE to estimate densities and find that weighting using energies improves the coverage towards the target coverage 1 -α in most cases. We notice the highest improvement in the Fingerprint splitting scenario for the AMES(Y=1) category, where the coverage improved from 0.63 to 0.88 (target coverage 0.9). Note that the coverage is 'improved' if it is closer to 1 -α - improvement does not always mean higher coverage, because an unusually high coverage often indicates unpredictable behavior of the underlying model.

While our energy-weighting approach generally improves coverage, there are rare instances where it underperforms compared to the baseline. A prominent example is the ClinTox dataset with Y = 1 , which sees limited improvement or even a reduction in coverage. This is due to the constraints of the density estimation procedure, which relies on the quantity of available data. Notably, this dataset is the smallest and most imbalanced, with only 19 points in the calibration set for class Y = 1 .

Additionally, our results show that using energy weighting leads to better overall coverage than directly weighting the features. This is likely because the energy values are two-dimensional, while the features are an eight-dimensional vector: As the dimension of the feature input to KDE increases, one typically requires more samples to get a high-quality density estimate.

5.2.3 Ablation studies

In this section, we present an analysis of the importance of various components in the CoDrug pipeline - KDE, energy regularization term ( Eq. (16)), and covariate shift correction. In addition to CoDrug (Energy) and Unweighted (baseline) reported in the previous section, we also compare with:

CoDrug (NoEnergy): We use the same protocol as CoDrug (Energy) but the models are trained without the energy regularization term L energy in Eq. (16).

Logistic (Energy): In this experiment, the features are same as CoDrug (Energy), but KDE is not used to estimate densities. Instead, the weights w ( x i ) in Eq. (10), are given by ˆ p ( x i ) / 1 -ˆ p ( x i ) , where ˆ p ( x i ) is obtained by fitting a classifier to features in calibration and test sets (suggested by [11]).

The results are depicted in Table 3. In the table, we compile the "mean absolute coverage deviation" across all the datasets and different random runs from all the experiments reported in Table 2 at different values of α (i.e Mean of | Observed_Coverage -(1 -α ) | across the experimental runs). The results reveal that CoDrug (Energy), the proposed method, is closest to the target coverage in almost all different values of α . We paid close attention to the "Tail 25%", where we presented the metrics for the worst 25% performing experiments and CoDrug (Energy) outperforms all the other variants in comparison by a substantial margin inducting that all the different components in the CoDrug pipeline are helpful. The mean absolute coverage deviation from the target at α = 0 . 1 for CoDrug (Energy) is 0.052, a relative improvement of about 35% over that of Unweighted (0.081).

5.3 Application in de novo drug design.

In this section, we examine CoDrug 's application in de novo drug design models, which navigate a large chemical space to find optimized molecules using a computational oracle. After molecule sampling, validating their experimental properties, such as ADMET (Absorption, Distribution, Metabolism, Excretion, Toxicity), is crucial for safety and efficacy. When a machine learning model trained on such properties is available, assessing the uncertainty associated with the predictions before experimental validation is critical. However, note that the distribution of sampled molecules may substantially deviate from the training data, affecting the prediction sets' target coverage from CP.

In this section, we demonstrate the application of CoDrug on molecules generated by a de novo drug design model. We consider the de novo drug design model as a black box, that has been optimized w.r.t a certain objective. We experiment with two models - GraphGA [30] and Reinvent [14]. To predict properties, we compiled a dataset of logP values, as it can be computed cheaply with a computational oracle. We note that in reality, this dataset could correspond to experimental properties like ADMET. However, since it is not feasible to validate these properties for molecules generated from de novo drug design models, we use logP to demonstrate the method. The results of our experiments are depicted in Table 4 at a target alpha value of 0.1. Our proposed method consistently enhances coverage in all cases and exhibits a substantial improvement over the unweighted version. For example, in the "gsk3b+qed" objective, the unweighted version has a coverage of 0.44 against a target of 0.9, whereas our proposed method improves coverage substantially. The mean absolute coverage deviation from the target at α = 0 . 1 for CoDrug (Energy) is 0.05, a relative improvement of over 60% on the Unweighted version (0.14).

6 Limitations

While we demonstrated that KDE can provide asymptotic coverage guarantees, this may not necessarily translate to improved performance in scenarios with limited sample sizes. In our experiments, we do acknowledge that there are a few instances where the improvement in coverage is limited, where the availability of data is limited. As such, a direction for further research is to explore ways to obtain likelihood ratios that are more data-efficient, particularly in scenarios with smaller calibration sets.

It is worth noting that our current work focuses on addressing the coverage gap in classification tasks, and regression tasks were not included in this study. However, we recognize the importance of uncertainty quantification in regression settings, especially for various critical drug properties represented as regression problems, where our proposed framework can be extended with modifications. Furthermore, our current work primarily focuses on small molecules, yet covariate shift is a common phenomenon in various chemical and biological contexts. While this means that our framework could be more generally applied, obtaining high-quality feature vectors for computing likelihood in different applications remains a challenge that warrants further research.

7 Conclusion

We present a new method for uncertainty quantification in drug discovery, CoDrug , that effectively addresses the problem of co-variate shifts in test data. The proposed method involves a combination of three key steps, training an energy-based model for feature extraction and base classification, performing density estimation using KDE, and use the KDE to correct for covariate shift in conformal prediction to recover valid coverage. The results obtained in this study demonstrate the effectiveness of CoDrug in predicting valid conformal prediction sets and its utility in de novo drug design experiments. Our current work is limited to classification tasks in small molecules, but exploring its application to other chemical and biological tasks with covariate shifts is interesting for future work, including adapting the framework for regression tasks.

Acknowledgments and Disclosure of Funding

This work was supported by NSF award SCH-2205289, SCH-2014438, and IIS-2034479. This project has been funded by the Jump ARCHES endowment through the Health Care Engineering Systems Center.

References

Zhaoping Xiong, Dingyan Wang, Xiaohong Liu, Feisheng Zhong, Xiaozhe Wan, Xutong Li, Zhaojun Li, Xiaomin Luo, Kaixian Chen, Hualiang Jiang, et al. Pushing the boundaries of molecular representation for drug discovery with the graph attention mechanism. Journal of medicinal chemistry , 63(16):8749-8760, 2019.

Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning , pages 1263-1272. PMLR, 2017.

Tianfan Fu, Wenhao Gao, Connor W Coley, and Jimeng Sun. Reinforced genetic algorithm for structure-based drug design. arXiv preprint arXiv:2211.16508 , 2022.

Jiaxuan You, Bowen Liu, Zhitao Ying, Vijay Pande, and Jure Leskovec. Graph convolutional policy network for goal-directed molecular graph generation. Advances in neural information processing systems , 31, 2018.

Rafael Gómez-Bombarelli, Jennifer N Wei, David Duvenaud, José Miguel Hernández-Lobato, Benjamín Sánchez-Lengeling, Dennis Sheberla, Jorge Aguilera-Iparraguirre, Timothy D Hirzel, Ryan P Adams, and Alán Aspuru-Guzik. Automatic chemical design using a data-driven continuous representation of molecules. ACS central science , 4(2):268-276, 2018.

Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Junction tree variational autoencoder for molecular graph generation. In International conference on machine learning , pages 2323-2332. PMLR, 2018.

Vladimir Vovk, Alexander Gammerman, and Glenn Shafer. Algorithmic learning in a random world . Springer US, 2005.

Jin Zhang, Ulf Norinder, and Fredrik Svensson. Deep learning-based conformal prediction of toxicity. Journal of chemical information and modeling , 61(6):2648-2657, 2021.

Isidro Cortés-Ciriano and Andreas Bender. Concepts and applications of conformal prediction in computational drug discovery. arXiv preprint arXiv:1908.03569 , 2019.

Isidro Cortés-Ciriano and Andreas Bender. Deep confidence: a computationally efficient framework for calculating reliable prediction errors for deep neural networks. Journal of chemical information and modeling , 59(3):1269-1281, 2018.

Ryan J. Tibshirani, Rina Foygel Barber, Emmanuel J. Candes, and Aaditya Ramdas. Conformal prediction under covariate shift, 2020.

Mario Krenn, Florian Häse, AkshatKumar Nigam, Pascal Friederich, and Alan Aspuru-Guzik. Self-referencing embedded strings (selfies): A 100% robust molecular string representation. Machine Learning: Science and Technology , 1(4):045024, 2020.

Yu Rong, Yatao Bian, Tingyang Xu, Weiyang Xie, Ying Wei, Wenbing Huang, and Junzhou Huang. Self-supervised graph transformer on large-scale molecular data. Advances in Neural Information Processing Systems , 33:12559-12571, 2020.

Marcus Olivecrona, Thomas Blaschke, Ola Engkvist, and Hongming Chen. Molecular de-novo design through deep reinforcement learning. Journal of cheminformatics , 9(1):1-14, 2017.

Lior Hirschfeld, Kyle Swanson, Kevin Yang, Regina Barzilay, and Connor W Coley. Uncertainty quantification using neural networks for molecular property prediction. Journal of Chemical Information and Modeling , 60(8):3770-3780, 2020.

Fredrik Svensson, Natalia Aniceto, Ulf Norinder, Isidro Cortes-Ciriano, Ola Spjuth, Lars Carlsson, and Andreas Bender. Conformal regression for quantitative structure-activity relationship modeling-quantifying prediction uncertainty. Journal of Chemical Information and Modeling , 58(5):1132-1140, 2018.

Jiangming Sun, Lars Carlsson, Ernst Ahlberg, Ulf Norinder, Ola Engkvist, and Hongming Chen. Applying mondrian cross-conformal prediction to estimate prediction confidence on large imbalanced bioactivity data sets. Journal of chemical information and modeling , 57(7):15911598, 2017.

Yuanfeng Ji, Lu Zhang, Jiaxiang Wu, Bingzhe Wu, Long-Kai Huang, Tingyang Xu, Yu Rong, Lanqing Li, Jie Ren, Ding Xue, et al. Drugood: Out-of-distribution (ood) dataset curator and benchmark for ai-aided drug discovery-a focus on affinity prediction problems with noise annotations. arXiv preprint arXiv:2201.09637 , 2022.

Shurui Gui, Xiner Li, Limei Wang, and Shuiwang Ji. Good: A graph out-of-distribution benchmark. arXiv preprint arXiv:2206.08452 , 2022.

Kehang Han, Balaji Lakshminarayanan, and Jeremiah Liu. Reliable graph neural networks for drug discovery under distributional shift. arXiv preprint arXiv:2111.12951 , 2021.

Clara Fannjiang, Stephen Bates, Anastasios N Angelopoulos, Jennifer Listgarten, and Michael I Jordan. Conformal prediction under feedback covariate shift for biomolecular design. Proceedings of the National Academy of Sciences , 119(43):e2204569119, 2022.

Glenn Shafer and Vladimir Vovk. A tutorial on conformal prediction. Journal of Machine Learning Research , 9(3), 2008.

Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Fujie Huang. A tutorial on energybased learning. Predicting structured data , 1(0), 2006.

Weitang Liu, Xiaoyun Wang, John Owens, and Yixuan Li. Energy-based out-of-distribution detection. Advances in Neural Information Processing Systems , 33:21464-21475, 2020.

José E. Chacón and Tarn Duong. Multivariate Kernel Smoothing and its Applications . Chapman and Hall, 2018.

Bharath Ramsundar, Peter Eastman, Patrick Walters, Vijay Pande, Karl Leswing, and Zhenqin Wu. Deep Learning for the Life Sciences . O'Reilly Media, 2019. https://www.amazon.com/ Deep-Learning-Life-Sciences-Microscopy/dp/1492039837 .

David Rogers and Mathew Hahn. Extended-connectivity fingerprints. Journal of chemical information and modeling , 50(5):742-754, 2010.

Kexin Huang, Tianfan Fu, Wenhao Gao, Yue Zhao, Yusuf Roohani, Jure Leskovec, Connor W Coley, Cao Xiao, Jimeng Sun, and Marinka Zitnik. Therapeutics data commons: Machine learning datasets and tasks for drug discovery and development. arXiv preprint arXiv:2102.09548 , 2021.

John J Irwin and Brian K Shoichet. Zinc- a free database of commercially available compounds for virtual screening. Journal of chemical information and modeling , 45(1):177-182, 2005.

Jan H Jensen. A graph-based genetic algorithm and generative model/monte carlo tree search for the exploration of chemical space. Chemical science , 10(12):3567-3572, 2019.

Wenhao Gao, Tianfan Fu, Jimeng Sun, and Connor Coley. Sample efficiency matters: a benchmark for practical molecular optimization. Advances in Neural Information Processing Systems , 35:21342-21357, 2022.

Y Gao, C Gesenberg, and W Zheng. Oral formulations for preclinical studies: principle, design, and development considerations. In Developing solid oral dosage forms , pages 455-495. Elsevier, 2017.

Kexin Huang, Tianfan Fu, Wenhao Gao, Yue Zhao, Yusuf Roohani, Jure Leskovec, Connor W Coley, Cao Xiao, Jimeng Sun, and Marinka Zitnik. Artificial intelligence foundation for therapeutic science. Nature Chemical Biology , 18(10):1033-1036, 2022.

Mufei Li, Jinjing Zhou, Jiajing Hu, Wenxuan Fan, Yangkang Zhang, Yaxin Gu, and George Karypis. Dgl-lifesci: An open-source toolkit for deep learning on graphs in life science. ACS omega , 6(41):27233-27238, 2021.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Guy W Bemis and Mark A Murcko. The properties of known drugs. 1. molecular frameworks. Journal of medicinal chemistry , 39(15):2887-2893, 1996.

A Proofs

A.1 Proofs for Theorem 4.1

Denote T i = 1 -p f Y i ( X i ) as a new random variable. Because of our i.i.d. assumption (and that f is not trained on the calibration set), T 1 , . . . , T N +1 are also i.i.d., which means that for m = 0 , 1 , . . . , N :

The left-hand-side probability is ≤ (instead of < ) the right-hand-side due to the case when p f y is not continuous. Note also that

which means

Finally, note that

we thus have

A.2 Proofs for Theorem 4.2

This is a direct result of Corollary 1 in [11], with our nonconformity score 1 -p f Y ( X ) plugged in.

B Data and implementation details

B.1 Training Details of the base classifier:

Deep Learning Frameworks: We use the Pytorch framework for the implementation of the models. The Graph Neural Network backbone is obtained from the DGL-LifeSci library [34].

Training hyperparameters: We train the model using the PyTorch Lightning Framework for training. We use the ADAM Optimizer [35]. The batch size is set to 64, and the learning rate is set to 0.001.

Architecture Details: The model architecture consists of a GNN layer (Attentive FP [1]), , a readout layer, 2 hidden FCNN layers, and an output layer. The hidden state size in GNN is set to 512 dimensions. The linear layers have 256, and 8 dimensions respectively.

Cheminformatics Processing : We use the RDKit library for handling molecular entities in Python. We use the DeepChem library for generating dataset splits.

Energy Regularization hyperparameters : The parameters m in and m out in Eq. (15) are set to -5 and -35 respectively, and the parameter λ in Eq. (16) is set to 0.01. All of these hyperparameters are obtained from the reference implementation in [24].

Splitting Ratio: The datasets are split in the ratio of 70:15:15, for training, calibration and testing the CP model.

Error bars: All the experiments reported are over 5 random runs. The mean and the standard deviation across the random runs is reported in the table.

B.2 Splitting strategies:

In this section, we discuss the two strategies employed for creating calibration and test sets to ensure their dissimilarity. The implementation of these strategies was based on the DeepChem library [26]. We provide detailed explanations of the algorithms used for splitting below.

Scaffold Splitting: The core structure of a molecule is represented by scaffolds [36]. Scaffold splitting aims to generate train and test sets that do not share any common scaffolds, thereby creating a challenging yet realistic scenario of distribution shift. This strategy is commonly used to evaluate out-of-distribution prediction algorithms [20, 18]. The scaffold splitting procedure is as follows:

-First, the scaffolds of all molecules in the datasets are identified.

-The scaffolds are sorted by their frequency in the dataset.

-The least frequent scaffolds are added to the test set until the desired number of test data points is reached.

-The remaining points are randomly divided into training and calibration sets.

Fingerprint Splitting: A molecular fingerprint[27] is a compact binary representation of a molecule's structural features, capturing important information about its chemical composition and spatial arrangement. Fingerprint splitting utilizes molecular fingerprints to create distinct train and test sets. The primary objective is to include data points in the test set that exhibit the least maximum pairwise Jaccard similarity of fingerprints. The fingerprint-splitting procedure is as follows:

-First, the molecular fingerprints are computed for all the molecules in the dataset using Extended Connectivity Fingerprints (ECFP) [27].

-Pairwise Jaccard similarity is calculated for each data point in the dataset, considering its similarity with all other points. The Jaccard similarity between two fingerprints is determined by dividing the size of their intersection by the size of their union.

-The test set construction begins by selecting the data point with the least maximum Jaccard similarity to any other point in the dataset. This point is added to the test set.

-The iterative process continues until the desired number of test data points is reached. The remaining data points are assigned to the training/calibration set.

B.3 Property prediction datasets

B.3.1 Labelled data

We used four commonly used classification benchmarking datasets from ADMET properties. The datasets are obtained from Therapeutics Data Commons (TDC) [28]. We include the statistics of all the datasets used for property prediction in Table 5.

AMES Mutagenicity : Mutagenicity is a vital toxicity measure that measures the ability of the drug to induce genetic mutations. The dataset consists of toxicity classes for over 7000 compounds.

Tox21 : A data challenge that consists of qualitative toxicity measurements on 12 different targets. For the scope of this work, we picked the largest assay in the collection with over 6000 compounds.

ClinTox : A collection of compounds that includes drugs that have failed in clinical trials for toxicity reasons and the ones with successful outcomes. This dataset contains about 1500 drugs.

HIV Activity : The dataset from screening results published by Drug Therapeutics Program (DTP) AIDS Antiviral Screen. It measures the ability to inhibit HIV replication for over 40,000 compounds.

B.3.2 Unlabelled data

To train the EBM, we obtained the unlabelled data from the ZINC-250k dataset [29]. This is a subset of the ZINC database, typically used for pre-training generative models, and covers a large chemical space. For each dataset and split type, we removed the molecules that are similar to the training (and calibration) set.

In the case of scaffold splitting, we remove the molecules containing scaffolds present in the training set.

In the case of Fingerprint splitting, we compute the Tanimoto similarity with respect to the molecules in the training set and include only those molecules with similarities less than the minimum pairwise similarity in the training set.

B.4 Hyperparameters for Kernel Density Estimation(KDE):

Determining the bandwidth( h ), for Kernel Density Estimation (KDE) is crucial for accurate density estimation. To find the optimal value of h , we employ K-fold cross-validation (CV) using the scikit-learn library, with k = 10 folds. The following procedure is applied for each dataset:

The dataset is divided into k splits or folds.

We fit the KDE model using a range of h values, specifically choosing 25 uniformly spaced intervals from 10 -1 . 3 to 10 1 .

The fitted KDE models are evaluated by computing the log probability on the holdout split for each h value.

The h value that yields the highest average log probability across the k folds is selected as the optimal bandwidth for fitting the KDE model.

B.5 Ablation models:

As discussed in Section 5.2.3, we perform the following ablations:

CoDrug (NoEnergy): The main objective of this variant of the method is to understand the significance of training with energy regularization described in 4.3. In this study, we follow the same protocol as CoDrug (Energy), but during the training, we do not use the regularization term L energy in 16.

Logistic (Energy): Logistic (Energy): In this variation, the training procedure for the model remains the same, but instead of employing Kernel Density Estimation (KDE) for estimating molecular density, we utilize a logistic classifier. This approach, initially utilized by Tibshirani et al. (2020) [11] in their experiments, involves training the logistic classifier using the same input features as those used in KDE.

For training the classifier, the samples in the calibration set are labeled as 0, while the samples in the test set are labeled as 1. The weight assigned to a point x in the calibration set is determined by the estimate ˆ p ( x ) of P ( C = 1 | X = x ) obtained from the classifier, and is calculated as ˆ p ( x ) / 1 -ˆ p ( x ) . To implement the logistic classifier, we utilized the scikit-learn library, employing the default hyperparameters for the classifier.

B.6 Details of de novo drug design experiments:

In this section, we present details of the experiments described in Section 5.3. We first describe the de novo drug design models used for the experiments in Appendix B.6.1. We sample 200 points from these generative models on properties described in Appendix B.6.2. Our objective is to construct valid prediction sets on the sampled molecules pertaining to the property discussed in Section B.6.2.

B.6.1 De novo drug design models

As our main objective is to show that CoDrug is effective in estimating uncertainties on molecules sampled from de Novo drug design models, we experiment on molecules sampled from two different de novo drug design models - Reinvent and Graph GA. These two models are top-ranked on the MolOpt benchmark [31]. We used the implementation provided by MolOpt to run our experiments.

Reinvent: Reinvent[14] is a reinforcement learning-based de novo drug design model. The model uses an RNN to generate SMILES strings.

GraphGA : GraphGA [30] is a genetic algorithm based de novo drug design model that generates molecular graphs.

B.6.2 Objective functions used for optimization

In our experiments, we used molecule sets obtained from optimizing the molecules on the following properties. All the oracles are obtained from TDC [28].

QED: QED stands for Quantitative Estimate of Drug-likeness. It is a computational metric used in drug discovery to assess the "drug-likeness" of a compound. QED provides a quantitative measure of how likely a molecule is to possess drug-like properties based on its chemical structure.

QED+JNK3activity: JNK3 activity refers to the activity of a molecule against a c-Jun N-terminal Kinases-3 (JNK3) protein. This is a common Oracle function used in benchmarking de novo drug design models. The oracle is built using a random forest classifier using ECFP6 fingerprints using the ExCAPE-DB dataset. In addition, we also add QED score to the Oracle output to restrict molecule search to a "drug-like" region.

QED + GSK3b activity : GSK3b, which stands for glycogen synthase kinase 3 beta, is an enzyme encoded by the GSK3b gene in humans. Dysregulation and abnormal expression of GSK3b have been linked to a heightened vulnerability to bipolar disorder. Similar to previous case, the oracle is built using a random forest classifier that utilizes ECFP6 fingerprints from the ExCAPE-DB dataset. We also add QED score to the Oracle output.

B.6.3 Details of the property prediction experiments

As discussed in the main paper, we choose logP as our target property, i.e. the property for which we wish to obtain uncertainty estimates. LogP, also known as the logarithm of the partition coefficient, is a property used to quantify the lipophilicity of a drug molecule. The lipophilicity of a drug molecule, as determined by LogP, plays a crucial role in its pharmacokinetic properties. It is accepted that a drug-like molecule would have logP in the range of [1.0, 4.0] [32] and hence, in our experiments, we assign a label of Y=1 for logP in [1.0, 4.0]; Y=0 otherwise.

logP can be computed cheaply using a computational oracle, which obtained it from TDC[28]. Note that this property in reality would be an experimentally determined property (such as ADMET properties), but the obvious challenge in validating our method on such properties is that it is not possible to obtain ground truth values for novel molecules obtained from de novo drug design models. Nevertheless, since CP is agnostic to the underlying prediction model, we deem that the performance would remain robust across different properties and hence would potentially be beneficial in real-world drug discovery campaigns.

For the curation of the training set, we randomly pick a set of 20 scaffolds from the ZINC250k dataset [29], and pick 500 points from each scaffold (making a total of 10000 points). This is our training and calibration data. We assign labels to this set based on the above-mentioned criteria and train the model using the same procedure as other property predictors as described in section 5. Note that the test set for this exercise is the molecules sampled from de novo drug design models.

C Results

In section Section 5, we have provided results for experiments at α = 0 . 1 . Here, we provide additional results for the experiments at α = 0 . 05 and α = 0 . 2 .

C.1 Property prediction results

C.2 De Novo Drug design experiment results

D List of commonly used notations and terms

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper concerns the problem of conformal drug property prediction, where set valued predictions that should cover with high probability the true class are made. The authors argue that due to sampling design covariate shift tends to occur where test data has different input distribution than the training data, breaking down standard validation set based calibration methods leading to incorrect coverage for predictions. A semi-supervised co-training method and KDE are proposed for obtaining reliable density estimates from validation and test data, in order to re-calibrate the conformal predictions. Theoretical analysis establishes coverage guarantee asymptotically, and experiments performed on drug discovery and drug design data sets show how basic approaches can easily break and that the proposed methods can provide more proper coverage predictions.\n\nSTRENGTHS:\nDrug property prediction applications (as well as many other similar molecular prediction tasks) routinely break basic i.i.d. assumptions, and thus there is clear need for ML methods that remain valid even under such covariate shift. Though closely related works exist, the exact formulation of the problem appears novel, and the proposed solution is both theoretically justified and empirically validated on two real-world data sets. The paper is fairly clearly written, recent literature well covered, and the Figure well illustrates the overall workflow.\n\nWEAKNESSES:\nSignificance of the work: It is not clear to me, how useful this kind of coverage guarantee would in practice be for researchers working on drug property prediction tasks, and the problems solved in the experiments are somewhat artificial as a split leading to covariate shift is simulated by the authors rather than based on some real world application where it would naturally appear. Contribution would be stronger if the authors could demonstrate the applicability of the approach to more broad range of applications.\n\nQUESTIONS:\nCan you open up a bit the basic motivation of the work: \"Furthermore, de novo drug design models or screening datasets sample\nmolecules from a large chemical space, making the exchangeability assumption invalid.\" - why and how exactly is this assumption violated in typical study designs?\n\nLIMITATIONS:\n-", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies conformal drug property prediction under covariate shift, i.e. the calibration data and the test data are not exchangable.\nThe proposed method, CoDrug, uses an energy-based model trained on both labeled and unlabeled data to model the conditional $p(y|x)$, and KDE to estimate the likelihood ratio between the calibration and test dataset, which is subsequently used to rectify for the distribution shift.\n\nExperiments demonstate CoDrug effectively reduces coverage gap compared to unweighted models. A demo is also given where CoDrug is applied to estimate the logP values of a set of de novo designed molecules.\n\nSTRENGTHS:\n**Originality**\n\nThis work addresses uncertainty quantification in molecule property prediction under covariate shift. The authors proposed a novel method to estimate the likelihood ratio, and applied it to the molecular property prediction task.\n\n**Quality**\n\nThe overall paper quality is good. Contributions are clearly stated. The math is solid. Most model choices are reasonable or justified. The experimental setting is fair. And a nice, realistic demo is given as a dessert.\n\n**Clarity**\n\nThe writing is easy-to-follow in general.\n\n**Significance**\n\nThe task is important and realistic in drug discovery.\n\nWEAKNESSES:\n**Clarity**\n\nPerhaps Section 4.2 (implementation) could be put after 4.3 (the big picture).\nCurrently Section 5.3 is not clear enough. At first I thought you are using this model to generate new molecules. Also, I believe this section can be shortened to free up some spaces in the current somewhat overcrowded paper.\nIf you have the time, figure 1 could use a little bit of improvement.\n\n\n\\* See questions for other points to be clarified.\n\nQUESTIONS:\n1. How important is the unlabeled ZINC dataset to the experimental results?\n2. Could you explain, theoretically and/or empirically, why I couldn't use a generative model trained on ZINC to estimate the likelihood ratio?\n3. Why is ReLU not applied in the provided code when getting classification features?\n```python\nclass ClassificationFeatureModel(nn.Module):\n    def __init__(self, model_config):\n        super().__init__()\n        # several lines omitted\n        self.op_1 = nn.Linear(feature_len, int(feature_len / 2))\n        self.op_2 = nn.Linear(int(feature_len/2), 8)\n        self.op = nn.Linear(8, self.num_outputs)\n\n    def forward(self, mol_graph_batch):\n        node_features = self.gconv(\n            mol_graph_batch, mol_graph_batch.ndata[\"h\"], mol_graph_batch.edata[\"e\"]\n        )\n        gfeat = self.readout(mol_graph_batch, node_features)\n        \n        op1 = F.relu(self.op_1(gfeat))\n        op2 = F.relu(self.op_2(op1))\n        \n        outputs = self.op(op2)\n        return outputs\n\n    def get_features(self, mol_graph_batch):\n        node_features = self.gconv(\n            mol_graph_batch, mol_graph_batch.ndata[\"h\"], mol_graph_batch.edata[\"e\"]\n        )\n        gfeat = self.readout(mol_graph_batch, node_features)\n\n        op1 = self.op_1(gfeat)\n        op2 = self.op_2(op1)\n        return op2\n```\n\nLIMITATIONS:\nI believe the authors should have adressed more to the limitations of this paper. I suspect this is due to the page limit.\nI don't see any significant unreported negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper is concerned with the problem of conformal drug property prediction by formulating and solving the problem as a conformal prediction problem under covariate shift. To this end, the paper proposed training an energy-based model using both labeled and unlabeled data.  To derive conformal prediction under covariate shift, the paper makes use of kernel density estimation to estimate the covariate distribution in training time and in test time. The paper then proceeds to use the estimated density to counteract covariates shift. Theoretical guarantees are given about the proposed methods. Experiments are conducted on four drug toxicity datasets and de novo drug development experiments. The proposed method is compared to a baseline in the experiments.\n\nSTRENGTHS:\n* The paper is concerned with an important practical problem in drug design. It provides a good application of conformal drug prediction under covariate shift.\n* Experiments are conducted on a wide variety of datasets that are relevant to the problem of drug design.\n* The presentation of the paper is clear.\n\nWEAKNESSES:\n* One of my major concerns is the originality of the paper. While the paper makes reasonable choices to develop the proposed method to deal with conformal prediction under covariate shifts for conformal drug properties predictions, many of these choices appeared to be straightforward applications of existing works. It is not clear to me whether the original contribution of the paper is sufficient for acceptance.\n\n* In terms of empirical evaluation, while I think the methods are compared on a fair amount of various datasets related to drug design, I am not quite sure if the proposed method has been compared to enough alternative methods to understand the effectiveness of the proposed method. As such, the paper can benefit from a more exhaustive comparison of related work in the literature.\n\n* While there are theoretical justifications mentioned in the paper. They are mostly straightforward derivations from existing work. Although this is fine, it does not seem to contribute to the strength of the paper.\n\nQUESTIONS:\n* The authors may further explain why the proposed method has sufficient novelty.\n* The authors can also consider introducing more alternative methods for comparison empirically.\n\nLIMITATIONS:\nThe authors mention the limitation of their work in the conclusion section.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: SUMMARY:
The authors apply conformal prediction under covariate shift to molecular property prediction. They propose using KDE to estimate the densities for the calibration and test sets, and demonstrate that doing so is better than plain vanilla CP.

STRENGTHS:
### Originality

This paper puts together existing ideas that provide uncertainty estimates for molecular property prediction and related tasks.

### Quality

Quality is high, the method is evaluated on many different prediction and generative tasks.

### Clarity

The paper is well-written and easy to understand.

### Significance

Proper uncertainty estimation is important for many molecular prediction tasks, and this paper provides a CP method with guarantees that deals with the central issue of such tasks, namely covariate shift.

WEAKNESSES:
- Should show under what conditions using KDE fails.

QUESTIONS:
No questions

LIMITATIONS:
There is no discussion of limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07