PETA: evaluating the impact of protein transfer learning with sub-word tokenization on downstream applications

Abstract

Protein language models (PLMs) play a dominant role in protein representation learning. Most existing PLMs regard proteins as sequences of 20 natural amino acids. The problem with this representation method is that it simply divides the protein sequence into sequences of individual amino acids, ignoring the fact that certain residues often occur together. Therefore, it is inappropriate to view amino acids as isolated tokens. Instead, the PLMs should recognize the frequently occurring combinations of amino acids as a single token. In this study, we use the byte-pair-encoding algorithm and unigram to construct advanced residue vocabularies for protein sequence tokenization, and we have shown that PLMs pre-trained using these advanced vocabularies exhibit superior performance on downstream tasks when compared to those trained with simple vocabularies. Furthermore, we introduce PETA, a comprehensive benchmark for systematically evaluating PLMs. We find that vocabularies comprising 50 and 200 elements achieve optimal performance. Our code, model weights, and datasets are available at https://github.com/ginnm/ProteinPretraining.

Scientific contribution

This study introduces advanced protein sequence tokenization analysis, leveraging the byte-pair-encoding algorithm and unigram. By recognizing frequently occurring combinations of amino acids as single tokens, our proposed method enhances the performance of PLMs on downstream tasks. Additionally, we present PETA, a new comprehensive benchmark for the systematic evaluation of PLMs, demonstrating that vocabularies of 50 and 200 elements offer optimal performance.

Graphical Abstract

Introduction

Proteins play a pivotal role in sustaining life forms and have found extensive applications in human endeavors, including gene editing [1, 2], drug discovery [3], and enzymatic catalysis [4]. Furthermore, gaining insights into protein properties or enhancing their functionality holds significant practical value, such as enhancing the function of the original protein [5] or annotating an unknown sequence [6]. Protein engineering typically follows two common approaches: laboratory-based experiments and computation-based methods. The former includes structural analysis [7], expression purification [8] and direct evolution [9], while valuable, are time-consuming and heavily reliant on domain-specific knowledge. This limitation falls short of meeting the evolving demands of both the scientific community and industry. Conversely, the computation-based modeling strategy relies on machine-learning or physics-based methods that are often not particularly accurate but are cost-effective and time-saving. Thanks to the advancements in protein sequencing technology [10], new avenues have opened up for training large-scale protein models capable of capturing a more comprehensive understanding. For instance, ESM series [11,12,13] to leverage the UniProt database [14] which contains over 200 million protein sequences or its subsets for training purposes.

In recent years, there has been significant development in pre-trained protein language models [11,12,13, 15,16,17,18,19]. A protein will be tokenized into a sequence of bio-tokens (per-AA or multi-aa) and then use a pre-trained Transformer to convert this sequence into dense vectors, which serve as representations of the protein. Typically, these models tokenize the protein sequence by naturally dividing it into its constituent amino acids. The vocabulary size for amino acids is approximately 20, contrasting with natural language models, which often encompass thousands of words or sub-words. Many tokenization algorithms [20, 21] have been effectively employed in language models to replace character-level tokenization by grouping frequently co-occurring characters into words and show different performance in pre-trained human language models [22, 23]. However, we discovered that, in the domain of proteins, there had been no systematic research evaluating how different tokenization algorithms impact protein language models. In this paper, we draw inspiration from [24] and aim to develop a universal amino acid coding approach capable of delivering robust performance across various protein-related tasks, while harnessing the benefits of knowledge sharing and transfer as shown in Fig 1. A recent study [25] has shown reducing vocabulary size will decrease the model’s performance and distort evolutionary information, and we conduct a more comprehensive study on increasing the alphabets based on pre-training. To facilitate a thorough assessment and take cues from the success of benchmark datasets in domains like computer vision and natural language processing, e.g., ImageNet [26] and GLUE [27], we have meticulously curated a collection of 33 datasets categorized into 15 distinct tasks. These datasets are integral to advancing the realm of deep learning in protein comprehension. Our PETA benchmark encompasses five groups of tasks, including protein fitness prediction, protein localization prediction, protein-protein interaction prediction, protein solubility prediction, and protein fold prediction. For each individual dataset, we evaluate the performance of three types of tokenizers, two new residue-pair tokenizers are used to train five models with distinct vocabularies respectively and one per-amino-acid (Per-AA) model acts as baseline. Two different pooling mechanisms and three random seeds are employed in downstream tasks to mitigate potential classification biases. We anticipate that our comprehensive analysis of protein tokenizers and the PETA benchmark will serve as a pivotal milestone for the continued advancement of protein language model pre-training.

Fig. 1

The protein sequence is formed into a new discrete token sequence through different word segmentation methods. As the size of the vocabulary increases, the amino acid composition of a single token becomes more complex

Our contributions are as follows:

• Creation of the PETA Benchmark: We meticulously curate the PETA benchmark, a comprehensive collection of 33 datasets categorized into 15 distinct protein-related tasks. This benchmark spans 5 diverse aspects of protein research. It provides a standardized evaluation framework for protein language models.

• Protein Tokenization Analysis: We summarize how the amino acid coding approach enhances the effectiveness of protein language models across diverse protein-related tasks. By addressing the influence of amino acid combinations, the research offers valuable insights into the optimization of protein language models.

• Comprehensive Experiments: We have pre-trained 13 protein language models with 3 types of tokenizers, and thousands of downstream experimental evaluations are conducted to ensure the validity of the results. The model weights, code, etc. are completely open-source in the community.

Related work

Protein representation learning

Representation learning harnesses knowledge acquired from large-scale corpora to generalize across various tasks. Early approaches primarily employed machine learning techniques from natural language processing, such as word2vec [28] and doc2vec [29], to extract features from protein sequences [30,31,32]. Recently, deep learning has exhibited tremendous potential by enabling models with increased capacity and deeper encoders, capable of handling millions or billions of protein sequences. ESM-1V [12], SESNet [33], and ECNet [34], which focus on predicting mutation fitness. Additionally, ESM-1b [13] and ESM-2 [11] employ mask language modeling. ProtTrans [18] pre-trains language models under various architectures [35,36,37,38,39], while XTrimo [40] aligns its pre-trained architecture with GLM [41]. Ankn [19] uses an asymmetric encoder and decoder framework and different mask probabilities to improve the pre-training performance. CPCProt [42] leverages a contrastive predictive coding loss, whereas ProGen [15, 16], UniRep [43], ProXLNet [18], ProtGPT2 [17], and Tranception [44] are pre-trained using next amino acid prediction tasks. Although many of these approaches share common objectives with natural language processing, there are also innovations like ProteinBERT [45] and CARP [46] which employ convolutional networks for downstream tasks. Some works delve into protein multiple sequence alignments (MSAs) [44, 47, 48], while others take structure-based approaches to extract topology information for inverse folding [49,50,51], protein design [52, 53], and protein engineering [54]. Notably, LM-GVP [55], MIF-ST [50], and ProtSSN [56] integrate sequence and structural information to enhance the learning of effective protein representations [57, 58]. In this benchmark, our primary focus revolves around evaluating the performance of language models utilizing different tokenization strategies.

Protein modeling benchmarks

A comprehensive benchmark has shown great influence in the traditional computer science community and driven the research direction of different works [26, 27, 59,60,61]. However, it is worth noting that the field of computing protein engineering still lacks a large-scale benchmarking framework. In contrast, the biennial Critical Assessment of Protein Structure Prediction (CASP) [62] has emerged as a gold standard for assessing advancements in protein structure prediction. In tandem with CASP, the Critical Assessment of Functional Annotation (CAFA) challenge [63] has been established to evaluate the prediction of protein functions. Several notable works, such as DeepSequence [64], Envision [65], and ProteinGym [44], focus on measuring very different functional fitness variations in response to diverse protein modifications, including substitutions and insertions/deletions. Techniques like deep mutational scanning (DMS) [66] and other protein engineering methods are used to build up these datasets. On the other hand, works like SoluProtMutDB [67], SKEMPI [68], and ProThermDB [69] concentrate on assessing specific properties concerning single amino acid variations (SAVs). Additionally, FLIP [70] offers various data partitioning methods across three protein landscapes for evaluating fitness prediction. The TAPE benchmark [71] encompasses five tasks, with three focusing on structure prediction and the remaining two targeting fitness prediction. PEER [72] encompasses seventeen biologically relevant tasks spanning five aspects of protein understanding. ProteinGLUE [73] comprises seven downstream tasks designed for self-supervised protein representation learning. DeepLoc [74, 75] provides datasets for subcellular localization classification. The STRING database [76] annotates protein-protein interactions (PPIs) with seven types of interactions. TDA [77] generates protein-related datasets and tasks tailored for drug discovery. ESOL website [78] aggregates solubility scores for ensemble E.coli proteins.

Methods

Fig. 2

The framework of PETA. (a) Pre-trained models use rotary position embedding, which possesses favorable theoretical properties and is an absolute positional encoding applicable to linear Attention. (b) We employed two distinct classification heads, namely mean pooling and attention1d pooling. The former is the most commonly used method at present, while the latter is relatively more complex. (c) Our benchmark comprises 15 downstream tasks, which can be categorized into five groups. Some of these downstream tasks include multiple datasets or data splitting methods, amounting to a total of 33 datasets

We designed PETA to answer two important questions:

• Is residue-wise tokenizer good enough for protein language model pre-training?

• How do different vocabulary sizes influence the representation ability on downstream tasks?

Most of the works choose one tokenizer aligned with previous research without much concern, to answer the first question, we utilize three amino acid segmentation strategies including residue-wise and sub-word tokenizers. For the second question, we design larger vocabulary size arranged from {50, 100, 200, 800, 1600, 3200} for the Unigram and BPE segmentation methods. The model trained under per-AA is the baseline, it has a vocabulary size of 20 common amino acids and many works adopt this [11,12,13, 18, 49]. It is worth noting that these vocabulary sizes do not contain special tokens, such as $<mask>$ or $<pad>$. In general, we utilize three tokenization methods, two types of classification heads, and two model pipelines to solve different tasks in the PETA benchmark. The framework of PETA is shown in Fig 2.

