We measure the impact of data amount on performance through three data sets (Table 1, SOM Fig. 9, which can be found on the Computer Society Digital Library at http://doi. ieeecomputersociety.org/10.1109/TPAMI.2021.3095381): Uniref50 [ 32 ], UniRef100 [ 32 ], and BFD (Big Fantastic Database) [ 24 ], [ 33 ]. Merging UniProt [ 23 ] and proteins translated from multiple metagenomic sequencing projects, BFD has become the largest collection of protein sequences: about eight times larger than the largest sets used previously for pLMs [ 34 ]. The 8-fold increase in data increased the number of tokens 5-fold (Table 1), as proteins were 1.6-fold longer in UniRef100. Without a clear mapping for LMs from NLP to proteins, i.e., the concept of words can be related to single amino acids, a window of amino acids (k-mer motifs [ 35 ]) or functional units (domains [ 36 ]), we decided to interpret single amino acids as input tokens/words. Thereby, protein databases contain several orders of magnitude more tokens than corpora used in NLP, e.g., Google’s Billion Word data set [ 37 ] top in NLP with about 829m (million) tokens (words), compared to BFD with 393b (billion) tokens (amino acids). Interpreting domains as words, would cut the number of tokens in BFD roughly by a factor of 100 (average domain length [ 38 ]) still leaving 5-times more tokens in BFD than the Billion Word corpus. Uniref50, UniRef100 and BFD were tokenized with a single space (indicating word-boundaries) between each token. Each protein sequence was stored on a separate line, with lines/proteins representing the equivalent of ”sentences”. Additionally, an empty line was inserted between each protein sequence in order to indicate the ”end of a document”; however, this is only essential for models with auxiliary task (Bert and Albert). Non-generic or unresolved amino acids ([BOUZ]) were mapped to unknown (X). For training ProtTXL and ProtT5, the data was transformed to pytorch and tensorflow tensors, respectively on the fly. For ProtBert, ProtAlbert, ProtXLNet and ProtElectra, the data was pre-processed and stored as tensorflow records. Given tensorflow records with terabytes, data sets had to be chunked into 6000 files for thousands of parallel workers.

We extracted the information captured by the protein LMs through embeddings, i.e., vector representations from the last hidden state of the protein LM (Fig. 1) and analyzed it by projecting the high-dimensional representations down to two dimensions (2D) using t-SNE [ 39 ]. Toward this end, we took annotations from several sources. First, a non-redundant (PIDE<40%) version of the SCOPe database [ 40 ] (release 2.07 with 14,323 proteins). Second, we mapped proteins into the three major domains of life (archaea, bacteria, or eukarya) or to viruses (removing all proteins with missing classifications). The number of iterations for the t-SNE projections was set to 3,000 and the perplexity to 30 for all plots with the exception of the amino acid plot for which we used a perplexity of 5. All visualizations used the same random seed (42).

We also assessed the information captured during selfsupervised pre-training of our protein LMs by using embeddings extracted from those models as sole input for supervised training. Although we mostly relied on previously published data sets to ease comparisons to other methods, for the supervised training, we also added a novel test set to refine the evaluation.

Per-Residue Prediction /Single Tokens. To predict properties of single tokens (here: single amino acids, dubbed residues when joined in proteins), we used the training set published with NetSurfP-2.0 [ 15 ] describing secondary structure in 3and 8-states (class distribution for all data sets in SOM Tables 4, 3, available online). We also included other public test data sets, namely CB513 [ 41 ]), TS115 [ 42 ], and CASP12 [ 43 ]. Each of those has severe limitations (CASP12: too small, CB513 and TS115 redundant and outdated). Therefore, we added a new test set using only proteins published after the release of NetSurfP-2.0 (after Jan 1, 2019). We included proteins from the PDB [ 44 ] with resolutions 2:5 A and 20 residues. MMSeqs2 [ 45 ] with highest sensitivity (-s 7.5) removed proteins with >20% PIDE to either the training set or to itself. On top, PISCES [ 46 ] removed any protein considered by its procedure to have >20% PIDE. These filters reduced the number of new proteins (chains) from 18k to 364 (dubbed set NEW364).

Per-Protein Prediction/Embedding Pooling. For the prediction of features for entire proteins (analogous to the classification of whole sentences in NLP), the DeepLoc [ 16 ] data set was used to classify proteins into (i) membrane-bound versus water-soluble and (ii) ten classes of subcellular localization (also referred to as cellular compartments).

We trained six successful LMs in NLP (T5 [ 47 ], Electra [ 48 ], BERT [ 49 ], Albert [ 50 ], Transformer-XL [ 51 ] and XLNet [ 11 ]) on protein sequences. BERT was the first bidirectional model in NLP which tried to reconstruct corrupted tokens, and is considered the de-facto standard for transfer learning in NLP. Albert reduced BERT’s complexity by hard parameter sharing between its attention layers which allows to increase the number of attention heads (64 chosen here). Electra tries to improve the sampling-efficiency of the pretraining task by training two networks, a generator and a discriminator. Instead of only reconstructing corrupted input tokens, the generator (BERT) reconstructs masked tokens, potentially creating plausible alternatives, and the discriminator (Electra) detects which tokens were masked. This enriches the training signal as the loss can be computed over all tokens instead of the subset of corrupted tokens (usually only 15 percent). T5 uses the original transformer architecture proposed for sequence translation, which consists of an encoder that projects a source language to an embedding space and a decoder that generates a translation to a target language based on the encoder’s embedding. Only later, models used either the encoder (BERT, Albert, Electra) or the decoder (TransformerXL, XLNet), but T5 showed that this simplification might come at a certain price as it reaches state-of-the-art results in multiple NLP benchmarks. Additionally, it provides the flexibility to apply different training methodologies and different masking strategies, e.g., T5 allows to reconstruct spans of tokens instead of single tokens.

