EquiPocket: an E(3)-Equivariant Geometric Graph Neural Network for
  Ligand Binding Site Prediction
Predicting the binding sites of target proteins plays a fundamental role in
drug discovery. Most existing deep-learning methods consider a protein as a 3D
image by spatially clustering its atoms into voxels and then feed the voxelized
protein into a 3D CNN for prediction. However, the CNN-based methods encounter
several critical issues: 1) defective in representing irregular protein
structures; 2) sensitive to rotations; 3) insufficient to characterize the
protein surface; 4) unaware of protein size shift. To address the above issues,
this work proposes EquiPocket, an E(3)-equivariant Graph Neural Network (GNN)
for binding site prediction, which comprises three modules: the first one to
extract local geometric information for each surface atom, the second one to
model both the chemical and spatial structure of protein and the last one to
capture the geometry of the surface via equivariant message passing over the
surface atoms. We further propose a dense attention output layer to alleviate
the effect incurred by variable protein size. Extensive experiments on several
representative benchmarks demonstrate the superiority of our framework to the
state-of-the-art methods.
Yang Zhang, Zhewei Wei, Ye Yuan, Chongxuan Li, Wenbing Huang
# EquiPocket: an E(3)-Equivariant Geometric Graph Neural Network for Ligand Binding Site Prediction

###### Abstract.

Predicting the binding sites of the target proteins plays a fundamental role in drug discovery. Most existing deep-learning methods consider a protein as a 3D image by spatially clustering its atoms into voxels and then feed the voxelized protein into a 3D CNN for prediction. However, the CNN-based methods encounter several critical issues: 1) defective in representing irregular protein structures; 2) sensitive to rotations; 3) insufficient to characterize the protein surface; 4) unaware of data distribution shift. To address the above issues, this work proposes EquiPocket, an E(3)-equivariant Graph Neural Network (GNN) for binding site prediction. In particular, EquiPocket consists of three modules: the first one to extract local geometric information for each surface atom, the second one to model both the chemical and spatial structure of the protein, and the last one to capture the geometry of the surface via equivariant message passing over the surface atoms. We further propose a dense attention output layer to better alleviate the data distribution shift effect incurred by the variable protein size. Extensive experiments on several representative benchmarks demonstrate the superiority of our framework to the state-of-the-art methods.

Binding Site Prediction, Graph Neural Network, Drug Discovery +
Footnote †: journal: Computer Science and

ACM Research in Computing and Networking, 2018

+
Footnote †: journal: Computer Science and

## 1. Introduction

Nearly all biological and pharmacological processes in living systems involve interactions between receptors (_i.e._ target proteins) and ligands (_i.e._ small molecules or other proteins) (Steintein et al., 2017). The places where such interactions occur are known as the binding sites/pockets of the ligands on the target protein structures, which are essential to determine whether or not the ligands are druggable and functionally relevant. Moreover, the knowledge of the binding site is able to facilitate many downstream tasks, such as docking (Steintein et al., 2017), the design of drug molecules (Stein et al., 2018), etc. Therefore, predicting the binding sites of the target proteins via in-silico algorithms forms an indispensable and even the first step in drug discovery.

Through the past years, plenty of computational methods have been proposed to detect binding sites, which can be roughly classified into three categories (Stein et al., 2017): the geometry-based (Stein et al., 2017; Wang et al., 2018; Wang et al., 2018; Wang et al., 2018; Wang et al., 2018), the probe-energy-based (Wang et al., 2018; Wang et al., 2018; Wang et al., 2018; Wang et al., 2018) and the template-based (Stein et al., 2017; Wang et al., 2018). The computational methods exploit hand-crafted algorithmic procedures guided by domain knowledge or external templates, leading to insufficient expressivity in representing complicated proteins. Recently, with the accumulation of labeled data and the development of machine learning techniques, the learning-based approaches have been developed (Wang et al., 2018), which manage to analyze and extract the underlying patterns of the input data that eventually align with the assign labels through the iterative process of learning. Although the learning-based methods have exhibited clear superiority over the classical computational counterparts, their performance and flexibility are still limited by the hand-crafted input features and the insufficiently-expressive models they used (Wang et al., 2018).

More recently, motivated by the breakthrough of deep learning in a variety of fields, Convolutional Neural Networks (CNNs) have been applied successfully for the prediction of binding sites (Wang et al., 2018). Typical works include DeepSite (Wang et al., 2018), DeepPocket (Chen et al., 2018), DeepSurf (Wang et al., 2018), etc. The CNN-based methods consider a protein as a 3D image by spatially clustering its atoms into the nearest voxels, and then model

Figure 1. Illustrative comparison between previous CNN-based methods and our EquiPocket.
the binding site prediction as a object detection problem or a semantic segmentation task on 3D grids. Thanks to the multi-layer and end-to-end learning paradigm, the CNN-based methods are observed to outperform traditional learning-based approaches and generally achieve the best performance on several public benchmarks (Shen et al., 2017).

In spite of the impressive progress, existing CNN-based models still encounter several issues as below:

**Issue 1.** Defective in leveraging regular voxels to model the proteins of irregular shape. First, a considerable number of voxels probably contain no atom since the protein atoms are unevenly distributed in space, which yields unnecessary redundancy in computation and memory. Moreover, the voxelization is usually constrained within a fixed-size space (_e.g._\(70\AA\times 70\AA\times 70\AA\)) (Shen et al., 2017; Wang et al., 2018). The atoms beyond this size will be directly discarded, resulting in incomplete and inaccurate modeling particularly for large proteins. Besides, although the voxelization process is able to encode certain spatial structure of the protein, it overlooks the irregular chemical interactions (_i.e._ the chemical bonds) between atoms and the topological structure upon that is also useful for binding site detection.

**Issue 2.** Sensitive to rotations. To discretize the protein into 3D grids, the CNN methods fix the three bases of the coordinates beforehand. When rotating the protein, the voxelization results could be distinct, and the predicted binding sites will change, which, however, conflicts with the fact that any rotation of the protein keeps the binding sites invariant. While it can be alleviated by the local grid (Shen et al., 2017) or augmenting training data with random rotations (Shen et al., 2017; Wang et al., 2018), which yet is data-dependent and unable to guarantee rotation invariance in theory.

**Issue 3.** Insufficient to characterize the geometry of the protein surface. The surface atoms comprise the major part of the binding pocket, which should be elaborately modeled. In the CNN-based methods, the surface atoms are located in the voxels that are surrounded by empty voxels, which somehow encodes the surface geometry. Nevertheless, such information is coarse to depict how the surface atoms interact and what their local geometry is. Indeed, the description of the surface atoms is purely driven by the geometric shape of the solvent-accessible surface of the protein (Shen et al., 2017) (Figure 1(b)), which, unfortunately, is less explored in current works.

**Issue 4.** Unaware of data distribution shift. In practical scenarios, the size of the proteins varies greatly across different datasets. It requires the deep learning model we apply to be well generalizable and adaptive, so that it is able to overcome the distribution shift incurred by the variable protein size. However, this point is not seriously discussed previously.