Amino acid segmentation

In this study, we utilize three classic sequence segmentation methods: per-amino-acid encoding (Per-AA), byte pair encoding (BPE) [20], and unigram language modeling (Unigram) [21] as shown in Fig 1. Per-AA focuses on individual amino acid units, enabling high-resolution analysis of subtle variations. BPE offers flexibility by segmenting sequences into subunits, effectively capturing structural information, while Unigram, based on character-level statistics, captures global sequence characteristics. These diverse methods collectively enhance our comprehensive analysis of protein sequences, each serving a unique role in addressing specific analytical requirements.

Byte pair encoding

BPE segments sequences into the most frequent subunits or tokens. This method iteratively merges the most frequently adjacent pairs of characters into a single token, thus reducing the vocabulary size and simplifying the model’s complexity. In the context of protein sequences, BPE helps in identifying and encoding common motifs and structural domains that are crucial for functional characterization. By applying BPE, we can efficiently manage and interpret large datasets, as it bridges the gap between amino acid-level granularity and holistic sequence representation. This method is widely used in various natural language models such as GPT [79] and BERT [36].

Unigram language modeling

Unigram modeling simplifies text segmentation by independently calculating the likelihood of each word based on how often it appears in the data. Unlike the BPE method, which looks at the frequency of pairs of adjacent characters, Unigram creates a list of words by finding the most likely combination of tokens to form a language model. This approach is especially useful for analyzing protein sequences because it can identify rare, yet important, amino acids or patterns that other methods might miss. Additionally, Unigram’s probabilistic approach allows it to adjust the vocabulary flexibly according to the context of the sequence, making it adaptable to new or uncommon variations. This flexibility makes it suitable for working with multiple languages or in situations with limited language data [80].

Pre-training protein language models

Model architecture

Our pre-training architecture employs RoFormer [81], an autoencoding model that adopts a BERT-like structure augmented with rotary positional embeddings, as illustrated in Fig. 2 (a). These rotary positional embeddings effectively harness positional information within sequences. Detailed hyperparameter configurations are delineated in "Experimental Setups" section. Initially, protein sequences are tokenized and transformed into one-hot encoded representations. These representations are subsequently fed into RoFormer’s encoder, which generates sets of hidden states that maintain the length consistency with the input tokenized sequence. Finally, these hidden states are transformed into a vector with a dimensionality corresponding to the vocabulary size, upon which a softmax function is applied to yield the reconstruction probability density distribution.

Pre-training objective

We employ the mask language modeling (MLM) objective for pre-training our models [35]. Given an input sequence, a subset of tokens is selected at random and replaced with a special mask token. The model is then trained to predict these masked tokens based on the unmasked context tokens. The loss function for this objective can be defined as:

$$\begin{aligned} L_{MLM}=E_{x \sim X}E_{M}\sum _{x}-\log p(x_{i}|x_{/M}) \end{aligned}$$

(1)

here, x is a sequence from the dataset X, and $x_{/M}$ represents the sequence with masked tokens removed. $p(x_{i}|x_{/M})$ is the conditional probability of predicting the correct token $x_i$ given the context $x_{/M}$. The aim is to minimize the negative log-likelihood of the true token at each masked index i, which in this case are amino acids, given the unmasked sequence as context. Intuitively, the model must learn to identify the dependencies between the masked and unmasked tokens to successfully predict the masked positions.

Language modeling perplexity

We use Perplexity (PPL) to evaluate the performance of the MLM, computed as:

$$\begin{aligned} \text {PPL}=\exp (-\frac{1}{N}\sum _{i=1}^N\log p(x_i|x)) \end{aligned}$$

(2)

where N is the number of masked tokens, as well as $x_i$ is the ith token of sequence. To account for potential unfair comparisons arising due to varying vocabulary sizes across different models, we introduce the metric of Normalized Perplexity (NPPL) range from 1 to positive infinity. The formula for Normalized Perplexity is as follows:

$$\begin{aligned} \text {NPPL}=\exp {(-\frac{1}{N\times V}\sum _{i=1}^N\log {p(x_i|x)})} \end{aligned}$$

(3)

where V is the vocabulary size.

Pre-training data

We train all models on UniRef90 [82], a comprehensive protein sequence database that contains approximately 138 million sequences from diverse life forms. This large-scale database serves as a robust training set for capturing the underlying patterns and intricacies associated with protein sequences. For model validation, a subset of 100,000 sequences is reserved from the UniRef90 database. The validation set serves multiple purposes: monitoring the training process by observing perplexity during pre-training to ensure correct model behavior, and selecting hyperparameters such as batch size and learning rate via grid search, where we split the validation set into subsets A and B (in an 8:2 ratio) for training and evaluation respectively. Additionally, it is used to compute the metric of normalized perplexity to assess the effectiveness of pre-training results. This setup ensures that the trained models are subjected to a variety of sequence patterns, thereby facilitating a more robust understanding of protein sequences. By reserving a significant number of sequences for validation, we also ensure an unbiased assessment of the model performance.