As self-attention is a set-operation and thus order-independent, Transformers require explicit positional encoding. Models trained with sinusoidal position signal like BERT, Albert or Electra, can process only sequences shorter or equal to the length of the positional encoding which has to be set before training. Due to the huge memory requirement of Transformer-models, this parameter is usually set to a value lower than the longest proteins, e.g., Titin with 33k residues. Here, we trained models that were affected by this limitations (ProtBERT, ProtAlbert, ProtElectra) first on proteins of length 512, then on proteins 1024. Only setting the length of the positional encoding to 40k after pre-training allowed the models to process protein sequences up to a length of 40k. In contrast to this, TransformerXL introduced a memory that allows it to process sequences of arbitrary length. Still, the model cuts sequences into fragments but allows for flow of information between fragments for longer proteins by re-using hidden states of fragments which have already been processed. While its memory is uni-directional as fragments are processed sequentially, TransformerXL captures only uni-directional context within one memory fragment (auto-regressive) while XLNet, which uses a similar memory mechanism to process sequences of arbitrary length, allows to gather bidirectional context within one memory fragment.

In contrast to this, T5 learns a positional encoding for each attention head that is shared across all layers. This way, the model learned to combine the relative offset between residue pairs of lower layers, enabling the model to make predictions beyond the actual length of the positional encoding. No auxiliary tasks like BERT’s next-sentence prediction were used for any model described here.

ProtTXL, ProtBert, ProtXLNet, ProtAlbert and ProtElectra were trained on UniRef100, ProtT5 on UniRef50, and ProtTXL, ProtBert & ProtT5, on BFD (Table 2). Largely, we transferred configurations successfully from NLP to protein sequences [ 52 ], [ 53 ], [ 54 ], with the exception of the number of layers that was increased to optimize memory utilization.

ProtTXL. The Transformer-XL1 was trained using both UniRef100 and BFD-100 datasets (referred to as ProtTXL and ProtTXL-BFD, respectively; Table 2). Both models used a dropout rate of 15 percent, a memory length of 512 tokens and using mixed precision. The number of layers, number of heads, batch size, learning rate, weight decay, training steps and warm-up steps were adjusted according to training set size as well as GPU utilization. The number of warm-up steps was set to cover at least one epoch for each data set. We tested initial learning rates between 0.001 and 0.005 which were increased linearly at every training step over the warmup period. To avoid model divergence during training, the 1. https://github.com/NVIDIA/DeepLearningExamples/ learning rate had to be (i) reduced along with the warm-up steps (for BFD), or (ii) increased for both (for Uniref100). Even after increasing the warm-up steps to two epochs, the maximum learning rate remained at 0.0025 for both data sets. Beyond this point, the training diverged. Using weight decay to regularize the network increased the GPU memory usage as it required to compute the norm of all weight vectors on our models, thus reducing the batch size. ProtTXLBFD was trained for 40k steps in total, with 13.6k warm-up steps using a learning rate of 0.0005, while ProtTXL was trained for 31k steps with 5k warm-up steps using a learning rate of 0.002. The Lamb optimizer was able to handle the resulting batch sizes of 44k and 22k for ProtTXL-BFD and ProtTXL, respectively, without divergence.

ProtBert. BERT2 was trained using both UniRef100 and BFD-100 datasets (referred to as ProtBert and ProtBert-BFD, respectively; Table 2). Compared to the original BERT publication, the number of layers was increased. Unlike Transformer-XL which was trained on Nvidia GPUs, mixedprecision was not used to train other models because those were trained on TPUs. Similar to the BERT version trained in the Lamb paper [ 55 ], ProtBert was first trained for 300k steps on sequences with a maximum length of 512 and then for another 100k steps on sequences with a length of a maximum length of 2k. While ProtBert-BFD was trained for 800k steps, then for another 200k steps for sequences with maximum length of 512 and 2k, respectively. This allows the model to first extract useful features from shorter sequences while using a larger batch size, rendering training on longer sequences more efficient.

ProtAlbert. We trained Albert3 on UniRef100 (ProtAlbert; Table 2). We used the configuration from the official GitHub repository for Albert (version: xxlarge v2). For Albert the number of layers is increased through the number of times, Albert stacks its single layer. Compared to the original 2. https://github.com/google-research/bert 3. https://github.com/google-research/albert publication, we achieved increasing the global batch size from 4096 to 10752 on the same hardware. The reason for this counter-intuitive effect is the reduced vocabulary size in proteins: the entire diversity of the protein universe is realized by 20 different amino acids, compared to tens of thousands of different words. Similar to ProtBert, ProtAlbert was first trained for 150k steps on sequences with a maximum length of 512 and then for another 150k steps on sequences with a maximum length of 2k.