In this paper, to address the above issues, we propose to apply Graph Neural Networks (GNNs) (Grover et al., 2016; Wang et al., 2018; Wang et al., 2018) instead of CNNs to represent proteins. By considering atoms as nodes, interactions as edges, GNNs are able to encode the irregular protein structures by multi-layer message passing. More importantly, a recent line of researches (Grover et al., 2016; Wang et al., 2018; Wang et al., 2018) has enhanced GNNs by encapsulating E(3) equivariance/invariance with respect to translations/rotations; in this way, equivariant GNNs yield outputs that are independent of the choice of the coordinate systems, leading to improved generalization ability. That being said, trivially applying equivariant GNNs for the binding site prediction task is still incapable of providing desirable performance, and even achieves worse accuracy than the CNN-based counterparts. By looking into the design of the architecture, equivariant GNNs naturally cope with the first two issues as mentioned above, yet leave the other two unsolved. To this end, we make the contributions as follows:

* To the best of our knowledge, we are the first to apply an E(3)-equivariant GNN for ligand binding site prediction, which is dubbed **EquiPocket**. In contrast to conventional CNN-based methods, EquiPocket is free of the voxelization process, able to model irregular protein structures by nature, and insensitive to any Euclidean transformation, thereby addressing Issue 1 and 2.
* EquiPocket consists of three modules: the first one to extract local geometric information for each surface atom with the help of solvent-accessible surface (Shen et al., 2017), the second one to model both the chemical and spatial structure of the protein, and the last one to capture the comprehensive geometry of the surface via equivariant message passing over the surface atoms. The first and last module are proposed to tackle Issue 3, while the second module attempts to involve both the chemical and spatial interactions, as presented in Issue 1.
* To resolve Issue 4, namely, alleviating the effect by data distribution shift, we further propose a novel output layer called _dense attention output layer_ in Equipocket, which enables us to adaptively balance the scope of the receptive field for each atom in accordance to the density distribution of the neighbor atoms.
* Extensive experiments on serveral representative benchmarks demonstrate the superiority of our framework to the state-of-the-art methods in prediction accuracy. The design of our model is sufficiently ablated as well.

## 2. Related Work

### Binding Site Prediction

**Computational Methods.** The computational methods for binding site prediction include geometry-based (Beng et al., 2016; Wang et al., 2018; Wang et al., 2018; Wang et al., 2018; Wang et al., 2018; Wang et al., 2018), probe- and energy-based (Wang et al., 2018; Wang et al., 2018; Wang et al., 2018) and template-based (Beng et al., 2016; Wang et al., 2018) methods: 1) Since most ligand binding sites occur on the 3D structure, geometry-based methods (POCKET (Wang et al., 2018), CriticalFinder (Wang et al., 2018), LigSite (Wang et al., 2018), Focket (Wang et al., 2018), etc. ) are designed to identify these hollow spaces and then rank them using the expert design geometric features. 2) Probe-based methods (SURFNET (Wang et al., 2018), Q-SiteFinder (Wang et al., 2018), etc. (Wang et al., 2018)), also known as energy-based methods, calculate the energy resulting from the interaction between protein atoms and a small-molecule probe, whose value dictates the existence of binding sites. 3)Template-based methods (FINDSITE (Beng et al., 2016), LIBRA (Wang et al., 2018), etc.) are mainly to compare the required query protein with the published protein structure database to identify the binding sites.

**Traditional Learning-based Methods.** PRANK (Kang et al., 2018) is a learning-based method that employs the traditional machine learning algorithm random forest(RF) (Beng et al., 2016). Based on the pocket points and chemical properties from Fpocket (Wang et al., 2018) and Concavity (Beng et al., 2016), this method measures the "ligandability" as the binding ability of a candidate pocket using the RF model. However, those methods require the manual extraction of numerous features with limit upgrading.

**CNN-based Methods.** Over the last few years, deep learning has surpassed far more traditional ML methods in many domains.
For binding site prediction task, many researchers (Han et al., 2017; Wang et al., 2018; Wang et al., 2019; Wang et al., 2019; Wang et al., 2019) regard a protein as a 3D image, and model this task as a computer vision problem. DeepSite (Wang et al., 2019) is the first attempt to employ the CNN architecture for binding site prediction, which like P2Rank (Wang et al., 2019) treats this task as a binary classification problem and converts a protein to 3D voxelized grids. The methods FRSite (Wang et al., 2019) and Kalasundy (Klass and Kalasundy, 2019) adhere to the principle of deepsite, but the former regards this task as an object detection problem, and the latter regards this task as a semantic segmentation task. Deeppocket (Han et al., 2017) is a method similar to p2rank, but implements a CNN-based segmentation model as the scoring function in order to more precisely locate the the binding sites. The recent CNN-based method DeepSurf (Wang et al., 2019) constructs a local 3D grid and updates the 3D-CNN architecture to mitigate the detrimental effects of protein rotation.

### Graph Neural Networks for Molecule Modeling

There are multi-level information in molecules including atom info, chemical bonds, spatial structure, physical constraints, etc. Numerous researchers view molecules as topological structures and apply topological-based GNN models (like graph2vec (Wang et al., 2019), GAT (Wang et al., 2019), GCN (Chen et al., 2019), GCN2 (Chen et al., 2019), GIN (Yang et al., 2019) and etc. (Wang et al., 2019)) to extract the chemical info, which achieve positive outcomes. With the accumulation of structure data for molecules, spatial-based graph models (DimeNet (Wang et al., 2019), DimNet++ (DimeNet et al., 2019), SphereNet (Sandel et al., 2019), SchNet (Wang et al., 2019), Egnn (Yang et al., 2019), (Wang et al., 2019) and etc.) are proposed for molecule task which aggregates spatial and topological information. However, these models may not be adequate for macro-molecules due to their high calculation and resource requirements.

## 3. Notations and Definitions

Protein GraphA protein such as the example in Figure 1(b) is denoted as a graph \(\mathcal{G}_{P}=(\mathcal{V}_{P},\mathcal{E}_{C},\mathcal{E}_{D})\), where \(\mathcal{V}_{P}=\{v_{0},...,v_{N}\}\) forms the set of \(N\) atoms, \(\mathcal{E}_{C}\) represents the chemical-bond edges, and \(\mathcal{E}_{D}\) collects the spatial edges between any two atoms if their spatial distance is less than a cutoff \(\theta>0\). In particular, each node (_i.e._ atom) is associated with a feature \((\mathbf{x}_{i},\mathbf{c}_{i})\), where \(\mathbf{x}_{i}\in\mathbb{R}^{3}\) denotes the 3D coordinates and \(\mathbf{c}_{i}\in\mathbb{R}^{5}\) is the chemical feature.

