SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Substitutions Benchmark On the substitutions benchmark, PoET performs comparably to TranceptEVE L, the best performing baseline method (Table 1). PoET improves average Spearman correlation  $\Delta \bar{\rho} = 0.013$ , but the difference falls just short of statistical significance  $p = 0.05029$ . Although the methods perform similarly, PoET offers a number of advantages. Inference with PoET is substantially faster when considering on the order of tens to hundreds of thousands of variants (Appendix J), which allows users to more quickly accomplish common tasks such as computing the fitness of all single site substitutions. By avoiding the costly step of training a new model for each protein family, users can more easily experiment with different ways of defining the protein family. For example, users can retrieve the homologous sequences using different sequence search programs and settings, and prior knowledge can be incorporated by selecting the subset of most relevant homologous sequences e.g. sequences from the taxon of interest, or that are known to be functional (Appendix K). Since PoET is able to generalize across protein families, whereas the EVE component of TranceptEVE L cannot, it is less critical that the protein family of interest is defined by a large number of homologous sequences (Appendix C).
On the other hand, if one is willing to spend more time on inference and the protein family is well defined a priori, better fitness predictions can be obtained by ensembling PoET with TrancaptEVE L (Results: Table 1, Methods: Appendix G). An ensemble of PoET with TrancaptEVE L performs significantly better than TrancaptEVE L alone ( $\Delta \bar{\rho} = 0.021$ ,  $p < 7\mathrm{e - }6$ ). The performance improvement is consistent across subsets of the substitutions benchmark broken down by MSA depth (Table 1), mutation depth (Appendix Table S2), and taxon (Appendix Table S3).
We also considered ensembles of PoET with other baselines methods (Appendix Table 7). The ensemble of PoET with GEMME is also notable as it performs similarly to the ensemble with TrancaptEVE L on the substitutions benchmark, and requires very little additional time to compute; GEMME is thousands of times faster than either PoET or TrancaptEVE L [12]. The main disadvantage is that GEMME is unable to score indels.
Indels Benchmark On the indels benchmark, the best performing baseline model is TrancaptEVE M, a variation of TrancaptEVE L that uses a protein language model with fewer parameters  $(\Delta \bar{\rho} = 0.05)$ . PoET and the ensemble of PoET with TrancaptEVE L both outperform TrancaptEVE L ( $\Delta \bar{\rho} \in [0.044, 0.055]$ ), and perform comparably to TrancaptEVE M ( $\Delta \bar{\rho} \in [-0.006, 0.005]$ ). However, no differences between the aforementioned models are statistically significant. One advantage of PoET is that it is able to not only score indels, but also generate sequences with indels (§6). It is also able to score and generate insertions not present in the MSA.
5.2 Characterizing the PoET architecture
5.2.1 Training Distribution and Model Size
We train variations of PoET with different training distributions and model sizes and investigate their effect on fitness prediction (Figure 3).
Context Length We trained 57M parameter versions of PoET for up to 3 days on 7 x A100 GPUs with three context lengths: 4K, 8K, and 16K. Notably, these context lengths far exceed the 1K or 2K context lengths typically used to train other protein and natural language models [5, 11, 24, 36, 37]. We reason that longer context lengths are needed for the model to observe a sufficient number of homologs; even with a context length of 8K, for a sequence of length 500, which is the length of the typical human protein in ProteinGym, the model would only observe  $\sim 16$  of possibly thousands of homologs. Interestingly, while a context length of 8K performs better than a context length of 4K, increasing the context length further to 16K had a neutral or negative effect on variant effect prediction. It does improve perplexity of heldout sequences.
Sequence Set Weighting Next, we trained an 8K context length model using the naive strategy of sampling UniRef50 sequence sets uniformly instead of using our "inverse count" strategy that weights sequence sets based on their size (§3.2). We find that the "inverse count" sampling strategy substantially outperforms the naive strategy. Other methods that also train on sequence sets like MSA Transformer [24] may also benefit from this sampling strategy.
Model Size Lastly, using the optimal parameters found thus far, we trained 57M, 201M, and 604M parameter models to explore the effect of model size. Models were trained until validation performance plateaued. Increasing the model size to 201M parameters significantly improves performance, but increasing the size further to 604M appears to have a neutral effect. We chose the 201M version as the standard PoET model.
5.2.2 Model Architecture
To determine the importance of using TieredTransformerDecoderLayers (§3.1.2) instead of regular Transformer layers, we trained, on the same task, a regular 201M parameter RoPE-based Transformer that ignores the sequence-of-sequences structure of the input. We compared the generative capabilities of PoET and the Transformer by measuring the perplexity at context lengths both within and beyond the training the context length (Figure 4).
We trained both models with a context length of 8K for 500K steps. As expected based on previous studies[38], the RoPE-based Transformer does not generalize to context lengths beyond the training context length. In contrast, PoET generalizes to  $\sim 8x$  the training context length (61K), which is the longest we tested. PoET also performs better on context lengths within the training context length,
Figure 3: Performance of PoET on the ProteinGym validation set when trained with (left) various training distributions and (right) model sizes.
Figure 4: Comparison of the perplexity of a regular Transformer and PoET when generating a protein sequence conditioned on a fixed number of tokens from other sequences in the same protein family. Protein families consist of UniRef50 sequences, and both models perform much better than a profile HMM baseline (Appendix Figure 12).
but by a smaller margin. Next, we finetuned both models with 16K and 32K context lengths for 50K and 30K steps respectively. While the perplexity of the Transformer improved significantly at the 16K and 32K context lengths, the Transformer still does not generalize to context lengths beyond the training context length, and underperforms all variants of PoET at all context lengths. The Transformer also underperforms PoET for variant fitness prediction (Appendix H).
6 PoET generates novel sequences that preserve structure within a protein family
PoET's state-of-the-art performance on variant fitness prediction demonstrates that it assigns higher likelihood to regions of sequence space with high fitness variants. This suggests that PoET can be used to directly generate high fitness variants belonging to a protein family by conditioning PoET on sequences from the protein family and sampling from the resulting conditional distribution. Direct generation of variants makes the exploration of higher order mutants computationally tractable as sequence space grows exponentially large with number of mutations, making it impossible to explore this space by scoring all such variants.
To evaluate PoET's ability to generate sequences from specific protein families, we sampled 1,000 sequences each from PoET conditioned on homologs from two protein families, phage lysozymes (PF00959) and chorismate mutases with activity in E. coli (Appendix K.2), and examined the novelty of the generated sequences and their predicted structural conservation. Sequence novelty was measured as the max sequence identity to any natural sequence (of the same protein family). Predicted
Figure 5: Sequence novelty and predicted structural conservation of (A) functional chorismate mutases generated by PoET and (B) phage lysozymes generated by PoET and ProGen. PoET generates diverse sequences (50-100% seq id to a natural sequence) while preserving 3D structure within the protein family (TM-score  $>0.8$  to a structure of a natural sequence and pLDDT  $>90$ ).
structural conservation was measured as the max TM-score between the predicted AlphaFold2 [39] structure of the generated sequence and the solved structure of any natural sequence. For both protein families, PoET generates high diversity sequences that are predicted with high confidence to be structurally similar, supporting the quality of sequences generated by PoET (Figure 5).
Furthermore, when tasked to generate phage lysozymes, we find that sequences generated by PoET are more diverse and better preserve structure at similar levels of sequence novelty than sequences generated by ProGen [8] (Figure 5B), a 1.2B parameter model fine-tuned specifically for generating phage lysozymes. As many sequences generated by ProGen have been experimentally confirmed to be functional phage lysozymes, these results suggest that PoET generates functional sequences with even greater diversity even without fine-tuning. Since fine-tuning ProGen is critical for achieving good variant effect prediction [8], and a version of ProGen fine-tuned on chorismate mutases is not available, we did not compare PoET and ProGen for generating novel chorismate mutases.
PoET is a Transformer-based autoregressive generative model of whole protein families. By framing family generation as a sequence-of-sequences generation problem, we are able to train across tens of millions of protein sequence clusters to encode fundamental rules of protein evolution into the PoET model. This enables the model to generalize to protein families unseen during training and extrapolate from small numbers of conditioned-upon sequences. The sequence-of-sequences generative framework allows PoET to be used as a retrieval-augmented language model, generating new sequences conditioned on a set of sequences representing the family or other properties of interest. We demonstrate that PoET improves over other protein language models and evolutionary sequence models for variant fitness prediction across a wide range of deep mutational scanning datasets. PoET also enables efficient sequence generation and the generative distribution can be controlled via conditioning. Phage lysozyme- and chorismate mutase-like sequences sampled from PoET are novel and predicted to fold with high confidence. PoET can be backed by other sequence databases and naturally improves as databases grow without the need for retraining. We anticipate that PoET will become a fundamental component of ML-enabled protein design in the future.
[48] Nadav Brandes, Dan Ofer, Yam Peleg, Nadav Rappoport, and Michal Linial. ProteinBERT: a universal deep-learning model of protein sequence and function. Bioinformatics, 38(8): 2102-2110, 02 2022. ISSN 1367-4803. doi: 10.1093/bioinformatics/btac020. URL https://doi.org/10.1093/bioinformatics/btac020.
[49] Lin Li, Esther Gupta, John Spaeth, Leslie Shing, Tristan Bepler, and Rajmonda Sulo Caceres. Antibody representation learning for drug discovery, 2022.
[50] Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. In International Conference on Learning Representations, 2020. URL https://openreview.net/forum?id=rygGQyrFvH.
A Detailed description of TieredTransformerDecoderLayer
Algorithm 1 TieredTransformerDecoderLayer
Require: representations for a sequence of sequences  $h_{i,j} \in \mathbb{R}^d, i \in 1..N, j \in 1..L_i$  First, apply causal self-attention to each sequence individually  
 $f \gets \text{LayerNorm}(h)$ $q_{i,j} \gets \text{RoPE(Linear}(f_{i,j}), j) \forall i, j$ $k_{i,j} \gets \text{RoPE(Linear}(f_{i,j}), j) \forall i, j$ $v \gets \text{Linear}(f)$ $f_i \gets f_i + \text{CausalAttention}(q_i, k_i, v_i) \forall i$   
