Enhancing Protein Mutation Effect Prediction through a Retrieval-Augmented Framework

Authors: ruihan  guo
rui  wang
ruidong  wu
zhizhou  ren
jiahan  li
shitong  luo
zuofan  wu
qiang  liu
jian  peng
jianzhu  ma
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 17:24:55

Abstract: Predicting the effects of protein mutations is crucial for analyzing protein functions and understanding genetic diseases. 
However, existing models struggle to effectively extract mutation-related local structure motifs from protein databases, which hinders their predictive accuracy and robustness. To tackle this problem, we design a novel retrieval-augmented framework for incorporating similar structure information in known protein structures. We create a vector database consisting of local structure motif embeddings from a pre-trained protein structure encoder, which allows for efficient retrieval of similar local structure motifs during mutation effect prediction. 
Our findings demonstrate that leveraging this method results in the SOTA performance across multiple protein mutation prediction datasets, and offers a scalable solution for studying mutation effects.

Main paper: 
 introduction: 
 Protein fitness plays significant roles in diverse applications in pharmaceutical industry [Amara, 2013] , drug design [De Carvalho, 2011] , biofuel production [Huang et al., 2020] , and environmental bioremediation [Lu et al., 2022] . Deciphering mutation effects on protein fitness is crucial for understanding their functional dynamics and yet remains a central challenge in molecular biology. Most recent breakthroughs on computational predictions of mutation effects are driven by coevolutionary information [Riesselman et al., 2018 , Luo et al., 2021 , Notin et al., 2022] . The mutations on contacting residue pairs would become correlated under the evolutionary pressure to maintain protein stability and optimize functional efficiency within cellular environments. Consequently, conserved patterns within protein sequences and structures typically signify their stability and functionality. The predominant methodology to exploit such coevolutionary patterns is to perform multiple sequence alignments (MSA) [Thompson et al., 1994 [Thompson et al., , 1997] ] and fit either statistical [Seemayer et al., 2014] or machine learning models [Rao et al., 2021] . In addition to sequence-level alignments, performing domain-level structure clustering [Orengo et al., 1997 , Dong et al., 2018] is also a promising approach to extract evolutionary information from a protein family.In this paper, we explore an alternative perspective to retrieve information for mutation effect prediction. In contrast to the global protein representation considered by MSA and domain-level structure alignments, we extract coevolutionary information in the scope of local microenvironments. We focus on the alignment of local structure motifs, i.e., a central amino-acid with a few contacting neighbors. Such a local representation of coevolutionary information is specialized to the scenarios of protein engineering, where a common practice involves introducing a few point mutations to enhance a desired function [Shroff et al., 2020] . It is widely observed that the effects of point mutations are mainly given by the alteration of local biochemical microenvironments [Kim et al., 2011 , Lu et al., 2022] . Given these empirical insights, we propose to retrieve local structure fragments with similar backbone positions as auxiliary information for mutation effect prediction, while withdrawing the constraint on global sequence/structure similarity. This microscopic retrieval mechanism enables us to extract atom-level information from the whole protein universe rather than restricting to moleculelevel instances within a certain protein family.To elucidate the intricate details of local coevolutionary patterns in proteins, we employed a structurebased embedding approach, ESM-IF [Hsu et al., 2022] , to encode local structure motifs into latent embeddings, assuming the metric space of such embeddings measures the similarity of motif structures. We preprocess the entire Protein Data Bank (PDB) [Berman et al., 2003] and build a database, we call Structure Motif Embedding Database (SMEDB), to support fast information retrieval by GPUaccelerated k-nearest neighbors (kNN) search. This retrieval procedure, we call Multiple Structure Motif Alignment (MSMA), is designed to extract coevolutionary information from protein fragments with similar local structure. We rigorously evaluate these extracted embeddings, focusing on their structural similarity and predictive accuracy for mutation effects. Remarkably, the embeddings derived from local coevolutionary patterns demonstrate superior performance compared to those obtained from traditional Multiple Sequence Alignments (MSA). Furthermore, the distribution of these embeddings was found to be complementary to those derived from MSA profiles, indicating that they capture distinct aspects of coevolutionary information, thereby enhancing our understanding of protein structure and function dynamics.In addition, we introduce a novel model architecture, called Multi-Structure Motif Invariant Point Attention (MSM-IPA), to aggregate the retrieved coevolutionary motif information for predicting the structural fitness of proteins. This model, which effectively incorporates mutational information, has an outstanding capability to generalize across various mutations. It is trained to predict the change in binding free energy (∆∆G) on protein surfaces, an essential element for assessing protein-protein interactions. We extensively evaluate our model on a suite of widely-used protein stability and binding affinity benchmarks, including S669 [Pancotti et al., 2022] , cDNA [Tsuboyama et al., 2023] , and SKEMPI [Jankauskaitė et al., 2019] and demonstrate substantial improvement over baseline methods.Our contributions are as follows:• We develop Structure Motif Embedding Database (SMEDB), a comprehensive local structure alignment database encompassing all structures from the Protein Data Bank (PDB), which is organized to accelerate GPU-based kNN search.• We propose Multiple Structure Motif Alignment (MSMA) to retrieve local coevolutionary motifs based on embeddings of protein structure encoders, which is shown to be complementary to classical sequence-level retrievals.• Our model, Multi-Structure Motif Invariant Point Attention (MSM-IPA), pretrained on our novel database and retrieval mechanism, demonstrates superior performance in predicting ∆∆G, surpassing other models on benchmark datasets, i.e., S669 and SKEMPI. 
preliminaries: 
 Definitions A protein consists of multiple residues, possibly from different chains. For each residue i, we represent the residue as its residue type a i ∈ {1, . . . , 20} and its positions of backbone heavy atoms with p i,C , p i,Cα , p i,N ∈ R 3 . A structure motif M is a fragment of the whole protein structure with N residues. For the raw structure motif M raw , we choose N raw = 256 to match the pretraining settings and keep enough information about raw structure. For the retrieved structure motif (denoted as {M i } L i=1 , where L is the size of the retrieved motif set), we choose N retr = 16 to ensure sufficient interactions from similar environments are captured.Our approach uses ESM-IF as a pretrain model and CUHNSW for vector database implementation.Here we make a brief introduction to these two methods. [Hsu et al., 2022] is a model designed for protein sequence prediction using backbone structures. This approach treats inverse folding as a sequence-to-sequence problem with an autoregressive encoder-decoder architecture, allowing the model to recover native sequences from backbone atom coordinates. By incorporating a large number of sequences with predicted structures as additional training data, ESM-IF effectively learns even when experimental structures are unavailable. The augmented data and backbone-only input make ESM-IF a suitable model to encode the nearby backbone structure of a residue.CUHNSW CUHNSW [Malkov and Yashunin, 2020] is a CUDA implementation of Hierarchical NSW (HNSW), a novel approach of vector database that can do approximate K-nearest neighbor (K-NN) search utilizing navigable small world graphs with a controllable hierarchy. Unlike traditional methods, CUHNSW is fully graph-based and eliminates the need for additional search structures. It incrementally builds a multi-layer structure of proximity graphs, with elements randomly assigned to layers using an exponentially decaying probability distribution. This method enhances performance by starting the search from upper layers and leveraging scale separation, resulting in logarithmic complexity scaling. Performance evaluations show that CUHNSW outperforms previous state-ofthe-art vector-only approaches and its similarity to the skip list structure allows for straightforward distributed implementation.  