ProtXLNet. XLNet4 was trained on UniRef100 (ProtXLNet) using the original NLP configuration [ 11 ] (Table 2) except for the number of layers that was increased to 30 layers which reduced the global batch size to 1024. Due to the relatively small batch-size, we used the original optimizer: Adam with a learning rate of 0.00001. The model was trained through more steps, i.e., 20k warm-up and 847k steps to compensate for the smaller batch-size of this model.

ProtElectra. Electra5 consists of two models, a generator and discriminator (same number of layers, generator 25 percent of the discriminator’s hidden layer size, hidden layer intermediate size, and number of heads). We copied Electra’s NLP configuration with two changes: increasing the number of layers to 30 and using Lamb optimizer. Again, we split the training into two phases: the first for proteins 512 residues (400k steps at 9k global batch size), the second for proteins 1024 (400k steps at 3.5k global batch size). While ProtTXL, ProtBert, ProtAlbert and ProtXLNet relied on pre-computed tensorflow records as input, Electra allowed to mask sequences on the fly, allowing the model to see different masking patterns during each epoch.

ProtT5. Unlike the previous LMs, T56 uses an encoder and decoder [ 10 ]. We trained two model sizes, one with 3B (T5-XL) and one with 11B parameters (T5-XXL). T5-XL was trained using 8-way model parallelism, while T5-XXL was trained using 32-way model parallelism. First, T5-XL and T5-XXL were trained on BFD for 1.2M and 920k steps respectively (ProtT5-XL-BFD, ProtT5-XXL-BFD). In a second step, ProtT5-XL-BFD and ProtT5-XXL-BFD were fine-tuned on UniRef50 for 991k and 343k steps respectively (ProtT5XL-U50, ProtT5-XXL-U50). Contrary to the original T5 model which masks spans of multiple tokens, we adopted BERT’s denoising objective to corrupt and reconstruct single tokens using a masking probability of 15 percent. All T5 models used the AdaFactor optimizer with inverse square root learning rate schedule for pre-training. Like ProtElectra, T5 masks each sequence on the fly. In our hands, the encoder outperformed the decoder on all benchmarks significantly and running the model in half-precision during inference instead of full-precision had no effect on performance but allowed to run the model on on a single Nvidia TitanV (12GB vRAM). Thus, we dropped the decoder from further analysis which cuts model size by half during inference. For completeness, we made weights for encoder and decoder publicly available.

4. https://github.com/zihangdai/xlnet 5. https://github.com/google-research/electra 6. https://github.com/google-research/text-to-text-transfertransformer

In the transfer-learning step, embeddings from our pretrained protein LMs served as sole input to subsequent supervised training, thereby transferring knowledge acquired during self-supervised pre-training to other tasks. To best analyze the impact of transfer learning, we deliberately kept the supervised models using the embeddings from the protein LMs as input minimal. In particular, compared to SOA solutions such as NetSurfP-2.0, all our experiments used the pre-trained LMs as feature extractors without finetuning, i.e., without gradient back-propagating to the LMs. Throughout, we extracted the embeddings from the last hidden state of the pre-trained LMs as described in detail elsewhere [ 56 ]. To briefly summarize (Fig. 1): we applied tasks on two different levels, namely on the level of single tokens (perresidue) and whole sentences through pooling (per-protein) predictions. For the per-residue prediction, we input the embeddings into a two-layer convolutional neural network (CNN). The first CCN layer compressed the embeddings to 32 dimensions using a window size of 7. The compressed representation was fed into two different CNNs (each with window size 7). One learned to predict secondary structure in 3-states, the other in 8-states. The network was trained on both outputs simultaneously by adding their losses (multitask learning). For ProtBERT-BFD embeddings we additionally trained three other models: logistic regression, FNN and LSTM. Similar to the CNN, the two-layer FNN first compressed the output of the language model down to 32 dimensions which the second FNN-layer used to predict 3- and 8states simultaneously. The bi-directional LSTM compressed the embeddings down to 16 dimensions. Concatenating both directions, the resulting 32 dimensional representation was used by a FNN layer to predict 3- or 8-states. As the CNN performed best (SOM Table 8, available online), we used CNNs throughout. For the per-protein prediction, we also extracted the embeddings from the last layer of the protein LMs. However, then we pooled the representations over the lengthdimension resulting in a fixed-size representation for all proteins. Using ProtBERT-BFD embeddings, we compared alternative pooling strategies (SOM Table 8, available online) and chose mean-pooling for all further experiments. The resulting vector was used as an input to a single feed forward layer with 32 neurons which compressed information before making the final predictions for both per-protein tasks, i.e., the prediction of subcellular localization and the differentiation between membrane-bound and water-soluble proteins, simultaneously (multi-task learning).

HPC hardware is advancing both through infrastructure of supercomputers, such as Fugaku [ 57 ], Summit [ 1 ] or the SuperMUC-NG [ 58 ], and through its components, such as TPU pods [ 2 ], specifically designed to ease large scale neural network training for users. Concurrent software improvements in form of more efficient libraries such as Horovod [ 6 ] allow executing general purpose code on large distributed clusters with minor code changes. In this section we give details on the hardware used for training language models on large protein sequence databases.

ORNL Summit & Rhea. The Oak Ridge National Laboratory (ORNL) provides several clusters for researchers who need computational resources not provided by research facilities such as universities. Here, we used Summit and Rhea. Summit was used to train the deep learning models, while Rhea was used for the pre-processing of data sets including the distributed generation of tensorflow records.

