SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Siddhartha Laghuvarapu
Department of Computer Science  
University of Illinois Urbana-Champaign  
Urbana, IL 61801  
sl160@illinois.edu
Zhen Lin
Department of Computer Science  
University of Illinois Urbana-Champaign  
Urbana, IL 61801  
zhenlin4@illinois.edu
Jimeng Sun
Department of Computer Science  
Carle Illinois College of Medicine  
University of Illinois Urbana-Champaign
Urbana, IL 61801 jimeng@illinois.edu
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.
Drug discovery is a challenging and complex task, with a high failure rate and limited understanding of the chemical and biological processes involved. These contribute to making drug discovery an extremely costly and time-consuming endeavor. Recently, advances in deep learning have aimed to reduce the cost of drug discovery by proposing AI methods for developing accurate property prediction models and De Novo drug design models:
- Property prediction models aim to aid the laborious and expensive stages of drug discovery by building accurate supervised learning models that take in a drug representation as input and output a target property [1, 2].
The code associated with the paper is available at https://github.com/siddharthal/CoDrug/
- De novo drug design models, on the other hand, aim to discover new drug molecules that satisfy a set of pharmaceutical properties [3, 4, 5, 6].
With the high cost and significance of drug discovery, it is essential to have accurate and reliable uncertainty estimates in supervised learning models for property prediction. By providing set-valued or interval-valued estimates instead of solely relying on point estimates, uncertainty estimation enables more informed decision-making and reduces the risk of failures, making the drug discovery process more efficient. Conformal prediction (CP), pioneered by [7], offers a solution to uncertainty quantification for complex models like neural networks, by constructing provably valid prediction sets<sup>1</sup> in supervised learning models. Its application to drug property prediction has also been explored for various drug discovery tasks [8, 9, 10].
A crucial assumption in the CP framework is that the test samples are exchangeable with the holdout set used to calibrate the algorithm. In drug discovery, there is often a limited amount of available training or validation data. Furthermore, de novo drug design models or screening datasets sample molecules from a large chemical space, making the exchangeability assumption invalid.
In this paper, we deal with a situation where the training data originates from a distribution,  $P(X)$ , while the test data comes from a different distribution,  $P^{test}(X)$ . However, in both cases, the molecular properties, determined by the conditional distribution  $P(Y|X)$ , remain the same as they are governed by nature and unaffected by shifts in the input distribution, assuming testing parameters are stable. This is referred to as covariate shift. Although recent research in Conformal Prediction (CP) [11] suggests a method for correcting covariate shift, accurately estimating the precise level of covariate shift remains a practical challenge.
This paper proposes a novel and practical method for Conformal Drug property prediction, dubbed as CoDrug, to improve coverage in the conformal prediction framework under covariate shift. We address the problem of non-exchangeability by quantifying the underlying covariate shift at test time and leverage recent advances in conformal prediction to obtain prediction sets. Further, we demonstrate applying CoDrug to obtain valid uncertainty estimates w.r.t. a target property on molecules sampled from de novo drug design models. We summarize our main contributions below:
- We propose a novel approach to create prediction sets for drug property prediction, dubbed as CoDrug. Using kernel density estimates (KDE) and recent advances in CP, CoDrug corrects for covariate shift at test time and creates prediction sets whose coverage rate is closer to the target.
- We show that the kernel density estimates are consistent, which means that asymptotically, the covariate shift is precisely adjusted for, and the coverage guarantee is recovered.
- We demonstrate the loss of coverage in property prediction tasks induced by two forms of distribution shifts - molecular scaffold splitting and molecular fingerprint splitting. Our experiments show that CoDrug effectively reduces the gap between actual and target coverage for prediction sets, with an average enhancement of up to  $35\%$  compared to the conformal prediction method without covariate shift adjustment. Additionally, in our experiments on molecules generated by de novo drug design models, we observe a  $60\%$  reduction in the coverage gap on average.
Recently, deep learning techniques have been extensively studied for their potential in drug discovery, specifically in developing accurate predictive and generative models. This led to various architectures for predicting drug properties from SMILES/SELFIES strings [12], molecular graph representations[2, 1] and self-supervised learning [13]. Another area of research focused on building generative models to discover novel molecules using variational autoencoder [5, 6] and reinforcement learning [3, 4, 14].
Furthermore, several methods have been proposed for addressing uncertainty quantification in molecule property prediction, utilizing various Bayesian techniques [15]. Recently, conformal prediction methods have gained increasing attention for drug property prediction [8, 9, 10, 16, 17]. However, these studies primarily focus on generating efficient conformal predictors, without taking into account distribution shifts. Although several benchmarking datasets [18, 19] and methods [20] have been developed for drug property prediction under distribution shift, the problem of uncertainty quantification under distribution shift is still open.
Recent advancement in conformal prediction recovers the coverage guarantee for conformal prediction under known covariate shift [11]. [21] built upon [11] and proposed the Feedback Covariate Shift
<sup>1</sup> Prediction intervals can be viewed as prediction sets, with each interval being a subset of  $\mathbb{R}$ .
(FCS) method for the task of protein design. In practice, one cannot know the exact densities to measure the covariate shift. Like [21], we also leverage [11], but a key difference is that the training density is well-defined in [21] but unknown in ours, requiring us to estimate it. Additionally, our focus diverges from [21] as we concentrate on molecule property prediction rather than protein design.
Reliable estimation of drug properties is crucial for identifying potential drug candidates. Many essential drug properties, such as toxicity, efficacy, drug-drug interactions etc. are formulated as classification problems. Consider a classification task, with each data point  $Z = (X,Y) \in \mathbb{R}^d \times [K]$  ( $[K] = \{0,1,2,\dots,K - 1\}$ ). For instance, in Fig. 1(a), we seek to construct prediction sets for the problem of solubility classification. (Note that in practice, most drug discovery tasks are formulated as binary classification problems, with  $K = 2$ , but we present the general form of the methodologies.) While building an accurate base classifier ( $f$ ) is important, we usually would like more than a point estimate of the solubility of the molecule, but also some "confidence level". This could be encoded in the form of a prediction set denoted as  $\hat{C}(X) \subseteq [K]$ .
The main goal we seek in such prediction sets is valid coverage: Given a target (e.g.  $90\%$ ), we would like to construct a set-valued prediction (Fig. 1(a)) such that, if a molecule is water soluble, this prediction set will include the label "water soluble" with at least  $90\%$  probability. Formally, given  $1 - \alpha \in (0,1)$ , and a new test molecule  $(X_{N + 1},Y_{N + 1})$ , we would like  $\hat{C}$  to be  $1 - \alpha$  valid:
$$
\mathbb {P} \left\{Y _ {N + 1} \in \hat {C} \left(X _ {N + 1}\right) \right\} \geq 1 - \alpha . \tag {1}
$$
Conformal Prediction(CP) framework enables us to achieve such validity in Eq. (1). We will expand the details in Section 4.2. Remarkably, the only requirement of CP is a hold-out calibration set where the base classifier  $f$  is not trained on<sup>2</sup>.
One critical assumption for typical CP methods is that the test and calibration data are i.i.d (or exchangeable) which is rarely realistic in drug discovery tasks. On the other hand, although the distribution of molecules  $X$  changes from calibration to test time, the conditional distribution  $Y|X$  is unlikely to change as the molecular properties are determined by nature and remain the same under similar experimental conditions. Formally, if we denote our calibration set as  $\{(X_i,Y_i)\}_{i=1}^N$  and the test point as  $(X_{N+1},Y_{N+1})$ , we have:
$$
\forall i \in [ N ], \left(X _ {i}, Y _ {i}\right) \stackrel {i. i. d} {\sim} P ^ {c a l} = P _ {X} ^ {c a l} \times P _ {Y | X} ^ {c a l} \tag {2}
$$
$$
\left(X _ {N + 1}, Y _ {N + 1}\right) \sim P ^ {\text {t e s t}} = P _ {X} ^ {\text {t e s t}} \times P _ {Y | X} ^ {\text {c a l}}. \tag {3}
$$
It is important to note that the test distribution  $P^{test}$  maintains the same conditional distribution  $P_{Y|X}^{cal}$  as the calibration distribution, a phenomenon known as covariate shift. This shift is prevalent in de novo drug design models, which require navigating a vast chemical space to pinpoint optimal molecules for a specific goal. However, in many drug discovery tasks, the datasets typically contain only a few thousand data points, representing a limited chemical space. Thus, when models trained on these smaller datasets are used on molecules drawn from the broader molecular space, they inevitably encounter covariate shift. Next we will lay out the exact details of constructing prediction sets with the presence of covariate shifts for supporting drug discovery applications.
4 CoDrug Method
In the subsequent subsections, we describe the three primary components of CoDrug. In Section 4.2, we first provide a brief overview of inductive conformal prediction, presenting a method for constructing valid prediction sets in scenarios both without and with distribution shifts, presuming oracle access to the unknown distributions  $P_X^{test}$  and  $P_X^{cal}$ . Next, in Section 4.3, we present the details of the training aspects of the base energy-based classifier, emphasizing additional regularization using unlabeled data to enhance its capability to model varying molecule distributions. Finally, in Section 4.4, we employ kernel density estimation (KDE) on the embeddings or logits of the energy model trained in Section 4.3 to estimate the unknown distributions  $P_X^{test}$  and  $P_X^{cal}$ , and rectify covariate shift using
The assumption indicates that since the model  $f$  was not trained on the calibration set, whatever over-fitting happens on  $P_{Y|X}^{train}$ , but not  $P_{Y|X}^{cal}$ . Roughly speaking, we assume the classifier's performance on the calibration set is similar to that on an unseen test set.
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  $(1 - \alpha)$  quantile serves as the threshold for the conformal prediction set. (c),(d),(e) describe the CoDrug pipeline. (c) Training an 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_{i}$  are computed from the densities in (d). Then, Quantile is computed in a weighted fashion. Note how the quantile at  $\alpha = 0.3$  is shifted from 0.64 (in (b)) to 0.71 to account for the distribution shift.
Section 4.2. As KDE is consistent, we regain the coverage guarantee asymptotically. Together, these elements constitute the pipeline depicted in Fig. 1.
4.2 Conformal Prediction Set
In this section, we will explain how to use conformal prediction to construct valid prediction sets. We will start with the case without covariate shift, and then explain how to correct for covariate shift.
4.2.1 Conformal Prediction without Covariate Shift
Conformal prediction, pioneered by [7], is a powerful framework to construct prediction sets with the guarantee in Eq. (1). We assume that a base classifier  $f$  is trained on a training set  $\mathcal{D}_{train}$ , and we have a hold-out calibration set  $\mathcal{D}_{cal}$ . To simplify notation, we will denote the calibration set as  $\{Z_i\}_{i=1}^N$  and the test point of interest as  $Z_{N+1}$ . We will also abuse the notation to use  $\mathcal{D}$  to denote both the empirical calibration/test set as well as the underlying distribution. Note that we ignored the training samples because they are no longer used after the classifier  $f$  is trained.
We first introduce some useful definitions: (empirical) CDF and quantile function.
The cumulative distribution function (CDF)  $F$  of a set of values  $\{v_{i}\}_{i = 1}^{N}$  is defined as:
$$
F _ {\{v _ {i} \} _ {i = 1} ^ {N}} := 1 / N \sum_ {i = 1} ^ {N} \delta_ {v _ {i}}, \text {w h e r e} \delta_ {v} (x) := \mathbb {1} \{x \geq v \} \tag {4}
$$
The quantile function with respect to a CDF  $F$  is:
$$
\operatorname {Q u a n t i l e} (\beta ; F) := \inf  \{x: F (x) \geq \beta \} \tag {5}
$$
Given a target coverage level  $1 - \alpha \in (0,1)$ , the Mondrian inductive conformal prediction set (Mondrian ICP) is given by:
$$
\hat {C} \left(X _ {N + 1}\right) := \left\{y: 1 - p _ {y} ^ {f} \left(X _ {N + 1}\right) \leq t \right\} \tag {6}
$$
$$
\text {w h e r e} t := \operatorname {Q u a n t i l e} \left(1 - \alpha ; F _ {\{1 - p _ {Y _ {i}} ^ {f} (X _ {i}) \} \cup \{\infty \}}\right). \tag {7}
$$
where,  $p_{y}^{f}(x)$  corresponds to the softmax output of class  $y$  from model  $f$ . Here,  $\{v_{i}\}_{i=1}^{N}$  are defined by  $v(x_{i},y_{i}) = 1 - p_{y_{i}}^{f}(x_{i})$ , which are called "nonconformity scores" [22] and measure how "anomalous"
a point  $z = (x, y)$  is with respect to other points from this distribution. Intuitively, we assign to each molecule a score using the same rule using  $f$ , which is trained on a separate data split  $\mathcal{D}_{train}$ . Now, we choose a threshold  $t$  that is larger than  $1 - \alpha$  (e.g.,  $90\%$ ) of the molecules. Because of our i.i.d. assumption, if we sample another molecule  $Z_{N+1}$  from the same distribution, we expect its score to be lower than this threshold with a probability of  $1 - \alpha$ . For eg, in Fig. 1(b), notice how the threshold  $t$ , is computed as the value of the Quantile function at  $\alpha = 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_i,Y_i)\}_{i = 1}^{N + 1}$ . The  $\hat{C}$  in Eq. (6) satisfies:
$$
\mathbb {P} \left\{Y _ {N + 1} \in \hat {C} \left(X _ {N + 1}\right) \right\} \geq 1 - \alpha . \tag {8}
$$
Remarks: Theorem 4.1 is a result of classical Mondrian inductive conformal prediction [7]. In fact, in the classification setting, instead of the i.i.d. assumption, one could make a slightly milder assumption that data are exchangeable within each class.
4.2.2 Conformal Prediction with Covariate Shift
While Theorem 4.1 provides a nice first step, the i.i.d. assumption poses a significant limitation in drug discovery. As mentioned in Eq. (3), the distributions of  $X$  on  $\mathcal{D}^{test}$  and  $\mathcal{D}^{cal}$  can differ. In Eq. (7), we used the empirical CDF  $F$  (Eq. (4)) to choose the threshold  $t$ . This is because of our i.i.d. assumption: a particular molecule type appears with equal probability/density in both the calibration and test sets. This is no longer the case with covariate shift, which means our  $F$  needs to account for such difference in  $P_{X}$ .
Formally, recall that  $P_X^{cal}$  and  $P_X^{test}$  represent the density of the molecule  $X$  for the calibration and test sets. We will assign a weight to each molecule  $x$  that is proportional to the density/likelihood ratio  $\frac{dP_X^{test}}{dP_X^{cal}}$  in the empirical CDF, leading to:
$$
F _ {x _ {N + 1}} ^ {w} := w \left(x _ {N + 1}\right) \delta_ {\infty} + \sum_ {i \in [ N ]} w \left(x _ {i}\right) \delta_ {1 - p y _ {i} \left(x _ {i}\right)} / W \tag {9}
$$
$$
w \left(x ^ {\prime}\right) := d P _ {X} ^ {t e s t} \left(x ^ {\prime}\right) / d P _ {X} ^ {c a l} \left(x ^ {\prime}\right), \forall x ^ {\prime} \tag {10}
$$
$W = \sum_{i=1}^{N+1} w(x_i)$  is just a normalizing factor. The subscript  $x_{N+1}$  is used to highlight that our updated CDF now depends on the test molecule  $x_{N+1}$  through the weights. Here,  $w(x')$  could be viewed as a likelihood ratio, and is crucial in adjusting for the covariate shift. For eg, in Fig. 1(e). Notice how the values of  $w(x_i)$  depend on the densities  $P_X^{cal}$  and  $P_X^{test}$ . In the figure, the value of weighted Quantile at  $\alpha = 0.3$  or the threshold  $t$  is shifted from 0.64 to 0.71 to account for the shift. We formally state the modified theorem from [11] that recovers the coverage guarantee under covariate shift for Mondrian ICP:
Theorem 4.2. [11] Assume that  $\tilde{P}_X$  is absolutely continuous with respect to  $P_X$ . For any  $\alpha \in (0,1)$ , let  $F^w$  be defined as in Eq. (9), and
$$
\hat {C} (x) = \left\{y: 1 - p _ {y} ^ {f} (x) \leq \text {Q u a n t i l e} \left(1 - \alpha ; F _ {x} ^ {w}\right) \right\} \tag {11}
$$
$$
\mathbb {P} \left\{Y _ {N + 1} \in \hat {C} \left(X _ {N + 1}\right) \right\} \geq 1 - \alpha . \tag {12}
$$
However, in practice, both  $P_X^{test}$  and  $P_X^{cal}$  are unknown, rendering Theorem 4.2 impractical. In section 4.4, we will provide a viable way to estimate  $P_X^{test}$ , and recover the guarantee asymptotically (namely with large calibration and test sets). In the following section, we will delve into our training methodology, which harnesses unlabeled data to effectively model varying molecular distributions.
4.3 CoDrug Training Methodology
CoDrug handles distribution shift by proposing an energy-based model formulation [23]. The core idea behind an energy-based model is to construct a function  $E$  that maps an input  $x$  to a scalar value, known as energy. A collection of energy values can be transformed into a probability density function  $p(x)$  through the Gibbs distribution
$$
p (y | x) = e ^ {- E (x, y) / T} / e ^ {- E (x) / T} \tag {13}
$$
Consider a discriminative neural network  $f$  used in a  $K$  class classification setting.  $f(x)$  maps an input  $x$  into  $K$  real-valued scalars, which are used to derive a conditional class-wise probability:
$$
p (y | x) = e ^ {f _ {y} (x) / T} / \sum_ {y ^ {\prime}} ^ {K} e ^ {f _ {y ^ {\prime}} (x) / T} \tag {14}
$$
where,  $f_{y}(x)$  refers to the  $y^{th}$  logit of the classifier  $f(x)$ . In this setting, the energy function  $E(x)$  can be expressed in terms of the denominator of the softmax probabilities in Eq. (14).
$$
E (x; f) = - T \cdot \log \sum_ {y ^ {\prime}} ^ {K} e ^ {f _ {y ^ {\prime}} (x) / T} \tag {15}
$$
Directly using embeddings from a model trained on labeled data may not yield reliable density estimates, as the model lacks knowledge of data outside the training distribution. To overcome this, we co-train the model with unlabeled molecule data. This aids the model  $f$  in effectively mapping molecules with distribution shifts to a distinct embedding space. We follow [24], and use an extra regularization term in the loss function to ensure energy separation between in-distribution and out-of-distribution data. The objective function is defined as follows:
$$
\min  _ {\theta} \mathbb {E} _ {(x, y) \sim \mathcal {D} _ {i n}} [ - \log (p _ {y} ^ {f} (x)) ] + \lambda \cdot L _ {\text {e n e r g y}} \tag {16}
$$
where  $p_y^f(x)$  refers to the softmax outputs of the classification model  $f$  for class  $y$ , and  $\mathcal{D}_{in}$  is the in-distribution training data for which labels are available. The training objective is combined with an additional term  $L_{energy}$  given by:
$$
L _ {\text {e n e r g y}} = \mathbb {E} _ {\left(x _ {i n}, y\right) \sim \mathcal {D} _ {i n}} \left(\max  \left(0, \left(E \left(x _ {i n}\right) - m _ {i n}\right)\right) ^ {2}\right) + \mathbb {E} _ {\left(x _ {o u t}, y\right) \sim \mathcal {D} _ {o u t}} \left(\max  \left(0, \left(m _ {o u t} - E \left(x _ {o u t}\right)\right)\right) ^ {2}\right) \tag {17}
$$
where  $\mathcal{D}_{out}$  refers to the subset of the unlabelled that is out-of-distribution (OOD). The objective of this term is to enforce a margin of separation between the training samples and the OOD data using the hyper-parameters  $m_{in}, m_{out}$ . Particularly, one term penalizes the model if the energy values for in-distribution data are higher than a certain value, while the other term penalizes if the OOD samples have an energy lower than a certain value. In the next section, we will explain how to use either the latent embedding of  $f$  or the logits to estimate the density and correct for covariate shift.
4.4 Density Estimation
As discussed earlier, we need to estimate  $P_X^{test}$  and  $P_X^{cal}$  to correct for the covariate shift. We resort to Kernel Density Estimation (KDE), a classical nonparametric method, to estimate the density of arbitrary distributions of molecules. For a set of data  $X_1, \ldots, X_n \stackrel{i.i.d.}{\sim} \mathcal{D}$ , KDE is given by:
$$
\hat {p} _ {h} (x; \mathcal {D}) = (n h) ^ {- 1} \sum_ {i = 1} ^ {n} K \left(x - x _ {i} / h\right), \tag {18}
$$
where  $K$  is a fixed non-negative kernel function, and  $h > 0$  is a smoothing bandwidth. Such KDE estimates have nice asymptotic convergence properties, as formally stated in the following theorem:
Theorem 4.3. [25] Assume that the true density  $p$  is square-integrable and twice differentiable and that its second-order partial derivatives are bounded, continuous, and square-integrable. If  $K$  is spherically symmetric on  $\mathbb{R}^d$ , with a finite second moment, and we choose the bandwidth  $h$  such that
$$
\lim  _ {m \rightarrow \infty} h ^ {d} m \rightarrow \infty a n d \lim  _ {m \rightarrow \infty} h \rightarrow 0 \tag {19}
$$
then as  $m\to \infty$
$$
\left\| \hat {p} _ {h} (x) - p (x) \right\| _ {2} \stackrel {P} {\rightarrow} 0, \tag {20}
$$
where  $\stackrel{P}{\to}$  means convergence in probability.
Note that commonly used kernels, such as Gaussian kernel, satisfy the requirements.