multiple structure motif alignment: 
 We use per-residue embeddings extracted by ESM-IF model [Hsu et al., 2022] for each of the protein chains in PDB database and perform local structure search by HNSW. To obtain embeddings for ESM-IF, we utilize the encoder module named GVPTransformerEncoder in ESM-IF to generate a 512-dimensional embedding for each residue. As we exclusively use the encoder module of ESM-IF, the computational cost remains manageable, requiring approximately 3 days on 32 A100 GPUs.The PDB dataset comprises more than 130 million residues, making it impractical to load the entire dataset into memory for querying. Traditional databases lack the capability to leverage GPU acceleration for efficiently querying the top k nearest neighbors. To address this, we integrate CUHNSW as a module to obtain the k-nearest neighbors in the ESM-IF embedding space. The structure motifs are retrieved for their similar local interactions between residues that have similar backbone positions. The interaction between residues occurs only when they are in close proximity. Therefore, when processing the retrieved structure motifs, a large motif size is not necessary. For simplicity, in this paper, we always use N retr = 16. With a highly efficient CUDA implementation [Yoon, 2021] the time consumption querying the top 10 5 neighbors of a result is about 8 seconds in 8 A100 GPUs and only 0.5 seconds for the top 10 3 neighbors.The advantage of using an inverse folding model as a structure encoder, compared to traditional methods that rely on retrieving from a database based on continuous structural fragments, lies in the model's inherent ability to generate embeddings tailored for predicting the surrounding environment based on a given backbone structure. The inverse folding model encodes positions within the spatial context that significantly influence the central amino acid type, rather than merely contiguous positions along the sequence.Since we use embeddings from a structure encoder (ESM-IF in this case) instead of directly search on structures, we would like to validate the performance of the embedding retrieval method by checking the matched motif size N matched of the search result comparing with that of structure search. We first define the matched motif size asN matched = i∈Nraw 1 min j∈Nretr (dist(R raw,i , R motif,j )) < 2.0Å(1)The left panel of Figure 2 illustrates the distribution of alignment sizes for the retrieved structures, demonstrating that numerous similar spatial structures can be found centered on each position. The other panel of Figure 2 shows the relationship between local similarity (number of matching amino acids) and overall similarity (TM-score) with different motif size. The figure show that the similar motif may have low TM-score, indicating that our search approach can extract analogous local structures motif from structurally unrelated proteins to aid in prediction.Figure 2 : CUHNSW combined with embeddings from ESM-IF successfully retrieved motifs with large matched regions, and possibly motifs distant in sequence identity, where matched residues are defined by the indicator function in Eq. (1). Moreover, lower TM-Score is obtained for motifs with high number of locally matched residues, indicating that our search approach can extract analogous local structures motif from structurally unrelated proteins.In this section, we introduce how we utilize the retrieved structure motifs to help predict the effects of mutations. We begin by discussing the filtering of retrieved structure motifs, followed by the superposition of retrieved data, which enables the model to learn information from diverse structures. Finally, we introduce our module, MSM-IPA, which extracts similarity information from the PDB database.Retrieved Structure Motif Filter After extraction, we obtain a large volume of neighbor data (on the order of 10 3 ). For efficiency reasons, we only reserve L filter = 16 data points with most valuable information. This is done in the following manner: Firstly, retrieved structure motifs with central amino acid different from query are discarded. After that, We rank them with a scoring function and retain the top results. The scoring formula can be expressed as follows:i∈Nraw 1 min j∈Nretr (dist(R raw,i , R motif,j )) < 2.0Å • exp(-∥p i,Cα -p 0,Cα ∥ 2 ) (2)where 1[•] is the indicator function and the distance function dist(R 1 , R 2 ) = ∥p 1,Cα -p 2,Cα ∥ 2 . Intuitively, we use the weighted term exp(-∥p i,Cα -p 0,Cα ∥ 2 ) to rewrite Eq. (1) to retain as many contacts around the central amino acid as possible.When dealing with multiple mutations, which involves several central amino acids, we ensure a balanced selection of retrieved structure motif information by using each central amino acid as a query. This approach guarantees an even distribution of the retrieved information across all central amino acids.In the context of mutation analysis, we conduct separate selection processes for both the pre-mutation (wild type) and post-mutation states. The key distinction between these two processes is the amino acid type in the initial step-one being the wild type and the other the mutated form.Superposition The choice of superimposition method depends on the nature of the extracted information. We compare two superimposition approaches: alignment based on the central frame and alignment based on the overall structure. Experimental results indicate that the central frame-based alignment method provides better alignment quality for the retrieved structures. Consequently, we select the central frame-based alignment method for subsequent analyses.For each residue, we can construct the frame from its backbone atom positionsp C , p Cα , p N . v N,Cα = (p N -p Cα )/∥p N -p Cα ∥ v C,Cα = (p C -p Cα )/∥p C -p Cα ∥ R = [v N,Cα , v C,Cα , v N,Cα × v C,Cα ] t = p Cα ,where × is the cross product between two vectors.Hence with the given raw structure motif M raw and any other structure motif M retr , we firstly take out the central residue R C raw and R C retr . Then we can calculate the frame of both the residues, called (R C , t C ) and (R retr , t retr ) respectively. Then we can define the alignment function as:F align-atom (p) = R C R T retr (p -t retr ) + t C .The alignment procedure involves applying the alignment function to each atom of M retr to obtain the aligned motif M aligned retr . For simplicity, in the following sections, we will omit the term "aligned" and assume that all motifs are aligned with the raw motif by default.MSM-IPA To retrieve information from structure motifs, we propose Multi-Structure Motif Invariant Point Attention (MSM-IPA), which is inspired by and similar to the IPA module in Al-phaFold2 [Jumper et al., 2021] . The MSM-IPA module takes the raw structure motif M raw and the processed retrieved structure motifs M 1 , . . . , M L as input. For simplicity, we merge the raw motif with the retrieved motifs into a single motif, denoted as M comb = {r ∈ M i , ∀i ∈ {1, . . . , L}}∪ {r ∈ M raw }. M comb is then used as key motifs in the cross attention mechanism. The size of this combined motif is represented as N comb .To extract information from the motif, we define two encoders, Enc s (M) for single node-wise representations and Enc z (M 1 , M 2 ) for pair edge-wise representations. The single encoder encodes the residue types and positions of atoms into a single representation, s. The pair encoder encodes information such as relative positions in the sequence, spatial relative positions, and amino acid type pairs into a pair representation, z.In MSM-IPA, we extract scalar and vector representations from s. The attention weights are updated and aggregated from three sources of information: single, pair, and predicted points. This approach ensures that the model efficiently utilizes the information while maintaining invariance to overall rotation and translation. The updated single representation of the raw structure motif is obtained by concatenating the weighted sums of each data representation. The detailed implementation of this algorithm is shown in Algorithm 1.The MSM-Mut model mainly consists of two parts: a series of MSM-IPA to fusion the information and a mutation effect predictor consisted of a series of MLP. For both the wild-type and mutated structures, we retrieve structure motifs independently, merge the information using MSM-IPA, and model it with IPA to obtain the single representations for both structures. These representations are then fed into the mutation effect predictor to get the final impact on the structure.Our model training comprises two phases. In the pretraining phase, we utilize a meticulously curated dataset from the Protein Data Bank-REDO (PDB-REDO) [Joosten et al., 2014] for our pretraining data. The dataset is split into training, validation, and test sets in a ratio of 95%:0.5%:4.5%. The pretraining process involves an initial 200,000 steps without the inclusion of retrieved structure motifs, followed by an additional 30,000 steps incorporating retrieved structure motifs. During training, a random amino acid was selected, and its 256 nearest amino acids are extracted with their amino acid types and backbone atom positions. The type of the central amino acid was masked since the model was tasked with predicting it.In the finetuning phase, to ensure comparability of our model, we follow the settings of StableOracle and RDE. For the stability mutation effect prediction task, we use the CDNA120K [Tsuboyama et al., 2023] subset of the training set, which was deduplicated against s669 in StableOracle. For the PPI surface mutation effect prediction task, we perform 3-fold cross-validation on the SKEMPI dataset, partitioned by PDBID. 
