ProSST: Protein Language Modeling with Quantized Structure and Disentangled Attention
Authors: Mingchen Li
Yang Tan
Xinzhu Ma
Bozitao Zhong
Huiqun Yu
Ziyi Zhou
Wanli Ouyang
Bingxin Zhou
Pan Tan
Liang Hong
Venue: NeurIPS 2024 poster (2024)
Keywords: Protein
Protein Language Modeling
Protein structure
Disentangled Attention
Primary area: machine_learning_for_other_sciences_and_fields

Abstract: Protein language models (PLMs) have shown remarkable capabilities in various protein function prediction tasks. However, while protein function is intricately tied to structure, most existing PLMs do not incorporate protein structure information. To address this issue, we introduce ProSST, a Transformer-based protein language model that seamlessly integrates both protein sequences and structures. ProSST incorporates a structure quantization module and a Transformer architecture with disentangled attention. The structure quantization module translates a 3D protein structure into a sequence of discrete tokens by first serializing the protein structure into residue-level local structures and then embeds them into dense vector space. These vectors are then quantized into discrete structure tokens by a pre-trained clustering model. These tokens serve as an effective protein structure representation. Furthermore, ProSST explicitly learns the relationship between protein residue token sequences and structure token sequences through the sequence-structure disentangled attention. We pre-train ProSST on millions of protein structures using a masked language model objective, enabling it to learn comprehensive contextual representations of proteins. To evaluate the proposed ProSST, we conduct extensive experiments on the zero-shot mutation effect prediction and several supervised downstream tasks, where ProSST achieves the state-of-the-art performance among all baselines. Our code and pre-trained models are publicly available.

ProSST: Protein Language Modeling with Quantized Structure and Disentangled Attention
Mingchen L223,5+
Yang Tan 23,5,+ Xinzhu Ma
Bozitao Zhong 1.4
Huiqun Yu
Ziyi Zhou!
Wanli Ouyanga4 Bingxin Zhou!
Pan Tan Liang Hong!.2.
Shanghai Jiao Tong University, China
(zy-zhou, bingxin wwaingiin.zoou.he.. , www.hhogggnnnneeeee..eee euu.uuuuuuuuuuu.u.....u.uuuuuu.. edu. en, tpan1039@gmail.cmm com,
2 Shanghai Artificial Intelligence Laboratory, China [ouyang-wanli,
maxinzhh]@@jjabb org. en
3 East China University of Science and Technology, China
(1me, tyang) @mail..a.a....m @mail.ecust. ecust. edu, cn, yhq@ecust. eduecu.steeu..co cn The Chinese University of Hong Kong, China
zbztzhz0gmail. com;
5 Chongqing Artificial Intelligence Research Institute Shanghai Jiao Tong University, China
Abstract
Protein language models (PLMs) have shown remarkable capabilities in various protein function prediction tasks. However, while protein function is intricately tied to structure, most existing PLMs not incorporate protein structure information. To address this issue, introduce ProSST, Transformer-basee protein language model that seamlessly integrates both protein sequences and structures. ProSST incorporates structure quantization module and Transformer architecture with disentangled attention. The structure quantization module translates a 3D protein structure into sequence of discrete tokens by first serializing the protein structure into residue-level local structures and then embeds them into dense vector space. These vectors are then quantized into discrete structure tokens by pre-trained clustering model. These tokens serve as an effective protein structure representation. Furthermore, ProSST explicitly learns the relationship between protein residue token sequences and structure token sequences through the sequence-structur disentangled attention. We pre-train ProSST on millions of protein structures using a masked language model objective, enabling it to learn comprehensive contextual representations of proteins. To evaluate the proposed ProSST, we conduct extensive experiments on the zero-shot mutation effect prediction and several supervised downstream tasks, where ProSST achieves the saae-of-thee-an performance among all baselines. Our code and pre-trate pre-trained models are publicly available
Introduction
Predicting the functions of proteins is one of the critical areas in life sciences (1] In recent decades, protein sequence databases have experienced exponential growth [2), making it possible learn the fundamental representations of protein sequences with large-scale models in data- driven manner. Inspired by pre-trained language models in natural language processing [3] 4]. many pre-trained Protein Language Models (PLMs) have emerged 5 6
Benefiting from
* 'Thee authors contributed equally this work, and this work done during their internship the Shanghai Artificial Intelligence Laboratory. Corresponding authors: Bingxin Zhou (bingxin. zhou@sj tu. edu. en), Pan Tan 1033gmaml com) and Liang Hong [oongl3liangee ju tu. edu. en) httpps//gghhbb.com/aaapppoonnnee
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

remarkable protein representation capabilities, they have become fundamental tools for bioinformatics in protein-related tasks.
The function of a protein is determined by its structure 10] However, most PLMs mainly focus on modeling protein sequences, neglecting the importance of structural information, and one significant reason for this phenomenon is the lack of structural data. Fortunately, some excellent works, such as AlphaFold [T] and RoseTTAFold [12), are proposed, which can accurately predict protein structures. These works significantly expand the protein structure dataset [13 to millions and enable the pre-training of sarge-scale structure aware PLMs. After that, the major challenge is how to effectively integrate protein structure information into PLMs. Specifically, existing structure- aware PLMs [14 [5] first use Foldseek [16] convert protein structures into discrete structure tokens and then integrate these structural data into the Transformer architecture. However, despite achieving promising performance on several tasks, this approach still faces two main issues. First, Foldseek encodes the structure of residue within a protein by considering only the features of its previous and next residues. This representation insufficient and may overlook subtle differences in the local structure of proteins, such as catalytic sites or binding pockets, which are crucial for protein function [17]. Second, the naive Transformer architecture lacks the ability to explicitly model the relationship between protein sequences and structure token sequences, making it challenging to effectively leverage structural cues.
In this paper, we develop ProSST (Protein Sequence-Structurg Transformer). structure- pre-trained protein language model. Specifically, ProSST mainly consists of two modules: a structure quantization module and a Transformer with sequence-structure disentangled attention. The structure quantization module based on GVP (Geometric Vector Perceptron) [[8] encoder, which can encode a residue structure along with its neighborhoods in its local structure and quantize the encoding vectors into discrete tokens. Compared to Foldseek, which only considers individual residues, this encoder can take into account more information from the micro-environment of residue. The sequence-structure disentangled attention module replaces the self-attention module in the Transformer model. This can make Transformer model explicitly model the relationship between protein sequence tokens and structure tokens, enabling it to capture more complex features of protein sequences and structures. To enable ProSST to learn the contextual representation comprchensively, we pre-train our model with the Masked Language Modeling (MLM) objective on a large dataset containing 18.8 million protein structures. To summarize, main contributions
follows:
We propose a protein structure quantizer, which can convert protein structure into sequence of discrete tokens. These token sequences effectively represent the local structure information of residues within protein.
We propose disentangled attention mechanism explicitly learn the relationship between protein structure and residue, facilitating more efficient integration of structural token sequences and amino acid sequences.
To evaluate the proposed ProSST, we conduct extensive experiments on zero-shot mutation effect prediction and multiple supervised downstream tasks, where the proposed model achieves state-of- the-art results among baselines. Besides, we also provide detailed ablations to demonstrate the effectiveness of each design in ProSST.
2
Related Work
2.1 Protein Representation Models
Based on the input modality, protein representation models can be divided into three categories: sequence-based models, structure-based models, and structure-sequence hybrid models.
Sequence-based models. Sequence- based models treat proteins as a sequence of residue tokens, using the Transformer model [19] for unsupervised pre-training on extensive datasets of sequence. According to the pre-training objective, current models can be further divided into BERT-based models 4. GPT-based models [3], and span-mask based models. Specifically, BERT-style models, including ESM-series models [5] 6. 7), ProteinBert[9] and TAPE (20], aim to recover the masked tokens in the training phase. The GPT-style models, such as Tranception 21, ProGen2 (22), and

---

ProtGPT2 [23], progressively generate the token sequences in an auto-regressive manner. Lastly, models that use span-mask as the training objective include Ankh [24]. ProtT5 8 and xTrimo 25
Structure-based models. Protein structures play a dominant role in protein functionality. Therefore, models leveraging structure information generally get more accurate predictions. Recently, various techniques have been applied in learning protein structure representation, including CNN-based models (26) and GNN-based models [18 27
28 29 30), where the GNN-based ones have demon- strated significant versatility in integrating protein-specific features through node or edge attributes. Moreover, the recent advancements in protein folding models 7 111 31 enable the structure- based models access to extensive datasets of protein structures. This led to growing interest in developing
PLMs that leverage protein structure cues
15 32].
Structure-sequence hybrid models. Hybrid models, which incorporate both sequence and structure information of proteins, offer more effective representations of proteins. For example, the LM- GVP model employs ProtBERT-BFD [9] embeddings as input features for the GVP (18] model, while ESM-GearNet [34] investigates various methods of integrating ESM- lb [5] representations
with GearNet 332. Similarly, the recent ProtSSN (35) model leverages ESM-2 7 embeddings as input for the EGNN [36] model, resulting in notable advancements. Both ESM-IFI
k77 and
MIF-ST (38] target inverse folding, utilizing the structure to predict corresponding protein residues,
whereas ProstT5 [5] focuses on the transformation between residue sequences and their structure token sequences [16] as pre-training objective. SaProt [14] constructs structure-wwar vocabulary using structure tokens generated by foldseck [16]. Both SaProt and ProstT5 extensively utilize large structure databases [3] for their pre-training datasets. ProSST is also a hybrid structure-sequence model. Compared to previous work, ProSST develops an advanced structure quantization method and better attention formulation to leverage the structure cues.
2.2 Protein Structure Quantization
The most intuitive way to represent a protein structure using continuous features, such as coordi- nates, dihedral angles and distance map. However, directly using these continuous features in the pre-training may lead to overfitting [14]. This issue arises from the mismatched representations of the structure between the training set (derived from model predictions) and the test set (measured by wet-lab experiments). As the bridge to eliminate this gap, structure quantization has been investigated by few works. These methods can be divided into two groups based on the way to generate the discrete secondary structure, including the methods based on physical computing, such as DSSP 339, and the methods based deep learning, such Foldseek (16), which have been successfully applied to structure-aware PLMs [14] [3]. The structure quantization module of ProSST also relies on learning-basee approaches but provides more detailed residue structure representation than Foldseek.
3 Method
In this section, we introduce the architecture of ProSST. ProSST mainly contains two modules: structure quantization (Section 3.T) module and a-transformer-bassed model with sequence-strructur disentangled attention. (Section 3.2).
3.1 Structure Quantization Module
The structure quantization module aims to transform residue's local structure into a discrete token. Initially, the local structure is encoded into a dense vector using a pre-trained structure encoder. Subsequently, a pre-trained k-means clustering model assigns category label to the local structure based on the encoded vector. Finally, the category label is assigned to the residue as the structure token. The pipeline of structure quantization is shown in Figure] I
Structure representation. We categorize protein structures into two distinct levels: protein structure and local structure. Protein structure denotes the complete architecture of a protein, including all its residues. The local structure focuses on specific individual residues. It describes the local environment of a residue by centering 40 residues surrounding it in three-dimensional space [8] Compared to protein structure, local structures are in finer granularity, which allows for more accurate description of the structure of
specific residue and including it along with the nearest
3

---

A. Structure encoder
B. Local structure code book
centroids
Cauesian
(0.()
Noised protein structure
Protain
GND
strusture
CUP
MLP
FK
Glaster Model (k [k-maans]
Local structures
Encoder
Decoder
Local structure codebook
Oateaal DATHH 58K alrucuuee)
(-5 Millinn structures)
C. Protein structure serialization quantization
GVP
434
2
3
Structure tokens
Protein structure residues)
Residue-level local structures
vectors
Encoder
Figure 1: The pipeline of structure quantization. (A) Training of the structure encoder. (B) Local structure clustering and labeling. (C) Converting protein structure structure token sequence.
residue. Therefore, a protein containing L residues has one protein structure and L local structures. Despite the different levels of structure, we can use graphs represent it. Formally, we represent a structure using graph G = (V, E), where V and E denote the residue-level nodes and edges, respectively. For any given node without any residue type information of the residue itself. This ensures that the structure encoder is solely focused on the structure cues. The edge set E {e.y} } includes all i for which vj one of the top-40 nearest neighbors of u; determined by the distance between their Co atoms.
V, it contains only the structure information of the residue,
Structure encoder. Based the lobove-mentiones definition of structure, we use geometric vector perceptrons (GVP) [88] as the (local) structure encoder. In particular, the GVP can be represented as structure feature extraction function т&a (G) Ixd where is the number of nodes, is the embedding dimension, and is trainable parameters. We integrate GVP with a decoder that includes position-wise multi-layer perceptron (MLP) form an auto-encoder model. The entire model is trained using a de-noising pre-training objective. In this process, we perturb Ca coordinates with 3D Gaussian noise (Figure TA) and use Brownian motion on the manifold of rotation matrices. according to RF-Diffusion [40]. The model is then tasked with recovering the structure to its original, noise-free state. After being trained on the C.A.T.H dataset 41 (see Appendix A..)) we exclude the decoder and utilize solely the mean pooled output of the encoder the final representation of structures. Although the structure encoder is trained on protein structures, can effectively encode local structures. Therefore, for graph G of protein structure, the encoding is: r == E (g.).. where 9; represents the graph of the local structure associated with the i-th residue in the graph G, and ne (9:) E [dd is the output of the encoder for the i-th node. Here, E IRd is the mean pooled output of the encoder and the vectorized representation of the local structure.
Local structure codebook. The structure code book quantizes dense vectors representing protein structure into discrete tokens (Figure BB. To build this, we employ structure encoder to embed the local structures of residues from the C.A.T.H dataset (See in Appendix (A.2) into a continuous latent space. Then apply the k-means algorithm to identify K centroids within this latent space, denoted as {e These centroids constitute the structure codebook, as shown in Figure 13. For any
embedding, it is quantized by the nearest vector e; within the codebook and j serving as the structure token. this paper, the clustering number K also referred to as the structure vocabulary size.
Protein serialization and quantization. In general, for residue at position in a protein sequence, first build graph g, only based its local structure, and then use the structure encoder to embed

---

it into continuous vector Then we use the codebook to assign structure token E 21,2.... K} this vector the structure token of the residue. Overall, the entire protein structure can be serialized and quantized into sequence of structure tokens (Figure
3.2 Sequence-Structtre Disentangled Attention
Un-masked Residue tokens
Decoder
Disentangled Attention
L
X
Residue
Position
Feed-forward
+
Norm
+
Disentangled Multi-head Attention
to R
R to R
Attention
Norm
Norm
Residue Embedding
Structure Embedding
Masked Residue tokens
Structure tokens
Relative positions
Structure quantization module
Protein
Figure 2: Model architecture of ProSST. ProSST is Transformer-style model and the difference is that ProSST uses disentangled attention instead of self-attention [19
use expanded form of disentangled attention to combine the attention of residual sequences and structure sequences as well as relative positions. Specifically, for a residue position in protein sequence, it can be represented by three items: R, denotes its residue token hidden state, S, represents the embedding of residue-levell local structure, and Pitj is the embedding of relative position with the token at position j. The calculation of the cross attention A: between residue i and residue j can be decomposed into nine components by:
Inspired by DeBerta 22
A.j (R,, S, Pay x (R,, P
R.R. + R.S. + R.P
(I)
+ S.R + SSS +SPP] j|b
R. SJ PPpP
As formulated Equation the attention weight of residue pair can be calculated by separate matrices, including residue tokens, structure tokens, and relative positions. These matrices are utilized for various interactions such residue-to-resiiue, residue-to-strrctuue, residue-to-position structure-tor to-rsiduee structure, and cintiion-to-ooiiiion Since our model concentrates learning contextual embeddings for residues, the terms struccureeto-strrueuree (S,S)), structure-to-position [S,P). position-to-
structure-to ponecture-t--posiii to-position, position-to-residdr position-to-
structure (P)).S))) ), and position-to-possiion (P j)j P PJD do not provide relevant information about

