Lightweight Equivariant Graph Representation Learning for Protein Engineering
Authors: Bingxin Zhou
Outongyi Lv
Kai Yi
Xinye Xiong
Pan Tan
Liang Hong
Yu Guang Wang
Venue: Submitted to ICLR 2023 (2023)
Keywords: graph neural networks

Abstract: This work tackles the issue of directed evolution in computational protein design that makes accurate predictions of the function of a protein mutant. We design a lightweight pre-training graph neural network model for multi-task protein representation learning from its 3D structure. Rather than reconstructing and optimizing the protein structure, the trained model recovers the amino acid types and key properties of the central residues from a given noisy three-dimensional local environment. On the prediction task for the higher-order mutants, where many amino acid sites of the protein are mutated, the proposed training strategy achieves remarkably higher performance by 20% improvement at the cost of requiring less than 1% of computational resources that are required by popular transformer-based state-of-the-art deep learning models for protein design.

Under review as a conference paper at ICLR 2023
LIGHTWEIGHT EQUIVARIANT GRAPH REPRESENTA- TION LEARNING FOR PROTEIN ENGINEERING
Anonymous authors Paper under double-blind review
ABSTRACT
This work tackles the issue of directed evolution in computational protein design that makes accurate predictions of the function of protein mutant. We design lightweight zero-shot graph neural network model for multi-task protein rep- resentation learning from its 3D structure. Rather than reconstructing and op- timizing the protein structure, the trained model recovers the amino acid types and key properties of the central residues from given noisy three- ftree-dimensional local environment. On the prediction of higher-order mutations where multiple amino acid sites of the protein are mutated simultaneously. the proposed strat- egy achieves remarkably higher performance by 20% improvement at the cost requiring less than 1% of computational resources that are required by popular transformer-basee state-of-the. -he-ttttr deep learning models for protein design.
INTRODUCTION
Mutation is biological process where the amino acid (AA) of one multiple sites of specific protein is changed. While the wild-type proteins' functions do not always meet the demand of bio-engineering, is vital manually optimize the functionality, namely fitness, with favorable mutations so that they are applicable in designing antibodies (Wu et al. [0019] Pinheiro et al. 2021
Shan et 2022 or enzymes (Sato & Ishida 2019 Wittmann al. (021).
A protein usually constitutes hundreds to thousands of AAs, where each residue belongs to one of twenty AA types. To optimize a protein's 's functional fitness, a greedy search is usually conducted in the local sequence, where AA sites are mutated to proper AA types to render protein mutant with the highest gain-of-function (Rocklin et al. 2017 Such process called directed evolution Arnold (1998). To obtain mutant with great fitness, multiple AA sites (>5-10) of the protein need to be mutated, namely deep mutations (see Figure D It, however, requires enormous experimental costs, as the total number of potential combinations of mutations for deep mutants is astronomical.
Since impossible to conduct systematic experimental tests on all possible deep mutations, in silico examination of protein variants' fitness becomes highly desirable. A handful of deep learning methods have been developed to accelerate the discovery of advantageous mutants. For instance, Lu al. 2022) applied 3DCNN to identify new polymerase with advantageous single-site mutation and enhanced the speed of degrading PET, i.e., type of solid waste, by 7-8 times at 50%C. Luo al, (202I) proposed ECNET that predicts functional fitness for protein engineering with evolution- ary context. The model guides the engineering of TEM-1 3-lactamase and identifies variants with improved ampicillin resistance. [Thean al. (2022) enhanced SVD with deep learning to predict nuclease variants' activities in muuli--iie-saturaann mutagenesis libraries from and identified Cas9 nuclease variants that possess higher editing activity of derived base editors human cells.
Due to the scarcity of labeled protein data, researchers often pre-train an encoder for unsupervised learning with protein sequences or structures, and use the learned protein representations to train specific tasks, such as de novo protein design (Hsu et al. (022), mutation effect prediction (Ingra- ham et al. 2019 Jing et al. 2020 Meier et al. 2021 Notin et al. 2022). and higher-level structure prediction (Enngggar al. 2021). In the context of fitness prediction of mutation effect, existing methods usually transform the problem to mini-de novo design, which infers a specific AA type from its microenvironment, based protein learning methods rely heavily on multiple sequence alignment (MSA; Riesselman et al. (2018) Frazer al. [2021]: Rao al. (2021)) and protein language models (Elnaggar al.
analogously its neighboring AA types. Current state-of-the-as art sequence-

---

Under review as a conference paper at ICLR 2023
CAPED AM725
1C04 HUSAN
2006
2000
30000
10003
120000
10000
20003
30000
70000
140000
23003
330030
200030
Figure 1: Mutating more sites frequently results in a higher score, i.e., smaller rank value.
2021
Rives et al. 2021 Nijkamp et al. 2022 Brandes et al. 2022). While MSA helps capture im- portant evolutionary properties of the protein family, it nevertheless multiplies the requirements of computing resourees. The latter protein language models derived from natural language processing (NLP) encode sequence semantics and often need hundreds of GPU cards to train on hundreds of millions of proteins. Meanwhile, an autoregressive inference process is usually required along the entire protein sequence to score a mutation on a single site, which further slows down the inference
speed (Sato & Ishida 2019 Liu et al. 2022 Hsu et al. 2022 Notin et 2022). More importantly, when predicting the fitness of the higher-order mutants, most of these models made a crude assump- tion that the mutations on different sites happen sequentially or individually, which incorrect in most cases (Lehner [001]: Breen et al.) 2012). The ignored epistatic effects between different sites are potentially a key factor hindering the acquisition of favorable high-order mutants in directed evolution (Sarkisyan et al., 2016 Rollins et al. 2019).
Mutation of AA sites also occurs in nature, where an AA site might be mutated to any of the other 19 AA types in random manner. It is suggested by natural selection that only the mutants that exhibit the best fitness and fit the environment survive. As protein's 's functionality is determined by its structure, we encode the folded protein by a protein graph with AAs being graph nodes to provide an elegant 3D spatial description of the protein. The first-level information, such as AA types, spatial coordinates of Ca, and C-N angles between neighboring AAs, are embedded in node features. Altering AA types of a protein in nature can then be viewed as adding corruptions to the node features of the protein graph, and denoising the graph makes a remedy to search for mutants with the best fitness. We model the protein mutation effect prediction as a denoising problem with equivariant graph neural networks (Satorras et 2021). For given protein, the recovered predictions can be leveraged to forecast the fitness of deep mutational effects and discover favorable mutants.
Compared to existing state-of-the-art deep learning methods for mutation effect prediction, such as ESM-1 Iv (Meier et al. 2021) and ESM-IF1 (Hsu et al.) 2022), the designed lightweight equivariant graph neural network (LGN) stands out in three perspectives.
First, LGN improves generalization ability through the multi-task learning strategy and biological prior knowledge. The pre-trained model encodes the chemical and physical properties of given AA's microenvironmen with domain knowledge for practically meaningful representations.
assumptions by generating the probabilities of all the amino acid residues at time, which implements the joint distribution of all variations. In literature, the higher-order mutation effect usually approached by summing up log-odd- ratio scores of the corresponding individual single-site mutants. The linear combination over separately assigned predictions is unsubstantiated, as the independent mutations neglect the epistatic effect.
Secondly, LGN avoids the andependent-mutatian
Thirdly, LGN is efficient in both the training and inference phases. The spatial graph inputs portray the topological properties of proteins, which circumvents data augmentation that is typically required by sequence or grid representations. Equivariant message passing, alternatively, provides feature distillation unit with translation and rotation equivariance and encodes AAs' microenviron- ment defined by the protein graph's geometry.
2

---

Under review as a conference paper at ICLR 2023
RANK
layers
downstream
AA type
bloch kory
node enccddnn
SASA
Layer norm
layer
B-factor
posttoo
edge
dihedral
eistance
Figure 2: An illustration of the proposed LGN. The model is pre-trained with a set of protein graphs that are featured by (perturbed) node attributes and 3D positions with a multi-task learning strategy. A stack of EGC layers encodes rotation and translation equivariant structural representations for each node on individual graphs. Next, funne-con layers are employed to learn different labels, where the AA type prediction is used for suggesting top-ranked mutations.
2 ZERO-SHOT LEARNING FOR PROTEIN RECOVERY
The excessive cost in laboratory results in scarce mutation scanning data, especially deep mutant results. It thus favorable to pre-train zero-shot protein prediction model that can be generalized directly to an unseen task without any further supervision to specialize the model.
2.1 GRAPH REPRESENTATION OF PROTEIN STRUCTURE
For a given protein, we create a k-nearest neighbor (kNN) graph g = -V. to describe its 3D structure and molecular properties. Here each node v; V represents amino acid residue with X E R.34 node attributes constituting biochemical properties and geometric properties of amino acids. The former includes 20-dimensional attributes of one-hot encoded amino acid types (Xm), two scalars for each residue, i.e., solvent-accessible surface area (SASA) and the standardized crys- tallographic B-factor, and normalized ourface-aware node features. The geometric properties in- clude the direction position pos) of each residue by 3D coordinates of its q-carbon and the relative position of the amino acid in the protein chain (Xel) by the dihedral angles (sin, cos) o {Ø) vt computed from the backbone atom positions. For specific of the ith amino acid in the protein sequence, the dihedral angles are measured from Caa-1, N,, Cai, N1+1-
To build edge connections, we first define symmetric adjacency matrix A with the kNN-graph to capture the nodes" microenvironment, i.e., each node connected to up to other nodes the graph that has the smallest Euclidean distance over other nodes, and the distance is smaller than certain cutoff A3).A. Consequently, v, and f are connected to each other, we have A.j = Ayi * 0. The inter-atomin distances, 12
edge attributes E R3 feature the connected edges in E, including local N-C positions, and the relative position in the protein sequence in 66-dimensions.
2.2 PRE-TRAINING WITH PRIOR DOMAIN KNOWLEDGE FOR BETTER PROTEIN FITNESS
The wild-type proteins suffer from random perturbations or mutations that not every AA site has the best AA type (Liu et al., 2022). To this end, pre-train model with multitask learning strategy, which removes the natural corruptions and predicts key protein properties to help encode the microenvironment of the stabilized proteins of interest.
AA Type Denoising
We refine Lea the AA type node, to E.a with Bernoulli noise, i.c.,
{ pr ).a) + pm, m T2, T2,..., Tn),
where the confidence level is a tunable parameter that controls the proportion of residues that are "noise-free". The probability for the residue to become partieular type depends the distribution
3