Surface Point SetThe surface geometry of a protein is of crucial interest for binding site detection. Here we define the set of surface points, by \(\mathcal{S}=\{s_{0},...,s_{M}\}\), \(M\gg N\). Each surface point \(s_{i}\) is NOT necessarily an atom of the protein, and it corresponds to \((\mathbf{x}_{i},p_{i})\), where \(\mathbf{x}_{i}\in\mathbb{R}^{3}\) represents the 3D coordinates of \(s_{i}\) and \(p_{i}\in\mathcal{V}_{P}\) indicates the index of the nearest protein atom in \(\mathcal{V}_{P}\) to \(s_{i}\). We employ the open source MSMS (Wang et al., 2019) to derive surface points.

Protein Surface GraphBy referring to the surface points defined above, we collect all the nearest protein atoms \(p_{i}\) of the surface points, giving rise to the surface graph \(\mathcal{G}_{S}=(\mathcal{V}_{S},\mathcal{E}_{S})\), and clearly \(\mathcal{G}_{S}\subseteq\mathcal{G}_{P}\). _We call the atoms in the surface graph as surface atoms, which are distinguished from surface points defined in the last paragraph_. Notably, the edges of the surface graph, _i.e._, \(\mathcal{E}_{S}\) is only composed of spatial edges from \(\mathcal{E}_{D}\), since those chemical edges are mostly broken among the extracted atoms.

Equivariance and InvarianceIn 3D space, the symmetry of the physical laws requires the detection model to be equivariant with respect to arbitrary coordinate systems (Klass and Kalasundy, 2019). In form, suppose \(\mathbf{X}\) to be 3D geometric vectors (positions, velocities, etc) that are steerable by E(3) group (rotations/translations/reflections), and \(\mathbf{h}\) non-steerable

Figure 2. An illustration of the scheme of our EquiPocket framework.
features. The function \(f\) is E(3)-equivariant, if for any transformation \(g\in\mathbb{E}(3)\), \(f(g\cdot\mathbf{X},\mathbf{h})=g\cdot f(\mathbf{X},\mathbf{h})\), \(\forall\mathbf{X}\in\mathbb{R}^{3\times m},\mathbf{h}\in\mathbb{R}^{d}\). Similarly, \(f\) is invariant if \(f(g\cdot\mathbf{X},\mathbf{h})=f(\mathbf{X},\mathbf{h})\). The group action \(\cdot\) is instantiated as \(g\cdot\mathbf{X}:=\mathbf{X}+\mathbf{b}\) for translation \(\mathbf{b}\in\mathbb{R}^{3}\) and \(g\cdot\mathbf{X}(t):=\mathbf{OX}\) for rotation/reflection \(\mathbf{O}\in\mathbb{R}^{3\times 3}\).

Problem StatementGiven a protein \(\mathcal{G}_{P}\) and its surface points \(\mathbb{S}\), as well as the constructed surface graph \(\mathcal{G}_{S}\), our goal is to learn an E(3)-invariant model \(f(\mathcal{G}_{P},\mathbb{S},\mathcal{G}_{S})\) to predict the atoms of the binding site: \(\mathcal{V}_{B}\subseteq\mathcal{V}_{P}\).

## 4. The Proposed Methodology

Figure 2 illustrates the overall framework of our EquiPocket, which consists of three modules: the _local geometric modeling module_ SS 4.1 that focuses on extracting the geometric information of each surface atom, the _global structure modeling module_ SS 4.2 to characterize both the chemical and spatial structures of the protein, and the _surface message passing module_ SS 4.3 which concentrates on capturing the entire surface geometry based on the extracted information by the two former modules. The training losses are also presented. We defer the pseudo codes of EquiPocket to Appendix 1.

### Local Geometric Modeling Module

This subsection presents how to extract the local geometric information of the protein surface \(\mathcal{G}_{S}\), with the help of surface points \(\mathbb{S}\). The local geometry of each protein atom closely determines if the region nearby is appropriate or not to become part of binding sites. We adopt the surrounding surface points of each protein surface atom to describe the local geometry.

To be specific, for every surface atom \(i\in\mathcal{V}_{S}\), its surrounding surface points are returned by a subset of \(\mathbb{S}\), namely, \(\mathbb{S}_{i}=\{s_{j}=(\mathbf{x}_{j},p_{j})\in\mathbb{S}\mid p_{j}=i\}\), where \(p_{j}\), as defined before, indicates the nearest protein atom. We now construct the geometric information based on \(\mathbb{S}_{i}\). We denote the center/mean of all 3D coordinates in \(\mathbb{S}_{i}\) as \(\bar{x}_{i}\). For each surrounding surface point \(s_{j}\in\mathbb{S}_{i}\), we first search its two nearest surface points from \(\mathbb{S}\) as \(s_{j_{1}}\) and \(s_{j_{2}}\), and then calculate the following relative position vectors:

\[\left\{\begin{array}{l}\mathbf{x}_{jj_{1}}=\mathbf{x}_{j}-\mathbf{x}_{j_{1}},\\ \mathbf{x}_{jj_{2}}=\mathbf{x}_{j}-\mathbf{x}_{j_{2}},\\ \mathbf{x}_{j,\text{center}}=\mathbf{x}_{j}-\bar{\mathbf{x}}_{i},\\ \mathbf{x}_{j,\text{protein}}=\mathbf{x}_{j}-\mathbf{x}_{i},\\ \mathbf{x}_{\text{center,protein}}=\bar{\mathbf{x}}_{i}-\mathbf{x}_{i}.\end{array}\right. \tag{1}\]

We further derive the following scalars upon Eq. 1:

\[\begin{split}&\mathbf{g}(s_{j}):=[\|\mathbf{x}_{jj_{1}}\|_{2},\|\mathbf{x}_{ jj_{2}}\|_{2},\angle_{1},\\ &\|\mathbf{x}_{j,\text{center}}\|_{2},\|\mathbf{x}_{j,\text{protein}}\|_{2},\|\mathbf{x}_{ \text{protein,center}}\|_{2},\angle_{2}],\end{split} \tag{2}\]

where the angels are computed by \(\angle_{1}=\frac{\mathbf{x}_{jj_{1}}\cdot\mathbf{x}_{jj_{2}}}{\|\mathbf{x}_{jj_{1}}\| \frac{\mathbf{x}_{jj_{1}}}{\|\mathbf{x}_{jj_{2}}\|_{2}}}\) and \(\angle_{2}=\frac{\mathbf{x}_{j,\text{center}}\cdot\mathbf{x}_{\text{center,protein}}}{\|\mathbf{x}_{j, \text{center}}\|_{2}}\|\mathbf{x}_{\text{center,protein}}\|_{2}\); here the operator \(\cdot\) defines the inner-product between two vectors. Basically, as displayed in Figure 3, the first three quantities in \(\mathbf{g}(s_{j})\) depict how the nearby surface points are arranged around \(s_{j}\), and the last four ones describe where \(s_{j}\) is located within the global region of \(\mathbb{S}_{i}\).

We aggregate the geometric information \(\mathbf{g}(s_{j})\) over all surface points in \(\mathbb{S}_{i}\) and obtain a readout descriptor for surface atom \(i\) as

\[\begin{split}&\mathbf{g}_{i}=[\text{Pooling}(\{\text{MLP}(\mathbf{g}(s_{j}))\}_{s_{j}\in \mathbb{S}_{i})},\\ &\text{MLP}(\text{Pooling}(\{\mathbf{g}(s_{j})\})_{s_{j}\in\mathbb{S}_ {i}\})]\end{split} \tag{3}\]

Here, MLP denotes multi-layer perceptron, and the function Pooling is implemented as a concatenation of mean pooling and max pooling throughout our experiments. The front part in Eq. 3 is used to gather local geometric features, while the latter part attempts to compute the global size of surrounding surface points. Notably, the geometric descriptor \(\mathbf{g}_{i}\) is E(3)-invariant.

### Global Structure Modeling Module

This module aims at processing the information of the whole protein \(\mathcal{G}_{P}\), including atom type, chemical bonds, relevant spatial positions, etc. Although the binding pocket is majorly comprised of surface atoms, the global structure of the protein in general influences how the ligand is interacted with and how the pocket is formulated, which should be modeled. We fulfil this purpose via two concatenated processes: chemical-graph modeling and spatial-graph modeling.

The chemical-graph modeling process copes with the chemical features \(\{\mathbf{c}_{i}\}_{i\in\mathcal{V}_{P}}\) and the chemical interactions \(\mathcal{E}_{C}\) of the protein graph. For each atom in the protein, its chemical type, the numbers of electrons around, and the chemical bonds connected to other atoms are important clues to identify the interaction between the protein and the ligand (Stein and Grinner, 2017). We employ typical GNNs (Stein and Grinner, 2017; Grinner, 2018; Grinner, 2019) to distill this type of information. Formally, we proceed:

\[\{\mathbf{c}^{\prime}_{i}\}_{i\in\mathcal{V}_{P}}=\text{GNN}(\{\mathbf{c}_{i}\}_{i\in \mathcal{V}_{P}},\mathcal{E}_{C}), \tag{4}\]

where \(\mathbf{c}^{\prime}_{i}\) is the updated chemical feature for atom \(v_{i}\). While various GNNs can be used in Eq. 4, here we implement GAT (Grinner, 2019) given its desirable performance observed in our experiments.

The spatial-graph modeling process further involves the 3D coordinates \(\{\mathbf{x}_{i}\}_{i\in\mathcal{V}_{P}}\) to better depict the spatial interactions \(\mathcal{E}_{D}\) within the protein. Different from chemical features \(\mathbf{c}^{\prime}_{i}\), the 3D coordinates provide the spatial position of each atom and reflect the pair-wise distances in 3D space, which is helpful for physical interaction modeling. We leverage EGNN (Stein and Grinner, 2019) as it conforms to E(3) equivariance/invariance and achieves promising performance on modeling spatial graphs. Specifically, we process EGNN as follows:

\[\{\mathbf{c}^{\prime\prime}_{i}\}_{i\in\mathcal{V}_{P}}=\text{EGNN}(\{\mathbf{x}_{i}, \mathbf{c}^{\prime}_{i}\}_{i\in\mathcal{V}_{P}},\mathcal{E}_{D}). \tag{5}\]

Here, we only reserve the invariant output (_i.e_., \(\mathbf{c}^{\prime\prime}_{i}\)) and have discarded the equivariant output (_e.g_. updated 3D coordinates) of EGNN, since the goal of this module is to provide invariant features.

Figure 3. An illustration of local geometric features.
We select the updated features of the surface atoms \(\mathcal{V}_{\text{S}}\), which will be fed into the module in SS 4.3.

### Surface Message Passing Module.

Given the local geometric features \(\{\mathbf{g}_{i}\}_{i\in\mathcal{V}_{\text{S}}}\) from SS 4.1, and the globally-encoded features of the surface atoms \(\{\mathbf{c}^{\prime\prime}_{i}\}_{i\in\mathcal{V}_{\text{S}}}\) from SS 4.2, the module in this subsection carries out equivariant message passing on the surface graph \(\mathcal{G}_{\text{S}}\) to renew the entire features of the protein surface. We mainly focus on the surface atoms here, because firstly the surface atoms are more relevant to the binding sites than the interior atoms, and secondly the features \(\{\mathbf{c}^{\prime\prime}_{i}\}_{i\in\mathcal{V}_{\text{S}}}\) that are considered as the input have somehow encoded the information of the interior structure via the processes in 4.2.

**Surface-EGNN.** During the \(l\)-th layer message passing, each node is associated with an invariant feature \(\mathbf{h}^{(l)}_{i}\in\mathbb{R}^{m_{l}}\) and an equivariant double-channel matrix \(\mathbf{X}^{(l)}_{i}\in\mathbb{R}^{3\times 2}\). We first concatenate \(\mathbf{c}^{\prime\prime}_{i}\) with \(\mathbf{g}_{i}\) as the initial invariant feature:

\[\mathbf{h}^{(0)}_{i}=[\mathbf{c}^{\prime\prime}_{i},\mathbf{g}_{i}]. \tag{6}\]

The equivariant matrix \(\mathbf{X}^{(0)}_{i}\) is initialized by the 3D coordinates of the atom and the center of its surrounding surface points, that is,

\[\mathbf{X}^{(0)}_{i}=[\mathbf{x}_{i},\bar{\mathbf{x}}_{i}]. \tag{7}\]

We update \(\mathbf{h}^{(l)}_{i}\in\mathbb{R}^{ld}\) and \(\mathbf{X}^{(l)}_{i}\in\mathbb{R}^{3\times 2}\) synchronously to unveil both the topological and geometrical patterns. Inspired from EGNN (Steintein and Tschur, 2017) and its multi-channel version GMN (Garvin et al., 2018), we formulate the \(l\)-th layer for each surface atom \(i\in\mathcal{V}_{\text{S}}\) as:

(8) \[\mathbf{m}_{ij} =\phi_{m}\left(\mathbf{h}^{(l)}_{i},\mathbf{h}^{(l)}_{j},f_{x}(\mathbf{X}^{(l) }_{i},\mathbf{X}^{(l)}_{j}),e_{ij}\right),\] (9) \[\mathbf{h}^{(l+1)}_{i} =\phi_{h}\left(\mathbf{h}^{(l)}_{i},\sum\nolimits_{j\in\mathbb{N}(i )}\mathbf{m}_{ij}\right),\] (10) \[\mathbf{X}^{(l+1)}_{i} =\mathbf{X}^{(l)}_{i}+\frac{1}{|\mathbb{N}(i)|}\sum\nolimits_{j\in \mathcal{N}(i)}(\mathbf{x}^{(l)}_{i}-\mathbf{x}^{(l)}_{j})\phi_{x}(\mathbf{m}_{ij}),\] (11)

where the functions \(\phi_{m},\phi_{h},\phi_{x}\) are all MDs, \(\mathbb{N}(i)\) denotes the neighbors of node \(i\) in terms of the spatial edges \(\mathcal{E}_{d}\), \(|\cdot|\) counts the size of the input set, and the invariant message \(\mathbf{m}_{ij}\) from node \(j\) to \(i\) is employed to update the invariant feature \(\mathbf{h}^{(l+1)}_{i}\) via \(\phi_{h}\) and the equivariant matrix \(\mathbf{X}^{(l+1)}_{i}\) via the aggregation of the relative position \(\mathbf{x}^{(l)}_{i}-\mathbf{x}^{(l)}_{j}\) multiplied with \(\phi_{x}\).

As a core operator in the message passing above, the function \(f_{x}(\mathbf{X}_{i},\mathbf{X}_{j})\) is defined as follows:

\[f_{x}(\mathbf{X}_{i},\mathbf{X}_{j})\coloneqq\{\|\mathbf{x}_{ij}\|_{2},\|\mathbf{x}_{ci}\|_{2 },\|\mathbf{x}_{ej}\|_{2},\angle_{c1,ij,i},\angle_{c2,ij,ij},\angle_{c1,cj}\}, \tag{12}\]

where, the relative positions are given by \(\mathbf{x}_{ij}=\mathbf{x}_{i}-\mathbf{x}_{j},\mathbf{x}_{ci}=\bar{\mathbf{x}}_{i}-\mathbf{x}_{i}\) and \(\mathbf{x}_{cj}=\bar{\mathbf{x}}_{j}-\mathbf{x}_{j}\); the angles \(\angle_{c1,ij}\), \(\angle_{c2,ij,j}\), \(\angle_{c1,c2,j}\) defined as the inner-products of the corresponding vectors denoted in the subscripts, _e.g._, \(\angle_{c1,ij}=\frac{\mathbf{x}_{ci}\cdot\mathbf{x}_{ij}}{\|\mathbf{x}_{ci}\|_{2}\|_{2}\|_{ 2}}\). Through the design in Eq. 12, \(f_{x}(\mathbf{X}_{i},\mathbf{X}_{j})\) elaborates the critical information (including relative distances and angles) around the four points: \(\mathbf{x}_{i},\bar{\mathbf{x}}_{i},\mathbf{x}_{j},\bar{\mathbf{x}}_{j}\), which largely characterizes the geometrical interaction between the two input matrices. Nicely, \(f_{x}(\mathbf{X}_{i},\mathbf{X}_{j})\) is invariant, ensuring the equivariance of the proposed Surface-EGNN.

**Dense Attention Output Layer.** Conventionally, we can apply the output of the final layer, _i.e._, \((\mathbf{h}^{(L)}_{i},\mathbf{X}^{(L)}_{i})\) to estimate the binding site. Nevertheless, such flat output overlooks the discrepancy of size and shape between different proteins. As showed in Figure 5(b), for small or densely-connected proteins, the receptive field of each node will easily cover most nodes after a small number of message-passing layers, and excessive message passing will lead to over-smoothing (Krause et al., 2018) that will incurs performance detriment. For large or sparsely-connected proteins, on the contrary, insufficient message passing can hardly attain the receptive field with a desirable scope, which will also decrease the performance. It thus requires us to develop an adaptive mechanism to balance the message passing scope between different proteins. We propose the _dense attention output layer_ to achieve this goal.

Intuitively, for each target atom, the spatial distribution of the neighbors is able to reflect the density of the spatial connections around. This motivates us to calculate the proportion of the atoms with different distance ranges. As \(\theta\) is the cutoff to create the spatial graph, we use it as the distance unit. We compute by:

\[n^{(l)}_{i}=\frac{|\{j\in\mathcal{V}_{\text{P}}\mid 0\leq\|\mathbf{x}_{i}-\mathbf{x}_{j}\|_{2 }<l\theta\}|}{N_{\text{P}}}, \tag{13}\]

where, the proportion is evaluated within the distance range \([0,l\theta]\), \(N_{\text{P}}=|\mathcal{V}_{\text{P}}|\), and the neighbor hop \(l\in\mathbb{Z}^{+}\). We collect the proportions of all hops from 0 to \(L\), yielding the proportion vector \(\mathbf{n}_{i}=[n^{(0)}_{i},n^{(1)}_{i},\cdots,n^{(L)}_{i},N_{\text{P}}]\in\mathbb{R }^{L+2}\) with \(N_{\text{P}}\) plus to emphasize the total number of the protein atoms. Clearly, \(\mathbf{n}_{i}\) contains rich information of the spatial density, and we apply it to determine the importance of different layers, by producing the attention as:

\[\mathbf{a}_{i}=\text{Sigmoid}(\phi_{a}(\mathbf{n}_{i})). \tag{14}\]

Here, \(\phi_{a}\) is an MLP with the number of output channels as \(L+1\), the Sigmoid function1 is applied for each channel, implying that \(\mathbf{a}_{i}\in(0,1)^{L+1}\). We then multiply the hidden feature of the corresponding layer with each channel of the attention vector, and concatenate them into a vector:

Footnote 1: Note that the sum of all channels of \(\mathbf{a}_{i}\) is unnecessarily equal to 1, since the Sigmoid function instead of the previously-used SoftMax function is applied here.

\[\mathbf{h}^{\text{out}}_{i}=\text{Concat}(a_{i0}\mathbf{h}^{(0)}_{i},...,a_{iL}\mathbf{h}^{(L )}_{i}),\]

Figure 4. An illustration of Dense Attention in a Protein.
where \(a_{il}\) is the \(l\)-th channel of \(\mathbf{a_{i}}\). By making use of Eq. 14, the learnable attentions enable the model to adaptively balance the importance of different layers for different input proteins. We will illustrate the benefit of the proposed strategy in our experiments. As for the coordinates, we simply compute the mean of all layers to retain translation equivariance:

\[\mathbf{X}_{i}^{\text{out}}=\frac{1}{L+1}\sum_{l=0}^{L}\mathbf{X}_{i}^{(l)}. \tag{15}\]

### Optimization Objective

We set \(y_{i}=1\) if a surface atom \(i\) is within \(4\)A to any ligand atom (Kang et al., 2017). We predict the probability \(\hat{y_{i}}\) of being a part of binding site according its dense embedding \(h_{i}\).

\[\hat{y_{i}}=\textbf{Sigmoid}(\text{MLP}(\mathbf{h}_{i}^{\text{out}})). \tag{16}\]

Following (Kang et al., 2017; Li et al., 2018), Dice loss is used:

\[\mathcal{L}_{b}=1-\frac{2\cdot\sum(\hat{y_{i}}\cdot y_{i})}{\sum(\hat{y_{i}}) +\sum(y_{i})+\epsilon}, \tag{17}\]

where \(\epsilon>0\) is a small value to maintain numeric stability.

**Predict the relative direction of nearest ligand atom.** Beyond the CNN-based methods, our EquipPocket is an E(3)-equivariant model, which can not only output the embedding \(\mathbf{h}_{i}^{\text{out}}\) but also the coordinate matrix \(\mathbf{X}_{i}^{\text{out}}\) (with initial position vector \(\mathbf{x_{i}}\)). We further leverage the position vector to predict the relative direction \(\mathbf{d_{i}}\) of its nearest ligand atom (with position vector \(\mathbf{m_{i}}\)), in order to enhance our framework to gather local geometric features.

\[\mathbf{d_{i}}=\frac{\mathbf{m_{i}}-\mathbf{x_{i}}}{\|\mathbf{m_{i}}-\mathbf{x_{i}}\|_{2}},\quad \hat{\mathbf{d_{i}}}=\frac{\mathbf{x}_{i}^{\text{out}}-\mathbf{x_{i}}}{\|\mathbf{x}_{i}^{ \text{out}}-\mathbf{x_{i}}\|_{2}}. \tag{18}\]

The cosine loss function is used for the direction loss \(\mathcal{L}_{d}\):

\[\mathcal{L}_{d}=\sum(1-\cos(\hat{\mathbf{d_{i}}},\mathbf{d_{i}})). \tag{19}\]

The eventual loss is \(\mathcal{L}=\mathcal{L}_{b}+\mathcal{L}_{d}\). We train the parameters of all the three modules end to end.

## 5. Experiments

In this section, we will conduct experiments on multiple datasets to evaluate the performance of our framework in comparison to baseline methods and investigate the following tasks:

* **Task 1.** Can our framework's performance match or surpass that of existing methods?
* **Task 2.** Can different modules of our framework bring significant improvement for binding site prediction?
* **Task 3.** Can our framework mitigate the detrimental effects of the data distribution shift?
* **Task 4.** How do the hyperparameters (the cutoff \(\theta\) and depth of surface-egnn) affect the performance and computational cost?

### Experimental Settings

#### 5.1.1. Dataset

We conduct experiments based on the following datasets:

* **scPDB**(Kang et al., 2017) is the famous dataset for binding site prediction, which contains the protein structure, ligand structure, and 3D cavity structure generated by VolSite (Kang et al., 2017). The 2017 release of scPDB is used for training and cross-validation of our framework, which contains 17,594 structures, 16,034 entries, 4,782 proteins, and 6,326 ligands.
* **PDBbind**(Kang et al., 2017) is a well-known and commonly used dataset for the research of protein-ligand complex. It contains the 3D structures of proteins, ligands, binding sites, and accurate binding affinity results determined in the laboratory. We use the release of v2020, which consists of two parts: general set (14, 127 complexes) and refined set (5,316 complexes). The general set contains all protein-ligand complexes. The refined set contains better-quality compounds selected from the general set, which is used for the test in our experiments.
* **COACH 420 and HOLO4K** are two test datasets for the binding site prediction, which are first introduced by (Kang et al., 2017). Consistent with (Beng et al., 2017; Kang et al., 2017; Kang et al., 2017), we use the mlig subsets of each dataset for evaluation, which contain the relevant ligands for binding site prediction.

**Data Distribution Shift.** As depicted in the Figure 5(a) and Table 1 that after data processing, there is a significant gap in protein size and protein distribution between the training dataset (scPDB) and the test dataset (COACH420, HOLO4K, PDBbind). The number of atoms within a protein ranges from hundreds to tens of thousands. As for protein distribution in datasets, scPDB has the longest average structure, followed by HOLO4k and PDBbind, with COACH420 having the shortest average protein structure. This fact will hurt model learning and generalization, as discussed in SS 5.2.3.

#### 5.1.2. Target of Binding Sites

The CNN-based methods (Beng et al., 2017; Li et al., 2018; Li et al., 2018) label the subgrid as positive if its geometric center is closer than \(4\)A to the binding sites geometric center. In our experiment, consistent with (Kang et al., 2017), we set the protein atoms within \(4\)A of any ligand atom as positive and negative otherwise. After obtaining the probability that an atom is a candidate binding site, we use the mean-shift algorithm (Beng et al., 2017) to predict the binding site center, which can determine the number of clusters on its own (details in Appendix A.2.4).

#### 5.1.3. Data Preparation

We perform the following four processing steps: i) Cluster the structures in scPDB by their Uniprot IDs, and

