SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Qizhi Pei $^{1,5*}$ , Kaiyuan Gao $^{2*}$ , Lijun Wu $^{3\dagger}$ , Jinhua Zhu $^{4}$ , Yingce Xia $^{3}$ , Shufang Xie $^{1}$ , Tao Qin $^{3}$ , Kun He $^{2}$ , Tie-Yan Liu $^{3}$ , Rui Yan $^{1,6\dagger}$
<sup>1</sup> Gaoling School of Artificial Intelligence, Renmin University of China
$^{2}$  School of Computer Science and Technology, Huazhong University of Science and Technology  
 $^{3}$  Microsoft Research AI4Science
$^{4}$  School of Information Science and Technology, University of Science and Technology of China
<sup>5</sup> Engineering Research Center of Next-Generation Intelligent Search and Recommendation, Ministry of Education
$^{6}$  Beijing Key Laboratory of Big Data Management and Analysis Methods
{qizhipei, shufangxie, ruiyan}@ruc.edu.cn, {im_kai, brooklet60}@hust.edu.cn
{lijuwu,yinxia,taoqin,tyliu}@Microsoft.com,teslazhu@mail.ustc.edu.cn
Modeling the interaction between proteins and ligands and accurately predicting their binding structures is a critical yet challenging task in drug discovery. Recent advancements in deep learning have shown promise in addressing this challenge, with sampling-based and regression-based methods emerging as two prominent approaches. However, these methods have notable limitations. Sampling-based methods often suffer from low efficiency due to the need for generating multiple candidate structures for selection. On the other hand, regression-based methods offer fast predictions but may experience decreased accuracy. Additionally, the variation in protein sizes often requires external modules for selecting suitable binding pockets, further impacting efficiency. In this work, we propose FABind, an end-to-end model that combines pocket prediction and docking to achieve accurate and fast protein-ligand binding. FABind incorporates a unique ligand-informed pocket prediction module, which is also leveraged for docking pose estimation. The model further enhances the docking process by incrementally integrating the predicted pocket to optimize protein-ligand binding, reducing discrepancies between training and inference. Through extensive experiments on benchmark datasets, our proposed FABind demonstrates strong advantages in terms of effectiveness and efficiency compared to existing methods. Our code is available at  $\mathrm{Gib}\mathbf{u}\mathbf{b}^{3}$ .
Biomolecular interactions are vital in the human body as they perform various functions within organisms [44, 2]. For example, protein-ligand binding [35], protein-protein interaction [17], protein-DNA interaction [18], and so on. Among these, drug-like small molecules (ligands) binding to protein is important and widely studied in interactions as they can facilitate drug discovery. Therefore, molecular docking, which involves predicting the conformation of a ligand when it binds to a target protein, serves as a crucial task, the resulting docked protein-ligand complex can provide valuable insights for drug development.
*Equal contribution. This work was done during their internship at Microsoft Research AI4Science.
†Corresponding authors: Rui Yan (ruiyan@ruc.edu.cn) and Lijun Wu (lijuwu@microsoft.com).
<sup>3</sup>https://github.com/QizhiPei/FABind
Though important, fast and accurately predicting the docked ligand pose is super challenging. Two families of methods are commonly used for docking: sampling-based and regression-based prediction. Most of the traditional methods lie in the sampling-based approaches as they rely on physics-informed empirical energy functions to score and rank the enormous sampled conformations [8, 35, 43], even with the use of deep learning-based scoring functions for conformation evaluation [34, 49], these methods still need a large number of potential ligand poses for selection and optimization. DiffDock [7] utilizes a deep diffusion model that significantly improves accuracy. However, it still requires a large number of sampled/generated ligand poses for selection, resulting in high computational costs and slow docking speeds. The regression-based methods [9, 32, 11] that use deep learning models to predict the docked ligand pose bypass the dependency on the sampling process. For instance, TankBind [31] proposes a two-stage framework that simulates the docking process by predicting the protein-ligand distance matrix and then optimizing the pose. In contrast, EquiBind [41] and E3Bind [51] directly predict the docked pose coordinates. Though efficient, the accuracy of these methods falls behind the sampling-based methods. Additionally, the variation in protein sizes often requires the use of external modules to first select suitable binding pockets, which can impact efficiency. Many of these methods rely on external modules to detect favorable binding sites in proteins. For example, TankBind [31] and E3Bind [51] take P2Rank [23] as priors for generating the pocket center candidates, which results in the need for a separate module (e.g., affinity prediction in TankBind and confidence module in E3Bind) for pocket selection, increasing the training and inference complexity.
To address these limitations, we propose FABind, a Fast and Accurate docking framework in an end-to-end way. FABind unifies pocket prediction and docking, streamlining the process within a single model architecture, which consists of a series of equivariant layers with geometry-aware updates, allowing for either pocket prediction or docking with only different configurations. Notably, for pocket prediction, we utilize lightweight configurations to maintain efficiency without sacrificing accuracy. In contrast to conventional pocket prediction, which only uses protein as input and forecasts multiple potential pockets, our method incorporates a specific ligand to pinpoint the unique pocket that the ligand binds to. Integrating the ligand into pocket prediction is crucial as it aligns with the fundamental characterization of the docking problem. In this way, we achieve a quick (without the need for external modules such as P2Rank) and precise pocket prediction (see Section 5.1).
Several strategies are additionally proposed in FABind to make a fast and accurate docking prediction. (1) Our pocket prediction module operates as the first layer in our model hierarchy and is jointly trained with the subsequent docking module to ensure a seamless, end-to-end process for protein-ligand docking prediction. We also incorporate a pocket center constraint using Gumbel-Softmax. This way assigns a probabilistic weighting to the inclusion of amino acids in the pocket, which helps to identify the most probable pocket center and improve the precision of the docking prediction. (2) We incorporate the predicted pocket into the docking module training using a scheduled sampling approach [4]. This way ensures consistency between the training and inference stages with respect to pocket leverage, thereby avoiding any mismatch that may arise from using the native pocket during training and the predicted pocket during inference. In this way, FABind is trained on a variety of possible pockets, allowing it to generalize well to new docking scenarios. (3) While directly predicting the ligand pose [41] and optimizing the coordinates based on the protein-ligand distance map [31] are both widely adopted in the sampling-based methods to ensure efficiency, we integrate both predictions in FABind to produce a more accurate pose prediction.
To evaluate the performance of our method, we conducted experiments on the binding structure prediction benchmark and compared our work with multiple existing methods. Our results demonstrate that FABind outperforms existing methods, achieving a mean ligand RMSD of 6.4. This is a significant improvement over previous methods and demonstrates the effectiveness of our approach. Remarkably, our approach also demonstrates superior generalization ability, performing surprisingly well on unseen proteins. This suggests that our model can be applied to a wide range of docking scenarios and has the potential to be a useful tool in drug discovery. In addition to achieving superior performance, our method is also much more efficient during inference (e.g.,  $170\times$  faster than Diff-Dock), making it a fast and accurate binding framework. This efficiency is critical in real-world drug discovery scenarios, where time and resources are often limited.
1 1 1 1 1 1 1 1 1 1 1 1
Figure 1: An overview of FABind. Left: The pocket prediction module takes the whole protein and the ligand as input and predicts the coordinates of the pocket center, where the ligand is randomly placed at the center of the protein. After determining the pocket center, a pocket is defined as a set of amino acids within a fixed radius around the center. Subsequently, the docking module moves the ligand to the pocket center and the ligand-pocket pair iteratively goes through the FABind layers to obtain the final pose prediction.  $M$  and  $N$  are the number of layers in pocket prediction and docking. Right: Architecture of FABind layers. Each layer contains three modules: independent message passing takes place within each component to update node embeddings and coordinates; cross-attention captures correlations between residues and ligands and updates embeddings only; and interfacial message passing focuses on the interface, attentively updating coordinates and representations.
Protein/Molecule Modeling. Learning effective protein and molecule representations is fundamental to biology and chemistry research. Deep learning models have made great progress in protein and molecule presentation learning. Molecules are usually represented by graph or string formats, and the graph neural networks (e.g., GAT [46], GCN [48]) and sequence models (e.g., LSTM [12], Transformer [45]) are utilized to model the molecules. For proteins, the sequence models are the common choices since the FASTA sequence is the most widely adopted string format [38]. Considering the 3D representations, incorporating the geometric information/constraint, e.g., SE(3)-invariance or equivariance, is becoming promising to model the protein/molecule structures, and there are multiple equivariant/invariant graph neural networks proposed [39, 20, 3, 16]. Among them, AlphaFold2 [19] is almost the most successful geometry-aware modeling that achieves revolutionary performance in protein structure prediction. In our work, we also keep the equivariance when modeling the protein and ligand.
Pocket Prediction. Predicting binding pocket is critical in the early stages of structure-based drug discovery. Different approaches have been developed in the past years, including physical-chemical-based, geometric-based, and machine learning-based methods [40]. Typical geometric-based methods are SURFNET [26] and Fpocket [27], and many alternatives have been proposed [47, 6, 42]. Recently, machine learning-based, especially deep learning-based methods, have been promising, such as DeepSite [15] and DeepPocket [1] which directly predict the pocket using deep networks, and some works utilize the deep scoring functions (e.g., affinity score) to score and rank the pockets [36, 50]. Among them, P2Rank [23] is an open-sourced tool that is widely adopted in existing works [51, 31] to detect potential protein pockets. In this work, we distinguish ourselves from previous methods through the introduction of a novel equivariant module and the combination of two separate losses. Besides, we jointly train the pocket prediction module and the docking module, leveraging the knowledge gained from the docking module to improve the performance of pocket prediction.
Protein-Ligand Docking. Protein-ligand docking prediction is to predict the binding structure of the protein-ligand complex. Traditional methods usually take the physics-informed energy functions to score, rank, and refine the ligand structures, such as AutoDock Vina [43], SMINA [21], GLIDE [8]. Recently, geometric deep learning has been attractive and greatly advances docking prediction. There are two distinct approaches in docking research. Regression-based methods, such as EquiBind [41], TankBind [31], and E3Bind [51], directly predict the docked ligand pose. On the other hand, sampling-based methods, like DiffDock [7], require extensive ligand pose sampling and optimization, but often yield more accurate predictions. Our work lies in the research of regression-based methods with much less computational costs, but achieves comparable performance to sampling-based approaches.
An overview of our proposed FABind is presented in Fig. 1. We first clarify the notations and problem definition in Section 3.1. In Section 3.2, we illustrate our FABind layer in detail. The specific design of the pocket prediction module is explained in Section 3.3, while the docking module is introduced in Section 3.4. Furthermore, the training pipeline is comprehensively described in Section 3.5.
Notations. For each protein-ligand complex, we represent it as a graph and denote it as  $\mathcal{G} = (\mathcal{V} := \{\mathcal{V}^l, \mathcal{V}^p\}, \mathcal{E} := \{\mathcal{E}^l, \mathcal{E}^p, \mathcal{E}^{lp}\})$ . Specifically, the ligand subgraph is  $\mathcal{G}^l = (\mathcal{V}^l, \mathcal{E}^l)$ , where node  $v_i = (\mathbf{h}_i, \mathbf{x}_i) \in \mathcal{V}^l$  is an atom,  $\mathbf{h}_i$  is the pre-extracted feature by TorchDrug [52], and  $\mathbf{x}_i \in \mathbb{R}^3$  is the corresponding coordinate. The number of atoms is denoted as  $n^l$ .  $\mathcal{E}^l$  is the edge set that represents chemical bonds in the ligand. The protein subgraph is  $\mathcal{G}^p = (\mathcal{V}^p, \mathcal{E}^p)$ , where node  $v_j = (\mathbf{h}_j, \mathbf{x}_j) \in \mathcal{V}^p$  is a residue,  $\mathbf{h}_j$  is initialized with the pre-trained ESM-2 [28] feature following DiffDock [7], and  $\mathbf{x}_j \in \mathbb{R}^3$  is the coordinate of the  $C_\alpha$  atom in the residue. The number of residues is denoted as  $n^p$ .  $\mathcal{E}^p$  is the edge set constructed by a cut-off distance  $8\AA$ . The input of the pocket prediction module is  $\mathcal{G}$ . We further denote the pocket subgraph as  $\mathcal{G}^{p*} = (\mathcal{V}^{p*}, \mathcal{E}^{p*})$ , with  $n^{p*}$  residues in the pocket. The pocket and ligand form a new complex as the input to the docking module:  $\mathcal{G}^{lp*} = (\mathcal{V} := \{\mathcal{V}^l, \mathcal{V}^{p*}\}, \mathcal{E} := \{\mathcal{E}^l, \mathcal{E}^{p*}, \mathcal{E}^{lp*}\})$ , where  $\mathcal{E}^{lp*}$  defines edges in the external contact surface. Detailed edge construction rule is in Appendix Section A.4. For clarity, we always use indices  $i, k$  for ligand nodes, and  $j, k'$  for protein nodes.
Problem Definition. Given a bounded protein and an unbounded ligand as inputs, our goal is to predict the binding pose of the ligand, denoted as  $\{\mathbf{x}_i^{l*}\}_{1\leq i\leq n_l}$ . Following previous works, we focus on blind docking scenario in which we have zero knowledge about protein pocket.
3.2 FABind Layer
In this section, we provide a comprehensive description of the FABind layer. For clarity, we use FABind in the pocket prediction module for a demonstration.
Overview. Besides node-level information, we explicitly model pair embedding for each protein residue-ligand atom pair  $(i,j)$ . We follow E3bind [51] to construct a pair embedding  $\mathbf{z}_{ij}$  via an outer product module (OPM):  $\mathbf{z}_{ij} = \mathrm{Linear}\left(\left(\mathrm{Linear}(\mathbf{h}_i) \otimes \mathrm{Linear}(\mathbf{h}_j)\right)\right)$ . For the initial pair embedding, OPM is operated on the transformed initial protein/ligand node embedding  $\mathbf{h}_i / \mathbf{h}_j$ . As illustrated in Figure 1, each FABind layer conducts three-step message passing: (1) Independent message passing. The independent encoder first passes messages inside the protein and ligand to update node embeddings and coordinates. (2) Cross-attention update. This block operates to exchange information across every node and updates pair embeddings accordingly. (3) Interfacial message passing. This layer focuses on the contact surface and attentively updates coordinates and representations for such nodes. The fundamental concept behind this design is the recognition of distinct characteristics between internal interactions within the ligand or protein and external interactions between the ligand and protein in biological functions. After several layers of alternations, we perform another independent message passing for further adjustment before the output of pocket prediction/docking modules. Notably, the independent and interfacial message passing layers are E(3)-equivariant, while cross-attention update layer is E(3)-invariant since it does not encode structure. These ensure each layer is E(3)-equivariant.
Independent Message Passing. We introduce a variant of Equivariant Graph Convolutional Layer (EGCL) proposed by EGNN [39] as our independent message passing layer. For simplicity, here we only illustrate the detailed message passing of ligand nodes, while the protein updates are in a similar way. With the ligand atom embedding  $\mathbf{h}_i^l$  and the corresponding coordinate  $\mathbf{x}_i^l$  in the  $l$ -th layer, we perform the independent message passing as follows:
$$
\mathbf {m} _ {i k} = \phi_ {e} \left(\mathbf {h} _ {i} ^ {l}, \mathbf {h} _ {k} ^ {l}, \left\| \mathbf {x} _ {i} ^ {l} - \mathbf {x} _ {k} ^ {l} \right\| ^ {2}\right),
$$
$$
\mathbf {h} _ {i} ^ {l} = \mathbf {h} _ {i} ^ {l} + \phi_ {h} \left(\mathbf {h} _ {i} ^ {l}, \sum_ {k \in \mathcal {N} (i | \mathcal {E} ^ {l})} \mathbf {m} _ {i k}\right), \quad \mathbf {x} _ {i} ^ {l} = \mathbf {x} _ {i} ^ {l} + \frac {1}{| \mathcal {N} (i | \mathcal {E} ^ {l}) |} \sum_ {k \in \mathcal {N} (i | \mathcal {E} ^ {l})} \left(\mathbf {x} _ {i} ^ {l} - \mathbf {x} _ {k} ^ {l}\right) \phi_ {x} \left(\mathbf {m} _ {i k}\right),
$$
where  $\phi_e, \phi_x, \phi_h$  are Multi-Layer Perceptions (MLPs) [10] and  $\mathcal{N}(i|\mathcal{E}^l)$  denotes the neighbors of node  $i$  regarding the internal edges  $\mathcal{E}^l$  of the ligand.
Cross-Attention Update. After independent message passing, we enhance the node feature with cross-attention update over all protein/ligand nodes by passing the messages from all ligand/protein nodes. The pair embeddings are also updated accordingly. We also take the ligand embedding update for clarity. Given the ligand atom representations and the pair embeddings, we first perform multi-head cross-attention over all protein residues:
$$
a _ {i j} ^ {(h)} = \mathrm {s o f t m a x} _ {j} \left(\frac {1}{\sqrt {c}} \mathbf {q} _ {i} ^ {(h) ^ {\top}} \mathbf {k} _ {j} ^ {(h)} + b _ {i j} ^ {(h)}\right), \quad \mathbf {h} _ {i} ^ {l} = \mathbf {h} _ {i} ^ {l} + \mathrm {L i n e a r} \left(\mathrm {c o n c a t} _ {1 \leq h \leq H} \left(\sum_ {j = 1} ^ {n ^ {p ^ {*}}} a _ {i j} ^ {(h)} \mathbf {v} _ {j} ^ {(h)}\right)\right),
$$
where  $\mathbf{q}_i^{(h)},\mathbf{k}_j^{(h)},\mathbf{v}_j^{(h)}$  are linear projections of the node embedding, and  $b_{ij}^{(h)} = \mathrm{Linear}(\mathbf{z}_{ij}^l)$  is a linear transformation of pair embedding  $\mathbf{z}_{ij}^l$ . The protein embeddings  $\mathbf{h}_j^l$  are updated similarly. Based on updated node embeddings  $\mathbf{h}_i^l$  and  $\mathbf{h}_j^l$ , the pair embeddings are further updated by  $\mathbf{z}_{ij}^l = \mathbf{z}_{ij}^l +\mathrm{OPM}(\mathbf{h}_i^l,\mathbf{h}_j^l)$ .
Interfacial Message Passing. With the updated protein and ligand representations, we perform an interfacial message passing to update the included node features and the coordinates on the contact surface. Our interfacial message passing derives from MEAN [22] with an additional attention bias. The detailed updates are as follows:
$$
\alpha_ {i j} = \frac {\exp \left(\mathbf {q} _ {i} ^ {\top} \mathbf {k} _ {i j} + b _ {i j}\right)}{\sum_ {j \in \mathcal {N} (i | \mathcal {E} ^ {l p *})} \exp \left(\mathbf {q} _ {i} ^ {\top} \mathbf {k} _ {i j} + b _ {i j}\right)},
$$
$$
\mathbf {h} _ {i} ^ {l + 1} = \mathbf {h} _ {i} ^ {l} + \sum_ {j \in \mathcal {N} (i | \mathcal {E} ^ {l p *})} \alpha_ {i j} \mathbf {v} _ {i j}, \quad \mathbf {x} _ {i} ^ {l + 1} = \mathbf {x} _ {i} ^ {l} + \sum_ {j \in \mathcal {N} (i | \mathcal {E} ^ {l p *})} \alpha_ {i j} \left(\mathbf {x} _ {i} ^ {l} - \mathbf {x} _ {j} ^ {l}\right) \phi_ {x v} (\mathbf {v} _ {i j}),
$$
where  $\mathbf{q}_i = \phi_q(\mathbf{h}_i^l)$ ,  $\mathbf{k}_{ij} = \phi_k\left(\| \mathbf{x}_i^l -\mathbf{x}_j^l\| ^2,\mathbf{h}_j^l\right)$ ,  $\mathbf{v}_{ij} = \phi_v\left(\| \mathbf{x}_i^l -\mathbf{x}_j^l\| ^2,\mathbf{h}_j^l\right)$ , and  $b_{ij} = \phi_b(\mathbf{z}_{ij}^l)$ , and  $\phi_q,\phi_k,\phi_v,\phi_b,\phi_{xv}$  are MLPs.  $\mathcal{E}^{lp*}$  denotes the external edges between ligand and protein contact surface constructed by cut-off distance  $10\AA$
3.3 Pocket Prediction
In the pocket prediction module, given the protein-ligand complex graph  $\mathcal{G}$ , our objective is to determine the amino acids of the protein that belong to the pocket. Previous works such as TankBind and E3Bind both use P2Rank [23] to produce multiple pocket candidates. Subsequently, either affinity score (TankBind) or self-confidence score (E3Bind) is utilized to select the most appropriate docked pose. Though P2Rank is faster and better than previous tools, it is based on numerical algorithms and traditional machine learning classifiers. Furthermore, the incorporation of P2Rank necessitates the selection of candidate poses following multiple poses docking. These factors could potentially restrict the performance and efficiency of fully deep learning-based docking approaches.
In our work, we propose an alternative method by treating pocket prediction as a binary classification task on the residues using the FABind layer, where each residue in the protein is classified as belonging to the pocket or not. Hence, the pocket prediction is more unified with the deep learning docking. Specifically, we use a binary cross-entropy loss to train our pocket classifier:
$$
\mathcal {L} _ {p} ^ {c} = - \frac {1}{n _ {p}} \sum_ {j = 1} ^ {n _ {p}} \left[ y _ {j} \log (p _ {j}) + (1 - y _ {j}) \log (1 - p _ {j}) \right],
$$
where  $n_p$  is the number of residues in the protein, and  $y_j$  is the binary indicator for residue  $j$  (i.e., 1 if it belongs to a pocket, 0 otherwise).  $p_j = \sigma(\mathrm{MLP}(\{\mathrm{FABind~layer}(\mathcal{G})\}_j))$  is the predicted probability of residue  $j$  belonging to a pocket,  $\mathcal{G}$  is the protein-ligand complex graph, and  $\sigma$  is sigmoid function.
Besides the direct classification of each residue to decide the pocket, the common practice of leveraging P2Rank pocket prediction is to first predict a pocket center coordinate, then a sphere near the pocket center under a radius  $20\mathring{\mathrm{A}}$ . Therefore, we add a constraint about the pocket center to make a more accurate prediction.
Constraint for Pocket Center. To constrain the predicted pocket center, we introduce a pocket center regression task over the classified pocket residues. Given  $n^{p'}$  predicted pocket residues  $\mathcal{V}^{p'} = \{v_1^{p'}, v_2^{p'}, \dots, v_{np'}^{p'}\}$  from our classifier, the pocket center coordinate is  $\mathbf{x}^{p'} = \frac{1}{n^{p'}} \sum_{j=1}^{n^{p'}} \mathbf{x}_j^{p'}$ . Then we can add a distance loss between the predicted pocket center  $\mathbf{x}^{p'}$  and the native pocket center  $\mathbf{x}^{p*}$ . This pocket center computation inherently involves discrete decisions – selecting which amino acids contribute to the pocket. Hence, we apply Gumbel-Softmax [14] to produce a differentiable approximation of the discrete selection process. It provides a probabilistic "hard" selection, which more accurately reflects the discrete decision to include or exclude an amino acid in the pocket.
$$
\gamma_ {j} ^ {p} = \frac {\exp \left(\left(\log \left(p _ {j}\right) + g _ {j}\right) / \tau_ {e}\right)}{\sum_ {k ^ {\prime} = 1} ^ {n ^ {p}} \exp \left(\left(\log \left(p _ {k ^ {\prime}}\right) + g _ {k ^ {\prime}}\right) / \tau_ {e}\right)},
$$