---

Under review as a conference paper at ICLR 2023
of the 20 types (nn, R1, T2, T2, Tx), which involves prior knowledge in molecular biology. This paper defines the distribution by the observed probability density of amino acid types in wild-type proteins [ See Appendix E to better understand the influence of different confidence levels.
Geometric Properties Denoising For the continuous-valle valuued features, such as 3D coordinates and dihedral angles, an ii.d Gaussian noise is introduced, learning to remove which corresponds to approximating the data- data-generating force field of molecules (Zaidi et al. 2022). To be specific,
pas pos aee where (N, 13).
(2)
The noise effect is determined by a, which tunable to the scale of the noiseless raw feature.
Bio-chemistry Properties Recovery Aside from denoising the perturbed residues type and ge- ometric properties, other auxiliary tasks are introduced to help establish an expressive hidden mi- croenvironment representation. Specifically, SASA is known strongly influence AA type prefer- ences, and B-factors are associated with the conformations and mobility of the neighboring AA. We thus introduce inductive biases to the model by predicting these properties in the output.
Label Smoothing with Amino Acid Substitution Matrices Protein sequence alignments provide important insights for understanding gene and protein functions. The similarity measurement of an alignment protein sequence reflects the favors of all possible exchanges one amino acid with another. We employ BLOOSUM (Henikoff & Henikoff 1992), substitution matrix, account for the relative substitution frequencies and chemical similarity of AAs. The matrix is derived from the statistics for every conserved region of protein families in BLOCKS database. As AA sites more likely to be mutated to the AA type within the block of high similarity scores in the BLOSUM table, we hereby modify our loss function that mutation an AA type with higher similarity score accumulates smaller penalty than to the one with lower similarity score.
2.3 PROTEIN STRUCTURE REPRESENTATION WITH EQUIVARIANT GNNS
Proteins are structured in the 3-dimensional space, and is vital for the model to predict the same binding complex matter how the input proteins positioned and oriented. Instead of practicing expensive data augmentation strategies, we follow [Satorras et al. (202T) and construct SE(3)-equivariant neural layers for graph embedding. At the (th layer, an Equivariant Graph Con- volution (EGC) inputs a set of n hidden node properties embedding H' = { h²} as well the node coordinate embeddings XL {=1 . = for graph of nodes. The at- tributed edges are denoted as E (.... ei..... eij......... J. The target of EGC layer is to output transformation on the node feature embedding H and coordinate embedding X(+1, Concisely: H x EGC [ X! pos
....,,,,)
pos E. To achieve this, an EGC layer defines
my = de [ni, 11,, | x511² .ews)
of x Σ (xt x²) dx (mij)
jti
(3)
m. Σ mij
jfi
h Φn (h/. m.)
where de, db are respectively the edge and node propagation operations, such multi-layer percep- trons (MLPs). The on is an additional operation that projects the vector embedding mij to scalar value. The EGC layer preserves equivariance to rotations and translations on the set of 3D node coordinates Xposs while simultaneously performing invariance to permutations on the set nodes V in the same fashion GNNs.
'Retrieved from the folded protein dataset by AlphaFold2 Varadi alphafold. uk/
al. 2022
https:
4

---

Under review as a conference paper at ICLR 2023
2.4 MODEL OVERVIEW
Our model is depicted in Figure 2 We take a set of protein graphs with attributed nodes and edges, as well as each node's 3D coordinates, as the input to pre-train a zero-shot model. stack of EGC layers is trained to extract rotation and translation equivariant representations for each node on individual graphs. The hidden representation then sent to fully-connected layers to establish multiple outputs, such as AA type classification, SASA and B-factor prediction, and 3D coordinates denoising. The total loss for the multitask learning task is given by
Luval
Lau + A Caasa + A2 + A3 + A4
(4)
to balance different losses on auxiliary regres- sion tasks. These losses are measured by mean squared error (MSE) loss. For AA type classification, we measure its loss Caa by cross- cross-entropy with label smoothing technique (Szegedy et al.) 2016). The classification loss on an arbitrary node reads
where A; 1,...,4 4 are tunable
20
20
L. (1-E) Σ P( X;, E,) log go (Daa X, E.) +e +e Σ
u(Yall X, E,) log (Daml X, E:)
where p(yml X,, E,) denotes the ground-truth distribution and qu (bmm) x, E,) is the distribution of predicted labels following softmax function. In order improve the generalization and respect the prior biological knowledge, we modify the ground truth label distribution p(yas| E,) from the hard one-hot encoding to (1 =))(ya X,, ; E.) + Eu(yaa [x, E,) when the predicted jaa yea and zu eu(Vaa x,, E,) otherwise with some tolerance factor E. In particular, we define the distribution of ul uly(2,) by the BLOSUM substitution matrix.
3 RESULTS
3.1 EXPERIMENTAL SETUP
We train LGN on CATH v4.3.0 (Orengo et al. 1997> with artificial noise to predict AA type, 3D coordinates, dihedral angles, and chemical properties (SASA and B-factor). The hidden embeddings of amino acids are learned by by zero-shot prediction task for the fitness of mutation prediction with deep mutational scanning (DMS; Fowler & Fields (2014) datasets. The model performance is compared against popular state-of-the-arr language models and structure-enhancee models.
graph convolutions. The performance is validated
Baseline Models We compare with a diverse of state-of-the-ar models on the fitness of mutation effects prediction. In particular, DEEPSEQUENCE [Riesselman et al. 2018) trains VAE on protein- specific MSAS to capture higher-order interactions from the distribution an AA sequence. MSA TRANSFORMER Rao et al. (202T) is language model with aligned protein sequences of interest; ESM-1V (Meier et al. 2021) make zero-shot mutation predictions with masked language modeling; and ESM-IFI (Hsu et al.] 2022) predicts protein sequence with GVP (Jing et al. 2020). graph representation learning methods for vector and scalar features of protein graphs. Furthermore, both
TRANCEPTION (Notin et al. 2022 and PROGEN2 (Nijkamp et al. 2022 leverages autoregressive
language models to retrieve AA sequence without family-ppoctn MSAs.
Lightweight Equivariant Graph Neural Networks (LGN) To train our LGN framework, we first generate protein graphs for the sequences in CATH. See Appendix A for a detailed introduc- tion to the generation, and Appendix/[[oor a summary of the generated dataset. For the total number of 31, 848 protein graphs of 150 nodes on average, we randomly pick 500 graphs for validation and leave the remaining for model fitting. During the learning phase, we assign random perturbations to AA types and other features we mentioned earlier in Section 2 The noises are fixed to guaran- tee stable and comparable measurements at the validation step. The specific influence of different choices on the hyper- pyper-parameters (e.g.. the noise level) will be discussed later in this section and Ap- pendix E The main architecture constitutes stack of 6 EGC layers following fully- connecco layer to make predictions on the different learning tasks. On each node, the output is a vector rep- resentation consisting of 20 probabilities of the masked amino acid, 1 predicted SASA, 1 B- B-factors,
5

---

Under review as a conference paper at ICLR 2023
Table 1:
Performance Comparison of Baseline Models DMS assays prediction. Results with higher Spearman's 's correlation are preferred.
DEEPSEQUENCE
TRANCEPTION
PROGEN2
MSA TRANSFORMER
ESM-IV
ESM- IIFI
LGN (ours)
CAPSD
0.4419
0.2212
0.3589
DLG4 HUMAN
0,5013
0.6200
0.5712
0.4654
0.4654
0.6164
0.6197
F7YBW8
0,39.9
0,4280
0.3231
0.3483
0.2865
0.3714
0,4223
FYYBW8.EESOW
0.4208
0.4036
0.3231
0.3631
0.2522
0.3690
0.4210
GFF
0.6331
0.6647
0.6459
0,7069
0.7208
0.6203
0.6455
GRB2 HUMAN
0.3886
0.4441
0.5211
0.2808
0.3211
0.7033
0.6044
average correlation
0.4698
0.4177
0.4186
0.4511
0.3756
0.4714
0,5071
f The top three highlighted by First, Second, Third.
and 4 dihedral values (when applicable). The 3D coordinates derived directly from EGC out- puts. The loss function by Equation guides the backward propagation with ADAM (Kingma & Baj 2015) optimizer. The model is trained with 300 epochs with the initial rate set 0.001 and weight decay to 0.01. The learning rate is dampened to 0.0001 after 150 epochs.
Evaluation All the models are evaluated with deep mutational scanning (DMS) assays that assess a diverse set of 15 proteins, where 9 of them only contains single-site mutation scores, and 6 of them have both single-site and higher-order mutational records (see Appendix B The protein structures are folded from the provided sequence information with ALPHAFOLD2 (Jumper et al. 2021). fol- lowing the exact same pre-processing steps as CATH for generating protein graphs. The only difference is that we do not append artificial noises onto the test proteins, as assume they are already noisy. We then send the unmutated test proteins to the pre-trained LGN model and use the log-odd-ratio in Equation [5 of the predicted probabilities of AA types for suggesting the rank of deep mutations. The prediction performance evaluated on Spearman's 's correlation coefficient between the computational and experimental scores all the mutation combinations.
0.7
Spearman's Carrelation
0.6
0.5
Transformar
ESM-1v
0.4
DeepSerguence
0.3
Tranception
ProCun2
02
LGN
0.1
DLGA HUMAN
CPP
GR82 RUMAN
FrYDie MESOOW
CAPSO AAV2S
Bi STRSO
BLAT ECOLX
PLU PIEN RUMAN
31783 HAFARSTABILIZZE
HG. FLU
KKA2 KLEPN
CALMI HUMAN
BASH, HUMAN
Figure 3: Per task performance on the fitness of deep mutant prediction with pre-trained models. Each point indicates Spearman's correlation coefficients the corresponding protein. The left proteins contain higher-order mutations, and the right 9 proteins record shallow mutants.
3.2 FITNESS OF DEEP MUTANTS PREDICTION
The first experiment evaluates the fitness of proteins' mutation effects prediction, where the fitness scores are inferred directly from a pre-trained model without supervision on a task-specific model. We visualize the overall performance comparison on protein-wise Spearman's correlation coeffi cients in Figure 3 The deep mutation scores are reported in Table where LGN outperforms baseline methods and achieves at least comparable results in single-site mutant tests. Overall, our model achieves 0.5071 weighted average correlations on deep mutant effect predictions. While DEEPSEQUENCE achieves superior performance over the majority rest, it should be noticed that the model has to be trained every new protein, and it cannot be generated for other proteins. Also, the training speed and performance of the learned model rely heavily on the quality of the available MSA information, which can very lot different proteins.

---

Under review as a conference paper at ICLR 2023
Conteston Matris Mcbooe
1.6
U-focter
SASA
34
50
2
32
CST AAGDENNUURRE KN INELLVO LVWYT
30
-1
Figure 4: The three plots (from left right) display the confusion matrix of predicted AA types, and linear regression on the predicted SASA and B-factor, respectively.
3.3 PROTEIN RECOVERY
This experiment investigates the auxiliary learning tasks of the pre-trained model, including the learning performance in predicting AA types, SASA, and B-factor. In specific, we visualize the confusion matrix of the predicted AA types with respect to the ground-truth AA types to eval- uate the model's capability to recover from noisy sequences to the original sequence, i.e., the large values are accumulated the diagonal of the confusion matrix. For the SASA and B-factor predictions, we examine the R² of the predicted and the ground-truth values on CATH. The results are visualized in Figure 4 with denoised AA type, well as the predicted SASA and B-factor as the output tasks. The AA type prediction achieves high accuracy with the majority of predictions accumulated on the diagonal line. For the two regression tasks, we the true value and the pre- dicted value with linear regression. The estimated coefficients are 1.008 and 0.989 for B- B-factor and SASA, respectively. The p-value for both coefficients is 0.001. In addition, Pearson's s correlation coefficients for the two variants are respectively 0.884 and 791.
8.50
0,50
0.45
Spcarmin Carrelatio
(AA)
(AA- (ASA)
Ohd 5153-0 factor)
0:3
0.40
3Deddbeirall
0.30
0.35
0.30
0.7
0.9
01.5
Jamieln
Figure 5: Average Performance with different ps (left) and As (right) on the auxiliary tasks.
3.4 SELECTION CRITERTA ON THE ADDITIONAL.
As LGN introduces two additional hyper- pyrem-parameters in training the model, this section examines the influence of selecting different ps and Äs. Figure 5 visualizes a selection of model performance on the Spearman's 's correlation with variant ps and Às with wild-type noise distribution.
While the choice of p can be determined by prior knowledge regarding the quality of wild-type proteins, here treat p as a data-driven hyper-parameter be optimized during the model training. We exclude extremely small ps to avoid drastic perturbation rates and search for the optimal p (0.3, 0.4, 0.9, 1). The different choices on ps are validated with different learning tasks on the left side of Figure for higher-order mutants. In general, moderate p between 0.3 and 0.6 best suits the majority selection of learning modules and noise distribution. Based on the overall performance, suggest p- 0.6 the default value of the confidence level (See Appendix E for more results different types of perturbation noise and proteins).
We also investigate wide range of the choices of Às. For simplicity, let \ X₂ X (0.05, 0.1, 0.2, 0.5, 0.8, 5, 10) and fix X = 0.5. All the results are conducted under the recom- mended 0.6 with wild-type noise. We report the average performance deep mutants in the

---

Under review as a conference paper at ICLR 2023
Table 2: Comparison of baseline models. The train and inference speed tested on GFP.
model
DEEFSRQUENCE
MSA TRANS.
ESM-IV
ESM-IFI
TRANCEPTION
PROGEN2
LGN (ours)
ipput
sequence
sequence
sequence
sqqucnoetstucture
SeQuence
sequenee
structure
MSA
train pretein
miining
Uniref50 (2018-08)
Unire190
Uniref90-BFFD
(202-403)
v4.3.l
minnngg size (M)
45
98
12
249
>,,0000
0.03
max. Input
1,024
2, 2,77
# parameters (M)
4.3
100
650
142
700
2: 700
1.3
layers
1. 60000
12
32
82
dim.
100 2, 000
512
2,048
512
(training day)
18a V100
653
~100
0.17
(train)
32xV100
2xTPU-x
1x3090
preparing speed (sce) inference speod (sec)
6, 36i 25, 0G0
6, 300
6,900
A08
927
1,920
1,4
25
right plot of Figure which demonstrates a relatively flat and steady trend with a mild peak at 0.2, 0.5. Additional results are provided in Table 8f Appendix with various model setups.
3.5 INFERENCE SPEED
LGN consumes significantly fewer computational resources in training and inference. We compare the model scale, inference time, and prediction performance in Figure Table
The model size and the required resources with the baseline methods are provided by the authors. As the majority of models are pre-trained, we record the inference speed on single 3090 GPU. While time cost is significantly lower than experimental methods, we measure it to indicate the cost of forward propagation in one iteration, which can be viewed as the indirect empirical evidence of the training cost. As each protein requires independent inference progress, we hereby take GFP as an example protein sample. The protein constitutes 236 amino acid residues, and it has over 50, 000 mutant records (see Table 3 in Appendix B for more details). Note that: 1). The 2, 687 input token length only refers the maximum protein length used during training. In fact, the model itself can process large protein graphs containing over tens of thousands of amino acids. 2). The training speed and required resources for MSA TRANSFORMER are retrieved from Meier et al. (2021). The original work by Rao et al. (2021) only reports that they used 32 V100 GPUs for training, without revealing the training time.
Spearman Correlation
(LAMAAA)))
0.5
MSA Tranderaer
DeepSequenee
#Promaterss
ProGee3
Trvoception
1.5M
0.4
100M
1596-It
500M
0.2
5000
10000
15000
20000
25000
30000
Inference Time (second)
Figure 6: Comparison of Inference Efficiency. The area of the ball indicates the number of network parameters of model. Our model (in blue) achieve SOTA performance (y-axis) with minimum inference time (x-axis) and 1% number of parameters the ESM.
RELATED WORK
Protein Sequence and Structure Representation Due to the enormous experimental cost of mea- suring protein structures, the number of known protein sequences is thousands of times larger than
8

---

Under review as a conference paper at ICLR 2023
protein structures (Eswar al. 2008 Hsu et al. 2022 Meanwhile, the protein sequences repre- sentation is highly similar to human language, which naturally promotes the fast development of natural language processing (NLP), especially transformer- based methods for encoding protein quences (Coin et al. 2003 Meier et al.
2021 Ofer et al., 2021] Castro et al. (022). However, the geometry of proteins also suggests higher-level structures and topological relationships that are vital to protein functionality. Structure prediction of proteins always attracts great attention in the field (Chi & Liberles 2016 Jumper et al. 2021 Baek et al. [022]] Varadi al. (022). The breakthrough progress in protein folding also enriches structured proteins. For instance, Hsu et al. (2022) and and ALPHAFOLD prediccee for model traming. which greatly eases the data shortage problem and achieves significant performance gain.
Ma al. (2022) mixed experimentally-teste
Structural encoding for Protein Graphs According to the laws of physics, the atomic dynamics do not change matter how a protein is translated or rotated from one place another (Han jet al., 2022). Therefore, the inductive bias of symmetry should be incorporated into the design of protein structure-based models. To this end, research work has been proposed respect the spatial relationship of amino acids (Torng & Altman] 2017] [Sato & Ishida] (019). Such CNN-based methods aggregate the local structure of each residue and integrate estimated local qualities into the whole protein properties. However, these methods neglect geometric equivariance, which can usually be captured by equivariant graph neural networks (Ganea et al. 2021 Stärk et al. 2022)
Protein Representation As existing protein language models require high computational costs and are difficult to train, finding an effective feature representation of protein data is important for downstream tasks TTompson et al. 2012). Contrastive learning and self self-prediction (Elnaggar et al. 2021 Zhang et al., 2022 Hsu et al. 2022 used self-supervised pre- training methods to extract good representation for reducing computational resources. Despite only applying classical representation learning methods on protein, some researchers designed sophisticated encoders for expressive pro- tein representation. For instance, Li et al. (2022) proposed W- W-GNN variants that efficiently interact with scalar-vector features. [Somnath et al. (202T) introduced HOLOPROT to connect different modalities of proteins, including surface, structure, and sequence representation.
Mutation Effeet Prediction Multiple sequence alignment (MSA) is an essential ingredient for many of the existing state-of of-the-art methods to predict the effect of single amino acid substitutions such as DEEPSEQUENCE (Riesselman et al. 2018), ALPHAFOLD2 (Jumper et al., 202I). MSA
2022 The MSA for protein sequence or domain captures meaningful information on the evolutionary information of the protein within its family the cost of bringing severe limitations-not all proteins are alignable, such as CDRs of antibody variable domains (Shin et al. [222]) ). and not all the alignments are deep enough to train models sufficiently large to learn the complex interactions between residues. To deal with this issue, ESM 1v (Meier et al. (02II) trains a zero-shot model on large set of unaligned sequences to secure a scalable and bias-free training procedure, and TRANCEPTION (Notin al., 2022) leverages autoregressive predictions and retrieval of homologous sequences at inference.
TRANSFORMER (Rao al., (021). and LM-GVP Wang all
5
CONCLUSION
Designing directed evolution proteins, especially with deep mutants for functional fitness, is of enormous engineering and pharmaceutical importance. However, existing experimental methods economically costly, and in silico methods require significant computational resources. This paper proposed a lightweight zero-shot model for mutant effect prediction on arbitrary numbers of AAs by transferring the problem to denoising protein graph. Our model is trained to recover AA types and other important properties (e.g., B-factor, SASA, and the spatial position of Ca) from observed noisy proteins. We employ translation and rotation equivariant neural message passing layers to extract geometric-aware representation for the microenvironment of central AAs and thus grasp rich information for efficiently learning protein function. The model achieves state-of-the-art performance on PDB datasets in deep mutant tests with significantly fewer computational resources than existing SOTA models.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Frances H Arnold. Design by directed evolution. Accounts of chemical research, 3153)125-131, 1998.
Minkyung Back, Frank DiMaio, Ivan Anishchenko, Justas Dauparas, Sergey Ovchinnikov, Gyu Rie Lee. Jue Wang. Qian Cong. Lisa N Kinch, R Dustin Schaeffer, et al. Accurate prediction of protein structures and interactions using three-track neural network. Science, 373(6557)8871- 876, 2021.
Nadav Brandes, Dan Ofer, Yam Peleg. Nadav Rappoport, and Michal Linial. Proteinbert: A uni- versal deep-learnig model of protein sequence and function. Bioinformatics, 3809)22002 2022.
Michael S Breen, Carsten Kemena, Peter K Vlasov, Cedric Notredame, and Fyodor A Kondrashov. Epistasis as the primary factor in molecular evolution. Nature, 490(7421):535 -388, 2012.
Egbert Castro, Abhinav Godavarthi, Julian Rubinfien, Kevin Givechian, Dhananjay Bhaskar, and Smita Krishnaswamy. Transformer-based protein generation with regularized latent space opti- mization. Nature Machine Intelligence, 4(10))840-851, 2022.
Peter B Chi and David A Liberles. Selection on protein structure, interaction, and sequence. Protein Science, 25(7):1 1160)11688-18888 2016.
Lachlan Coin, Alex Bateman, and Richard Durbin. Enhanced protein domain discovery by using language modeling techniques from speech recognition. Proceedings of the National Academy of Sciences, 100(8):4516-455 2003.
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 Pattern Analysis and Machine Intelligence, pp. 1-1, 2021. doi: 10.1 1..10009PAMM..222333053338
Narayanan Eswar, David Eramian, Ben Webb, Min-Y Yi Shen, and Andrej Sali. Protein structure modeling with modeller. In Structural proteomics, pp. 145-159. Springer, 2008.
Douglas Fowler and Stanley Fields. Deep mutational scanning: new style of protein science. Nature methods, 11(8):801-80) 2014.
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.
Octavian-Eugen Ganea, Xinyuan Huang, Charlotte Bunne, Yatao Bian, Regina Barzilay, Tommi S Jaakkola, and Andreas Krause. Independent se (3)-equivariant models for end-to-end rigid protein docking. In International Conference Learning Representations, 2021.
Jiaqi Han, Yu Rong, Tingyang Xu, and Wenbing Huang. Geometrically equivariant graph neural networks: A survey, arXiv:2202.072300 2022.
Rhys Heffernan, Kuldip Paliwal, James Lyons, Abdollah Dehzangi, Alok Sharma, Jihua Wang, Abdul Sattar, Yuedong Yang, and Yaoqi Zhou. Improving prediction of secondary structure, local backbone angles and solvent accessible surface area of proteins by iterative deep learning. Scientific reports, 5(1):1-11, 2015.
Steven Henikoff and Jorja G Henikoff. Amino acid substitution matrices from protein blocks. Pro- ceedings of the National Academy of Sciences, 89(22): 10915-1099,1 1992.
Chloe Hsu, Robert Verkuil, Jason Liu, Zeming Lin, Brian Hie, Tom Sercu, Adam Lerer, and Alexan- der Rives. Learning inverse folding from millions of predicted structures. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.). Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 8946-8970. PMLR, 17-23 Jul 2022.
10

