SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Since  $x$  here refers to molecular entities (e.g. SMILES strings), we cannot use a Gaussian kernel directly. Instead, we use embeddings or prediction logits produced by a trained model  $f$  as the input to the kernel. Under the assumption that KDE accurately reflects the true density of the underlying distribution, we could construct kernel density estimators for both the calibration set and test sets (remember that we do not have access to the test labels but have access to the input  $X$ ), and use
$$
\hat {w} (x) := \hat {p} _ {h _ {\text {t e s t}}} (x; \mathcal {D} _ {\text {t e s t}}) / \hat {p} _ {h _ {\text {c a l}}} (x; \mathcal {D} _ {\text {c a l}}) \tag {21}
$$
to replace the unknown  $w$  in Eq. (10), giving us the final prediction set:
$$
\hat {C} ^ {\mathrm {C o D r u g}} (x) = \left\{y: 1 - p _ {y} ^ {f} (x) \leq \text {Q u a n t i l e} \left(1 - \alpha ; F _ {x} ^ {\hat {w}}\right) \right\} \tag {22}
$$
$$
F _ {x _ {N + 1}} ^ {\hat {w}} := \hat {w} (x _ {N + 1}) \delta_ {\infty} + \sum_ {i \in [ N ]} \hat {w} (x _ {i}) \delta_ {1 - p y _ {i} ^ {f} (x _ {i})} / \hat {W} \tag {23}
$$
where  $\hat{W} = \sum_{i=1}^{N+1} \hat{w}(x_i)$  is a normalizing factor. Here,  $\hat{p}_{h_{test}}(\cdot; \mathcal{D}_{test})$  is constructed using samples from the test data with an optimal bandwidth  $h_{test}$  chosen on the test data via cross-validation, and  $\hat{p}_{h_{cal}}(\cdot; \mathcal{D}_{cal})$  is constructed similarly but on the calibration data. It is clear that, as the number of samples from  $\mathcal{D}_{cal}$  and  $\mathcal{D}_{test}$  increases,  $\hat{w}$  converges to  $w$  in Eq. (10), and  $\hat{W}^{\mathrm{CoDrug}}$  recovers the coverage guarantee asymptotically. In practice, recovering asymptotic coverage on a finite amount of data is challenging. However, the coverage tends to approach the target value as we observe in our experiments. The overall procedure for density estimation is depicted in Fig. 1(d).
Algorithm 1 summarizes all the components in Section 4. In Section 5, we will verify the efficacy of CoDrug in property prediction tasks, and molecules sampled from de novo drug design models.
Algorithm 1 Procedure for Property Prediction
Split the dataset into training set  $\mathcal{D}_{train}$  and calibration set  $\mathcal{D}_{cal} = \{z_i\}_{i=1}^N$ .
Train a neural net classifier  $f$  on  $\mathcal{D}_{train}$  by minimizing Eq. (16).
Compute the KDE  $\hat{p}_{h_{cal}}(\cdot ;\mathcal{D}_{cal})$  for all points in  $\mathcal{D}_{cal}$  using Eq. (18).
Test Time, for a test set  $\mathcal{D}_{test}$ :
Compute KDE  $\hat{p}_{h_{test}}(\cdot ;\mathcal{D}_{test})$  for all points in  $\mathcal{D}_{cal}$  using Eq. (18).
For any  $x_{N + 1}\in \mathcal{D}_{test}$ , compute  $\hat{w} (x)$  and  $\hat{w} (x_i)$  for  $x_{i}\in \mathcal{D}_{cal}$ .
Construct the prediction set  $\hat{C}^{\mathrm{CoDrug}}(x_{N + 1})$  using Eq. (22).
In this section, we put our proposed method, CoDrug, to the test on various drug discovery tasks. Section 5.1 describes the datasets used and key implementation details. Section 5.2.1 empirically demonstrates the loss of validity in conformal prediction sets on different drug discovery datasets. Section 5.2.2 shows how the setup improves the validity of the conformal prediction sets. Section 5.3 confirms the utility of CoDrug in de novo drug design. We include additional details on implementation, datasets, and hyperparameters in the appendix.
5.1 Data and Implementation Details
- Splitting Strategies: To demonstrate the effectiveness of CoDrug under covariate shift, we use two different strategies when creating calibration/test splits. In both strategies, we try to create calibration and test splits that are dissimilar to each other, which is a challenging but realistic setting in drug discovery. We used the DeepChem[26] library for splitting. In scaffold splitting, the dataset is grouped based on chemical scaffolds, representing core structures of molecules. The test set and train set consist of different scaffolds. In fingerprint splitting, the dataset is partitioned based on Tanimoto similarity of molecular fingerprints [27]. Molecules with the highest dissimilarity in terms of Tanimoto similarity are included in the test set.
- Datasets: We use four binary classification datasets for toxicity prediction (AMES, Tox21, ClinTox) and activity prediction (HIV activity), obtained from TDC [28]. To train the Energy based model, we obtained the unlabelled data from the ZINC-250k dataset [29], a subset of the ZINC that covers a large chemical space. For each dataset and split type, we removed the molecules that are similar to the training (and calibration) set from the unlabelled dataset.
- Classification Model: The architecture of our classifier  $f$  is AttentiveFP [1], a graph neural network-based model. We chose AttentiveFP as it has state-of-the-art results in several drug property prediction tasks. It is trained using the objective function described in Eq. (16).
- De Novo Drug Design Experiments: In Section 5.3, we perform experiments to construct conformal prediction sets on molecules sampled from de novo drug design models. As generative models, we use REINVENT[14] and GraphGA[30] (top-ranked methods in MolOpt [31] benchmark). The models are optimized to sample molecules w.r.t. three popularly used computational oracles - QED (quantitative estimate of drug-likeness), JNK3 activity, and GSK3B activity. For building the conformal prediction sets, we chose logP as our target property, assigning values in the range of
[1.0,4.0] a class of  $\mathrm{Y} = 1$ , and  $\mathrm{Y} = 0$  otherwise (representing the drug-like range [32]). We obtain the computational oracles from TDC [33], and generative models from MolOpt package[31].
5.2 Property Prediction Results
5.2.1 Unweighted conformal prediction (baseline)
In this section, we demonstrate the unpredictable behavior of the unweighted CP method without proper correction under distribution shift. Table 1 shows the results of conformal prediction under various distribution shift conditions. "Random" refers to the ideal/unrealistic scenario where the test and calibration samples are split randomly (aka. no distribution shift). "Scaffold" and "Fingerprint" denote scenarios in which there is a distribution shift between the test and training data outlined in the Methods section. In all scenarios,  $15\%$  training set is held out for calibration, and prediction sets are calculated using the algorithm described in Algorithm 1 without any correction.
From Table 1, we observe that the Random configuration demonstrates little loss in coverage and coverage decreases under distribution shifts (Scaffold and Fingerprint). But for fingerprint and scaffold split, unweighted CP failed to provide target coverage and exhibit unpredictable behavior. For instance, at  $\alpha = 0.2$ , under fingerprint split, Unweighted has a coverage of 0.34 against a target coverage of 0.8 for the AMES dataset, while achieving a very different coverage of 0.77 with scaffold split on the same dataset.
Random	Fingerprint	Scaffold	Random	Fingerprint	Scaffold
Dataset	α=0.1	α=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
ClinTox(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.89	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 closely with the target  $1 - \alpha$ . However, in most datasets with fingerprint and scaffold splits—reflective of more realistic scenarios—the baseline method falls short. Often, substantial deviations in coverage confirm the unpredictability of unweighted CP when exchangeability ceases to apply. Here, values not significantly deviating from  $1 - \alpha$  at a p-value of 0.05 are highlighted in bold, indicating desirable performance.
5.2.2 Weighted conformal prediction improves coverage
In Table 2, we present the benefits of using weighted CP via CoDrug. The table depicts results from conformal prediction using 3 different schemes.
- CoDrug (Energy): This variant of CoDrug uses weights computed from KDE on the prediction logits of the trained EBM, as described in Section 4.3.
- CoDrug (Feature): This variant of CoDrug builds the KDE instead on the features extracted from the penultimate layer of the trained EBM.
- Unweighted: Refers to the unweighted prediction conformal prediction (baseline).
In both weighting schemes of CoDrug, we use KDE to estimate densities and find that weighting using energies improves the coverage towards the target coverage  $1 - \alpha$  in most cases. We notice the highest improvement in the Fingerprint splitting scenario for the AMES(Y=1) category, where the coverage improved from 0.63 to 0.88 (target coverage 0.9). Note that the coverage is "improved" if it is closer to  $1 - \alpha$ -improvement does not always mean higher coverage, because an unusually high coverage often indicates unpredictable behavior of the underlying model.
While our energy-weighting approach generally improves coverage, there are rare instances where it underperforms compared to the baseline. A prominent example is the ClinTox dataset with  $Y = 1$ , which sees limited improvement or even a reduction in coverage. This is due to the constraints of the density estimation procedure, which relies on the quantity of available data. Notably, this dataset is the smallest and most imbalanced, with only 19 points in the calibration set for class  $Y = 1$ .
Additionally, our results show that using energy weighting leads to better overall coverage than directly weighting the features. This is likely because the energy values are two-dimensional, while the features are an eight-dimensional vector: As the dimension of the feature input to KDE increases, one typically requires more samples to get a high-quality density estimate.
Dataset	Fingerprint Splitting	Scaffold Splitting
CoDrug (Energy)	CoDrug (Feature)	Unweighted (baseline)	CoDrug (Energy)	CoDrug (Feature)	Unweighted (baseline)
AMES(Y=0)	0.93(0.03)	0.87(0.03)	1.00(0.00)	0.85(0.02)	0.89(0.02)	0.82(0.01)
AMES(Y=1)	0.88(0.03)	0.90(0.03)	0.63(0.05)	0.83(0.01)	0.79(0.01)	0.85(0.03)
ClinTox(Y=0)	0.86(0.04)	0.76(0.02)	0.78(0.02)	0.90(0.03)	0.83(0.01)	0.84(0.00)
ClinTox(Y=1)	0.73(0.00)	0.69(0.08)	0.80(0.00)	0.85(0.03)	0.83(0.00)	0.97(0.04)
HIV(Y=0)	0.89(0.06)	0.87(0.07)	0.93(0.04)	0.82(0.08)	0.82(0.04)	0.91(0.01)
HIV(Y=1)	0.92(0.05)	0.95(0.03)	0.84(0.07)	0.90(0.01)	0.90(0.05)	0.87(0.03)
Tox21(Y=0)	0.90(0.02)	0.80(0.02)	0.77(0.03)	0.91(0.03)	0.83(0.05)	0.89(0.05)
Tox21(Y=1)	0.97(0.00)	0.96(0.01)	0.97(0.00)	0.86(0.05)	0.91(0.05)	0.93(0.03)
Table 2: Coverage of CoDrug and baseline unweighted CP, under different datasets and distribution shifts at  $\alpha = 0.1$ . The realized coverage rate closest to the target coverage  $1 - \alpha$  (best) is 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. Results for different  $\alpha$  values are 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 ( Eq. (16)), and covariate shift correction. In addition to CoDrug (Energy) and Unweighted (baseline) reported in the previous section, we also compare with:
- CoDrug (NoEnergy): We use the same protocol as CoDrug (Energy) but the models are trained without the energy regularization term  $\mathcal{L}_{energy}$  in Eq. (16).
- Logistic (Energy): In this experiment, the features are same as CoDrug (Energy), but KDE is not used to estimate densities. Instead, the weights  $w(x_{i})$  in Eq. (10), are given by  $\hat{p}_{(x_i)} / 1 - \hat{p}_{(x_i)}$ , where  $\hat{p}_{(x_i)}$  is obtained by fitting a classifier to features in calibration and test sets (suggested by [11]).
The results are depicted in Table 3. In the table, we compile the "mean absolute coverage deviation" across all the datasets and different random runs from all the experiments reported in Table 2 at different values of  $\alpha$  (i.e. Mean of  $|\mathrm{Observed\_Coverage} - (1 - \alpha)|$  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  $\alpha$ . We paid close attention to the "Tail  $25\%$  ", where we presented the metrics for the worst  $25\%$  performing experiments and CoDrug (Energy) outperforms all the other variants in comparison by a substantial margin inducing that all the different components in the CoDrug pipeline are helpful. The mean absolute coverage deviation from the target at  $\alpha = 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 absolute coverage deviation (tail 25%)
α = 0.3	α = 0.2	α = 0.1	α = 0.3	α = 0.2	α = 0.1
Unweighted (baseline)	0.157 (0.14)	0.12 (0.12)	0.081 (0.07)	0.347 (0.14)	0.276 (0.13)	0.176 (0.08)
Logistic(Energy)	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 (Energy)	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  $\alpha$ , 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 a computational oracle. After molecule sampling, validating their experimental properties, such as ADMET (Absorption, Distribution, Metabolism, Excretion, Toxicity), is crucial for safety and efficacy. When a machine learning model trained on such properties is available, assessing the uncertainty associated with the predictions before experimental validation is critical. However, note that the distribution of sampled molecules may substantially deviate from the training data, affecting the prediction sets' target coverage from CP.
In this section, we demonstrate the application of CoDrug on molecules generated by a de novo drug design model. We consider the de novo drug design model as a black box, that has been optimized w.r.t a certain objective. We experiment with two models - GraphGA [30] and Reinvent [14]. To predict properties, we compiled a dataset of logP values, as it can be computed cheaply with a computational oracle. We note that in reality, this dataset could correspond to experimental properties like ADMET. However, since it is not feasible to validate these properties for molecules generated from de novo drug design models, we use logP to demonstrate the method. The results of our experiments are depicted in Table 4 at a target alpha value of 0.1. Our proposed method consistently enhances coverage in all cases and exhibits a substantial improvement over the unweighted version. For example, in the "gsk3b+qed" objective, the unweighted version has a coverage of 0.44 against a target of 0.9, whereas our proposed method improves coverage substantially. The mean absolute coverage deviation from the target at  $\alpha = 0.1$  for CoDrug (Energy) is 0.05, a relative improvement of over  $60\%$  on the Unweighted version (0.14).
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	1	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)
GSK3b+QED	1	0.79 (0.08)	1.0 (0.0)	0.98 (0.0)	1.0 (0.0)
QED	0	0.96 (0.0)	0.83 (0.04)	0.96 (0.0)	0.69 (0.1)
QED	1	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  $\alpha = 0.1$ . The realized coverage rate closest to the target coverage  $(1 - \alpha)$  are 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 is logP  $(1.0 < \log P < 4.0$  is considered  $\mathrm{Y} = 1$ ;  $\mathrm{Y} = 0$  otherwise [32]). Using the proposed method improves coverage in almost all scenarios.
While we demonstrated that KDE can provide asymptotic coverage guarantees, this may not necessarily translate to improved performance in scenarios with limited sample sizes. In our experiments, we do acknowledge that there are a few instances where the improvement in coverage is limited, where the availability of data is limited. As such, a direction for further research is to explore ways to obtain likelihood ratios that are more data-efficient, particularly in scenarios with smaller calibration sets.
It is worth noting that our current work focuses on addressing the coverage gap in classification tasks, and regression tasks were not included in this study. However, we recognize the importance of uncertainty quantification in regression settings, especially for various critical drug properties represented as regression problems, where our proposed framework can be extended with modifications. Furthermore, our current work primarily focuses on small molecules, yet covariate shift is a common phenomenon in various chemical and biological contexts. While this means that our framework could be more generally applied, obtaining high-quality feature vectors for computing likelihood in different applications remains a challenge that warrants further research.
We present a new method for uncertainty quantification in drug discovery, CoDrug, that effectively addresses the problem of co-variate shifts in test data. The proposed method involves a combination of three key steps, training an energy-based model for feature extraction and base classification, performing density estimation using KDE, and use the KDE to correct for covariate shift in conformal prediction to recover valid coverage. The results obtained in this study demonstrate the effectiveness of CoDrug in predicting valid conformal prediction sets and its utility in de novo drug design experiments. Our current work is limited to classification tasks in small molecules, but exploring its application to other chemical and biological tasks with covariate shifts is interesting for future work, including adapting the framework for regression tasks.
A.1 Proofs for Theorem 4.1
Denote  $T_{i} = 1 - p_{Y_{i}}^{f}(X_{i})$  as a new random variable. Because of our i.i.d. assumption (and that  $f$  is not trained on the calibration set),  $T_{1},\ldots ,T_{N + 1}$  are also i.i.d., which means that for  $m = 0,1,\dots ,N$ :
$$
\mathbb {P} \{| \{i \in [ N ]: T _ {N + 1} > T _ {i} \} | = m \} \leq \frac {1}{N + 1} \tag {24}
$$
$$
\Rightarrow \mathbb {P} \left\{\left| \left\{i \in [ N ]: T _ {N + 1} > T _ {i} \right\} \right| \geq (N - m) \right\} \leq \frac {m + 1}{N + 1} \tag {25}
$$
The left-hand-side probability is  $\leq$  (instead of  $<$ ) the right-hand-side due to the case when  $p_y^f$  is not continuous. Note also that
$$
1 - p _ {Y _ {N + 1}} ^ {f} \left(X _ {N + 1}\right) > t \Rightarrow | \{i \in [ N ]: T _ {N + 1} > T _ {i} \} | \geq \lceil (1 - \alpha) (N + 1) \rceil \tag {26}
$$
$$
\mathbb {P} \{1 - p _ {Y _ {N + 1}} ^ {f} \left(X _ {N + 1}\right) > t \} \leq \frac {\lfloor \alpha (N + 1) \rfloor}{N + 1} \leq \alpha . \tag {27}
$$
Finally, note that
$$
Y _ {N + 1} \notin \hat {C} \left(X _ {N + 1}\right) \Longrightarrow 1 - p _ {Y _ {N + 1}} ^ {f} \left(X _ {N + 1}\right) > t, \tag {28}
$$
we thus have
$$
\mathbb {P} \left\{Y _ {N + 1} \in \hat {C} \left(X _ {N + 1}\right) \right\} \geq 1 - \alpha . \tag {29}
$$
A.2 Proofs for Theorem 4.2
This is a direct result of Corollary 1 in [11], with our nonconformity score  $1 - p_{Y}^{f}(X)$  plugged in.
B Data and implementation details
B.1 Training Details of the base classifier:
- Deep Learning Frameworks: We use the Pytorch framework for the implementation of the models. The Graph Neural Network backbone is obtained from the DGL-LifeSci library [34].
- Training hyperparameters: We train the model using the PyTorch Lightning Framework for training. We use the ADAM Optimizer [35]. The batch size is set to 64, and the learning rate is set to 0.001.
- Architecture Details: The model architecture consists of a GNN layer (Attentive FP [1]), a readout layer, 2 hidden FCNN layers, and an output layer. The hidden state size in GNN is set to 512 dimensions. The linear layers have 256, and 8 dimensions respectively.
- Cheminformatics Processing : We use the RDKit library for handling molecular entities in Python. We use the DeepChem library for generating dataset splits.
- Energy Regularization hyperparameters : The parameters  $m_{in}$  and  $m_{out}$  in Eq. (15) are set to -5 and -35 respectively, and the parameter  $\lambda$  in Eq. (16) is set to 0.01. All of these hyperparameters are obtained from the reference implementation in [24].
- Splitting Ratio: The datasets are split in the ratio of 70:15:15, for training, calibration and testing the CP model.
- Error bars: All the experiments reported are over 5 random runs. The mean and the standard deviation across the random runs is reported in the table.
B.2 Splitting strategies:
In this section, we discuss the two strategies employed for creating calibration and test sets to ensure their dissimilarity. The implementation of these strategies was based on the DeepChem library [26]. We provide detailed explanations of the algorithms used for splitting below.
- **Scaffold Splitting:** The core structure of a molecule is represented by scaffolds [36]. Scaffold splitting aims to generate train and test sets that do not share any common scaffolds, thereby creating a challenging yet realistic scenario of distribution shift. This strategy is commonly used to evaluate out-of-distribution prediction algorithms [20, 18]. The scaffold splitting procedure is as follows:
- First, the scaffolds of all molecules in the datasets are identified.
- The scaffolds are sorted by their frequency in the dataset.
- The least frequent scaffolds are added to the test set until the desired number of test data points is reached.
- The remaining points are randomly divided into training and calibration sets.
- Fingerprint Splitting: A molecular fingerprint[27] is a compact binary representation of a molecule's structural features, capturing important information about its chemical composition and spatial arrangement. Fingerprint splitting utilizes molecular fingerprints to create distinct train and test sets. The primary objective is to include data points in the test set that exhibit the least maximum pairwise Jaccard similarity of fingerprints. The fingerprint-splitting procedure is as follows:
- First, the molecular fingerprints are computed for all the molecules in the dataset using Extended Connectivity Fingerprints (ECFP) [27].
- Pairwise Jaccard similarity is calculated for each data point in the dataset, considering its similarity with all other points. The Jaccard similarity between two fingerprints is determined by dividing the size of their intersection by the size of their union.
- The test set construction begins by selecting the data point with the least maximum Jaccard similarity to any other point in the dataset. This point is added to the test set.
- The iterative process continues until the desired number of test data points is reached. The remaining data points are assigned to the training/calibration set.
B.3 Property prediction datasets
B.3.1 Labelled data
We used four commonly used classification benchmarking datasets from ADMET properties. The datasets are obtained from Therapeutics Data Commons (TDC) [28]. We include the statistics of all the datasets used for property prediction in Table 5.