Next, apply causal self-attention to all sequences together  
 $g_{i,j} \gets f_{i,j} \forall i, j$ $g \gets \text{LayerNorm}(g)$ $q_{i,j} \gets \text{RoPE(Linear}(g_{i,j}), j) \forall i, j$ $k_{i,j} \gets \text{RoPE(Linear}(g_{i,j}), j) \forall i, j$ $v \gets \text{Linear}(g)$ $g \gets g + \text{CausalAttention}(q, k, v)$   
Finally, the feedforward layer  
 $g' \gets \text{LayerNorm}(g)$ $h' \gets \text{GELU(Linear}(g'))$ ,  $h' \in \mathbb{R}^{4d}$ $g' \gets g + \text{Linear}(h')$   
output  $g' \in \mathbb{R}^d$
B Additional Training Details
To find homologs in UniRef50 using Diamond, we used the following command, which defines an all-against-all search:
diamond blastp -q uniref50.fasta -d diamond/uniref50 -f 6 -header -k 200000  
-max-hsps 1 -e 0.001 -p 96 -o output.tab
The command returns, for each sequence in UniRef50, a set containing all its putative homologs in UniRef50. We call each such set a "Diamond-UniRef50 Cluster". Diamond was used over other homology search tools due to its high performance ( $>100\mathrm{x}$  speed of BLAST).
To form a training example, sequences are sampled without replacement from a "Diamond-UniRef50 Cluster" until the total number of tokens reaches a predetermined limit, and then concatenated to form a sequence-of-sequences. Each sampled UniRef50 sequence is replaced with a UniRef100 sequence by sampling a random UniRef100 sequence from the same UniRef50 cluster as the UniRef50 sequence being replaced. The UniRef100 sequences are randomly sampled with weight inversely proportional to the size of the UniRef90 clusters they belong to. Each UniRef100 sequence is sampled at most once in each sequence-of-sequences.
As a final data augmentation, the order of the tokens in a sequence-of-sequences is reversed with probability  $50\%$  (i.e. all sequences in a sequence-of-sequences are ordered from either N-terminus to C-terminus only or C-terminus to N-terminus only). This augmentation has been shown to improve the performance of other protein language models [11].
Following this sampling procedure, the order of sequences in a sequence-of-sequences is random, which promotes order invariance. To evaluate the sensitivity of PoET to the order of sequences, we compare heldout sequence perplexities after conditioning on 10 sequences ordered either 1) randomly as described above, 2) shortest-to-longest, or 3) longest-to-shortest. We also compare the full joint log likelihoods of these sequences of 10 sequences with the same orderings (Table 2).
We find that sequence ordering has little-to-no effect on next sequence perplexity and joint log likelihoods. Ordering sequences shortest-to-longest changes next sequence perplexity by only 0.00753 on average and has a relative difference of  $< 0.5\%$  from random for both next sequence perplexity and joint log likelihood.
Table 2: Comparison of (a) next sequence perplexities and (b) joint log likelihoods  $\log p(s_1,s_2,\ldots ,s_n)$ , of sequences of 10 sequences with different orderings from heldout Uniref50 sequence clusters. We report mean perplexity and joint log likelihood as well as mean difference from random and relative difference from random for shortest-to-longest and longest-to-shortest orderings with the standard deviation across clusters in parenthesis. Relative difference is calculated as  $(a - r) / r$  where  $a$  is the ordered value and  $r$  is the value with random ordering.
(a) Perplexity of Next Sequence N=9184
Sequence Order	Mean Perplexity	Mean Difference from Random	Mean Relative Difference from Random
Random	7.44075	-	-
Shortest-to-Longest	7.44828	+0.00753 (0.07686)	+0.00109 (0.01033)
Longest-to-Shortest	7.44410	+0.00336 (0.07608)	+0.00055 (0.01030)
(b) Joint Log Likelihood N=9184
Sequence Order	Mean Joint Log Likelihood	Mean Difference from Random	Mean Relative Difference from Random
Random	-8099.74973	-	-
Shortest-to-Longest	-8123.18167	-23.43149 (55.76313)	-0.00335 (0.00838)
Longest-to-Shortest	-8091.89462	+7.85511 (52.92428)	+0.00105 (0.00746)
B.2 Optimizer and learning rate
We used the AdaFactor optimizer [40] with initial learning rate 1e-2, square root learning rate decay, and otherwise default parameters.
Hyperparameters for PoET variations used in ablation experiments (§5.2.1) are summarized in Table 3.
Table 3: Hyperparameters for PoET models of different sizes.
# Parameters	NLayers	Hidden Dim	NHead	Batch Size (# Tokens)
57M	6	768	12	500k
201M	12	1024	16	250k
604M	16	1536	24	250k
C Homologous Sequence Retrieval Methods
We experimented with two methods for retrieving homologous sequences from UniRef100, JackHMMer [41] and MMseqs2 [42]. Both programs return an MSA of homologs. For JackHMMer, we used the MSAs provided by ProteinGym [11], which were constructed from UniRef100 Version 2104. For MMseqs2, we used the protocol defined by ColabFold [35], and retrieved homologs from UniRef100 Version 2202. The main advantage of the ColabFold protocol is that it is substantially faster than the ProteinGym protocol (minutes vs up to hours), whereas the ProteinGym protocol may be more sensitive.
We evaluated all relevant baselines and PoET with homologs retrieved by both protocols (Table 4). Note that in the main result table (Table 1), we report the results of each method with the best protocol based on validation set performance. We replicate the findings of Laine et al. [12], who show that GEMME performs similarly with both protocols, and also find that most models perform similarly regardless of the protocol. The notable exception is EVE, which performs significantly worse. This is
Figure 6: Redundancy, as measured by  $N / N_{\mathrm{eff}}$ , of homologous sequence sets retrieved by the ProteinGym and ColabFold retrieval methods for target proteins in ProteinGym.
most likely because the MSAs returned by the ColabFold protocol are substantially less redundant (Figure 6). Thus they do not contain enough sequences to properly train EVE models, which are trained from scratch for each protein family and can contain millions of parameters. Removing redundancy filters in the protocol could improve results for EVE; we leave this to further research.
Given the poor performance of EVE with the ColabFold protocol, we did not also evaluate the EVE ensemble or TranceptEVE M/L with ColabFold, presuming that the performance of both would also be worse with ColabFold because they both depend on EVE. Furthermore, we did not evaluate the performance of Tranception with ColabFold because Tranception relies on homologs via the site independent model, which performs similarly with both protocols.
Table 4: Average Spearman correlation on ProteinGym subsets of baseline models and PoET using different methods for retrieving homologs.
Model	Retrieval Method	Val	Subst.
Site independent	ProteinGym	0.386	0.408
ColabFold	0.371	0.404
GEMME	ProteinGym	0.416	0.457
ColabFold	0.418	0.463
EVE	ProteinGym	0.423	0.443
ColabFold	0.268	0.342
MSA Transformer (ens.)	ProteinGym	0.427	0.431
ColabFold	0.440	0.423
PoET (ens.)	ProteinGym	0.463	0.474
ColabFold	0.481	0.484
D Subsampling and Filtering of Retrieved Homologous Sequences
Even with their high sensitivity and specificity, homologous sequence retrieval methods are imperfect, and inevitably run into at least the following issues:
1. some sequences will be retrieved that do not share function with the target sequence, either due to false positives (i.e. sequences found that do not in fact share common ancestry), or because sequences have diverged in function despite their common ancestry [43]
2. the set of sequences retrieved will be biased towards those that have been chosen to be sequenced by humans, or by natural processes that may not reflect relative fitness [17]
To resolve the second issue, we adopt the sequence weighting scheme used by Hopf et al. [17], which has been shown to work well in many previous studies [4, 11, 13, 40]. The weighting scheme is detailed in Appendix D.5.
To address the first issue, existing methods have used sequence identity as a simple heuristic for filtering out irrelevant or misleading sequences. Laine et al. [12] apply a max dissimilarity filter of 0.8 sequence identity to remove sequences that are less likely to be homologous, and a max similarity filter of 0.98 sequence identity to remove sequences that are redundant with the target. Meier et al. [4] explore using a variety of max similarity filter thresholds and find that a threshold of 0.9 or 0.75 tends to be optimal, depending on the target. Both Laine et al. [12] and Meier et al. [4] also limit the total number of sequences in a subsample, which helps to limit the computational cost of inference, and ensemble the results from applying their model to multiple subsamples to obtain better predictions.
To determine the best way to subsample and filter homologous sequences for PoET, we evaluate the performance of the model on the ProteinGym validation set with (1) various max dissimilarity and (2) similarity thresholds for filtering, and (3) various context lengths (i.e. total number of tokens used to compute the conditional log-likelihood (Equation 3), including the tokens of the variant sequence being scored) for limiting the computational cost of inference.
D.1 Maximum Sequence Dissimilarity
First, we explored the impact of the maximum sequence dissimilarity threshold. We evaluate PoET with (1) homologous sequences retrieved using either the ProteinGym or the ColabFold protocol, (2) a context length of 8192, and (3) various max sequence dissimilarity thresholds (Figure 7). When retrieving homologous sequences using the ProteinGym protocol, a threshold of 0.7 is optimal for both the validation and substitution benchmarks, and is significantly better than no filtering i.e. using a threshold of 1.0 (substitutions benchmark:  $\Delta \bar{\rho} = 0.024, p < 0.01$ ). Interestingly, filtering was particularly helpful for prokaryotic datasets, suggesting that the ProteinGym protocol may be prone to finding false positives in this taxon. In contrast, for the ColabFold protocol, no filtering based on max dissimilarity is necessary; there is essentially no difference in performance when filtering with a threshold of 0.7.
D.2 Maximum Sequence Similarity and Context Length
Next, using the optimal maximum sequence dissimilarity thresholds, we evaluated PoET while jointly varying the max similarity threshold  $\in \{1.0, 0.95, 0.90, 0.70, 0.50\}$  and context length  $\in \{6144, 12288, 24576\}$  (Figure 8). When averaged across datasets, we found that PoET generally performed better when some sequences are filtered i.e. max similarity threshold  $< 1.0$ , and with smaller context lengths (6144 or 12288). However, when looking at individual datasets (Figure 9), we find a wide variety of behaviors. On some datasets, increasing the context length consistently improves performance across max similarity thresholds and the optimal max similarity threshold is 1.0 (no filtering), whereas for others the opposite is true, and yet for others the behavior lies somewhere in between.
From these observations, we draw two main conclusions. First, although the average performance is lower at higher context lengths, the ability of PoET to improve in performance with increased context length on some datasets suggests that PoET is able to generalize to context lengths much longer than the training context length, and that decreased average performance at higher context lengths may be due to factors other than the model's ability to handle long context lengths. Weinstein et al. [44] suggest that optimal modeling of the sequences in an observed protein family may be detrimental to protein fitness prediction due to biases in the sequences we observe. Consequently, protein fitness prediction based on evolutionary models perform better when they do not or cannot model the family exactly e.g. via model misspecification. In the context of PoET then, limiting the max sequence similarity and context length when sampling sequences from the full set of retrieved homologs may