---

Under review as a conference paper at ICLR 2023
John Ingraham, Vikas Garg, Regina Barzilay, and Tommi Jaakkola. Generative models for graph- based protein design. Advances neural information processing systems, 32, 2019.
Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael John Lamarre Townshend, and Ron Dron. Learning from protein structure with geometric vector perceptrons. In International Conference on Learning Representations. 2020.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool. Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-5588
2021.
Diederik P Kingma and Jimmy Ba. ADAM: A method for stochastic optimization. In Proceedings of International Conference on Learning Representation (ICLR), 2015.
Ben Lehner. Molecular mechanisms of epistasis within and between genes. Trends in Genetics, 27 (8):323-331, 2011.
Jiahan Li, Shitong Luo, Congyue Deng. Chaoran Cheng, Jiaqi Guan, Leonidas Guibas, Jian Peng, and Jianzhu Ma. Directed weight neural networks for protein structure representation learning. arXiv:2201. 132., 2022.
Yufeng Liu, Lu Zhang. Weilun Wang. Min Zhu, Chenchen Wang, Fudong Li, Jiahai Zhang, Houqiang Li, Quan Chen, and Haiyan Liu. Rotamer-free protein sequence design based on deep learning and self-consistency. Nature Computational Science, 2022.
Hongyuan Lu, Daniel Diaz, Natalie J Czarnecki, Congzhi Zhu, Wantae Kim, Raghav Shroff, Daniel Acosta, Bradley Alexander, Hannah O Cole, Yan Zhang, et al. Machine learning-aided engineering of hydrolases for pet depolymerization. Nature, 604(7907))662--665 -67, 2022.
Yunan Luo, Guangde Jiang, Tianhao Yu, Yang Liu, Lam Vo, Hantian Ding, Yufeng Su, Wesley Wei Qian, Huimin Zhao, and Jian Peng. Ecnet is evolutionary context-integrated deep learning framework for protein engineering. Nature communications, 12(1):1-14, 2021.
Wenjian Ma, Shugang Zhang, Zhen Li, Mingjian Jiang, Shuang Wang, Weigang Lu, Xiangpeng Bi, Huasen Jiang, Henggui Zhang, and Zhiqiang Wei. Enhancing protein function prediction performance by utilizing and Modeling, 2022.
protein structures. Journal of Chemical Information
Joshua Meier, Roshan Rao, Robert Verkuil, Jason Liu, Tom Sercu, and Alex Rives. Language models enable zero-shot prediction of the effects of mutations protein function. In Advances in Neural Information Processing Systems, volume 34, pp. 29287-29303, 2021.
Erik Nijkamp, Jeffrey Ruffolo, Eli N Weinstein, Nikhil Naik, and Ali Madani. Progen2: exploring the boundaries of protein language models. 2rXiv:2206.13517, 2022.
Pascal Notin, Mafalda Dias, Jonathan Frazer, Javier Marchena Hurtado, Aidan Gomez, Debora Marks, and Yarin Gal. Tranception: Protein fitness prediction with autoregressive transformers and inference-time retrieval. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepes- vari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Ma- chine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 16990-17017. PMLR, 17-23 Jul 2022.
Dan Ofer, Nadav Brandes, and Michal Linial. The language of proteins: Nlp, machine learning & protein sequences. Computational and Structural Biotechnology Journal, 19: 19:1750-1758. 2021.
CA Orengo, AD Michie, S Jones, DT Jones, MB Swindells. and JM Thortoon. Cath a hierarchic classification of protein domain structures. Structure, 588) 1997. ISSN 0969-2126. doi: https://doi.oog/1...r wtt.:////....0000000999991110100
Fernanda Pinheiro, Omar Warsi, Dan Andersson, and Michael Lässig. Metabolic fitness landscapes predict the evolution of antibiotic resistance. Nature Ecology & Evolution, 5(5):677-687, 2021.
11