Summit is the world’s second fastest computer, consisting of approximately 4618 nodes. Each node has two IBM POWER9 processors and six NVIDIA Volta V100 with 16GB of memory each [ 1 ]. Every POWER9 processor is connected via dual NVLINK bricks, each capable of a 25GB/s transfer rate in both directions. A single node has 0.5 TB of DDR4 main memory and 1.6TB of non-volatile memory that can be used as a burst buffer. Summit is divided into racks with each rack having 18 nodes. In all of our experiments we reserved 936 nodes for training. As having nodes on the same rack decreases the communication overhead, we reserved entire racks.

The smaller cluster (Rhea) contains two partitions: Rhea and GPU. The Rhea partition has 512 node, each with 128 GB of memory and two Intel Xeon E5-2650. The GPU partition has only 9 nodes, each with 1 TB of memory and two Intel Xeon E5-2695. Reha reduced the time needed for creating tensorflow records for the BFD dataset from 7.5 months (!) to fewer than two days, by converting the original sequential script to distributed processing using MPI. The generation script used two nodes of the GPU partition, with a total of 112 parallel threads.

Google TPU Pod. In 2016, Google introduced tensor processing unit (TPU) as its application-specific integrated circuit optimized for training neural networks. TPUs can be accessed through Google Cloud. Training the protein LMs used both older TPU generation (V2) with 256 cores, and the latest TPU generation (V3) with 512 and 1024 cores. These cores are divided into hosts with each host having access to 8 cores. Consequently, we had access to 32, 64 and 128 hosts for V2/V3-256, V3-512 and V3-1024, and each core had 8 GiB and 16 GiB of high-bandwidth memory for V2 and V3. Training on the TPUs required access to a virtual machine on Google Cloud and storage on Google Bucket [ 59 ].

Summit integrates several pre-configured modules which include the most popular libraries and tools required for simulation, deep learning, distributed training and other purposes. We used the IBM Watson Machine Learning module versions 1.6.0 and 1.6.2 for our deep learning training. In contrast to this, the Google Cloud server, which we used for the TPU Pod training, had to be configured manually because only the operating system was installed.

Pytorch was used to train ProtTXL, tensorflow to train ProtBert, ProtAlbert, ProtXLNet, ProtElectra and ProtT5.

Both libraries used the Horovod framework [ 6 ] to train the models on distributed clusters such as Summit. Horovod supports distributed GPU training with minimal change in the code. It supports different backends including MPI, NCCL and IBM PowerAI distributed deep learning (DDL).

We tested all three backends and found DDL to be the fastest for our training purpose on Summit. The time needed to finish a single batch with ProtTXL-BFD increased from one to two nodes due to the communication overhead (Fig. 2). After two nodes the communication overhead plateaued, even when scaling up to 936 nodes with 5616 GPUs. Summit has integrated DDL in their Watson Machine Learning module which comes with most DDL libraries including pytorch, tensorflow, apex, DDL and horovod. However, Summit has only a license for using DDL up to 954 nodes. Contrary to Summit, training on TPU Pods did not require any changes in the Tensorflow code to use either a single TPU host or to distribute workload among multiple TPU hosts.

Mixed precision allows to fit larger models and batch sizes into GPU memory by using 16-bit precision only or a mix of 16-bit and 32-bit precision. Nvidia’s APEX library [ 60 ] was used for mixed precision training of ProtTXL, because APEX supports pytorch. As ProtTXL training became instable when training with 16 Bit precision, we switched to almost half precision training (more details in SOM - 1.2 Software, available online). We did not use mixed-precision for models trained on TPUs.

