Multi-Scale Representation Learning for Protein Fitness Prediction

Authors: zuobai  zhang
pascal  notin
yining  huang
aurélie  lozano
vijil  chenthamarakshan
debora  marks
payel  das
jian  tang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 16:47:51

Abstract: Designing novel functional proteins crucially depends on accurately modeling their fitness landscape. Given the limited availability of functional annotations from wet-lab experiments, previous methods have primarily relied on self-supervised models trained on vast, unlabeled protein sequence or structure datasets. While initial protein representation learning studies solely focused on either sequence or structural features, recent hybrid architectures have sought to merge these modalities to harness their respective strengths. However, these sequence-structure models have so far achieved only incremental improvements when compared to the leading sequence-only approaches, highlighting unresolved challenges effectively leveraging these modalities together. Moreover, the function of certain proteins is highly dependent on the granular aspects of their surface topology, which have been overlooked by prior models.
To address these limitations, we introduce the Sequence-Structure-Surface Fitness (**S3F**) model — a novel multimodal representation learning framework that integrates protein features across several scales. Our approach combines sequence representations from a protein language model with Geometric Vector Perceptron networks encoding protein backbone and detailed surface topology. The proposed method achieves state-of-the-art fitness prediction on the ProteinGym benchmark encompassing 217 substitution deep mutational scanning assays, and provides insights into the determinants of protein function.
Our code is at https://github.com/DeepGraphLearning/S3F.