---

Under review as a conference paper at ICLR 2023
Roshan M Rao, Jason Liu, Robert Verkuil, Joshua Meier, John Canny, Pieter Abbeel, Tom Sercu, and Alexander Rives. Msa transformer. International Conference on Machine Learning, pp- 8854. PMLR, 2021.
Adam J Riesselman, John Ingraham, and Debora S Marks. Deep generative models of genetic variation capture the effects of mutations. Nature methods, 15(10):8 8600))886-888 122,88222222 2018.
Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Jason Liu, Demi Guo, Myle Ott, C Lawrence Zitnick, Jerry Ma, al. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. Proceedings of the National Academy of Sciences, (18(5)):2011231111 2021.
Gabriel Rocklin, Tamuka M Chidyausiku, Inna Goreshnik, Alex Ford, Scott Houliston, Alexander Lemak, Lauren Carter, Rashmi Ravichandran, Vikram K Mulligan, Aaron Chevalier, et al. Global analysis of protein folding using massively parallel design, synthesis, and testing. Science, 357 (6347):168-4 175, 2017.
Nathan J Rollins, Kelly P Brock, Frank Poelwijk, Michael A Stiffler, Nicholas P Gauthier, Chris Sander, and Debora S Marks. Inferring protein 3d structure from deep mutation scans. Nature genetics, 51(7):1170-11722 170-1176, 2019.
Karen S Sarkisyan, Dmitry A Bolotin, Margarita V Meer, Dinara R Usmanova, Alexander S Mishin, George V Sharonov, Dmitry N Ivankov, Nina Bozhanova, Mikhail Baranov, Onuralp Soyle- mez, et al. Local fitness landscape of the green fluorescent protein. Nature, 533(7603):397-441, 2016.
Rin Sato and Takashi Ishida. Protein model accuracy estimation based on local structure quality assessment using 3d convolutional neural network. PloS one, 14(9)xe2211447, 2019.
Victor Gareia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural net- PMLR, 2021.
works. In International conference on machine learning, pp.
Sisi Shan, Shitong Luo, Ziqing Yang, Junxian Hong. Yufeng Su, Fan Ding, Lili Fu, Chenyu Li, Peng Chen, Jianzhu Ma, et al. Deep learning guided optimization of human antibody against sars-cov-2 variants with broad neutralization. Proceedings of the National Academy of Sciences, 119(11):
e2122954119, 2022.
Jung-Eun Shin, Adam J Riesselman, Aaron W Kollasch, Conor McMahon, Elana Simon, Chris Sander, Aashish Manglik, Andrew C Kruse, and Debora Marks. Protein design and variant prediction using autoregressive generative models. Nature communications, 12(1):1-11, 2021.
Vignesh Ram Somnath, Charlotte Bunne, and Andreas Krause. Multi-scale representation learning on proteins. Advances in Neural Information Processing Systems, 34:25244-25255. 2021.
Hannes Stärk, Octavian Ganea, Lagnajit Pattanaik, Regina Barzilay, and Tommi Jaakkola. Equibind: Geometric deep learning for drug binding structure prediction. In International Conference Machine Learning, pp. 20503-20521. PMLR, 2022.
Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jon Shlens, and Zbigniew Wojna. Rethink- ing the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pallern recognition, pp. 2818-2826, 2016.
Dawn GL Thean, Hoi Yee Chu, John HC Fong, Becky KC Chan, Peng Zhou, Cynthia Kwok, Yee Man Chan, Silvia YL Mak, Gigi CG Choi, Joshua WK Ho, et al. Machine leaming-coupled combinatorial mutagenesis enables resource-efficient engineering of crispr-cas9 genome editor activities. Nature Communications, 13(1):1-14, 2022.
Andrea D Thompson, Amanda Dugan, Jason E Gestwicki, and Anna Mapp. Fine-tuning multi- protein complexes using small molecules. ACS chemical biology. 78)):3111-1320, 2012.
Wen Torng and Russ B Altman. 3d deep convolutional neural networks for amino acid environment similarity analysis. BMC bioinformatics, 18(1):1-23, 2017.
12

