CoDrug: Conformal Drug Property Prediction with Density Estimation under Covariate Shift
Authors: Siddhartha Laghuvarapu
Zhen Lin
Jimeng Sun
Venue: NeurIPS 2023 poster (2023)
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
Siddhartha Laghuvarapu
Zhen Lin
Department of Computer Science University of Illinois Urbana-Chane Champpaign
Department of Computer Science University of Illinois Urbana-Champann Urbana, IL 61801 zhenlin4G @ illinois.edu
Urbana, IL 61801
sl160@ illinois.edu
Jimeng Sun
Department of Computer Science Carle Illinois College of Medicine University of Illinois Krbana-Champan Urbana, IL 61801 jimeng @ illimois.edu
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 subse- quent experimental validation. Conformal Prediction (CP) is a promising tool for creating such prediction sets for molecular properties with 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 method called CoDrug that employs energy-based based model leveraging both training data and unlabelled data, and Kernel Density Estimation (KDE) to assess the densities of 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-molecne drug discovery tasks, we demonstrate the ability of CoDrug 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.
Introduction
[Drug discovery 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 deep learning have aimed to reduce the cost of drug discovery by proposing Al methods for developing accurate property prediction models and De Novo drug design models:
Property prediction models aim 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
The code associated with the paper is available at a/thg///ghhhaccmmmiiiaaaallooo
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

De novo drug design models, on the other hand, aim to discover new drug molecules that satisfy a set of pharmaceutical properties 3 5
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-makinn and reduces the risk of failures, making the drug discovery process more efficient. Conformal prediction (CP), pioncered by 7. offers solution to uncertainty quantification for complex models like neural networks, by constructing provably valid prediction set in supervised learning models. Its application to drug property prediction has also been explored for various drug discovery tasks 88 [ [0].
crucial assumption in the CP framework that the test samples are exchangeable with the holdout set used to calibrate the algorithm. In drug discovery, there is often 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 situation where the training data originates from a distribution, P(X), while the test data comes from a different distribution, plest (X). However, in both cases, the molecular properties, determined by the conditional distribution P(Y remain the same as they are governed by nature and unaffected by shifts in the input distribution, assuming testing parameters stable. This is referred to as covariate shift. Although recent research in Conformal Prediction (CP) [1I] suggests a method for correcting covariate shift, accurately estimating the precise level of covariate shift remains a practical challenge.
This paper proposes novel and practical method for Conformal Drug property prediction, dubbed as CoDrug, to improve coverage in the conformal prediction framework under covariate shift. We address the problem of won-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. target property 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 consistent, which means that asymptoticallyy, 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 distri- bution 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
I 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 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 [ 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 [1]] and proposed the Feedback Covariate Shift
Predicttin intervals can be viewed prediction sets, with each interval being subset of H.

---

(FCS) method for the task of protein design. In practice, one cannot know the exact densities to measure the covariate shift. Like [2T]. we also leverage [TT] but key difference is that the training density is well-defined in [1] but unknown in ours, requiring to estimate it. Additionally, our focus
diverges from 21 as we concentrate on molecule property prediction rather than protein design.
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 classification task, with each data point Z (X, Y) E R x [K] [K]= {0. 1, 1,2.... K - 1)). For instance, in Fig. F 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 point estimate of the solubility of the molecule, but also some "confidence level". This could be encoded in the form of prediction set denoted as (()) S [K].
The main goal we seek such prediction sets is valid coverage: Given a target (e.g. 90%). we would like construct set-valued prediction (Fig. [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 - a E (0, 1), and a new test molecule (XN+1, 1N+1), we would like C' to be 1 - a valid:
P(YN+1 E C(XN+1)) 2 -aa a.
(1)
Conformal Prediction(CP) framework enables to achieve such validity in Eq. [T. 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 or
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,, Y:) and the test point as (XN+1, YN+1), we have:
Vi E [M1, (X,, Yi) kidd poal = Prea x Py
(2)
KN41,YN+H plest plest x peat
important to note that the test distribution plest maintains the same conditional distribution peat as the calibration distribution, phenomenon known as covariate shift. This shift is prevalent in de novo drug design models, which require navigating vast chemical space to pinpoint optimal molecules for a specific goal. However, in many drug discovery tasks, the datasets typically contain only few thousand data points, representing 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.
CoDrug Method
4.1 Overview
In the subsequent subsections, we describe the three primary components of CoDrug. In Section 4.2 first provide brief overview of inductive conformal prediction, presenting a method for con- structing valid prediction sets in scenarios both without and with distribution shifts, presuming oracle access the unknown distributions prest and Pea Next, in Section 4.3 we present the details of the training aspects of the base energy-based classifier, emphasizing additional regularization using unla- beled data enhance its capability 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 plest and peal and rectify covariate shift using
'The assumption indicates that since the model f was not trained on the calibration set, whatever over wett-ftting happens on FIX"" but not fyl. Roughly speaking, we assume the classifier's 's performance on the calibration set is similar to that unseen test
3

---

[A] Conformal Prediction Framework
shift
Training Enorgy model
Anter
seranesiiall X
W5 (4)
0.2
Sametriipal
Traiiin Set
0.25
6:32
Melre videe]
0.89
66:
12-0..)
687
KEAI Pel
0:21
oul heghtyb))
Sabuuiity Conditaatinn
(4) Density Estimation
(c) Calibration under covoriate shift
Sorted stos
Pos (K
W.002
0.35
Done sukeaticn
Par
0.1
Tot
18-0.3)
07/033
0.71
pevard
Irivvill
Driration
BBC
W,- Pmn 0K) (Q))))
Figure 1: CoDrug overview: (a) A depiction of the conformal prediction (CP) framework. A valid prediction set includes the true label of the input molecule. (b) Standard procedure for computing quantiles from the calibration set when the test set is exchangeable. The calibration set's "non- conformity" scores are sorted, and the (I-a) quantile serves as the threshold for the conformal prediction set. (c) describe the CoDrug pipeline. (c) Training Energy-based model using labeled and unlabeled data. (d) Density estimation: The model from (c) is used to estimate the density of the calibration and test sets. (e) Calibration under covariate shift: First, likelihood ratios w, computed from the densities (d). Then, Quantile is computed a weighted fashion. Note how the quantile at 0.3 is shifted from 0.64 (in (b)) 0.71 to account for the distribution shift.
Section As KDE is consistent, we regain the coverage guarantee asymptotically. Together, these elements constitute the pipeline depicted in Fig.
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, pioncered by [7]. is powerful framework to construct prediction sets with the guarantee in Eq. [T. We assume that a base classifier f is trained on a training set Dtrain and we have hold-out calibration set Deat To simplify notation, will denote the calibration set as { and the test point of interest as Z We will also abuse the notation to use D to denote both the empirical calibration/test set well as the underlying distribution. Note that 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) of set of values (v; J is defined as:
N
F(O == I/N Σ
duss where (x) = 1 2 w}
(4)
i=1
The quantile fuuccton with respect to a CDF F is:
Quantile(B; F) := inf(z: F(x) 2 3}
(5)
Given target coverage level (Mondrian ICP) is given by:
(0. 1), the Mondrian inductive conformal prediction set
Č(XN+1) = (v 1 1/ (( St) where t: == Quantile(1 a F ((1-n(((a))00)))
(6)
where, (xx) corresponds to the softmax output of class from model f. Here, {ui) are defined by v(21, v1) 1 pl, (zz), which are called "nonconformity scores" [22] and measure how "anomalous"

---

point 2 = (z,3) 3) 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 Dirain Now, we choose threshold / that is larger than 1 - a (e.g., 90%) of the molecules. Because of our i.i.d. assumption, if we sample another molecule ZN+1 from the same distribution, we expect its score to be lower than this threshold with prohability of 1 a. For eg. in Fig. [b), notice how the threshold t, is computed as the value of the Quantile function at a = 0.7 (t = 0.64 in this case). We formally state the coverage guarantee without covariate shift, in the following theorem:
Theorem 4.1. Assume i.i.d.
X.) The O in Eq.
6 satisfies:
(8) Theorem 4.1 is 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.
P(YN+1 E C((N+1)) 2 21-a. a.
Remarks:
4.2.2 Conformal Prediction with Covariate Shift
While Theorem 4.1 provides nice first step, the i.i.d. assumption poses significant limitation in drug discovery. As mentioned in Eq. (3, the distributions of X on Diest and Deal can differ. In Eq. we used the empirical CDF F (Eq. ( to choose the threshold This is because of our i.i.d. assumption: 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 Px.
Formally, recall that Px" and Piaat represent the density of the molecule for the calibration and test sets. We will assign weight to each molecule x that proportional to the density/likllihood ratio dPX /apgal in the empirical CDF, leading to:
F 28-1 == (((=N+))6<< + Eie m(a.) 11 (x,/
(9)
w w(xx) dPk (/aps" Va!
(10)
W=E: w(x;)) is just a normalizing factor. The subscript IN+1 is used to highlight that our updated CDF now depends the test molecule NN+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. Te). Notice how the values of w(xx;) depend on the densities posl and Plest, In the figure, the value of weighted Quantile at a = 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 [1] that recovers the coverage guarantee under covariate shift for Mondrian ICP:
Theorem 4.2. [T]] Assume that Px is absolutely continuous with respect Px. For any a E (0, 1), let F be defined as in Eq. (9) and
Ö(x) by: 10:1-p((22 px) S Quantile(1-an - a; F1')
(11)
Then,
P(YN+1 E ((XN=)))
(12)
impractical. In section 4.4 we will provide a viable way to estimate plest and recover the guarantec 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.
However, in practice, both ptest and Pea are unknown, rendering Theorem 4.2
4.3 CoDrug Training Methodology
CoDrug handles distribution shift by proposing an energy-based model formulation idea behind an energy- based model is to construct a function E that maps an input x to scalar value, known as energy. A collection of energy values can be transformed into a probability density function p(x) through the Gibbs distribution
23
The core
p(y|=) = Els.u)/T/ /c- E(e)/T
(13)
Consider a discriminative neural network f used in K class classification setting. f(x) maps an input æ into K real-valued scalars, which are used to derive conditional class-wise probability:
p(y|z) = abucn/ /> df (0//T
(14)
5

---

where, fz) refers to the yth logit of the classifier fzz). In this setting, the energy function E(x)
can be expressed in terms of the denominator of the softmax probabilities in Eq. (14)
K
E(;; f) = -T log Σ C
(15)
3
Directly using embeddings from model trained labeled data may not yield reliable density estimates, as the model lacks knowledge of data outside the training distribution. To overcome this, co-train the model with unlabeled molecule data. This aids the model 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:
min log(i/
Lenergy
(16)
where (x) in-distribution training data for which labels are available. The training objective is combined with an additional term Lenergy given by:
(x)) refers the softmax outputs of the classification model for class y, and Din. is the
Lenergy E (e (max(0, (E(za.) m))))²) P (max(0, (Mout ((
(17)
where Dout refers to the subset of the unlabelled that is out-of-distribution (OOD). The objective of this term to enforce a margin of separation between the training samples and the OOD data using the hyper-pann paremetersttssss Min, Mout. 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 energy lower than 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, need estimate Plest and pool 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 set of data "i,....,,,, ~ D, KDE is given by:
i.dd
(r D) (nh) Σ K(o-ws/h),
(18)
where K is fixed non-negative kernel function, and h is 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-integrabe and twice differentiable and that its second-order partial derivatives are bounded, continuous, and square intter-iiteg.able If K is spherically symmetric on Rdd with a finit second moment, and we choose the bandwidth h such that
lim hem + oo and lim h + 0 m-000
(19)
1-0000
then as m - Dc,
- ( P o.,
(20)
where Ps means convergence in probability.
Note that commonly used kemnlls, such Gaussian kernel, satisfy the requirements.
Since here refers to molecular entities (e.g. SMILES strings), we cannot use 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
w (x)) == Prices (w:Drees) //pp (x::D.nt)
(21)

---

to replace the unknown w in Eq. (T), giving us the final prediction set:
Čoooro (x) {y 1- pz S Quantile ax; F2))
(22)
F #N+1 == (x EEEIN 1 -oh,
(23)
where W- -DiT (244 is normalizing factor. Here, Phease (:; Drest) is constructed using samples from the test data with an optimal bandwidth freat chosen on the test data via cross-validation, and Phoat Deaaa)) is constructed similarly but on the calibration data. It is clear that, as the number of samples from Deat and Dteet increases, i converges coverage guarantee asymptotically. In practice, recovering asymptotic coverage on 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. [T).
N+1
in Eq. 10). and VirCaDrug recovers the
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 Derain and calibration set Peat {
Train a neural net classifier on Dtrain by minimizing Eq. (6).
Compute the KDE Pheai (-i (eaeatt for all points in Doat using Eq.
18
Test Time, for a test set Drest:
Compute KDE Pircese (:; Drest) for all points in Deat using Eq. 18
For any IN+1 E Dtest, compute (xx) and (x;;) for x; Deal- Construct the prediction set (CoDrus (x)D using Eq. (2)).
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. 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. Sec- tion 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. both strategies, we try to create calibration and test splits that dissimilar each other, which is 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 Tanimoto similarity of molecular fingerprints [27]. Molecules with the highest dissimilarity in terms of Tanimoto similarity are included in the test set.
Datasets: We four binary classification datasets for toxicity prediction (AMES, Tox21, Clin'Tox) 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]. subset of the ZINC that covers
a large chemical space. For each dataset and split type, we removed the molecules that are similar the training (and calibration) set from the unlabelled dataset.
Classification Model: The architecture of our classifier f is AttentiveFP [.] a graph neural network-based model. We chose AttentiveFP as has state-of-the results in several drug property prediction tasks. It is trained using the objective function described in Eq. (T6).
De Novo Drug Design Experiments: In Section 5.3] we perform experiments to construct confor- mal prediction sets molecules sampled from de novo drug design models. As generative models, we use REINVENT(4] 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 target property, assigning values in the range of
7

---

[1.0,4.0] class of Y=1, and Y=0 otherwise (representing the drug-like range [32] We obtain the computational oracles from TDC [3]. and generative models from MolOpt package
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 shows the results of conformal prediction under various distribution shift conditions. "Random" refers to the idcal/unrcalistic scenario where the test and calibration samples are split randomly (aka. no distribution shift). "Scaffold" and "Fingerprint" denote scenarios in which there 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 without any correction.
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 coverage of 0.34 against target coverage of 0.8 the AMES dataset, while achieving a very different coverage of 0.77 with scaffold split the same dataset.
From Table
Random
Fingorprint
Scaffold
Random
Fingerprint
Scaffold
Dataset
a=0.1
a=0.2
AMES(Y=0)
0.94
1.00
0.82
0.85
1.00
0.66
AMES(Y=1)
0.87
0.63
0.85
0.78
0.34
0.77
TinTox(Y=0)
0.88
0.78
0.84
0.77
0.58
0.75
ClinTox(Y=1)
0.82
0,80
0.97
0.78
0.73
0.81
HIV(Y=0)
0.90
0.93
0.91
0,80
0.89
0.81
HIV(Y=1)
0.89
0.84
0.87
0,80
0.72
0.73
Tox21(Y=0)
0.90
0.77
0.80
0.65
0.75
Tox21(Y=1) 0.86
0.97
0.93
0.72
0.97
0.82
Table 1: Unweighted CP's (baseline) coverage under various distribution shifts (or absence thereof)
should ideally align elosely with the target a. However, in datasets with fingerprint and seaf- fold splits-reflectie 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 p-value of 0.05 are highlighted in bold, indicating desirable performance.
5.2.2 Weighted conformal prediction improves coverage
In Table 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 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=I) category, where the coverage improved from 0.63 to 0.88 (target coverage 0.9). Note that the coverage is "improved" if it closer to 1 - a improvement does not always mean higher coverage, because an unusually high coverage often indicates unpredictable behavior of the underlying model.
While our mergy-weigghtee approach generally improves coverage, there are rare instances where it underperforms compared to baseline. A prominent example is the ClinTox dataset with 1, which sees limited improvement even 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 the smallest and most imbalanced, with only 19 points in the calibration set for class Y 1.
8

---

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 wwodiimensional, 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 high-quality density estimate.
Dataset
Fingerpiit Splitring
Scaffold Spiiting
CoDrugg (Energy)
Cabrug (Feature)
Unweighted (haseline)
CaDrug (Energy)
CoDrrg (Feature)
Unweighted (haseline)
AMES/Y=0)
0.93(0.03)
0,87/0.03)
1,00(0.00)
0.85(0.02)
0.89(0.02)
0.52(0.01)
AMES/Y-1)
0.88(0.0.3)
0.90:0.03)
0.63(0.05)
0.83/0.01)
0.79(0.91)
(.85(0.03)
Clin Y-0)
0.86(0.11)
0.7660.02)
0.78(0.012)
0.90(0.03)
0,83(0001)
0.84(0.00)
Clin (o1
0.73(0.00)
0.6910.08)
0.80)
0.85(0.03)
0.83(0.00)
0.97(0.04)
HIV(Y=D)
0.89)0.06
0.87(0.07)
0.93(0.04)
0.8200.08)
0.82(0.94)
0.91(0.01)
HIVY=-I)
0.92(0.05)
0.95(0.03)
0.8410.07)
0.9610.01)
0,90(0.05)
0.87(0.03)
Tox21(Y-0) Tox21(Y=1)
0.90(0.02)
0.50:0.02)
0.7700.03)
0.91(0.03)
0.83(0.05)
0.8900.05)
0.97(0.00)
0.96/0.01)
0.97(0.00)
0.86(0.05)
0.91(0.05)
0.93(0.33)
Table 2: Coverage of CoDrug and baseline unweighted CP, under different datasets and distribution shifts at 0.1. The realized coverage rate closest to the target coverage 1 1 (best) is marked in bold. The second best coverage (in case better than unweighted) is marked in bold and gray. Results averaged over 5 random runs. Results for different a values available in appendix.
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 ( EEq. (16)), and covariate shift correction. In addition to CoDrug (Energy) and Unweighted (baseline) reported in the previous section, we also compare with:
CoDrug (NoFnergy): We use the same protocol as CoDrug (Energy) but the models are trained without the energy regularization term Lenergy 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 wzx;) in Eq. (10), are given by Ptea)/1-0x where P(z,) is obtained by fitting a classifier to features in calibration and test sets (suggested by [l].
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 a (i.e Mean of OOseevedd Coverage (1 - a)l 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 a. We paid close attention to the "Tail 25%", where we presented the metrics for the worst 25% performing experiments and CoDrug (Fnergy) 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).
Method
Mean absolute coverage deviation
Mean ubsolute coverage deviation (tuil 25%)
Q 0.3
0.2
0.1
0.3
C 0.2
0.1
Unweighted (baseline)
0.157 (0.14)
0.12 (0.12)
0.081 (0.07)
0.347 (0.14)
0.276 (3.)3)
0.176 (0.08)
(Egist))(Enerry)
0.123 (0.13)
0.106 (0.13)
0.083 (0.14)
0.315 (0.11)
0.263 (0.17)
0.222 (0.22)
CoDrug (NoEnergy)
0.112 (0.11)
0.083 (0.09)
0.047 (0.05)
0.288 (0.05)
0.215 (0.08)
0.112 (0.03)
CoDrug (Rnergy)
0.104 (0.09)
0.079 (0.07)
0.052 (0.05)
0.233 (0.05)
0.179 (0.07)
0.11 (0.04)
Table 3: Ablations: Results comparing various versions of the proposed framework. At each a, the
mean of deviations from target coverage across all the experiments and random seeds is computed (Smaller is better). CoDrug (Energy) has the least deviation from coverage and a substantial difference when only the worst performing 25% of the experiments are considered.
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 computational oracle. After molecule sampling, validating their experimental properties, such ADMET (Absorption, Distribution, Metabolism, Excretion, Toxicity), 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 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 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 at target alpha value of 0.1. Our proposed method consistently enhances coverage in all cases and exhibits substantial improvement over the unweighted version. For example, in the "gsk3b+qed" objective, the unweighted version has a coverage of 0.44 against target of 0.9, whereas our proposed method improves coverage substantially. The mean absolute coverage deviation from the target at a = 0.1 for CoDrug (Energy) is 0.05, a relative improvement of over 60% on the Unweighted version 10.1
REINVENT
GraphGA
Objective
Y
CoDrug (Energy)
Unweighted
CoDrug (Energy)
Unweighted
JNK3+QED
0
0.95 (0.01)
0.62 (0.12)
0.86 (0.0)
0.84 (0.01)
JNK3+QED
I
0.91 (0.01)
0.99 (0.0)
0.93 (0.01)
0.89 (0.02)
GSK3b+QED
0
0.81 (0.04)
0.44 (0.16)
0.87 (0.0)
0.75 (0.01)
GSK3++EED
0.79 (0.08)
1.0 (0.0)
0.98 (0.0)
1.0 (0.0)
QED
0.96 (0.0)
0.83 (0.04)
0.96 (0.0)
0.69 (0.1)
QED
0.92 (0.01)
0.84 (0.05)
0.98 (0.0)
0.99 (0.0)
Table 4: Observed coverages on molecules sampled by generative models at a = 0.1. The realized coverage rate closest to the target coverage(1 - a) marked in bold. For each experiment, a set of 200 points optimized w.r.t. the "Objective" using the generative models GraphGA and REINVENT are sampled. The target property for prediction logP (1.0 logP < 4.0 is considered Y=1; Y=0 otherwise 32)) ). Using the proposed method improves coverage in almost all scenarios.
6 Limitations
While we demonstrated that KDE can provide asymptotic coverage guarantees, this may not necessar- ily translate improved performance in scenarios with limited sample sizes. In our experiments, do acknowledge that there are few instances where the improvement in coverage is limited, where the availability of data is limited. As such, direction for further research is to explore ways obtain likelihood ratios that are more data-efficient, particularly in scenarios with smaller calibration sets.
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 modifi- cations. Furthermore, our current work primarily focuses on small molecules, yet covariate shift is 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 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 masee model for feature extraction and base classification, performing density estimation using KDE, and use the KDE to correct for covariate shift 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 interesting for future work, including adapting the framework for regression tasks.
10

