SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Bowen Gao $^{1*}$ , Bo Qiang $^{2*}$ , Haichuan Tan $^{1}$ , Minsi Ren $^{3}$ , Yinjun Jia $^{4}$ , Minsi Lu $^{5}$ , Jingjing Liu $^{1}$ , Wei-Ying Ma $^{1}$ , Yanyan Lan $^{1,6\dagger}$
$^{1}$ Institute for AI Industry Research (AIR), Tsinghua University
$^{2}$ Department of Pharmaceutical Science, Peking University
$^{3}$ Institute of Automation, Chinese Academy of Sciences
$^{4}$ School of Life Sciences, Tsinghua University
$^{5}$ Department of Pharmaceutical Science, Tsinghua University
$^{6}$ Beijing Academy of Artificial Intelligence
Virtual screening, which identifies potential drugs from vast compound databases to bind with a particular protein pocket, is a critical step in AI-assisted drug discovery. Traditional docking methods are highly time-consuming, and can only work with a restricted search library in real-life applications. Recent supervised learning approaches using scoring functions for binding-affinity prediction, although promising, have not yet surpassed docking methods due to their strong dependency on limited data with reliable binding-affinity labels. In this paper, we propose a novel contrastive learning framework, DrugCLIP, by reformulating virtual screening as a dense retrieval task and employing contrastive learning to align representations of binding protein pockets and molecules from a large quantity of pairwise data without explicit binding-affinity scores. We also introduce a biological-knowledge inspired data augmentation strategy to learn better protein-molecule representations. Extensive experiments show that DrugCLIP significantly outperforms traditional docking and supervised learning methods on diverse virtual screening benchmarks with highly reduced computation time, especially in zero-shot setting. The code for DrugCLIP is available at https://github.com/bowen-gao/DrugCLIP.
Virtual screening is a crucial computer-aided drug discovery (CADD) technique that uses computational methods [26] to search for candidate drug structures from compound libraries, aiming to identify molecules most likely to bind to a target (e.g., a protein receptor or enzyme). Common practice follows the wisdom of "bigger is better" [25]: the larger the search library, the better chance to find a matching drug candidate. For example, statistics show that increasing the data size from  $10^{5}$  to  $10^{8}$  leads to a significant jump in the number of true ligands among top-1000 results [19].
Molecular docking [11, 43, 41] is currently the dominant virtual screening method, which models protein-molecule interaction by a quantitative score correlated with the free energy of binding. The computation of such scores heavily depends on molecule orientations and conformations sampling, which is time-consuming and impractical when dealing with large libraries in the magnitude of billions [14]. As estimated by [32], given a standard computing rate of 10 seconds per compound on
†Correspondence to lanyanyan@air.tsinghua.edu.cn
a single CPU core, it takes 3000 years to complete screening 10 billion compounds using commercial docking methods (costing over 800k dollars). Consequently, docking methods are gravely limited by their high computational cost and slow inference speed.
As molecule libraries continue to grow, new high-throughput virtual screening methods are in pressing demand. Supervised learning algorithms [20] such as regression [22] and classification [54] have been investigated for binding-affinity prediction. However, these supervised methods usually require carefully labeled data samples for training, thus struggling with poor generalization [44]. In addition, the scarcity of reliable negative samples imposes restrictions on model performance. As a result, they still underperform current docking methods.
Rethinking the problem of virtual screening, we find that the key issue is to identify which molecules are likely to bind with a protein pocket, instead of determining the accurate binding-affinity score (the goal of prediction in regression/classification models) or binding pose (the goal of docking). Following this thought, we recast virtual screening as an information retrieval task, i.e., given a protein pocket as the query, we aim to retrieve from a large-scale molecule library the most relevant molecules with the highest probability of binding to the target pocket. In this new perspective, virtual screening is boiled down to a similarity matching problem between proteins and molecules.
To this end, we introduce DrugCLIP, a dense retrieval approach (inspired by CLIP [31]) that computes a contrastive loss between two separate pre-trained encoders to maximize the similarity between a protein-molecule pair, if they have a binding affinity, and minimize it otherwise. Compared with supervised learning methods, our contrastive learning approach enjoys several advantages. Firstly, the objective of finding the matching relations between proteins and molecules is naturally in accordance with the formulation of the virtual screening task. Secondly, the designed contrastive loss relieve the dependency on explicit labeling of binding affinity, and facilitates the usage of large-scale unlabeled data beyond densely annotated small datasets such as PDBBind [45]. Thus, we extend the ability of our model with a large pool of protein-molecule pairs by utilizing BioLip [47] and ChEMBL [8] datasets for training. We further introduce a biological-knowledge inspired augmentation method, HomoAug, which creates protein-molecule pairs based on protein homology evolutions. Lastly, the dense retrieval setting allows for offline pre-computation of protein and molecule encodings, bringing high efficiency to online inference and promising high-throughput virtual screening on billions of molecules.
Experiments on two challenging virtual screening benchmarks, DUD-E [28] and LIT-PCBA [40], demonstrate that zero-shot performance of our model surpasses most deep learning baselines that carefully finetune on labeled data. We also conduct a human evaluation to compare DrugCLIP with Glide, a commercial docking system widely used by pharmacology experts. In  $80\%$  cases, judges prefer the selection of top-10 molecules from our method over Glide. Furthermore, since DrugCLIP indiscriminately models the similarity between protein and molecule, it can be extended to other important tasks in drug discovery such as 'target fishing', where protein candidates are ranked for a given molecule. DrugCLIP also outperforms docking methods on these benchmarks.
- To our best knowledge, this is the first effort to position large-scale virtual screening as a dense retrieval problem, which enables ultra-fast screening over billion-scale chemical libraries for candidate search by storing pre-computed molecule embeddings offline.
- We propose a novel contrastive learning framework that learns a generic joint representation of proteins and molecules, which can be applied to molecule-pocket pairing tasks. Novel data augmentation strategy and training techniques are also introduced.
- DrugCLIP, with its impressive zero-shot performance on virtual screening benchmarks, well addresses poor generalization and low efficiency issues faced by docking and learning-based screening methods.
There are mainly two schools of virtual screening methods, molecular docking and supervised learning. Molecular docking is a computational technique that predicts the binding energy, optimal orientation, and conformation of a small molecule ligand within a protein binding site [19]. It uses sampling algorithms such as genetic algorithms [41] and Monte Carlo [11, 43] to generate a set of
candidate ligand poses, by exploring the conformational space of the ligand and the protein receptor. These candidate poses are then evaluated by molecule-protein scoring functions such as empirical force fields [21] to assess their binding affinity. This iterative process continues until convergence, which is computationally demanding.
To accelerate the prediction process, supervised learning methods have emerged as an alternative to the iterative refining process in docking. A recent work [13] proposes pocket pretraining to find ligands for similar pockets by only utilizing information from one side of the pocket-ligand pairs. By training on given binding-affinity labels, regression models such as DeepDTA [30], OnionNet [52], GraphDTA [29] and SG-CNN [16] learn the mapping between protein-molecule representations by first predicting the binding affinity for every protein-molecule pair, then ranking them to determine top candidates. However, these models suffer from high false-positive rates, due to the lack of negative binding-affinity data. Another way is to use predefined rules (e.g., DUD-E [28]) to obtain negative samples and train a classifier to discriminate positive and negative protein-molecule pairs (DrugVQA [54], AttentionSiteDTI [24]). Previous work has observed poor generalization in these methods. For example, Wang and Dokholyan [44] suggests that models trained on DUD-E [28] cannot be transferred to other classification benchmarks.
Concurrently, several studies have employed contrastive learning objectives in the context of virtual screening. Singh et al. [36] utilized a protein language model and rule-based molecule fingerprints for representation, while CoSP [7] opted for chemical similarity in their negative sampling approach. In the meantime, we propose a more holistic and flexible solution by leveraging an equivariant 3D model, efficient dense retrieval, and novel data augmentation techniques.
3 DrugCLIP Framework
To formulate the problem, we denote the protein pocket of interest as  $p$ , and a set of  $n$  small molecules is represented by  $\mathcal{M} = \{m_1, m_2, \dots, m_n\}$ . The objective of virtual screening is to identify the top  $k$  candidates with the highest probability of binding to the target pocket. This selection process is typically guided by a scoring function  $s(\cdot, \cdot)$ , which assesses the pairwise data between the pocket  $p$  and each candidate molecule  $m_i$ . The scoring function can be derived from techniques such as docking simulations or supervised learning models, which perform ranking and selection of the most promising candidates based on their likelihood of binding to the target pocket.
To view virtual screening as a dense retrieval task, we treat the pocket as the query to retrieve relevant molecules from the given library. The overall framework is illustrated in Figure 1. First, two separate encoders are trained to learn the representations of protein pockets (abbreviated as proteins hereafter) and molecules. Then, the similarity between each protein-molecule pair is computed and a contrastive learning objective is utilized to discriminate between positive and negative pairs. All the parameters in the encoders and similarity functions are trained jointly.
3.2 Protein and Molecule Encoders
Diverse representation learning methods can be used as protein and molecule encoders. In this paper, we follow the encoder architecture of UniMol [55], a powerful 3D encoder pre-trained with large-scale unsupervised data. Here we briefly introduce the encoding process.
Firstly, both molecules and protein pockets are tokenized to atoms. A molecule with  $L$  tokens is denoted as a feature vector  $x^{m} = \{c_{m}, t_{m}\}$ , where  $c_{m} \in \mathbb{R}^{L \times 3}$  represents the atom coordinates and  $t_{m} \in \mathbb{R}^{L}$  represents the atom types. The same setting is applied to obtain pocket features, denoted as  $x^{p} = \{c_{p}, t_{p}\}$ .
As described in UniMol [55], the encoder is a SE(3) 3D transformer that accepts tokenized atom features as input. To preserve SE(3) invariance for embedding the molecular structure, the 3D coordinate features are utilized as geometric distances. Specifically, the pairwise representation  $q_{ij}^{0}$  is initialized based on the distance between each pair of atoms. For each transformer layer  $l$ , the self-attention mechanism for learning atom representation is defined in Equation 1. The pairwise representation serves as a bias term in the attention mechanism, encoding 3D features into atom representations. The update rules between adjacent transformer layers are also defined in Equation 1.
Figure 1: An illustration of the training procedure. Molecule conformations are generated by RDkit simulation, and pocket data are augmented with HomoAug. At each training iteration, sampled 3D molecule and 3D pocket representations are learned with a contrastive objective.
$$
\operatorname {A t t e n t i o n} \left(Q _ {i} ^ {l}, K _ {j} ^ {l}, V _ {j} ^ {l}\right) = \operatorname {s o f t m a x} \left(\frac {Q _ {i} ^ {l} \left(K _ {j} ^ {l}\right) ^ {T}}{\sqrt {d}} + q _ {i j} ^ {l}\right) V _ {j} ^ {l}, \quad \text {w h e r e} q _ {i j} ^ {l + 1} = q _ {i j} ^ {l} + \frac {Q _ {i} ^ {l} \left(K _ {j} ^ {l}\right) ^ {T}}{\sqrt {d}}. \tag {1}
$$
Inspired by BERT[5], we randomly mask atom types and pretrain our model by predicting the masked atom types. Besides, we introduce random noise to corrupt atom coordinates and pretrain the model to reconstruct the original coordinates. Specifically, uniform noises of [-1, 1] are added to  $15\%$  of atom coordinates, where the pair-distance prediction heads estimate uncorrupted distances and the SE(3)-equivariant head directly predicts correct coordinates. Detailed implementations are described in the Appendix B. A special atom [CLS] with coordinates at the center of all atoms is added to output the representation of the corresponding protein and molecule. Specifically, we denote the protein encoder as  $g_{\phi}$  and molecule encoder as  $f_{\theta}$ . Then, the representation of protein  $x^{p}$  and molecule  $x^{m}$  are defined as  $g_{\phi}(x^{p})$  and  $f_{\theta}(x^{m})$ , correspondingly.
3.3 Contrastive Learning Objective
To conduct the contrastive learning process, we first need to obtain the similarity measurements between each protein-molecule pair. Following previous work, both dot product and cosine similarity can be adopted as the similarity functions. For example, when using dot product, the similarity score of  $(x_i^p,x_j^m),\forall i,j\in [1,N]$  can be written as:
$$
s \left(x _ {i} ^ {p}, x _ {j} ^ {m}\right) = g _ {\phi} \left(x _ {i} ^ {p}\right) ^ {T} \cdot f _ {\theta} \left(x _ {j} ^ {m}\right), \tag {2}
$$
while cosine similarity can be obtained by normalization.
In the field of virtual screening, positive pairs of binding protein and molecule are usually provided, with limited true negative pairs. Therefore, we need to construct negative pairs for the contrastive objective. Here we use an in-batch sampling strategy similar to CLIP [31]. Specifically, given a batch of paired data  $\{(x_k^p,x_k^m)\}_{k = 1}^N$  with batch size  $N$ , we extract a list of proteins  $\{x_k^p\}_{k = 1}^N$  and a list of corresponding molecules  $\{x_k^m\}_{k = 1}^N$ . Combining them together results in  $N^2$  pairs  $(x_i^p,x_j^m)$  where  $i,j\in [1,N]$ . When  $i = j$  it is a positive pair, and when  $i\neq j$  it is a negative pair.
Please note that the in-batch negative construction is intrinsically based on a simple assumption that, if a certain pair of protein and molecule has been tested as having a binding relation, it is likely that they have a negative binding relation with other molecule/protein. This assumption is reasonable as the true distribution of positive and negative molecules exhibits a sharp contrast, with a proportion significantly smaller than  $0.1\%$  [25].
Formally, we introduce two losses: Pocket-to-Mol loss and Mol-to-Pocket loss. The former describes the likelihood of ranking its binding molecules before other molecules for a given protein  $x_{k}^{p}$ :
$$
\mathcal {L} _ {k} ^ {p} \left(x _ {k} ^ {p}, \left\{x _ {i} ^ {m} \right\} _ {i = 1} ^ {N}\right) = - \frac {1}{N} \log \frac {\exp \left(s \left(x _ {k} ^ {p} , x _ {k} ^ {m}\right) / \tau\right)}{\sum_ {i} \exp \left(s \left(x _ {k} ^ {p} , x _ {i} ^ {m}\right) / \tau\right)}, \tag {3}
$$
while the latter is the likelihood of ranking its binding targets for a given molecule  $x_{k}^{m}$ , and is defined as:
$$
\mathcal {L} _ {k} ^ {m} \left(x _ {k} ^ {m}, \left\{x _ {k} ^ {p} \right\} _ {i = 1} ^ {N}\right) = - \frac {1}{N} \log \frac {\exp \left(s \left(x _ {k} ^ {p} , x _ {k} ^ {m}\right) / \tau\right)}{\sum_ {i} \exp \left(s \left(x _ {i} ^ {p} , x _ {k} ^ {m}\right) / \tau\right)}. \tag {4}
$$
In the above two equations,  $\tau$  represents the temperature parameter that controls the softmax distribution, which has been widely utilized in previous representation learning methods[3, 10, 46, 31].
Combining the two losses, we obtain the final loss for a mini-batch:
$$
\mathcal {L} = \frac {1}{2} \sum_ {k = 1} ^ {N} \left(\mathcal {L} _ {k} ^ {p} + \mathcal {L} _ {k} ^ {m}\right). \tag {5}
$$
Virtual screening with DrugCLIP contains two phases. In the offline phase, embeddings of each molecule are obtained by DrugCLIP encoders  $f_{\theta}$ . These embedding vectors are then stored in memory for later-stage online retrieval. Specifically, for a given query protein pocket, it is first encoded into an embedding vector using the trained protein encoder  $g_{\phi}$ . We then measure the similarity between the encoded pocket vector and all the embedding vectors of candidate molecules (dot product or cosine similarity). Finally, we proceed to select the top- $k$  molecules from the candidate pool based on their similarity scores. Notably, our method offers a distinct advantage compared to other supervised learning screening frameworks. While other methods involve complex neural network computations for the scoring functions during the online screening phase, our approach capitalizes on the precomputed and cached candidate embedding vectors. Consequently, the only computation required is the high-speed dot product calculation. This novel design allows for rapid screening of a large number of candidates, without incurring additional computational overhead. For a detailed time analysis, please refer to Section 4.3.
Constructing Training Data
We use three datasets for training: PDBBind [45], BioLip [47], and ChEMBL [8]. PDBbind is a standard database used in docking and binding-affinity prediction. It consists of experimentally measured protein-ligand complex structures along with their binding-affinity labels, from which true positive protein-molecule pairs with accurate structures can be extracted. We use PDBBind 2019, which includes over 17,000 protein-molecule complexes with binding-affinity data covering a wide range of chemical space and protein families. We use the general set for training and the refined set for validation.
BioLip is a dataset updated weekly by a standard data mining workflow that extracts complex structure data from PDB. We filter out all complexes that contain peptides, DNA, RNA, and single ions, and obtain 122861 protein-molecule pairs, much larger than PDBBind.
DrugCLIP model can also use known receptor-ligand pairs without their binding structures. From the ChEMBL [8] dataset, we filter out proteins with only one known binding pocket. Then we pair the pocket with all positive binders in the ChEMBL database. We hypothesize (supported by domain experts [9, 18]) that assayed ligands dominantly bind to the known pocket in solved structures, and our model can tolerate introduced noise of this filtration protocol.
Since our deep-learning-based method is insensitive to occasional inaccuracies in coordinates and can be trained with raw element types. Therefore, only minimal cleaning-ups are performed for the protein structures in the above mentioned datasets to remove irrelevant molecules like water.
Biological Data Augmentation
Directly applying common data augmentation techniques to augment biological data is infeasible, as introducing noise or perturbations to pocket or molecule data can result in unstable or chemically incorrect structures, rendering the augmented data unreliable, especially for virtual screening [34]. To address this challenge, we propose a new augmentation method called HomoAug that takes into account the biological significance of the data. It utilizes the concept of homologous proteins in biology, to combine ligands from PDBbind [45] with homologous proteins corresponding to their pockets, thereby generating new training data (Figure 2).
Figure 2: Illustration of HomoAug's pipeline. The pocket protein instances from PDBBind are searched for homologous counterparts in the AlphaFold Protein Structure Database. Then the TMalign method is employed to achieve structural alignment between the homologous protein and the original protein. Following a filtering process, the homologous protein is combined with the ligand to create an augmented pocket-ligand pair.
Training-test Inconsistency
There exists a training-test inconsistency problem in virtual screening. This is because holo structures in training data are depicted after binding, which differs from their apo structures. Previous methods usually use docking software to obtain estimated binding structure. For better efficiency, we propose to use the chemical simulation package [23], RDkit to generate noisy data as input. For a given protein  $x^{p} = \{c_{p},h_{p}\}$  and molecule  $x^{m} = \{c_{m},h_{m}\}$ , we denote  $c$  as coordinates and  $h$  as atom types. Inaccurate molecule coordinates  $c_{m} + \delta$  simulated by RDkit are used as noisy input  $\tilde{x}_{m} = \{c_{m} + \delta ,h_{m}\}$ . It should be noted that our model adopts a dual-tower architecture, while previous 3D supervised learning methods are mainly based on a single tower [52, 16]. The following Proposition proves that the scoring function of DrugCLIP is more robust than those obtained by supervised learning methods. In a more general setting, we suggest that dual-tower models based on 3D encoders enjoy more generalization by decoupling the relative 3D distances from the prediction. By learning molecule-protein interactions from noisy unbind molecule conformations, our model addresses the consistency issue between training and testing.
Proposition 1. Denote the scoring function of a supervised learning method as  $k_{\gamma}$ , we have,
$$
\lim  _ {\delta \rightarrow 0} \left\{s \left(\tilde {x} _ {m}, x _ {p}\right) - s \left(x _ {m}, x _ {p}\right)\right\} = 0, \tag {6}
$$
$$
\lim  _ {\delta \rightarrow 0} \left\{k _ {\gamma} \left(h _ {p}, h _ {m}, c _ {p}, \tilde {c} _ {m}\right) - k _ {\gamma} \left(h _ {p}, h _ {m}, c _ {p}, c _ {m}\right)\right\} \neq 0. \tag {7}
$$
The proof is provided in the Appendix A
We first introduce evaluation metrics. Although AUROC (area under the receiver operating characteristic curve) is commonly used for classification, it has been criticized for being unsuitable for virtual screening. [51] This is because the target of virtual screening is to select a small fraction of molecules from a large pool, resulting in a significantly low false positive rate (FPR) in this scenario. However, AUROC is calculated by averaging the FPR from 0 to 1. To overcome this, we also use BEDROC (Boltzmann-enhanced discrimination of ROC), Enrichment Factor(EF) and ROC enrichment metric (RE) for evaluation. BEDROC incorporates exponential weights that assign greater importance to early rankings. EF and RE are two widely used metrics for virtual screening (detailed definitions in Appendix B).
4.1 Evaluation on DUD-E Benchmark
DUD-E [28] is one of the most popular virtual screening benchmarks. It contains 102 proteins with 22,886 bio-active molecules, for which 50 topological dissimilar decoys that possess matched physicochemical properties are retrieved from the ZINC database.
Table 1: Results on DUD-E in zero-shot setting.
AUROC (%)	BEDROC (%)	EF
0.5%	1%	5%
Glide-SP [11]	76.70	40.70	19.39	16.18	7.23
Vina [41]	71.60	-	9.13	7.32	4.44
NN-score [6]	68.30	12.20	4.16	4.02	3.12
RFscore [1]	65.21	12.41	4.90	4.52	2.98
Pafnucy [38]	63.11	16.50	4.24	3.86	3.76
OnionNet [52]	59.71	8.62	2.84	2.84	2.20
Planet [49]	71.60	-	10.23	8.83	5.40
DrugCLIPZS	80.93	50.52	38.07	31.89	10.66