Another optimization technique/library crucial for our training on Summit was IBM’s large model support (LMS) [ 61 ]. Similar to gradient checkpointing [ 62 ], LMS virtually extends the GPU memory by outsourcing parts of the model from GPU to main memory. This allows training models larger than the GPU memory. The obvious drawback of LMS is the increase in training time due to shuttling data between CPU and GPU and back. However, the reduced memory consumption of the model allows to increase the batch size, potentially compensating for the communication overhead. ProtTXL was used to evaluate the effect of Pytorch’s implementation of LMS while ProtBert was trained for a few steps BFD using Summit to evaluate tensorflow’s implementation of LMS. Training ProtBert for a few steps was sufficient to assess the effect of LMS on batch-size, model size as well as an estimate of training time. In the end, we used LMS only for ProtTXL to strike a balance between model size and training time. The number of LM parameters could be increased by about 15.6 percent for ProtTXL-BFD and to 6.6 percent for ProtBert (SOM Fig. 10a, available online). Additionally, we could increase the batch size by 700 percent for ProtTXL-BFD (SOM Figs. 10b Fig. 3. Protein LMs learned constraints. t-SNE projections visualized information extracted by the unsupervised protein LMs (here best-performing ProtT5-XL-U50; upper row: before training (Random), and lower row: after pre-training on BFD & UniRef50. (A) The left-most column highlights single amino acids by biophysical features. (B) The middle column annotates protein structural class (taken from SCOPe). (C) The right-most column distinguishes proteins according to the kingdom of life in which it is native. Although the random projections on top may suggest some adequacy of the clusters, the trained models shown on the lower row clearly stood out. Incidentally, the comparison of the two also highlighted the potential pitfalls of using t-SNE projections from many dimensions onto 2D: although random, the human might see some correct patterns even in the top row. Most impressive might be the fine-grained distinctions of biophysical features of amino acids (A), however, more surprising are the classifications of entire proteins according to structural class (B) and organism (C). For these, we created embeddings through global average pooling over the representations extracted from the last layer of ProtT5-U50 (average over protein length, i.e., per-protein embeddings; Fig. 1). and 10c, available online). The NV-Link between CPU and GPU on Summit-nodes, reduced the training time for ProtTXL by 60 percent while it increased by 72 percent for ProtBert (SOM Fig. 10d, available online). 3

Embeddings extract constraints about protein function and structure learned by the protein LMs during self-supervised pre-training on raw (unlabeled) protein sequences. t-SNE plots [ 39 ] visualized this information by projecting the embeddings onto 2D and by overlaying annotations (labels) of structural, functional or evolutionary features. Using attention maps, we analyzed the DNA-binding zinc-finger motif well conserved in evolution in more detail.

Capturing Biophysical Features of Amino Acids. Applying tSNE to the uncontextualized token embedding layer visualized information extracted by the pLMs for individual amino acids independent of their context (residues next to it). As previously established for another pLM [ 53 ], the tSNE projections (e.g., ProtT5-XL-U50 SOM Fig. 14a, available online, or ProtBert-BFD SOM Fig. 15a, available online) suggested that all pLMs captured essential biophysical amino acid features, including charge, polarity, size, hydrophobicity, even to the level of aliphatic ([AILMV]) versus aromatic ([WFY]).

We compared the embedding projection with a randomly initialized model of identical architecture to ascertain that the observed effects did not originate from coincidental signals originating from projecting highdimensional data (Fig. 3 A) or some inductive bias of neural networks [ 63 ]. The random projection clearly did not carry biophysical information, while the embeddings projection did.

Capturing Protein Structure Classes. We averaged over the length-dimension of the representations derived from the last layer of each pLM (Fig. 1) to derive fixed size representations for each protein and superposed structural class from SCOPe [ 40 ]. ProtBert-BFD and especially ProtT5-XL-U50 embeddings visually separated the proteins best (SOM Figs. 14-19, available online). Although sequence length was not explicitly encoded and our pooling squeezed proteins into a fixed vector size, all pLMs separated small from long proteins (brown, e.g., ProtT5-XL-U50 SOM Fig. 14D, available online). All models also distinguished between water-soluble and transmembrane proteins (light blue, e.g., ProtT5-XLU50 SOM Fig. 14D, available online) and, to some extent, between proteins according to their secondary structure composition(e.g., all-alpha (dark blue) versus all-beta (dark green) ProtT5-XL-U50 Fig. 14D, available online). While having much higher entropy, even the random model clustered small proteins from long proteins (brown, Fig. 3B).

Capturing Domains of Life and Viruses. The analysis distinguished three domains of life: archaea, bacteria, and eukarya, along with viruses - typically not considered as life. We used the same proteins and per-protein pooling as above for SCOPe. All pLMs captured some organism-specific aspects (e.g., ProtT5-XL-U50 SOM Fig. 14E, available online). Eukarya and bacteria clustered better than viruses and archaea. Comparing different pLMs revealed the same trend as for protein structure classes: ProtTXL (SOM 19E, available online) and ProtBert (SOM 16E, available online) produced higher entropy clusters, while ProtAlbert (SOM 17E, available online), ProtXLNet (SOM 18E, available online), ProtBERT-BFD (SOM Fig. 15E, available online) and ProtT5-XL-U50 (SOM Fig. 14E, available online) produced visually easier separable clusters.

Capturing Protein Function in Conserved Motifs. A similar overall per-protein analysis as for structural classes and domains of life also suggested some clustering according to protein function as proxied by enzymatic activity (EC-numbers [ 64 ] and subcellular location (SOM -1.3 pLMs unsupervised, available online). We focused in more detail on the attention mechanism [ 65 ] at the core of each Transformer model [ 10 ] providing some limited understanding [ 66 ], [ 67 ]. We visualized [ 68 ] the attention weights of ProtAlbert to analyze the structural motif of a zinc-finger binding domain (SOM Fig. 11, available online) crucial for DNA- and RNAbinding and conserved across diverse organisms. Some of the attention heads of ProtAlbert (SOM Fig. 11, available online; line thickness resembles attention weight) focused mostly on the four residues involved in zinc-binding (residues highlighted in the left part of SOM Fig. 11, available online) which is essential for function.

The acid test to prove that pLM embeddings extracted important constraints is to exclusively use embeddings as input to supervised training. We proxied this through predictions on two different levels, namely on the per-residue or token level (secondary structure) and on the per-protein or sentence level through pooling over entire proteins (location, and classification into membrane/non-membrane proteins). The pLMs remained unchanged, i.e., both approaches (per-residue/per-protein) used only embeddings derived from the hidden state of the last attention layer of each pLM (Fig. 1), i.e., pLMs were only used as static feature extractors. 3.2.1

We compared the time needed to generate representations for EI/MSA-based prediction methods and pLMs by generating MSAs and embeddings for each protein in the human proteome (20,353 proteins with median length of 415). We used the fastest method available, namely MMseqs2 [ 45 ], Fig. 5. Effect of MSA size. We used our new test set (NEW364) to analyze the effect of the size of an MSA upon secondary structure prediction (Q3) for the two top methods (both reaching Q3=84.3%): NetSurfP-2.0 (using MSA) and ProtT5-XL-U50 (not using MSA). As proxy for MSA size served Neff, the number of effective sequences [ 74 ] (clustered at 62 percent PIDE): leftmost bars: MSAs with Neff=1, middle: Neff& 10, right: Neff>10. As expected ProtT5-XL-U50 tended to reach higher levels than NetSurfP-2.0 for smaller families. Less expected was the almost on par performance for larger families. with parameters established by NetSurfP-2.0 to generate MSAs (SOM for more details, available online). MMseqs2 was about 16 to 28-times slower than the fastest pLMs (ProtElectra and ProtBert), and about 4 to 6-times slower than our best model (ProtT5-XL; Fig. 8. ProtT5-XL, required on average 0.12 seconds to generate embeddings for a human protein, completing the entire human proteome (all proteins in an organism) in only 40 minutes.

We also investigated the cross-effect of sequence length and batch-size (SOM Table 10, available online) on the inference speed of different pLMs. When using a single Nvidia Fig. 6. Number of pre-training correlated with performance. We plotted 3state secondary structure prediction performance (Q3) on NEW364 for all pLMs trained here against the number of samples seen during pretraining (training steps times global batch-size in K). For simplicity, we dropped the ” Prot” prefix from all models. Despite the vastly different pLMs, the high Spearman’s r of 0.62 indicated a common trend: more pre-training samples: better prediction.

Fig. 7. Per-protein (sentence-level) performance: The prediction of location in 10 states (lower bars in cyan: Q10: percentage of proteins with 1 of 10 classes correctly predicted) and the classification of membrane/ other (higher bars in magenta: Q2: percentage of proteins correctly classified in either of two classes). Embeddings were derived from pLMs by mean-pooling, i.e., averaging over the length of the entire protein (Fig. 1). Abbreviations as in Table 4 except for one method using neither pLMs nor MSA (no LM no MSA: DeepLoc-BLOSUM62 [ 16 ]), and two methods using MSAs (MSA input evolutionary information): the current state-of-the-art (SOA) method (performance marked by horizontal thin lines in magenta and cyan) DeepLoc [ 16 ], and LocTree2 [ 77 ]. Almost all pLMs outperformed LocTree2 and a version of DeepLoc not using MSAs (DeepLoc-BLOSUM62). Only, ProtT5-XXL-U50 and ProtT5-XL-U50 outperformed the SOA. A recent method optimized location prediction from ProtT5 embeddings through a light-attention mechanism; it clearly outperformed the SOA without using MSAs (LA_ProtT5 & LA_ProtT5-U50 [ 78 ]). The top row shows the complete range from 0-100, while the lower row zooms into the range of observed differences.

Quadro RTX 8000 with half precision on varying batch-sizes (1,16,32) as well as sequence lengths (128, 256, 512), ProtBert and ProtElectra provided the fastest inference with an average of 0.007 seconds per protein when using a batch size of 32, followed by ProtT5-XL and ProtAlbert (0.025s). The batch-size of most pLMs could have been increased on the same hardware but was limited to allow a direct comparison between all models, due to large memory requirements for ProtT5-XXL. The script for running this benchmark is freely available as part of our github repository. 4 4.1

HPC Supercomputers such as Summit [ 1 ] and Google’s cloud TPU Pod [ 2 ], combined with optimized libraries such as IBM’s DDL [ 7 ] and Horovod [ 6 ] set the stage for training LMs with billions of free parameters on terabytes of data in hours or days. Increasing model size improves performance for some NLP applications [ 12 ], although the massive data challenges the communication between thousands of nodes and divergence between large batches during training. Here, we presented some solutions to overcome these challenges for training protein LMs (pLMs) by fully utilizing 20 percent of Summit for TransformerXL [ 51 ], one TPU Pod V3-512 for Bert [ 49 ], Electra [ 48 ], Albert [ 50 ] and XLNet [ 11 ], and a mix of TPU Pod V3-256, V3-512, and V3-1024 for ProtT5-XL and ProtT5-XXL [ 47 ]. This implied the parallel use of 5616 GPUs on Summit or 256/512/1024 TPU cores on a TPU Pod, while avoiding training divergence with specialized optimizers such as LAMB [ 55 ] up to a global batch size of 44K samples (here: protein sequences).

Mostly, we proxied the relevance of information extracted by pre-trained pLMs through the performance of the subsequent supervised tasks trained exclusively on embeddings from pre-trained pLMs. Using pLMs only as static feature extractors, we considered the pLM to improve (or to be better) when the supervised task using embeddings from this pLM as input reached higher performance.

BFD - Biggest Data to Pre-Train. We trained our pLMs on the largest protein database ever used for this purpose, namely BFD [ 33 ], more than an order of magnitude larger than the standard UniProt [ 23 ]. Although bigger did not equate better supervised performance, several pLMs improved through pre-training on more data (UniRef100 versus BFD, Table 1). Nevertheless, the performance increase appeared small given the 10-fold larger data set (e.g., DQ3(ProtBert-BFD/UniProt)=1.3%). In contrast, the pre-training protocol by which we first trained on the larger but more noisy (more mistakes in sequences) and redundant BFD and then continued pre-training using the smaller, less redundant UniRef50, improved performance significantly for both ProtT5 versions (DQ3(ProtT5-XLBFD/U50)=2.8% and DQ3(ProtT5-XXL-BFD/U50)=1.4%). The improvement through refined pre-training for ProtT5XL (3B parameters) exceeded that for ProtT5-XXL (11B parameters), presumably, because it saw more samples when continuing pre-training for a similar amount of time (limited by resources).

This highlighted a remarkable trend common across the wide diversity of pLMs and data sets: the performance of supervised downstream tasks using embeddings from pretrained pLMs increased with the number of samples presented during pre-training (Fig. 6; Spearman’s r=0.62). We could not observe a similarly consistent trend for model size. However, this might be attributed to some trade-off: training for more steps might require models with sufficient capacity to absorb the information of the corpus. For instance, while ProtBERT-BFD (420M parameters) saw around 27B proteins during pre-training, it fell short compared to ProtT5-XL-BFD (3B parameters) which saw only around 5B proteins (Figs. 4, SOM Tables 6, 7, available online). This finding aligned with results from NLP suggesting that larger models absorb information faster and need less training to achieve similar performance [ 79 ]. However, the comparison between, e.g., ProtT5-XL and ProtT5-XXL suggested a possible cap to this trend, as larger models see fewer samples in the same amount of computing power. The clear correlation between performance and samples combined with the need for sufficient model size spotlighted the crucial role of substantial computing resources (HPC, TPUs, and GPUs): big data needs large models need loads of computational resources.

Some rudimentary information about how proteins are formed, shaped, and function has been learned by the pLMs during pre-training: all our pLMs (ProtT5, ProtBert, ProtElectra, ProtAlbert, ProtTXL, ProtXLNet) extracted valuable information as revealed by visualizing embeddings without further supervised training on labeled data. The comparisons to randomly initialized pLMs highlighted two other aspects. First, how easy it is to incorrectly imagine patterns when projecting from high-dimensional spaces into 2D: although the random pLMs contained no information, some annotations might have suggested the opposite (Fig. 3 top row untrained). Second, the pLMs extracted important global constraints about protein structure and function (lower row in Fig. 3). These span from the most local (individual token level) biophysical features of amino acid building blocks (e.g., hydrophobicity, charge, and size, Fig. 3 A), over global classifications into structural classes (Fig. 3 B), to the macroscopic domains of life (Fig. 3 C). Global structural (e.g., overall secondary structure content, Fig. 3 B) and biochemical (e.g., transmembrane, SOM Fig. 14B, available online) properties appeared most distinctive. In contrast, local features relying on short motifs were less separated (EC-numbers: Fig. 14F, location: Fig. 14C, available online) but still clustered, e.g., for secreted/extracellular proteins or hydrolases.

On a more fine-grained level, the visual analysis of the attention mechanism at the core of each Transformer, confirmed the pLMs to have even picked up more subtle signals of short functional motifs. Specifically, a few attention heads of ProtAlbert zoomed into the four residues most important to coordinate zinc-binding (SOM Fig. 11, available online). Although limited in scope [ 66 ], such an analysis provided some explanation about the inner workings of Transformers not needing large sets of experimental annotations (labels). On top, the resulting interpretations of the AI/ML might be less biased than experimental annotations. For instance, databases with annotations of protein function such as SwissProt [ 80 ] and of protein structure such as PDB [ 44 ] are extremely biased by today’s experiments [ 75 ], [ 81 ], [ 82 ].

The t-SNE and UMAP analyses revealed the pLMs to have extracted some understanding of the language of life. As prediction is the acid test for understanding, we extracted the pLM embeddings as input to predicting aspects of protein structure. Overall, the results confirmed [ 56 ] that evolutionary information (EI, i.e., methods using multiple sequence alignments MSAs) significantly outperformed most pLMs not using such information, except for ProtT5-XL (on all per-residue and per-protein tasks, Figs. 4, 7 and SOM Tables 7, 6, 9, available online). ProtT5-XL eliminated this gap from embeddings-only input: on some tasks/data sets, it outperformed the current state-of-the-art (SOA) MSA-based method, on others it remained slightly behind. Newer pLMs using context improved over both previous pLMbased approaches [ 56 ] (8-9 percentage points in Q3), other transformers [ 53 ] (2-4 percentage points in Q3), and over non-contextualized word2vec-type approaches [ 83 ], [ 84 ] (18-22 percentage points in Q3).

The performance differences between two data sets (CASP12 and NEW364) highlighted another problem. While it is clear that we need to reduce redundancy within the evaluation set and between it and all development sets, it is less clear how to exactly do this. In focusing on CASP12 and NEW364, we approached two different assumptions. CASP12 mostly measures how well predictions will be for proteins with unseen structures. A comprehensive rigorous realization of data sets following this perspective has recently been published [ 85 ]. NEW364, on the other hand, builds on the assumption that the maximal redundancy is defined by sequence similar to protein in the PDB. In this sense, we interpreted results for CASP12 as a lower and those for NEW364 as an upper limit. The NEW364 comparisons also highlighted the importance to constantly create up-to-date data sets with enough non-redundant proteins never used for development by any of the methods assessed.

pLMs so Powerful to Render Simple Baseline Effective. While pre-training pLMs is computationally costly, training supervised models using embeddings as input requires much less energy. For instance, the logistic regression trained on top of ProtBERT-BFD was already competitive with substantially more complex CNNs or LSTMs in predicting secondary structure (SOM Table 5, available online). In another example, a parameter-free nearest neighbor lookup using distances from pLM embeddings sufficed to outperform homology-based inference for predicting protein function [ 86 ]. This suggested that pLMs are particularly suitable when so few experimental annotations (labels) are available that the complexity of modern deep learning becomes prohibitive. In fact, all supervised solutions presented here that reached the SOA were much less complex (fewer free parameters) than the EI-based methods they reached. We focused on the development of pLMs using performance on supervised tasks to rank without optimizing particular supervised solutions. Others have already begun surmounting EI-based methods by custom-designing pLM-based solutions [ 78 ] possibly even through end-to-end systems [ 30 ], [ 87 ]. Such solutions can even compete when MSAs are absolutely essential, e.g., for predicting effects of single amino acid variants (SAVs) [ 88 ].

Bi-Directionality Crucial for pLMs. In NLP uni-directional (auto-regressive) and bi-directional (auto-encoding) models perform on par [ 12 ], [ 89 ]. In contrast, the bi-directional context appeared crucial to model aspects of the language of life. While auto-encoding models such as Albert [ 50 ] utilize context to both sides during loss calculation, auto-regressive models such as TransformerXL [ 51 ] consider only context to one side. Performance increased substantially from unidirectional ProtTXL to bi-directional ProtXLNet (Fig. 4, SOM Tables 6, 7, available online). This might be compensated for by first pre-training on sequences and their reverse and then concatenating the output of uni-directional LMs applied on both directions. While this does not allow the LM to use bidirectional context during training, it allows supervised networks to combine context derived independently from both sides. For instance, ELMo [ 8 ] concatenates the embeddings derived from a forward and a backward LSTM. The protein LM version of ELMo (SeqVec) outperformed the uni-directional ProtTXL but not the bi-directional ProtXLNet. The difference in model size (SeqVec = 93M versus ProtXLNet = 409M) and in pre-training data (SeqVec = 30M versus ProtAlbert=224M) might explain some of this effect. Nevertheless, pure uni-directionality as used in TransformerXL appeared detrimental for modeling protein sequences.

Short answer: clearly not. For instance, since first submitting this work (July 2020), the mark has been pushed substantially higher. More generally, adopting techniques from NLP to proteins opens new opportunities to extract information from proteins in a self-supervised, data-driven way. New protein representations may complement existing solutions, most successful when combining evolutionary information7 and machine learning [ 21 ], [ 22 ], [ 31 ], [ 90 ]. Here we showed for the first time that pLM embeddings input to relatively simple supervised learning models can reach similar levels of performance without using EI and without optimizing the supervised training pipeline much. However, the gain in inference speed for pLMs compared to traditional solutions using EI/MSAs is so significant that largescale predictions become, for the first time since 30 years, feasible on commodity hardware. For instance, the best-performing model ProtT5-XL-U50 can run on a Nvidia TitanV with 12GB vRAM. Nevertheless, given the pLMs described here and elsewhere [ 34 ], [ 52 ], [ 53 ], [ 56 ], [ 91 ], [ 92 ], [ 93 ], we might expect an upper limit for what pLMs can learn 7. Throughout this work, we used evolutionary information (EI) as synonymous for using multiple sequence alignments (MSAs). Whether pLMs do not implicitly extract EI will have to be proven in separate publications. through masked language modeling (or auto-regressive pre-training). This work could establish three findings. (1) Less noisy and less redundant corpora (e.g., UniRef50) improved over larger but more noisy and redundant corpora (e.g., BFD). (2) In our perspective of limited resources, it was most important to use the resources for long-enough training because the number of samples seen during pretraining correlated with the prediction performance of downstream tasks. Ultimately, this seemed to originate from a trade-off between sufficient model size and sample throughput. (3) The bi-directional outperformed the unidirectional models tested. However, given the advances of protein LMs over the course of the reviewing of this work, we have seen no evidence for having reached a limit for pLMs, yet.

Many Open Questions. Answers to the following questions might advance the status-quo. (1) Would the addition of auxiliary tasks such as next-sentence or sentence-order prediction offered by BERT or Albert suit protein sequences? A suggestion might be incorporating structure information [ 94 ] or evolutionary relations [ 92 ], [ 95 ]. (2) Might the efficiency of training transformer pLMs improve through sparse transformers [ 96 ] or attention optimized with locality-sensitive hashing (LSH) [ 97 ] as introduced recently by the Reformer model [ 98 ] or more recent work of linear Transformers [ 99 ]? (3) Which data set, pre-processing, redundancy reduction and training batch sampling should be used to improve? (4) How much will it improve to tailor the supervised training pipeline to particular tasks? We treated secondary structure or location prediction more as proxies to showcase the success of protein LMs than as an independent end. (5) Will the combination of EI and AI [ 95 ] bring the best protein predictions of the future, or will the advantages of single-protein predictions (speed, precision) win out? In fact, single-protein predictions also have the advantage of being more precise in that they do not provide some implicit average over a protein family.

Overall, our results established that the combination of HPC solutions for training protein LMs and subsequent training of supervised prediction methods scaled up to the largest data sets ever used in the field. Only the combination of these different domains conclusively demonstrated that pLMs reach up to or even above the performance of the best methods in the field combining EI and AI without ever exploiting multiple sequence alignments. 5