SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Timothy F. Truong Jr
OpenProtein.AI
NY, USA
ttruong@openprotein.ai
Tristan Bepler
tbepler@openprotein.ai
Generative protein language models are a natural way to design new proteins with desired functions. However, current models are either difficult to direct to produce a protein from a specific family of interest, or must be trained on a large multiple sequence alignment (MSA) from the specific family of interest, making them unable to benefit from transfer learning across families. To address this, we propose Protein Evolutionary Transformer (PoET), an autoregressive generative model of whole protein families that learns to generate sets of related proteins as sequences-of-sequences across tens of millions of natural protein sequence clusters. PoET can be used as a retrieval-augmented language model to generate and score arbitrary modifications conditioned on any protein family of interest, and can extrapolate from short context lengths to generalize well even for small families. This is enabled by a unique Transformer layer; we model tokens sequentially within sequences while attending between sequences order invariantly, allowing PoET to scale to context lengths beyond those used during training. In extensive experiments on deep mutational scanning datasets, we show that PoET outperforms existing protein language models and evolutionary sequence models for variant function prediction across proteins of all MSA depths. We also demonstrate PoET's ability to controllably generate new protein sequences.
Proteins carry out most of the biological functions at the molecular level of life. The function of a protein is encoded by its specific amino acid sequence and the three-dimensional structure that the sequence folds into. Engineering proteins for novel and enhanced function is a key problem in pharmaceuticals and biotechnology and involves designing novel sequences or modifying existing natural proteins for these purposes. Deep mutational scans and directed evolution experiments have been used to successfully design novel proteins [1-3], but can be costly and difficult to implement, which makes these experimental methods inapplicable for many proteins and functions of interest. Accurate computational models of sequence-function relationships can narrow down the protein sequence search space, reduce the need for expensive experiments, and enable the design of more novel proteins and functions.
Protein language models have emerged as promising methods for understanding and designing protein sequences [4-7]. In particular, generative models offer a natural way to produce new protein designs. By training on large corpuses of natural protein sequences, these models learn evolutionary constraints on sequence space. They can then be used either to generate realistic sequences directly by sampling [8, 9], or to identify promising protein sequence variants by predicting the relative fitness of the variants of interest using the sequence likelihoods as a proxy [4, 10, 11].
$^{1}$ Code and pre-trained model weights are available at https://github.com/OpenProteinAI/PoET.
Traditionally, family-specific models learn evolutionary constraints specific to the protein family of interest by training on a multiple sequence alignment (MSA) of homologous sequences [12, 13]. However, this is ineffective for protein families with few sequences due to the lack of sufficient training data and inability to exploit information across families. These models also assume that MSAs are accurate, and cannot model novel insertions or deletions (indels) not present in the training MSA. More recently, unconditional large protein language models [5, 10, 11] have been developed. Trained on all known natural protein sequences, unconditional protein language models generalize across protein families. However, these unconditional single-sequence models cannot be easily directed to generate a protein from a specific family of interest, and underperform family-specific models for relative fitness prediction. Hybrid models such as Tranception [11] and TranceptEVE [14] combine unconditional language models with family-specific models to enable specialization to protein families. Nonetheless, it is unclear how to use these models to generate sequences with novel indels, and predictions from the family-specific models do not directly benefit from transfer learning across families.
Here, we propose the Protein Evolutionary Transformer (PoET), an autoregressive generative model of whole protein families that addresses these limitations. By learning to generate sets of related proteins as sequences-of-sequences across tens of millions of natural protein sequence clusters, PoET is able to generalize about evolutionary processes across protein families, and avoids issues related to conditioning on MSAs. In order to capture conditioning between sequences in an order independent manner (the order of sequences within a family is arbitrary) and to generalize to large context lengths, we propose a novel Transformer layer (§3.1.2) that models order-dependence between tokens within sequences and order-independence between sequences. PoET has the following properties:
- PoET can be used as a retrieval-augmented protein language model by conditioning the model on sequences from any family of interest. This also allows PoET to be used with any sequence database and to incorporate new sequence information without retraining.
- PoET is a fully autoregressive generative model, able to generate and score novel indels in addition to substitutions, and does not depend on MSAs of the input family, removing problems caused by long insertions, gappy regions, and alignment errors.
- By learning across protein families, PoET is able to extrapolate from short context lengths allowing it to generalize well even for small protein families.
- PoET can be sampled from and can be used to calculate the likelihood of any sequence efficiently.
We demonstrate these properties and show that PoET outperforms existing protein language models and evolutionary sequence models for variant effect prediction in extensive experiments on the 94 deep mutational scanning datasets in ProteinGym. PoET improves substitution effect prediction across proteins of all MSA depths, and also improves effect prediction of indels. A simple weighted average ensemble of PoET with existing methods further improves performance both across MSA depths and in sequences with large numbers of mutations. Furthermore, when used for generation, PoET controllably produces diverse, novel sequences that more closely match natural statistical constraints, resulting in better folding pLDDTs, than other generative protein language models. We expect PoET to become integral to future protein mutagenesis efforts.
Evolutionary sequence models are well established methods in biological sequences analysis. To model protein families, these models search large protein sequence databases for homologs, align the positions of these homologs in an MSA, and then fit statistical sequence models to the MSA. Common models include site independent models [15], profile HMMs [16], and coupling models [17]. Newer variants incorporate higher order correlations between sequence positions by training a VAE [13] or by building phylogentic trees [12]. These approaches are often referred to as "alignment-based" and must be fit on a family-by-family basis, requiring large numbers of members to generalize. A significant limitation of these models is that they assume the MSA is an accurate model of the evolutionary process generating the sequences, when in fact, MSA algorithms inevitably make alignment errors; regions with long insertions or lots of gaps can be particularly problematic. Furthermore, they cannot model novel insertions or deletions (indels) that are not present in the training MSA.
Unconditional protein language models that do not condition on homologs at inference have emerged as powerful methods for understanding and generating protein sequences. Both bidirectional models [4, 6, 18-20] and autoregressive generative models [5, 10, 11] have demonstrated competitive performance for variant function prediction. The latter type of model has the advantage of being able to score indels, but both cannot integrate evolutionary context not present in the trained model parameters. In contrast to family-specific evolutionary sequence models trained on sequences derived from a specific protein family [13, 17, 21], these models are trained on large protein databases [22, 23] that span all known sequences. This enables them to learn evolutionary constraints that generalize across families to improve predictions for small families with few homologs, but they generally underperform family-specific models for larger families.
Conditional protein language models fit between the unconditional and family-specific paradigms. Only a few works have explored this direction to date. Masked language models of whole MSAs are able to integrate evolutionary context directly for conditioning site predictions [24], but are unable to model insertions in the alignment. Ram and Bepler [25] use an encoder-decoder framework to generate new sequences conditioned on an MSA, which removes the insertion limitation of Rao et al. [24], but still requires conditioning on aligned input sequences. Notin et al. [11] combine predictions from an unconditional language model and an alignment-based model and show that integrating retrieval-based methods with protein language models can improve variant function prediction performance. However, the reliance on an alignment-based model means that the combined model is still limited by the constraints of MSAs.
Retrieval-augmented language models have shown impressive results in natural language processing, especially on Question Answering tasks. These models incorporate a database search as part of the text generation process in order to generate new text conditioned on prototypes found in the database [26-32]. In this way, they are conceptually similar to the conditional protein language models above. Retrieval-augmented approaches have the advantage of not requiring the entire training corpus to be encoded within the model parameters and the ability to easily integrate new data without retraining by simply adding it to the retrieval database. We adopt a similar approach with PoET in order to realize similar benefits. However, we build on well-established, fast, and accurate protein search tools for retrieval and frame the protein sequence generation problem as a sequence-of-sequences problem to incorporate retrieved-sequence conditioning, a fundamentally different paradigm than that employed by current retrieval-augmented models in natural language processing.
3 The Protein Evolutionary Transformer (PoET)
PoET is an autoregressive generative model of the distribution over protein families, where each family is generated as a sequence-of-sequences. Specifically, it models the distribution P(X=x), where x = s1, s2, ..., sn is the concatenation of n sequences si from the same family, and each sequence si = si,1, si,2, ..., si,Li is a sequence of Li amino acids padded by a start and end token. For example, below is a sequence of three protein sequences of lengths 4, 6, and 5 with start token denoted by $ and stop token denoted by*:
When referring to a sequence-of-sequences,  $s_i$ , which has one index  $i$ , refers to a sequence of tokens - the  $i$ th sequence in the sequence-of-sequences, whereas  $s_{i,j}$ , which has two indices  $i,j$ , refers to one token, the  $j$ th token of the  $i$ th sequence. We use  $x$  to denote the full sequence-of-sequences.
PoET generates each token in a sequence  $x$  one at a time, decomposing the probability of  $x$  as
$$
P (x) = P \left(s _ {1}, s _ {2},..., s _ {n}\right) = \prod_ {i = 1} ^ {n} P \left(s _ {i} \mid s _ {<   i}\right) = \prod_ {i = 1} ^ {n} \prod_ {j = 1} ^ {L _ {i}} P \left(s _ {i, j} \mid s _ {<   i}, s _ {i, <   j}\right) \tag {1}
$$
The order of the individual sequences in a sequence-of-sequences is arbitrary, and we propose a novel Transformer-based architecture to exploit this order invariance.
$$
p _ {3} = p (s _ {3} | s _ {1}, s _ {2}) = \prod_ {j = 1} ^ {4} P (s _ {3, j} | s _ {1}, s _ {2}, s _ {3, <   j})
$$
(a) PoET generates sets of homologous proteins as sequences-of-sequences. The equation in red demonstrates how to compute the conditional probability of a sequence given homologous sequences.
Figure 1: PoET Architecture
Tokens Associated With Position
(b) Inter-sequence relative position encoding scheme used in sequence-of-sequences self-attention that is invariant to sequence ordering. (Left) Relative positions between tokens in a sequence-of-sequences composed of two sequences. (Right) Tokens associated with absolute positions.
We propose a variant of the common Transformer decoder layer (Figure 1a) to capture order invariance between sequences while preserving order-dependence between tokens within sequences. We accomplish this using two attention modules: (1) a within-sequence module in which the representation at each position of each sequence is updated based on attending only to the other tokens within this sequence, and (2) a between-sequence module in which the representation at each position of each sequence is updated based on attending to other sequences within the sequence-of-sequences ( $\S 3.1.2$ ). This tiered approach is critical for capturing long-range dependencies between sequences and uniquely allows our model to extrapolate to much longer context lengths than used during training ( $\S 5.2.2$ ), improving sequence generation and performance on downstream tasks (Appendix H). Our full PoET model is a stack of these layers with causal self-attention.
3.1.1 Input Embedding
The input sequence  $x = s_{i,j}, i \in 1..n, j \in 1..L_i$  of amino acids and start/stop tokens is first converted into a sequence of continuous embeddings  $h_{i,j}$  by mapping each token to its learned embedding:
$$
h _ {i, j} = W _ {s _ {i, j}}, s _ {i, j} \in \mathbb {A A} \cup \{\text {S T A R T}, \text {S T O P} \} \tag {2}
$$
$$
W \in \mathbb {R} ^ {\left| \mathrm {A A} \cup \{\text {S T A R T , S T O P} \} \right| \times d}
$$
where AA is the set of 20 standard amino acids, and  $W$  is a matrix of learnable embeddings of dimension  $d$ .
3.1.2 Tiered Transformer Decoder Layers
Next, the embeddings  $h_{i,j}$  are transformed by  $N$  layers of TieredTransformerDecoderLayers (Appendix Algorithm 1), a novel layer for processing a sequence-of-sequences that is invariant to the order of the individual sequences, and extrapolates to context lengths substantially longer than the training context length.
The TieredTransformerDecoderLayer is composed of two phases. In the first phase, causal self-attention is applied independently to each sequence  $h_i$  of the input sequence-of-sequences, transforming them into new sequences  $f_i = \text{PerSequenceSelfAttn}(h_i)$ . Relative positional information is encoded by applying Rotary Positional Encodings (RoPE) [33] to the queries and keys before applying self-attention in the standard manner; the absolute position for  $f_{i,j}$  is  $j$ .
The second phase applies causal self-attention to the entire sequence-of-sequences by concatenating the individual  $f_{i}$  from the previous layer into one sequence before applying self-attention:  $g_{i,j} = \text{SequenceOfSequencesSelfAttn}([f_{<i}; f_{i, < j}])$ . In order to make self-attention in this phase invariant to sequence order, we adopt a simple but novel inter-sequence relative positional encoding scheme (Figure 1b): for  $g_{i,j}$  the absolute position is  $j$ . Just as in the first phase, the absolute position for tokens in the  $i$ th sequence  $g_{i}$  is independent of the position  $i$  of the sequence in the sequence-of-sequences. Thus, the positional information encoded by RoPE in this layer alone does not distinguish between the positions of tokens in different sequences. For example, the relative position between the first token of the first sequence  $f_{1,1}$  and the first token of the second sequence  $f_{2,1}$  is 0. The fact that these two tokens come from two different sequences is encoded by the first phase, which operates on the two sequences independently. This inter-sequence relative positional encoding scheme has two useful properties:
1. it encodes the fact that amino acids at similar absolute positions in homologous proteins are more likely to be drawn from the same distribution
2. it limits the maximum relative position encoding needed to the number of tokens in an individual protein sequence $^{2}$ , rather than the number of tokens in a sequence-of-sequences, allowing the model to generalize to longer sequences-of-sequences than seen during training
3.1.3 Decoded Probabilities
Lastly, the output from the last TieredTransformerDecoderLayer,  $g_{i,j}$ , is decoded into token probabilities by applying a linear transformation  $P(s_{i,j}|s_{<i}, s_{i,<j}) = p_{i,j}(s_{i,j}) = \text{Linear}(g_{i,j})$ . Here,  $p_{i,j}$  is a vector of probabilities, one for each distinct token  $\in \mathbb{A}\mathbb{A} \cup \{\text{START}, \text{STOP}\}$ , and  $p_{i,j}(s_{i,j})$  is the probability of the token  $s_{i,j}$  according to  $p_{i,j}$ .
3.2 Training Data
Models were trained on 29 million sets of homologous sequences. Each set corresponds to a sequence in UniRef50 Version 2103, and contains all its homologs in UniRef50 found using Diamond [34]. We removed any sets with fewer than 10 homologs. To avoid overfitting on promiscuous sequences which may belong to a large number of sets, we sample each set with weight inversely proportional to the size of the set ("inverse count" sequence weighting). The order of sequences within a sequence-of-sequences is randomized to promote order invariance. See Appendix B for more details.
4 Protein Variant Fitness Prediction
Protein variant fitness prediction is the task of assigning a score to each sequence in a set of variants  $\{v_{1}, v_{2}, \ldots, v_{n}\}$  of a target sequence  $t$  that accurately reflects the relative fitness of the variants. A protein variant  $v_{i}$  can be any sequence with a limited number of substitutions, insertions, and/or deletions relative to the target  $t$  that the experimenter believes may have improved fitness. Fitness
2While the length of a protein is unbounded (but finite), in practice, the vast majority  $(>99\%)$  of proteins are less than 1024 amino acids long [11].
Figure 2: Illustration of evaluating PoET for variant fitness prediction on a DMS dataset
refers to the value of any property of a protein sequence related to function that the experimenter is interested in optimizing e.g. thermostability, enzymatic activity, etc.
Benchmarking using deep mutational scans Deep mutational scanning (DMS) is a method for measuring the fitness properties of thousands to millions of variants of a protein using next-generation sequencing techniques [1]. Data from these experiments have been used to evaluate fitness prediction methods. We evaluate PoET on ProteinGym [11], the largest collection of such data yet, containing 87 datasets with substitution variants and 7 datasets with indel variants. We use the same validation set as Notin et al. [11] for tuning hyperparameters.
Fitness Prediction with PoET PoET predicts the fitness of a variant as the conditional log-likelihood of the variant  $v_{i}$  given a set of sequences  $S$  homologous to the target  $t$ :
$$
\hat {F} _ {i} \left(S = \left\{s _ {1}, \dots , s _ {m} \right\}\right) = \log P \left(v _ {i} \mid s _ {1}, s _ {2}, \dots , s _ {m}\right) = \sum_ {j = 1} ^ {L _ {i}} \log P \left(v _ {i, j} \mid s _ {1}, s _ {2}, \dots , s _ {m}, v _ {i, <   j}\right) \tag {3}
$$
The set  $S$  is retrieved by searching a large database of proteins such as UniRef100 [23] for sequences homologous to  $t$ . The homologs form a diverse set of sequences that define the protein family of interest. By conditioning on these sequences, PoET can infer evolutionary constraints on the protein family to improve fitness prediction. The sequences are conditioned on in an arbitrary order. Figure 2 illustrates the evaluation of PoET for variant fitness prediction on one DMS dataset.
Based on validation set performance, we selected the best method for (1) retrieving homologous sequences (Appendix C), (2) subsampling and filtering the homologous sequences to a reasonable context length for efficient inference (Appendix D), and (3) assembling conditional log-likelihoods computed from different subsamples of homologous sequences. We present results with the best settings in the main results ( $\S 5$ ). Our final approach uses the ColabFold protocol [35] for retrieving homologs, and the final fitness prediction score is obtained by averaging the conditional log-likelihoods across subsamples of the full set of retrieved homologous sequences:
$$
\hat {F} _ {\text {e n s e m b l e}, i} (S) = \frac {1}{N _ {\text {e n s e m b l e}}} \sum_ {j = 1} ^ {N _ {\text {e n s e m b l e}}} \hat {F} _ {i} \left(S _ {j} \subset S\right) \tag {4}
$$
The subsets  $S_{j}$  are drawn by sampling with sequence weights [17], maximum sequence identity to the target  $t$  in  $\{1.0, 0.95, 0.90, 0.70, 0.50\}$ , and total tokens (i.e. context length) in  $\{6144, 12288, 24576\}$ . All combinations of these parameters are used, for a total of  $N_{\text{ensemble}} = 15$ .
Table 1: Average Spearman correlation between model scores and experimental measurements on ProteinGym by MSA depth, and # of parameters in language models. MSA depth measures the amount of sequence information the MSA of the homologous sequences defining a protein family contains about the target protein (Appendix E). N/A indicates not applicable, whereas a dash (-) indicates applicable, but not computed.
Model Type	Model name	#
Param	Substitutions by MSA Depth
Low	Medium	High	All	Indels
Align-
ment-
based	Site independent	N/A	0.417	0.404	0.411	0.408	N/A
GEMME	N/A	0.445	0.449	0.522	0.463	N/A
EVE (ensemble)	N/A	0.414	0.441	0.498	0.448	N/A
Uncond-
ditional	ESM-1v (ensemble)	3.25B	0.356	0.372	0.510	0.398	N/A
ProGen2 (ensemble)	10.8B	0.357	0.416	0.448	0.411	0.407
PLM	Trancection L (no retrieval)	700M	0.377	0.399	0.429	0.401	0.430
Con-
ditional	MSA Transformer (ens.)	100M	0.372	0.421	0.477	0.423	N/A
PoET (ensemble)	201M	0.476	0.466	0.542	0.484	0.510
Hybrid	Trancection L	700M	0.441	0.437	0.472	0.445	0.464
TranceptEVE M	300M	-	-	-	-	0.516
TranceptEVE L	700M	0.454	0.463	0.508	0.471	0.466
PoET (ensemble)
+ TranceptEVE L	901M	0.479	0.480	0.537	0.492	0.521
5 PoET achieves state-of-the-art performance in variant fitness prediction
We evaluate PoET's ability to predict the relative fitness of protein variants by measuring the average Spearman correlation  $\bar{\rho}$  between the predicted fitness and the measured fitness across the DMS datasets in ProteinGym. We compare PoET to the best existing alignment-based models, unconditional and conditional protein language models, and hybrid models that combine multiple approaches (Results Summary: Table 1, Per Dataset: Figures S1, S2). A brief overview of the baseline methods is provided in Appendix F. Statistical significance was assessed using a paired t-test. For a fair comparison, we ran the relevant baselines with all homologous sequence retrieval methods considered for use with PoET (Appendix C), and present the results for each baseline method paired with its best retrieval method. We do not tune parameters for assembling the same method with different subsamples of the set of retrieved homologs (Equation 4) because these parameters need to be considered and defined on a per method basis, and many baselines already have their own assembling strategy as defined by their authors [4, 12, 14].
5.1 Comparison to baselines