---

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
[1] 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(1)):8749933 8760, 2019.
[2] Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. International conference on machine learning, pages 1272-127. PMLR, 2017.
[3] Tianfan Fu, Wenhao Gao, Connor W Coley, and Jimeng Sun. Reinforced genetic algorithm for structure-based drug design. arXiv preprint arXiv:2211.16500 2022.
[4] 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.
[5] Rafael Gómez-Bombarell Jennifer Wei, David Duvenaud, José Miguel Hernández-Lobow Benjamín Sinchez--ennelling Dennis Sheberla, Jorge gguilera-Iparraguiin Timothy D Hirzel, Ryan Adams, and Alán Aspuru-Guzik. Automatic chemical design using a data-driven continuous representation of molecules. ACS central science, 4(2):268-217 2018.
[6] 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
[7] Vladimir Vovk, Alexander Gammerman, and Glenn Shafer. Algorithmic learning in a random world. Springer US 2005.
[8] Jin Zhang, Ulf Norinder, and Fredrik Svensson. Deep learning- based conformal prediction of toxicity. Journal of chemical information and modeling, (1(6):264-26655 2021.
[9] Isidro
and Andreas Bender. Concepts and applications conformal prediction in computational drug discovery. arXiv preprint arXiv:1908.03550 2019.
[10] Isidro Cortés-Ciriano and Andreas Bender. Deep confidence: computationally efficient framework for calculating reliable prediction errors for deep neural networks. Journal of chemical information and modeling, (2(3):126-11233 2018.
[11] Ryan J. Tibshirani, Rina Foygel Barber, Emmanuel J. Candes, and Aaditya Ramdas. Conformal prediction under covariate shift, 2020.
[12] Mario Krenn, Florian Häse, AkshatKumar Nigam, Pascal Friederich, and Alan Aspuru-Guzi Self-referencing embedded strings (selfies): A 100% robust molecular string representation. Machine Learning: Science and Technology, 144):045024, 2020.
[13] 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 Neural Information Processing Systems, 33: 13:11259-11577 2020.
[14] Mareus Olivecrona, Thomas Blaschke, Ola Engkvist, and Hongming Chen. Molecular de-novo design through deep reinforcement learning. Journal of cheminformatics, 9()::-114, 2017.
[15] 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, (0(8):3770-3780 2020.
[16] Fredrik Svensson, Natalia Aniceto, Ulf Norinder, Isidro Cortes-Ciriano, Ola Spjuth, Lars Carls- son, and Andreas Bender. Conformal regression for quantitative cannture-activvin relationship modeling moodelnng-quantiffing prediction uncertainty. Journal of Chemical Information and Modeling, 58(5): 1:2-111321110 140, 2018

---

[17] Jiangming Sun, Lars Carlsson, Ernst Ahlberg, Ulf Norinder, Ola Engkvist, and Hongming Chen. Applying mondrian cross-conformal prediction estimate prediction confidence on large imbalanced bioactivity data sets. Journal of chemical information and modeling, 57(7):1591- 1598, 2017.
[18) 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.09633 2022.
[19] Shurui Gui, Xiner Li, Limei Wang. and Shuiwang Ji. Good: A graph out-of-distribution benchmark, arXiv preprint arXiv:2206.0888 08452, 2022.
[20] Kehang Han, Balaji Lakshminarayanan, and Jeremiah Liu. Reliable graph neural networks for drug discovery under distributional shift. arXiv preprint arXiv:211 2021.
[21] Clara Fannjiang, Stephen Bates, Anastasios N Angelopoulos. Jennifer Listgarten, and Michael I Jordan. Conformal prediction under feedback covariate shift for biomolecular design. Proceed- ings of the National Academy of Sciences, 119(4))::2204569119 2022.
[22] Glenn Shafer and Vladimir Vovk. A tutorial conformal prediction. Journal of Machine Learning Research, 9(3), 2008.
[23] Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Fujie Huang. A tutorial on energy- based learning. Predicting structured data, 1(0), 2006.
[24] Weitang Liu, Xiaoyun Wang, John Owens, and Yixuan Li. Energy- out-of-distributin detection. Advances Neural Information Processing Systems, 33:21464 21475, 2020.
[25] José E. Chacón and Tarn Duong. Multivariate Kernel Smoothing and its Applications. Chapman and Hall, 2018
[26] Bharath Ramsundar, Peter Eastman, Patrick Walters, Vijay Pande, Karl Leswing, and Zhenqin Wu. Deep Learning for the Life Sciences. O'Reilly Media, 2019. httpp://wwwr amazon. com/ Deep-Learning Life-
[27] David Rogers and Mathew Hahn. Extended-connectivne fingerprints. Journal of chemical information and modeling, 50(5):442-754, 2010.
(28) 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 learn- ing datasets and tasks for drug discovery and development. arXiv preprint 2021.
[29] John J Irwin and Brian K Shoichet. Zinc- free database of commercially available compounds for virtual screening. Journal of chemical information and modeling, 45(1):177-182, 2005.
[30] Jan H Jensen. graph- based genetic algorithm and generative model/monte carlo tree search for the exploration of chemical space. Chemical science, 10(12):3567-3572. 2019.
[31] 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-213511 2022.
[32] 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.
[33] 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.
[34] Mufei Li, Jinjing Zhou, Jiajing Hu, Wenxuan Fan, Yangkang Zhang, Yaxin Gu, and George Karypis. Dgl- Dgl-lifesci: An open-sourse toolkit for deep learning on graphs in life science. ACS omega, 6(41 (141):27233-222)
2021.
[35] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412.690,, 2014.
[36] Guy W Bemis and Mark A Murcko. The propertics of known drugs. 1. molecular frameworks. Journal of medicinal chemistry, 39(15):2887-299) 1996.
12

---

A Proofs
A.1 Proofs for Theorem 4.1
Denote T; = -1-py (X,) as a new random variable. Because of our i.dd. assumption (and that f not trained the calibration set), ......, TN+1 also i.i.d., which means that for m 0, ,......00 N:
Pilie E [M] TN+1 > TH = m) <
(24)
N +1
m +1
P(lf E [N] TN+1 > T.H 2 (N m)) -
(25)
N
+1
The left- probability I (instead of <) the right- hand-side due to the case when p.y is not continuous. Note also that
1 1 PYPy >t = [N] TN+1 > T.l 2 [(1 +
(26)
which means
P(1 PYNni (XN+1) > 21s [a(N + 1)]
N+1 aa.
(27)
Finally, note that
YN-I & Ō(XN+1) = 1 plans (XN+1) > b
(28)
we thus have
P(YN+1 E C(XN+1)) 21-2
(29)
A.2
Proofs for Theorem 4.2
This is direct result of Corollary in [1] with our nonconformity score - (X) plugged in.
13

---

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 [. Training hyperparameters: We train the model using the PyTorch Lightning Framework for training. We use the ADAM Optimizer [35] The batch size is to 0.001.
64, and the learning rate is set
The model architecture consists of a GNN layer (Attentive FP M.. readout layer, 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.
Architecture Details:
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 Man and Mout in Eq. (15) are set to -5 and -35 respectively, and the parameter
Eq. (16) is set 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, 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 molecule is represented by scaffolds [36]. Scaffold
splitting aims to generate train and test sets that do not share any common scaffolds, thereby creating challenging yet realistic scenario of distribution shift. This strategy is commonly used to
evaluate -dittt--isttibuttin prediction algorithms 20 18] follows:
The scaffold splitting procedure is as
- First, the scaffolds of all molecules the datasets are identified.
The scaffolds 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 reached.
The remaining points are randomly divided into training and calibration sets.
Fingerprint Splitting: A molecular fingerprint (7] is compact binary representation of molecule's 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 angerprint-splitting procedure 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 T
AMES Mutagenicity: Mutagenicity is vital toxicity measure that measures the ability of drug to induce genetic mutations. The dataset consists of toxicity classes for over 7000 compounds.
14

---

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. Clin'Tox: 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.
Dataset
#Positive
#Negative
#Total
Task
Tox21
309
6,956
7,265
prediction
ClinTox
112
1,366
1.478
prediction
AMES
3,974
3,304
7,278
prediction
HIV
1,443
39,684
41,127
prediction
ZINC
0
0
250,000
pre-training
Table 5: Dataset statistics
B.3.2 Unlabelled data
To train the EBM, we obtained the unlabelled data from the ZINC-250K dataset 29 of the ZINC database, typically used for pre-training generative models, and covers a large chemical space. For each dataset and split type, removed the molecules that similar the training (and calibration) set.
This is a subset
the case of scaffold splitting, we remove the molecules containing scaffolds present in the
training set.
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- cross-validee validation (CV) using the scikit-learn library, with k = 10 folds. The following procedure is applied for each dataset:
The dataset is divided into splits or folds.
We fit the KDE model using range of l values, specifically choosing 25 uniformly spaced intervals from 10 1.3 to 10º
The fitted KDE models are evaluated by computing the log probability 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 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
Lenergy in 16
Logistic (Energy): Logistic (Energy): 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 are labeled O, while the samples in the test set are labeled as 1. The weight assigned to point x in the calibration set is determined by the estimate ((x)) of P(C 1|x x) obtained from the classifier, and calculated as p(x)/1-((u) To implement the logistic classifier, we utilized the scikit-learn library, employing the default hyperparameters for the classifier.
15

---

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 [.6. We sample 200 points from these generative models on properties described in Appendix [B.62] 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, experiment 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 our experiments.
Reinvent [4] is reinforcement learning-basee de novo drug design model. The model uses an RNN to generate SMILES strings.
Reinvent:
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 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-likenn of how likely a molecule is to possess drug-like properties based on its chemical structure.
of a compound. QED provides quantitative measure
+ activity: JNK3 activity refers to the activity of molecule against c-Jun N-terminal Kinases-3 (JNK3) protein. This is 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 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 wish to obtain uncertainty estimates. LogP, also known as the logarithm of the partition coefficient, is a property used to quantify the lipophilicity of drug molecule. The lipophilicity of 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.
target property, i.e. the property for which
logP can be computed cheaply using a computational oracle, which obtained it from TDC[2)] Note that this property in reality would be an experimentally determined property (such 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 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.
16

---

C Results
In section S we have provided results for experiments 0.1. Here, we provide additional results for the experiments at = 0.05 and a = 0.2.
C.1
Property prediction results
Dataset
Fingerprint split (a 0.05)
Fingerprint split (cr = 0.2)
CoDrug(energies)
CoDrugefestures)
Unweightedibaslime)
CoDrug(energies)
(nweightellbuseine)
AMES(Y=0)
0.96(0.07)
6.93(0.09)
1,0000.01)
(.85(0.066
0.78(0.06)
1,0000.00)
AMES(Y=1)
0.94(0.08)
0.95().08)
0.78(0.19 19)
0.79(0.05)
0.82(0.06)
0.34(0.14)
Clin TinTox(Y-00
0.94(0.08)
0.83(0.04)
0.91(0.07)
0.68(0.07)
0.67(0.04)
0,58(0.04)
Clin (Y-1
0.8800.11
0.73(0.0)))
0.93(0.00)
0.73(0.00)
(.4(00.03)
0.7310.00)
HIV(Y-0)
0.94(0.10)
0.94(0.14)
0.96(0.08)
0.81(0.88)
0.74(0..3))
0.89(0.07)
HIV(Y=l)
0.95(0.08)
0.997(..07)
0.9000 14)
(.85(0006)
0.50Q0004)
0.72(0.11)
Tox21(==0)
0.87(0.03)
0.71(0.03)
Tox21(Y-1)
0.97(0.01)
0.97(0.02)
0.97(0.01)
0.95(0.01)
0.93(0.02)
0.97(0.00)
Datesct
Scalfold split (or 0.05)
Seaffold split 0.2)
AMES/Y-0)
0.93(0.04)
0.90x0.04)
0.80/0.02)
0.6660.04)
AMES(Y==1
0.9000.02)
0.87(0.115)
0.91(0.02)
0.72(0.01)
0.66(0.02)
0.77(0.92)
(o=()
0.95(0.05)
0.8910.01>
0,94(0.02)
0.80(0.03)
0.74(0.01)
0.75(0.01)
CiiTTxu-Y->>
0.95(0.02)
0.77(0.03)
0.81(0.02)
HIV(Y-0)
0.8900.05)
0.88(D.177)
0.96(0.01)
0.72(0.07)
0.71(00.06)
0.81(0.91)
HIV(Y=I)
0.95(0.03)
0.94(0.07)
0.94(0.06)
0.81(0.01)
0.81(0.05)
0.73(0.05)
Tox21(Y=0)
0.95:0.07)
0.890.04)
(194(0.13)
0.81(0.05)
0.73(0.05)
0.75(0.09)
Tox21(Y=1)
0.94(0.06)
0.96(0.05)
0.97(0.04)
0.77(0.05)
0.80(0.04)
0.82(0.06)
Table 6: Coverage of CoDrug and baseline unweighted CP, under different datasets and distribution shifts at a = 0.05 and a = 0.2. The realized coverage rate closest to the target coverage 1 a (best) marked in bold. The second best coverage (in case better than unweighted) is marked in bold and gray, Results are averaged over 5 random runs.
C.2 De Novo Drug design experiment results
REINVENT ta (.16))
GraphGiA ta - (.15)
REINVENT - 0.2)
GraphGA ta (1.2)
Objective
Y
Colrag (Energy)
Unweighted
Callrug (Energy)
Unweighted
Colrrag (Energy)
Unweighted
Callrug (Eneryy)
Unweighted
QEI
0
0.97 (0.02)
0.94 (0.06)
0.97 (0.02)
0.81 (0.23)
0.87 (0.1)
0.53 (0.25)
0.75 (0.09)
0.41 (0.26)
QED
0.97 (0.02)
0.96 (0,07)
1.0 (0.01)
1.0 (0.01)
0.81 (0.18)
0.78 (0.26)
0.9 (0.08)
0.88(0.13)
INKC+QED
0
0,89 (0.15)
0.83 (0.34)
(0.02)
0,72
0.2:0.4)
0.73 (0.06)
(0.15)
INK3+QEE
0,98
(0.0)
(0.09)
GSK
(0.05)
OSK36+QFI
(8.15)
(0.0)
(0.02)
(0.00)
(0.37)
(0.04)
(8.07)
(0.04)
Table 7: Observed coverages on molecules sampled by generative models at & = 0.05 and a 0.2.
The realized coverage rate closest to the target coverage (1 Q.) is marked in bold. For each experiment, a set of 200 points optimized w.r.t. the "Objective" using the generative models GraphGA and REINVENT are sampled, similar the procedure in Appendix B.6. Using the proposed method improves coverage in almost all scenarios.
17