results: 
 In this section, we show that our model outperforms other models on a series of tasks with the retrieved MSM information. To begin with, in section 4.1.1, we demonstrate that the retrieved information provides an advantage in predicting mutation effects on protein-protein interfaces. Subsequently, in section 4.1.2, we present a case study on antibody engineering for SARS-CoV-2, illustrating our model's strong performance on real-world targets. In addition, in section 4.2.1, we conducted an experiment to show that our model outperforms others on the protein stability change dataset. To further validate our model's capability in stability prediction, in section 4.2.2, we tested it on a novel enzyme dTM dataset and achieved excellent results.The baselines referenced in this study encompass two main categories: unsupervised models, and semi-supervised/supervised models. The unsupervised models comprise PSSM (position-specific scoring matrix), ESM-1v [Meier et al., 2021] , MSA Transformer [Rao et al., 2021] , Tranception [Notin et al., 2022] , ESM-IF [Hsu et al., 2022] , ProteinMPNN [Dauparas et al., 2022] , Rosetta [Park et al., 2016 , Alford et al., 2017] , and FoldX [Delgado et al., 2019] . The semi-supervised category includes MIF-Net [Yang et al., 2023] , RDE-Net [Luo et al., 2023] , and DiffAffinity [Liu et al., 2024] , while supervised models encompass DDGPred [Shan et al., 2022] and the End-to-End method.Evaluation metrics In this section, we employ six metrics to evaluate our model: Pearson and Spearman correlations to quantify overall trends, per-structure metrics to assess the model's ability to identify beneficial and detrimental mutations within each structure, and RMSE and MAE to measure numerical differences in predicted energy values. Among these, the most critical metric is the per-structure correlation, as practical applications often involve ranking different mutations within the same structure.Following the settings in RDE [Luo et al., 2023] , we partitioned the dataset into three folds based on PDB IDs. Two of these folds were further divided into training and validation sets in a 95:5 ratio based on PDB IDs, while the remaining fold was used as the test set. By training in this manner, we obtained three distinct sets of model parameters, each corresponding to the performance measured on the entire dataset used as the test set.This experiment primarily demonstrates two key points. Firstly, we establish a simple baseline that directly utilizes the top 100 retrieved structure motifs obtained from our search to construct a profile, i.e., simple statistics of the amino-acid types according to the retrieval results. This baseline, called MSM-profile, achieves significantly better performance compared to MSA, and even surpass the majority of unsupervised language models. Note that MSM-profile is built upon the embeddings obtained from the ESM-IF model, and we observed that the per-structure Pearson and Spearman correlation coefficients extracted directly are comparable to those predicted by the ESM-IF model. This indicates that our search method effectively retains the information from ESM-IF, allowing us to explicitly extract local structure motifs corresponding to the ESM-IF distribution.Secondly, the semi-supervised model trained with the enhanced information achieved state-of-the-art performance across almost all metrics. Among the metrics, the per-structure correlation is the most crucial, as it reflects the model's ranking capability for mutations on PPI surfaces. We found that the use of the information we extracted resulted in a significant improvement in the model's prediction correlation, underscoring the value of the retrieved information. In Shan et al. [2022] , five single mutations are identified that enhance the neutralization effectiveness of antibodies against SARS-CoV-2. Within the three CDR regions in the heavy chain, spanning a total of 26 positions, there are 494 possible mutations. Our task is to select these five beneficial mutations from the pool. We compare our results with a subset of the baseline methods that performed well in Section 4.3. Experimental results indicate that our model excels in predicting the ranking of RH103M, while also maintaining high ranks for the other four mutations. This improvement is attributed to our model's ability to extract similar structure motifs from the database. Case Study: RH103M Given that this task focuses on antibody optimization, we incorporated a post-processing step in the retrieval process. We prioritize those retrieved data corresponding to an antibody or nanobody and if the relevant position was located in a loop region. At position H103, out of the top 1000 candidates, only 5 structure motifs remained after filtering, with one having the central residue type as methionine (M). Experimental results indicated that including this motif significantly influenced the outcome. Analysis revealed that this structure corresponds to residue 576 of chain T in the 5KOV structure. The aligned structures are depicted in Figure 3 . This demonstrates that our data retrieval method can effectively assist in extracting relevant local structures for antibody design, potentially offering an interpretable approach to the antibody engineering process.Training MSM-Mut(or other model name) for protein engineering To train the model's affinity prediction module, we fine-tuned the stability prediction component using the cDNA dataset. To fully utilize the mutation data and reduce bias, ensuring a more balanced representation of mutation types, we followed the settings described in Stability Oracle . Following Diaz et al.[2023], we use Thermodynamic Permutations (TP) to perform data augmentation to balance the mutation type distribution of dataset cDNA120K [Tsuboyama et al., 2023] and train on 2M thermodynamically valid ddG measurements. TP achieves this by exploiting properties of Gibbs free energy and dataset characteristics. Specifically, the original 19 mutations at each position are augmented to cover all 380 (19 × 20) pairwise mutations.In this study, we utilized a newly curated dataset (S669) [Pancotti et al., 2022] derived from the latest version of ThermoMutDB. This dataset comprises 669 protein variants not found in commonly used training sets, offering a robust basis for evaluating prediction models. Although our model does not match Stability Oracle in predicting stability changes, the integration of retrieved structure motifs enabled our model to surpass the such methods, thereby establishing a new state-of-the-art in this task.To demonstrate our model's robust generalization capability on new data, we tested it on a novel enzyme thermostability dataset provided by Novozymes [Pultz et al., 2022] . This dataset includes experimental measurements of melting temperature of point mutations on a novel enzyme sequence that has no high-similarity match in PDB database with sequence identity higher than 30%. An AlphaFold prediction of wild-type protein structure is released with the dataset as the reference to perform structure-based methods. We adopt this predicted structure to retrieve local motifs and feed to MSM-Mut. As the results shown in Table 3b , our method significantly outperforms both classical force fields and machine-learning-based baselines. It suggests that our method can be applied to novel proteins and predicted structures. Mutation effect prediction Mutation effect prediction plays a pivotal role in in-silico protein engineering. There are three major categories of methods for mutation effect predictionbiophysical [Schymkowitz et al., 2005 , Park et al., 2016 , Alford et al., 2017 , Steinbrecher et al., 2017] , statistical [Geng et al., 2019 , Zhang et al., 2020] , and deep learning-based methods [Rao et al., 2021 , Liu et al., 2021 , Shan et al., 2022 , Yang et al., 2023 , Luo et al., 2023] .Different information is provided as inputs to different deep learning models. Models based on protein language models (PLMs) usually require only the primary sequence of both wild type and mutation [Meier et al., 2021 , Notin et al., 2022] . MSAs also serve as inputs to mutation prediction models [Hopf et al., 2017 , Riesselman et al., 2018 , Rao et al., 2021 , Luo et al., 2021 , Frazer et al., 2021] based on the observation that co-evolutionary information correlates with the information needed. In addition to sequences, structures of protein also provide additional information.Along this line, a number of models are pretrained on protein structure data to encode the structure information [Hsu et al., 2022 , Yang et al., 2023 , Zhang et al., 2023] . However, MLStrA-IPA is, to our knowledge, the first model that employs multiple structures for single mutation effect prediction.Retrieval-based deep learning Retrieval-based deep learning has made its splash in language modeling recently [Guu et al., 2020 , Wu et al., 2022b ] that achieves higher performance-per-parameter.Retrieval-based predictions have been long employed in protein structure prediction where people searches for evolutionarily related sequences of the sequence of interest and gather them into one MSA input to a model, which has been an important building block for models including AlphaFold [Jumper et al., 2021] and RoseTTAFold [Baek et al., 2021] , while replacing MSA modules with language model demonstrates a small but noticeable drop in performance [Wu et al., 2022a , Lin et al., 2023] .Retrieval mechanism other than MSA has also been investigated previously. Retrieved Sequence Augmentation [Ma et al., 2023] augments PLMs with a collection of related sequences, either in sequence or in structure, and claims substantial improvements over MSA Transformer [Rao et al., 2021] both in performance and in prediction speed. Wang et al. [2023] use retrieval techniques to improve the performance of controllable molecule generation. 