---

residues and thus do not contribute significantly. Consequently, these terms are removed from our implementation of the attention weight calculation. As shown in Figure 2 our sequence-strrctun disentangled attention mechanism includes 5 types of attention.
In the following part, sequence-structurs disentangled attention. To begin, we define the relative position of the i-th to the j-th residue, denoted j(;,j):
single-head attention example demonstrate the operation of
if - S -Lmar
5( = 2T
if j 2 Iimaz
i- + Lmaz
otherwise
where, Lmar represents the maximum relative distance we consider, which is set to 1024 in the implementation. Similar to standard self-attention operation [19], the computation of query, key for structure, residue and relative position, and the value for residue is as follows:
Q' RW K" RW
V RW
Q' SW K SW
(3)
Q' PW K SW
and the the attention score A from residue to residue j can be calculated as follows:
A., = Q
K
Q.K. ((55))
+
K KJQ
(4)
(2) residae residue
(b) resadue stracture
(c) residto--o-postioon
(d) structure residue
(c) positiientt--rssdiee
where Q represents the i-th row of the matrix Q", and Kj denotes the j-th row of K". Q, and K are the i-th and j-th rows of Q" and K respectively. The term K 5(ä.f) refers to the row in K indexed by the relative distance ((,,)), and Q% refers to the row Q" indexed by the relative distance 8(j, i). To normalize the attention scores, a scaling factor of -1 is applied to A. This scaling is crucial for ensuring the stability of model training [42], particularly when dealing with large-scale language models. All the Ay form the attention matrix, and the final output residue hidden state is R.:
vdd
À
R, = softmax
V
(5)
V5d
which is used as the input for the hidden state of the next layer.
3.3 Pre-Traie Pre-Training Objective
ProSST pre-trained with the structure- conditioone masked language modeling. In this approach, each input sequence æ is noised by substituting a fraction of the residues with special mask token or other residues. The objective ProSST is to predict the original tokens that have been noise the input sequence, utilizing both the corrupted sequence and its structure token sequence context:
CMLM Exax EM Σ log
iEM
We randomly select 15% indices from the M for nosing and computing loss for back mack-propagation. At each selected index i, there is an 80% chance of substituting the residue with mask token, a 10% chance of replacing with random residue token, and the remaining residues are unchanged. The training objective is to minimize the negative log-likelihood for each noised residue æ,, based on partially noised sequence a æ// and the un-noised structure tokens, serving as contextual cues. Therefore, to accurately predict the noised tokens, this objective enables the model not only to learn the dependencies between residues but also the relationship between residues and structures. The details of pre-training dataset and paramee configuration be found in Appendix A.2

---

Model
Model Type
Pa
NDCG
Top-recall f
EVE
0.439
0.781
0.230
EVmutation
X
0.395
0.777
0.222
DeepSequence
S
0,407
0.774
0.225
Evolution-based
WaveNet [50
0.373
0.761
0.203
GEMME 47
0.457
0.777
0.211
MSA-Transformer
48
0.434
0.779
0.217
Tranception
0.434
0.779
0.220
RITA [4
0,372
0.751
0.193
UniRep 45
0.190
0.647
0.139
ESM-1v 6
Sequence-based
0.374
0.732
0.211
ESM-2
0.414
0.747
0.217
ProGen2
0.391
0.767
0.199
VESPA
0.394
0.759
0.201
ESM-IF 37
0,422
0.748
0.223
Inverse-folding
MII-ST 38
0.401
0.765
0.226
Trancepiton EVE 122
0.457
0.786
0.230
ESM-Iv
Ensemble Models
0.407
0.749
0.211
DeepSequence" [1
0,419
0.776
0.226
SaProt [14
0,457
0.768
0,233
Sequence-Struce models
ProSST
0.504
0.777
0.239
Table Comparison of zero-shot mutation prediction performance ProteinGYM benchmark between ProSST and other models. Ps is the Spearman rank correlation.
4 Experiments
In this section, comprehensively evaluate the representation ability of ProSST in several bench- marks, covering zero-shot mutant effective prediction tasks (Section 4.1) and various supervised function prediction tasks (Section 4.2). Additionally, also provide ablation studies and discussions to further show the effectiveness of the detailed designs in our model (Section 4.3
4.1 Zero-Shot Mutant Effect Prediction
Datasets. evaluate the effectiveness of ProSST in zero-shot mutant effect prediction, we conduct experiments on ProteinGym 43 and utilize AlphaFold2 [1] to generate the structures of wild-type sequences. See Appendix A.2 for the details of the dataset and Appendix A for scoring method.
Baselines. We compare ProSST with the current state-of-the-ar models, including sequence-based
models 6 7 evolutionary models 47
22
46 sequence-structune model [14], inverse folding models (37 38)
49
50
51]. and ensemble models 6 52 51]
Results. Table shows the performance of zero-shot mutant effect prediction ProteinGYM, Based on the results, we draw several noteworthy conclusions:
ProSST outperforms all baselines on zero-shot mutant effect predictions of ProteinGYM. We used the non-pan non-parametric bootstrap method to calculate the standard error of the difference in Spearman performance between each model and ProSST. The results showed that all standard errors were less than 0.01. This calculation was based 10,000 bootstrap samples extracted from proteins the ProteinGym benchmark. Furthermore, ProSST was compared against other models subsets of ProteinGYM categorized by function, such stability, activity, binding, and expression. ProSST achieves state-of-the-art (SOTA) performance in the stability, binding, and expression subsets, detailed in Appendix A.4 Notably, ProSST achieves the best performance in predicting stability. aligning with the previous findings that models incorporating structure information typically perform better in stability predictions 43
The degraded version of ProSST (without structure) gets results similar other sequence- based models. This demonstrates that the performance improvement of our model stems from the efficient modeling of structure information, rather than other factors such as more powerful backbones.

---

DeepLoe ore
Metal Ton Binding
Thermostability
GO-MF
GO-BP
GO-CC
Model
Params
Acc
Acc ?
Px
FI-Max
FI-Max 1
FI-Max 1
ESM-2
650M
91.96
71.56
0.680
0.670
0.473
0.470
ESM-1b
650M
92.83
73.57
0.708
0.656
0.451
0.466
MIF-ST
643M
91.76
75.08
0.694
0.633
0.375
0.322
GearNet
42M
89.18
71.26
0.571
0.644
SaProt-35M
91.97
74.29
0.692
0.642
SaProt-650M
650M
93.55
75.75
0.724
0.682
0.486
0.479
ESM-GeerNee
690M
93.55
74.11
0.651
0.676
0.516
0.507
ProSST
110M
94.32uam
76.37
0.726
0.682
0.492
0.501
Table 2: Comparison of supervised fine-tuning on downstream tasks. Ps denotes the Spearman correlation coefficient.
4.2 Supervised Fine-Tuning Tasks
Downstream tasks. For supervised learning, we choose four protein downstream tasks, including thermostability prediction, Metal Ion Binding prediction, protein localization prediction (DeepLoc) and GO annotations prediction (three settings including MF, BO, and CC). More details of the tasks, datasets, and metrics can be found in Appendix A.2
Baselines. compared ProSST with other PLMs including ESM-2 7,, ESM-Ib 5. and
struetce-ttructure model SaProt [14] (two parameter versions, 35M and 650M), MIF-ST 38)
as
well as the protein structure representation model GearNet [32] and ESMM-Gertie 34
Results. The results of the supervised fine-tuning tasks are shown in Table 4.2 and we can get the following conclusions:
ProSST gets the best results among all models with 4 firsts in all 6 settings. For the tasks (settings) of DeepLoc, Metal Ion Binding, ProSST largely surpasses other methods, and ESM-GeerNet gets comparable (or slightly better) results for thermostability and GO-BP and GO-CC with PrOST, the price of more than 6x model size.
The sequence struuccuee models, ESM-GcarNet, SaProt and ProSST, show better results than other counterparts, which suggests the importance of the structure cues in protein modeling. Furthermore, ProSST is more capable of integrating sequence and structure information of proteins than SaProt, which confirms the effectiveness of our designs.
Combined with the results in Section
ProSST exhibits powerful ability in multiple settings.
4.3 Ablation Study
In this section, provide additional ablation studies and discussions to show the necessity and effec- tiveness of the detailed designs in ProSST. Specifically, we use zero-shot mutant effect prediction on ProteinGYM, supervised downstream task DeepLoc, and the perplexity in the pre-training validation set to conduct corresponding experiments.
A Ablations on quantized structure. The ablation results of quantized structure are shown in Table and Figure ((a). and we can get the following findings:
We can find, as the increases of (the size of local structure vocabulary), the performance of ProSST shows upward trend all metrics, and most metries achieve the best results with K = 2048. Based on that, we set = 2048 as our default setting.
As the increase of K, the convergence of ProSST improves progressively (Figure (aa). which suggests incorporating structure cues can improve the representation capabilities of models.
Based on the same network architecture, the proposed structure quantization method (with an appropriate hyper-parameter K)) performs better than Foldseek [66 and DSSP [39]. which shows the effectiivnees of our design.

---

DeepLoc
ProteinGYM
Pretraining
Acc% 1
NDCG
Top-Recall
Perplexity +
(OoST((K=0990
93.88 (#0.15)
0.498
0.773
0.233
8.880
ProSST (K=2048)
94.32 (=0.10)
0.504
0.777
0.239
9.033
ProSST (K=1024)
93.43 (#0.15)
0.485
0.760
0.231
9.333
ProSST (K=5112)
93.70 (#0.16)
0.471
0.759
0,223
9.577
ProSST (K=128)
93.14 (#0.04)
0.469
0.753
0.228
10.021
ProSST (K=20)
93.05 (#0.
0.438
0.744
0.210
10.719
ProSST (K=1)
89.48 (=0.24)
0.390
0.738
0.181
12.182
ProSST (K=0)
89.77 (#0.26)
0.392
0.741
0.184
12.
ProSST (Foldseek)
93.08 (#0.22)
0.468
0.759
0.228
10.049
ProSST (DSSP)
93.16 (#0.16)
0.439
0.760
0.204
10.009
Table Ablation studies on quantized structure. We first show the performance of our models with K centroids of local structures. ProSST (K=0) refers to the model without structure token sequence. We also replace the proposed quantization method with existing Foldseck and DSSP, and show the results of these variants.
DeepLoc
ProteinGYM
Pretraining
Acc%
Ps
NDCG t
Top-Recall f
Perplexity $
ProSST
94.32 (#0.10)
0.504
0.777
0.239
9.033
ProSST (- P2R)
91.31 (#0.14)
0.478
0.778
0.227
9.173
ProSST (- R2P)
92.17 (#0.32)
0.466
0.772
0.216
9,410
ProSST (- R2S)
90.48 (#0.41)
0.438
0.766
0.208
12.142
ProSST (- S2R)
91.27 (#0.20)
0.475
0.779
0.226
9.355
ProSST (- PE)
86.05 (+0.65)
0.095
0.634
0.126
13.885
ProSST (self-attention)
90.37 (+0.21)
0.401
0.728
0.189
12.346
Table 4: Ablation studies on disentangled attention. The term "-S2R" denotes the removal of structure- to-residue in our attention formulation, similar to other terms, and "- PE" denotes the removal of positional encoding. ProSST (self-attention) refers to the model trained with standard attention (with structure cues).
Perplexity Curves
Perplexity Curves
14
Perpleat
Perpleat
12
11
11
9
100K
200K
300K
400K
500K
100K
200K
300K
400K
500K
Training Stop
Training Step
ProtST IK- 109r0
PreSST or 512)
FreSST x
PrefiT RP51
PreOCT : 728)
PreSST conteeteeiionn PruSST
PrussT 2048)
ProsS OK 1280
ProsST X 0)
PruST ( S2R1
PruSTT ( PE)
PruSST i 10240
ProSST 0K 200
ProSST C R29)
(a) Ablations on quantized structure.
(b) Ablations on disentangled attention.
Figure 3: Perplexity curves of ProSST under different settings. We ablate the components of quantized structure and disentangled attention, and show their perplexity curves the validation set.

---

ProSST (Foldseek), ProSST (DSSP), and all ProSST (K>0) models significantly surpass ProSST (K=0) in all metrics, which confirms the importance of the structure cues again.
There is almost no difference in performance between ProSST (K=1) and ProSST (K=0), indicating that the improvement does not come from the rise in parameters of disentangled attention.
Ablations on disentangled attention. Here show detailed ablations and comparisons of disentan- gled attention in Table and Figure 3(b) and we can get the following observations:
All items in Equation 4 are necessary to our attention formulation. Also note that 'P2R' attention has the least impact model capacity, with the Perplexity slightly increasing from 9.033 to 9.173, suggesting that positional attention to amino acids is relatively less critical than other items. Conversely, removing 'R2S' item results in a significant increase in Perplexity from 9.033 to 12.142, underscoring the important role of structure information in enhancing the model's representation capability.
Compared with standard self-attention, our attention formulation gets better results for all metrics, indicating that explicitly modeling structure cues crucial for integrating such information. Besides, positional encoding is also necessary in our design.
As we have mentioned in the Section 2 our disentangled should learn the connections between structure and residue sequence. To valid these, we conduct further experiments to analyze disentangled attention in the Appendix A.S
5 Conclusion and Limitations
This paper introduces ProSST, protein sequence-stuee struunce-ttuuture transformer for PLM. ProSST includes two key techniques, protein structure quantization and sequence-structure disentangled attention. The structure quantization module contains an encoder and k-means clustering model. The encoder is trained with a denoising objective and is utilized for encoding protein structures. Leveraging this encoder, we embed the local structures of each residue within every protein in the C.A.T.H dataset into a continuous latent space. Then we utilize k-means clustering algorithm to obtain K (default setting 2048) centroids. These centroids are then utilized to discretize the local structures of residues based on the index of the nearest centroid of its structure embedding vectors. A protein structure can be transformed into sequence of discrete numbers (or referred to tokens) and each token representing the corresponding local structure of residue. The sequence-structurg attention enhances standard self-attention by not only considering self-attention residues but also incorporating attention between residues and structures, and vice versa. This enables the model to learn the relationships between residues and structures, thereby acquiring improved adequate contextual representations of residues. Furthermore, we pre-train ProSST with 18.8 million protein structures using a MLM objective. Experimental results show that ProSST can outperform existing models in ProteinG benchmark and other supervised learning tasks. Despite of this, there some limitations of ProSST. For example, the local structure construction and encoding requires heavy computations. In the future work, aim speed up the protein structure quantization process. Another threat that the structural and sequential data are required for ProSST to derive the final protein representations, since the amount of available structural data is lower than that of sequence data. We provide solutions in the Appendix Section A.6 Additionally, we plan to enhance ProSST by training it with larger structure datasets and expanding its parameter, which may further improve its performance.
Acknowledgements
This work was supported by the grants from the National Natural Science Foundation of China (Grant Number 12 12104295), the Computational Biology Key Program of Shanghai Science and Technology Commission 231S1400600), Shanghai Jiao Tong University Scientific and Technological Innovation Funds (21 X01000044),, and Science and Technology Innovation Key R&D Program of Chongqing (CSTIB002TIADDSSTXIO))
the Student Innovation Center at Shanghai Jiao Tong University, and Shanghai Artificial Intelligence Laboratory.
10

---

References
[T] William Jencks. Catalysis in chemistry and enzymology. Courier Corporation, 1987.
[2] The UniProt Consortium. UniProt: the Universal Protein Knowledgebase in 2023. Nucleic Acids Research, 51(DI):D53-DD53) 11 2022.
[3] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9 2019.
[4] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04800 2018.
[5] Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Jason Liu, Demi Guo, Myle Ott, C. Lawrence Zitnick, Jerry Ma, and Rob Fergus. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. Proceedings of the National Academy of Sciences, (18(15)e201623911 18, 2021.
[6] Joshua Meier, Roshan Rao, Robert Verkuil, Jason Liu, Tom Sereu, and Alex Rives. Language models enable zero-shot prediction of the effects of mutations protein function. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances Neural Information Processing Systems, volume 34, pages 29287-29303. Curran Associates, Inc., 2021.
[7] Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, al. Evolutionary-scale prediction of atomic-level protein structure with language model. Science, 379(6637) 2023.
[8] Ahmed Elnaggar, Michael Heinzinger, Christian Dallago, Ghalia Rehawi, Yu Wang. Llion Jones, Tom Gibbs, Tamas Feher, Christoph Angerer, Martin Steinegger, et al. Prottrans: Toward understanding the language of life through self-supervised learning. IEEE transactions pattern analysis and machine intelligence, 44(10):711-7112) 2021.
[9] Nadav Brandes, Dan Ofer, Yam Peleg, Nadav Rappoport, and Michal Linial. Proteinbert: universal deep-leamnng model of protein sequence and function. Bioinformatics, 38(8):2102- 21 10, 2022.
[10] Hedi Hegyi and Mark Gerstein. The relationship between protein structure and function: comprehensive survey with application to the yeast genome Hedited by g. von heijne. Journal of Molecular Biology, 288(1) 14:-111, 1999.
[11] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek. Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-589 2021.
[12] 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).871-876 2021.
[13] Mihaly Varadi, Stephen Anyango, Mandar Deshpande, Sreenath Nair, Cindy Natassia, Galabina Yordanova, David Yuan, Oana Stroe, Gemma Wood, Agata Laydon, al. Alphafold protein structure database: massively expanding the structural coverage of protein-sequence space with high-aceuracy models. Nucleic acids research, 50(D1): D444, 2022.
[14] 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.
[15] Michael Heinzinger, Konstantin Weissenow, Joaquín Gomez Sanchez, Adrian Henkel, Martin Steinegger, and Burkhard Rost. Prost1:: Bilingual language model for protein sequence and structure. bioRxiv, pages 2023-07, 2023.

---

[16] 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, pages 1-4,201 2023.
[17] Hongyuan Lu, Daniel J Diaz, Natalie J Czarnecki, Congzhi Zhu, Wantae Kim, Raghav Shroff, Daniel Acosta, Bradley R Alexander, Hannah Cole, Yan Zhang, et al. Machine learning- aided engineering of hydrolases for pet depolymerization. Nature, 604(7907):662-666) 6677 2022.
[18) Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael John Lamarre Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. In International Conference Learning Representations, 2021.
[19] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[20] Roshan Rao, Nicholas Bhattacharya, Neil Thomas, Yan Duan, Peter Chen, John Canny, Pieter Abbeel, and Yun Song. Evaluating protein transfer learning with tape. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bus Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[21] 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 Machine Learning, pages 16990-17017. PMLR, 2022.
[22] Erik Nijkamp. Jeffrey A Ruffolo, Eli N Weinstein, Nikhil Naik, and Ali Madani. Progen2: exploring the boundaries of protein language models. Cell systems, (4(11):968-9) 2023.
[23] Noelia Ferruz, Steffen Schmidt, and Birte Höcker. Protgpt2 is a deep unsupervised language model for protein design. Nature communications, 13(1):4348, 2022.
[24] Ahmed Elnaggar, Hazem Essam, Wafaa Salah-Eldin, Walid Moustafa, Mohamed Elkerdawy, Charlotte Rochereau, and Burkhard Rost. Ankh: Optimized protein language model unlocks general-purpose modelling. arXiv preprint arXiv:2301. 06568,, 2023.
[25] Bo Chen, Xingyi Cheng, Pan Li, Yangli-ao Geng, Jing Gong, Shen Li, Zhilei Bei, Xu Tan, Boyan Wang, Xin Zeng, et al. xtrimopglm: unified 100b-scale pre-trained transformer for deciphering the language of protein. arXiv preprint 1xXi::2401006999 2024.
[26] Vladimir Gligorijević, 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):3168, 2021.
[27] Yang Tan, Jia Zheng, Liang Hong, and Bingxin Zhou. Protsolm: Protein solubility prediction with modal features. PrXiv:2406.19744 2024.
[28] Bingxin Zhou, Lirong Zheng, Banghao Wu, Kai Yi, Bozitao Zhong, Yang Tan, Qian Liu, Pietro Liö, and Liang Hong. A conditional protein diffusion model generates artificial programmable endonuclease sequences with enhanced activity. Cell Discovery, 10(1):95, 2024.
[29] Bingxin Zhou, Lirong Zheng, Banghao Wu, Yang Tan, Outongyi Lv, Kai Yi, Guisheng Fan, and Liang Hong. Protein engineering with lightweight graph denoising neural networks. Journal of Chemical Information and Modeling, 64(9):3650-3666 2024.
[30] Yang Tan, Lirong Zheng, Bozitao Zhong, Liang Hong, and Bingxin Zhou. Protein representation learning with sequence information embedding: Does it always lead to a better performance? arXiv:2406. /9755, 2024.
[31] Ruidong Wa, Fan Ding, Rui Wang, Rui Shen, Xiwen Zhang. Shitong Luo, Chenpeng Su, Zuofan Wu, Qi Xie, Bonnie Berger, et al. High-resolution de novo structure prediction from primary sequence. BioRxiv, pages 2022-07, 2022.
12

---

[32] Zuobai Zhang, Minghao Xu, Arian Rokkum Jamasb, Vijil Chentammarkkshan, Aurelie Lozano, Payel Das, and Jian Tang. Protein representation learning by geometric structure pretraining. In The Eleventh International Conference on Learning Representations, 2023.
[33] 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-gyp: an extensible sequence and structure informed deep learning framework for protein property prediction. Scientific reports, 12(1))6832, 2022.
[34] Zuobai Zhang, Minghao Xu, Aurelie Lozano, Vijil Chenthamarakshan, Payel Das, and Jian Tang. Enhancing protein language model with structure- based encoder and pre-training. In ICLR 2023
Machine Learning for Drug Discovery workshop, 2023.
[35] Yang Tan, Bingxin Zhou, Lirong Zheng, Guisheng Fan, and Liang Hong. Semantical and topologiaal protein encoding toward enhanced bioactivity and thermostability. bioRxiv, pages 2023-12, 12, 2023.
[36] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International conference on machine learning, pages
PMLR, 2021.
[37] 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 1946-8970. PMLR, 2022.
[38] Kevin K Yang, Niccolò Zanichelli, and Hugh Yeh. Masked inverse folding with sequence transfer for protein representation learning. Protein Engineering, Design and Selection, 36:g7ad015, 2023.
[39] Wolfgang Kabsch and Christian Sander. Dictionary of protein secondary structure: pattern recognition of hydrogen-bonded and geometrical features. Biopolymers: Original Research on Biomolecules, 22(12):2577-2633 1983.
[40] Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L. Trippe, Jason Yim, Helen Eisenach, Woody Ahem, Andrew J Borst, Robert J Ragotte, Lukas F Milles, al. De novo design of protein structure and function with rfdiffusion. Nature, 62097976):1089-11000 2023.
[41] Ian Sillitoe, Nicola Bordin, Natalie Dawson, Vaishali P Waman, Paul Ashford, Harry M Scholes, Camilla SM Pang, Laurel Woodridge, Clemens Rauer, Neeladri Sen, et al. Cath: increased structural coverage of functional space. Nucleic acids research, 49(D1):D266-D27) 2021.
[42] Pengcheng He, Xiaodong Liu, Jianfeng Gao, and Weizhu Chen. Deberta: Decoding enhanced bert with disentangled attention. In International Conference on Learning Representations. 2021.
[43] Pascal Notin, Aaron Kollasch, Daniel Ritter, Lood van Niekerk, Steffanie Paul, Han Spinner, Nathan Rollins, Ada Shaw, Rose Orenbuch, Ruben Weitzman, Jonathan Frazer, Mafalda Dias, Dinko Franceschi, Yarin Gal, and Debora Marks. Proteingym: Large-scale benchmarks for protein fitness prediction and design. In A. Oh, T. Neumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 64331-64378 Curran Associates, Inc., 2023.
[44] Daniel Hesslow, Niccoló Zanichelli, Pascal Notin, Iacopo Poli, and Debora Marks. Rita: study on scaling up generative protein sequence models. arXiv preprint arXiv:2205.05789, 2022.
[45] 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-133 2019.
[46] Céline Marquet, Michael Heinzinger, Tobias Olenyi, Christian Dallago, Kyra Erckert, Michael Bernhofer, Dmitrii Nechaev, and Burkhard Rost. Embeddings from protein language models predict conservation and variant effects. Human genetics, -699-16411 2022.
13

---

[47] Elodie Laine, Yasaman Karami, and Alessandra Carbone. Gemme: simple and fast global epistatic model predicting mutational effects. Molecular biology and evolution, 36:64) 2619, 2019.
[48] Roshan M Rao, Jason Liu, Robert Verkuil, Joshua Meier, John Canny, Pieter Abbeel, Tom Sercu, and Alexander Rives. Msa transformer. In International Conference on Machine Learning. pages
PMLR, 2021.
[49] 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);11-95, 2021.
[50] 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.
[51] Adam Riesselman, John Ingraham, and Debora Marks. Deep generative models of genetic variation capture the effects of mutations. Nature methods, 15(10))816-82) 2018.
[52] Pascal Notin, Lood Van Nickerk, Aaron W Kollasch, Daniel Ritter, Yarin Gal, and Debora Susan Marks. TranceptEVE: Combining family-specific and family-agnostic models protein sequences for improved fitness prediction. In NeurlPS 2022 Workshop on Learning Meaningful Representations of Life, 2022.
[53] Thomas A Hopf, John B Ingraham, Frank J Poelwijk, Charlotta PI Schärfe, Michael Springer, Chris Sander, and Debora Marks. Mutation effects predicted from sequence co-variation. Nature biotechnology, 35(2):128-135 2017.
[54] Raghav Shroff, Austin W Cole, Daniel Diaz, Barrett R Morrow, Isaac Donnell, Ankur Annapareddy, Jimmy Gollihar, Andrew Ellington, and Ross Thyer. Discovery of novel gain-of-function mutations guided by structure-based deep learning. ACS synthetic biology, 9111) 2020.
[55] Christian Dallago, Jody Mou, Kadina E Johnston, Bruce Wittmann, Nick Bhattacharya, Samuel Goldman, Ali Madani, and Kevin K Yang. FLIP: Benchmark tasks in fitness landscape inference for proteins. In Thirty-fftth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
[56] José Juan Almagro Armenteros, Casper Kaae Sønderby, Søren Kaae Sønderby, Henrik Nielsen, and Ole Winther. Deeploc: prediction of protein subcellular localization using deep learning. Bioinformatics, (3(21))3387-333 2017.
[57] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference Learning Representations, 2019.
[58] Minghao Xu, Zuobai Zhang, Jiarui Lu, Zhaocheng Zhu, Yangtian Zhang, Ma Chang, Runcheng Liu, and Jian Tang. Peer: comprehensive and multi-task benchmark for protein sequence understanding. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 35156-355 35173. Curran Associates, Inc., 2022.
14

---

A Appendix
A.1 Zero-Shot Scoring
Previous studies have demonstrated that PLMs, when trained on extensive and varied protein sequence databases, are capable of predicting experimental measurements protein mutants function without further supervision 6 [4] For those PLMs that are trained with masked language modeling objective, the calculation of mutation scores can be formalized as follows:
Score F) - Σ log P(ap, = f(( - log P(rp, = =wi(2))
Here F is a single or multi-point mutant, and {pis 1, FI is a set of triplets, where Pi E N represents the mutation position. f, is the residue after the point mutation, and w, is the original residue of the point mutation. the sequence of residues the wild type. We slightly modify the formula above to adapt to ProSST, where the structure sequence is an additional condition to score mutants:
IF| Score F) log P@ap, = =f,120 8) logp( = =wi22,0
Σ
(8)
Here, is the structure token sequence of the wild type.
A.2 Details of the Datasets and Metrics
Dataset for pre-training. The pre-training data is collected from AlphaFoldDB 13 which contains more than 214 million structures predicted by AlphaFold [T.. We downloaded the 90% reduced version, containing 18.8 million structures I From this collection, we randomly select 100,000 structures for validation (sequences with a similarity of over 30 to the training set will be removed for data deduplication.), enabling us to monitor the perplexity in the training phase. During pre-training, proteins with more than 2048 resiudes (594 samples) are removed for training efficacy.
Dataset for training structure encoder. The dataset used for training the structure encoder originates from CATH43-S40 | This dataset is manually annotated and comprises protein crystal structural domains that have been deduplicated for sequence similarity by 40%. The original dataset contains 31,885 structures. After removing structural domains missing atoms such as Ca and N, the dataset is reduced to 31,270 entries. From this, 200 structures were randomly selected to serve as a validation set. The auto-encoder model was then trained using the configuration that yielded the lowest loss on this validation set.
Dataset for training structure codebook. The dataset for training the structure codebook consists of local structures extracted from CATH43-S40. Given a protein structure, slide along the residue sequence to select a segment with chosen residue as the anchor. Connect up 40 residues within 10 Å (54) to form star-shaped graph. As shown in Figure A4 local structures with more than 40 neighbors account for only 0.00052%, indicating that our choice covers most cases. For pairwise amino acid pairs this graph, if the Euclidean distance is less than 10 Å, a link will be assigned to them. This process yields number of protein local structures equal to the length of the protein multiplied by the total number of proteins, resulting in 4,735,677 local structures from the protein structures in CATH43-SA4 These sub-structures are fed into structural encoder to obtain embeddings. By setting various quantities for K, different structure codebooks are obtained using the k-means clustering algorithm.
Dataset and metries for zero-shot mutant effect prediction.
We utilize the ProteinGYM benchmark 43 to assess the zero-shot mutant effect prediction capabili- ties of ProSST. ProteinGYM offers comprehensive benchmarks specifically collected for predicting
hhtpps//cluutterfolldeecc. com/ http://wwwmiiaaaaatttiiiiaaaalae
Wnoo-redundann-dttt--ettt
15

---

175000
150000
Requency
125000
100000
75000
50000
23000
Number
Figure A4: The distribution of the number of residues within 10 Å distance of a local structure.
protein fitness. It contains a wide range of deep mutational scanning assays with millions of mutated sequences. ProSST is evaluated using the most extensively utilized datasets for substitution mutations, which include 217 experimental assays. Each assay incorporates both the sequence and structure of the protein, with a particular emphasis on 66 datasets that focus on thermo-stability. The evaluation metrics employed are the Spearman coefficient, Top-recall, and NDCG, where higher values signify superior model performance. These metries computed using scripts provided by ProteinG Y.M.
Datasets and metrics for downstream tasks. The downstream datasets have the same train, valid, and test splits as SaProt's and are downloaded from SaProt. Data statisties are provided in Table A5
Thermostability. The task is to predict the thermostability values of proteins using the "Human-cell" divisions from the Thermostability task in FLIIP (55). For this regression task, the Spearman correlation coefficient is utilized as the evaluation metric to evaluate the prediction results.
DeepLoc (Protein Sub-cellular Localization). The task is to output a probability distribution across two sub-cellule localization categories for a protein. This is binary classification task, and we utilize accuracy as the metric to evaluate the predictions. This dataset was the original data split.
introduced by DeepLoc [56] and
Metal Ion Binding. The task is to predict whether metal ion-binding sites exist within protein. This is also a binary classification task, and we utilize accuracy as the metric to evaluate the predictions. This dataset was introduced by TAPE 20 and we use the original data split.
GO annotations prediction. This task is to predict Gene Ontology terms to evaluate the model's ability to predict protein functions. This task was introduced by DeepFRI (26). and three types of GO labels: Molecular Function (MF), Biological Process (BP), and Cellular Component (CC). This is multi-label classification task, and we evaluate the model using the Max FI-Score.
A.3 Details of Implementations
Structure encoder. We describe a structure with the graph (V, E), adopting the characteriza- tions of V and E as outlined in the GVP framework [18]. The GVP encoder includes six-layer bossaae-ppassing graph neural network in which geometric perceptron replaces the MLP to ensure translational and rotational invariance of the input structure. Our GVP encoder is consistent with the original GVP-GNN [18], except that we removed the residue type information. The GVP encoder is trained from scratch. The dimensions for node and edge representations set at 256 and 64,
16

---

Dataset
Training
Valid
Test
Total
Termostability
5,056
639
1,336
7,031
DeepLoc
5,477
1,336
1,731
8,544
Metal Ion Binding
5,067
662
665
6,394
Go annotations prediction
26,224
2,904
3,350
32,478
Table A5: Downstream datasets split statistics.
respectively, with the encoder comprising six layers. For optimization, we employ the Adam opti- mizer mini-batch gradient descent approach. To manage computational load, batches formed by grouping structures of similar sizes, with each batch containing no more than 3000 nodes. The learning rate set to 1.0 10-4 The dropout probability set to 0.01. And The number of graph layers is set at 6. The training and validation curves of the structure encoder are shown in Figure ((a)
Training and Validation Curves of Structure Encoder
Inference Speed
10
2.5
135H4
2.0
5 1.5
10-1
1.0
0.5
10-2
10k.
19k
20k
1750
2000
Step
Sequence Length
(a)
(b)
Figure A5: (a) Training and validation curves of the local-structure auto-encoder. (b) Inference speed of ProSST on different sequence lengths. (batch size 16).
Pre-training. All ProSST models is trained DGX-A800 GPU (8x80G) server in BF16 precision for about a month. The model has 12 transformer layers, 12 attention heads, and 768 embedding dims with 3172 feed-forward embedding dimensions with the GELU activation function. We train with 8192 tokens per mini-batch for 500,000 steps. We use AdamW [57] as our optimizer with 31 and B2 set 0.9 and 0.999, and weight decay value of 0.001. We warm up the learning rate from 0 to 0.0002 over the first 2000 steps, then decay it by a cosine schedule to the O. We use dropout rate of 0.1 and clip gradients using a clipping value of 1.0. For the tokenization of the protein data, we use the residue-level tokenizer which is adopted in several PLMs 7 6 To make the structure sequence the same length as the amino acid sequence, we also added special [SOS], [EOS), and [PAD] token for the structure sequences.
Fine-tuning. To ensure fair comparisons, fine-tuned ProSST using fixed of hyper-parameters. We use for the Adam optimizer with 31 set to 0.9, 82 to 0.98, and applied an 1.2 weight decay of 0.001. The batch size was maintained at 64 (If 64 causes the GPU memory explode, we will reduce the batch size and then use gradient accumulation to achieve the same batch size.) and the learning rate was
0.00003, except for Go annotation prediction, where was adjusted to 0.00001. We fine-tuned all model parameters for 200 epochs, and we choose the best checkpoints based on validation set performance. Following SaProt 14 we downloaded all protein structures identified by Uniprot IDs from AFDB 13], and any proteins not found in AFDB were excluded.
Inference Speed. We computed the inference speed of ProSST, SaProt (650M) and SaProt (35M) on proteins of different lengths using a batch size of 16 on a server equipped with two Intel 6248R processors and 3090 GPU and the results shown in Table (bb)
17

---

Model
Activity
Binding
Expression
Organismal Fitness
Stability
EVE
0.464
0.386
0.408
0.447
0.491
EVmutation
0.440
0.317
0.378
0.411
0,430
DeepSequence
0.455
0.363
0.390
0.413
0,476
WaveNet
0.379
0.325
0.350
0.365
0.449
GEMME
0.482
0.383
0.438
0.452
0.519
MSA-Tramslormer
0.469
0.337
0.446
0,421
0.495
Tranception
0.465
0.349
0.450
0,436
0,471
RITA
0.366
0.302
0.414
0.381
0.398
UniRep
0.182
0.202
0.216
0.141
0.210
ESM-1v
0.396
0.268
0.405
0.362
0,437
ESM-2
0.425
0.337
0,415
0.369
0.523
ProGen2
0.402
0.302
0.418
0.387
0,445
VESPA
0.429
0.347
0.326
0.404
0,461
ESM-IF
0.368
0.389
0.407
0.324
0.624
MIF-ST
0.390
0.321
0.438
0.366
0.485
Trancepiton-EVE
0.487
0.376
0.457
0.460
0.500
ESM-1v (ensemble)
0.420
0.320
0,429
0.387
0.477
DeepSequence (ensemble)
0.455
0.363
0.390
0.413
0.476
SaProt
0.458
0.378
0.488
0.367
0.592
ProSST
0.448
0.477
0.506
0.415
0.674
Table A6: Spearman's rank correlation of baseline models and ProSST on the ProteinGym, separated into five functional categories (Activity, Binding, Organismal Fitness, Stability and Expression).
Structure token
ProteinGYM (p.)
Perplexity
Original
0.504
9.033
All-zero
0.112
14.524
Random
0.182
14.024
Table A7: Performance of ProSST (K=2048) using different structure tokens inputs.
A.4 Performance of models the ProteinGYM benchmark separated by functional categories
Table A.4 shows Spearman's 's rank correlations on ProteinGYM, categorized by five function types. ProSST achieves tfte-offthheart (SOTA) performance in the Stability, Binding, and Expression subsets.
A.5 Additional experiments on disentangled attention.
We conducted additional experiments to analyze the relationship between disentangled attention.
We replaced all structure tokens in the ProteinGY' and the validation set with zeros or random numbers from uniform distribution and re-evaluated ProSST. The results are shown in Table A7 The results show that the incorrect structure tokens harmed the performance of ProSST
Experiment 1.
18

---

R
to 5
S
Query (R)
Qutrry (5)
Query
R to
to R
Attention
Quer 6
Quary
Key
Figure A6: Different types of attentions on Green Fluorescent Protein (GFP). average of each head in the final layer of the Transformer.
These attentions are the
Model
DeepLoc (Ace%)
ProteinGYM (p.)
Perplexity
ProSST (K=2048)
94.32 (+0.10)
0.504
9.033
ProSST (K=1)
89.48 (+0.24)
0.390
12.182
ProSST (K=0)
89.77 (+0.26)
0.392
12.190
Table A8: Performance comparison of ProSST with special K values.
significantly, suggesting that disentangled attention learned sequence-structurr relationship. Otherwise, the performance would only have a minimal impact.
Experiment To verify disentangled attention can indeed enhance performance, we train ProSST (K=1) by substituting the structural tokens with a constant value of 1. This configuration maintains the disentangled attention mechanism even though the structure input is useless. If ProSST (K=1) still improves performance, it indicates that the improvement is solely due to the disentangled attention. We evaluate this model on DeepLoc, ProteinGYM and the valid set. The results are shown in Table A8 There is little difference between K=1 and K==0 since their perplexity curves (refer to Figure 3(a) nearly overlap. This suggests that disentangled attention alone cannot enhance performance without correct structure tokens.
Experiment We visualize the learned different types of attentions on Green Fluorescent Protein (GFP, Unipro ID:P42212 ) including 238 residues, in Figure A6 We can see that disentangled attention learns different attention patterns, with notable differences between "R2S" and "S2R".
httpp:: :tttps:://ww..oo wnwpuot. http:///rwr.u.oooooo///////////.
19

---

Model
Structure Source
ProteinGYM (p.)
BLP (Acc%)
Perplexity
PrOSST (K=2048)
AlphaFold2
0.504
94.32
9.033
ProSST (K=2048)
ESMFold
0,471
92.73
9.144
ProSST (MST)
Missing
0.438
91.84
10.325
ProSST (MST)
AlphaFold
0.456
92.31
9.447
ProSST (K=0)
Missing
0.392
89.65
12.190
Table A9:
Performance comparison of ProSST with different special K.
We can conclude from Experiment and Experiment that the disentangled attention and structure tokens sequence are inseparable. The disentangled attention mechanism cannot function without correct structure tokens. This also indicates that the performance improvement of our model stems from the design the Experiment 3 shows that our disentangled attention actually learned different patterns of attentions.
model rather than increasing its parameters of the attention layer. Furthermore,
A.6 Solutions to Sequence-only Datasets
In conclusion, we offer two solutions for obtain representations of the sequence-only protein datasets:
Utilize AlphaFold or ESMFold 7 for structure prediction as they are highly reliable methods.
Use ProSST (MST), which is trained with structure masking and supporting sequence-only inputs. The MST denotes "Masked Structure Training (MST)", which means that during pre-training, each sample's structure sequence has 50% probability of being replaced by a fully masked sequence [I.I,I,I,I,...II]] simulating missing protein structure. Therefore. when applying ProSST to sequence-only datasets, need to the masked sequence [I,I,I,I,I,,... l] as a substitute for the structure token sequence.
We have evaluated the two approaches on the ProteinGym benchmark, binary localization prediction (BLP) from sequence-only benchmark, PEER [58), and perplexity on the validation set. The results show in Table A9) In the Table A9) the first rows show the performance differences between AlphaFold and ESMFold. Rows 3-4 show the performance of the new model ProSST(MST). And row 5 shows the performance of the sequence-only model.
A.7 AlphaFold PLDDT versus Zero-shot mutant effect performance
(a)
(b)
(c)
Figure A7: (a) PLDDT vs. Spearman of ProSST on ProteinGYM. (b) pLDDT vs. Spearman of SaProt ProteinGYM. (c) PLDDT Spearman of ESM-IFI on ProteinGYM.
Protein structures containing disorder region may not be accurately predicted by AlphaFold 2, potentially leading to reduced performance of structure-aware models. Here, we test the relationship
20

---

between AlphaFold pLDDT scores and the performance of structure-awar models including ProSST, SaProt, and ESM-IF on ProteinGYM, as illustrated in Figures A7 Our findings reveal positive correlation between pLDDT values and model performance: correlation coefficient of 0.30 for ProSST, 0.31 for SaProt, and 0.42 for ESM-IFF, where the correlation coefficient (0p) represents the strength of the relationship. These results suggest that structure-aware models may exhibit limitations accurately predicting the structures of disordered proteins.
21

---

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] We state clearly the scope and contribution of both communities of protein
language modeling and computer science in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes] See Section 5
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-texi system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
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] This paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
22

---

All the theorems, formulas, and proofs in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results
paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes] This paper has disclosed all the necessary information to reproduce the main experimental results, including the dataset (See Appendix A.2 network architecture (See
Section 3.2 and network hyper-parameters (See Appendix A.3
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c)
If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes] Our code and pre-trained models available https::///e com/ ai @i4protein/ProSST
23

---

Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (httpss:/nnip..ce)
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
Policc for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer:
[Yes] We provide implementation details of our model Section A.3 Baseline methods are pre-trained by their authors, which requires no further supervision or learning procedures.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
[NA] Measurements used to this work return fixed results on a determined model, and pre-trained model barely trains set of models under the same dataset or hyper-parameter settings.
Answer:
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes] Details for the running machine are provided at the beginning of Ap- pendix A.3
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics ptttps://neurips.. httpp/////errrggs..c//bbbb//////
[Yes] The general ethical conduct has been carefully gone through to make sure the listed concerns avoided or not applicable in this research.
Answer:
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA] There is societal impact of the work performed.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinfornation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
25

