Hashing for Protein Structure Similarity Search

A protein is represented as a graph $\mathcal{G(V,E)}$ with $\mathcal{V}$ denoting the set of nodes and $\mathcal{E}$ denoting the set of edges. Each amino acid corresponds to a node in the graph, and the edges are constructed by connecting each node with its $k$ nearest neighbors ($k$NN). The distance between nodes is measured by the Euclidean distance between the raw node features of the $C_{\alpha}$ atoms. The raw node features for distance computation are introduced in the following content.

Although the TM-score of proteins is typically computed based on the coordinates of $C_{\alpha}$ atoms, we also incorporate other backbone atoms to construct node features. Because $C_{\alpha}$ atoms and the other backbone atoms are constrained by the peptide plane and form the 3D protein shape together, it is reasonable to include them in node features. In particular, we calculate the bond and dihedral angles for each residual as our node features. As shown in Figure 2, the bond angle is the angle formed between two covalent bonds that share a common atom, and the dihedral angles, also known as torsion angles, refer to the angles between planes defined by four consecutive atoms in the protein backbone. For residual $i$, we denote the bond angles of $N_{i}\mbox{-}C_{\alpha_{i}}\mbox{-}C_{i},C_{i-1}\mbox{-}N_{i}\mbox{-}C_{\alpha_{i}},C_{\alpha_{i}}\mbox{-}C_{i}\mbox{-}N_{i+1}$ as $\alpha_{i},\beta_{i},\gamma_{i}$, and the dihedral angles of $N_{i-1}\mbox{-}C_{\alpha_{i-1}},C_{\alpha_{i-1}}\mbox{-}C_{i-1},C_{i}\mbox{-}N_{i}$ as $\phi_{i},\psi_{i},\omega_{i}$. The final node features are computed by $\{sin,cos\}\circ\{\alpha_{i},\beta_{i},\gamma_{i},\phi_{i},\psi_{i},\omega_{i}\}$. We use $\bm{V}\in R^{n\times d_{v}}$ to denote these hand-crafted node features (also called raw node features), where $n$ is the number of nodes and $d_{v}$ is the dimension of the raw node features.

Unlike existing learning-based methods which do not use edge features, we additionally extract edge features for better description of the structure. We consider the following five common types of atoms in the protein chain: $C$, $C_{\alpha}$, $N$, $O$, and $C_{\beta}$. If a $C_{\beta}$ atom is missing, we employ a technique in [26] to add a virtual $C_{\beta}$ atom based on the constraint of the other backbone atoms. To calculate the distance between the $i$th node and the $j$th node, we compute the Euclidean distance $\|X_{i}-Y_{j}\|$, where $X_{i}$ is chosen from the set $\{C_{i},C_{\alpha_{i}},N_{i},O_{i},C_{\beta_{i}}\}$, and $Y_{j}$ is chosen from the set $\{C_{j},C_{\alpha_{j}},N_{j},O_{j},C_{\beta_{j}}\}$. Here, $X_{i}$ and $Y_{j}$ represent the respective coordinates of the atoms. Subsequently, we encode the distances using a Gaussian radial basis function (RBF). The final representation of the edge features is computed by $\text{RBF}_{k}(\|X_{i}-Y_{j}\|)$, where $\text{RBF}_{k}$ denotes the $k$th RBF employed in the encoding process. We use $\bm{E}\in R^{m\times d_{e}}$ to denote these hand-crafted edge features (also called raw edge features), where $m$ is the number of edges and $d_{e}$ is the dimension of the raw edge features.

For each node $i$ in layer $l$, we update its representation $\bm{h}_{i}^{l}$ by aggregating features from its neighboring nodes and adjacent edges. This message passing mechanism is defined as follows:

	$\displaystyle\bm{u}_{i}^{l}$	$\displaystyle=\bm{h}_{i}^{l}+\frac{1}{|\mathcal{N}_{i}|}\sum_{j\in\mathcal{N}_{i}}NodeMLP(\bm{h}_{i}^{l}\|\bm{h}_{j}^{l}\|\bm{e}_{ij}^{l}),$		(1)
	$\displaystyle\tilde{\bm{h}}_{i}^{l}$	$\displaystyle=\Phi(\bm{h}_{i}^{l}+MLP(\bm{u}^{l}_{i})),$

where $\mathcal{N}_{i}$ denotes the index set of neighboring nodes of node $i$ and $|\mathcal{N}_{i}|$ denotes the number of nodes in $\mathcal{N}_{i}$. $\|$ denotes the concatenation operation. $NodeMLP$ is a two-layer feed-forward neural network. The updated node representation $\bm{u}_{i}^{l}$ is further updated by a $MLP$, which also denotes a two-layer feed-forward neural network and a residual connection from $\bm{h}_{i}^{l}$ is then added. $\Phi$ is a follow-up batchnorm function. This straightforward message passing mechanism is memory efficient so that we can sample more negative samples in a single mini-batch, the detail of which will be introduced in Section 3.4.