open access to data and code: 
 Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: All datasets we used in evaluation are open-sourced, and we will release our model upon acceptance.Guidelines:• The answer NA means that paper does not include experiments requiring code. • At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable). • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
conclusion: 
 In this study, we addressed the challenge of predicting the effects of protein mutations by introducing a novel retrieval-augmented framework that leverages local structure motifs. By creating the Structure Motif Embedding Database (SMEDB) and using Multiple Structure Motif Alignment (MSMA), we efficiently retrieved and utilized local coevolutionary information. Our Multi-Structure Motif Invariant Point Attention (MSM-IPA) model demonstrated superior performance on benchmark datasets, proving the value of focusing on local structural environments. Our model offers a scalable and robust method for studying protein mutations, with important implications for protein engineering and genetic disease research.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: All details of training and evaluation procedures are included in Appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
acknowledgements: 
 This work is supported by the National Key RD Program of China No.2021YFF1201600Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: We plot the complete distribution in Figure 2 to justify the design of our method.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: We state the amount of computation resources to use when we introducing our method in section 3.Guidelines:• The answer NA means that the paper does not include experiments. Justification: As we discussed in the introduction section, our work may have positive societal impact since the proposed method can be applied to drug discovery and enzyme engineering.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)? Answer: [NA] Justification: This paper does not involve crowdsourcing nor research with human subjects. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: The code is available at https://github.com/guoruihan/MSM-MutThe conservation of Multiple Sequence Alignment (MSA) in sequences is relatively strong, but when dealing with surface mutation data and predicting the impact of mutations on binding affinity, sequence-based profiles tend to be weaker. Our MSM-Profile addresses this limitation by leveraging local structure motifs, which can share information regardless of whether they are intra-chain or inter-chain. In this study, we use ESM-IF embeddings to construct a database and retrieve Multi-Structure-Motifs (MSM) to support mutation effect predictions. We demonstrate that using only the top-1 neighbor provides interpretability benefits, as retrieval with ESM-IF embeddings effectively identifies information valuable for mutation effect prediction.To assess how varying the number of retrieved neighbors impacts predictive performance, we conducted experiments on the s669 dataset. Results indicate that increasing the number of neighbors improves the model's predictive ability. This finding suggests that the top-ranked neighbors contain diverse, biologically relevant information that enhances prediction accuracy.The table below presents an ablation study comparing the performance differences when retrieval is omitted, pre-training is omitted, or both are omitted. We observed that pre-training an IPA module is crucial. Without pre-training, the model lacks a proper initial distribution for the 20 types of amino acids at masked positions, which negatively impacts subsequent tasks. To illustrate the impact of retrieval dataset size on model performance, we present results based on random selections of approximately one-tenth, one-hundredth, and one-thousandth of the database. Specifically, we randomly selected 16 motifs from the top 100, 1000, and 10,000 motifs, simulating a reduction of the database to approximately one-tenth, one-hundredth, and one-thousandth of its original size. Our observations indicate that within the existing high-quality structure database, the size of the dataset is indeed critical, as it significantly influences the availability of high-quality data. To provide a better comparison, we implemented a simpler encoding method called Continuous Backbone Angle Embedding (CBAE). For this approach, we defined the ϕ and ψ angles corresponding to amino acid i as ϕ i and ψ i . For each amino acid, we created an 18-dimensional embedding vector[ϕ i-4 , ψ i-4 , . . . , ϕ i+4 , ψ i+4 ],where each value ranges from [-π, π] . This vector consists of the angles for the amino acid itself and the four consecutive amino acids before and after it in the sequence. We calculated the distance between two embeddings using the Manhattan distance, ensuring that retrieved structures have similar local sequence backbone structures.Using this retrieval method, we tested model performance on the SKEMPI2.0 and S669 datasets, as shown in Tables 9 and 10 . Results indicate that, while structures retrieved using ESM-IF embeddings tend to perform better, CBAE still maintains reasonable backbone similarity.Our experiments showed that structures retrieved using ESM-IF embeddings tend to perform better than CBAE in capturing relevant mutation information, as ESM-IF embeddings can capture nonadjacent backbone atom positions to some extent. However, the differences between these methods To further illustrate the retrieval process, we provide a figure with two subfigures detailing the alignment and visualization of highly similar local antibody structures obtained through Multiple Structure Motif Alignment. Ensure: Merged motif representation si 1: Combine raw and retrieved motifs:M comb = {residue ∈ M i , ∀i ∈ {1, . . . , L}} ∪ {residue ∈ M raw } 2:Encode single representations:s raw = Enc s (M raw ) s comb = Enc s (M comb ) 3: Encode pair representation: z = Enc z (M raw , M comb )4: Linear transformations:q h i , ⃗ q h i = LinearNoBias(s raw ) k h i , v h i , ⃗ k h i , ⃗ v h i = LinearNoBias(s comb ) b h ij = LinearNoBias(z ij ) 5:Compute attention weights:a h ij = softmax k w L 1 √ c q h ⊤ i k h j + b h ij - γ h w C 2 p T i • - → q hp i -T j • - → k hp j 2 6: Compute outputs: õh i = j a h ij z ij o h i = j a h ij v h j - → o hp i = T -1 i • j a h ij T j • - → v hp j 7:Compute final representation:s updated-raw = Linear concat h,p õh i , o h i , - → o hp i , - → o hp i• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.38th Conference on Neural Information Processing Systems (NeurIPS 2024). 


Decision: This paper presents a retrieval-augmented framework for protein mutation prediction. Structural motifs are first retrieved according to the structural similarity with the query structural, measured by a protein structure encoder. Experimental results on multiple protein mutation prediction tasks prove the effectiveness of the proposed approach. 

Overall, leveraging structural similar motif for mutation prediction is a very interesting and novel idea. Increasing the interpretability of the proposed approach and also validate the model on more benchmarks (e.g., ProGym) will further strengthen the paper.