---

D
List of commonly used notations and terms
Table 8: List of key notations used in the paper
Symbol
Description
Refers the user-defined target coverage level in conformal prediction. It determines the confidence level of the prediction regions.
((x))
Corresponds to a prediction set of an input 2; obtained from conformal prediction method.
X,
Refers to the input features corresponding to data point in a machine learning model.
Y;
Refers the label corresponding data point in machine learning model. Refers to the base classifier of the model. It is the underlying algorithm or model used to make predictions on the data.
f
fx(x)
Refers the yt logit of the classifier on a data point æ without applying the softmax layer.
(x)
denotes the classwise probability scores of the data point æ with respect to the model after the application of the softmax layer on the logits fy(x)) It represents the probability of the data point belonging to class y.
E(x))
data point x in an energy-based model. The energy is computed based on the logits of the classifier.
Refers to the energy
F (or1)2
Refers to an (unweighted) cumulative distribution function computed from a set of values (v: |
w(x))
Refers to the likelihood ratio or weight assigned to data point in weighted conformal prediction. It quantifies the importance of the data point during calibration.
Fin #N+1
Weighted cumulative distribution function computed using the likelihood ratios w from weighted conformal prediction. It represents the distribution of the weighted values.
D Dirain
Refers to the distribution of molecules from which the training is sampled. It represents the underlying data distribution used to train a machine-lcarning model.
Dent
Refers to the distribution of molecules from which the calibration set is sampled. represents the underlying data distribution used to calibrate a conformal prediction model.
Poul
Refers to the true density of an input molecule X in the calibration set distribu- tion.
Plest
Refers to the true density of input molecule X in the test set distribution.
Phoaet
Refers to the density of an input molecule X computed from kernel density estimation (KDE) on the test distribution. is an estimation of the probability density function of the input molecule in the test set distribution.
Pheat
Refers to the density of an input molecule computed from KDE on the calibration set distribution. It is estimation of the probability density function of the input molecule in the calibration set distribution.
Table List of commonly used terms the paper
Term
Description
Activity (property)
Refers to the ability of drug to bind to specific target molecule and produce a biological effect. It is an important property to consider in drug discovery.
ADMET properties
Refers to the absorption, distribution, metabolism, excretion, and toxicity of drug candidate. These properties play a critical role in determining the safety and efficacy of drug.
18