As proved in [27, 28], neglecting the update of edge features could lead to suboptimal performance. Hence, we also perform edge message passing in our method. For each edge $(i,j)$, we aggregate features from the connected node $i$ and node $j$, the result of which is then added to the original edge representation. Our edge message passing mechanism is defined as follows:

$\displaystyle\bm{e}_{ij}^{l+1}=\Phi(\bm{e}_{ij}^{l}+EdgeMLP(\tilde{\bm{h}}_{i}^{l}\|\tilde{\bm{h}}_{j}^{l}\|\bm{e}_{ij}^{l})).$

(2)

Here, the updated node representation $\tilde{\bm{h}}_{i}^{l}$ and $\tilde{\bm{h}}_{j}^{l}$ are concatenated with the edge representation $\bm{e}_{ij}^{l}$. The concatenated result will then be updated with an $EdgeMLP$, which is also a two-layer feed-forward neural network. For the next layer, we set $\bm{h}_{i}^{l+1}=\tilde{\bm{h}}_{i}^{l}$.

We first calculate the pairwise TM-score [16] for protein structures using TM-align [1], serving as the ground-truth measurement of structural similarity. In the training phase, we define two protein structures, denoted as $(P_{a},P_{b})$, to be similar based on the following criterion: if the TM-score between $P_{a}$ and $P_{b}$ is larger than a threshold $\rho$ multiplied by the maximum TM-score between $P_{a}$ and any other structure $P_{i}$ in the training set $\mathcal{D}$ excluding $P_{a}$, as shown below:

$$TM(P_{a},P_{b})\geq\rho\cdot max(\{TM(P_{a},P_{i})|{P_{i}\in\mathcal{D}\setminus P_{a}}\}),$$

(7)

where $TM(\cdot)$ denotes the TM-score, $\rho$ is a hyperparameter in (0, 1). In each time, we sample a mini-batch of data comprising one protein structure as query, one positive sample, and $K$ negative samples. The positive sample means structurally similar to the query, while the negative samples are dissimilar.

The above sampling strategy might cause an issue. We observe that some protein structures have few (less than 5) similar protein structures (positive samples), and hence, the same structure will be repeatedly sampled during training. This could result in a potential risk of overfitting. This is caused by those structure pairs with high TM-score, which improves the threshold of similar pairs in Eq.(7). To address this issue, we further propose a substructure sampling strategy to enhance the diversity of positive samples. This strategy has been used in the pretraining of protein structures [30, 28]. However, directly applying the empirical sampling length or ratio in these existing methods will lead to a significant decline in performance for our method.

Due to the sensitivity of positive samples to protein structure variations, casually sampling a substructure could make a positive sample no longer similar to the query. In this paper, we propose a novel substructure sampling method. More specifically, our substructure sampling method starts from a randomly selected amino acid within the protein, and then we traverse along the protein chain towards both ends until the desired length is attained. To make sure our sampled substructure will not deviate from our original structure too much, we use TM-score as a guidance. Assume the original positive sample is denoted as $P$, and the sampled substructure is denoted as $P_{s}$. We aim to satisfy the following constraint:

$$TM(P,P_{s})\geq\alpha.$$

(8)

Here $\alpha$ is set to 0.9, which means the two structures are extremely similar according to the meaning of TM-score [31]. Performing on-the-fly TM-score calculation during the training process would incur significant computation cost. Hence, we calculate the minimum sampling length for each structure to satisfy the above constraint before training. Through this finely controlled substructure sampling strategy, we achieve an increase in the diversity of positive examples while avoiding the failure of training. As shown in Figure 1, positive and negative samples are first sampled based on the query in the training phase. Then, the positive substructures are sampled based on the positive sample and pre-computed substructure length.

We employ three datasets in our experiment, which are SCOPe v2.07 [7] , ind_PDB [15], and Alphafold database [11]. SCOPe is a database to study protein structural relationships, which has been utilized in existing works [14, 15]. To ensure a fair comparison, we employ the same filtering criteria as in [15], and the resulting dataset contains 14,215 protein structures. ind_PDB has also been utilized in [15], which consists of 1,900 protein structures collected from the Protein Data Bank [32]. AlphaFold database contains over 200 million protein structures predicted by Alphafold 2. Following existing work [14, 15], we first conduct 5-fold cross-validation on the SCOPe dataset and then evaluate the performance on the ind_PDB dataset. AlphaFold database is mainly used for comparing memory cost.