Reviews: review: paper_summary: paper_summary: - The paper presents a novel retrieval-augmented framework to efficiently retrieve similar local structure motifs in protein sequences for mutation effect prediction
- Current methods to understand coevolutionary patterns include MSA and domain-level structure clustering, which serve as a global representation of the protein and its evolutionary couplings
- The paper introduces Structure Motif Embedding Database (SMEDB), constructed from the ESM-IF structure-based embedding approach, to enable rapid GPU-accelerated kNN search and retrieve local structure motifs similar to those affected by mutations
- The retrieval method Multiple Structure Motif Alignment (MSMA) leverages embeddings that capture local coevolutionary patterns
Multi-Structure Motif Invariant Point Attention (MSM-IPA) model aggregates retrieved coevolutionary motif information to predict changes in binding free energy (G) on protein surfaces to assess protein-protein interactions
strength_weakness: strength_weakness: - The main weakness of the work is in the evaluations.
- It is difficult from Table 2 to conclude the competitive performance of MSM-Mut as other DL-based methods are performing better?
- In Table 2, I did not understand the rationale for making some numbers bold. If the lowest number is better, then why are all numbers bold?
- It would make reading the paper much easier if Tables were referred to in the paper.
- Several important mutation prediction baselines, such as DeepSeqeice and EVE, are missing from Table 1. These and others are outlined and now standardized in the ProteinGym benchmark paper.
- Without a large-scale evaluation, it would be hard to assess the utility of the method and fully evaluate the contributions 
- The interpretability of MSM-IPA’s predictions might be obscured because it incorporates multiple structure motifs
- Figure 3 could benefit from some labels for the reader to identify T chain 576, 5KOV, H chain 103, and 7FAE
questions: questions: 1. Given that ESM-IF is trained on "global" protein structure and not local, why is it intuitive to use ESM-IF as a tool to embed local structures?
2. Would a stand-alone embedding module trained entirely on local structures perform better as an alternative to ESM-IF within the MSM-IPA framework?
limitations: limitations: As discussed earlier in the weakness section, the main limitation of the paper is evaluation which currently do not fully support the main claim.
paper_summary: paper_summary: The paper presents a novel retrieval-augmented framework for enhancing the prediction of protein mutation effects, which is essential for analyzing protein functions and understanding genetic diseases. The authors design a system that incorporates similar structure information from known protein structures into the mutation effect prediction process. Central to this framework is the creation of a vector database, the Structure Motif Embedding Database (SMEDB), which stores embeddings of local structure motifs derived from a pre-trained protein structure encoder. This allows for efficient retrieval of similar local structure motifs during the prediction of mutation effects.
strength_weakness: strength_weakness: 1. In the field of bioinformatics, it is crucial not only to make accurate predictions but also to understand the reasons behind certain predictions. Enhancing models with interpretability features can help explain the importance of retrieved motifs for each prediction, which could be valuable.

2. A more thorough comparative analysis with existing methods, especially those that also focus on local structure motifs, could be beneficial for this paper. Expanding the discussion on how the proposed framework differs from and improves upon these methods will strengthen the novelty claims of the paper.
questions: questions: 1. Regarding the Retrieved Structure Motif Filter, could you explain why 16 motifs were chosen and whether this number is sensitive to performance variations?

2. Can the authors elaborate on how the performance and limitations of the ESM-IF and CUHNSW modules impact the overall framework? How were these modules selected, and were alternative options considered?

3. How does the model handle proteins with low sequence similarity or sparse structural motifs?

4. In Equation 2, why is the distance calculated with the 0th alpha carbon atom? Does this represent the alpha carbon atom of the central amino acid?

5. What impact does pre-training have on the predictive ability of the model? How would this method perform without pre-training?
limitations: limitations: The paper acknowledges certain limitations, but it could expand on these by discussing potential weaknesses in greater detail and suggesting more specific directions for future research. This might include exploring different protein structure encoders, expanding the database, or adapting the method for different types of mutations.
paper_summary: paper_summary: This paper presents a novel approach for mutation effect prediction using a retrieval-augmented framework. The main contributions involved: 
* Structure Motif Embedding Database (SMEDB): a vector database storing ESM-IF local structure motif embeddings from experimentally determined protein structures in Protein Data Bank. 
* Multiple Structure Motif Alignment (MSMA): an approach to efficiently retrieve and filter local coevolutionary motifs from SMEDB.
* MSM-IPA: a new model architecture leveraging retrieved coevolutionary motif information to predict the effect of mutations. 
* The approach was validated on multiple benchmarks for mutation effect prediction and a case study on SARS-COV-2 Antibody Optimization.
strength_weakness: strength_weakness: *  **Scalability**: The authors claim their approach is scalable but the paper lacks experimental evidence on the impact of increasing the number of encoded residues in the retrieval database. There's no analysis in increasing/decreasing the database size or including predicted structures from AlphaFoldDB. Predicted structures are used for encoding input proteins in the Novozymes benchmark but not for retrieval. If database size is a significant bottleneck for retrieval, this should be discussed.
*  **Choice of Embedding Method**: The reliance on ESM-IF for embedding structural motifs is not well justified, and alternative methods are not explored. Discussing how to overcome performance limitations imposed by ESM-IF would strengthen the paper.
*  **Exploration of Retrieval Hyper-parameters**: The paper lacks details on hyper-parameter choices for the retrieval mechanism. The retrieved motif length (Nretr=16) and the number of retrieved motifs (Lfilter=16) could significantly impact performance but their effect is not studied. The discussions comparing the proposed approach against MSAs should also consider MSA depth vs the number of retrieved local motifs. 
*  **Applicability to Other Mutation Effect Benchmarks**: I found it surprising that the methodology is not tested on ClinVar or ProteinGym which are common mutation effect prediction benchmarks. It should be discussed if the lack of a training set for these datasets is a limitation.

 **Minor**:
*  MSM-Mut Description: The paper describes the classification head as a series of MLPs but does not specify the number of layers. Including this information would help reproducibility.
*  Training and Fine-tuning Details: Limited information on hyper-parameters and procedures for pre-training and fine-tuning. Including specifics like learning rate, batch size, and training time would help other researchers replicate the work.
*  Extending the Figure 1 legend with more details on the methodology would help understanding.
questions: questions: * Why were predicted structures not considered to extend the database? Could the authors extend their discussion on scalability and/or include analyses on the effect of increasing/decreasing the size of the underlying dataset used for retrieval?
* What was the maximum sequence length used during the encoding of the database and during mutation effect prediction? How did you handle proteins longer than this maximum length?
* Did you consider alternative underlying residue embedding methods such as ProteinMPNN? 
* In the Introduction, you claim that the distribution of local motif embeddings provides complementary information to MSA profiles. Could you clarify the basis for this claim? 
* Why was the method not benchmarked on ClinVar and ProteinGym datasets? Are there specific challenges or limitations that prevented these evaluations?
limitations: limitations: The paper discusses some limitations imposed by third-party modules such as ESM-IF and CUHNSW. However, some areas could benefit from further discussion or stating the limitations on:
* Scalability: Discuss how scalable the method is and any potential challenges or limitations on retrieval when increasing the database size. 
* Benchmarking on Clinvar/ProteinGym: providing reasons for not benchmarking on those datasets and discussing any potential challenges or limitations would give a more comprehensive view of the method's applicability.

Comments: title: Thanks for the notes
More Rebuttal Information
Rebuttal Information
More Rebuttal Information
Rebuttal Information
More Rebuttal Information
comment: Thanks for the detailed rebuttal. I have raised the score.
During this period, we have conducted additional experiments specifically addressing some of the weaknesses you mentioned. We hope these efforts help resolve the concerns raised. We look forward to hearing your feedback.

**1. Exploring different protein structure encoders**

To provide a better comparison, we tried a simpler encoding method: we defined the $\phi$ and $\psi$ angles corresponding to the amino acid i as $\phi_i$ and $\psi_i$. For an amino acid, we defined its embedding as [$\phi_{i-4}, \psi_{i-4}, ..., \phi_i, \psi_i, ..., \phi_{i+4}, \psi_{i+4}$], an 18-dimensional vector where each value ranges from [-pi, pi]. This vector consists of the angles corresponding to this amino acid and the four consecutive amino acids before and after it in the sequence. We defined the distance between two embeddings as their Manhattan distance. This method is referred to as the Continuous Backbone Angle Embedding (CBAE). This embedding method ensures that the retrieved structures have similar local sequence backbone structures. Using such retrieving methods, we tested the performance on SKEMPI and S669 datasets as follows:

Task: SKEMPI 2.0

  | Method               | Pearson (P.S.) | Spearman(P.S.) | Pearson | Spearman | RMSE   | MAE    |