---

Under review as a conference paper at ICLR 2023
Mihaly Varadi, Stephen Anyango, Mandar Deshpande, Sreenath Nair, Cindy Natassia, Galabina Yordanova, David Yuan, Oana Stroe, Gemma Wood, Agata Laydon, et al. Alphafold protein structure database: massively expanding the structural coverage of rooein-sequenne space with high-aceuracy models. Nucleic acids research, 50(D1)) 2022.
Zichen Wang, Steven A Combs, Ryan Brand, Miguel Romero Calvo, Panpan Xu, George Price, Nataliya Golovach, Emmanuel Salawu, Colby J Wise, Sri Priya Ponnapalli, al. Lm-gyp: an extensible sequence and structure informed deep learning framework for protein property predic- tion. Scientific reports, 12(1):1-12, 2022.
Bruce Wittmann, Kadina E Johnston, Zachary Wu, and Frances H Arnold. Advances in machine learning for directed evolution. Current opinion in structural biology, 69:11-18, 2021.
Zachary Wu, SB Jennifer Kan, Russell D Lewis, Bruce J Wittmann, and Frances H Arnold. Ma- chine learning-assisted directed protein evolution with combinatorial libraries. Proceedings of the National Academy of Sciences, (885188:885288858 2019.
Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Banchez-Gonzalee Peter Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via de- noising for molecular property prediction. arXiv:2206.00133, 2022.
Zuobai Zhang, Minghao Xu, Arian Jamasb, Vijil Chenthamarakshan, Aurelie Lozano, Payel Das, and Jian Tang. Protein representation learning by geometric structure pretraining. arXiv:2203.06/22 2022.
13


Reviews:
[{"review_id": "xJVLySIcWJR", "reviewer": "Reviewer_LEve", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose a ‘lightweight graph neural network’ (LGN) for fast zero-shot fitness prediction of mutational effect from protein structures. Noisy perturbations of node and edge features computed from the protein structure are fed into Equivariant Graph Convolution (EGC) layers to extract hidden representations. These representations are used for multitask learning of amino acid type, microenvironment properties (solvent accessible surface area and crystallographic B-factor) and denoising of residue coordinates. Using the learned representations, the authors demonstrate highly efficient SOTA performance in mutation effect prediction.", "summary_of_the_review": "The general aim of leveraging structure information to improve fitness prediction performance in the context of protein design is important. However the manuscript and its approach are not well motivated or clearly presented. Moreover, the manuscript contains errors that impact both the strength and the readability of the work presented. The captions of Figures are sparse and make it difficult to interpret results and one figure is not mentioned anywhere in the main manuscript. ", "strength_and_weaknesses": "Strengths:\n\nThe general aim of leveraging structure information to improve fitness prediction performance in the context of protein design is important. \n\nWeaknesses: \n\nI am not sure if the overall premise of treating mutants as perturbations to node signals and using a denoising autoencoding approach works for their goals. The premise of denoising, is that one brings the perturbed data back to the original data manifold, so likely perturbations will simply be denoised to the training distribution (consisting primarily of wild type proteins). This will not allow for novel protein generation. Indeed there is no quantification in the manuscript for the ability to generate fitter more novel mutants.  The authors state “Altering AA types of a protein in nature can be viewed as adding corruptions to the node features of the protein graph, and denoising the graph makes a remedy to search for mutants with the best fitness”  The connection between denoising and fitness is assumed but never fleshed out. \n\n\nAdditional comments for improvement are below:\n\n1. Sato & Ishida (2019) are cited twice in the introduction. I believe the first citation in the following sentence is incorrect: “Sato & Ishida (2019) applied 3DCNN to identify a new polymerase with advantageous single-site mutation and enhanced the speed of degrading PET by 7-8 times at 50”. The correct citation should be Lu et al., Nature, vol. 604, pgs 662–667, 2022. Note that the temperature of 50 degrees is measured in Celsius, and some background information (i.e. PEG is a waste product that can recycled by degradation) would be helpful to contextualize this work. The second citation of this paper, regarding autoregressive inference for scoring single-site mutations, is also irrelevant. Please re-check all citations, taking care to only include works that are relevant to this paper.\n\n2. What does “capture the micro-frame property in the protein graph geometry” mean? Does it refer to local geometric properties (e.g., alpha helix, beta sheets) of a contiguous subset of residues in the protein? The term ‘micro-frame’ is never clearly defined.\n\n3. The usage of terms ‘invariance’ and ‘equivariance’ is inconsistent throughout.\n\n4. The 93-dimensional edge attributes include a 66-dimensional encoding of the ‘relative position in the protein sequence’. This comprises of a 65-dimensional 1-hot encoding of sequence distance between pairs of nodes, and a scalar contact signal describing ‘if the two residues are in contact in space’. Here, the threshold value of 65 (possible typo in Appendix A.3 states 64) is chosen according to the sequence distance distribution. Is this threshold fixed for all proteins in your dataset?\n\n5. In the introduction, the authors state that their model predicts mutation scores without assuming independence of individual mutations. Specifically, the authors point to the shortcomings of summing up log-odd ratios. However, in Appendix B.2, Eqn 5, the fitness score for multi-site mutations is computed by summing log-odd ratios for single- site mutations. Can you please clarify?\n\n6. “The p-value for both coefficients is 0.000.” This is likely due to rounding or truncation. Perhaps say that p-values for both coefficients is < 0.001?\n\n7. I cannot find any discussion on the limitations of this work. The proposed model is ‘lightweight’ and efficient to train because it relies on the availability of prior knowledge, i.e., tertiary protein structures. This is an important distinction as the models argued against in the introduction are large due to their attempt to model evolutionary-scale information about protein sequence composition in the absence of structure information.\n \t\t\t\t\t\t\n8. Can you comment on the application of this model towards predicting protein stability changes (i.e., changes in the Gibbs free energy of unfolding) between wild type and mutant proteins?\n\n9. Consider citing related work on fitness prediction using regularized latent space optimization (Castro et al., Nature Machine Intelligence, vol. 4, pgs 840–851, 2022), and structure-informed sequence embedding (Wang et al., Scientific Reports, vol. 12, 6832, 2022). \n\n10. “The latter protein language models derived from natural language processing (NLP) encode sequence semantics and often need hundreds of GPU cards to train on billions of protein sequences” This claim needs a citation as most protein datasets are upper bounded at 100s of Millions. ESM-1 (Rives et al 2021) was trained on 250M sequences and AlphaFold 2 (Jumper et al 2021) trained on 350M.\n\n11. Figure 2 seems to be an important result however it is never mentioned in the manuscript and therefore is difficult to interpret.\n\n12. The ablations in Figure 3 are obscured by their same coloring, especially for datasets where their y-positions overlap.\n\n13. I would also suggest adding quantifications of the ability to perform directed protein evolution (i.e. generate newer fitter molecules). \n\n\n\nAfter reading the rebuttal: \nI raised the score from \"reject, not good enough\" to \"marginally below the acceptance threshold\":\nThey have made all the corrections that we and the other reviewers pointed out.\nThey have included additional comparisons to DeepSequence, Tranception, ProGen2 in Figure 3, Figure 6, and Table 1. In cases where a comparison could not be performed (EVE and Potts model) due to technical issues, the authors provide a reasonable explanation.\nThe wording around \"independence of single mutations\" (Appendix B.2) that was confusing for us and the other reviewers has been fixed.\nI hesitate to recommend acceptance because:\nVAE type models can generate high fitness protein sequences, and they do not suffer from the independence of single mutations assumption. Our main critique, i.e. \"This will not allow for novel protein generation. Indeed there is no quantification in the manuscript for the ability to generate fitter more novel mutants.\" was not addressed in the rebuttal. Ultimately their method is good for predicting fitness and other properties efficiently, but is not practically useful for generating new sequences.\nThe authors admit that computing MSA for pre-training is computationally expensive. However they assume that this info is given ahead of time and they do not include it in their runtime. At the same time they argue: \"we doubt the speed of inference time for Potts Model, as the efforts in preparing the MSA information should also be included.\"\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "M9EIkYQI2X", "reviewer": "Reviewer_G8po", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors studied the problem of amino-acid mutation prediction in protein by regarding this problem as solving a denoising problem with lightweight GNN. Several large-scale pre-trained models have previously been proposed and applied to this problem, such as ESM-1v and MSA transformer. The authors proposed a new lightweight graph neural network approach for this problem. Three major contributions are claimed.\n1. Improvement in generalization ability.\n2. Improvement in efficiency.\n3. Improvement in performance on high-order mutations.", "summary_of_the_review": "The authors introduced a new approach to the protein design problem but did not present sufficient evidence to support the claimed advantages. The results also lack further explanation and interpretation in detail. As a result, it is unclear to judge whether the contributions are valid. In addition, no clear explanation of how the model could be used is provided. In my opinion, the paper does not meet the standard of ICLR at this moment and may need further major revisions with a focus on the results part. ", "strength_and_weaknesses": "# Strength\n1. The authors provided a clear formulation of the problem of interest.\n2. The authors provided a clear statement of their claimed major contributions.\n3. The authors provided a very detailed description in the method section, including intuition and mathematical formulation of the loss function.\n# Comments\n1. The authors did not offer a very strong intuition on applying zero-shot or multi-task learning to the mutation fitness prediction problem.\n2. The authors claimed that the model improved the generalization ability, but no clear evidence is given to support how the generalization ability has been improved. There are also no explicit explanations for this supposedly most significant advantage of the model.\n3. There is not any reference or explanation to Figure 2 in the main text, which puts the reader in a quite confusing situation of wondering about its purpose. \n4. The authors tried various LGN configurations and presented the results in Figure 3. However, there isn’t a conclusion or hypothesis on which configuration may result in better performance under certain conditions. The authors could have conducted more experiments and provided more insights into these results. \n5. Figure 3 is not very illustrative, e.g., ESM-1v is missing/hidden for F7YBW8, and the labels for the LGN models are not distinguishable.\n6. The comparison in section 3.5 is somehow not clear enough. First, the authors only tested on one protein sample, which does not offer sufficient evidence to prove the claimed advantage. Second, the training time usage comparison is ambiguous since different models are trained on datasets of various sizes, as indicated in Table 1 of the paper. \n7. The authors did not mention how the model might be used for downstream tasks, particularly mutation prediction. I would personally be interested to know whether there would be a web tool or a pre-trained model for users to try.\n# Minor\n1. Some grammar mistakes (although without affecting the understanding of the text), such as “approach … sum up” and “results are conducted”. \n2. Since Table 2 is only included in the supplementary material, the “Table 2” reference in the main text’s section 3.5 might need be changed.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "CfIOK8C2BQn", "reviewer": "Reviewer_cBgt", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper focuses on the task of predicting the effects of mutations in protein sequences. To that end it leverages a graph neural network with : 1) Biochemical (eg., AA type, SAS, B-factor) and geometric (eg., 3D coordinates of alpha-carbons, dihedral angles) features of amino acids (nodes in the graphs) and edge features (eg., interatomic distances, local N-C positions and position encoding) 2) A stack of Graph Equivariant Convolution (EGC) layers to learn node and edge embeddings 3) A multi-task pre-training framework that combines AA type classification, SAS and B-factor prediction and 3D-coordinates denoising objectives (positions and angles). The architecture is trained on CATH 4.3.0 and mutation effect prediction is assessed in the zero-shot setting against several baselines.", "summary_of_the_review": "The modeling approach introduced in this paper is very promising. There are a couple residual issues -- in particular regarding task interference during the multi-task pre-training and evaluation / baselines chosen in experiments -- that would greatly benefit the work if discussed further / corrected. Willing to increase my score if these are addressed during rebuttal.\n\n-------------------------------------------------------------------------------------------------------------------------------------\n[Update post rebuttal]\nBased on my discussion with the authors and the additional results provided during rebuttal, I believe that the paper in its current form is still not ready for publication and hence maintain my original score (ie., below the acceptance threshold). I provided clear and actionable recommendations to improve the paper in my last response to the authors.", "strength_and_weaknesses": "**Strengths**\n- The GCN layers help bypass data augmentations / be more data efficient\n- The multi-task pre-training is a very sensible idea to learn protein embeddings that combine the different modalities characterizing protein sequences (ie., primary, secondary, tertiary structure) and in turn improve downstream task performance (eg., mutation effects prediction)\n\n**Weaknesses**\n- As evidenced by the results in section 3.4 and figure 5, there seems to be severe task interference when pre-training simultaneously on all tasks. Including the AA recovery and SAS prediction tasks captures the bulk of the performance lift, while other tasks are either not adding substantial value beyond that, or seem to actually be destructive (eg., dihedral prediction). It is thus not clear whether one should keep all these objectives when pre-training? What is the final training objective / set of pre-training tasks that is recommended to use by the authors?\n- The current empirical evaluation lacks several important baselines with respect to mutation effects prediction (eg., DeepSequence, EVE, Tranception) without which claims about “SOTA performance” in abstract/conclusion seem unsubstantiated.\n- The method introduced in this work does seem to achieve relatively high fitness prediction performance Vs inference time tradeoff. To drive this point home though would require comparing with Potts models which also achieve a very favorable tradeoff as well.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject