SlowGuess/NeurIPS_2023_OCR huggingface dataset:
(a) Performance using ProteinGym protocol for retrieving homologous sequences
Figure 7: Optimal threshold for filtering homologous sequences by max dissimilarity to target sequence differs depending on retrieval method Left: performance as a function of max dissimilarity threshold when using the ProteinGym protocol (top) and the ColabFold protocol (bottom). For ProteinGym, a threshold of 0.7 is optimal and significantly better than no filtering i.e. using a threshold of 1.0 (paired t-test on full substitutions benchmark:  $p < 0.01$ , effect size:  $\Delta \bar{\rho} = 0.024$ ). For ColabFold, no filtering is necessary and there is no essentially no difference in performance when filtering with a threshold of 0.7. Right: comparison of Spearman correlations on datasets when filtering homologous sequences to 0.7 max dissimilarity vs not filtering, broken down by taxon. For ProteinGym (top), filtering improves performance particularly on Prokaryotic datasets, whereas for ColabFold (bottom), there is no clear pattern.
(b) Performance using ColabFold protocol for retrieving homologous sequences
help avoid some of these biases while still providing the model information about the specific protein family of interest. Second, the large variance in the optimal max similarity threshold and context length on a per dataset level suggests that better sampling approaches, possibly dataset specific, may be an avenue for further improving protein fitness prediction.
D.3 Ensemble over subsamples
Given the variance in best max sequence similarity and context length parameters on the validation datasets, we explored ensembling (Equation 4) over subsamples where the subsamples are drawn with different values for max similarity and context length. Table 5 compares (1) not ensembling, (2) ensembling with  $N_{\text{ensemble}} = 15$  subsamples where the max sequence similarity and context length parameters are fixed to the best values on the validation set (Figure 8), and (3) ensembling subsamples with varied parameters (max similarity threshold ∈ {1.0, 0.95, 0.90, 0.70, 0.50} and context length ∈ {6144, 12288, 24576}; 15 combinations in total). On the validation and substitutions benchmarks, we find that both types of ensembling improve performance. On the validation benchmark, the two ensembling methods perform similarly; fixed has a slight edge when using the ProteinGym protocol
(a) Performance using ProteinGym protocol for retrieving homologous sequences
Figure 8: Averaged across datasets, PoET performs best with some filtering based on max similarity and at moderate context lengths. Each line plot shows the performance of PoET on ProteinGym as the max similarity threshold and context length is varied. Left: Performance on the validation set. Right: Performance on the full substitutions benchmark.
Table 5: Average Spearman correlation between model scores and experimental measurements on ProteinGym for PoET when the homologous sequence subsampling parameters are fixed or varied across the ensemble, broken down by retrieval and sampling methods.
Retrieval Method	Max Similarity and Context Length Parameters	Val	Substitutions	Indels
ProteinGym	Fixed (no ensemble)	0.442	0.445	0.474
Fixed (15x ensemble)	0.464	0.467	0.521
Varied (15 combinations)	0.463	0.474	0.498
ColabFold	Fixed (no ensemble)	0.451	0.445	0.485
Fixed (15x ensemble)	0.478	0.469	0.512
Varied (15 combinations)	0.481	0.484	0.510
for retrieving homologs, and varied has a slight edge when using the ColabFold protocol. On the full substitutions benchmark, the varied method consistently performs better, and by a larger margin; this result may be less biased than the validation benchmark result because the parameters for the fixed method are tuned on the validation benchmark.
D.4 Summary
Based on validation benchmark performance of the PoET variants explored in the above sections, for our final PoET model, we choose the following parameters
Figure 9: Optimal max similarity threshold and context length varies across ProteinGym datasets. Each line plot shows the performance on a validation dataset as the max similarity threshold and context length is varied. Datasets are grouped based on the effect of increasing context length for that dataset. A: Datasets for which increasing the context length tends to improve performance. B: Datasets for which increasing the context length tends to have a neutral effect on performance. C: Datasets for which increasing the context length tends to have a negative effect on performance.
1. Retrieve homologous sequences using the ColabFold protocol
2. Do not filter homologous sequences by max dissimilarity when subsampling the set of homologous sequences
3. Ensemble the predictions of PoET across subsamples of the full set of homologous sequences with max similarity thresholds  $\in \{1.0, 0.95, 0.90, 0.70, 0.50\}$  and context lengths  $\in \{6144, 12288, 24576\}$  (15 combinations total)
D.5 Sequence Weighting During Homologous Sequence Sampling
We sampled homologous sequences using the sequence weighting scheme described in Hopf et al. [17]. Specifically, sequences are sampled from a probability distribution where a sequence's probability is inversely proportional to its number of neighbors in the MSA of the homologous sequences. A target sequence is considered a neighbor of a query sequence if the sequence identity to the query is greater than 0.8. The sequence identity is computed based on the alignment in the MSA, and gaps in the query sequence are ignored.
D.6 Full length sequences vs aligned fragments
In the subsampling and filtering of homologous sequences experiments above, we condition on only the amino acids in the homologous sequences that are aligned to a position of the target sequence in the MSA of the homologous sequences i.e. a sequence in a subset  $S_{j}$  (Equation 4) may not be the full length sequence as recorded in UniRef100, rather it is potentially a non-contiguous subsequence of the UniRef100 sequence. We find that using only these "aligned fragments" performs similarly to using the full length sequences (Table 6), and use the former over the latter as the former can be derived directly from the MSA.
Table 6: Average Spearman's rank correlation on the validation and substitutions benchmarks using full length sequences vs aligned fragments sampled from MSAs of homologous sequences generated by the ColabFold protocol.
Method	Val	Substitutions
Full Length Sequences	0.464	0.474
Aligned Fragments	0.466	0.479
E MSA Depth
MSA depth measures the amount of sequence information the MSA contains about the target protein. We use the same definition and low/medium/high categorization as Notin et al. [11]. At a high level, the MSA depth  $N_{\mathrm{eff}} / L$  is computed as the number of non-redundant sequences  $N_{\mathrm{eff}}$  (at 0.8 sequence identity for non-viral proteins, and at 0.99 sequence identity for viral proteins) divided by the sequence length  $L$ . The MSA depth categories thresholds are - Low:  $N_{\mathrm{eff}} / L < 1$ ; Medium:  $1 < N_{\mathrm{eff}} / L < 100$ ; High:  $N_{\mathrm{eff}} / L < 100$ .
Although the MSA depth changes based on the method used to retrieve homologous sequences, we categorize proteins/datasets using the ProteinGym MSAs (1) in order to maintain consistency with the categorizations in the ProteinGym benchmark, which allows for easier comparison with existing work, and (2) because the MSA depth is highly correlated between the retrieval methods considered in this paper (Figure 10).
F Protein Variant Fitness Prediction Baselines
F.1 Overview of Baselines
We compare PoET to the following baseline models for protein variant fitness prediction, which includes the best existing alignment-based models, unconditional and conditional protein language models, and hybrid models that combine multiple approaches models:
Figure 10: Comparison of MSA depth ( $N_{\mathrm{eff}}$ ) of MSAs of homologous sequences retrieved by the ProteinGym and ColabFold protocols for target proteins in ProteinGym ( $\rho = 0.4, p < 0.001, n = 72$ ).
F.1.1 Alignment-based Models
Alignment-based models first compute a MSA of homologous sequences of a given protein family, and compute sequence likelihoods based on MSA statistics or models trained on the MSA.
Site Independent Model The simplest alignment-based model; sequence likelihoods are computed from the frequency of amino acids at each position of the MSA, considered independently from other positions of the MSA.
GEMME [12] is a fast, scalable, and simple algorithm for predicting variant effect that models the interdependencies between different positions in the MSA by considering the evolutionary tree relating the sequences in the MSA. It is among the best performing baseline models while also requiring substantially less compute than comparably performing alternatives, but it cannot score indels.
EVE [13] (ensemble) is an ensemble of five Bayesian variational autoencoder trained on the MSA.
F.1.2 Unconditional PLMs
Unconditional PLMs are large language models trained on protein databases spanning many protein families e.g. UniRef [23], using either the causal or masked language modeling objective.
ESM-1v [4] (ensemble) is an ensemble of five 650M parameter Transformer-based masked language models trained on UniRef90.
ProGen2 [5] (ensemble) is an ensemble of five Transformer-based causal language models ranging between 150M and 6.4B parameters and trained on UniRef90 and BFD [45].
Trancection L (no retrieval) [11] is a 700M parameter Transformer-based causal language model trained on UniRef100.
F.1.3 Conditional PLMs
Conditional PLMs are large language models trained on sets of homologous protein sequences spanning many protein families.
MSA Transformer [24] (ensemble) is a 100M parameter masked language model trained on MSAs of all UniRef50 sequences. MSAs are generated using HHblits [46].
F.1.4 Hybrid
Hybrid models combine predictions from different types of models.
Trancection L [11] combines predictions from Trancection L (no retrieval) and a site independent model in a manner that allows novel indels to be scored by aligning the sequence to be scored to the MSA used by the site independent model.
TranceptEVE [14] combines predictions from Tranception and EVE in a manner that allows novel indels to be scored by using EVE to "define a family-specific prior distribution over amino acids at each sequence position that is independent from the particular protein sequence". It is the best performing baseline method.
F.2 Methodology
We run all baselines using the same methodology as Notin et al. [11, 14], with two exceptions:
1. Homologous Sequence Retrieval Method We explored using two methods, the ProteinGym and ColabFold protocols, for retrieving the homologous sequences used by baselines (Appendix C, Table 4), and compare all methods using the best homologous sequence retrieval method (Table 1).
2. Site Independent Model We present results for the site independent model using sequence weighting (Appendix D.5) and filtering sequences with more than 0.8 (in terms of sequence identity) dissimilarity to the target sequence. These are common preprocessing steps used to improve the performance of other models for protein variant fitness prediction (Appendix D), and are the same parameters used to define the retrieval-based site independent prior in Tranception [11] and TranceptEVE [14]. We believe that this variation of the site independent model better reflects the model class's ability to predict protein variant fitness.
G Ensembling PoET with Baseline Methods for Protein Variant Fitness Prediction
To ensemble PoET with a baseline method for protein variant fitness prediction, we compute a weighted sum of the score predicted by PoET and by the baseline method i.e.
$$
\hat {F} _ {\text {P o E T + B a s e l i n e}, i} = \hat {F} _ {\text {P o E T}, i} + \alpha \times \hat {F} _ {\text {B a s e l i n e}, i} \tag {5}
$$
where  $\hat{F}_{\mathrm{Baseline},i}$  is the score predicted by the baseline method,  $\hat{F}_{\mathrm{PoET},i}$  is the score predicted by PoET, and  $\hat{F}_{\mathrm{PoET} + \mathrm{Baseline},i}$  is the ensemble prediction score.
The weight  $\alpha$  is determined by optimizing performance over the ProteinGym validation set using ternary search over the range [1e-3, 1e4].
Table 7: Average Spearman correlation between model scores and experimental measurements on ProteinGym for various ensembles of PoET and baseline methods. Performance of TrancptEVE L, the best baseline method, and PoET alone are also shown for reference.
Model A	Model B	Val	Substitutions	Indels
TrancaptEVE L	N/A	0.470	0.471	0.466
PoET (ens., ProteinGym Protocol)	N/A	0.463	0.474	0.498
PoET (ens., ColabFold Protocol)	N/A	0.481	0.484	0.510
PoET (ens., ProteinGym Protocol)	GEMME	0.477	0.488	N/A
EVE (ens.)	0.485	0.487	N/A
Tranception L (no retrieval)	0.471	0.473	0.512
MSA Transformer (ens.)	0.476	0.472	N/A
Tranception L	0.480	0.475	0.510
TrancaptEVE L	0.488	0.486	0.504
PoET (ens., ColabFold Protocol)	GEMME	0.487	0.496	N/A
EVE (ens.)	0.495	0.493	N/A
Tranception L (no retrieval)	0.485	0.484	0.525
MSA Transformer (ens.)	0.487	0.484	N/A
Tranception L	0.489	0.483	0.525
TrancaptEVE L	0.494	0.492	0.521
H Effect of Architecture on Protein Variant Fitness Prediction
Figure 11 compares the performance of the regular RoPe-based Transformer and PoET on the ProteinGym validation benchmark. We use homologous sequences retrieved using the ColabFold protocol, and explore various max similarity thresholds and context lengths as in Appendix D.2. PoET consistently outperforms the Transformer. It also extrapolates to context lengths longer than the training context length (8K), whereas the Transformer does not.
(a) Performance using regular Transformer.
(b) Performance using PoET.
Figure 11: Performance of regular RoPE-based Transformer and PoET on ProteinGym validation benchmark with homologous sequences retrieved using the ColabFold protocol and various max similarity thresholds and context lengths.
I Comparison of Models for Generating Sequences From a Protein Family
To benchmark PoET's ability to generate sequences from a protein family, we evaluated the mean perplexity of a protein sequence conditioned on a fixed number of tokens from other sequences in the same protein family across all "Diamond-UniRef50 clusters" (Appendix B) in the validation set. PoET significantly outperforms two baseline models, a profile HMM and a regular Transformer (Figure 12).
Figure 12: Comparison of the perplexity of a profile HMM, regular Transformer trained with 8K context length, and PoET trained with 8K context length 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. The profile HMM is trained on MSAs inferred using mafft [47].
PoET effectively models indel rich protein families The UniRef50 clusters in the validation set for perplexity evaluation are highly indel rich. An alignment for an exemplar cluster shows that there are large insertions and high diversity between sequences (Figure 13a). A histogram of the columns by percent gap shows that highly gappy columns are the most common (Figure 13b, left). Across all of the validation clusters (Figure 13b, right), this trend remains true, with alignments becoming
more indel rich when more sequences (longer context lengths) are considered. PoET achieves low perplexities on these heldout sequences, showing that it is a good generative model of indel rich families.
(a) Typical example of a MSA of a UniRef50 cluster. Bar plot above the MSA shows the percentage of gaps in each corresponding MSA column.
(b) (Left) Distribution of percentage of gaps in MSA columns for the example MSA above. (Right) Same distribution for all MSAs of UniRef50 clusters in the validation set with  $< 4096$  tokens (top) and  $< 61440$  tokens (bottom). Larger MSAs have more novel indels.
Figure 13: MSAs of UniRef50 protein clusters reveal that these families contain many novel indels
J Inference Speed of PoET vs TrancaptEVE L
In the section, we compare the time it takes for PoET and TrancaptEVE to predict the fitness of a set of protein variants on A100 GPUs. We begin by outlining the major steps required to perform inference for each method. For PoET, there are two main steps:
1. Retrieve homologs of the target protein using the ColabFold protocol
2. Compute the conditional log-likelihood of each variant given subsamples of the homologs. Do this for 15 subsamples and ensemble the results.
For TrancaptEVE, there are five main steps:
1. Retrieve an MSA of homologs of the target protein using the ProteinGym protocol
2. Use the MSA to compute the site independent log prior
3. Train 5 EVE models on the MSA
4. Use the 5 EVE models to compute the EVE log prior
5. Compute the likelihood of each variant by combining scores from the Trancection protein language model, the site independent log prior, and the EVE log prior
Since TranceptEVE requires training new models and PoET does not, PoET can score many variants before TranceptEVE even finishes training. For simplicity of analysis, we establish a lower bound on the number of variants that can be scored by PoET before a single variant can be scored by TranceptEVE by computing the number of variants that can be scored by PoET in the time it takes to train one of the EVE models required for TranceptEVE. The homolog retrieval time is irrelevant for this analysis since both methods require a homolog retrieval step, although we note that the ColabFold
protocol used to retrieve homologs by PoET is substantially faster than the ProteinGym protocol used by TrancaptEVE.
We find that independent of sequence length, on a single A100 GPU, PoET is able to score approximately 50,000 variants in the time it takes to train one EVE model. Furthermore, since variant scoring with PoET can be parallelized whereas training an EVE model cannot, inference with PoET can be sped up further by a factor of the number of A100 GPUs available. This allows up to hundreds of thousands of variants to be scored by PoET with a reasonable amount of compute before TrancaptEVE is able to score a single variant.
K Function-specific variant fitness prediction and sequence generation via prompt engineering
K.1 Function-specific variant fitness prediction
Thus far, we have evaluated the ability of PoET to predict the "general" fitness of a protein variant, where "general" fitness can refer to any property that is related to the function of the protein. But in practice, we are generally interested in optimizing specific properties of a protein. Since many properties of interest are correlated and together contribute to the general fitness of a protein (e.g. a more thermostable variant is also more likely to have higher expression or enzyme activity), optimizing general fitness is likely to optimize the true properties of interest, albeit indirectly. PoET is a useful tool for protein engineering because it is able to successfully predict general fitness by conditioning on and inferring evolutionary constraints from a diverse set of homologs of the target protein. This raises a natural question - can PoET be used to optimize specific properties of interest by conditioning on only the subset of relevant homologs that are known to or are predicted to display the specific properties of interest? It is difficult to answer this question in its full generality, since the process of selecting the subset of relevant homologs must be tailored to the specific target protein of interest, and there is an inevitable trade-off between learning from a smaller set of homologs that are more relevant, and a larger set of homologs from which there is more data to learn.
As a proof of concept, we demonstrate that PoET is able to learn function-specific evolutionary constraints for the target protein and property of interest in the chorismate mutase indels dataset [9] from ProteinGym. The dataset contains measurements of the catalytic activity, in E. coli, of 1130 natural chorismate mutase sequences, and 1618 designed chorismate mutase variants. The natural sequences are comprised of the target protein, a chorismate mutase found in E. coli, and homologs of the target protein found using the PSI-BLAST program for sequence search. The designed variants were selected by Monte Carlo sampling from a Potts model trained on an MSA of the natural sequences. This data presents an ideal scenario for selecting the subset of most relevant homologs; we can simply select the subset of natural homologs that are measured to be functional. In the absence of such data, we could instead use predictions from another model, or other relevant known attributes of the sequences e.g. if we would like to optimize for activity at high temperatures, we may select only the homologs from thermophiles.
On the chorismate mutase dataset, we find that the catalytic activity of designed chorismate mutase variants is better predicted when PoET is conditioned on only the subset of functional natural sequences rather than all natural sequences ( $\Delta \rho = 0.2$ , Figures 14b, 14d). In fact, PoET conditioned on functional natural sequences even outperforms fully supervised methods, including a Gaussian process trained on mean embeddings from a BERT-like protein masked language model ( $\Delta \rho = 0.06$ , Table 8). Such embeddings have been shown to be highly predictive of a variety of protein properties [20, 48, 49] and provide a strong baseline. These fully supervised methods are trained on more data than PoET because they train on the measured catalytic activities of all the natural sequences, whereas PoET is simply conditioned on positively labeled natural sequences and does not have access to the measured activities. This enables PoET to be used with assays that only measure binary endpoints rather than continuous values.
Table 8: Spearman correlation between measured and predicted fitness of designed chorismate mutase variants.
Method	Spearman Correlation
Ridge regression trained on one hot encodings	0.45173
Gaussian process trained on mean embeddings	0.53658
from a BERT-like masked protein language model
PoET conditioned on functional natural sequences	0.56772
(a) Predicted fitness using Potts model trained on natural sequences from Russ et al. [9].
(b) Predicted fitness using PoET conditioned on natural sequences. Natural sequences were retrieved using PSI-BLAST to search for homologs of the target protein [9].
(c) Predicted fitness using PoET conditioned on homologs of the target protein retrieved by the ColabFold protocol.
Figure 14: Measured fitness vs predicted fitness of (left) natural and (right) designed chorismate mutase enzymes from Russ et al. [9].
(d) Predicted fitness using PoET conditioned on functional natural sequences. Natural sequences with measured fitness  $>0.42$  are considered functional by Russ et al. [9].
K.2 Function-specific sequence generation
Conditioned on the functional natural sequences, we used PoET to generate 1000 novel putative chorismate mutases using nucleus sampling [50] with  $p = 0.9$ .
Figure 17 visualizes the sequence space spanned by natural chorismate mutases, chorismate mutases functional in E. coli, and sequences sampled from PoET when conditioned on the latter sequences. We find that the sequences sampled from PoET are concentrated in clusters around the known functional sequences, suggesting that the generated sequences sample from the functional subspace. The generated sequences are highly diverse and not simply a recapitulation of the natural sequences; the maximum sequence identity to any natural sequence is between 0.4 and 1.0, with a mode around 0.55 (Figure 15a). Nevertheless, the entropy at each position of an MSA of the functional natural sequences and an MSA of the generated sequence are closely matched (Figure 15b), indicating that the generated sequences reflect at least the first order evolutionary constraints of functional chorismate mutases.