|----------------------|---------|----------|---------|----------|--------|--------|
| MSM-Mut (w/o retrieval) | 0.4325  | 0.4031   | 0.6233  | 0.4954   | 1.6076 | 1.2155 |
 | MSM-Mut (CBAE retrieval) | 0.4619   |  0.4262   |  0.6524   |  0.5158   |  1.5531   |  1.1622 
 MSM-Mut              | 0.4736  | 0.4354   | 0.6814  | 0.5786   | 1.4703 | 1.0212 |

Task: S669

  | Method               | Pearson | RMSE|
|----------------------|---------|----------|
| MSM-Mut (w/o retrieval) | 0.45  | 1.73   | 
| MSM-Mut (CBAE retrieval) | 0.52  | 1.55   | 
| MSM-Mut | 0.54  | 1.51   | 

Our experiments showed that structures retrieved using ESM-IF performed better than this simpler method. ESM-IF can capture non-adjacent backbone atom positions to some extent. However, the differences between these methods are limited because we observed that structures retrieved using ESM-IF embeddings also ensure a high degree of backbone similarity.

**2.Expanding the database**

Given the constraints of time, it is challenging to incorporate the entire set of predicted structures into our analysis. We are currently working on reducing the dimensionality of the embeddings through contrastive learning. However, the size of the entire AFDB dataset is substantial. To illustrate the impact of dataset size on model performance, we present results based on random selections of approximately one-tenth, one-hundredth, and one-thousandth of the database. Specifically, we randomly selected 16 motifs from the top 100, 1000, and 10,000 motifs, simulating a reduction of the database to approximately one-tenth, one-hundredth, and one-thousandth of its original size. Our observations indicate that within the existing high-quality structure database, the size of the dataset is indeed critical, as it significantly influences the availability of high-quality data.

Task:S669

  | Method               | Pearson | RMSE|
|----------------------|---------|----------|
| MSM-Mut (w/o retrieval) | 0.45  | 1.73   | 
| MSM-Mut (random in top 10000) | 0.43  | 2.03   | 
| MSM-Mut (random in top 1000) | 0.52  | 1.57   | 
| MSM-Mut  (random in top 100) | 0.53  | 1.59   | 
| MSM-Mut (top 16 neighbor)| 0.54  | 1.51   | 

**3. Adapting the Method for Different Types of Mutations**

In the manuscript, we have already tested our method on several datasets, including SKEMPI2.0 (general surface mutation ΔΔG prediction), SARS-CoV-2 (antibody binding affinity prediction), S669 (mutation effects on stability), and novo-enzyme mutation effects on stability. These datasets cover a broad range of mutation effect prediction scenarios. If there are additional datasets or settings where you would like to see our method evaluated, we would be more than happy to provide those results.

We have addressed the key weaknesses you pointed out in our paper. If these responses adequately resolve your concerns, we kindly ask you to consider slightly adjusting your evaluation. We are also open to further discussing any remaining issues you might have.
We sincerely appreciate the reviewer’s valuable suggestions. Regarding your concerns about interpretability and the lack of baseline comparisons, we would like to provide the following clarifications:

**Interpretability**: First, we would like to explain why we chose ESM-IF, a global embedding, to model local information. Our motivation stems from the fact that the Protein Data Bank (PDB) is sufficiently large, leading us to believe that local structures of proteins are relatively dense. Thus, we aimed to directly construct a local-structure-based profile that could model the effects of mutations, particularly those on the protein surface, something that conventional MSA profiles cannot achieve. ESM-IF is trained on an inverse folding task, where the internal embeddings for each amino acid align well with our task. Therefore, we hypothesized that the similarity between these embeddings could, to some extent, quantify the similarity between local structures.

Additionally, we have provided evidence of the validity of the information retrieved by our method. In Rebuttal Fig. 2, we display the retrieved structures, showing that the local motifs with the most similar embeddings indeed share highly similar backbone structures.

**Baseline Comparisons**: We have already added several baseline experiments on the SKEMPI and S669 datasets, as mentioned earlier in the rebuttal. We also explained the reasons for not selecting the two large-scale datasets you mentioned. Currently, we are conducting tests of our method and baselines on a sub-dataset of ProteinGym that has no overlap with our fine-tuning training set. We will include these results in the appendix of the manuscript.

We hope that these clarifications address your concerns. If you have any further questions, we would be more than happy to continue the discussion with you.
Thank you very much for your positive feedback and for recognizing the additional experiments and detailed analysis we provided in the rebuttal. We are grateful for your thoughtful suggestions and are pleased to hear that our revisions have addressed your concerns and improved the overall quality of the paper.

We appreciate your recommendation to include the content of this discussion in the main manuscript or Appendix. In response, we will thoughtfully incorporate the key points and additional details from the rebuttal into the Appendix of the manuscript. This will ensure that the insights we discussed are thoroughly documented and accessible to readers, thereby further enhancing the contribution of our work.

Once again, we sincerely thank you for your valuable suggestions and guidance throughout this process.
Thank you for your comprehensive rebuttal. The revisions have addressed various concerns and enhanced the manuscript quality with the additional content. I encourage the authors to add the content of this discussion to the main manuscript or Appendix.
I thank the authors for the notes. My main criticism regarding choosing a global embedding (ESM-IF) to obtain local information is the difficulty of interpretability and lack of comparison with baseline stands. At this point, I cannot be more positive than a score of 5. I hope the comments help improve the paper.
**Question 5: Did you consider alternative underlying residue embedding methods such as ProteinMPNN?**

ESM-IF can potentially be replaced by other methods based on C-alpha point cloud matching or ProteinMPNN. In our future development, we will attempt to train an independent local structure motif encoder to reduce the embedding dimension centered on each amino acid from 512 dimensions to 32 or even lower. This reduction will enable us to incorporate a larger amount of data, including predicted data.

**Question 6: In the Introduction, you claim that the distribution of local motif embeddings provides complementary information to MSA profiles. Could you clarify the basis for this claim?**

The conservation of MSA in sequences is relatively strong, but when dealing with surface mutation data and predicting the impact of mutations on binding affinity, the profile from sequences tends to be weaker. Our MSM-Profile does not have this issue because, whether it is intra-chain or inter-chain, the local structure motifs are similar and can share information. Below are the performance results of our method and the MSA profile on the SKEMPI2.0 dataset. It can be seen that our profile has a natural advantage in this task.

Task: SKEMPI 2.0

| Category       | Method               | Pearson (P.S.) | Spearman(P.S.) | Pearson | Spearman | 
|----------------|----------------------|---------|----------|---------|----------|
| Profile        |MSA-Profile                 | 0.0826  | 0.0822   | 0.0159  | 0.0666   | 
|                | MSM-Profile          | 0.1551  | 0.1766   | 0.1433  | 0.1739   | 
|                | MSM-Profile(Filtered)          | 0.1905  | 0.1886   | 0.1653  | 0.2063   | 

Additionally, we found that on the s669 dataset, although the distributions of our MSM-Profile and MSA-Profile are similar, simply adding the two profiles results in a significant increase in Pearson correlation. This indicates that there is a portion of the information in both profiles that is independent. Therefore, we can enhance current tasks that use MSA by incorporating MSM-Profile, which may improve the model's performance.

Task: S669

  | Method               | Pearson | 
|----------------------|---------|
| MSA-Profile | 0.17  |
| MSM-Profile | 0.19  |
| MSA-Profile + MSM-Profile | 0.23  |

**Question 7: Why was the method not benchmarked on ClinVar and ProteinGym datasets? Are there specific challenges or limitations that prevented these evaluations?**

In this paper, we followed the setting of the Stability Oracle for benchmarking models predicting mutation effects on stability. Within the Stability Oracle test sets, we selected the latest and most challenging dataset, s669, as our benchmark dataset.

