Multi-Scale Representation Learning for Protein Fitness Prediction

Thank you for your insightful comments and suggestions. We have replied to each of your questions below and included additional analyses that we believe greatly strengthen our submission.

C1: Missing of some closely related baselines and relevant work.

Thank you for pointing out additional recent baselines we could have considered. We follow your suggestion and compare against 4 additional baselines: ESM-GearNet, ESM-GearNet-Edge, RDE and PPIFormer. We could not include ProteinINR, since no public code is currently available for this work. We report the performance of these new baselines in the following table.

Table A. Average spearman correlation over 217 ProteinGym assays

Note that all the new baselines significantly underperform our suggested methods. RDE and PPIFormer are specifically designed for predicting the mutational effects on protein-protein interactions, a task that fundamentally differs from protein fitness prediction. We had initially considered GearNet as our structure backbone in earlier stages of development, but eventually switched to GVP as it provided superior fitness prediction performance, as can be seen in the table above.

C2: 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].

Thank you for the suggestion -- we will include this work in the background section of the final version.

C3: There are two doubts about dmasif surface generation: (1) ignoring side-chain atoms makes the surface smaller. (2) randomness in surface generation. Has the author adopted software like PyMol?

Several great points! We answer your questions as follows:

• Ignoring side-chain atoms indeed makes the surface smaller, but results in significant computational efficiency for surface generation and message passing. Our results demonstrate that we can achieve significant performance gains with limited computational overhead, making our implementation appealing to practitioners. If we set aside any computational consideration, an ideal modeling strategy would likely involve all-atom mutant structures and surfaces. We plan to explore such approaches in future work.
• Randomness in surface generation is indeed unavoided. To mitigate this effect, we use a high-resolution surface (resolution = 1.0Å, with 6K-20K points) and construct a dense surface-to-backbone correspondence graph (using 20 nearest surface points). This approach ensures that our learned surface representations remain robust across different surfaces. Given your feedback, we tested S3F using surfaces generated with five different seeds and find that the standard deviation of the overall performance is approximately 0.001, thereby confirming the robustness of our learned surface features.
• There are several good software options for surface generation, including PyMol. We chose dMaSIF as it offered efficient GPU implementation and could easily be integrated into our codebase, but other software suites could have been similarly used.

C4: Can you please explain the details of the alignment part?

Since the model output from EVE are unnormalized delta ELBOs, model predictions from EVE and S3F are on unrelated scales. To facilitate model ensembling we thus first standard normalize model predictions separately, then take their arithmetic average. We will add this clarification in the revision.

C5: "Multi-modality" would be more appropriate than “multi-scale” in the title.

Thank you for the suggestion. We will change the title accordingly in the final version.

C6: 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?

Thank you for bringing this work to our attention. The idea looks very interesting and related. We will reference this paper in the final version and consider the idea in the future work.

C7: I would recommend the author use "mutant effect prediction" instead of "fitness prediction" to better depict the target problem.

We used the particular terminology used in the ProteinGym [1] work for consistency, although the two expressions are used synonymously in the relevant literature.

[1] ProteinGym: Large-Scale Benchmarks for Protein Fitness Prediction and Design

Thank you for your insightful comments and suggestions. We have responded to each of your questions below and included additional analyses that we believe significantly improve our submission.

C1: The idea of integrating surface into protein representation is not novel, for example [1], but is neither acknowledged nor benchmarked in the paper.

We kindly refer the reviewer to our clarifications regarding the novelty and contributions from our work in our response to all reviewers above. Our background section (section 2) referenced several prior works that had already leveraged structure and surface features for general protein representation learning -- AtomSurf belongs to the same category, and we will include it as an additional reference in the revision. However, we would like to reiterate that we did not benchmark against any of these prior works leveraging surface features since none of them, including AtomSurf, supports (zero-shot) fitness prediction, which is the core task we focus on. The novelty of our work comes through the careful architecture design to best leverage structure and surface features to achieve state-of-the-art protein fitness prediction performance.

C2: The advantage brought by surface representation is more incremental.

Please review our clarifications regarding the significance of our performance lift in our response to all reviewers above (point 2). The advantage brought by incorporating surface features in terms of fitness prediction performance is substantial: +0.02 Spearman between S2F and S3F across the 217 assays from ProteinGym. This is as much as the performance lift conferred by incorporating structural features (S2F vs ESM2). Together, structural and surface features provide a performance lift (+0.04 Spearman) that is comparable to factoring in epistasis vs not (Potts vs PSSM), which can hardly be characterized as incremental to fitness prediction performance.

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