---

Alpha
In conformal prediction, alpha refers to the user-defined confi- dence level used to construct the prediction sets. The parameter determines the amount of error that the user is willing to tolerate in the predictions and is typically set to a small value, such as 0.1 or 0.2. A smaller alpha typically results in a wider prediction set. In the Mondrian Inductive Conformal Prediction framework, cal- ibration refers to the procedure of using the calibration set to determine the threshold for each class. The objective is that the proportion of true labels across prediction sets matches the desired confidence level, as specified by the alpha parameter A calibration set is a labeled subset of the dataset held out from the training set used to estimate the threshold for each class (in Mondrian ICP).
Calibration( of conformal prediction)
Calibration set
Conformal Prediction (CP)
A framework for constructing reliable prediction intervals or sets at desired confidence level for a given machine learning model. The framework can be used with any machine learning model. Coverage of a conformal predictor the proportion of times that the true label falls within the prediction sets produced by the predictor, over all the inputs.
Coverage
Covariate shift
A phenomenon that occurs when the distribution of the input data changes between the training and testing phases of machine learning model. It is assumed that the conditional distribution of the target variable given the input features (P(YIX)) remains the same across the training and test sets.
Cumulative Distribution Function (CDF)
CDF gives the cumulative probability of the random variable taking on value less than or equal to a particular value.
De novo Drug design model
A machine learning model used to generate novel drug candidates with desired properties. These models are based on generative models such as Variational Autoencoders, Reinforcement Learn- ing, or Genetic algorithm and explore large chemical space.
Energy-basee based model
A type of model that learns a function that assigns low energy scores to data points that are similar to the training data and high energy scores data points that are dissimilar.
Exchangeability
Exchangeability refers to the property of sequence of random variables such that the joint distribution of any permutation of the variables is the same as the joint distribution of the original sequence. Independent and identically distributed (IID) implies exchangeability. Exchangeability is an important consideration in the Conformal prediction framework.
Fingerprint splitting
A method used to divide dataset of molecules into training and testing sets based on the similarity of their molecular fingerprints.
Generative model
A type of machine learning model that learns the distribution of dataset and can be used to generate new data points (in this case drug molecules) with similar properties.
Kernel Density Estimation
Kernel density estimation (KDE) is non-parametric method for estimating the probability density function of a random variable based on set of observations. It involves placing a kernel at each data point and summing the kernels to obtain a smoothed estimate of the density function.
Mondrian ICP
Mondrian Inductive Conformal Prediction (Mondrian ICP) is variant of the conformal prediction framework that provides class-wise coverage guarantees for multi-class classification prob- lems. The prediction sets are constructed to provide class-wise coverage guarantees, meaning that they are guaranteed to contain the true class label with a certain probability (determined by user-defined confidence level) for each class.
19