During stability fine-tuning, a significant portion of our training set, the cDNA dataset, overlaps with the ProteinGym test set. Removing these overlapping data points may weaken our benchmark results on ProteinGym; hence, we did not conduct tests on this dataset. To demonstrate our model's capability, we show significant improvements on the s669 dataset, which has no overlap with the cDNA dataset. This indicates that our model performs well.

The data on ClinVar is more focused on binary classification, determining whether a mutation is pathogenic. Since both our pre-training and fine-tuning datasets do not match the ClinVar setting, it is challenging to compare our model's performance on the ClinVar dataset. In the future, we will attempt to create a new dataset based on our current model. However, our method does not consider sequence co-evolution information (MSA), so a more valuable comparison would be between MSA Profile and our MSM-Profile. We plan to conduct this experiment in future work.
**Question 1:MSM-Mut Description: The paper describes the classification head as a series of MLPs but does not specify the number of layers. Including this information would help reproducibility.**

To predict the mutation effect, we first pass the information before and after the mutation, along with the corresponding retrieved motifs, through an MSM-IPA to obtain embeddings for the mutation site. Then, we compute the difference between the features before and after the mutation. To ensure the symmetry of the mutation (i.e., the ΔΔG of type 1 to type 2 should be the opposite of type 2 to type 1), we concatenate $feat_{raw} - feat_{mutate}$​ and $feat_{mutate} - feat_{raw}$ ​, and then pass them through a four-layer MLP to predict the mutation result.

**Question 2: Training and Fine-tuning Details: Limited information on hyper-parameters and procedures for pre-training and fine-tuning. Including specifics like learning rate, batch size, and training time would help other researchers replicate the work.**

During the pre-training on the PDB database, we used the PDB database up to December 31, 2021. In this process, we used the AdamW optimizer with a learning rate of 3e-4, beta1=0.9, beta2=0.999, and a batch size of 128. The pre-training process was conducted on 8 A100 GPUs over five days (300k steps).

For the cDNA fine-tuning, we followed the same optimizer and parameters. The fine-tuning process took approximately two hours (4000 steps) on 8 A100 GPUs to achieve the best performance. Training for a longer period led to overfitting on the cDNA dataset, which negatively affected the model's generalizability.

**Question 3: Extending the Figure 1 legend with more details on the methodology would help understanding.**

Subgraph (a) illustrates the process of Multiple Structure Motif Alignment. First, we use ESM-IF to obtain embeddings for all amino acids in the PDB database. The size of this library is approximately 500 GB, making it impractical to use standard retrieval methods. To improve retrieval efficiency, we index this data using CUHNSW, constructing a Motif Embedding Database. For each query, we similarly use ESM-IF to obtain embeddings for the query structure, then quickly retrieve the approximate top-k nearest neighbors using HNSW based on the corresponding embeddings. After filtering and superposition, we obtain structure motifs suitable for downstream tasks. For efficiency, we have pre-searched the retrieved motifs from the entire PDB, which also facilitates future tasks.

Subgraph (b) illustrates how MSM-IPA uses retrieved structure motifs to enhance the prediction of mutation effects. In the upper part of this figure, the retrieved structure motifs are aligned with the central amino acids of the query structure motif through preprocessing. These features are concatenated and input into MSM-IPA for cross-attention, thereby enhancing the model's ability to predict mutation effects.

**Question 4: Why were predicted structures not considered to extend the database? Could the authors extend their discussion on scalability and/or include analyses on the effect of increasing/decreasing the size of the underlying dataset used for retrieval?**

There are two considerations for not including predicted structures: quality and efficiency. In terms of quality, we assume that the local structure motifs in the RCSB PDB database are already sufficiently dense, and additional local motifs from predicted structures are also learned from the existing database. In terms of efficiency, the embeddings we currently use with ESM-IF have 512 dimensions, meaning the computational load for retrieval is quite large. Therefore, it is not feasible to increase the dataset from the existing 200k structures to 60M. However, this approach can be attempted in the future by modifying the embedding method and reducing the dimensionality. If the data balance can be managed well, theoretically, better results could be achieved.

Here, we present results using a random selection of one-tenth, one-hundredth, and one-thousandth of the database. By randomly selecting 16 motifs from the top 100, 1000, and 10000, we simulate reducing the database to approximately one-tenth, one-hundredth, and one-thousandth of its original size. We observe that in the existing high-quality structure database, the size of the database is quite important as it determines the quantity of high-quality data.

Task:S669

  | Method               | Pearson | RMSE|
|----------------------|---------|----------|
| MSM-Mut (w/o retrieval) | 0.45  | 1.73   | 
| MSM-Mut (random in top 10000) | 0.43  | 2.03   | 
| MSM-Mut (random in top 1000) | 0.52  | 1.57   | 
| MSM-Mut  (random in top 100) | 0.53  | 1.59   | 
| MSM-Mut (top 16 neighbor)| 0.54  | 1.51   |
**Question 4: Can the authors elaborate on how the performance and limitations of the ESM-IF and CUHNSW modules impact the overall framework? How were these modules selected, and were alternative options considered?**

**ESM-IF:**

ESM-IF uses the GVP-Transformer to embed all backbone atoms, and this embedding is ultimately used for amino acid type prediction, which is why we chose this approach. We observed that structures retrieved using ESM-IF embeddings mostly resemble the backbone frames of amino acids that are close in the sequence. Please refer to the  figure Rebuttal Fig. 2.

One issue with ESM-IF is that, although we observed that the neighbors retrieved by ESM-IF have high local backbone similarity, directly using the distance between embeddings for retrieval is a method not defined during training. Therefore, in the future, we plan to train an independent local structure motif encoder using contrastive learning to give this distance metric a clearer meaning.
ESM-IF can potentially be replaced by other methods based on C-alpha point cloud matching or ProteinMPNN. In our future development, we will attempt to train an independent local structure motif encoder to reduce the embedding dimension centered on each amino acid from 512 dimensions to 32 or even lower. This reduction will enable us to incorporate a larger amount of data, including predicted data.

**CUHNSW:**

CUHNSW is an implementation of the classic HNSW (Hierarchical Navigable Small World) algorithm on GPUs, designed to perform approximate nearest neighbor search through a pre-built graph. This approach leverages GPU acceleration to enhance the performance of the search process. CUHNSW can be replaced by other scalable and parallelizable approximate nearest neighbor algorithms, such as LSH (Locality-Sensitive Hashing), KD-trees, or similar methods, depending on the specific task requirements (such as embedding dimensions), database size, and available resources.

**Question 5: How does the model handle proteins with low sequence similarity or sparse structural motifs?**

In the worst-case scenario, the model may converge to a state with no neighbors, relying entirely on information obtained from pretraining without any retrieval augmentation. We have observed that such data accounts for a relatively small proportion in real PDB databases. However, to enhance the model's capabilities, we will incorporate high-quality predicted structures (suck as AFDB) into the database. To reduce the overall database size, we will perform retrieval for each newly added motif and only add those motifs that do not have very similar neighbors in the database.

**Question 6: In Equation 2, why is the distance calculated with the 0th alpha carbon atom? Does this represent the alpha carbon atom of the central amino acid?**

The 0th amino acid represents the alpha carbon atom of the central amino acid. In Equation 2, our formula consists of two parts. The first part $1\left[ \min_{j \in N_{\text{retr}}}(\text{dist}(R_{\text{raw}, i}, R_{\text{motif}, j})) < 2.0\text{A}\right]$ attempts to find the closest match for each amino acid in the query motif within the retrieved motif. The second part $\exp(-\|p_{i, C\alpha} - p_{0, C\alpha}\|_2)$ assigns a weight to each amino acid, with the weight decreasing as the distance increases.



**Question 7: What impact does pre-training have on the predictive ability of the model? How would this method perform without pre-training?**

