Title: Multi-Scale Representation Learning for Protein Fitness Prediction
Venue: NeurIPS 2024
Accepted: yes
Keywords: Protein representation learning; protein fitness prediction; self-supervised pre-training
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.

Multi-Scale Representation Learning for Protein Fitness Prediction

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 .

1 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-9], subsequent methods have sought to learn general functional patterns across protein families, giving rise to 'protein language models' or 'familyagnostic models' [10-12]. Recently, hybrid models have achieved state-of-the-art fitness prediction performance by leveraging the relative strengths of both types of approaches [13-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-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-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 S equenceS tructure F itness 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 S equenceS tructureS urface F itness 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).

2 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-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 Method

3.1 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 by S = [ s 1 , s 2 , · · · , s n r ] , where s i ∈ { 1 , ..., 20 } represents the type of the i -th residue. The structure is represented by X = [ x 1 , x 2 ..., x n a ] ∈ R n a × 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

YNYNSHNVYIMADKQKI

Protein Structure Database

S

H

N

*

*

*

M

A

S

H

N

*

*

*

M

A

S

equence-

S

tructure-

S

urface

F

itness Model (

S3F

)

S

equence-

S

tructure

F

itness Model (

S2F

)

YNYNSHNVYIMADKQKI

WYRFHMEKQDNLICTVPSAG

Active

Deleterious

Mutation Preference

̸

a new one. For a mutant ( S mt , X mt ) with multiple mutations T , the sequence changes such that s 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.

3.2 Protein Language Models for Mutational Effect Prediction

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 :

where 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.

3.3 Sequence-Structure Model for Fitness Prediction (S2F)

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 S equenceS tructure F itness 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 ′ × 3 with 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,

where 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 ) i of 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.

3.4 Sequence-Structure-Surface Model for Fitness Prediction (S3F)

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 S equenceS tructureS urface F itness 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 { ˜ x 1 , ˜ x 2 , ..., ˜ x n s } ∈ 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 ˜ f i , 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:

with 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 ˜ e ( j, i ) is initialized using RBF kernels and represented as ˜ e ( j, i ) = ( rbf (˜ x j -˜ x i ) , ˜ x j -˜ x i ) .

Similar to S2F, we employ GVP to perform message passing on the scalar and vector representations:

where 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:

where ← 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.

3.5 Pre-Training and Inference

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 ESM2-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.

4 Experiment

4.1 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], CARP640M [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.

4.2 Benchmark Result

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.

4.3 Breakdown Analysis for Multi-Scale Learning

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:

Overall: Consistent improvements are observed across all types of assays in different categories when structure, surface, and alignment information are added to the model.

Function type: Introducing 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

Activity

Binding

Expression

OrganismalFitness

Stability

0.3

0.4

0.5

0.6

SpearmanR

.425

.337

.415

.369

.523

.458

.384

.456

.402

.568

.468

.404

.471

.411

.593

.502

.440

.479

.476

.581

(a)

Function Type

Method

ESM

S2F

S3F

S3F-MSA

Low

Medium

High

0.3

0.4

0.5

SpearmanR

.335

.406

.515

.394

.458

.528

.412

.473

.547

.469

.509

.546

(b)

MSA Depth

Human

Eukaryote

Prokaryote

Virus

0.2

0.3

0.4

0.5

0.6

SpearmanR

.456

.471

.476

.238

.471

.505

.506

.380

.484

.525

.520

.397

.500

.550

.530

.496

(c)

Taxon

1

2

3

4

5+

0.2

0.3

0.4

0.5

SpearmanR

.421

.247

.220

.162

.218

.468

.278

.283

.241

.297

.484

.326

.323

.281

.337

.498

.331

.369

.330

.383

(d)

Mutation Depth

Low

Medium

High

0.30

0.35

0.40

0.45

0.50

0.55

SpearmanR

.348

.426

.457

.377

.468

.496

.381

.485

.514

.402

.512

.542

(e)

Structure Quality

0

1

2

3

4

0.43

0.44

0.45

0.46

0.47

SpearmanR

.454

.452

.451

.447

.445

.470

.464

.463

.458

.455

(f)

PLDDT Rank

All (217 assays)

OOD (23 assays)

0.35

0.40

0.45

0.50

SpearmanR

.414

.373

.454

.413

.470

.421

.496

.485

(g)

Overall Performance

structure-disrupting mutations. This observation is consistent with previous studies [24], while our methods also maintain competitive performance across other function types.

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.

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.

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.

4.4 Impact of Structure Quality

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

(a) ESM

(b) S2F

(c) S3F

(d) SPG1_STRSG

SpearmanR

228

228

228

233

233

233

238

238

238

243

243

243

248

248

248

253

253

253

258

258

258

263

263

263

268

268

268

273

273

273

278

278

278

228233 238 243 248 253 258 263 268 273 278

228 233238 243 248 253258 263 268 273 278

228233 238 243 248253 258 263 268 273 278

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.

4.5 Generalization Ability to Unseen Protein Families

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.

4.6 Case Study: Epistatic effects in the IgG-binding domain of protein G

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 effects in 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)).

5 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 different 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.

Acknowledgments

The authors would like to thank Jiarui Lu and Chenqing Hua for their helpful discussions.

This project is supported by AIHN IBM-MILA partnership program, the Natural Sciences and Engineering Research Council (NSERC) Discovery Grant, the Canada CIFAR AI Chair Program, collaboration grants between Microsoft Research and Mila, Samsung Electronics Co., Ltd., Amazon Faculty Research Award, Tencent AI Lab Rhino-Bird Gift Fund, a NRC Collaborative R&D Project (AI4D-CORE-06) as well as the IVADO Fundamental Research Project grant PRF-2019-3583139727.

P.N. was supported by a Chan Zuckerberg Initiative Award (Neurodegeneration Challenge Network, CZI2018-191853). D.M. holds a Ben Barres Early Career Award from the Chan Zuckerberg Initiative as part of the Neurodegeneration Challenge Network (CZI2018-191853) and is supported by a NIH Transformational Research Award (TR01 1R01CA260415).

References

Philip A. Romero and Frances H. Arnold. Exploring protein fitness landscapes by directed evolution. Nature Reviews Molecular Cell Biology , 10:866-876, 2009.

Pascal Notin, Nathan J. Rollins, Yarin Gal, Chris Sander, and Debora Marks. Machine learning for functional protein design. Nature Biotechnology , 42:216-228, 2024.

Surojit Biswas, Grigory Khimulya, Ethan C Alley, Kevin M Esvelt, and George M Church. Low-n protein engineering with data-efficient deep learning. Nature methods , 18(4):389-396, 2021.

Thomas A Hopf, John B Ingraham, Frank J Poelwijk, Charlotta PI Schärfe, Michael Springer, Chris Sander, and Debora S Marks. Mutation effects predicted from sequence co-variation. Nature biotechnology , 35(2):128-135, 2017.

Joshua Meier, Roshan Rao, Robert Verkuil, Jason Liu, Tom Sercu, and Alex Rives. Language models enable zero-shot prediction of the effects of mutations on protein function. Advances in neural information processing systems , 34:29287-29303, 2021.

Adam J Riesselman, John B Ingraham, and Debora S Marks. Deep generative models of genetic variation capture the effects of mutations. Nature methods , 15(10):816-822, 2018.

Jonathan Frazer, Pascal Notin, Mafalda Dias, Aidan Gomez, Joseph K Min, Kelly Brock, Yarin Gal, and Debora S Marks. Disease variant prediction with deep generative models of evolutionary data. Nature , 599(7883):91-95, 2021.

Elodie Laine, Yasaman Karami, and Alessandra Carbone. Gemme: a simple and fast global epistatic model predicting mutational effects. Molecular biology and evolution , 36(11):26042619, 2019.

Jung-Eun Shin, Adam J Riesselman, Aaron W Kollasch, Conor McMahon, Elana Simon, Chris Sander, Aashish Manglik, Andrew C Kruse, and Debora S Marks. Protein design and variant prediction using autoregressive generative models. Nature communications , 12(1):2403, 2021.

Ethan C Alley, Grigory Khimulya, Surojit Biswas, Mohammed AlQuraishi, and George M Church. Unified rational protein engineering with sequence-based deep representation learning. Nature methods , 16(12):1315-1322, 2019.

Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Jason Liu, Demi Guo, Myle Ott, C Lawrence Zitnick, Jerry Ma, et al. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. Proceedings of the National Academy of Sciences , 118(15), 2021.

Ali Madani, Bryan McCann, Nikhil Vijay Naik, Nitish Shirish Keskar, Namrata Anand, Raphael R. Eguchi, Po-Ssu Huang, and Richard Socher. Progen: Language modeling for protein generation. bioRxiv , 2020.

Roshan M Rao, Jason Liu, Robert Verkuil, Joshua Meier, John Canny, Pieter Abbeel, Tom Sercu, and Alexander Rives. Msa transformer. In Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 8844-8856. PMLR, 18-24 Jul 2021.

Pascal Notin, Mafalda Dias, Jonathan Frazer, Javier Marchena Hurtado, Aidan N Gomez, Debora Marks, and Yarin Gal. Tranception: protein fitness prediction with autoregressive transformers and inference-time retrieval. In International Conference on Machine Learning , pages 16990-17017. PMLR, 2022.

Pascal Notin, Lood Van Niekerk, Aaron W Kollasch, Daniel Ritter, Yarin Gal, and Debora S Marks. Trancepteve: Combining family-specific and family-agnostic models of protein sequences for improved fitness prediction. bioRxiv , pages 2022-12, 2022.

Faruck Morcos, Andrea Pagnani, Bryan Lunt, Arianna Bertolino, Debora S. Marks, Chris Sander, Riccardo Zecchina, José Nelson Onuchic, Terence Hwa, and Martin Weigt. Directcoupling analysis of residue coevolution captures native contacts across many protein families. Proceedings of the National Academy of Sciences , 108:E1293 - E1301, 2011.

John Ingraham, Vikas K. Garg, Regina Barzilay, and T. Jaakkola. Generative models for graph-based protein design. In DGS@ICLR , 2019.

G Baumann, Cornelius Frömmel, and Chris Sander. Polarity as a criterion in protein design. Protein engineering , 2 5:329-34, 1989.

Bowen Jing, Stephan Eismann, Pratham N. Soni, and Ron O. Dror. Learning from protein structure with geometric vector perceptrons. In International Conference on Learning Representations , 2021.

Zuobai Zhang, Minghao Xu, Arian Jamasb, Vijil Chenthamarakshan, Aurelie Lozano, Payel Das, and Jian Tang. Protein representation learning by geometric structure pretraining. In The Eleventh International Conference on Learning Representations , 2023.

Chloe Hsu, Robert Verkuil, Jason Liu, Zeming Lin, Brian Hie, Tom Sercu, Adam Lerer, and Alexander Rives. Learning inverse folding from millions of predicted structures. In International conference on machine learning , pages 8946-8970. PMLR, 2022.

Justas Dauparas, Ivan Anishchenko, Nathaniel Bennett, Hua Bai, Robert J Ragotte, Lukas F Milles, Basile IM Wicky, Alexis Courbet, Rob J de Haas, Neville Bethel, et al. Robust deep learning-based protein sequence design using proteinmpnn. Science , 378(6615):49-56, 2022.

Kevin K Yang, Niccolò Zanichelli, and Hugh Yeh. Masked inverse folding with sequence transfer for protein representation learning. Protein Engineering, Design and Selection , 36:gzad015, 2023.

Steffanie Paul, Aaron Kollasch, Pascal Notin, and Debora Marks. Combining structure and sequence for superior fitness prediction. In NeurIPS 2023 Generative AI and Biology (GenBio) Workshop , 2023.

Matteo Cagiada, Sergey Ovchinnikov, and Kresten Lindorff-Larsen. Predicting absolute protein folding stability using generative models. bioRxiv , 2024.

Jin Su, Chenchen Han, Yuyang Zhou, Junjie Shan, Xibin Zhou, and Fajie Yuan. Saprot: Protein language modeling with structure-aware vocabulary. In The Twelfth International Conference on Learning Representations , 2024.

Yang Tan, Bingxin Zhou, Lirong Zheng, Guisheng Fan, and Liang Hong. Semantical and topological protein encoding toward enhanced bioactivity and thermostability. bioRxiv , pages 2023-12, 2023.

Jun Cheng, Guido Novati, Joshua Pan, Clare Bycroft, Akvil˙ e Žemgulyt˙ e, Taylor Applebaum, Alexander Pritzel, Lai Hong Wong, Michal Zielinski, Tobias Sargeant, et al. Accurate proteomewide missense variant effect prediction with alphamissense. Science , 381(6664):eadg7492, 2023.

Pablo Gainza, Freyr Sverrisson, Frederico Monti, Emanuele Rodola, D Boscaini, Michael M Bronstein, and Bruno E Correia. Deciphering interaction fingerprints from protein molecular surfaces using geometric deep learning. Nature Methods , 17(2):184-192, 2020.

Natalie L Dawson, Tony E Lewis, Sayoni Das, Jonathan G Lees, David Lee, Paul Ashford, Christine A Orengo, and Ian Sillitoe. Cath: an expanded resource to predict protein function through structure and sequence. Nucleic acids research , 45(D1):D289-D295, 2017.

Pascal Notin, Aaron Kollasch, Daniel Ritter, Lood Van Niekerk, Steffanie Paul, Han Spinner, Nathan Rollins, Ada Shaw, Rose Orenbuch, Ruben Weitzman, et al. Proteingym: largescale benchmarks for protein fitness prediction and design. Advances in Neural Information Processing Systems , 36, 2023.

Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, et al. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science , 379(6637):1123-1130, 2023.

Vignesh Ram Somnath, Charlotte Bunne, and Andreas Krause. Multi-scale representation learning on proteins. Advances in Neural Information Processing Systems , 34:25244-25255, 2021.

Roshan Rao, Nicholas Bhattacharya, Neil Thomas, Yan Duan, Xi Chen, John Canny, Pieter Abbeel, and Yun S Song. Evaluating protein transfer learning with tape. In Advances in Neural Information Processing Systems , 2019.

Ahmed Elnaggar, Michael Heinzinger, Christian Dallago, Ghalia Rehawi, Wang Yu, Llion Jones, Tom Gibbs, Tamas Feher, Christoph Angerer, Martin Steinegger, Debsindhu Bhowmik, and Burkhard Rost. Prottrans: Towards cracking the language of lifes code through self-supervised deep learning and high performance computing. IEEE Transactions on Pattern Analysis and Machine Intelligence , pages 1-1, 2021.

Vladimir Gligorijevi´ c, P Douglas Renfrew, Tomasz Kosciolek, Julia Koehler Leman, Daniel Berenberg, Tommi Vatanen, Chris Chandler, Bryn C Taylor, Ian M Fisk, Hera Vlamakis, et al. Structure-based protein function prediction using graph convolutional networks. Nature communications , 12(1):1-14, 2021.

Pedro Hermosilla, Marco Schäfer, Matˇ ej Lang, Gloria Fackelmann, Pere Pau Vázquez, Barbora Kozlíková, Michael Krone, Tobias Ritschel, and Timo Ropinski. Intrinsic-extrinsic convolution and pooling for learning on 3d protein structures. International Conference on Learning Representations , 2021.

Can (Sam) Chen, Jingbo Zhou, Fan Wang, Xue Liu, and Dejing Dou. Structure-aware protein self-supervised learning. Bioinformatics , 39, 2022.

Zuobai Zhang, Minghao Xu, Aurelie Lozano, Vijil Chenthamarakshan, Payel Das, and Jian Tang. Pre-training protein encoder via siamese sequence-structure diffusion trajectory prediction. In Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Freyr Sverrisson, Jean Feydy, Bruno E Correia, and Michael M Bronstein. Fast end-to-end learning on protein surfaces. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 15272-15281, 2021.

Vincent Mallet, Souhaib Attaiki, and Maks Ovsjanikov. Atomsurf: Surface representation for learning on protein structures. arXiv preprint arXiv:2309.16519 , 2023.

Zichen Wang, Steven A Combs, Ryan Brand, Miguel Romero Calvo, Panpan Xu, George Price, Nataliya Golovach, Emmanuel O Salawu, Colby J Wise, Sri Priya Ponnapalli, et al. Lm-gvp: an extensible sequence and structure informed deep learning framework for protein property prediction. Scientific reports , 12(1):6832, 2022.

Zuobai Zhang, Chuanrui Wang, Minghao Xu, Vijil Chenthamarakshan, Aurelie Lozano, Payel Das, and Jian Tang. A systematic study of joint representation learning on protein sequences and structures. arXiv preprint arXiv:2303.06275 , 2023.

Youhan Lee, Hasun Yu, Jaemyung Lee, and Jaehoon Kim. Pre-training sequence, structure, and surface features for comprehensive protein representation learning. In The Twelfth International Conference on Learning Representations , 2024.

Fang Wu, Lirong Wu, Dragomir Radev, Jinbo Xu, and Stan Z Li. Integration of pre-trained protein language models into geometric deep learning networks. Communications Biology , 6(1):876, 2023.

Kevin K Yang, Zachary Wu, and Frances H Arnold. Machine-learning-guided directed evolution for protein engineering. Nature methods , 16(8):687-694, 2019.

Sam Gelman, Sarah A Fahlberg, Pete Heinzelman, Philip A Romero, and Anthony Gitter. Neural networks to learn protein sequence-function relationships from deep mutational scanning data. Proceedings of the National Academy of Sciences , 118(48):e2104878118, 2021.

Christian Dallago, Jody Mou, Kadina E Johnston, Bruce J Wittmann, Nicholas Bhattacharya, Samuel Goldman, Ali Madani, and Kevin K Yang. Flip: Benchmark tasks in fitness landscape inference for proteins. bioRxiv , pages 2021-11, 2021.

Pascal Notin, Ruben Weitzman, Debora Marks, and Yarin Gal. Proteinnpt: improving protein property prediction and design with non-parametric transformers. Advances in Neural Information Processing Systems , 36, 2023.

Erik Nijkamp, Jeffrey A Ruffolo, Eli N Weinstein, Nikhil Naik, and Ali Madani. Progen2: exploring the boundaries of protein language models. Cell systems , 14(11):968-978, 2023.

Kevin K Yang, Nicolo Fusi, and Alex X Lu. Convolutions are competitive with transformers for protein sequence pretraining. Cell Systems , 15(3):286-294, 2024.

Timothy Truong Jr and Tristan Bepler. Poet: A generative model of protein families as sequences-of-sequences. Advances in Neural Information Processing Systems , 36, 2023.

Michel Van Kempen, Stephanie S Kim, Charlotte Tumescheit, Milot Mirdita, Jeongjae Lee, Cameron LM Gilchrist, Johannes Söding, and Martin Steinegger. Fast and accurate protein structure search with foldseek. Nature Biotechnology , 42(2):243-246, 2024.

Douglas M Fowler and Stanley Fields. Deep mutational scanning: a new style of protein science. Nature methods , 11(8):801-807, 2014.

Carl Edward Rasmussen and Hannes Nickisch. Gaussian processes for machine learning (gpml) toolbox. The Journal of Machine Learning Research , 11:3011-3015, 2010.

Yueqi Cao, Didong Li, Huafei Sun, Amir H Assadi, and Shiqiang Zhang. Efficient curvature estimation for oriented point clouds. stat , 1050:26, 2019.

Jian Sun, Maks Ovsjanikov, and Leonidas Guibas. A concise and provably informative multiscale signature based on heat diffusion. In Computer graphics forum , volume 28, pages 1383-1392. Wiley Online Library, 2009.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

Jeffrey I Boucher, Daniel NA Bolon, and Dan S Tawfik. Quantifying and understanding the fitness effects of protein mutations: Laboratory versus nature. Protein Science , 25(7):1219-1226, 2016.

C. Anders Olson, Nicholas C. Wu, and Ren Sun. A comprehensive biophysical description of pairwise epistasis throughout an entire protein domain. Current Biology , 24(22):2643-2651, November 2014.

A Broader Impact

The main objective of this project is to model protein fitness landscape more accurately with multiscale representation learning. Our approach utilizes structural information in the CATH dataset to build a multi-scale encoder with both sequence, structure and surface features. This advantage allows for more comprehensive analysis of protein research and holds potential benefits for various real-world applications, including protein engineering, sequence and structure design. It is important to acknowledge that powerful fitness prediction models can potentially be misused for harmful purposes, such as the design of dangerous drugs. We anticipate that future studies will address and mitigate these concerns.

CBPA2_HUMAN

NPC1_HUMAN

B2L11_HUMAN

S22A1_HUMAN

CUE1_YEAST

SC6A4_HUMAN

F7YBW7_MESOW

A0A247D711_LISMN

RCD1_ARATH

VKOR1_HUMAN

MLAC_ECOLI

NCAP_I34A1

I6TAH8_I68A0

F7YBW8_MESOW

ACE2_HUMAN

CD19_HUMAN

CAS9_STRP1

0.0

0.2

0.4

0.6

0.8

SpearmanR

ESM

S2F

S3F

S3F-MSA

B Additional Experimental Results

B.1 Performance on OOD test sets

In Section 4.5, we select 23 out-of-distribution 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.

B.2 Statistical Significance Analysis

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.

B.3 Ablation Study

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.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The main claims about the experimental results of our methods are supported in Section 4.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The limitations of this paper are discussed in Section 5.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: This paper does not include theoretical results.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The implementation details are provided in Section 3.

5. 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: [No]

Justification: The code, data and pre-trained model weights will be released upon acceptance.

6. Experimental Setting/Details

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: The evaluation setup are provided in Sections 3.5 and 4.1.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: The statistical significance will be reported when the results are released on the ProteinGym benchmark.

8. 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: The information about compute resources is provided in Section 3.5.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research conform with the NeurIPS Code of Ethics.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: The broader impacts are discussed in Section A.

11. 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)?

Answer: [NA]

Justification: The paper poses no such risks.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: The license of used datasets are provided.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper does not release new assets.