Classification head

To test the potential bias caused by different classification methods, the Mean Pooling and Attention1d Pooling are adopted under our evaluation, as shown in Fig 2 (b). The former is trained on the average of features aligned with the first dimension, using MLP and ReLU activation to make a prediction. The latter is trained on an attention mechanism and a 1D convolution layer to map different weights to embed and predict the label.

Model pipeline

Protein-wise tasks Leaning a function $y=f_{\theta }(x)$ that maps protein x to a label y, where $f_\theta$ is parameterized by a sequence-based encoder and a classification head defined upon the residue-wise or residue-pair protein embedding.

Protein-pair tasks Leaning a function $y=f_{\theta }(x)$ that maps a pair of proteins ($x_1, x_2$) to a label y, where $f_\theta$ is parameterized by a pair of siamese sequence-based encoders and a classifier defined upon the sum of the embeddings of two proteins.

Table 1 Benchmark task details. Each task, along with its task name, category, the count of datasets or splits, the source of the dataset, and evaluation metric are shown below

Benchmark tasks

The PETA benchmark includes 15 tasks within 5 groups and 33 datasets in total, mainly focusing on protein-wise and protein-pair tasks. We have curated tasks from influential protein engineering applications and made updates to certain datasets to ensure their relevance and accuracy, a summary of the downstream dataset statistics is shown in Table 1.

Fitness prediction

This set of tasks aims to forecast functional attributes of proteins, which may be either discrete or continuous in nature.

GB1 fitness prediction assesses fitness scores among mutations within the GB1 landscape from [86]. Given a protein sequence x, we map it to a regression value $y \in \textrm{R}$, where a fitness score of 1 represents the wild-type and 0 indicates non-binding affinity. In our analysis, we utilize all the dataset splits proposed in FLIP [70], encompassing “one-vs-rest”, “two-vs-rest”, “three-vs-rest”, “low-vs-high” and “Sampled”. For example, “one-vs-rest” indicates the wild type and single mutants are assigned to train and validation, while the rest are assigned to test. More details about dataset splits can be found in supplementary materials Section.

Impact: GB1 serves as the binding domain of protein G [87], an immunoglobulin binding protein found in Streptococcal bacteria [88]. This task stands as a gold standard for investigating epistatic interactions.

AAV fitness prediction entails the evaluation of fitness values associated with Adeno-associated virus (AAV) capsid proteins [89]. Given a protein sequence x, we establish a mapping to a regression value $y \in \textrm{R}$, focusing on the mutational window spanning positions 561 to 588 from [90]. We adopt all the dataset splits from FLIP [70], which includes “Mut-Des”, “Des-Mut”, “one-vs-rest”, “two-vs-rest”, “seven-vs-rest”, “low-vs-high” and “Sampled”. More details about dataset splits can be found in supplementary materials Section 2.1.2

Impact: AAV proteins are responsible for facilitating the integration of a DNA payload into a target cell by the virus [91]. This task specifically addresses the prediction of fitness for an extended sequence subjected to mutations at select positions.

Thermostability prediction involves the analysis of protein melting curves, which are acquired through a mass spectrometry-based assay and meticulously sourced from [92]. In this endeavor, we focus on a protein sequence x, which is drawn from a vast pool of 48,000 proteins spanning 13 diverse species. Our objective is to predict a thermostability score $y \in \textrm{R}$. For this analysis, we have employed the dataset splitting strategies “Human”, “mixed_split”, and “Human_cell” as outlined in FLIP [70]. More details about dataset splits can be found in the supplementary materials Section 2.1.3.

Impact: Thermostable proteins [93, 94] demonstrate an ability to endure higher temperature conditions for extended periods or function at an accelerated rate. This task aligns closely with applications in protein engineering, particularly within industrial settings, where the enhanced stability of proteins can yield substantial benefits.

Fluorescence prediction primarily focuses on forecasting the fitness of mutants of the green fluorescent protein (GFP) [95], as documented by [96]. In this context, we are presented with a GFP mutant sequence x and aim to predict the corresponding fluorescence intensity $y \in \textrm{R}$. We leverage the dataset and split methodology derived from TAPE [71], which involves training the model on lower-order mutants and subsequently evaluating it on higher-order mutants.

Impact: Green fluorescent protein can mark particular proteins in an organic structure by its green fluorescence [97], this makes it easier for researchers to observe. This task bears significance in uncovering mutational patterns that either enhance or diminish such vital biological properties.