The table below presents an ablation study comparing the performance differences when retrieval is omitted, pre-training is omitted, or both are omitted. We observed that pre-training an IPA on PDB is crucial. Without pre-training, the model lacks a proper initial distribution for the 20 types of amino acids at masked positions, which negatively impacts subsequent tasks.

Task: S669

  | Method               | Pearson | RMSE|
|----------------------|---------|----------|
| MSM-Mut (w/o retrieval, w/o pretrain) | 0.37  | 2.82   | 
| MSM-Mut (w/o pretrain, 16 neighbor) | 0.43  | 2.15   | 
| MSM-Mut (w/o retrieval) | 0.45  | 1.73   | 
| MSM-Mut (16 neighbor)| 0.54  | 1.51   |

**Reference**

[1] Runze Dong, Zhenling Peng, Yang Zhang, Jianyi Yang, mTM-align: an algorithm for fast and accurate multiple protein structure alignment, Bioinformatics, Volume 34, Issue 10, May 2018, Pages 1719–1725

[2] Sieg J, Rarey M. Searching similar local 3D micro-environments in protein structure databases with MicroMiner[J]. Briefings in Bioinformatics, 2023, 24(6): bbad357.
**Question 1: In the field of bioinformatics, it is crucial not only to make accurate predictions but also to understand the reasons behind certain predictions. Enhancing models with interpretability features can help explain the importance of retrieved motifs for each prediction, which could be valuable.**

Thank you for your valuable suggestion. In our paper, we primarily used ESM-IF embedding to create a database and then searched for Multi-Structure-Motifs (MSM) to help predict mutation effects.

From the perspective of interpretability, we have demonstrated that using only the top-1 neighbor is beneficial, indicating that retrieving with ESM-IF embeddings can indeed find information useful for predicting mutation effects. We also trained models with different numbers of neighbors and tested them on the s669 dataset. We found that as the number of neighbors increased, the model's predictive ability improved. This suggests that the top-ranked neighbors contain diverse information that can help the model make better predictions. Based on this conclusion, we can infer that the retrieved structure motifs are biologically relevant to mutations and can serve as an effective data source.

Task: S669

  | Method               | Pearson | RMSE|
|----------------------|---------|----------|
| MSM-Mut (w/o retrieval) | 0.45  | 1.73   | 
| MSM-Mut (1 neighbor) | 0.49  | 1.62   | 
| MSM-Mut (2 neighbor) | 0.51  | 1.57   | 
| MSM-Mut (4 neighbor) | 0.53  | 1.53   | 
| MSM-Mut (8 neighbor) | 0.53  | 1.55   | 
| MSM-Mut (16 neighbor)| 0.54  | 1.51   | 
| MSM-Mut (32 neighbor)| 0.54  | 1.52   | 
| MSM-Mut (1024 neighbor)| 0.51  | 1.63   |

**Question 2: A more thorough comparative analysis with existing methods, especially those that also focus on local structure motifs, could be beneficial for this paper. Expanding the discussion on how the proposed framework differs from and improves upon these methods will strengthen the novelty claims of the paper.**

Of traditional methods, mTM-Align[1] introduced the concept of Multiple Structure Alignment (MStrA). As the name suggests, this method excels in traditional template searches as TM-Align. mTM-align is an extension of the pairwise structure alignment program TM-align. However, this method cannot handle local structure information.

The most recent work on extracting local structure motifs is MicroMiner[2]. Searching similar local 3D micro-environments in protein structure databases with MicroMiner]. MicroMiner converts protein sequences into multiple k-mers and uses a method similar to MMseqs2 for initial filtering of k-mers from both the original and mutated sequences in the database. Finally, candidates are superimposed and structurally filtered.

In comparison, our method has two main advantages: First, MicroMiner heavily relies on sequence similarity, which can be too strict. Our method can identify local structure motifs with similar backbone structures but different sequences, greatly enriching the candidates. Second, MicroMiner can only search continuous sequence segments and requires separate searches for complexes. Our method can directly search surface motifs in the database and use intra-chain motifs to enhance inter-chain mutation prediction.

**Question 3: Regarding the Retrieved Structure Motif Filter, could you explain why 16 motifs were chosen and whether this number is sensitive to performance variations?**

As shown in the table below, we chose 16 motifs because further increasing the number did not significantly improve performance and might introduce more noise, leading to worse predictions.

Task: S669

  | Method               | Pearson | RMSE|
|----------------------|---------|----------|
| MSM-Mut (w/o retrieval) | 0.45  | 1.73   | 
| MSM-Mut (1 neighbor) | 0.49  | 1.62   | 
| MSM-Mut (2 neighbor) | 0.51  | 1.57   | 
| MSM-Mut (4 neighbor) | 0.53  | 1.53   | 
| MSM-Mut (8 neighbor) | 0.53  | 1.55   | 
| MSM-Mut (16 neighbor)| 0.54  | 1.51   | 
| MSM-Mut (32 neighbor)| 0.54  | 1.52   | 
| MSM-Mut (1024 neighbor)| 0.51  | 1.63   |
**Question 7: Figure 3 could benefit from some labels for the reader to identify T chain 576, 5KOV, H chain 103, and 7FAE**

In Figure 3, blue represents 5KOV and green represents 7FAE. The red segment is 5KOV's T chain 574-578, and the yellow segment is 7FAE's H chain 101-105.

To help understand this, we extracted the local structure motifs (Rebuttal Fig. 1). Despite overall differences, their local backbone atoms and surrounding non-charged polar amino acids like serine and asparagine are very similar.

**Question 8: Given that ESM-IF is trained on "global" protein structure and not local, why is it intuitive to use ESM-IF as a tool to embed local structures?**

We chose ESM-IF for embedding because, although the embedding is calculated based on global structure, it is used for protein inverse-folding, which is sensitive to local structure. Preliminary case studies showed that ESM-IF embeddings help retrieve structures with similar nearby backbone frames. See Rebuttal Fig. 2 in the supplementary PDF.

**Question 9: Would a stand-alone embedding module trained entirely on local structures perform better as an alternative to ESM-IF within the MSM-IPA framework?**

Firstly, we aim to demonstrate that our method outperforms simple local backbone encoding. To compare, we tried a simpler method: for an amino acid, we defined its embedding as [$\phi_{i-4}, \psi_{i-4}, ..., \phi_{i+4}, \psi_{i+4}$], an 18-dimensional vector of angles for this amino acid and its four sequential neighbors.  We call this embedding Continuous Backbone Angle Embedding (CBAE). This method ensures similar sequential-local backbone structures. We tested performance on the SKEMPI and S669 datasets:

Task: SKEMPI 2.0

| Method| Pearson (P.S.) | Spearman(P.S.) | Pearson | Spearman | RMSE | MAE|
|----|---|----|-----|-----|-----|---|
| MSM-Mut (w/o retrieval) | 0.4325| 0.4031| 0.6233| 0.4954 | 1.6076 | 1.2155 |
 | MSM-Mut (CBAE retrieval) | 0.4619|0.4262 |0.6524 |0.5158 |1.5531 |1.1622 |
 |MSM-Mut | 0.4736| 0.4354 | 0.6814| 0.5786 | 1.4703 | 1.0212 |

Task: S669

| Method| Pearson | RMSE|
|--|---|--|
| MSM-Mut (w/o retrieval) | 0.45| 1.73 | 
| MSM-Mut (CBAE retrieval) | 0.52| 1.55 | 
| MSM-Mut | 0.54| 1.51 | 

Our experiments showed that ESM-IF performed better than this simpler method, capturing non-adjacent backbone atom positions. Training a good pre-trained local structure encoder is still an open question. Different tasks may benefit from task-specific embeddings. For our mutation effect prediction, we used embeddings from the commonly used ESM-IF model.