---

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e-g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
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, scraped datasets)?
Question:
Answer: [Yes] The paper poses such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (c.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
We cite all the papers for baseline methods. The baseline results of Prote- inGYM are copied from the latest release on their GitHub repositories, which we attached URLs in Appendix A.2 The baselines results of supervised fine-tuning are copied from the
Answer: [Yes]
latest reviewed version of SaProt.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, thcode thcode. com/ con/datasete
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online. the authors encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA] The paper does not release new assets.
26

---

Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the datasct/codc/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
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] The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
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: [NA] The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
27


Reviews:
[{"review_id": "1H2KqxTKoi", "reviewer": "Reviewer_t3Vw", "summary": "This paper presents ProSST, a new language model for protein data that captures both structural and sequential modalities of proteins. The protein's structure is tokenized using a graph-based auto-encoder architecture, where each residue's local structure is embedded into a vector in a high-dimensional embedding space and then classified using a $k$-means clustering algorithm within that embedding space. The two sequence-based and structure-based data streams are then passed to multiple disentangled attention layers, allowing the model to better capture the connections between structure and sequence data. Numerical results showcase the outperformance of ProSST over several baselines across a range of downstream tasks.", "questions": "- What is the reasoning behind choosing **40** residues as the neighborhood of each node in Section 3.1?\n\n- Is $L$ on line 128 the same as $l$ on line 137? \n\n- Could you please elaborate on how the GVP encoder is parameterized?\n\n- Have you studied what would happen if, instead of discrete structure tokens on the bottom-right of Figure 1, the centroid embedding of each cluster is used as the input structural vector? Also, what if the continuous local structure embeddings are directly used and the clustering is removed altogether?\n\n- Could you explain the average pooling that happens on line 146? In particular, what does $l$ precisely represent here? Is it the 40-node neighborhood of each residue's local graph?\n\n- What is the difference between the disentangled attention in Eq. (1) and a regular attention mechanism, where the tokenized representations of each residue are simply concatenated together?\n\n- What is $k$ in Eq. (2)? Is it the same as the number of clusters in $k$-means clustering of structural embeddings?\n\n- Why is the structure $s$ in Eq. (6) left unmasked? Shouldn't you also mask $s$ when you make the residue tokens in the sequence data? On a related note, shouldn't the structure of the variant sequence be changed in the first term of Eq. (8) when calculating the scores?\n\n- It would be helpful if you could also compare your approach with ESM-GearNet, and also provide the results of supervised downstream prediction with frozen embeddings in Section 4.2 (as opposed to fine-tuning the model).\n\n- Minor points: $e$ should be replaced by $e_i$ on line 154, and $m$ should be replaced by $M$ in Eq. (6).", "limitations": "As the authors allude to, the main limitation of the model is the assumption of the availability of structural data and the computational complexity of the structural quantization process.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Fusing protein structure and sequence data is an important and timely research topic, and the proposed architecture provides a novel way of combining these two modalities to derive more expressive protein representations.\n\n- The numerical results and ablation studies are quite comprehensive and show considerable gains over state-of-the-art benchmarks.\n\n- The paper is very well written and easy to follow for the most part.", "weaknesses": "The main weakness of this work, in my opinion, is that both structural and sequential data are required for the model to derive the final protein representations, since the amount of available structural data is lower than that of sequence data. I wonder if there is any way in which the model can operate on sequence input alone, with the structure knowledge embedded into the model parameters. Especially, how does the model work on protein sequences for which no structural data is available? Is there a \"sequence-only\" mode that the model can revert to (e.g., by feeding a \"default/noise\" structure input alongside the actual sequence)? I see the ProSST (-structure) model in Table 1, but that seems to be a model that was only trained on sequence data, so the training pipeline and its parameterization are completely different than the complete ProSST framework.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your thoughtful feedback on our work. We would like to address your questions and concerns as follows:\n\n**Weaknesses**\n\n**W1.** Feeding a \"default/noise\" structure input alongside the actual sequence was not included in the current version, but we have a solution for it, and this solution is not very sophisticated. We trained ProSST(MST), where MST stands for Masked Structure Training. During pre-training, each sample's structure sequence has a 50% probability of being replaced by a fully masked sequence [1,1,…,1]. This approach simulates the scenario of missing protein structures. Although ProSST(MST) outperformed ProSST(K=0) in sequence-only mode, it was inferior to both the original ProSST model with structure inputs and ProSST(MST) with structure inputs.\n\nWe offer two approaches for sequence-only proteins:\n\n- Approach 1: Use AlphaFold or ESMFold to predict structures.\n- Approach 2: Use the ProSST(MST).\n\nWe have evaluated these methods on the ProteinGym benchmark and perplexity on the validation set. The results are as follows:\n\n|Model|Structure|ProteinGym|Perplexity|\n|---|---|---|---|\n|ProSST (K=2048)|AlphaFold2|0.504|9.033|\n|ProSST (K=2048)|ESMFold|0.471|9.144|\n|ProSST (MST)|Missing|0.438|10.325|\n|ProSST (MST)|AlphaFold|0.456|9.447|\n|ProSST (K=0)|Missing|0.392|12.190|\n\nRows 1-2 show the performance of Approach 1. Rows 3-4 show the performance of the new model ProSST(MST). When using AlphaFold, ProSST (MST) is inferior to ProSST (K=2048). We will add a discussion on how to apply ProSST to the revised paper.\n\n**Questions:**\n\n**Q1.** The reason is that **40** covers most cases. We selected the nearest **40** residues within a distance of **10 Å**, according to [1]. As shown in **Figure R1**, local structures with more than **40** neighbors account for only 0.00052% (5.2 x 10^-6), indicating that our choice covers most cases.\n\n**Q2.** They are not the same. $L$ is the number of residues in a protein, also referred to as the length of the protein, while $l$ is the number of nodes in a graph. A protein contains $L$ residues, each corresponding to a local structure, which in turn corresponds to a graph $G$. For any arbitrary graph $G$, $l$ is the number of nodes in it.\n\n**Q3.** The GVP encoder includes a six-layer message-passing graph neural network in which a geometric perceptron replaces the MLP to ensure translational and rotational invariance of the input structure. Our GVP encoder is consistent with the original GVP-GNN [2], except that we removed the residue type information. The GVP encoder is trained from scratch. We will provide detailed descriptions and parameterizations of the GVP in the Appendix. \n\n**Q4.** Both choices are worse than discrete structure tokens. (1) Using the centroid embedding (K=2048) as input yields inferior results:\n\n|Structure Inputs|ProteinGym|DeepLoc|Perplexity|\n|---|---|---|---|\n|Centroid Embeddings|0.462|91.73(±0.24)|9.932|\n|Structure Tokens|0.504|94.32(±0.10)|9.033|\n\nWe believe it is because the Transformer requires a learned structure token embedding rather than a fixed embedding. \n(2) Directly using continuous local structure embeddings as structure inputs leads to overfitting, as shown in **Figure R2**. Similar results have also been observed in SaProt.\n\n**Q5.** $l$ represents the total number of nodes in a graph $G$. Average pooling refers to averaging the embeddings of all nodes within the graph. $l$ is not always equal to 40. Because when there are not more than 40 residues within a 10 Å distance, $l$ will be less than 40.\n\n**Q6.** The disentangled attention contains multiple regular attention mechanisms, allowing the model to learn different attention patterns for better contextual representation. The computation of the disentangled attention is also less than that of directly concatenated self-attention because the complexity of the attention mechanism increases quadratically with the sequence length. Although disentangled attention involves additional attention calculations, it does not increase the sequence length.\n\n**Q7.** The $k$ in Eq.(2) should actually be $L_{max}$, which is the cutoff of relative position. We will correct it in the paper.\n\n**Q8.**  (1) Because our model aims to learn the contextual representation of residues rather than structure tokens, structure $s$ in Eq.(6) is left unmasked. (2) When calculating scores of variants, we utilize the structure of wild sequences because mutants only slightly alter the structure of proteins [3], and the structure of mutants is difficult to predict[4]. Other structure models, including ESM-IF, SaProt, ProteinMPNN, etc., also use the wild-type structure when scoring variant sequences.\n\n**Q9.** We have compared ProSST to ESM-GearNet and ProSST(fixed parameters) in fine-tuning downstream tasks. The results are as follows:\n\n|  |  | DeepLoc | Metal Ion Binding | Thermostability | GO-MF | GO-BP | GO-CC |\n| --- | --- | --- | --- | --- | --- | --- | --- |\n| ESM-GearNet | 690M | 93.55 | 74.11 | 0.651 | 0.676 | 0.516 | 0.507 |\n| ProSST(fixed) | 110M | 92.36 (±0.24) | 74.27 (±0.15) | 0.697 (±0.06) | 0.651 (±0.013) | 0.479 (±0.013) | 0.482 (±0.009) |\n| ProSST | 110M | 94.32(±0.10) | 76.37(±0.02) | 0.726(±0.04) | 0.682 (±0.003) | 0.492 (±0.004) | 0.504 (±0.002) |\n\nThe scores of ESM-GearNet are derived from SaProt. Note that the slightly different scores compared to the manuscript are due to the updated random seed and dataset. We have re-evaluated ProSST with different seeds to compute the average performance. Additionally, the evaluation was conducted on the updated GO dataset by SaProt.\n\n**Q10**. We will correct them. Thank you.\n\nReferences\n\n[1] Discovery of Novel Gain-of-Function Mutations Guided by Structure-Based Deep Learning.\n\n[2] Learning from Protein Structure with Geometric Vector Perceptrons.\n\n[3] A folding space odyssey.\n\n[4] Can AlphaFold2 predict the impact of missense mutations on structure?", "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": "", "extra_scores": {}}, {"review_id": "9OxL3C40Cd", "reviewer": "Reviewer_cgxa", "summary": "The paper proposes ProSST (Protein Structure-Sequence Transformer), a novel PLM incorporating sequence and structure information. ProSST can be split into two parts: a modified version of the Transformer architecture and a quantization module. The quantization module consists of a structure encoder using a Geometric Vector Perceptron (GVP). The authors train the GVP on a large dataset of protein structures (18.8 million) in which they encode residues by contextualizing local structures with surrounding ones in a manner more robust than current methods. They then discretize local residues into tokens using a pre-trained k-means clustering model. ProSST adds relative positional information to the structure before feeding it into the Transformer architecture. Instead of standard self-attention, the authors propose using disentangled attention, a method that allows the model to process contextual information from sequence and structure tokens. They test their model on various downstream protein function tasks, for which they achieve state-of-the-art performance across almost all metrics.\nThe main contributions of the paper are as follows:\n● Introduced a novel method for protein structure quantization by incorporating contextual structural information into a PLM using a structure autoencoder\n● Proposed a nonstandard attention calculation that allows for sequence and structure information to be incorporated into protein function prediction", "questions": "Is there a reason why several of the baselines (e.g., EVmutation, DeepSequence, WaveNet, UniRep, RITA, ProGen2, and VESPA) and datasets (DMS data other than thermostability) from ProteinGym are not included in this study?", "limitations": "- Contribution is limited given that now ESM3 and many other PLMs consider structure in the language model\n- There has been no discussion on several protein functions that do not benefit from rigid 3D structure, as predicted by AF2/3, e.g., disordered protein functions with floppy regions: how would ProSST handle those cases?\n\nConclusion:\n\nThis paper provides interesting methods for improving PLM performance by intelligently incorporating structural information into protein function prediction problems. The two novelties of this paper come from their quantization method and disentangled attention application. Despite providing compelling methods and promising results, there are many details that are glossed over, which over the course of the paper accrue many concerns from the reader on the validity of the results. My suggestion to the authors is to provide more insight into their choice of disentangled attention (related works section) and analysis of the results (conclusion). Similarly, the experiment section and supplemental materials need to be updated with enough information for the reader to faithfully reproduce the results according to best practices. In general, there is enough concern from the experimental results that warrant justification and response from the authors.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper clearly defines their contributions at the beginning of the paper and justifies most of the design choices through a proper literature review. \n- The paper presents current gaps in research and proposes relevant methods to bridge this gap.\n- The methods section provides sufficient design details, allowing for the model architecture to be reproduced. \n- Similarly, the formulation for both disentangled attention and the training objective are well defined. \n- The tasks outlined in the experiments section range in variety, covering many different facets of protein function prediction. \n- Following this, the ablation study provided by the authors is robust in its analysis of architecture choices, justifying empirically the choices that are made. \n- Overall, some of the claims made in the introduction are justified by the experimental results.", "weaknesses": "- Despite having clearly formulated reasons for choosing their quantization methods, the authors need a similar justification for disentangled attention. There is little discussion on this design choice until after the related works section. Similarly, there is little experimental analysis of disentangled attention outside the ablation study. This is surprising as the departure from standard self-attention allows for a nearly four percent increase in performance, which is just as impactful as their quantization module. The paper talks about learning the connections between structural tokens and sequence tokens as the main contribution of ProSST but never demonstrates the connection.\n\n- The paper conveys the generalization power of their model by testing it on multiple downstream protein function prediction tasks. ProSST outperforms nearly every other model on these benchmarks, yet some experimental setups raise concerns. Most glaringly, the paper provides little information on their testing regimen across all training tasks. In the experiments that do outline testing, the authors describe extremely small validation sets. For example, in the training phase, the authors train ProSST on 18.8 million samples, with 100,000 being used for validation (~ 99:1 split). Similarly, the split for training the structure codebook is 31,070 train points with only 200 points for validation (~ 99:1 split). The paper needs to provide more information on the testing regimen for the fine-tuning tasks. Along with omitting these training details, the authors seemingly make many mistakes in reporting their results. First and foremost, the paper provides no error or significance analysis despite including empirical evaluations of ProSST (which is required in the NeurIPS paper checklist). \n\n- In Table 3 and Table 4 on page 8, the paper evaluates the effectiveness of ProSST, marking the best results in bold. When comparing the best results from both tables, it is clear that the ablation study was done with K=2048 being their clustering hyperparameter due to the matching accuracies (although this is not reported explicitly). However, almost all of the other metrics outside of accuracy for this experiment do not match, which is either a copying mistake by the authors or is indicative of a different experiment. If the latter explanation is true, then the variance in the results indicates that the paper should have been more thorough in its analysis of the model’s performance across initializations and data splits. Along with the extremely small validation set sizes and little testing information, this makes me skeptical of the empirical results section. In addition, in Table 1, the authors incorrectly bold their model’s performance on the NDCG metric when it is clear that five other models outperform theirs on this specific metric.\n\nAnother area for improvement of the paper is its analysis of the structure quantization module. The model has a four percent increase in performance when jumping from K=0 to K=20 and has diminishing returns from there onward. At K=0, there are no structure tokens passed to the model, meaning that disentangled attention is unable to work effectively in this experimental setting. This becomes an issue when you compare the results to the ablation study on disentangled attention; it is difficult to compare the effects of the quantization module to disentangled attention. Because of this, their results might suggest that the increase in performance is due to disentangled attention, not necessarily their quantization module.\n\nThe conclusion section is also a weakness of thenot give much insight into the methods proposed and instead reviews the methods section. It lacks a thorough ana paper. It does lysis of the paper's results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your meticulous feedback. Below are our responses.\n\n**Weakness:**\n\n**W1.**  We conducted additional experiments to analyze disentangled attention. \n\n**Experiment 1**: We replaced all structure tokens in the test set with zeros or random numbers from a uniform distribution and re-evaluated ProSST. The results are:\n\n|Structure|ProteinGym|Perplexity|\n|---|---|---|\n| Original|0.504|9.033|\n| All-zero|0.112|14.524|\n| Random|0.182|14.024|\n\nIncorrect structure tokens decreased performance, suggesting that **disentangled attention learned the sequence-structure relationship.** Otherwise, performance would have been less affected.\n\n**Experiment 2**: We train ProSST (K=1), where the structure tokens are replaced with a constant value of 1. This setup helps preserve the disentangled attention mechanism. If ProSST (K=1) still improves performance, it indicates that the improvement is solely due to the disentangled attention.\n\n| |DeepLoc|ProteinGym|Perplexity|\n|---|---|---|---|\n|ProSST(K=2048)|94.32(±0.10)|0.504|9.033|\n|ProSST(K=1)|89.48 (±0.24)|0.390|12.182|\n|ProSST(K=0)|89.77 (±0.26)|0.392|12.190|\n\nThere is little difference between K=1 and K=0, as their perplexity curves (see **Figure R3**) almost overlap. This indicates that **disentangled attention cannot improve performance without correct structure tokens**.\n\n**Experiment 3:**  Visualizing the learned attentions. We show different types of attentions on Green Fluorescent Protein (GFP), including 238 residues, in **Figure R4**. We can see that **disentangled attention learns different attention patterns**, with notable differences between R2S and S2R.\n\n**W2.**  \n\n(1) **Dataset Split** \n\n(1.1) There is no test set for Pre-training for the structure codebook and Transformer. We use only a small validation set (without a test set) for pre-training to maximize data usage, similar to pre-trained models like ESM-2, which use 0.5% of data for validation [1].\n\n(1.2) The fine-tuning datasets have train/valid/test splits like SaProt's and are downloaded from SaProt. Data statistics will be provided in the revised paper.\n\n(1.3) In this zero-shot mutant effect prediction, all data are in the test set.\n\n**(2) Error and Significance Analysis** \n\n(2.1) We will add error analysis to fine-tuning datasets. We use the same hyperparameters and repeated experiments five times with different seeds. The average performance served as the metric, with the standard deviation as the error. The results are:\n\n| |DeepLoc|Metal Ion Binding|Thermostability|GO-MF|GO-BP| GO-CC |\n|---|---|---|---|---|---|---|\n|ProSST|94.32(±0.10)|76.37(±0.02)|0.726(±0.04)|0.682(±0.003)|0.492(±0.004)|0.504(±0.002)|\n\nNote that the slightly different scores compared to the manuscript are due to the updated random seed. Additionally, the evaluation was conducted on the updated GO dataset by SaProt.\n\n(2.2) The error in the ablation study tables will be updated in the revised version, which would not change the performance ranking. \n\n(2.3) We used the non-parametric bootstrap method for zero-shot mutant fitness prediction to test differences between the baselines and ProSST. In all tests, the p-values were less than 0.01.\n\n**W3.** The inconsistency between Table 3 and Table 4 was due to a copying error. The data for ProSST (K=2048) in Table 3 is correct, but an error occurred when copying to Table 4. The performance of ProSST in Table 4 on ProteinGym should be 0.504 for Spearman, 0.777 for NDCG, and 0.239 for Top-recall. We will correct these and fix the black marking issues in all tables.\n\n**W4.** Disentangled attention cannot improve performance without correct structure tokens. We train ProSST (K=1), where the structure tokens are replaced with a constant value of 1. This setup helps preserve the disentangled attention mechanism. If ProSST (K=1) still improves performance, it indicates that the improvement is solely due to the disentangled attention. However, ProSST(K=1) is not better than ProSST(K=0). The results can be referred to response to **W1**.\n\n**W5.** We will emphasize the reasons for selecting discrete structure tokens and using disentangled attention, and we will highlight the analysis of the relationship between them, as discussed in W1.\n\n**Questions:**\n\n**Q1.** We select baselines with high Spearman correlations from the ProteinGym webpage [2]. We also compared the mentioned baselines to ProSST and will update it in our paper. The results are as follows:\n\n|Model|ρs|NDCG|Top-recall|\n|---|---|---|---|\n|EVmutation|0.395|0.777|0.222|\n|DeepSequence|0.407|0.774|0.225|\n|WaveNet|0.373|0.761|0.203|\n|RITA|0.372|0.751|0.193|\n|UniRep|0.19|0.647|0.139|\n|ProGen2|0.391|0.767|0.199|\n|VESPA|0.394|0.759|0.201|\n|ProSST|0.504|0.777|0.239|\n\nSince structure models often excel in the Stability subset, we additionally compared ProSST with other models in this subset. We will include other ProteinGYM subsets (Stability, Activity, Binding, Expression) in the Appendix. ProSST achieves state-of-the-art (SOTA) performance in the Stability, Binding, and Expression subsets.\n\n**Limitations:**\n\n**L1.** We believe our model is a valuable addition to structure-aware protein language models. Note that ESM-3 was introduced in June 2024, which was after our submission to NeurIPS in May 2024.\n\n**L2.** We utilize AlphaFold to predict structures of disordered proteins. We provide the relationship between AlphaFold pLDDT and the performance of structure models like ProSST, SaProt, and ESM-IF on ProteinGYM, as shown in **Figure R5-R7**. There is a positive correlation between pLDDT and model performance: for ProSST, $\\rho=0.30$; for SaProt, $\\rho=0.31$; and for ESM-IF1 $\\rho=0.42$, where $\\rho$ is the Pearson correlation coefficient. This may indicate that structure models may not perform well on disordered proteins. \n\nReferences\n\n[1] Evolutionary-scale prediction of atomic-level protein structure with a language model.\n\n[2] https://proteingym.org/benchmarks", "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": "", "extra_scores": {}}, {"review_id": "cARXk8wZs4", "reviewer": "Reviewer_1pAx", "summary": "This paper focuses on the protein representation task. ProSST introduces a quantized method to combine the information from protein structure. Additionally, the authors also propose a disentangled attention mechanism on top of the quantized structures to learn the relationship between residue and structure token sequences. They claim that the proposed method outperforms state-of-the-art methods in several tasks under both zero-shot and supervised learning settings.", "questions": "1. How does ProSST deal with extremely long protein sequences?\n2. How to deploy this method to sequence-only datasets like the PEER benchmark [2]? It seems ProSST requires both sequence and structure as inputs.\n\n[2] Peer: a comprehensive and multi-task benchmark for protein sequence understanding", "limitations": "My main concern is the experiment setting mentioned in the Weaknesses part. Authors are welcome to answer my questions and I am leaning to raise my rate.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The paper is easy to follow and the presentation is clear. \n2. The authors proposed an interesting approach to merge representations of two protein language models based on sequences and structures. The authors provide extensive experiments to confirm the effectiveness of the proposed methods.\n3. While cross-attention between sequence and structure has been widely explored before, the proposed disentangled attention seems novel. The introduction of the relative position encoding also seems beneficial according to the ablation.", "weaknesses": "1. The novelty is somewhat limited. I do believe such work is meaningful for computational biology, but this is not the first hybrid approach to protein language models (such as ESM-GearNet [1]).\n2. Regarding the experiment, I strongly suggest comparing ProSST to ESM-GearNet, which I think is a fairer comparison considering the utilization of both sequence and structure. \n3. Besides comparing the number of parameters, I suggest the authors provide experiments on inference speed. In my experience, clustering is often time-consuming, especially for long sequences and multiple cluster centers (K > 2000 in this paper).\n\n[1] A Systematic Study of Joint Representation Learning on Protein Sequences and Structures", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your thoughtful feedback on our work and provide a thorough response addressing your main concerns.\n\n**Weakness**\n\n**W1.**  We agree that our work is not the first hybrid approach to protein language models, and we have cited works such as ESM-GearNet and LM-GVP in Section 2.1. However, what sets our work apart is our disentangled attention mechanisms and structure quantization module. We redesigned the self-attention mechanism in the Transformer to accommodate multiple sequence inputs, including protein sequence, structure sequence, and relative position matrix sequence. Meanwhile, the structure quantization module can encode local structure of residues with more comprehensive information. Evaluation results demonstrate the efficacy of our model across various protein understanding tasks, especially in zero-shot mutant fitness prediction. We contend that the overall model, local structure encoding and quantization, the disentangled attention mechanism, and the evaluation results jointly represent a valuable contribution to the field of computational biology.\n\n**W2.** We evaluated ESM-GearNet on the fine-tuning task and compared it to ProSST.  The results are as follows:\n\n|  | DeepLoc | Metal Ion Binding | Thermostability | GO-MF | GO-BP | GO-CC |\n| --- | --- | --- | --- | --- | --- | --- |\n| ProSST | 94.32(±0.10) | 76.37(±0.02) | 0.726(±0.04) | 0.682 (±0.003) | 0.492 (±0.004) | 0.504 (±0.002) |\n| ESM-GearNet | 93.55 | 74.11 | 0.651 | 0.676 | 0.516 | 0.507 |\n\n\nThe scores of ESM-GearNet are derived from SaProt. Note that the slightly different scores compared to the manuscript are due to the updated random seed and dataset. We have re-evaluated ProSST with different seeds in order to compute the average performance. Additionally, the evaluation was conducted on the updated GO dataset by SaProt.\n\n**W3.** The clustering model is trained only once. The training and inference speeds of the clustering model are as follows:\n\n| K | Clustering Model Training Time | Clustering Model Inference Time |\n| --- | --- | --- |\n| 20 | 5m 22s | 5 ms |\n| 128 | 9m 39s | 5 ms |\n| 512 | 26m 41s | 8 ms |\n| 1024 | 51m 16s | 15 ms |\n| 2048 | 98m 4s | 26 ms |\n| 4096 | 185m 50s | 250 ms |\n\nWe also show the inference time of ProSST(110M), SaProt(35M), SaProt(650M) in **Figure R8**. ProSST(110M) is faster than SaProt (650M) and slower than SaProt (35M). The experiments are conducted on a server with one RTX 3090 GPU and two Intel 8468 CPUs. We will discuss inference speed in revised paper. \n\n**Questions**\n\n**Q1.** ProSST utilizes relative positional encoding, which supports inference without constraints on sequence length. For pre-training, we removed proteins longer than 2048 residues for training efficiency (extremely long sequences may cause an overflow of CUDA memory). However, we do not do any truncation during fine-tuning and zero-shot mutant effect prediction. Furthermore, we have evaluated the perplexity of ProSST on a long protein dataset containing 594 proteins longer than 2048 residues. The statistics of the sequence lengths and evaluation results are as follows:\n\n| Count | Length Mean. | Length Min. | Length Max. | Length Std. | Perplexity |\n| --- | --- | --- | --- | --- | --- |\n| 594 | 2313 | 2049 | 2699 | 180 | 9.013 |\n\nThe perplexity is **9.013**, which is similar to the validation set perplexity of **9.033**, suggesting that ProSST can effectively understand long protein sequences. Another notable point is that extremely long protein sequences are very scarce in the nature, as almost **99.7%** of proteins are shorter than 2048 [1]. We will add the discussion about long sequences to the manuscript.\n\n**Q2.:** The current implementation requires to input both sequence and structure. Extending the framework to accept sequence-only datasets is possible. Here, we provide three approaches:\n\n**Approach 1:** Use AlphaFold or ESMFold to predict structures. \n\n**Approach 2:** Use ProSST(MST), trained with structure masking and supporting sequence-only inputs. We will release a model ProSST(MST), an extension of the ProSST model that incorporates Masked Structure Training (MST). The MST means that during pre-training, each sample's structure sequence has a 50% probability of being replaced by a fully masked sequence [1,1,1,1,1,…,1], simulating missing protein structure. Therefore, when applying ProSST to sequence-only datasets, we need to use the masked sequence [1,1,1,1,1,…,1] as a substitute for the structure token sequence.\n\nWe have evaluated these methods on the ProteinGym benchmark, binary localization prediction (BLP) from the PEER benchmark, and perplexity on the validation set. The results are as follows:\n\n| Model | Structure Source | ProteinGym | BLP(Peer) | Perplexity |\n| --- | --- | --- | --- | --- |\n| ProSST (K=2048) | AlphaFold2 | 0.504 | 94.32 | 9.033 |\n| ProSST (K=2048) | ESMFold | 0.471 | 92.73 | 9.144 |\n| ProSST (MST) | Missing | 0.438 | 91.84 | 10.325 |\n| ProSST (MST) | AlphaFold | 0.456 | 92.31 | 9.447 |\n| ProSST (K=0) | Missing | 0.392 | 89.65 | 12.190 |\n\nRows 1-2 show the performance differences between AlphaFold and ESMFold. Rows 3-4 show the performance of the new model ProSST(MST). Row 5 shows the performance of the sequence-only model. A discussion on how to deploy ProSST on sequence-only datasets will be added to the revised version.\n\nReference\n\n[1] https://www.uniprot.org/uniprotkb/statistics", "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": "", "extra_scores": {}}, {"review_id": "sERR51PqRA", "reviewer": "Reviewer_jE9F", "summary": "In this paper the authors introduce ProSST, a method for training protein language models on both sequence and structure. They use a geometric vector perceptron trained with a structural denoising objective to obtain a structure encoder, and then perform k-means clustering on the embeddings of local residue neighborhoods in CATH to obtain a codebook of structural tokens. A masked language model is trained with disentangled attention, conditioned on the structural token inputs and the corrupted sequence.", "questions": "Can the authors give some indication about the increased time and computational cost of tokenizing structures using their approach?\n\nWhy is the two-step structural quantization approach used, rather than a typical VQ-VAE as used in FoldSeek? Are there advantages to training stability, simplicity, or inference time tokenization? \n\nHave the authors investigated using smaller neighborhood sizes than 40 nearest neighbors? Is this a route to reducing the computational burden of structural quantization?", "limitations": "Yes", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The approach follows a clearly promising and useful trend of incorporating explicit structural information into protein language models. ProSST improves on prior work by expanding the codebook size of structural tokenizers and showing the effects of disentangled attention for mutation effect prediction and improved pre-training perplexity. The authors ablate their design decisions to give evidence that these two changes are indeed improvements on traditional pLM design. The results on ProteinGym seem to significantly improve on pLMs without structural information, and even show noticeable improvements over SaProt.", "weaknesses": "It is not clear why the structural quantization is a two-step process, involving training a separate structure encoder with a denoising objective, and then clustering to form the local codebook. An advantage of SaProt is the simplicity of the codebook and almost negligible cost of structural tokenization. How does this compare to the pre-processing needed for ProSST?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for recognizing the novelty and contribution of our work. Your insightful comments helped us enrich the analysis a lot. With the response below, we hope to address your concerns properly.\n\n**Weakness:**\n\n**W1.** In fact, we do not need to train the structure encoder and the clustering model for every protein. Our structure quantization serves as a data preprocessing step. The codebook is only pre-trained on a pre-training structure encoder dataset (Appendix A.2, Dataset for training structure codebook). SaProt uses FoldSeek [1] as its structure codebook, which also needs pre-training. Since training a codebook is FoldSeek's task, it isn't explicitly detailed in SaProt's article. The primary difference between our structure quantization module and FoldSeek is that our local structure considers up to the nearest 40 residues within a 10 Å distance, whereas FoldSeek only incorporates structure information between preceding and succeeding residues.\n\n**Q1.**  Tokenizing a protein structure containing $L$ residues involves three steps:\n\n**Step 1.** For  $i = 1, 2, 3, \\ldots, L$, generate a local structure for residue $r_i$, converting it into a graph $G_i$. (Left side of Figure 2c) \n\n**Step 2.** For  $i = 1, 2, 3, \\ldots, L$, encode each graph $G_i$ into a continuous vector $e_i$  using  the trained structure encoder.\n\n**Step 3.** For  $i = 1, 2, 3, \\ldots, L$, utilize the trained clustering codebook to convert $e_i$ to a token $s_i$.\n\nWe also evaluated the time required to convert a protein into structure tokens for five proteins of different lengths using a server equipped with an RTX 1080 GPU and a Intel Xeon E5-2690 CPU. Step 2 was executed on the GPU, while steps 1 and 3 were executed on the CPU. The results are as follows.\n\n| Protein Name (Uniprot_ID) | Length (Local structures) | Step 1 | Step 2 + Step 3 |\n| --- | --- | --- | --- |\n| CCDB_ECOLI_Adkar_2012 | 101 | 0.29s | 4.43s |\n| ESTA_BACSU_Nutschel_2020 | 212 | 0.67s | 4.27s |\n| PTEN_HUMAN_Matreyek_2021 | 403 | 1.06s | 4.45s |\n| ENV_HV1B9_DuenasDecamp_2016 | 853 | 3.24s | 5.63s |\n\nWe believe the speed is acceptable. Nevertheless, one of our important future tasks is to use parallelization techniques to optimize the speed of subgraph segmentation (Step 1) and local structure encoding (Step 2). We are working on this development and will release it in the camera-ready version.\n\n**Q2.** Both our structure quantization module and FoldSeek involve two steps: training a structure encoder and a clustering model. We chose the denoising autoencoder for its focus on representation learning and simplicity, which better suits our needs. While VAE and VQ-VAE are generative models, we do not require a structure generation model. The training curve of our structure encoder is shown in **Figure R9**. The change of training loss and validation loss is stable.\n\n**Q3.** Smaller neighborhood sizes may ignore some important neighbors. Although using neighborhood sizes is a route to accelerate average pooling, it may mislead some neighboring nodes. We selected the nearest **40** residues within a distance of **10 Å**, according to [4]. As shown in Figure R1, local structures with more than **40** neighbors account for only 0.00052% (5.2e-6). Thus, 40 can cover most cases.\n\nReferences\n\n[1] Clustering predicted structures at the scale of the known protein universe.\n\n[2] Convolutions are competitive with transformers for protein sequence pretraining.\n\n[3] https://www.uniprot.org/uniprotkb/statistics\n\n[4] Discovery of Novel Gain-of-Function Mutations Guided by Structure-Based Deep Learning.", "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": "", "extra_scores": {}}]

Author rebuttal:
# Rebuttal

We thank all the reviewers for their detailed comments and insightful suggestions. We have incorporated additional experiments and analyses based on the recommendations.

Here, we present a concise overview of the major enhancements that have been universally implemented, focusing on:

- We provide approaches for deploying ProSST on sequence-only proteins, including using AlphaFold or ESMFold to predict structures, or using the new ProSST (MST) model trained on randomly masked structure sequences that support sequence-only input.
- We conducted additional analysis of disentangled attention: (1) The significant decline in ProSST's performance on incorrect and random structure tokens indicates that disentangled attention has learned to leverage structure tokens rather than ignore them. (2) Evidence is provided using ProSST (K=1) to show that the model's improvement is not solely due to disentangled attention. (3) The different types of attentions within disentangled attention are visualized using green fluorescent protein (GFP).
- We performed additional comparisons with more baselines, including adding ProSST (fixed) and ESM-GearNet to supervised fine-tuning downstream tasks, and EVmutation, DeepSequence, WaveNet, RITA, UniRep, ProGen2, and VESPA to zero-shot mutant effect prediction.
- We conducted additional speed analyses, including the training time of the clustering model, the time for structure tokenization, and the inference time of the Transformer model.
- Errors and spelling mistakes have been corrected.
- We reported experimental error for fine-tuning tasks with different seeds, as well as significance tests for zero-shot mutant effect prediction.

The figures are included in the attached PDF, with corresponding references provided in each of our responses. We believe that a brief overview of these additional results will provide clear context for understanding the significance of the updates we have made.

**Figure R1.** The distribution of the number of residues within 10 Å distance. 

This figure explains why we choose the nearest 40 residues as the neighborhood of each node. We only build edges for two residues when their distance is less than 10 Å. According to this figure, for an arbitrary residue, the number of residues within 10 Å of it almost never exceeds 40. Therefore, we chose 40 as the threshold, which covers almost all cases.

**Figure R2.** Perplexity curves of pre-training on continuous local structure embeddings.

This figure explains why we do not use continuous local structure embeddings as input: they can cause overfitting.

**Figure R3.** Perplexity curves of ProSST (K=1) and ProSST (K=0).    
We trained ProSST (K=1), where the structure tokens are replaced with a constant value of 1. This setup helps preserve the disentangled attention mechanism. Although ProSST (K=1) employs disentangled attention and ProSST (K=0) does not, their training curves show almost no difference. This result indicates that disentangled attention alone cannot enhance the model's performance without correct structure tokens.

**Figure R4.** Different types of attentions on Green Fluorescent Protein (GFP). These attentions are the average of each head in the final layer of the Transformer.

We visualize the attention learned by ProSST on GFP to investigate whether disentangled attention can learn different attention patterns. A significant difference between R2R and S2R can be observed.

**Figure 5.** pLDDT vs. Spearman of ProSST on ProteinGYM. 

**Figure 6.** pLDDT vs. Spearman of SaProt on ProteinGYM.

**Figure 7.** pLDDT vs. Spearman of ESM-IF1 on ProteinGYM. 

These figures show the relationship between the Spearman's performance in mutant effect prediction and the pLDDT predicted by AlphaFold. There is a positive correlation between pLDDT and model performance.

**Figure 8.** Inference speed of different sequence lengths. (Batch size=16)

We tested the inference speed of ProSST on proteins of different lengths using a batch size of 16 on a server equipped with two Intel 8468C processors and a 3090 GPU.

**Figure 9.** Training and validation curves of the structure encoder.

We present the training and validation loss during the training process of the structure encoder, showing that both losses are stable.

Meta-review: The work introduces ProSST, a novel protein language model that simultaneously leverages sequence and structure information to produce SOTA performance on various widely-used benchmarks. The novelty is in the use of (1) vector-quantized graph autoencoder to tokenize protein structure into local structure tokens and (2) disentagled attention to model the relationship between sequence tokens and structure tokens. The consensus is that the novelty, clarity, and performance introduced in this work will be of interest to NeurIPS audience.
Decision: Accept (poster)