14. Crowdsourcing and Research with Human Subjects

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: The paper does not involve crowdsourcing nor research with human subjects.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [No]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n(1) I advocate the research direction of pretraining on all feasible protein modalities, containing sequence, structure, and surfaces. Each representation form has its strengths and an incorporation is beneficial for the zero-shot protein fitness prediction. \n\n(2) The author conducts extensive ablation studies over function type, MSA depths, taxon, and mutant depth. This helps readers to better understand the impact of each component of their model. The OOD examination of unseen protein families is also encouraging to demonstrate the superiority. \n\n() In section 4.4, the author navigated the impact of structure quality and observed a clear performance drop when lower-quality structures were used. This phenomenon underscores the importance of accurate structures for fitness prediction. Notably, there have already been some studies that try to bridge the gap between representations in real and predicted structures [A], and I would recommend the author take a look if seeking further performance improvement. \n\n[A] Protein 3D Graph Structure Learning for Robust Structure-Based Protein Property Prediction. AAAI 2024. \n\n(3) The visualization and experimental analysis are elegant. I learned a lot from it.\n\nWEAKNESSES:\n(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.   \n\nBesides, 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. \n\nLast 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. \n\n[A] Protein Representation Learning by Geometric Structure Pretraining. ICLR 2023. \n\n\n[B] Pre-training Sequence, Structure, and Surface Features for Comprehensive Protein Representation Learning. ICLR 2024. \n\n[C] Rotamer Density Estimator is an Unsupervised Learner of the Effect of Mutations on Protein-Protein Interaction. ICLR 2023. \n\n[D] Learning to design protein-protein interactions with enhanced generalization. ICLR 2024.\n\n(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. \n\n[A] Integration of Pre-trained Protein Language Models into Geometric Deep Learning Networks. Communications Biology 2023. \n\n\n(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?\n\nQUESTIONS:\n(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? \n\n(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. \n\n(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? \n\n[A] Thermodynamics-inspired Structure Hallucination for Protein-protein Interaction Modeling. ICLR submission. \n\n(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.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-14"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n**Originality**\n\nThis is one of the first papers that integrates pLMs, protein structure and surface representations together to predict protein fitness and shows improvement on benchmarks. It is also novel to combine MSA information with it to further improve predictive power of the model.\n\n**Quality**\n\nThe submission is technically sound with a detailed description of their framework and extensive experimental results on benchmark datasets with appropriate reasoning and explanation.\n\n**Clarity**\n\nThe paper is clear to read and easy to follow. Presentation of results are simple yet effective. \n\n**Significance**\n\nThis method achieves SOTA performance on protein fitness prediction benchmark and provides a standardized way to integrate three modalities of protein together for a unified representation.\n\nWEAKNESSES:\n1. 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.\n\n2. 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.\n\n3. Statistical results should be provided for figure 2 when possible. How is the variation of performance on 217 assays? \n\nReferences:\n[1] https://arxiv.org/pdf/2309.16519\n\nQUESTIONS:\n1. line 195. ESM features are integrated into surface representations. Why not just use surface feature itself because it contains chemistry and geometry information already?\n\n2. 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?\n\n3. line 230. What percentage of results are output from ESM?\n\n4. Figure 3. Why do some residues have negative correlation to model scores? Is there any quantitative analysis of differences rather than just qualitative visualization?\n\n5. line 347. 'ignoring side-chain information'. If side chain is ignored, how to get surface information?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe paper is well written and presents a parameter-efficient, novel way to include protein surface information into a model for fitness prediction. The paper contains results that will be interesting to the protein modeling community, showing that modeling the surface explicitly improves function prediction beyond other baseline methods which include structure information. It seems very interesting to investigate why such surface information is not easily captured by other methods that do utilize structure information.\n\nWEAKNESSES:\n1. 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.\n2. 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.\n\nQUESTIONS:\n1. 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.\n2. What is your current intuition about why structure-aware fitness prediction models are unable to leverage surface information to its full extent?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nSignificance: Perhaps the strongest point of the paper is that it clearly shows they can leverage structure and surface data to improve the zero-shot performance over sequence-only methods. The improvement is relatively small compared to models that only use sequence however. \n\nClarity: Paper is well written, has a good flow and logical steps. \n\nSoundness: For the most part the right experiments are being done, the data benchmarks are correct and relevant models are tested (see weaknesses for a notable exception).\n\nWEAKNESSES:\nNovelty/Originality:\nI 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. \n\nYes 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.)\n\nFrom the abstract\n“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. “\n- 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? \n- 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.\n\nQUESTIONS:\nI think the authors really need to contrast this work with reference 33 and argue why it’s a meaningful conceptual or performance advance. \n\nI think the paper would be better if the authors would do at least two ablation studies where at least \n- structure is dropped\n- both structure and surface are dropped  (dropping sequence is more complex but it may also be interesting). \n\nIt 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.\n\nLIMITATIONS:\nConceptually 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.\n\n\n:::::\n\nUpdating my previous score to 5", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}