---

Prediction set
A prediction set is set of candidate labels (class values) for given input. A prediction set is considered valid if it contains the true class label of an input.
Quantile Function
A function that maps probability corresponding value in distribution.
the inverse of the cumulative distribution
function.
Scaffold splitting
A method used to divide a dataset of molecules into training and testing sets while ensuring that the two have similar scaffold diversity.
Toxicity
Refers to the potential drug to cause harm living organisms. It is an important ADMET property consider in drug discovery. A conformal predictor is said to be valid if its coverage level is equal to the user-defined significance level (usually denoted by alpha) used to construct the prediction sets.
Validity
20


Reviews:
[{"review_id": "wiD3C6tVYH", "reviewer": "Reviewer_SDHG", "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.", "questions": "No questions", "limitations": "There is no discussion of limitations.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "### Originality\n\nThis paper puts together existing ideas that provide uncertainty estimates for molecular property prediction and related tasks.\n\n### Quality\n\nQuality is high, the method is evaluated on many different prediction and generative tasks.\n\n### Clarity\n\nThe paper is well-written and easy to understand.\n\n### Significance\n\nProper 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the overall positive assessment of the paper. In our response, we would like to address the question about the applicability of KDE and elaborate on the limitations of the proposed framework. \n\n- While we demonstrated that KDE can provide asymptotic coverage guarantees, this may not necessarily translate to improved performance in scenarios with limited sample sizes. Although our experiments demonstrate that the convergence tends to approach the target value in most scenarios, we do acknowledge that there are a few instances where the improvement in coverage is limited. For example, in Table 2, we observe that the improvement in coverage for ClinTox (Y=1) is limited due to the relatively small size of the calibration set, consisting of only 19 points. As such, an important area for further research is to explore ways to obtain likelihood ratios that are more data-efficient, particularly in scenarios with smaller calibration sets. \n\n- 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. As part of our future work, we intend to explore and extend the proposed framework to encompass regression tasks.\n\n- Additionally, our current work primarily focuses on small molecules. Nevertheless, we recognize that covariate shift is a common occurrence in various chemical and biological phenomena, making the framework presented in this work applicable in broader contexts. Nonetheless, obtaining high-quality feature vectors for computing likelihood for different applications remains a challenge that warrants further research.\n\nIn our current version of the paper, we have outlined the general limitations in the conclusion section and specified the limitations concerning KDE in the experiments section. We will make necessary revisions to emphasize the limitations more prominently in the text."}, {"review_id": "1LJHeHvsnH", "reviewer": "Reviewer_e8Bh", "summary": "The 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.", "questions": "* 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.", "limitations": "The authors mention the limitation of their work in the conclusion section.", "rating": 3, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* 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.", "weaknesses": "* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback and would like to address the concerns regarding the novelty of our paper. Our work presents several contributions, which we believe are novel for various reasons:\n\n- We focus on conformal prediction (CP), a powerful and increasingly popular tool for obtaining distribution-free uncertainties. While CP has been applied to high-stakes applications including drug discovery ([8] [9] [10]), we identify and tackle the challenge of providing statistically rigorous uncertainty quantification in the presence of distribution shift—a common occurrence in drug discovery due to the vastness of the chemical space. Previous studies primarily concentrated on non-conformity measures, neglecting the critical consideration of distribution shift, which is a key aspect that significantly impacts practical usability.\n- To this end, we first show that the coverage guarantee offered by the conformal prediction framework is indeed violated in a practical drug discovery context. We devise experiments that simulate realistic drug discovery scenarios, employing fingerprint and scaffold splitting techniques. The experiments clearly demonstrate the violation of the CP coverage guarantee and shed light on the existing methods' limitations. Leveraging recent advances in CP under covariate shifts, we propose and construct a framework that effectively bridges the coverage gap. This framework is showcased through experiments on multiple property prediction datasets, showcasing its applicability in real-world scenarios.\n- We tackle the challenge of obtaining likelihood ratios in CP under covariate shift, which is not immediately clear from prior work ([11]). We propose a novel approach leveraging Kernel Density estimates to correct for covariates, and we empirically demonstrate its effectiveness. Further,  we show that the kernel density estimates are consistent, which means that the covariate shift is precisely adjusted for, and the coverage guarantee is recovered asymptotically.\n- We also develop a training framework that utilizes unlabeled data and an energy-based model formulation. This framework facilitates in extracting features that enable accurate computation of likelihoods using Kernel Density Estimation (KDE). Our results demonstrate the effectiveness of this end-to-end approach in addressing the coverage gap.\n-  In AI-based drug discovery, there has been a substantial focus on developing molecule-generative models ([3]-[6]). Despite their potential, their practicality is limited by the fact that only a small number of these molecules can be manufactured and experimentally tested. Thus, precise property estimates, accompanied by robust uncertainties, are crucial. We demonstrate that our proposed framework can be used with the molecules sampled from generative models, providing them with the ability to quantify uncertainties in a statistically robust manner.\n\n## Alternative methods for empirical comparisons: \nWe believe that we made our best efforts to conduct a comprehensive benchmark of the proposed method. In our experiments, we perform:\n1. Comparisons with vanilla CP without any correction, which is the most common approach used in various methods. \n2. We conducted ablation studies on different components of our approach - feature extraction and the energy-based training procedure.\n3. We also compared our method by using the likelihood estimation approach suggested by [11].\n\nConsidering the novel nature of uncertainty estimation for drug discovery under distribution shift, we welcome any reviewer's insights and suggestions for alternative methods of comparison that they may have in mind. \n\n[3] Fu, Tianfan, et al. \"Reinforced genetic algorithm for structure-based drug design.\" Advances in Neural Information Processing Systems 35 (2022): 12325-12338.\n[4] You, Jiaxuan, et al. \"Graph convolutional policy network for goal-directed molecular graph generation.\" Advances in neural information processing systems 31 (2018).\n[5] Gómez-Bombarelli, Rafael, et al. \"Automatic chemical design using a data-driven continuous representation of molecules.\" ACS central science 4.2 (2018): 268-276.\n[6] Jin, Wengong. et. al. Junction tree variational autoencoder for molecular graph generation. In ICML, pages 2323–2332.PMLR, 2018.\n[8] Zhang, Jin., et al. (2021). Deep learning-based conformal prediction of toxicity. Journal of chemical information and modeling, 61(6), 2648-2657.\n[9] Isidro Cortés-Ciriano and Andreas Bender. Concepts and applications of conformal prediction in computational drug discovery. arXiv preprint arXiv:1908.03569, 2019.\n[10] 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.\n[11] Tibshirani, Ryan J., et al. \"Conformal prediction under covariate shift.\" Advances in neural information processing systems 32 (2019)."}, {"review_id": "HjG7eoSKPo", "reviewer": "Reviewer_nCYs", "summary": "This 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.", "questions": "1. 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```", "limitations": "I 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.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**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.", "weaknesses": "**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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their overall positive assessment of the paper. We would like to address the questions/concerns raised by the reviewer. \n\n## Importance of ZINC unlabelled data\nWe report the importance of using ZINC-based unlabelled data in Table 3 and Section 5.2 of our paper. The \"CoDrug (NoEnergy)\", corresponds to the loss function trained without the energy regularization (i.e without the unlabelled ZINC data) as per Eq. (7). In our experiments, we observe a general improvement in bridging the coverage gap by including unlabelled ZINC dataset in our training procedure. Notably, we observed an average increase of 5\\% (at an alpha=0.3) and 4\\% (at an alpha=0.2) in coverage when analyzing the worst 25% of the experiments.\n\n## Estimating generative model trained on ZINC to estimate the likelihood ratio\nWe are unsure of the reviewer's intended meaning when mentioning \"using a generative model trained on ZINC to estimate likelihood ratio.\"  One way is to utilize features extracted from a generative model trained on ZINC, such as a VAE. In that case, it is indeed feasible to extract features from the generative model. However, a challenge here is that features often have hundreds of dimensions, making KDE more data-hungry, and requiring a larger number of data points for accurate density estimation. Nevertheless, it is still possible to use these features as inputs to train the EBM and extract low-dimensional features, which might yield similar results as our approach.\n\nAnother potential way is to use a generative model trained on ZINC data to directly compute the likelihood of samples.  That's an interesting direction to explore, which may be suitable for larger datasets. \nThis requires fine-tuning a generative model like VAE or normalizing flow for each specific dataset/task distribution, for both the test and calibration datasets. Learning this distribution can be challenging when the calibration and test sets consists of only a few hundred data points or lesser.\n\n\n## ReLU not applied in the provided code\nWe apologize to the reviewer for the confusion. In our experiments, we utilized features extracted from the \"EnergyFeaturesModel,\" which is also available in the same model.py file. This class has the ReLU function applied while extracting classification features. We will eliminate any such references when we release the final code version.\n\n## Regarding suggestions for improving the clarity of the paper\nWe thank the reviewer for their suggestions on enhancing the paper's clarity and understanding of the limitations imposed by space constraints. \nIn light of the suggestions provided, we will reorganize the paper by placing section 4.3 before section 4.2, and we will also make necessary adjustments to improve the coherence of the text. Additionally, in order to eliminate any potential ambiguity in section 5.3, we will carefully review the content to enhance its clarity. We will explicitly state that the molecules being tested originate from a distinct black box generative model, and our current method serves the purpose of providing uncertainty estimates on this set of molecules.\n\nWe also appreciate the reviewer's understanding of the challenge imposed by the space constraints. We have outlined the general limitations in the conclusion section and specified the limitations concerning KDE in the experiments section. We will make necessary revisions to emphasize the limitations more prominently in the text for the final version.  \n"}, {"review_id": "ypd4UcsM9B", "reviewer": "Reviewer_uhvu", "summary": "The 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. ", "questions": "Can 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? ", "limitations": "-", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Drug 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. ", "weaknesses": "Significance 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the author for their overall positive assessment of the paper. In our response, we would like to address the posed questions and elaborate on the stated weaknesses.\n\n## Clarification regarding the exchangeability argument\nA sequence of random variables is exchangeable if the joint probability distribution remains the same, regardless of the order in which the variables appear. If a sequence of random variables is drawn from i.i.d., then it is also exchangeable because permuting the order does not change the joint probability distribution.  In drug discovery experiments, this assumption is frequently violated because the molecules in the test set do not originate from the same distribution as the training set.  For instance, in various virtual screening and generative drug design models, the primary objective is to generate molecules optimized for a specific target property, like binding affinity or docking score. These screening libraries or generative methods explore an extensive chemical space to search for molecules that exhibit the desired properties. \n\nConsider the task of estimating interval or set properties of generated molecules based on drug properties using an Oracle trained on those attributes (like toxicity or PK/PD).  The distribution of the sampled set of molecules is distinct from the distribution of the molecules seen during training. This is due to the fact that a toxicity model, for example, might have been trained only on molecules with specific structural frameworks or features, violating the assumption of exchangeability.   In the current study, our focus is on developing techniques to produce robust and statistically valid set-valued or interval-valued predictions in such contexts. We use the conformal prediction framework in scenarios where the test and the calibration set are not exchangeable, but linked by covariate shifts.\n\n## Relevance of data-splitting in real-world\nAlthough it is challenging to perfectly emulate real-world settings for the application, we have designed experiments using scaffold and fingerprint splitting to mimic distribution shifts. In scaffold splitting, molecules are divided based on Bemis and Murcko scaffolds, effectively reducing a compound to its core components by removing all side chains. This approach closely resembles real-world data collection, where certain biological assays involve molecules with nearly identical scaffolds. Scaffold splitting has become a standard technique in several out-of-distribution (OOD) generalization benchmarks in drug discovery [18] [19]. In addition to scaffold splitting, we also utilize fingerprint splitting, obtained from the DeepChem library [26], to emulate scenarios where the train and test sets are sourced from different regions of the chemical space. This technique allows us to simulate situations commonly encountered in screening library searches and generative models, which explore a vast chemical space. We welcome the reviewer's suggestions for any other approaches to perform splitting with a more realistic distribution shift. \n\n## Broader applicability of the design setting\nWe present a few general examples where the coverage guarantee offered by the conformal prediction framework can be beneficial in drug property prediction:\n\nIn classification settings: The Mondrian conformal prediction setting, as demonstrated in this paper, provides class-wise guarantees. This is beneficial in scenarios where certain classes are rare, such as identifying rare active compounds (or non-toxic compounds) in a set of molecules. By setting a target coverage rate, for example, an alpha of 0.1 or a 90\\% target coverage, the prediction set will contain the class active (if the original class is active) with a marginal probability of 90\\%. This can potentially aid in cost-saving by preventing the omission of molecules with desired properties during experimentation.\n\nIn multi-label classification: Consider drug-drug interaction prediction, where, the objective is to generate a set of drug molecules with potential interactions for a given drug molecule or a set of potential interactions for a drug-drug pair. It is crucial to include all drug molecules with possible interactions confidently and the CP framework can help provide prediction sets with a coverage guarantee. This approach can also be applied similarly in drug-adverse reaction prediction tasks.\n\nIn regression: In drug discovery, predicting continuous properties like potency, affinity, and solubility is crucial. Conformal prediction in regression tasks provides prediction intervals at a specified coverage rate. For instance, a 90% target coverage rate implies the interval holds true for about 90% of cases. This helps researchers gauge uncertainty, prioritizing experimental efforts on molecules with the highest potential for desired properties.\n\nWe believe that the proposed framework can be adapted to a wide range of tasks and settings beyond the small molecule classification tasks discussed in this paper, and our future work will focus on exploring these possibilities.\nThis includes settings involving regression, multi-label classification, and utilization in other chemical/biological contexts like RNA sequence/structure analysis, protein property prediction, as well as peptide and antibody design. These tasks frequently encounter covariate shifts, making the adaptability of the proposed framework highly relevant.\n\n[18] Ji, Yuanfeng, 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).\n[19] Gui, Shurui, et al. \"Good: A graph out-of-distribution benchmark.\" Advances in Neural Information Processing Systems 35 (2022): 2059-2073.\n[26] Ramsundar, Bharath, et al. Deep learning for the life sciences. \" O'Reilly Media, Inc.\", 2019.\n"}]

Author rebuttal:

Meta-review: The paper puts forward a method that allows accurate conformal prediction of drug properties in the case of covariate shift, that is, when the test distribution does not match the training distribution, which often is the case in practice. The authors analyze the method theoretically proving asymptotic convergence and evaluate its behaviour empirically. The reviewers find the paper well written and technically sound. The reviewers posed some questions on the significance of the contribution, which were addressed by the rebuttal. On balance, the contribution appears to be sufficient.
Decision: Accept (poster)