Thank you for suggesting this analysis. As discussed in our overall response to all reviewers, we compute the non-parametric bootstrap standard error for the difference in Spearman performance between a given model and the best overall model, using 10k bootstrap samples. We do this given the inherent variation in experimental noise and data quality across assays, which leads all models to consistently perform lower on certain assays and higher on others. By focusing on the standard error of the difference in model performance, we abstract away this assay-dependent variability, and provide a quantity that better reflects the statistical significance of the performance lift between our best model and any other baseline in ProteinGym. Please see the detailed results in the attached pdf, which confirms that the performance lift of our best model (S3F-MSA) vs the prior best baseline (SaProt) is statistically significant.

C4: ESM features are integrated into surface representations. Why not just use surface feature itself because it contains chemistry and geometry information already?

We chose to combine ESM features with geometric features (such as Gaussian curvatures and Heat Kernel Signatures) to initialize the surface features. We believe that ESM can capture more informative chemical features by identifying residue types and co-evolutionary information, making it a strong complement to geometric surface features.

C5: How to ensure removing 20 closest surface points for each selected residue is sufficient?

We empirically chose this number of closest surface points to remove based on validation accuracy during pre-training. For instance, without removing these points, the pre-training accuracy quickly reaches 100%, suggesting likely information leakage. However, after removing these points, the S3F accuracy drops down to 52.4%, comparable to that of S2F (51.0%).

C6: What percentage of results are output from ESM?

Out of 2.46 million mutations, approximately 0.35 million (14%) rely on ESM predictions only due to low-quality structures. We also benchmarked the results of S2F and S3F without pLDDT filtering, yielding aggregate Spearman performance scores of 0.449 and 0.461, respectively. These results remain significantly higher than the ESM performance (0.414 Spearman).

C7: Why do some residues have negative correlation to model scores? Is there any quantitative analysis of differences rather than just qualitative visualization?

These negative correlations are already present in the ESM-based predictions, likely reflecting residue type preferences in ESM that do not align with experimental results. By introducing structural and surface features, our methods mitigate this effect to some extent. To quantitatively evaluate how much S2F and S3F improve over ESM methods, we calculated the average Spearman correlation in the regions of interest (residues 234-252 and 266-282). The results for ESM, S2F, and S3F are 0.301, 0.397, and 0.443 respectively, demonstrating that the introduction of structural and surface features can more effectively capture epistatic effects.

C8: If the side chain is ignored, how to get surface information?

Since point mutations can significantly alter the side-chain structure, we chose to only keep the backbone structure for surface generation resulting in a smaller surface graph that is broadly applicable to all mutated sequences for the same protein family. This approach significantly reduces the cost of surface generation and message passing, making model training and inference more efficient, while yielding the significant performance improvement discussed above and keeping computations tractable on the hardware we had access to.

Thank you for your insightful comments and suggestions. We provide detailed responses to your questions below.

C1: Lack of analysis for various hyperparameters of the model. It would be nice to see more in-depth experimentation and/or ablation of different parts of the S3F model..

During model development, we performed our main hyperparameter search based on the validation set accuracy for the S2F model (eg., number of layers, number of hidden dimensions), and used these same optimal values for S3F given the substantial compute costs required to process the surface features for pre-training. We provide below detailed ablation results for these model hyperparameters, as well as modality ablations. While the former tend to have a relatively marginal impact on the downstream performance, the latter clearly demonstrate the contributions of structural and surface information in S3F.

Table A. Ablation Study and Hyperparameter Analysis.

Note: we used 5 layers in our final S3F models, instead of the 8 layers that provide marginally better performance for S2F, due to GPU memory bottleneck when leveraging surface features.

C2: 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.

Surface message-passing is intended to capture fine-grained structural aspects that complement the coarse-grained features learned through structure message-passing. Additional ablations in Table A above clarify the respective benefits from including the different modalities.

C3: 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.

We report the ensembling results with these 3 baselines in the table below. We observe that S3F outperforms existing baselines in both regimes: When no MSA is available, S3F outperforms the original MIF-ST, ProtSSN and SaProt. When MSAs are available and we are willing to train protein-specific alignment-based models (eg. EVE) to increase fitness prediction performance, S3F-MSA leads to statistically significant higher performance over MIF-ST-MSA, ProtSSN-MSA and SaProt-MSA.