Stability prediction endeavors to assess the resilience of proteins within their natural environment. It involves taking a protein sequence, denoted as x, and predicting its corresponding experimental stability score, denoted as $y\in \textrm{R}$. In this pursuit, we leverage the dataset curated from [98] and employ the partitioning method introduced in TAPE [71]. The training dataset comprises proteins sourced from four rounds of experimental design, while the test dataset encompasses proteins that are Hamming distance-1 neighbors of the top candidate proteins.

Impact: The design of stable proteins in the face of mutations plays a pivotal role in the field of protein engineering [99]. This work is instrumental in various applications, such as ensuring the effective delivery of drugs before they degrade.

Protein-protein interaction prediction

Predicting protein-protein interactions (PPI) is pivotal for deciphering the intricate molecular networks underpinning cellular functions and disease mechanisms, guiding targeted therapeutic interventions.

Yeast PPI prediction involves the prediction of whether two yeast proteins interact with each other. When presented with two proteins, denoted as $x_1$ and $x_2$ from yeast, the classifier assigns a binary label $y \in {0,1}$ to signify the presence or absence of interaction between them. To accomplish this task, we utilize the yeast PPI dataset sourced from [100]. In this dataset, half of the instances represent positive cases, selected from the DIP_20070219 database of interacting proteins [101], with stringent criteria that exclude proteins with fewer than 50 amino acids or exhibiting $\ge$40% sequence identity on the full dataset. The negative cases are generated by randomly pairing proteins that lack evidence of interaction, and these pairs are further filtered based on their sub-cellular locations.

Impact: The yeast dataset serves as a widely recognized benchmark [102, 103] for assessing model performance, and yeast PPI prediction substantially enhances our comprehension of cellular processes by unveiling intricate protein interactions and providing crucial insights into the functional roles of proteins within yeast cells.

Human PPI prediction task predicts whether two human proteins interact or not. When provided with protein sequences $x_1$ and $x_2$ from humans, the predictor generates a binary label $y \in {0,1}$ to indicate the presence or absence of interaction between them. We adopt the dataset from [104], comprising positive protein pairs obtained from the Human Protein Reference Database (HPRD) [105] and negative protein pairs sourced from different cellular compartments with no documented interaction [106]. To ensure data quality, self-interactions and duplicate interactions were removed, resulting in the creation of two datasets, namely “AB” and “CD.” The “AB” dataset encompasses the entire dataset, while the “CD” dataset comprises selected proteins with identities below 25%. For evaluation, we exclusively employ the “AB” split strategy in this task.

Impact: Human PPI prediction holds immense practical significance in clinical research. Notably, insights into protein interactions linked to diseases enhance our understanding of human disease mechanisms, paving the way for innovative therapeutic strategies [107, 108].

SHS PPI prediction is to classify the type of interaction between a given protein pair. Given two protein sequences, $x_1$ and $x_2$, the model aims to predict a label y where $y \in \{0,1...,6\}$. These interaction types encompass categories such as “reaction”, “activation”, and “catalysis”, among others. Our analysis utilizes a dataset derived from interaction pairs specific to Homo sapiens, sourced from the STRING database [76]. We adopt preprocess strategies as recommended by [109] where the suboptimal health status (SHS) dataset is divided into two subsets: “SHS27k” and “SHS148k”. For computational efficiency, our study focuses solely on the “SHS27k” subset. The data is partitioned into training, validation, and test sets at a random split ratio of 8:1:1.

Impact: Understanding and categorizing the precise interactions between protein pairs is pivotal in unraveling intricate cellular mechanisms and shedding light on complex biological pathways. This knowledge not only aids in defining drug efficacy through network-based “drug-disease proximity measures” [110] but also plays a crucial role in interpreting the outcomes of genome-wide association screens [111].

Localization prediction

Identifying the localization or local-related biological mechanism of proteins within various cellular compartments is of paramount importance in the process of functional annotation.

Subcellular localization prediction aims to dig out the specific cellular location of a given natural protein. Given a protein sequence denoted as x, the model assigns it to multiple possible localizations $y \in {0, 1,..., 9}$, which may include designations such as ”Nucleus” and “Cytoplasm”, among others. To accomplish this task, we utilize two datasets from DeepLoc-1 [74] and DeepLoc-2 [75]. For the DeepLoc-1 dataset, we apply the split methodology introduced by [83]. Regarding the DeepLoc-2 dataset, its original split strategy involves 5-fold cross-validation from SwissProt. In our implementation, we employ the first three partitions as training data, the fourth as validation data, and ultimately evaluate the model’s performance on the last partition and the independent test dataset of human protein atlas (HPA) [112].

Impact: The subcellular localization of proteins plays a crucial role in deciphering the fundamental mechanisms of diseases linked to abnormal subcellular localization [113, 114]. Notably, some proteins are recognized for their ability to localize within multiple cellular compartments, underscoring the intricate and pertinent nature of this research domain.

Binary localization prediction constitutes a sub-problem of the aforementioned task. The model’s responsibility is to decide whether a given protein x should be categorized as “membrane-bound” or “soluble,” denoted as $y \in {0,1}$. The datasets for training and testing are drawn from DeepLoc-1 [74], which includes an additional label system where “S” represents soluble, “M” corresponds to membrane-bound, and “U” signifies unknown localization. To train the model, we employ the same data partitioning method as introduced by [83], while excluding data points labeled as “U”.