We employ three metrics to evaluate the accuracy: the area under the receiver operating characteristic curves (AUROC), the area under the precision-recall curves (AUPRC), and the Top-k hit ratio. For each query, we calculate the AUROC and AUPRC, and report the average value across all queries. The Top-k hit ratio is computed as $\frac{1}{N_{P_{Q}}}\sum_{i=1}^{N_{P_{Q}}}\frac{N_{\text{hit}}^{i}}{\min(k,N_{\text{nbr}}^{i})}$, where $N_{P_{Q}}$ represents the number of query structures, $N_{\text{hit}}^{i}$ denotes the number of correctly identified similar structures in the top-k rankings, and $N_{\text{nbr}}^{i}$ indicates the total number of similar structures for the given query $P_{Q}$. Our evaluation considers the Top-k values of Top-1, Top-5, and Top-10. We follow the settings of existing works [14, 15], considering structure pairs with a TM-score larger than or equal to $0.9\cdot\max\{{\text{TM}(P_{Q},P_{i})|{P_{i}\in\mathcal{D}}}\}$ as similar pairs.

To prove the efficiency of POSH, we also compare POSH with baselines in terms of time and memory cost.

We adopt four alignment-free methods as baselines in our experiment, which include SGM [12], SSEF [13], DeepFold [14] and GraSR [15]. SGM and SSEF are non-learning methods. DeepFold and GraSR are learning-based methods. The results of SGM and SSEF are obtained by running their provided scripts. For DeepFold, we train it on our dataset using the same settings described in the original paper. The results of GraSR are directly copied from its original paper for comparison since we utilize the same dataset and filtering criteria.

We compare the accuracy of POSH with baselines on SCOPe dataset. The results are shown in Table 1. We can find that POSH outperforms all other methods on all metrics except the Top-1 hit ratio. For the Top-1 hit ratio, POSH is comparable to the best baseline, and significantly outperforms other baselines.

We also evaluate the accuracy of POSH and baselines on ind_PDB. The results are shown in Table 2. Compared with the SCOPe dataset, the accuracy of all the methods on the ind_PDB dataset drops. The gap can be attributed to the fact that each protein structure in the ind_PDB dataset may consist of multiple domains, but in SCOPe each protein structure represents a single domain. From another perspective, ind_PDB is a dataset more related to real-world scenarios. The results show that POSH consistently outperforms all the other methods across all evaluation metrics on the ind_PDB dataset. We can find that POSH achieves about 3%-7% improvement on Top-k hit ratio compared to the state-of-the-art baselines.

We show the time cost of searching in databases of different sizes for SGM, SSEF, GraSR and POSH in Figure 4. We use ind_PDB as the query dataset, and the number of protein structures in the target database increases from 100,000 to the size of the Alphafold database. Because the search time of TM-align ranges from hours to days, we exclude it in the figure. We have also excluded the results of DeepFold because DeepFold embeds the protein structure into the same dimension as GraSR, which results in the same time cost. We can find that the time cost of POSH is significantly lower than that of the other methods. POSH is nearly four and ten times faster than SGM and SSEF, and more than four times faster than the other best method GraSR. The advantage of POSH becomes more pronounced with the increasing database size, ensuring POSH’s fast search of large-scale data.

The memory cost and memory compression ratio of TM-align, POSH, and the other alignment-free methods are shown in Table 3. The process of encoding a protein structure into a vector can be seen as a form of compression, so we also calculate the compression ratio for the Alphafold database. For TM-align, all protein structures must be stored on the computing device to calculate pairwise similarity. In contrast, alignment-free methods allow for memory saving by representing each protein structure as a vector. Hence, we can perform computation on the computing device with the vector representation, and then index the desired proteins on the storage device.

From the results in Table 3, we can find that the alignment-free methods can dramatically reduce the memory cost, compared with the alignment-based method TM-align. Compared with existing alignment-free methods, POSH achieves a memory saving of more than six times, by up to more than 100 times. Taking Alphafold database as an example. When calculating the TM-score by utilizing TM-align directly on the original protein structure data, the resulting memory cost is 21TB, which is unmanageable. Even if the protein structures are represented as real-valued vectors by using some alignment-free methods, the memory cost is still high. Specifically, the memory costs of SGM and SSEF are 72GB and 1196GB, which are 6.5 and 108.7 times larger than that of POSH, respectively. It takes over 300GB memory cost for the other two learning-based methods DeepFold and GraSR, which is about 32 times larger than that of POSH. The memory cost of POSH is only 11GB, which is acceptable for many real-world applications.