\begin{table}
\begin{tabular}{c c c c c} \hline \hline \multirow{2}{*}{DataSet} & \multicolumn{4}{c}{Average} \\ \cline{2-5}  & Atom Num & Atom in Surface & Surface Points & Target Atoms \\ \hline scPDB & 4205 & 2317 & 24010 & 47 \\ COACH420 & 2123 & 1217 & 12325 & 58 \\ HOLO4k & 3845 & 2052 & 20023 & 106 \\ PDBbind & 3104 & 1677 & 17357 & 37 \\ \hline \hline \end{tabular}
\end{table}
Table 1. Summary of Dataset

Figure 5. The protein distribution of datasets and spatial distribution of atom neighbors
select the longest sequenced protein structures from every cluster as the train data (Kalasanty et al., 2017). Finally, 5,372 structures are selected out. ii) Split proteins and ligands for the structures in COACH420 and HOLO4k, according to the research (Kalasanty et al., 2016). iii) Clean protein by removing the solvent, hydrogens atoms. Using MSMS (Kalasanty et al., 2016) to generate the solvent-accessible surface of a protein. iv) Read the protein file by RDKIT (Kalasanty et al., 2016), and extract the atom and chemical bond features. Remove the error structures.

#### 5.1.4. Evaluation Metrics

**DCC** is the distance between the predicted binding site center and the true binding site center. **DCA** is the shortest distance between the predicted binding site center and any atom of the ligand. The samples with DCC(DCA) less than the threshold are considered successful. The samples without any binding site center are considered failures. Consistent with (Becker et al., 2015; Kalasanty et al., 2016; Kalasanty et al., 2017; Kalasanty et al., 2018), threshold is set to 4 A. We use **Success Rate** and **Failure Rate** to evaluate experimental performance.

\[\text{Success Rate}(\text{DCC})=\frac{1(\{\text{Predicted sites}|\text{DCC}<\text{ threshold}\})}{1(\{\text{True sites}\})},\] \[\text{Success Rate}(\text{DCA})=\frac{1(\{\text{Predicted sites}|\text{DCA}<\text{threshold}\})}{1(\{\text{True sites}\})},\] \[\text{Failure Rate}=\frac{1(\{\text{Protein}|1(\text{predicted binding center})=0\})}{1(\{\text{Protein}\})}, \tag{20}\]

where \(1(\cdot)\) represents the cardinality of a set. After ranking the predicted binding sites, we take the same number with the true binding sites to calculate the success rate.

#### 5.1.5. EquiPocket Framework

We implement our EquiPocket framework based on (GAT (Kalasanty et al., 2016)+EGNN (Kalasanty et al., 2016)) as our global structure modeling module. The cutoff \(\theta\) and depth in our surface-egnn model are set to 6 and 4.

To indicate the EquiPocket Framework with different modules, we adopt the following symbol as follows: i) **EquiPocket-L**: Only contain the local geometric modeling module. ii) **EquiPocket-G**: Only contain the global structure modeling module. iii) **EquiPocket-LG**: Only contain both the local geometric and global structure modeling modules. iii) **EquiPocket**: Contain all the modules.

#### 5.1.6. Baseline Models

We compare our framework with the following models: 1) geometric-based method(Fpocket (Kalasanty et al., 2016)), 2) CNN-based methods (DeepSite (Kalasanty et al., 2016), Kalasanty et al. (2016) and DeepSurf (Kalasanty et al., 2016)), 3) topological graph-based models (GAT (Kalasanty et al., 2016), GCN (Kalasanty et al., 2016) and GCN2 (Kalasanty et al., 2016)), 4) spatial graph-based models (SchNet (Net et al., 2017), EGNN (Kalasanty et al., 2016)).

#### 5.1.7. Environment and Parameter

We implement our EquiPocket framework in PyTorch Geometric, all the experiments are conducted on a machine with an NVIDIA A100 GPU (80GB memory). We take 5-fold cross validation on training data scPDB and use valid loss to save checkpoint. The batch size is set to 8. For baseline models and models in global structure modeling module of EquiPocket, we use them suggested settings to get optimal performance. More details and related resources for our experiments can be found in **Appendix A.2.2**.

### Result analysis

#### 5.2.1. Model Comparison

In Table 2, we compared our EquiPocket framework with baseline methods mentioned above. As can be observed, the performance of the computational method Fpocket is inferior, with no failure rate, since it simply employs the geometric feature of a protein. The performance of CNN-based methods is much superior to that of the conventional method, with DCC and DCA metrics improving by more than 50 percent but requiring enormous parameter values and computing resources. However, these two early methods DeepSite and Kalasanty are hampered by data distribution shift (Issue 4) and their inability to process big proteins, which may fail prediction. The recently proposed method DeepSurf employs the local-grid concept to handle any size of proteins, although CNN architecture also still results in inevitable failures.

For graph models, the poor performance of topological-graph models (GCN, GAT, GCN2) is primarily due to the fact that they only

\begin{table}
\begin{tabular}{l c c c c c c c c c} \hline \hline \multirow{2}{*}{Methods} & \multirow{2}{*}{Type} & Param & Failure & \multicolumn{2}{c}{COACH420} & \multicolumn{2}{c}{HOLO4K} & \multicolumn{2}{c}{PDBbind2020} \\ \cline{4-10}  & & (M) & Rate \(\downarrow\) & DCC\(\uparrow\) & DCA\(\uparrow\) & DCC\(\uparrow\) & DCA\(\uparrow\) & DCC\(\uparrow\) & DCA\(\uparrow\) \\ \hline Fpocket\({}^{\text{b}}\) & Geometric-based & \(\backslash\) & **0.000** & 0.228 & 0.444 & 0.192 & 0.457 & 0.253 & 0.371 \\ \hline DeepSite\({}^{\text{b}}\) & & 1.00 & \(\backslash\) & \(\backslash\) & 0.564 & \(\backslash\) & 0.456 & \(\backslash\) & \(\backslash\) \\ Kalasanty\({}^{\text{b}}\) & 3D-CNN & 70.64 & 0.120 & 0.335 & 0.636 & 0.244 & 0.515 & 0.416 & 0.625 \\ DeepSurf\({}^{\text{b}}\) & & 33.06 & 0.054 & 0.386 & **0.658** & 0.289 & 0.635 & 0.510 & 0.708 \\ \hline GAT & & **0.03** & 0.11 & 0.039(0.005) & 0.130(0.009) & 0.036(0.003) & 0.110(0.010) & 0.032(0.001) & 0.088(0.011) \\ GCN & Topological & 0.06 & 0.163 & 0.049(0.001) & 0.139(0.010) & 0.044(0.003) & 0.174(0.003) & 0.018(0.001) & 0.070(0.002) \\ GAT + GCN & Graph & 0.08 & 0.31 & 0.036(0.009) & 0.131(0.021) & 0.042(0.003) & 0.152(0.020) & 0.022(0.008) & 0.074(0.007) \\ GCN2 & & 0.11 & 0.466 & 0.042(0.098) & 0.131(0.017) & 0.051(0.004) & 0.163(0.008) & 0.023(0.007) & 0.089(0.013) \\ \hline SchNet & Spatial & 0.49 & 0.14 & 0.168(0.019) & 0.444(0.020) & 0.192(0.005) & 0.501(0.004) & 0.263(0.003) & 0.457(0.004) \\ Egnn & Graph & 0.41 & 0.270 & 0.156(0.017) & 0.361(0.020) & 0.127(0.005) & 0.406(0.004) & 0.143(0.007) & 0.302(0.006) \\ \hline EquiPocket-L & & 0.15 & 0.552 & 0.070(0.009) & 0.171(0.008) & 0.044(0.004) & 0.138(0.006) & 0.051(0.003) & 0.132(0.009) \\ EquiPocket-G & Ours & 0.42 & 0.292 & 0.159(0.016) & 0.373(0.021) & 0.129(0.005) & 0.411(0.005) & 0.145(0.007) & 0.311(0.007) \\ EquiPocket-LG & & 0.50 & 0.220 & 0.212(0.016) & 0.443(0.011) & 0.183(0.004) & 0.502(0.008) & 0.274(0.004) & 0.462(0.005) \\ EquiPocket & & 1.70 & 0.051 & **0.423(0.014)** & 0.556(0.007) & **0.337(0.006)** & **0.662(0.007)** & **0.545(0.010)** & **0.721(0.004)** \\ \hline \hline \end{tabular}

* The standard deviation of each index is indicated in brackets. The result of 5-fold for EquiPocket is shown in Appendix A.2.5.
* We use their published pre-train models or published result, details in Appendix A.2.3.

\end{table}
Table 2. Experimental and ablation results of baseline models and our framework.a
consider atom attributes and chemical bond information, ignoring the spatial structure in a protein. The performance of spatial-graph models is generally better than that of topological-graph models. EGNN model utilizes not only the properties of atoms but also their relative and absolute spatial positions, resulting in a better effect. SchNet merely updates the information of atoms based on the relative distance of atoms. We attempt to execute the Diment++ (Diment++, 2018), which uses the angle info between atoms and atoms, but it requires too many computing resources, resulting in an OOM (Out Of Memory) error. However, the performance of the spatial-graph model is worse than that of the CNN-based and geometric-based methods because the former cannot obtain enough geometric features (Issue 3) and cannot address the data distribution shift (Issue 4).

As the above results indicate, geometric info of protein surface and multi-level structure info in a protein is essential for binding site prediction. In addition, it reflects the limitations of the current GNN models, where it is difficult to collect sufficient geometric information from the protein surface or the calculation resources are too large to apply to macromolecular systems like proteins. Consequently, our EquiPocket framework is not only able to update chemical and spatial information from an atomic perspective but also able to effectively collect geometric information without excessive computing expense, resulting in a 10-20% increase in effect over previous results. Case study based on different methods is showed in Appendix 5.2.5.

#### 5.2.2. **Ablation Study**

As shown in Table 2, we conduct ablation experiments on our EquiPocket framework with different modules.

**Local Geometric Modeling Module.** This module is used to extract the geometric features of protein atoms from their nearest surface points. EquiPocket-G consists solely of this module, and the performance is negligible. There are two primary causes for this result. First, geometric information can only determine part of the binding sites. Second, it can only reflect the geometric features over a relatively small distance and cannot cover an expansive area.

**Global Structure Modeling Module.** The primary purpose of this module is to extract information about the whole protein, such as atom type, chemical bonds, relevant spatial positions, etc. We implement EquiPocket-G based on (GAT + EGNN) models, which is E(3) equivariance/invariance and has a better effect than its predecessor, EquiPocket-L. In comparison, the value of DCC increased by about 10%, and DCA increased by about 20%. This demonstrates that structure information of the whole protein is necessary for binding site prediction. In addition, when the two modules are combined as the EquiPocket-LG, the prediction effect is significantly improved, proving the complementarity of surface geometric information and global structure information.

**Surface Message Passing Module.** In the previous model, EquiPocket-LG, information was extracted solely from atoms and their closest surface points. Nonetheless, the binding site is determined not only by the information of a single atom but also by the atoms surrounding it. Therefore, the surface message passing module is proposed to collect and update the atom's features from its neighbors. After adding this module, the performance of EquiPocket has been significantly enhanced, DCC and DCA have increased by approximately 20% on average, and the failure rate has been significantly reduced. Through the addition of multiple modules, we address the Issue 3 and the performance of our framework eventually surpasses that of the existing SOTA method, demonstrating the efficacy of our framework design.

#### 5.2.3. **Data Distribution Shift**

As shown in Figure 6, we calculate the average DCC with the distribution of various sizes proteins. The geometric-based method Fpocket only utilizes the geometric features of a protein surface. Therefore, its performance is superior to that of most other methods for proteins with fewer than 1,000 atoms, but its prediction effect decreases significantly as the size of the protein increases. Kalasanty is a CNN-based and learn-based method. As the number of atoms in the protein varies, the prediction effect exhibits an increasing and then a decreasing trend, which is not only influenced by the size of the protein but also has a significant correlation with the dataset's distribution. According to the train data (scPDB), the majority of proteins contain fewer than 2,000 protein atoms (as depicted in Figure 5(a)). Consequently, the model's parameters will be biased toward this protein size. In addition, for proteins with more than 8000 atoms, the prediction effect is not even as good as the geometric-based method. This is due to the fact that CNN methods typically restrict the protein space to 70A * 70A * 70A, and for proteins larger than this size, the prediction frequently fails. For our EquiPocket framework, we do not need to cut the protein into grids, and we utilize both geometric information from the surface points and global structure information from the whole protein, so the performance for proteins of varying sizes is significantly superior to that of other methods.

**Dense Attention.** The Dense Attention is introduced in SS 4.3 to reduce the negative impact caused by the data distribution shift (Issue 4). As shown in 6, when the number of atoms contained in a protein is less than 3000, the result of the EquiPocket (w/o Dense Attention) is weaker than that of the original EquiPocket, whereas between the protein is larger, there is no significant distinction between the two models. It simply reflects the role of Dense Attention, which, by weighting the surface-egnn layer at different depths, mitigates the detrimental effect of the data distribution shift.

**Direction Loss.** Direction loss is a novel task designed to improve the extraction of local geometric features. The result of the EquiPocket (w/o Direction Loss) in Figure 6 demonstrates conclusively that the prediction performance of small proteins with fewer than 3,000 atoms is diminished in the absence of this task, which reveals the importance of the task.

Figure 6. The performance of EquiPocket and baseline methods for proteins of various sizes.
#### 5.2.4. **Hyperparameters Analysis**

In our EquiPocket framework, the cutoff \(\theta\) and depth of surface-egnn are two crucial parameters that can impact performance and computational efficiency.

**Cutoff \(\theta\).** We set the depth of surface-egnn to 4 and implement various cutoff values (2, 4, 6, 8, 10). Figure 7(a) indicates that when the cutoff is set to 2, the average DCC of our framework is poor, and GPU memory consumption is relatively low (22GB). This is due to the fact that when the cutoff is small, the surface-egnn can only observe a tiny receptive field. As the cutoff increases, the performance and GPU memory continue to rise until the DCC reaches a bottleneck when the cutoff is 10, and the GPU memory reaches 62GB. Therefore, when selecting parameters for our framework, we must strike a balance between performance and efficiency.

**Depth.** The depth of surface-egnn has an immediate influences on the performance and computation cost. We set the cutoff to 6 and implement various depth (1, 2, 3, 4, 5, 6). Figure 7(b) demonstrates that as depth increases, performance steadily improves and becomes stable as GPU memory continues to expand. Because the prediction of binding sites is highly influenced by their surrounding atoms, therefore, an excessively large receptive field may not offer any benefits but will necessitate additional computing resources.

#### 5.2.5. Case Study

We also display two examples of our EquiPocket and other methods in Figure 8. We take two proteins, 1f8e (with 12,268 atoms) and 5ei3 (with 1,572 atoms), from the test dataset PDBbind. As can be seen from Figure 8: The binding sites predicted by the geometry-based method Fpocket are extremely distant from the actual binding sites. This is due to the fact that this method prioritizes local geometric information and disregards the multi-level structure information of proteins, resulting in limited scope and weak performance. The CNN-based method Kalasanty did not provide any predicted binding site for protein 1f8e. We conjecture that this method restricts the protein within a specific space size which is highly susceptible to failure with large proteins. The recently-proposed CNN-based method DeepSur takes local grids on the protein surface, which can address the issue of fixed space size. However, the prediction of binding sites in protein 5ei3 by DeepSur is far from the ground truth because the CNN-based methods are defective in obtaining geometric and chemical features. Our EquiPocket framework is unaffected by the shortcomings of the aforementioned methods, allowing it to achieve superior outcomes for both large and small proteins.

## 6. Conclusion

In this paper, concentrating on the ligand binding site prediction, we propose a novel E(3)-Equivariant geometric graph framework called EquiPocket, which contains the local geometric modeling module, global structure modeling module, and surface passing module to gather the surface geometric and multi-level structure features in a protein. Experiments demonstrate that our framework is highly generalizable and beneficial, and achieves superior prediction accuracy and computational efficiency compared with the existing methods.

### Future Work

#### 6.1.1. **Protein Surface**

As demonstrated by our experiments, the geometric information derived from the protein surface plays a significant role in the prediction of binding sites. In this work, we use MSMS (Song et al., 2017) to generate the protein surface, which may have an uncertain impact on the prediction results. Therefore, in the future, we will be able to establish a more efficient surface generation method or gather the geometric information of a protein without a fixed protein surface.

#### 6.1.2. **Global Structure Model**

In this work, we take the existing graph-based models (GAT + EGNN) to gather the multi-level structure information in a protein. However, their capabilities are limited because these models are not entirely tailored to complex structures such as proteins. In the future, we will be able to develop more effective models to gather information from proteins in order to improve the prediction performance of binding sites.

#### 6.1.3. **Computing Resources**

As the experimental results show, both our EquiPocket and CNN-based methods require a significant amount of computing resources to analyze protein data, which will have a negative impact on their actual deployment. Consequently, when applying our method to real-world problems, it is crucial to consider how to compress algorithm parameters, increase computational efficiency, and decrease resource consumption.