Impact: Predicting protein localization as either membrane-bound or soluble is vital for deciphering cellular functions, particularly in signal transduction and transport [115]. It plays a pivotal role in drug discovery, enabling the design of targeted therapies against membrane proteins.

Sorting signal prediction elucidates the intricate process of subcellular localization by identifying biological mechanisms within sorting signal sequences that guide proteins to specific subcellular structures or organelles. When presented with a short sequence x, the model assigns it to one of nine classes denoted as $y \in {0,1,...,8}$, encompassing designations such as “Signal Peptide (SP)” and “Mitochondrial Transit Peptide (MT)”, among others. This constitutes a multi-label classification task, and we employ the dataset sourced from DeepLoc-2 [75]. As the original dataset lacks an official split strategy, we perform a random split with a train/validation/test ratio of 8:1:1.

Impact: Protein sorting signals facilitate the precise intracellular localization of proteins, thereby sustaining cellular homeostasis and the integrity of subcellular compartments [116, 117]. These signals typically entail interactions with partner proteins or sorting complexes. It is significant to investigate protein sorting signals for comprehending the intracellular localization and functional intricacies of proteins.

Solubility prediction

This group of tasks is to predict the protein solubility, which is critical for optimizing protein expression, purification, and drug development processes.

Binary solubility prediction aims to determine whether a protein is soluble or insoluble. When presented with a protein denoted as x, the model assigns it to a binary label, $y \in {0,1}$. The dataset and data partitioning approach is drawn from DeepSol [84], where protein sequences exhibiting a sequence identity of $\ge$ 30% to any sequence in the test set are excluded from the training set. This task shares similarities with binary localization prediction but explicitly focuses on modeling solubility.

Impact: Protein solubility is pivotal for swiftly and efficiently selecting appropriate protein samples, saving resources and time, particularly in biotechnology, drug development, and laboratory protein purification [118, 119]. It improves experiment success rates and resource allocation while advancing scientific research.

E.coli solubility prediction involves forecasting the solubility value of E. coli proteins using an ensemble database, downloadable from the eSOL website [78]. When provided with a sequence from E. coli, the model predicts a regression value, denoted as $y \in \textrm{R}$. Solubility, in this context, is defined as the ratio of the supernatant fraction to the total fraction, as determined in physiochemical experiments referred to as PURE [120]. We utilize the training and validation datasets sourced from GraphSol [85] and further partition the validation dataset into separate validation and test sets.

Impact: E. coli, as a prevalent host organism for protein expression, demands precise solubility predictions to optimize recombinant protein production, purification, and subsequent functional studies [121]. Such predictions, based on experimental data and computational models, facilitate the selection of suitable protein candidates for diverse applications [122], ranging from structural biology to drug discovery and industrial processes.

Mutation solubility prediction measures the impact of mutations on protein solubility. Given a mutated protein sequence denoted as x, the model predicts the solubility change $y \in \textrm{R}$ relative to the wild-type sequence. This task encompasses three distinct protein mutation datasets, with mutations occurring at single points within proteins such as “Beta-lactamase TEM (blat)”, “Chalcone Synthase (cs)” and “Levoglucosan Kinase (lgk)”. These datasets were sourced from SoluProtMutDB [67] which provides manually curated and reliable data in the standardized format. Data points where recorded mutations did not align with the original sequence were excluded, and the training, validation, and test datasets were partitioned in an 8:1:1 ratio.

Impact: Low protein solubility is a significant hurdle in industrial processes and is implicated in numerous human diseases [123]. Investigating the impact of mutations on protein solubility not only sheds light on the mechanisms underpinning disease development but also enhances the application of protein engineering in various industrial domains.

Fold prediction

While AlphaFold [124] and RoseTTAFold [125] have made significant strides in structure prediction, the fold prediction task still remains a rigorous assessment to evaluate the representation quality of the sequence model.

Fold prediction is the automated classification of protein sequences into one of 1,195 known protein folds, facilitating the modeling of the sequence-structure relationship. Given any sequence x, the objective is to predict the fold label $y \in {0,1,...,1194}$, determined by the backbone coordinates of the corresponding protein structure. This task utilizes the dataset from [126], originally derived from the SCOP 1.75 database [127]. Notably, this dataset meticulously addresses homologous sequence redundancy between test and training datasets through two distinct strategies: a three-level redundancy reduction at fold/superfamily/family levels and sequence identity reduction.

Impact: Fold prediction is essential for unraveling the intricate relationship between a protein’s primary sequence and its three-dimensional structure, with profound implications for fields ranging from structural biology to drug design [128].

Experiments

Experimental setups

We perform the pre-training of our models on 8 A100-80GB GPUs, using a data-parallel distribution strategy. The global batch size is set to 1024 (local batch size is set to 32), and the maximum sequence length is constrained to 1024 tokens. We employ dropout regularization with a rate of 0.1 during the pre-training phase to mitigate overfitting. The architecture comprises 12 encoder layers, with each layer having a hidden size of 768 and an intermediate size of 3072. The multi-head attention mechanism contains 12 heads, each with a dimensionality of 64. For model optimization, we utilize the Adam optimizer, with a learning rate initialized at 1e-4. The maximum number of training steps is set to 530,000. The learning rate schedule involves a warm-up mechanism for the first 2000 iterations, following which the learning rate is linearly decayed to zero. The Adam hyperparameters are configured as follows: epsilon is 1e-8, $\beta _1=0.9$ and $\beta _2=0.98$. Gradient clipping is applied with a maximum value of 5.0 to prevent exploding gradients. Our implementation leverages the PyTorch framework in conjunction with the Hugging Face library, aligning with best practices for efficient and scalable training of language models.

In the case of supervised tasks, all pre-trained model weights are kept fixed to ensure a fair evaluation of their representation capabilities. Classifiers are trained using a batch size of 256, a learning rate of 0.001, and the Adam optimizer. Early stopping is employed with a patience threshold of 20 epochs, with a maximum of 100 epochs for training. It is important to note that these hyperparameters were adopted without adjustments, drawing reference from [70]. Each individual experiment underwent training three times using different random seeds, and the final results represent the average scores obtained.

Table 2 Perplexity and Normalized Perplexity on the validation set for the BPE model

Table 3 Perplexity and Normalized Perplexity on the validation set for the Unigram model

Pre-training results

Following the pre-training phase, all models achieved a reduction in loss to an acceptable level, demonstrating effective learning from the training data. Table 2 and Table 3 present the perplexity and normalized perplexity metrics calculated on the test set for both Byte Pair Encoding (BPE) and Unigram models, respectively. For the base model employing a per-amino-acid (Per-AA) tokenization strategy, the Perplexity value is 7.78, and the corresponding Normalized Perplexity is 1.06.

In the supplementary materials, Figures S1 to S30 show the loss curves, as well as the perplexity and normalized perplexity curves for the pre-trained models. It is important to note that evaluations were performed at intervals of 10,000 steps. These figures collectively demonstrate that all models have converged to a reasonable range, substantiating their effectiveness in learning the underlying data distribution.

Benchmark results overview

Table 4 The number of datasets or splits whose average score exceeds the baseline model of 20 vocabulary size

To provide researchers with insights into how the augmentation of vocabulary size in PLMs affects embedding quality, we conducted a systematic evaluation. The scores in Table 4 represent a given vocabulary size, how many datasets or splits, on average performance, surpassing the traditional Per-AA method (with a vocabulary size of 20). For example, We have 33 datasets or splits in total, and 22 of them outperform the baseline method using a vocabulary size of 20 when using a vocabulary size of 50. This counting method is obtained by comparing the average scores on each dataset, for the extended vocabulary models. The average score on each dataset was computed from the mean of 12 experiment results (2 tokenization methods x 2 classification heads x 3 random seeds). For the baseline models, the average score on each dataset was derived from the mean of 6 experiments (2 classification heads x 3 random seeds). More detailed results of each experimental setting can be found in the supplementary materials Table S1 to S54.

Our experimental findings have led to several key insights:

• Significant Impact of Vocabulary Size. Extensive experimentation has unequivocally demonstrated that vocabulary size profoundly influences protein representation, albeit with varying degrees of impact across different types of downstream tasks. Notably, in every dataset associated with fold prediction, an inverse relationship was observed wherein enhancements in vocabulary size correlated with negative optimization.

• Existence of an Optimal Vocabulary Threshold. Contrasting with language models utilized in NLP, PLMs with an excessively large vocabulary size can potentially exert detrimental effects on downstream tasks. Specifically, when the vocabulary size surpasses 800, the majority of tasks are performed suboptimally compared to the baseline model that employs per-AA segmentation.

Downstream tasks

Table 5 Performance on Fitness Prediction and Localization Prediction

Fig. 3

Detail performances of the GB1, Thermo, and AAV datasets across different vocabulary sizes

Fitness Prediction. Table 5 showcases results for five distinct tasks under the Fitness Prediction and the evaluation metric is Spearman correlation. It is worth highlighting that the datasets for GB1, AAV, and Thermo have different splits and the details are shown in Fig 3. A key observation from the data is the mixed effects of adding more words to the vocabulary. Compared to the per-AA tokenization, the performance of most splits is improved with an expanded vocabulary. However, AAV was an exception, experiencing a significant drop in performance as the vocabulary grew, with decreases ranging from $3\%$ to $7.4\%$. In contrast, the performance of Flu benefited from a larger vocabulary. Interestingly, the average performance for GB1 and Stab began to decline after the vocabulary size reached 200, even falling below the baseline set at a vocabulary size of 20. Thermo is not sensitive to vocabulary size changes, fluctuating approximately $0.5\%$ up or down.

