Modeling Protein Using Large-scale Pretrain Language Model

There are plenty of data in the computational proteomics field, however, current literature is fragmented in terms of unified datasets and evaluation metrics. The methods and models introduced by researchers are often evaluated on different datasets with different evaluation metrics. To solve this dilemma, TAPE(rao2019evaluating, ) put forward a set of five biologically related tasks, including secondary structure prediction ((ss, ), (ss_cc1, ), (ss_cc2, )), contact prediction ((rh_cc, ), (ss_cc1, ), (ss_cc2, )), remote homology detection ((rh_cc, )), fluorescence ((fluorescence, )), and stability ((stability, )). Besides, commonly used models, like LSTM(LSTM, ), Transformer(vaswani2017attention, ), and ResNet(resnet, ) are implemented for these tasks, serving as benchmarks for semi-supervised representation learning. One of their conclusions is that self-supervised training is beneficial for almost all models on all tasks, doubling performance in some downstream tasks. Our work is based on standardized datasets and evaluation metrics provided in TAPE(rao2019evaluating, ).

The success of pretraining makes researchers wonder whether the in language model scale can always bring about improved performance. ProtTrans(prottrans, ) is one of the representatives, the researchers trained a series of language models with tens of billions of parameters, the largest one ProtT5-XXL with 11B parameters, and achieved excellent performance on downstream tasks such as secondary structure prediction and solubility prediction.

Different from the usual pretraining, large-scale pretraining requires distributed training techniques, including model parallelism, data parallelism, memory optimization, data synchronization, etc. Fortunately, Megatron-LM(shoeybi2020megatronlm, ) provides us with an efficient training framework for language models. We have implemented and trained our protein language model within this framework, as well as downstream classification and regression tasks.

Description  
The goal for protein pretraining is modeling data patterns in massive unlabeled sequences. One closed-related model is BERT(bert, ) from natural language processing. We made some modifications to its loss function and model structure.

Dataset  
Our work takes the dataset put forward by TAPE(rao2019evaluating, ). So some date descriptions are inherited. PFAM(pfam2019, ) is a widely-used database consisting of more than 32 million protein sequences. Sequences in PFAM are clustered into evolutionarily related groups (protein families). Leveraging this family information, TAPE constructed a test set (about 1% of the data) of fully held out families. The remaining data are used for constructing training and test sets using a random 95%/5% split. We use preprocessed PFAM from TAPE as the pretrain corpus.

Training Objective  
BERT(bert, ) original loss consists of masked language model loss and next sentence prediction loss.

For protein pretraining, we inherited the masking strategy of the masked language model (MLM) in BERT(bert, ), randomly masking 15% of all the amino acid tokens, and then train the protein model to be able to predict the masked token from the rest of tokens. As for next sentence prediction (NSP), considering that the input sequences are randomly shuffled, we assume there is no evident semantic/biological correlation between sequences. So we discard the next sentence prediction loss, only keep the masked language model loss.

As for the training objective function, we modified the loss function of BERT: BERT’s loss function includes masked language model and next sentence prediction. Considering that there is no obvious contextual semantic relationship between protein and protein, we only retain masked language model loss.

In terms of model structure, Megatron-LM (shoeybi2020megatronlm, ) proposes that when the scale of the model grows huge, the position of the layernorm becomes critical. Therefore, the sublayers in the transformer layer have been reordered. The original layernorm is in the output layer, but now it is placed ahead of the input layer to prevent the input data from drifting.

There are three classification tasks, corresponding to token, sequence, and token-pair classification.

We pretrained two large models on a 480 GPUs (Tesla-V100-32GB) cluster for about three weeks. The MLM loss and PPL of the pretrained models can be found in Table 1.

The 3B parameters model reached language model loss $MLMloss_{3B}=1.318$, perplexity $PPL_{3B}=3.736$.

The 1.2B parameters model reached language model loss $MLMloss_{1.2B}=1.335$, perplexity $PPL_{1.2B}=3.802$.

In pretraining, although the overall training iteration for the 3 billion model is only half of that for the 1.2 billion model, it reached even smaller MLM loss and PPL. This phenomenon demonstrated that, when handled properly, the expansion in model scale can contribute to the accurate capture of patterns in the data.

Out of the five tasks, the results of four tasks have been improved.

1. (1)

   Contact Prediction: Table 2

2. (2)

   Remote Homology: Table 3

3. (3)

   Secondary Structure: Table 4

4. (4)

   Fluorescence: Table 5

5. (5)

   Stability: Table 6

We visualize the contact map of protein #TBM-T0889 in Figure 6, and we can intuitively see that there are many long-range contacts (contacts that are separated by at least 24 residues). This protein sequence can be used to distinguish the ability to capture long-distance dependence of different models.

Through the visualized predictions from TAPE (7), we can see that the TAPE model (small-scale transformer) can capture medium-range contacts (contacts that are separated by 12-23 residues). As for the long-range contact prediction, there are lots of missings on the anti-diagonal belt.

ProteinLM-3B model shows very good performance in contact prediction, and the visualized prediction map confirmed this. ProtrinLM-3B can capture medium-range and long-range dependencies, with lots of hits on the anti-diagonal belt.

With the limited amount of computing resources and computing time, the selection of hyper-parameters is critical to the model’s performance. A trade-off is necessary. The depth (transformer layers) of the model has a greater impact on the performance than the width (hidden size).

On the one hand, the model that is too flat ($\#transformer\_layers<8$) performs poorly, even though it has a large hidden size. We trained a model with $hidden\_size=8192$, $\#transformer\_layers=3$. Although its training speed (average time for each iteration) is the fastest among all models, it failed to converge after 5 days of training.

On the other hand, the model that is too deep is not a feasible choice in this scenario (limited training time). The training time of Figure 9 is 3.5 times that of Figure 10. And it takes about 25 days to train the model with 32 transformer layers Figure 9.

Our empirical conclusion is that the model parameters of $512<hidden\_size<3072$, $8<\#transformer\_layers<24$ are feasible and can well balance training efficiency and resource consumption.