Main paper: 
 introduction: 
 Proteins carry out a diverse range of functions in nature -from catalyzing chemical reactions to supporting cellular structures, transporting molecules or transmitting signals. These functions are uniquely determined by their amino acid sequences and three-dimensional structures. The ability to design these sequences and structures presents significant opportunities to tackle critical challenges in sustainability, new material, and healthcare. This optimization process typically begins by learning the relationship between protein sequences or structures and their function, referred to as a fitness landscape. This multivariate function describes how mutations impact protein fitness -the more accurately we model these landscapes, the better we can engineer proteins with desired traits [1, 2] .A significant challenge in modeling the fitness landscape is the scarcity of experimentally collected functional labels relative to the vastness of protein space [3] . As a result, self-supervised approaches to protein representation learning have become crucial for predicting mutation effect [4, 5] . While initial methods focused on learning a family-specific distribution over homologous protein sequences retrieved with a Multiple Sequence Alignment [4, [6] [7] [8] [9] , subsequent methods have sought to learn general functional patterns across protein families, giving rise to 'protein language models' or 'familyagnostic models' [10] [11] [12] . Recently, hybrid models have achieved state-of-the-art fitness prediction performance by leveraging the relative strengths of both types of approaches [13] [14] [15] .Although sequence-based methods are effective in recapitulating certain aspects of protein structure [16, 13] , several protein functions and tasks crucially benefit from using a more granular representation of the protein structures and surfaces [17, 18] . To bridge this gap, recent studies have exploited advances in protein structure representation learning [19] [20] [21] . For instance, inverse folding models that learn a distribution over protein sequences conditioned on a protein backbone have shown enhanced performance in stability prediction [21] [22] [23] [24] [25] . Recent efforts have also focused on integrating sequence-based and structure-based approaches [26, 27] . A prominent example, AlphaMissense, employs structural prediction losses to distill structural information into a hybrid model, highlighting the value from structural features [28] . However, these hybrid sequence-structure methods have thus far only achieved modest improvements over leading sequence-based models, or have not made their model weights publicly available. Furthermore, current methodologies fall short in effectively modeling protein surfaces, which are essential for deciphering protein interactions and capturing broader structural details [29] .In this work, we introduce a multi-scale protein representation learning framework that integrates comprehensive levels of protein information for zero-shot protein fitness prediction (Fig. 1 ). We begin with a Sequence-Structure Fitness Model (S2F) by combining a protein language model with a structure-based encoder. S2F utilizes the output of the protein language model as node features for a structure encoder, specifically a Geometric Vector Perceptron (GVP) [19] , which enables message passing among spatially close neighborhoods. Building on this, we develop a Sequence-Structure-Surface Fitness Model (S3F), which enhances S2F by adding a protein surface encoder that represents surfaces as point clouds and facilitates message passing between neighboring points. These multi-scale protein encoders are pre-trained using a residue type prediction loss on the CATH dataset [30] , enabling zero-shot prediction of mutation effects.Our methods are rigorously evaluated using the comprehensive ProteinGym benchmark [31] , which includes 217 substitution deep mutational scanning assays and over 2.4 million mutated sequences across more than 200 diverse protein families. Our experimental results show that S2F achieves competitive results with prior methods, while S3F reaches state-of-the-art performance after incorporating surface features. When further augmented with alignment information, our method improves the current state-of-the-art by 8.5% in terms of Spearman's rank correlation. Additionally, our methods have substantially fewer trainable parameters compared to other baselines, reducing pre-training time to several days on commodity hardware. Being both lightweight and agnostic of the model used to obtain initial node embeddings, they can be readily adapted to augment forthcoming, more advanced protein language models. To better understand the impact of multi-scale representation learning, we perform a breakdown analysis on different types of assays. Our results demonstrate the consistent improvements from multi-scale learning and show that incorporating structure and surface features can potentially correct biases in sequence-based methods, enhance accuracy in structure-related functions, and improve the ability to capture epistatic effects. We summarize our contributions as follows:• We develop a general and modular framework to learn multi-scale protein representations ( § 3);• We introduce two instances of this framework -S2F ( § 3.3) and S3F ( § 3.4), augmenting protein language model embeddings with structure and surface features for superior fitness prediction;• We thoroughly evaluate our methods on the 217 assays from the ProteinGym benchmark, demonstrating their state-of-the-art performance and fast pre-training efficiency ( § 4.2);• We perform a breakdown analysis for different types of assays to deep dive into the determinants of functions enabled by our multi-scale representation ( § 4.3). 
related work: 
 Protein Representation Learning. Previous research in protein representation learning has explored diverse modalities including sequences, multiple sequence alignments (MSAs), structures, and surfaces [32, 13, 20, 33] . Sequence-based methods treat protein sequences as a fundamental biological language, employing large-scale pre-training on billions of sequences to capture complex biological functions and evolutionary signals [34, 35, 11] . Alignment-based approaches, such as MSA Transformer [13] , enhance representations by incorporating MSAs, improving the capture of evolutionary relationships. Recent advancements in structure prediction tools have shifted focus toward explicitly using protein structures for representation learning [36, 20, 19, 37] . These methods employ diverse self-supervised learning algorithms like contrastive learning, self-prediction, denoising, and masked structure token prediction to train structure encoders [20, 38, 39, 26] . Additionally, extracting features from protein surfaces has shown promise in uncovering critical chemical and geometric patterns important for biomolecular interactions [29, 40, 41] . The growing availability of diverse protein data has also spurred the development of hybrid methods that combine multiple modalities for a holistic view [42, 43, 33, 44, 45] . Despite these advances, direct application of these models to zero-shot protein fitness prediction is still challenging and requires further design.Protein Fitness Prediction. Learning a fitness landscape has traditionally been approached as a discriminative supervised learning task, where models are trained to predict specific targets using labeled datasets [46] [47] [48] [49] . Recently, unsupervised fitness predictors has shown promise in surpassing these traditional methods by overcoming the limitations and biases associated with sparse labels. These unsupervised models, often designed as protein language models, are trained on vast evolutionary datasets comprising millions of protein sequences, aiming to capture a general distribution across all proteins [50, 10, 14, 11, 51] . In contrast, alignment-based models focus on specific protein families, learning from multiple sequence alignments (MSAs) to capture nuanced distribution patterns [4, 6, 8, 7, 9] . Additionally, hybrid sequence models integrate broad protein information with detailed family-specific data from alignments, enhancing the robustness of loglikelihood estimations and achieving top-tier performance [13, 15, 52] .The incorporation of structural information into protein fitness prediction has marked a promising direction in the field, inspired by recent advances in structure representation learning. Based on the concept of protein language models, SaProt utilizes structure tokens from Foldseek [53] to generalize sequence-based methods to structural data [26] . However, these structure-based methods only achieve limited improvement over the best sequence models. The latest method, AlphaMissense, integrates structural prediction losses into a hybrid model, thereby enhancing predictive accuracy [28] . Nevertheless, its high performance relies heavily on fine-tuning with weak supervision on human missense variants and it still lacks public accessibility to its model weights. Moreover, to the best of our knowledge, no existing work has employed surface-based methods for fitness prediction. To fill this gap, this work aims to integrate multi-scale protein information for fitness prediction.3 MethodIntroducing structure and surface features significantly enhances performance on binding and stability assays. This aligns with our intuition that binding and stability are closely related to structural features, suggesting that structure-based methods have an advantage in identifying structure-disrupting mutations. This observation is consistent with previous studies [24] , while our methods also maintain competitive performance across other function types.3. MSA Depth: As shown in Fig. 2 (b), protein language models perform poorly on assays with low MSA depths but excel on those with high MSA depths. This may be because proteins with low MSA depth are underrepresented in the ESM2 pre-training dataset, potentially reducing the diversity of specific families in protein language models. Introducing structure features partially mitigates this issue, and explicitly including family-specific training, like EVE, results in significant improvements. 4. Taxon: Leveraging structure and surface-based features consistently improves the fitness prediction performance across taxa. Critically, when the underlying protein language model is poor for a given taxon (eg., ESM on viral proteins [31] ), incorporating structure and surface components provides inductive biases that help overcome these limitations. Developing taxa-specific models with higher quality node representations would likely results in further performance gains. 5. Mutation Depth: Most methods perform better on single mutations than on multiple mutations, likely due to our simplified additive assumption between mutations. As mutation depth increases, the performance gains from structure and surface encoding become more pronounced. This suggests that structure-and surface-based models are better at capturing epistatic effects.In conclusion, multi-scale representation learning consistently improves performance. Incorporating structure and surface features can potentially correct biases in sequence-based methods, enhance accuracy in structure-related functions, and improve the ability to capture epistatic effects. 
preliminary: 
 Proteins. Proteins are macromolecules that form through the linkage of residues via peptide bonds. The three-dimensional (3D) structures of proteins are determined by the specific sequence of these residues. A protein with n r residues (amino acids) and n a atoms can be represented as a sequencestructure tuple (S, X). The sequence is denoted byS = [s 1 , s 2 , • • • , s nr ],where s i ∈ {1, ..., 20} represents the type of the i-th residue. The structure is represented byX = [x 1 , x 2 ..., x na ] ∈ R na×3, with x i specifying the Cartesian coordinates of the i-th atom. For simplification, we only consider the alpha carbon atoms and ignore the side-chain variations induced by mutations.Protein Fitness Landscape. The ability of a protein to perform a specific function, often referred to as protein fitness, is encoded by its sequence via spontaneous folding into structures. The effects of sequence mutations on protein function form a fitness landscape, which can be quantitatively measured through deep mutational scanning (DMS) experiments [54] . Modeling these landscapes is challenging due to the complicated relationship between sequences, structures and functions.Problem Definition. Unsupervised models that predict mutational effects are becoming fundamental tools in drug discovery, addressing the challenge of data scarcity. In this paper, we explore the task of zero-shot prediction of mutational effects using structural information. For simplicity, we focus solely on substitutions, leaving the study of insertions and deletions (indels) to future work.Formally, for each DMS assay, we start with a wild-type protein (S wt , X wt ) and generate a set of mutants by selecting specific mutation sites and randomly replacing the original residue type with During pre-training, protein sequences and structures are sampled from a database, with 15% of residue types randomly masked. These sequences are fed into a protein language model, ESM-2-650M. Then, the output residue representations are used to initialize node features in our structure and surface encoders. Through message passing on structure and surface graphs, our methods, S2F (blue) and S3F (green), accurately predict the residue type distribution at each masked position. This distribution is subsequently used for mutation preferences in downstream fitness prediction tasks.a new one. For a mutant (S mt , X mt ) with multiple mutations T , the sequence changes such thats mt t ̸ = s wt t if t ∈ T ; otherwise, s mt t = s wt t .We assume that the backbone structures remain unchanged post-mutation (X mt = X wt ). The objective is to develop an unsupervised model that can predict a score for each mutant to quantify the changes in fitness values relative to the wild-type.Protein language models trained using the masked language modeling objective are designed to predict the likelihood of a residue's occurrence at a specific position within a protein, based on the surrounding context [11, 32] . As demonstrated in [5] , this method can score sequence variations using the log odds ratio between mutant and wild-type proteins for given mutations T :EQUATIONwhere S \T denotes the input sequence masked at each mutated position in T and an additive model is assumed over multiple mutation sites. In the zero-shot setting, inference is performed directly on the sequence under evaluation using only forward passes of the model.Although these protein language models are effective in predicting mutational effects, they do not incorporate explicit structural information during pre-training, which is crucial for determining protein functions. Building on the ESM-2-650M model, we next propose Sequence-Structure Fitness Model (S2F), which integrates structural data into the sequence-based predictive framework.A major challenge in applying structure-based methods to fitness prediction is modeling how mutations impact protein structures. To bypass this problem, we adopt a simplified assumption: the backbone structures of proteins remain unchanged post-mutation. Additionally, we choose to omit side-chain information, which reveals residue types and could potentially lead to information leakage.Geometric Message Passing. We build a radius graph for each protein with nodes representing alpha carbons. Two nodes are connected if their Euclidean distance is less than 10Å. We use Geometric Vector Perceptrons (GVP) [19] to perform message passing across the graph. GVPs replace standard Multi-Layer Perceptrons (MLPs) in Graph Neural Networks, operating on scalar and geometric features that transform as vectors under spatial coordinate rotations.We represent the hidden state of residue i at the l-th layer by h(l) i ∈ R d × R d ′ ×3with d-dim scalar features and d ′ -dim vector features. The initial node feature of residue i is set using the protein language model embeddings, specifically, h (0) i = (ESM(s i |S \T ), 0). The edge features are given by e (j,i) = (rbf(x j -x i ), x j -x i ), where rbf(•) computes pairwise distance features with Radial Basis Function (RBF) kernels [55] . In the network, node and edge features are concatenated to facilitate message passing via the GVP module on the (scalar, vector) representations. Each message passing layer is followed by a feed-forward network. Formally,EQUATIONEQUATIONwhere N (i) represents the set of neighbors of node i. The GVP module ensures SE(3)-invariance for scalar features and SE(3)-equivariance for vector features. The scalar features at the last layer h(L) iof each node i are fed into a separate linear layer for predicting the residue type. Practically, we utilize L = 5 layers of GVP, with d ′ = 16 vector and d = 100 scalar hidden representations.Besides protein sequences and structures, protein surfaces-defined by neighboring amino acids-are characterized by distinct patterns of geometric and chemical properties. For instance, within a folded protein, hydrophobic residues tend to cluster inside the core, while hydrophilic residues are exposed to water solvent on its surface. These patterns provide crucial insights into protein function and potential molecular interactions. Now we integrate this aspect into our S2F model to propose a new model called Sequence-Structure-Surface Fitness Model (S3F).Surface Processing. We employ dMaSIF to generate the surface based on the backbone structure of each protein [40] . The surface is represented as a point cloud { x1 , x2 , ..., xns } ∈ R 3 , consisting of n s (6K-20K) points. These points are sampled based on the levels of a smooth distance function defined over each atom. Each surface point i is associated with geometric features fi , specifically Gaussian curvatures [56] and Heat Kernel Signatures [57] , providing a detailed characterization of the surface topology. In the sequel, we will use the notations • with tilde for surface points to distinguish with those for nodes in structure graphs.Surface Feature Initialization. After generating protein surfaces, we build a mapping between the structure and surface graphs, projecting residue features onto the surface points. For each surface point i, we identify its 3 nearest neighboring residues, denoted as N surf-res (i), each initialized with its ESM feature. These features are concatenated with their Euclidean distances to point i and processed through an MLP. The average features of these neighbors are then combined with the geometric features f i . Formally, the scalar feature for surface node i is initialized as follows:EQUATIONwith the vector feature is initialized as the zero vector 0.Surface Message Passing. To perform message passing on protein surfaces, we construct a surface graph with a k-nearest neighbor graph, wherein each surface point is linked to its 16 nearest neighbors on the surface. For each edge (j, i) in this graph, the edge feature ẽ(j, i) is initialized using RBF kernels and represented as ẽ(j, i)= (rbf( xj -xi ), xj -xi ).Similar to S2F, we employ GVP to perform message passing on the scalar and vector representations:EQUATIONEQUATIONwhere N surf (i) represents the neighbors of surface node i. Similar with the approach on structure graphs, we utilize another five layers of GVP, with 16 vector and 100 scalar hidden representations.Residue Representation Aggregation. After performing message passing on both the structure and surface graphs, we combine the structure representations, h (L) , with the surface representations, h(L) .For each residue i, we identify the 20 nearest surface points, denoted as N res-surf (i), and compute their mean representations to enhance the residue's representation. Here we use more neighbors for mapping the surface points back to residues, as there are typically much more surface points than residues. Specifically, we update each residue representation as follows:EQUATIONwhere ← indicates that the left-hand side is updated with the value from the right-hand side. The updated representation h (L) i is then input into a separate linear layer for predicting the residue type.Following the pre-training methodology in [58] , we select 15% of residues at random for prediction.If the i-th residue is selected, we manipulate the i-th token by replacing it with: (1) the [MASK] token 80% of the time, (2) a random residue type 10% of the time, and (3) leaving the i-th token unchanged 10% of the time. The final hidden state h (L) is then used to predict the original residue type using cross-entropy loss. For S3F, we avoid information leakage from surfaces by removing the top 20 closest surface points for each selected residue. During pre-training, the weights of the ESM-2-650M model are frozen, and only the GVP layers for structure and surface graphs are trainable. This strategy helps preserve sequence representations and improves pre-training efficiency. Our models are pre-trained on a non-redundant subset of CATH v4.3.0 dataset (CC BY 4.0 license) [30] , which contains 30,948 experimental structures with less than 40% sequence identity. S2F and S3F are trained with batch sizes of 128 and 8, respectively, for 100 epochs on four A100 GPUs. The pre-training time for S2F and S3F are 9 hours and 58 hours, respectively.During inference, we adopt a strategy similar to that used in ESM-1v [5] . For each mutation, we mask the residue type at all mutation sites and remove the corresponding surface points, similar as the pre-training process. We use AlphaFold2 to predict the wild-type structures. To mitigate the influence of low-quality structures, for mutations on residues with a predicted Local Distance Difference Test (pLDDT) score below 70, we use the output scores from the baseline model, ESM-2-650M. For residues with pLDDT score no less than 70, we use the outputs from our S2F or S3F models. 
setup: 
 Evaluation Dataset. To assess the performance of our method on zero-shot protein fitness prediction, we run experiments on the ProteinGym benchmark (MIT License) [31] . This benchmark contains a comprehensive collection of Deep Mutational Scanning (DMS) assays, which covers a variety of functional properties such as thermostability, ligand binding, viral replication, and drug resistance. Specifically, we use 217 substitution assays that include both single and multiple mutations.Metric. Given the complex, non-linear relationship between protein function and organism fitness [59] , we select the Spearman's rank correlation coefficient as a primary metric for evaluating model prediction against experimental measurements. Besides, to ensure a comprehensive assessment, we include other metrics from the official ProteinGym benchmark: the Area Under the ROC Curve (AUC) and the Matthews Correlation Coefficient (MCC), which compare model scores with binarized experimental data. We also report the Normalized Discounted Cumulative Gains (NDCG) and Top 10% Recall with the aim to identify the most functionally effective proteins.Baseline. We select a subset of baselines from the ProteinGym benchmark for comparison, categorizing them based on whether the model requires multiple sequence alignments (MSA) as input.In the category without MSAs, we include three protein language models-ProGen2 XL [50] , CARP-640M [51] , and ESM-2-650M [32] ; three inverse folding models-ProteinMPNN [22] , MIF [23] , and ESM-IF [21] ; and three sequence-structure hybrid models-MIF-ST [23] , ProtSSN [27] , and SaProt [26] . For models utilizing alignments, we choose three family-specific models-DeepSequence [6] , EVE [7] , and GEMME [8] , as well as three hybrid models that combine family-agnostic and specific approaches, including MSA Transformer [13] , Tranception L with retrieval [14] , and TranceptEVE [15] . All baseline results are taken directly from the ProteinGym benchmark [31] . * Since the numbers of trainable parameters for family-specific models depend on the length of protein sequences, we use the average length (400 AAs) of all ProteinGym sequences for estimation.We report the zero-shot performance of our methods, S2F and S3F. To benchmark against alignmentbased models, we further enhance our models by ensembling them with EVE predictions through the summation of their z-scores, resulting in two variants named S2F-MSA and S3F-MSA.We report the benchmark results and model information for all baselines in Table 1 . Among the methods that do not require MSAs, S2F achieves competitive results, surpassing protein language models and inverse folding models, and only slightly lagging behind SaProt. When augmented with surface features, S3F becomes the best model in this category, even outperforming the top alignmentbased model, TranceptEVE, by a significant margin in terms of Spearman's rank correlation. Despite relying on protein language models, our methods have substantially fewer trainable parameters compared to other baselines (20M in S3F v.s. 650M in SaProt). Hence, our methods finish pretraining within several days, whereas some large-scale baselines require months. Additionally, when enhanced with alignment information, S3F-MSA improve the current state-of-the-art method, SaProt, by 8.5% in terms of Spearman's rank correlation, further demonstrating the potential of our approach. These results highlight both the effectiveness and the parameter efficiency of our proposed methods.From Table 1 , we observe the progressive improvements achieved by gradually incorporating various protein aspects into the model, as demonstrated by comparing the results of ESM2, S2F, S3F, and S3F-MSA. To better understand the impact of this multi-scale representation learning, we conduct a breakdown analysis on different types of assays. In Fig. 2(a-d ), we report the performance of these four methods on assays grouped by function types, MSA depths, taxon, and mutation depths. Our analysis reveals the following key points:1. Overall: Consistent improvements are observed across all types of assays in different categories when structure, surface, and alignment information are added to the model.To facilitate the usage of our methods on proteins without experimentally determined structures, we employ AlphaFold2 to generate structures for ProteinGym assays. By default, five structures are generated for each assay, with the best one selected based on pLDDT scores. Now we study how the quality of these predicted structures influences performance through two experiments.First, we categorize the 217 assays into three groups based on the pLDDT of their predicted structures: 95 assays with pLDDT over 90 are classified as high-quality, 18 assays with pLDDT less than 70 as low-quality, and the remaining 104 assays as medium-quality. We report the average results for these categories in Fig. 2 (e). The results indicate that performance for all four baselines positively correlates with structure quality, even for the sequence-based method, ESM. This may suggest our limited prior knowledge about these proteins with low-quality structures for all models, including AlphaFold2 and ESM2. Additionally, while improvements from structure and surface features are observed across all assays, those with high-quality structures benefit more significantly, especially for S3F. This underscores the importance of accurate structures for fitness prediction.Next, we analyze the impact of structure quality using the five AlphaFold2-predicted structures. We perform fitness predictions five times, each time using a different group of structures ranked by their pLDDT scores. Specifically, we test S2F and S3F using the top 1, 2, 3, 4, and 5 predicted structures for all assays, respectively. The results, plotted in Fig. 2 (f), show clear performance drops when lower-quality structures are used. This further highlights the reliance on high-quality structures.In contrast to large-scale models pre-trained on UniRef100 or the AlphaFold Database, our method employs a much smaller dataset, CATH, which contains only 31,000 structures after clustering. This raises the question of whether the benefits of structure-and surface-based methods can generalize well to unseen protein families, despite the limited scale of the pre-training data. To address this, we select 23 out-of-distribution assays from ProteinGym, whose sequences have less than 30% similarity to the pre-training dataset. We plot the average performance of the four methods on these assays in Fig. 2(g ). Although the absolute performance of all four methods decreases compared to their overall performance, we still observe consistent improvements from the structure-and surface-based methods, as evidenced by the performance gains of S2F over ESM and S3F over S2F. This proves the generalization ability of our methods to unseen protein families.We illustrate the advantages of our method by focusing on the results from the high-throughput assay of the IgG-binding domain of protein G (GB1) introduced in [60] . The assay quantifies the effects of all single and double mutations, thereby providing detailed insights into the epistatic relationships between residue pairs in the domain. We report in Fig. 3 the pair-specific Spearman's rank correlation between the experimental measurements and predictions from ESM, S2F and S3F, respectively. Leveraging structure and surface features leads to a superior ability in predicting mutation effectsin particular the epistatic effects between residue pairs that are far in sequence space but close in the tertiary structure of GB1 (off-diagonal terms in Fig. 3(a-c ) and colored residues in Fig. 3(d) ). 
conclusion and limitations: 
 This work introduces S3F, a novel multi-scale representation learning framework for protein fitness prediction that achieves state-of-the-art performance on the ProteinGym benchmark with a lightweight model. The breakdown analysis provides insights into how protein modalities contribute to predicting various fitness landscapes. However, there are still several limitations in this work. First, our models are pre-trained on relatively small set of experimentally determined protein structures and may further benefit from leveraging more diverse structure sets, such as the AlphaFold database. Second, we make simplifying assumptions, such as ignoring side-chain information and assuming that backbone structures remain unchanged after mutations, which may limit the model's capacity. Third, our approach is limited to substitutions effects, and does not handle insertions nor deletions.Table 2 : Average Spearman and Std. Error of Difference to Best Score. The best performing model is S3F-MSA, highlighted in red. Spearman is the difference between the baseline models and the S3F-MSA. We followed the ProteinGym benchmarks to compute the std. error of difference to best score. In this table, it is the bootstrapped std. error of the differences between the Spearman of each model and the Spearman of S3F-MSA. 
 
 Appendix: In Section 4.5, we select 23 assays from ProteinGym by filtering with sequence similarity. Here we report the performance on these assays in Figure 4 . The results on assays with the same proteins are grouped together.To quantify the statistical significance of the performance, we follow the same methodology as in ProteinGym and compute the non-parametric bootstrap standard error of the difference between the Spearman performance of a given model and that of the best overall model (10k bootstrap samples).Our performance delta with prior methods are all statistically significant, as shown in Table 2 .We provide ablation study results in Table 3 , confirming the performance lift from the various modalities involved. The ablation dropping both structure and surface features corresponds to ESM2 in Table 1 . Note that surface message-passing is designed to capture fine-grained structural aspects that enhance the coarse-grained features learned by our S2F (sequence+structure) model. However, relying solely on these fine-grained features without the context from structural features, as we do in the ablation removing structural inputs, appears to be detrimental to performance. 38th Conference on Neural Information Processing Systems (NeurIPS 2024). 


Decision: The paper proposes a multimodal representation learning framework for protein by integrating structural and surface information, and applies it to zero-shot protein fitness prediction tasks. The paper addresses an important and practical task in the field of computational biology. Although the technical novelty of the method is relatively modest, this study could potentially benefit and inspire further research in protein studies. The reviewers have raised valuable comments, and the authors provide additional information and experiments in the rebuttal to address these concerns. Such content should be included in the revision to improve the paper quality and help better understand the contributions.

Reviews: review: paper_summary: paper_summary: This paper augments a protein language model with two additional features, structure and surface information. The authors show that they can effectively use this information to predict protein function slightly better (though SOTA-level on relevant benchmarks) than sequence-only PLMs.
strength_weakness: strength_weakness: Novelty/Originality:
I can’t tell how this paper is a real advance on (https://arxiv.org/pdf/2204.02337, ref 33 in this paper), with the title being nearly identical too, which is published 2 years ago. 

Yes the test set is different (evaluated on proteinGym), and some features of the architecture are also different, but the author choice of not benchmarking against Holoprot on the same dataset is problematic, if for nothing else because it’s the most clear way of showing the performance advance due to architectural improvement and no simply because of additional features (Because the “idea” of combining sequence, surface and structure is already executed.)

From the abstract
“Moreover, the function of certain 11 proteins is highly dependent on the granular aspects of their surface topology, 12 which have been overlooked by prior models. “
- Reference 33 (https://proceedings.neurips.cc/paper/2021/file/d494020ff8ec181ef98ed97ac3f25453-Paper.pdf) in your own citations is exactly about this, connecting sequence, structure and surface. Can you explain? 
- Relatedly I think the similarity (almost identical : "Multi-Scale Representation Learning on Proteins") in the paper title is a poor decision, if not for the appearance of plagiarism, for the fact that it makes the difference between the papers even less clear. It should be titled to embolden the difference. That paper is also doing fitness prediction.
questions: questions: I think the authors really need to contrast this work with reference 33 and argue why it’s a meaningful conceptual or performance advance. 

I think the paper would be better if the authors would do at least two ablation studies where at least 
- structure is dropped
- both structure and surface are dropped  (dropping sequence is more complex but it may also be interesting). 

It is surprising to me that in the non MSA case S2F basically does not improve on benchmarks, and it could be that the structure data is adding no performance on top of the surface. Comparing to the original sequence-only model in this architecture would also help clarify how much improvement is simply due to other factors than surface information.
limitations: limitations: Conceptually the paper is doing something that has been done before, with some small innovations. It is relevant however that it can indeed consistently improve on existing language models.


:::::

Updating my previous score to 5
paper_summary: paper_summary: This paper proposes a new protein fitness prediction model, the Sequence-Structure-Surface (S3F) model, which integrates protein sequence information from a protein language model embedding, protein structure information processed through a Geometric Vector Perceptron (GVP) module, and protein surface information processed through the dMaSIF model and a GVP module. After pre-training S3F on CATH (to be specific, pre-training the GVP modules while freezing the protein language model weights), the model outperforms state of the art baselines on zero-shot protein function prediction. The paper also analyzes the breakdown of these results across different fitness prediction settings, finding that all settings benefit from adding structure, surface, or MSA information on top of sequence information, with the largest gains arising in binding and stability assays and in settings where protein language models have limited training data or biased training data.
strength_weakness: strength_weakness: 1. The paper does not include results or analysis about the sensitivity of results for various hyperparameters of the model, such as the width, depth, and hidden dimension of GVP modules, the number of surface points to include in a neighborhood, and the choice of protein language model for embeddings. Since the paper focuses on the ProteinGym benchmark for its only evaluation, it would be nice to see more in-depth experimentation and/or ablation of different parts of the S3F model.
2. The central claim that surface information is important could also be strengthened with more experiments to understand how pre-processing structural information into surface features extracts useful information that is not as easily accessible with structure-based models that don’t do similar feature engineering/pre-processing.
questions: questions: 1. In Table 1, the bottom half of the table compares alignment-based models with S2F and S3F ensembled with EVE scores. It would be interesting to see results for each of the baselines in the top half of the table (e.g. MIF-ST, ProtSSN, SaProt) also ensembled with EVE scores.
2. What is your current intuition about why structure-aware fitness prediction models are unable to leverage surface information to its full extent?
limitations: limitations: The authors assess limitations of their work, including the fact that their method is limited to considering substitutions, and cannot handle insertions or deletions. I think this is a salient point to keep in mind, which limits the applicability of their method to real protein design tasks. See also weaknesses above.
paper_summary: paper_summary: The paper presents a multimodal framework that integrates protein sequence, structure, and surface information together to predict protein fitness. The task of protein fitness prediction is a critical quality assessment of protein embeddings. Protein language models (pLM) are used for sequence representation which is embedded in a graph representation of protein structure and processed with geometric vector perceptron (GVP). The dMaSIF, a protein surface representation model, is used to encode protein surface, together with pLM embeddings of nearest residues. Surface features and sequence-structure features are integrated by concatenation and linear layers to predict masked residue identities. The model is evaluated on ProteinGym, a golden protein fitness benchmark dataset, against various pLM models and multiple sequence alignment (MSA) -based models and achieves favorable results in terms of spearman correlation with mutational effects. Evaluations are further run on AlphaFold predicted structures and less representative proteins to prove generalization.
strength_weakness: strength_weakness: 1. There is not adequate reference of prior work on similar topics, e.g. multimodal fusion and protein surface representation learning. The idea of integrating surface into protein representation is not novel, for example [1], but is neither acknowledged nor benchmarked in the paper.

2. Even though the performance of proposed method is better than competing methods, it's marginal. As the author writes in abstract '...these sequence-structure models have so far achieved only incremental improvements when compared to the leading sequence-only approaches.' However, SaProt (sequence+structure) provides 0.35 performance gain of correlation compared to ESM2 (seq) while S3F (proposed) provides 0.13 compared to SaProt. The advantage brought by surface representation is more incremental.

3. Statistical results should be provided for figure 2 when possible. How is the variation of performance on 217 assays? 

References:
[1] https://arxiv.org/pdf/2309.16519
questions: questions: 1. line 195. ESM features are integrated into surface representations. Why not just use surface feature itself because it contains chemistry and geometry information already?

2. line 218. 'we avoid information leakage from surfaces by removing the top 20 closest surface points for each selected residue' How to ensure 20 is sufficient?

3. line 230. What percentage of results are output from ESM?

4. Figure 3. Why do some residues have negative correlation to model scores? Is there any quantitative analysis of differences rather than just qualitative visualization?

5. line 347. 'ignoring side-chain information'. If side chain is ignored, how to get surface information?
limitations: limitations: The majority of experiments in the paper is accomplished with experimental structure rather than predicted structures. And the results show that ordinary prediction quality could harm performance of the method. Thus, it is not clear if the model can be finetuned on AFDB to achieve better performance because of the quality of side-chain conformations in predicted structures.
paper_summary: paper_summary: This work targets at the protein fitness prediction and introduces a sequence-structure-surface multi-modality (aka. multi-scale) self-supervised learning scheme. The results show that S3F outperforms baseline algorithms and achieves SOTA in the ProteinGym benchmark. I like the inspiration and favor S3F's promising performance. However, the study is very close to some recently published papers at ICLR 2024 (i.e., ProteinINR and PPIFormer) but completely ignores them. A fair comparison and analysis of the difference is necessary and I am afraid this is a big reason for me to give a borderline score.
strength_weakness: strength_weakness: (1) Missing of some closely related baselines and relevant work, GearNet [A] is one of the earliest studies in structural pretraining. ProteinINR [B] also leverages sequence, structure, and surfaces in a self-supervised way. The only difference is that they are not specifically designed to solve zero-shot problems. Based on these facts, the so-called "multi-scale representation learning" of S3F is no longer novel to me. From my point of view, both GearNet and ProteinINR can be simply adjusted to realize zero-shot prediction. Thus, it would be more interesting to see whether S3F outpasses them in all categories of downstream tasks.   

Besides, there are also some appealing self-supervised methods to predict mutant effects. RDE [C] pretrain a structural encoder by masking and predicting side-chain angles. The author also did not discuss this line of research. 

Last but not least, PPIFormer [D] pretrained a structural encoder with simple MLM in complex structures and can conduct zero-shot fitness prediction. I suppose it should also be compared and mentioned. 

[A] Protein Representation Learning by Geometric Structure Pretraining. ICLR 2023. 


[B] Pre-training Sequence, Structure, and Surface Features for Comprehensive Protein Representation Learning. ICLR 2024. 

[C] Rotamer Density Estimator is an Unsupervised Learner of the Effect of Mutations on Protein-Protein Interaction. ICLR 2023. 

[D] Learning to design protein-protein interactions with enhanced generalization. ICLR 2024.

(2) Some important relevant works are missing. For instance, the author claims that the weights of ESM-2 are frozen and only the structural encoder is tuned. This practice is very similar to the one in [A]. I would recommend the author at least cite it. 

[A] Integration of Pre-trained Protein Language Models into Geometric Deep Learning Networks. Communications Biology 2023. 


(3) The author used dMaSIF to generate the surface based on the backbone structures of proteins. This raises two doubts. Firstly, the side-chain atoms are ignored for surface generation, therefore the surface should be smaller than its standard one. Secondly, the fast-sampling algorithm developed by dMaSIF has unavoidable randomness. Different random seeds produce different surfaces. Has the author taken these two potential negative effects into consideration? Has the author adopted software like PyMol to acquire the surfaces?
questions: questions: (1) In line 256, the author said "... by ensembling them with EVE predictions through the summation of their z-scores". I am not familiar with this alignment part. Can you please explain the details more? 

(2) The title used "mutli-scale" reprsentations of proteins. However, from my personal point of view, multi-scale refers to atom-scale, residue-scale, and protein-scale, As this study proposed to leverage sequence, structure, and surfaces, "multi-modality" would be more appropriate. 

(3) The author claimed that the backbone structures remain unchanged post-mutation. This is fine in most settings. However, did the author consider a more challenging circumstance, what if the structure varies significantly due to the mutation? [A] proposes a co-learning framework to simultaneously forecast the fitness change and the structural update. Does the author consider this factor? 

[A] Thermodynamics-inspired Structure Hallucination for Protein-protein Interaction Modeling. ICLR submission. 

(4) If I understand it correctly, the objective is to develop an unsupervised model that can predict a score for each mutant to quantify the changes in fitness values relative to the wild-type. Thus, it is more related to the mutant effect prediction task rather than the pure fitness prediction. I would recommend the author use "mutant effect prediction" instead of "fitness prediction" to better depict the target problem.
limitations: limitations: Yes

Comments: title: Concluding remarks
Concluding remarks
Concluding remarks
Concluding remarks
Reviewer response
Thank you
Thank you
Update
comment: Dear ethics reviewer,

Thank you very much for your thoughtful response. We concur with the other reviewers that our work does not present a particular ethical risk in itself, and acknowledge that our sentence about “future studies” may have unintentionally conveyed a different impression.

To clarify, our intention in this section was primarily to underscore the connection between our work -- which primarily focuses on predicting the effects of mutations in proteins, and protein engineering -- which is where potential misuse risks could arise.

Mitigating the risks stemming from ongoing progress in protein engineering, even though these risks are currently considered tenuous, is an active area of research. This is what we intended by our mention of “future studies,” and we will rephrase this in our revision. The general idea is that the same models that are developed for good (eg., to develop novel therapies) could, in theory, be repurposed for malignant applications (eg., to design dangerous biological agents). 

However, it is important to note that there are many steps involved in the creation of such a threat -- from conception, to in silico development (which is the only step that models could potentially facilitate), to experimental validation, to larger scale manufacturing. For the same reasons that “in silico development” is often not the most challenging aspect of drug discovery [1], it is not the primary barrier for malignant applications either. Consequently, current risk mitigation strategies rather focus on monitoring synthesis capabilities to detect and prevent potential misuse [2]. We will add these clarifications in the revision as well.

Kind regards,  
The authors

[1] Lowe. Why AlphaFold won’t revolutionise drug discovery? Chemistry World, 2022.  
[2] Baker & Church. Protein design meets biosecurity, Science, 2024.
Dear reviewer,

Thank you very much for reading through our responses and for raising your score. We will make sure to include the changes discussed above in the final revision.  
Please let us know if there are any other points we can clarify before the discussion period ends.

Kind regards,  
The authors
Dear reviewer,

Thank you very much for reading through our responses and the additional question. Our surface features are indeed derived from the backbone structure only, which is both an approximation that has been used in prior literature (see for example [1]) and which has led to strong empirical performance in our experiments. We will further clarify this point in the revised manuscript, as well as include the rationale we had provided above for adopting this approach.

Is there any other point of concern that we could help clarify before the discussion period ends?

Kind regards,  
The authors

[1] Hua et al. Effective Protein-Protein Interaction Exploration with PPIretrieval
Dear reviewer,

Thank you very much for reviewing our responses. As we understand it is the remaining point of concern, we would like to share a few concluding thoughts on the significance of our performance improvements. We consider the ProteinGym benchmark, with its 2.4 million mutation effect labels, to be the most critical benchmark for protein fitness prediction, akin to what ImageNet is for computer vision. While our experiments were focused on this benchmark, we believe it is the most relevant for the task we addressed. Regarding the improvement magnitude, we appreciate your feedback and would like to note the following: 
1) SaProt-MSA was not an existing baseline but rather one we specifically computed for this rebuttal based on your suggestion. This result may offer valuable insights for the community, and we plan to include it in our revision
2) Even a ~1-point increase in Spearman correlation across 217 assays from ProteinGym represents a meaningful advancement, which we believe underscores the importance of leveraging additional data modalities, such as surface information, as we have proposed in this work. 

Please let us know if there are any other points we can clarify before the discussion period ends.

Kind regards,  
The authors
Dear reviewer,

Thank you very much for your final feedback and for raising your score! 

As we near the end of the discussion period, we wanted to share a few concluding remarks regarding the novelty aspect of our work, as it appears to be the remaining point of concern. Firstly, we would like to emphasize that achieving the performance we obtained on ProteinGym required significant craftsmanship to optimally leverage the structural and surface features within our proposed architecture. In our view, this represents one of the many forms of novelty that NeurIPS aims to highlight. Secondly, our work is the first to explicitly demonstrate the value of these modalities for fitness prediction performance, offering another novel insight that we believe will be highly valuable for practitioners. Lastly, our method introduces a model-agnostic approach to augment protein language models with structural and surface features. This innovative aspect ensures that our approach can be seamlessly applied to enhance future protein language models as they continue to evolve and improve.  

Please let us know if there are any other points we can clarify before the discussion period ends.

Kind regards,  
The authors
Thank you for the response. If I understand correctly, the surface is calculated only from backbone structure? Then this definition is largely deviated from what people normally conceive as true protein surface which is mostly dictated by side chains. I believe the wording regarding surface is misleading and would cause confusion for audience interested in protein surface representation.
Thank you to the authors for their additional experiments and careful response. I think the paper presents nice empirical results improving upon state-of-the-art ProteinGym performance, although the improvements are somewhat small (e.g. SaProt-MSA compared to S3F-MSA), and limited to that benchmark. Based on the responses and other reviewer comments, I will keep my score.
I thank the authors for their responses. 

Having read their answers, together with the other reviewers, I'm comfortable increasing my score. I do still feel like the conceptual advance on the previous paper and performance advance over the best benchmarks (given the auxiliary data it needs) is not at the level of a clear accept.
Thanks for your response. I am satisfied with the additional experiments and the outstanding performance compared to those important baselines. However, I still believe the way to generate surface using dMaSIF is not a smart choice despite their effectiveness in GPU parallel. 

However, I appreciate the efforts in answering my question and would like to raise my score to 5. I hope the authors can incorporate new changes during the rebuttal period into the final revision.