Table B. Model Performance with EVE Ensembling.

C4: What is your current intuition about why structure-aware fitness prediction models are unable to leverage surface information to its full extent?

Explicitly including relevant information is a way to inject inductive bias into a task. For example, while structural information can be implicitly encoded in protein language models, we still want to use structure to augment PLM. Similarly, although surface information can be derived from protein structures, explicitly including it helps us better learn useful features for binding assays.

Thank you for the very thoughtful comments and suggestions. We address each of your questions below, including several additional analyses which we believe significantly strengthen our submission.

C1: The improvement is relatively small compared to models that only use sequence.

The best fitness prediction models using only sequence information are TranceptEVE [1] and GEMME [2] -- two methods leveraging multiple sequence alignments. Their aggregate Spearman performance on the 217 assays from the ProteinGym zero-shot substitution benchmark is of 0.456 for both (Table 1). In contrast, the performance of our best model (S3F-MSA) is 0.496. This is a massive performance lift (+0.04 Spearman), especially since averaged across 217 diverse DMS assays. To put things in perspective: This performance lift is larger than the performance lift between a PSSM and a Potts model (EVmutation) (0.359 vs 0.395) -- difference which can hardly be characterized as relatively small, given the critical importance of epistasis in protein fitness Over the past 5+ years, all 50+ fitness prediction baselines that were introduced after DeepSequence, the first deep learning-based approach for fitness prediction in 2018 [3], allowed to move from 0.419 to 0.457 (SaProt) aggregate Spearman. This also represents +0.04 Spearman performance lift -- our work extends over that best baseline by an extra +0.04 Spearman.

Finally, we note that the current best baseline on ProteinGym that combines sequence and structure features is SaProt, with an 0.457 aggregate Spearman. The corresponding performance lift over the best sequence-only methods (+0.001 Spearman) is marginal, illustrating the complexity of effectively integrating new data modalities to boost zero-shot fitness prediction performance.

C2: Novelty of the paper compared with HoloProt.

We kindly refer the reviewer to our clarifications regarding the novelty and contributions from our work in our response to all reviewers above. In particular, we would like to reiterate that, while our work is not the first to leverage structure and surface features for protein representation learning, it is the first to focus on and support zero-shot protein fitness prediction.

This is a significant difference as methods developed for general protein representation learning do not enable zero-shot fitness prediction or, if they do, typically underperform methods specifically developed for it [4]. For instance, in the work introducing HoloProt, Somnath et al. describe an approach for general protein representation learning using structure and surface features. The authors show that these representations can then be leveraged (via supervision) for several downstream tasks -- namely, ligand binding affinity prediction and an Enzyme-Catalyzed Reaction Classification (which the authors refer to as “function prediction” in their abstract). Nothing in the work from Somnath et al. covers fitness prediction, let alone in the zero-shot context, and it is not clear how one could use the corresponding architecture to address the task that we focus on in our work.

C3: Ablation study for the effect of structure and surface.

Thank you for the suggested ablation analyses. We provide the corresponding results in the table below, confirming the performance lift from the various modalities involved.

Table A. Ablation Study.

Notes: 1) The ablation dropping both structure and surface features was already in our original submission, and corresponds to just using the underlying pLM (ie., ESM2). 2) 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.

C4: 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.

We believe this statement is inaccurate as the performance lift of S2F over ESM2 (the underlying pLM) is of +0.04 aggregate Spearman (0.454 vs 0.414), which, as we argue in our response to your first comment above (C1), corresponds to a transformational performance lift, on par with factoring in epistasis vs not, or commensurate with the performance lift obtained from over 5 years of deep learning literature for protein fitness prediction.

[1] Notin, et al. "TranceptEVE: Combining family-specific and family-agnostic models of protein sequences for improved fitness prediction."

[2] Laine, et al. "GEMME: a simple and fast global epistatic model predicting mutational effects." Molecular biology and evolution.

[3] Riesselman, Adam J., John B. Ingraham, and Debora S. Marks. "Deep generative models of genetic variation capture the effects of mutations." Nature methods

[4] Notin, et al. "Proteingym: Large-scale benchmarks for protein fitness prediction and design." NeurIPS.