Localization Prediction. The right side of Table 5 presents results from five datasets under the category of Localization Prediction, and the monitored metric is Accuracy. Across all datasets and partitioning methodologies, the task of subcellular localization prediction consistently achieves a remarkable classification accuracy exceeding 90. This accuracy remains quite stable, with any changes in performance staying within $1\%$ despite differences in vocabulary size. Drawing from these experimental insights, it is evident that language models handle both multi-class and single-class prediction tasks for protein localization relatively easy. Moreover, the vocabulary size seems to have minimal influence on the prediction outcomes for protein localization tasks.

Protein-Protein interaction Prediction. Table 6 summaries the 3 datasets from PPI Prediction, and the metric is Accuracy. The table clearly shows that using more specific words helps in identifying the relationships between protein pairs. Additionally, datasets that are harder to classify show higher performance enhancements. For instance, in the Yeast dataset, the model with a vocabulary size of 1, 600 exhibited a $3.4\%$ average score increase compared to the model with a vocabulary size of 20. In SHS27k, every increase in vocabulary size resulted in performance improvements, with the biggest improvement being of $2.6\%$. In contrast, while the Human dataset showed improvements across the board, the maximum increase was a mere 0.2%.

Table 6 Performance on PPI Prediction, Solubility Prediction and Fold Prediction

Fig. 4

A detailed exposition is provided on the performance results of three distinct protein solubility mutation datasets: Beta-lactamase TEM (blat), Chalcone Synthase (cs), and Levoglucosan Kinase (lgk) across varying vocabulary sizes

Solubility Prediction. The middle section of Table 6 displays the findings for Solubility Prediction. While the evaluation metric for the three datasets in Solmut is the Spearman correlation, Esol employs MSE, BinSol and DeepSol use Accuracy as their metric. Predicting solubility regression values for Esol proves to be relatively straightforward, the lower MSE scores indicate better performance, and this remains consistent across models with different vocabulary sizes. BinSol observed the opposite situation to Esol, the average score increased by $2\%$ to $3\%$, which is relatively significant. An analysis of the Solmut datasets indicates that models with expanded vocabulary sizes have the potential to improve performance by $1\%$ to $5\%$ as shown in Fig. 4. Although most instances show enhancement, occasional instability is detected. Such variability could be attributed to the alterations in the inherent characteristics of proteins post-mutation.

Fold Prediction. The results for Fold Prediction are detailed on the right side of Table 6, and the metrics is Accuracy. In fold prediction, an important pattern is observed: as the vocabulary size enlarges, there is a clear drop in performance across all datasets. To illustrate, the superfamily split sees the most significant drop, falling by $1.3\%$ to $11.2\%$, while the family and fold splits experience declines of $5\%$ and $3.8\%$ respectively. This trend highlights a consistent drop in performance as the larger vocabulary size is. This decline could be due to the combination of multiple amino acid tokens during the encoding process, which might hide important details of the local structures. Consequently, this increase in vocabulary sizes could be negatively impacting the ability to predict structure-related tasks.

Ablation study

Table 7 The average results of different downstream task groups under the same vocabulary with varying tokenization methods are presented

Table 8 The average results of different downstream task groups under the same vocabulary with varying pooling heads are presented

From Tables 7 and 8, we can observe positive optimizations in most datasets and negative optimizations in some tasks due to the expansion of the vocabulary. Importantly, these optimizations, whether positive or negative, are independent of the tokenization method used, the type of classification head, and the random seed. The changes are solely attributed to the variations in vocabulary size. It is worth noting that, we remove Esol from Solubility Prediction due to the incomparable scale of the values.

Analysis of tokenizers. From Table 7, it can be observed that the discrepancies arising from different tokenization methods are minimal across various downstream tasks. The main source of performance variation stems from the impact of vocabulary size on model representation. Across different tasks, as the vocabulary size increases, the model performance exhibits a bell-shaped curve, showing an initial increase followed by a decline.

Impact of pooling heads. From Table 8, it can be observed that, when freezing the pre-trained model parameters and only tuning the pooling head, the performance is highly correlated with the choice of the classification head. When using the same vocabulary size, the attention1d pooling method outperforms the mean pooling method. Additionally, similar to the results in Table 7, as the vocabulary size increases, the model’s representational capacity across various downstream tasks tends to decline.

Conclusion

In this paper, we introduce PETA, a vocabulary study optimized for protein language models across a broad range of datasets. To mitigate potential biases arising from different tokenization methods, classification heads, and random seeds, for each fixed vocabulary size, we employed both BPE and Unigram tokenization methods, two classification heads (mean pooling and attention1d pooling), and experiments with three different random seeds on each dataset. Ultimately, we found that expanding the vocabulary size to some extent (50-200) generally enhances performance on downstream tasks. However, once the vocabulary size surpasses 800, the model’s representational power exhibits a broad decline across most tasks. We hope that this work and benchmark will influence the future protein language model community and contribute positively to human health, environmental development, and biomedicine.