PoET: A generative model of protein families as sequences-of-sequences
Authors: Timothy Fei Truong Jr
Tristan Bepler
Venue: NeurIPS 2023 poster (2023)
Keywords: protein fitness prediction
transformer
retrieval
language model
MSA
generative model
protein engineering
TL;DR: PoET is a generative model of protein families as sequences-of-sequences. It accurately predicts variant effects and can be sampled from to generate new sequences.

Abstract: 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 **P**r**o**tein **E**volutionary **T**ransformer (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.

PoET: A generative model of protein families as sequences-of=sequunces
Timothy Truong Jr
Tristan Bepler
OpenProtein.AI NY, USA ttruoog@openprotein..ai
OpenProtein.Al NY, USA
tbeplerropeeprooeiin..
Abstract
Generative protein language models are natural way to design new proteins with desired functions. However, current models either difficult to direct produce a protein from 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), autoregressive generative model of whole protein families that learns to generate sets of related proteins as across tens of millions of natural protein sequence clusters. PoET can be used as a retrieval-augmennte language model to generate and score arbitrary modifications conditioned any protein family of interest, and 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 scale context lengths beyond those used during training. 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.
sequences-of-sequeeee
Introduction
Proteins carry out most of the biological functions the molecular level of life. The function of a protein is encoded by its specific amino acid sequence and the dimenniinnsional structure that the sequence folds into. Engineering proteins for novel and enhanced function is 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 search space, reduce the need for expensive experiments, and enable the design of more novel proteins and functions.
relationships can narrow down the protein
Protein language models have emerged 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 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]. 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].
Code and pre-trained model weights are available at https: //github/ htt.e/hhb.ccoo/pppeeppteeeea//aa
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Traditionally, family-specific models learn evolutionary constraints specific to the protein family of interest by training on 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 specific family of interest, and underperform models for relative fitness prediction. Hybrid models such as Tranception [11] and TranceptEVE [14] combine unconditional language models with family-speciio protein families. Nonetheless, it is unclear how to use these models to generate sequences with novel indels, and predictions from the family-specifie models do not directly benefit from transfer learning across families.
models to enable specialization to
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 across tens of millions of natural protein sequence clusters, PoET
proteins as sequences-of-seqqueeee 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, propose novel Transformer layer ($3.1.2) that models grder-dependence between tokens within sequences and order-independence between sequences. PoET has the following properties:
PoET can be used as retrieval- augmeeted protein language model by conditioning the model on sequences from any family of interest. This also allows PoET to used with any sequence database and to incorporate new sequence information without retraining.
PoET 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 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 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.
2 Related Work
Evolutionary sequence models 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-bae bassed" and must be fit on 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.
2

---

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. contrast family-specific evolutionary sequence models trained on sequences derived from 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 Samily-specific models for larger families.
Conditional protein language models fit between the unconditional and family-specific paradigms. Only few works have explored this direction 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 al. [11] combine predictions from 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 still limited by the constraints of MSAs.
alignment-basse model means that the combined model is
Retrieval-augmentte language models have shown impressive results in natural language pro- cessing, especially on Question Answering tasks. These models incorporate database search as part of the text generation process in order to generate new text conditioned the database [26-32]. In this way, they are conceptually similar to the conditional protein lan- guage models above. Retrieval-augmenter 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 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 sequence- lequence-of-sequenee ofquennss problem to incorporate retrieved-sequencce conditioning, a fundamentally different paradig gm than that employed by current retrieval- augmentee models in natural language processing.
prototypes found
The Protein Evolutionary Transformer (PoET)
PoET is autoregressive generative model of the distribution over protein families, where each family is generated as sequence-of-seqqeenes. Specifically, models the distribution P(X = x). where 81, 82, ..., 8n is the concatenation of n sequences from the same family, and each sequence 8; = 81,1, 81,2 $s,L; sequence of L, amino acids padded by start and end token. For example, below is sequence of three protein sequences of lengths 4, 6, and with start token denoted by $ and stop token denoted by *:
M
K
$
M
H
$
M
P
V
811
812
813
814
521
822
823
824
825
826
831
832
833
834
835
S1
$2
$3
When referring to sequence-of-sequuenes 8,, which has one index i, refers to a sequence of tokens the ith sequence in the sequence-of of-eeunnecssss whereas $i,j, which has two indices i, refers one token, the jth token of the ith sequence. We use x to denote the full sequence-of-sequeeee of-seuuencss.......
PoET generates each token in sequence one time, decomposing the probability of r:: as
P(x) P(s1, $2 s,) II P(s
IIII P P(s.) 8i,z))
(I)
(=1==1
The order of the individual sequences sequence-of-sequenee arbitrary, and we propose a novel caansformer-bssee architecture to exploit this order invariance.

---

Py pasll (y
Phi
Piy
Pia
PIA
Pb1
Pha Pho Ps.
PMs
Pin
Ph
Ph
PAp
PM
Pis
Decoded Probabilitie
Sequence Suguonoes
fu
TTereTrransoomrrDsoodedde
Seo-Seeuunn Attention
Caueal Set--tttttt maatt-ttettoo
Caussi Get-Adttrtion
Embeddings
11
Inputxx
(a) PoET generates sets homologous proteins as sequeeeesss-- The equation demonstrates how to compute the conditional probability a sequence given homologous sequenees.
0
91
Tokens Associated
St:
$
K
With Position
8
52:
$
H
2
2
92
Poettion
Keys
H
(b) Inter-sequence relative position encoding scheme used in sequence-of-seqqeeees self-attention that is invariant sequence ordering. (Left) Relative positions hetween tokens in sequeniceoo-sequenees composed of two sequences. (Right) Tokens associated with absolute positions.
Figure 1: PoET Architecture
3.1 Model Architecture
We propose variant of the common Transformer decoder layer (Figure la) capture order invari- ance between sequences while preservering order-dependence between tokens within sequences. We accomplish this using tion at each position of each sequence is updated based on attending only to the other tokens within this sequence, and (2) between-sequencee module in which the representation each position of each sequence is updated based on attending to other sequences within the sequence-of-sqquenee 1.3...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 ($5.2.2), improving sequence generation and performance downstream tasks (Appendix H). Our full PoET model is stack of these layers with causal self-attention.
attention modules: (1) within-sequence module in which the representa-
3.1.1 Input Embedding
The input sequence Si.j. i L.m.j j E L..L, amino acids and start/stop tokens is first converted into a sequence of continuous embeddings h4.5 by mapping each token to its learned embedding:
ha = We.,y $i. E AA U EETART,STOR
(2)
W E IR [AAL (START,STOP) STOP [xd

---

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 ha.j are transformed by N layers of oieereefrannsooorrrreeeooeeeeeee (Appendix Algorithm 1), novel layer for processing sequence-of-sequeneee that invariant the order of the individual sequences, and extrapolates to context lengths substantially longer than the training context length.
The TieredTransf ormerrecoderlayer is composed of two phases. In the first phase, causal self-attention is applied independently to each sequence h, of the input sequence-of sauuence-of-seqeeness transforming them into new sequences fi PerSequenceSell Attn h2)) Relative positional information is encoded by applying Rotary Positional Encodings (RoPE) [33] the queries and keys before applying self-attention in the standard manner; the absolute position for Ji,j is
The second phase applies causal self-attention the entire sequence-of-sequencee by concate- nating the individual fi from the previous layer into one sequence before applying self-attention: gij SeeuencceffSequenceessee (ttreeeneeeeeeeet(((ffeeee fi,<y)). In order 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 gi.j the absolute position is j. Just as in the first phase, the absolute position for tokens in the ith sequence is independent of the position i of the sequence in the sequence-off 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 the first sequence S1,1 and the first token of the second sequence f2, is O. 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:
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", rather than the number of tokens in a sequence-of-sequeneee allowing the model to generalize to longer sequences-of-sequeecce than seen during training
3.1.3 Decoded Probabilities
Lastly, the output from the last orrredrrranfrrmmmreoddeeeeeeee probabilities by applying linear transformation P(st.y $1,45) (s.,( Linear (gr.)) Here, Pt. is a vector of probabilities, one for each distinct token E AAU AAU (START, STOP), and pc, ($1,) the probability of the token $1. according Pt.j-
9i.j, is decoded into token
3.2 Training Data
Models were trained on 29 million sets of homologous sequences. Each set corresponds to 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 a large number of sets, sample each 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 promote order invariance. See Appendix B for more details.
4 Protein Variant Fitness Prediction
Protein variant fitness prediction is the task of assigning score to each sequence in a set of variants {v1, U2, )x of a target sequence that accurately reflects the relative fitness of the variants. A protein variant us can be any sequence with limited number of substitutions, insertions, and/or deletions relative to the target that the experimenter believes may have improved fitness. Fitness
AWhile the length of protein is unbounded (but finite), in practice, the vast majority (> 99%) of proteins are less than 1024 amino acids long [11].
5

---

Retrieve Target Sequence Homologs
Evalluate correlation experimental
Homologs
Iumarce tareatas 2562/4
NP2PY 3IYAD
Sequence DB
PoET
Target Sequence
NELVS DRLAR
HAADS.
NETT?
Variants
Likelihoods
Compute likelihood of each in mutational landscape of interest
MAACO KOeeR
PNT Pronptt
Ptv, (ncoept)
NALOHO HSEPA
PIV, (Pcoopt)
MAA 22F2
PLY, Pronps]
Figure 2: Illustration of evaluating PoET for variant fitness prediction on a DMS dataset
refers to the value of any property 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 method for measuring the fitness properties thousands to millions of variants of a protein using next- gext-generation sequencing techniques [1]. Data from these experiments have been used to evaluate fitness prediction methods. We evalaute 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 al. [11] for tuning hyperparameters.
Fitness Prediction with PoET PoET predicts the fitness of variant the conditional log- likelihood of the variant os given a set of sequences homologous to the target I:
F((S (s....,, 8m)) log P(e./s,, ....... 8m) Σ log P(mry(s1, Sms 0,,4))
The set is retrieved by searching large database of proteins such as UniRef100 [23] for sequences homologous to t. The homologs form 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 conditioned in arbitrary order. Figure 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) ensembling conditional log-likelihoods computed from different subsamples of homologous sequences. We present results with the best settings in the main results ($5). Our final approach uses the ColabFold protocol [35] for retrieving ho- mologs, and the final fitness prediction score is obtained by averaging the conditional log liceeiosss across subsamples of the full set of retreived homologous sequences:
Nivembie
f Pensemble, (S)
RASS, C
Nensemble
The subsets Sj are drawn by sampling with sequence weights [17], maximum sequence identity to the target in (1.0, 0.95, 0.90, 0.70, 0.50), and total tokens (i.e. context length) (6144, 12288, 24576). All combinations of these parameters used, for total of N Nansemble
15.

---

Table 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 dash (-) indicates applicable, not computed.
Substitutions by MSA Depth
Model Type
#
Model name
Param
Low
Medium
High
All
Indels
Align-
Site independent
N/A
0.417
0.404
0.411
0.408
N/A
ment-
GEMME
N/A
0.445
0,449
0.522
0.463
N/A
based
EVE (ensemble)
N/A
0.414
0.441
0.498
0.448
N/A
Uncond-
ESM-Iv (ensemble)
3,25B
0.356
0.372
0.510
0.398
N/A
itional
ProGen2 (ensemble)
10.8B
0.357
0.416
0.448
0.411
0.407
PLM
Tranception (no retrieval)
700M
0.377
0.399
0.429
0.401
0.430
Cond- itional
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
Tranception L
700M
0.441
0.437
0.472
0.445
0.464
TranceptEVE VE M
300M
0.516
Hybrid
TranceptEVE
700M
0.454
0.463
0.508
0.471
0.466
PoET (ensemble) + TranceptEVE:
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 between the predicted fitness and the measured fitness across the DMS datasets 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 paired t-test. For a fair comparison, we ran the relevant baselines with all homologous sequence retrieval methods considered for with PoET (Appendix C), and present the results for each baseline method paired with its best retrieval method. We do not tune parameters for ensembling 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 ensembling strategy as defined by their authors [4, 12, 14].
5.1 Comparison to baselines
Substitutions Benchmark On the substitutions benchmark, PoET performs comparably Tran- ceptEVE L,, the best performing baseline method (Table 1). PoET improves average Spearman correlation Ap 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 substantially faster when considering on the order of tens 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 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 TranceptE cannot, is less critical that the protein family of interest defined by large number of homologous sequences (Appendix C).
7

---

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 TranceptEVE (Results: Table 1, Methods: Appendix G). An ensemble of PoET with TranceptEVE L. performs sig- nificantly better than TranceptEVE: alone (Ap = 0.021, < 7e-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 TranceptEVE L on the substitutions benchmark, and requires very little additional time to compute; GEMME is thousands of times faster than either PoET or TranceptEVE [12]. The main disadvantage is that GEMME is unable to score indels.
Indels Benchmark On the indels benchmark, the best performing baseline model is TranceptEVF M, a variation of TranceptEVE L. that uses a protein language model with fewer parameters (Ap = 0.05). PoET and the ensemble of PoET with TranceptEVE L both outperform TrancepLEVE (Ap [0.044, 0.055)), and perform comparably to TranceptEVE M (Ap E [-0.006, 0.005J). 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 -16 of possibly thousands of homologs. Interestingly, while context length of 8K performs better than 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.
Lastly, using the optimal parameters found thus far, trained 57M, 201M, and 604M parameter models to explore the effect of model size. Models were trained until validation perfor- mance plateaued. Increasing the model size 201M parameters signficantly improves performance, but increasing the size further to 604M appears to have neutral effect. We chose the 201M version as the standard PoET model.
Model Size
5.2.2 Model Architecture
To determine the importance of using TeereerTranssfforrrrreeeeee....e regular Transformer layers, we trained, on the same task, a regular 201M parameter RoPE-based Transformer that ignores the sequence-of-scquences structure of the input. We compared the genera- tive capabilities of PoET and the Transformer by measuring the perplexity at context lengths both within and beyond the training the context length (Figure 4).
($3.1.2) instead of
We trained both models with a context length of 8K for 500K steps. As expected based on previous (38) the RoPe-based Transformer does not generalize to context lengths beyond the training context length. In contrast, PoET generalizes to ~8x the training context length (61 K), which the longest we tested. PoET also performs better context lengths within the training context length,
8

---

Parass
Contect
Length
Weighting
Speaman 110
S7M
Inwerse Count
52M
4,896
Inwerse Count
57M
16,354
inverse Count
0.44
57M
8,102
Uh1fomn
231
8,102
Inverse Count
28
72
6
48
96
164
192 Training
260
268
316 354 432
Talning
Figure 3: Performance of PoET on the ProteinGym validation set when trained with (left) various training distributions and (right) model sizes.
Perplesity
7.0
Architecture
Regular Transformer
PoET Transtormer
Context
10K
32K
4096
8192
18384
32768
6140
Inference Context Length (logsscale)
Figure 4: Comparison of the perplexity of 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 profile HMM baseline (Appendix Figure 12).
but by smaller margin. Next, finetuned both models with 16K and 32K context lengths for 50K and 30K steps respectively. While the perplexity of the Transformer improved significantly 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).
POET generates novel sequences that preserve structure within a protein family
PoET's state-of-the-ar performance on variant fitness prediction demonstrates that 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 variants makes the exploration of higher order mutants computationally tractable as sequence space grows exponentially large with number of mutations, making 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 the max sequence identity to any natural sequence (of the same protein family). Predicted

---

A: Functional Chorismate Mutases
B: Phage Lysozymes
100
More
Confident
p.zor
Prediction
00
0.6
Max 1O ratural recuerce
Max sec id natural ORQUeECe
Model
imgen
ReT
1.0
TM compands
Higher Structural
0.00
Mar TRISOOR to af aaturi
Similarity
0.75
0.6
0.7
0.6
10
Max seQ natural sequence
More Novel Sequances
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 structure of natural sequence and pL.DDT > 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 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), 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.
7 Conclusion
PoET is a Transformer-basee autoregressive generative model of whole protein families. By framing family generation as a sequence-of-sequeene of millions of protein sequence clusters 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 retrieval-augee -augmentedd 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 fundamental component of ML-enabled protein design in the future.
generation problem, we are able to train across tens
10

---

References
[1] Douglas M. Fowler and Stanley Fields. Deep mutational scanning: a new style of protein science. Nature Methods, 1188):01-807, Aug 2014. ISSN 1548-7105, doi: 10.1038/mmeth.30277 URL https://don ottp/i...rr1.... 20003//mmme..000
[2] Ryan E Cobb, Ran Chao, and Huimin Zhao. Directed evolution: Past, present and future. AIChE
J. 595):1432--144, May 2013.
[3] Louis Tiefenauer and Sophie Demarche. Challenges in the development of functional assays of membrane proteins. Materials (Basel), 5(11):2205-2242, November 2012.
[4] Joshua Meier, Roshan Rao, Robert Verkuil, Jason Liu, Tom Sercu, and Alex Rives. Lan- guage models enable zero-shot prediction of the effects of mutations on protein function. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 29287-29303. Cur- ran Associates, Inc., 2021. URL https: httpttss:/pprooeedinnnn wwt.ig..neuiis.. Cearrpss.ccpappr/0000000 cc/papper/202///22e 153388r76600d222277777777---- Paper pdf.
[5] Erik Nijkamp, Jeffrey Ruffolo, Eli N. Weinstein, Nikhil Naik, and Ali Madani. Progen2: Exploring the boundaries protein language models, 2022. URL https://arxiv.orr arbaas// 2206. 13517.
[6] Tristan Bepler and Bonnie Berger. Learning the protein language: Evolution, structure, and function. Cell Systems, 12(6))654 669.e366.e3 2021. ISSN 2405-4712. doi: https:/doi.org/ 10. 00.1016f..els.2222.000000 URL https: //ww./w.ww/ooooo/// sciencedirect, www/sciince//aaeeeecaoooaae/eeee article/ pi1//2240647222100033.
[7] J. Dauparas, I. Anishchenko, N. Bennett, H. Bai, R. Ragotte, L. F. Milles, B. I. M. Wicky, A. Courbet, R. de Haas, N. Bethel, P. J. Y. Leung, F. Huddy, S. Pellock, D. Tischer, F. Chan, B. Koepnick, H. Nguyen, A. Kang, B. Sankaran, A. K. Bera, N. P. King, and D). Baker. Robust deep 200
protein sequence design using proteinmpnn. Science, 378(6615): 49-56, 2022. doi: 10.1 wttseiiecc..add2221 URL https: http::/wwww..ieeeeee org/doi/ /abs/ 10. 1tl1126/ssccin adt1//ecceceeeeec.11.... add2187.
[8] Ali Madani, Ben Krause, Eric R Greene, Subu Subramanian, Benjamin Mohr, James M Holton, Jose Luis Olmos, Jr, Caiming Xiong, Zachary Z Sun, Richard Socher, James S Fraser, and Nikhil Naik. Large language models generate functional protein sequences across diverse families, Nat. Biotechnol., January 2023.
[9] William P. Russ, Matteo Figliuzzi, Christian Stocker, Pierre Sarrat-Charlaix, Michael Socolich, Peter Kast, Donald Hilvert, Remi Monasson, Simona Cocco, Martin Weigt, and Rama Ran- ganathan. An evolution-based model for designing chorismate mutase enzymes. Science, 369 (6502):440-445, 2020. doi: 10.1 22///cienc..aaa...... URL https://wwwtt wwwp:wwwcinnc..e.. org/ doi/abs/10. 100 112/111/////ee/ science..ccccccccc aba3304.
[10] Daniel Hesslow, Niccoló Zanichelli, Pascal Notin, Iacopo Poli, and Debora Marks. Rita: a study on scaling up generative protein sequence models. arXiv preprint arXiv:2205.0577 055899 2022.
[11] Pascal Notin, Mafalda Dias, Jonathan Frazer, Javier Marchena Hurtado, Aidan N Gomez, Debora Marks, and Yarin Gal. Tranception: Protein fitness prediction with autoregressive transformers and inference-time retrieval. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 16990-17017. PMLR. 17-23 Jul 2022. URL https: https::/proceeeee pro.eeeiceeennneeeooe mlr. .presa/ v162/notin22: html.
[12] Elodie Laine, Yasaman Karami, and Alessandra Carbone. GEMME: A Simple and Fast Global Epistatic Model Predicting Mutational Effects. Molecular Biology and Evolution, 36 (11):2604-261, 08 2019. ISSN 0737-4038. doi: 10. 1093/mollee/mmm2179, URL https:
//doi. org 1093 htt.:/9mmmmm////s////ie. molbev/ms2179.

---

[13] Jonathan Frazer, Pascal Notin, Mafalda Dias, Aidan Gomez, Joseph K. Min, Kelly Brock, Yarin Gal, and Debora S. Marks. Disease variant prediction with deep generative models of evolutionary data. Nature, 599(7883):91-95, Nov 2021. ISSN 1476-4687. doi: 10. 10310 s4 1558-021-00403-8 URL URI. wtt.ttttp///oi.... orgg 1038/ s4 586 020-118000222100------
[14] Pascal Notin, Lood Van Niekerk, Aaron W Kollasch, Daniel Ritter, Yarin Gal, and Debora Susan Marks. TranceptEVE: Combining family-specific and family-agnostic models protein sequences for improved fitness prediction. In NeurIPS 2022 Workshop on Learning Meaningful Representations of Life, 2022. URL https: hppp:////ppeeeie//e http://pppneeeeeeeeeeeeeooeeeeee wwt..e/oooommddd//dooeeeoeeeooee
[15] P Ng and S Henikoff. Predicting deleterious amino acid substitutions. Genome Res., 11(5):
863-874, May 2001.
[16] R Eddy. Profile hidden markov models. Bioinformatics, 14(9):755-7633 1998.
[17]
Thomas A. Hopf, John B. Ingraham, Frank J. Poelwijk, Charlotta P. I. Schärfe, Michael Springer, Chris Sander, and Debora S. Marks. Mutation effects predicted from sequence co-variation. Nature Biotechnology, 35(2):228-135, Feb 2017. ISSN 1546-1696. doi: 10.1038/nb..376.
URL https: org/ 10. http///bb..66999
[18] Tristan Bepler and Bonnie Berger. Learning protein sequence embeddings using information from structure. In International Conference on Learning Representations, 2019. URL https:
/oopnreview
tttr:/oorrrmmttSSyyeecttttt d-SygLehCqtm
[19] Ethan C. Alley, Grigory Khimulya, Surojit Biswas, Mohammed AlQuraishi, and George M. Church. Unified rational protein engineering with soquence-basee deep representation learn- doi: 10.1038/ s41592-019-0598-1. URL https: http://doo.ooooo orrgiiorroo... 10. 1038/ 011-992---9-099------1
ing. Nature Methods, ((12)):3155-122) Dec 2019. ISSN 1548-7105.
[20] Roshan Rao. Nicholas Bhattacharya, Neil Thomas, Yan Duan, Xi Chen, John Canny, Pieter Abbeel, and Yun S Song. Evaluating protein transfer learning with TAPE. bioRxiv, June 2019.
[21] Adam J. Riesselman, John B. Ingraham, and Debora Marks. Deep generative models of genetic variation capture the effects of mutations. Nature Methods. 15(10):816 822, Oct 2018. ISSN 1548-7105. doi: 10.1038/s41599 00044411999-91--11 018-0138-4. URL attps://doi. org/10. 1038/
$41592- 018- 018--13--4--111111
[22] Jaina Mistry, Sara Chuguransky, Lowri Williams, Matloob Qureshi, Gustavo A Salazar, Erik L Sonnhammer, Silvio C E Tosatto, Lisanna Paladin, Shriya Raj. Lorna J Richardson, Robert D Finn, and Alex Bateman. Pfam: The protein families database in 2021. Nucleic Acids Research, (9(D)):D422DDD) D419, 10 2020. ISSN 0305-1048. doi: /0.10///ar//kka1113... URL
https://doi. http://oo..oo/////// org/10. 1093/nar /gkaa913.
[23] The UniProt Consortium. UniProt: the Universal Protein Knowledgebase in 2023. Nucleic Acids Research, 51(D1))D523-D533 11 2022. ISSN 0305-1048. doi: 10. 10.1093/nar/gkac 1052.
URL https://doi. org/10. 1ww///a0/////// 100////ar//aaaa0022222.
[24] Roshan M Rao, Jason Liu, Robert Verkuil, Joshua Meier, John Canny, Pieter Abbeel, Tom Sercu, and Alexander Rives. Msa transformer. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 8844-8856. PMLR, 18-24 Jul 2021. URL https:// proceedings mlr. press/v139/raoll11.. html.
[25] Soumya Ram and Tristan Bepler. Few shot protein generation, 2022. URL https://arxiv. org/abs//2004.111 01168.
[26] Kelvin Guu, Tatsunori B. Hashimoto, Yonatan Oren, and Percy Liang. Generating sentences by editing prototypes. Transactions of the Association for Computational Linguistics, 6:437-450, 2018. doi: 10. 10.1162/tacl a_00030. URL https://aclaathool htw.aaanthoooooooor//11110. wrg//18-10300
[27] Tatsunori B Hashimoto, Kelvin Guu, Yonatan Oren, and Percy S Liang. A retrieve-and-edee framework for predicting structured outputs. S. Bengio, H. Wallach, H. Larochelle, K. Grau- man, N. Cesa-Bianchi, and R. Gamett,, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL htt.:::tt:///ptt http:ttte///rrooediiiggee ingruuuisnggs..uuripp..... neurips. cc //0ppe//218/fiiiii
Paper pdf.
12

---

[28] Omar Khattab and Matei Zaharia. Colbert: Efficient and effective passage search via con- textualized late interaction over bert. In Proceedings of the 43rd International ACM SIGIR Conference on Research and Development in Information Retrieval, SIGIR '20, page 39-48, New York, NY, USA, 2020, Association for Computing Machinery. ISBN 9781 4503800164. doi: 10.1145/3397271 0t1./33377771.441177 1.301075.
URL attpss///oo 10rg//1..10 11.11.1113339771 3401075.
[29] Omar Khattab, Christopher Potts, and Matei Zaharia. Baleen: Robust multi hop reasoning at scale via condensed retrieval. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL https: / //openreview..om net/ htt/foormmiio di-mmiiichhhhhhh.......
[30] Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Ming-Wo Wei Chang. Realm: Retrieval-augmented language model pre-training. In Proceedings of the 37th International Conference on Machine Learning, ICMI'20. Jor..org. 2020.
[31]
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George Bm Van Den Driessche, Jean- Sann-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, Diego De Las Casas, Aurelia Guy, Jacob Memick, Roman Ring, Tom Hennigan, Saffron Huang, Loren Maggiore, Chris Jones, Albin Cassirer, Andy Brock. Michela Paganini, Geoffrey Irving. Oriol Vinyals, Simon Osindero, Karen Simonyan, Jack Rae, Erich Elsen, and Laurent Sifre. Improving language models by retrieving from trillions of tokens. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 2206-2240. PMLR, 17-23 Jul 2022. URL https:// proceedings... www.inngsmmm........mm.........m wttp:/sm.r..prees/11//2 prress/v16//borreadd222 html.
[32] Gautier Izacard, Patrick Lewis, Maria Lomeli, Lucas Hosseini, Fabio Petroni, Timo Schick, Jane Dwivedi- Yu, Armand Joulin, Sebastian Riedel, and Edouard Grave. Few-shot Learning with Retrieval Augmented Language Models. 2022.
//axiiv.//rxv.... http://wwwv...///.///////// org/abs/////eooeooooeo..ooeea abs//208. 03299.
[33] Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2021. URL https://arxii: org/ abs/2104. 09864..
[34] Benjamin Buchfink, Klaus Reuter, and Hajk-Georg Drost. Sensitive protein alignments at tree- of-life scale using diamond. Nature Methods, 18(4):366-36) Apr 2021. ISSN 1548-7105. doi: 10.
URL https: / //toi. org/10.. 1038/ /01////119222222 021---00---1012-
[35] Milot Mirdita, Konstantin Schütze, Yoshitaka Moriwaki, Lim Heo, Sergey Ovehinnikov, and Martin Steinegger. Colabfold: making protein folding accessible to all. Nature Methods, 19 (6),679688, Jun 2022. ISSN 1548-7105. doi: 10. 00.10388441592--02-001444
URL https:
//doi.org/ 1038/s41592 2t////441222222222222 000---22----000 01488-11
[36] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert- Voss. Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Ma- teusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCan- dlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Ad- vances in Neural Information Processing Systems, volume 33, pages 1877-1901, Curran Associates, Inc., 2020. URL 3157c0d6bfbb4967777 001.99661190bb0aa122466--------2 .dff.....ff.
neurips wttp:ccppc/p///////////.
[37] Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmmann, H. Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, Eliza Rutherford, Tom Hennigan, Jacob Menick, Albin Cassirer, Richard Powell, George van den Driessche, Lisa Anne Hendricks, Maribeth Rauh, Po-Sen Huang, Amelia Glaese, Johannes Welbl, Sumanth Dathathri, Saffron Huang, Jonathan Uesato, John Mellor, Irina Higgins, Antonia Creswell, Nat McAleese, Amy Wu, Erich Elsen, Siddhant M. Jayakumar, Elena Buchatskaya, David Budden, Esme Sutherland, Karen Simonyan, Michela Paganini, Laurent Sifre, Lena Martens, Xiang Lorraine Li, Adhiguna Kuncoro, Aida Nematzadeh, Elena Gribovskaya, Domenic Donato, Angeliki
13

---

Lazaridou, Arthur Mensch, Jean-Baptise Lespiau, Maria Tsimpoukelli, Nikolai Grigorev, Doug Fritz, Thibault Sottiaux. Mantas Pajarskas, Toby Pohlen, Zhitao Gong, Daniel Toyama, Cyprien de Masson d' Autume, Yujia Li, Tayfun Terzi, Vladimir Mikulik, Igor Babuschkin, Aidan Clark, Diego de Las Casas, Aurelia Guy, Chris Jones, James Bradbury, Matthew Johnson, Blake A. Hechtman, Laura Weidinger, lason Gabriel, William S. Isaac, Edward Lockhart, Simon Osindero, Laura Rimell, Chris Dyer, Oriol Vinyals, Kareem Ayoub, Jeff Stanway, Lorrayne Bennett, Demis Hassabis, Koray Kavukcuoglu, and Geoffrey Irving. Scaling language models: Methods, analysis & insights from training gopher. CoRR, abs/21 12. 1446, 2021. URL https://arxae org/ abs 2112 12. 11446.
[38] Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, 2022. URL https: /ttp://///eeeeeeeeeeeeee beet/forum/iddR8S@PPGGS
[39] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A. A. Kohl, Andrew J. Ballard, Andrew Cowie, Bernardino Ponera-Paredes, Stanislav Nikolov, Rishub Jain, Jonas Adler, Trevor Back, Stig Petersen, David Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger, Michalina Pacholska, Tamas Berghammer, Sebastian Bodenstein, David Silver, Oriol Vinyals, Andrew W. Senior, Ko- ray Kavukcuoglu, Pushmeet Kohli, and Demis Hassabis. Highly accurate protein structure prediction with alphafold. Nature, 596(7873) Aug 2021. ISSN 1476-4687. doi: (0.1038/s415866 000--0221-3888 03819-2...... URL https: //doi org/10. 1038/ s41586-022 021-03819-2
[40] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 3604. PMLR, 10-15 Jul 2018. URL shazeer18a. html.
mlr www.mmlrr.ss.....mmmm www.ssesssveeee.
[41] L. Steven Johnson, Sean R. Eddy, and Elon Portugaly. Hidden markov model speed heuristic and iterative hmm search procedure. BMC Bioinformatics, 11(1):431, Aug 2010. ISSN 1471-2105. URL https://doi.. org/ 10. 1186/1471- 11177112105--1 11-431.
doi: 10.
[42] Martin Steinegger and Johannes Söding. Mmseqs2 enables sensitive protein sequence searching for the analysis of massive data sets. Nature Biotechmology, (5(11))1022) 10:066-108,, Nov 2017. ISSN 1546- 1696. doi: 10.1038/nbt.398.. URL /atttps://oi.... org/ 110.///////. 10w8/10t....... .bt.
[43] Vineet Sangar, Daniel J. Blankenberg, Naomi Altman, and Arthur M. Lesk. Quantitative segueecce-function relationships in proteins based gene ontology. BMC Bioinformatics, 8(1): 294, Aug 2007. ISSN 1471-2105. doi: 10. 10. 11186/177-221--88
URL. https://////o org/
[44] Eli N Weinstein, Alan Nawzad Amin, Jonathan Frazer. and Debora Susan Marks. Non- identifiability and the blessings of misspecification in models of molecular fitness. In Al- ice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https: /ttpps://ppette /pperreeiee. net/forum? id= CuG-o0ind6t.
[45] Martin Steinegger and Johannes Söding. Clustering huge protein sequence sets in linear time. Nature Communications, 9(1):254,, Jun 2018. ISSN 2041-1723. doi: 10.1038/
11447-008--49644-- URL https: //doi. 11t../....0///1144.77..... 518-44446..-
[46] Michael Remmert, Andreas Biegert, Andreas Hauser, and Johannes Söding. Hhblits: lightning- fast iterative protein sequence searching by hmm-hmm alignment. Nature Methods, 9(2): 173-175, Feb 2012. ISSN 1548-7105. doi: 10.1038/mmeth.1111 URL https://doi. org/ 10. 10.103///meth....m 1818.
[47] Kazutaka Katoh and Daron M. Standley. MAFFT Multiple Sequence Alignment Software Version 7: Improvements in Performance and Usability. Molecular Biology and Evolution, 30(4):772-780, 2013. ISSN 0737-4038.
doi: 10. 1003///mlee////m0000 URI. https:
//doi.orr//// 10. 1093 htt.1//0mmmmm///mmm///00 molbee/mstt00.
14

---

[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, 100, 02 2022. ISSN 1367-4803. doi: 100100//iionfommati///aaa/////// URI.https. https: //doi. o00//... 1093 /tt.://oorrmmtiic//aaa/........
[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://openreviie
15

---

A Detailed description of oiierrdlrrannsforrrrreeedeeraa
Algorithm TTiareedraansfoorerrbocoooaaaa
Require:
representations for a sequence of sequences li E E i I..N,j 1..4,
First, apply causal self-attention to each sequence individually ʃ + LayerNorm(h)
4b.j RoPE(Linear Vi,j Next, apply causal self-attention all sequences together
k f- (oPE(Lineer((f.):)) V
f- Linear(f)
fr 1 fi + CausalAttention((;: k;, v.) Vi
9i, t fijj Vi.j
LayerNorr(()
4. f RoPE(Linear (e V
j f- (oPE(Linear(().)]) Vi,j
- Linear(g)
+ g CausalAttentionn(..
Finally, the feedforward layer
g' - LayerNorm(
h' + GELU((Lneaar(')) h E IRdd
g' + 9 + Linear(h")
output g' e Rd
B Additional Training Details
B.1
Data
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 diamon//uniref50 -f 6 -header -k 200000 -max-hsps 1 -e 0.001 -p 96 output. tab
The command returns, for each sequence UniRef50, a set containing all its putative homologs in UniRef50. We call each such set "Diamond- UniReff00 Cluster". Diamond was used over other homology search tools due to high performance (>100x speed of BLAST).
To form a training example, sequences are sampled without replacement from a Silamond-UniRefie Cluster" until the total number of tokens reaches predetermined limit, and then concatenated form a sequence-of-sqquenee by sampling a random UniRef100 sequence from the same UniRef50 cluster the UniRef50 sequence being replaced. The UniRef100 sequences are randomly sampled with weight inversely proportional to the size the UniRef90 clusters they belong to. Each UniRef100 sequence is sampled at most once in each sequence-of of sequenecss..
Each sampled UniRei50 sequence is replaced with UniRef100 sequence
As final data augmentation, the order of the tokens in souuuence-ofssqueenee is reversed with probability 50% (i.e. all sequences sequence-of-sequeeeee are ordered from either N-terminus to C-terminus only C-terminus N-terminus only). This augmentation has been shown improve the performance of other protein language models [11].
Following this sampling procedure, the order of sequences in sequence-of-sequreee 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- lo-goge or 3) sngest-to-shortes We also compare the full joint log likelihoods of these sequences of 10 sequences with the same orderings (Table 2).
is random.
We find that sequence ordering has little-to-no effect on next sequence perplexity and joint log likelihoods. Ordering sequences shortest-to-longer 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.
16

---

Table 2:
Comparison of (a) next sequence perplexities and (b) joint log likelihoods log p( $2..... s,). 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-longee and longest-to-shortest orderings with the standard deviation across clusters in parenthesis. Relative difference calculated as (a - r)/r /r where is the ordered value and r is the value with random ordering.
(a) Perplexity Next Sequence N=9184
Mean Perplexity
Mean Difference from Random
Mean Relative Difference from Random
Sequence Order
Random
7,44075
Shortest-to Mhortest-to-Loree to-Longest
7.44828
+0.00753 (0.07686)
+0.00109 (0.01033) +0.00055 (0.01030)
Longest-to-Shorter
7,44410
+0.00336 (0.07608)
(b) Joint Log Likelihood N=9184
Mean Joint Log Likelihood
Mean Difference from Random
Mean Relative Difference from Random
Sequence Order
Random
18899974977
Shortest-to-Longestt
-8123.18167
-23.43149 (55.76313) +7.85511 (52.92428)
-0.00335 (0.00838) +0.00105 (0.00746)
Longest-to-Shortest
-8091.89462
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.
B.3
Hyperparameters
Hyperparameters for PoET variations used in ablation experiments ($5.2.1) are summarized in Table
Table 3: Hyperparameters for PoET models of different sizes.
# Parameters
Ni Naayers
Hidden Dim
Nitest
Batch Size (# Tokens)
57M
6
768
12
500k
201M
12
1024
16
250k
604M
16
1536
24
250k
C Homologous Sequence Retrieval Methods
experimented with two methods for retrieving homologous sequences from UniRef100, JackHM- Mer [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, used the protocol defined by ColabFold [35], and retrieved homologs from UniRef100 Version 2202. The main advantage of the ColabFold protocol that it is substantially faster than the ProteinGym protocol (minutes 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 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
17

---

Retrieval
ProteinGym
60
ColabFold
50
Count
30
20
10
0.0
0.5
1.0
2.5
Redundancy log(NNNfff
Figure Redundancy, as measured by N/Na,, 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 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 VranceptEVE 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 homologs via the site independent model, which performs similarly with both protocols.
Table 4: Average Spearman correlation ProteinGym subsets of baseline models and PoET using different methods for retrieving homologs.
Retrieval Method
Model
Val
Subst.
Site
ProteinGym ColabFold
0.386
0.408
independent
0.371
0.404
ProteinGyym
0.416
0.457
GEMME
ColabFold
0.418
0.463
ProteinGym
0.423
0.443
EVE
ColabFold
0.268
0.342
MSA Trans-
ProteinGym
0.427
0.431
former (ens.)
ColabFold
0.440
0.423
PoET
ProteinGym
0.463
0.474
(ens.)
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:
18

---

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 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, 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-likelihode (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 dissmilarity threshold. We evaluate PoET with (1) homologous sequences retrieved using either the ProteinGym or the ColabFold protocol, (2) 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: AP 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 necessary; there 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 e [1.0, 0.95, 0.90, 0.70, 0.50) and context length e (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 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 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
19

---

Validation
Taxon
Spearman ro with Max Dissimilaniyy 0.7
0.450
Substitutions
Prokaryote
0.425
Human
0.6
Eukaryote
Spearman rho
0.400
0.4
0.375
0.350
0.2
0.300
1.0
0.8
0.6
0.4
0.2
0.0
0.2
0.4
0.6
0.8
Dissimilarity
Spearman Dissimilarity
(a) Performance using ProteinGym protocol for retrieving homologous sequences
0.475
Validation Substitutions
Taxon
Spearman sho with Max Dissimilarity 0.7
0.450
0,8
Virus
Eukaryote
0.425
0.6
Prokaryote
Spearman tho
0.400
0.4
0.375
0.2
0.325
0.300
1.0
0.8
0.6
0.4
0.2
0.0
0.2
0.4
0.8
Dissimilarity
Spearman rho with Max Dissimilarity 1.0
(b) Performance using ColabFold protocol 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 dissim- ilarity threshold when using the ProteinGym protocol (top) and the ColabFold protocol (bottom). For ProteinGym, a threshold of 0.7 optimal and significantly better than no filtering i.e. using threshold of 1.0 (paired t-test on full substitutions benchmark: p 0.01, effect size: Ap 0.024). For ColabFold, no filtering is necessary and there is no essentially no difference in performance when filtering with a threshold 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.
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 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, 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 Nensemble 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 E (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, the validation benchmark, the two ensembling methods perform similarly; fixed has slight edge when using the ProteinGym protocol
we find that both types of ensembling improve performance.
20

---

Max Similarity
0.46
0,95
SabatRuliorn Speerther 11a
0
0.45
0.40
0.5
0.44
0.43
0.45
Vallcatton
0.42
0.41
0.44
0.40
12268
24576
02.4
12289
24376
Length
Length
(a) Performance using ProteinGym protocol retrieving homologous sequences
Max Similarity
0.47
1.0
0.95
0.46
0.9
Suliiiuutirss Speetttan
0,5
0.45
Validation
C.44
0.43
6144
12268
24576
6244
12289
24576
Length
Length
(b) Performance using ColabFold protocol retrieving homologous sequences
Figure 8: Averaged across datasets, PoET performs best with some filtering based similarity and at moderate context lengths. Each line plot shows the performance PoET on ProteinGym 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 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
Fixed (no ensemble)
0.442
0.445
0.474
ProteinGym
Fixed (15x ensemble)
0.464
0.467
0.521
Varied (15 combinations)
0,463
0.474
0.498
Fixed (no ensemble)
0,451
0.445
0.485
ColabFold
Fixed (15x ensemble)
0.478
0.469
0.512
Varied (15 combinations)
0.481
0.484
0.510
for retrieving homologs, and varied has slight edge when using the ColabFold protocol. On the full
substitutions benchmark, the varied method consistently performs better, and by 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 the PoET variants explored in the above sections, for final PoET model, choose the following parameters
21

---

A: Increased context tends to improve performance
BLAT ECOLX Jacquier 2013
Q2NOS5 SHMV1 Haddox 2018
0.52
Max Similarity
0.72
0.1
0.70
0.50
0.95
0.68
Increased context has neutral effect
12288
24576
12288
24576
CALMI_HUMAN, 2017
ML Wu,
SPG1_STRSG Olson, 2014
0.5
0.225
0.4
0.220
0.3
02
0.210
01
Spearman
63.44
12288
24576
61.44
12288
24576
63.4
12288
24576
R.401 YEAST Roscoe__o00 2013
DL64 RAT Mctaagghin_ 2012
SPIKE SAR52 Starr, bind 2020
2622
0.56
0.53
0.54
0.62
0.52
0.60
0.24
0.57
0.2
63.64
12288
24576
01444
12208
26576
0144
32288
24576
TPOR_ VPOR_HUNAN_BRR Bridgford S505N 2020
(CO8_ECOU_ Brrp_ECOU_Ttaaahhuooo Tripthii2016
0.45
0.50
C: Increased context tends to decrease performance
0.48
6144
12288
24576
61.4
12288
24576
Context Length
Figure Optimal max similarity threshold and context length varies across ProteinGym
datasets. Each line plot shows the performance on validation dataset as the max similarity threshold
and context length 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 have a neutral effect on performance. C: Datasets for which increasing the context length tends to have negative effect on performance.
22

---

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 se- quences with max similarity thresholds {1.0, 0.95, 0.90, 0.70, 0.50} and context lengths E (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 al. [17]. Specifically, sequences are sampled from probability distribution where sequence's 's probability is inversely proportional its number of neighbors the MSA of the homologous sequences. A target sequence is considered a neighbor of query sequence if the sequence identity to the query is greater than 0.8. The sequence identity is computed based on the alignment 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, condition only the amino acids in the homologous sequences that are aligned to position the target sequence in the MSA of the homologous sequences i.e. sequence in a subset Sj (Equation 4) may be the full length sequence as recorded UniRef100, rather 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 Average Spearman's rank correlation on the validation and substitutions benchmarks using full length sequences 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 high level, the MSA depth Not/ L is computed as the number of non-redundant sequences Net (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: Not L 1; Medium:
1 < Nan/ L < 100; High: Nen/ L < 100.
Although the MSA depth changes based on the method used to retrieve homologous sequences, we categorize proteins/datasets using the ProteinGym MSAs (I) in order maintain consistency with the categorizations in the ProteinGiym 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:
23

---

4.0
3.5
ColabFold log(Neff)
3.0
2.5
2.0
1.5
ProteinGym log(Neff)
Figure 10: Comparison of MSA depth (Norr) of MSAs of homologous sequences retrieved by the
ProteinGym and ColabFold protocols for target proteins in ProteinGym (p = 0.4,p p < 0.001, n = 72).
F.1.1 Alignment-base Models
Alignment-basse -based models first compute MSA of homologous sequences given protein family, and compute sequence likelihoods based on MSA statistics or models trained on the MSA.
Site Independent Model The simplest alignment-ben 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 interpendencies 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 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-baset masked lan- guage models trained on UniRef90.
ProGen2 [5] (ensemble) an ensemble of five Transformer-basee causal language models ranging between 150M and 6.4B parameters and trained on UniRef90 and BFD [45].
Tranception L (no retrieval) [11] is 700M parameter Transformer- baasd 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 100M parameter masked language model trained on MSAs of all UniRef50 sequences. MSAs are generated using HHblits [46].

---

E11.4 Hybrid
Hybrid models combine predictions from different types of models.
Tranception L [11] combines predictions from Tranception L (no retrieval) and site independent model in 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 manner that allows novel indels be scored by using EVE to "define 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 Prote- inGym 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 the target sequence. These common preprocessing steps used to improve the performance of other models for protein variant fitness prediction (Appendix D), and are the same parameters used 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 predict protein variant fitness.
G Ensembling PoET with Baseline Methods for Protein Variant Fitness Prediction
To ensemble PoET with baseline method for protein variant fitness prediction, we compute weighted sum of the score predicted by PoET and by the baseline method i.e.
B AROETTE
Baseline,
where P Basollnee is the score predicted by the baseline method, FROET, is the score predicted by PoET,
and Fr is the ensemble prediction score.
The weight a is determined by optimizing performance over the ProteinGym validation set using ternary search over the range [le-3, 1e4].
25

---

Table 7: Average Spearman correlation between model scores and experimental measure-
ments on ProteinGym for various ensembles of PoET and baseline methods. Performance of TranceptEVE L, the best bascline method, and PoET alone are also shown for reference.
Model
Model
Substitutions
Indels
TranceptEVE
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
GEMME
0.477
0.488
N/A
EVE (cns.)
0,485
0.487
N/A
Tranception L (no retrieval)
0.471
0.473
0.512
PoET (ens., ProteinGym Protocol)
MSA Transformer (ens.)
0.476
0.472
N/A
Tranception
0.480
0.475
0.510
TranceptEVE
0.488
0,486
0.504
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
PoET (ens., ColabFold Protocol)
MSA Transformer (ens.)
0.487
0.484
N/A
Tranception
0.489
0.483
0.525
TranceptEVE
0.494
0.492
0.521
26

---

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.
Maz Simianty
Max Simiinity
0.48-
1.0
0.48
Volibaton Spcanman no
0.47-
C.9
6.47
C.7
C.7
0.46
0.46
3.85
0.4
614
8332
12288
24578
614
12288
24576
Contect Length
Langth
(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, 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- "Diamond-UniRef50 clusters" (Appendix B) in the validation set. PoET significantly outperforms two baseline models, profile HMM and a regular Transformer (Figure 12).
HMM
Regale Yo snttreee POET Thanffrmer
4090
8102
C144
Number of Tobne
Figure 12: Comparison of the perplexity of profile HMM, regular Transformer trained with 8K context length, and PoET trained with 8K context length when generating a protein sequence conditioned 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
27

---

more indel rich when more sequences (longer context lengths) are considered. PoET achieves low perplexities on these heldout sequences, showing that a good generative model of indel rich families.
Godc
Clpps 1ereteriet 1e SndBov LNDSLLIO
(a) Typical example of MSA of UniRef50 cluster. Bar plot above the MSA shows the percentage of gaps in each corresponding MSA column.
MSAs with <8826 tallens
MSAs with <61663 tokens
05 Popetimnc Courm
(b) (Left) Distribution of percentage of gaps in MSA columns for the example MSA above. (Right) Same distribution for all MSAs of UniRet50 clusters in the validation with <4096 tokens (top) and <61440 tokens (bottom). Larger MSAs more novel indels.
Figure 13: MSAs of UniRef50 protein clusters reveal that these families contain many novel indels
Inference Speed of PoET vs VV L.
In the section, compare the time it takes for PoET and TranceptEVE to predict the fitness of set of protein variants on A 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 TranceptEVE, 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 EVE models on the MSA
4. Use the EVE models compute the EVE log prior
5. Compute the likelihood of each variant by combining scores from the Tranception 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 lower bound on the number of variants that can 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 TranceptEVF. The homolog retrieval time is irrelevant for this analysis since both methods require homolog retrieval step, although note that the ColabFold
28

---

protocol used to retrieve homologs by PoET is substantially faster than the ProteinGym protocol used by TranceptEVE.
We find that independent of sequence length, on a single A100 GPU, PoET is able to score approxi- mately 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 hundreds of thousands of variants to be scored by PoET with a reasonable amount of compute before TranceptEVE able to score a single variant.
K Function-specifie 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 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 a protein (e.g. more thermostable variant is also more likely to have higher expression or enzyme activity), optimizing general fitness likely to optimize the true properties of interest, albeit indirectly. PoET is useful tool for protein engineering because it is able to successfully predict general fitness by conditioning on and inferring evolutionary constraints from diverse set of homologs of the target protein. This raises 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 inevitable trade-off between learning from smaller set of homologs that more relevant, and a larger set of homologs from which there is more data to learn.
As 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, 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 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.
MSA of
On the chorismate mutase dataset, 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 (Ap 0.2, Figures 14b, 14d). In fact, PoET conditoned on functional natural sequences even outperforms fully supervised methods, including a Gaussian process trained on mean embeddings from BERT-lik like protein masked language model (Ap = 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 simply conditioned on positively labeled natural sequences and does not have access the measured activities. This enables PoET to be used with assays that only measure binary endpoints rather than continuous values.
29

---

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 PoET conditoned on functional natural sequences
from BERT-like masked protein language model
0.56772
30

---

Natural (rho--0.0352)
Designed (rho- 0.40926)
0.52
Messured
0.25
0.00
-0.35
150
125
100
75
50
25
25
150
125
100
75
50
25
25
Predicted Atmeess
Predicted Athees
(a) Predicted fitness using Potts model trained on natural sequences from Russ al. [9].
Natural (fho=0.11743)
Designed 0h3-0.36388)
1.00
0.25
0.09
-0.31
250
225
200
175
150
125
100
75
250
225
200
175
150
125
100
75
Predicted Athoss
Prodicted Athoss
(b) Predicted fitness using PoET conditioned natural sequences. Natural sequenees retrieved using PSI-BLAST to search for homologs of the target protein [9].
Natural (tho=0.40608))
Designed (rho=0.50855)
1.00
Measured
0.25
0.00
-0.23
0..50
-230
-200
-150
-1.0
-50
-200
200
-150
-100
Predicted Aineess
Predicted fitness
(c) Predicted fitness using PoET conditioned on homologs of the target protein retrieved by the ColabFold protocol.
(rho-0.61857)
Designed Irho=0.56 (05-00))772)
1.00
0.23
0.00
-0.23
300
250
200
150
1.00
50
300
-50
200
150
1000
50
Predicted Atness
Predicted Ainees
(d) Predicted fitness using PoET conditoned functional natural sequences. Natural sequences with measured fitness >0.42 are considered functional by Russ al. [9].
Figure 14: Measured fitness vs predicted fitness of (left) natural and (right) designed chorismate mutase enzymes from Russ al. [9].
31

---

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 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 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.
High confidence (PLDDT>90) AlphaFold2 [35] predicted structures of generated sequences aligned to the experimentally solved structure of the target chorismate mutase (PDB: IECM) show that structure preserved despite significant divergences in sequence identity; two examples are shown in Figure 16 and the full distributions of pLDDTss and TM-scores are shown in Figure 5A.
45
0.20
Nabural SOQUADRCRS IeEn oonerrtedd
(a) Distribution of max seq id between gener- ated and natural sequences
(b) Per position entropy of natural and gener- ated sequences
Figure 15: Statistics of natural and generated sequences.
L Sampling methods for sequence generation experiments
When sampling sequences from PoET, we used nucleus sampling [50] with p = 0.90. When sampling sequences from ProGen, we tried nucleus sampling with E {0.25, 0.50, 0.75), which are the three sampling strategies used by the authors of ProGen for lysozyme sequence generation. We found that when p < 0.75, the resulting libraries had low diversity (> 50% of generated sequences were duplicates). Therefore, we only present results with p 0.75; under this setting, almost all generated sequences were unique.
M Limitations
Orphan Proteins While ProteinGym includes DMS data some proteins with low MSA depth, does not include any orphan proteins that have no homologs. Simulating orphan proteins by e.g. removing all homologs from the training set is computationally intractable it would require retraining large models and may be sensitive to homology search parameters. As a result, our experiments are inconclusive with respect to orphan proteins.
Exhaustive scoring of all high order mutants Some methods such as GEMME can achieve similar albiet lower performance than PoET with substantially lower compute. However, no method is efficient enough to exhaustively score all variants with more than couple of mutations. For
32

---

(b) Structure alignment with the highest scoring (according to PoET) generated sequence with maximum sequence identity of 0.5 natural sequence.
(a) Structure alignment with the highest scoring (according PoFT) generated sequence with maximum sequence identity of 0.8 to natural sequence.
Figure 16: AlphaFold2 predicted structures generated sequences aligned to the target chorismate mutase (PDB: IECM). Both alignments have RMSD 1.02 and TM Score 0.9.
exploring the space of higher order mutants, direct generation is far more effective, and is enabled by PoET via sampling.
Limited compute for ablation studies Our ability to conduct ablations is ultimately constrained by compute resources. Ablation models were trained until validation set performance plateaued, but further training may have changed the results.
33

---

Natural sequence space Functional subspace PoET Generated
Figure 17: UMAP of natural chorismate mutases (blue), chorismate mutases funetional in E. coli (orange), and sequences sampled from PoET when conditioned on the latter sequences (green). Sequences are compared using the Hamming distance metric.
34

---

Supplementary Tables
Table S1: Average Spearman correlation between model scores and experimental measurements on ProteinGym by MSA depth. This is extended version of the main result table (Table 1) that includes performance on the validation benchmark.
Substitutions by MSA Depth
Model type
Model name
Val
Medium
High
All
Indels
Alignment-
Site independent
0.386
0,417
0,404
0,411
0.408
N/A
based
GEMME
0.418
0.445
0.449
0.522
0.463
N/A
model
EVE (ens.)
0,435
0.414
0,441
0.498
0,448
N/A
Uncond-
ESM-1v (ens.)
0.355
0.356
0.372
0,510
0.398
N/A
itional
ProGen2 (ens.)
0.397
0.357
0,416
0.448
0,411
0.407
PLM
Tranception (no retrieval)
0.396
0.377
0.399
0.429
0.401
0.430
Cond-
MSA Transformer (ens.)
0.440
0.372
0,421
0.477
0.423
N/A
itional
PoET
0.481
0,476
0,466
0.542
0.484
0.510
Tranception L
0,447
0.441
0.437
0.472
0.445
0.464
TranceptEVE M
0.516
Hybrid
TranceptEVE
0.470
0.454
0,463
0,508
0.471
0,466
PoET
+TranceptEVELL
0.494
0.479
0,480
0,537
0.492
0.521
Table S2: Average Spearman correlation between model scores and experimental measurements on ProteinGym by mutation depth. Mutation depth is the number of substitutions variant sequence has compared to the target sequence. Bolded values indicate the highest value in each column, with values above and below the double-line considered separately.
Substitutions by Mutation Depth
Model type
Model name
Val
3
5+
All
Alignment-
Site independent GEMME
0.386
0.408
0.325
0.306
0.316
0.409
0.408
based
0.418
0.459
0.412
0.394
0.353
0.440
0.463
model
EVE (ens.)
0.435
0.449
0.409
0.405
0.351
0.429
0.448
Uncond-
ESM-Iv (cns.)
0.355
0.396
0.309
0.203
0.165
0.253
0.398
itional
ProGen2 (ens.)
0.397
0.409
0.312
0.233
0.228
0.282
0.411
P.M
Tranception L (no retrieval)
0.396
0.392
0.398
0.418
0.334
0.422
0.401
Cond-
MSA Transformer (ens.)
0.440
0.426
0.381
0,444
0.374
0.431
0.423
itional
PoET
0.481
0.483
0.456
0.473
0.404
0.449
0.484
Tranception L
0.447
0.441
0.427
0.440
0.370
0.461
0.445
TranceptEVE M
Hybrid
TranceptEVE L
0.470
0.471
0.445
0.456
0.389
0.464
0.471
PoET TranccppEVE
0.494
0.492
0.465
0.478
0.411
0.471
0.492
35

---

Table S3: Average Spearman correlation between model scores and experimental measurements on ProteinGym by taxon.
Substitutions by Taxon
Other Fukary.
Model type
Model name
Human
Prokary:
Virus
Alignment-
Site independent
0.386
0.383
0.471
0.402
0.414
0.408
based
GEMME
0.418
0.434
0,514
0.507
0.438
0.463
model
EVE (cns.)
0.435
0.403
0.499
0.500
0.435
0.448
Uncond-
ESM-Iv (cns.)
0.355
0.417
0.441
0.502
0.256
0.398
itional
ProGen2 (ens.)
0.397
0.394
0.459
0.462
0.364
0.411
PL.M
Tranception (no retrieval)
0.396
0.359
0.436
0.450
0.395
0.401
Cond-
MSA Transformer (ens.)
0.440
0.352
0.455
0.494
0.441
0.423
itional
PoET
0.481
0.445
0.524
0.526
0.479
0.484
Tranception L
0.447
0.413
0.503
0.478
0.429
0.445
TranceptEVE M
Hybrid
TranceptEVE
0.470
0.436
0.518
0.514
0.454
0.471
PoET + Tranncept.VF I.
0.494
0.450
0.538
0.536
0.486
0.492
36

---

Supplementary Figures
Figure S1: Model performance for each dataset in the ProteinGym substitutions benchmark.
Tanception
TancepteVEL 7ancoptEVE
NET TAncept
2.8
Spaarman tho
0.6
0.2
FINN mrNN Janoon pegger desellensa
HI67 SEAST Posuaave,, Jnddes
Sur per.o.aaoo www.ooooooooooo Connsine jnoeetts
con Jaozt sioe paresz
PEB. HUMAN joter dolatonns
miLOws gww.o Buno0
AGA MAALETT SHRO7 Doviioo
Figure S2: Model performance for each dataset in the ProteinGym indels benchmark.
37


Reviews:
[{"review_id": "K47D3Mk4dg", "reviewer": "Reviewer_n5MZ", "summary": "- This paper introduces PoET, a novel autoregressive transformer architecture that learns a distribution over protein families. \n- More specifically, PoET takes as input a concatenated set of homologous sequences, for example retrieved with a MSA. The model architecture consists of a stack of several TieredTransformerDecoderLayers (the main innovation from a modeling standpoint), where each layer applies successively: 1) causal self-attention within each sequence 2) causal self-attention across sequences in the homologous set 3) a standard feedforward NN.\n- From a practical standpoint, the model can be used at inference both to generate new proteins or to assess the effects of mutations on protein fitness\n- The former (new sequence generation) is partially covered (one interesting yet limited example in appendix for chorismate mutase)\n- The latter (fitness prediction performance) is the focus of experiments conducted in the paper. PoET achieves remarkable performance on the ProteinGym benchmarks using a 200M-param model", "questions": "- “Deep mutational scans and directed evolution experiments have been used to successfully design novel proteins [1, 2], but can be costly and difficult to implement,  which makes these experimental methods inapplicable for many proteins and functions of interest.” -- please elaborate on which proteins / functions DMS and directed evolutions are \"inapplicable\" (reference paper would be great to add here).\n- “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.” -- isn't your method subject to these errors as well since you retrieve homologous sets at inference via a MSA?\n- Any particular reason for using Diamond at training and then colabfold or jackhmmer at inference? Have you tried training on homologous sets retrieved with MSAs? (would make inference closer to training)\n- Given your position encoding (which requires coordinate systems to be ~ equal across homologous sequences), isn't your method also not adapted to scoring sequences with long insertions/deletions? \n- What are the limitations of your method? It seems very compelling for the majority of settings, except perhaps: 1) if one wants fast inference for good performance, GEMME is a better option 2) if one cares about deep mutants, Tranception / TranceptEVE (or ensembled with PoET) seem better 3) if one cares about disordered proteins / proteins with very few homologs (ie. less than 10), then it is not clear whether your model would handle well (since sets with fewer than 10 homologs were removed for training as per line 181). Would you agree with these limitations? How does the performance vary as the number of homologs goes to zero (note this would be a substantially different setting vs proteins in the \"low MSA depth\" bucket from ProteinGym)? \n- Figure S1/S2: there is a reference to Prots2prot? Please also provide the performance at the DMS level for PoET alone\n- Given the last point in the \"weaknesses\" box above, which avenue(s) do you see to further improve the performance of your model?", "limitations": "Not particularly discussed, would suggest adding a couple sentences about it, as per the questions above.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The performance of the architecture on the zero-shot fitness prediction tasks from ProteinGym is quite remarkable for a model of this size\n- The model architecture is simple and effective\n- The different ablations to curate the optimal prompt engineering strategy are thorough and packed with many interesting insights", "weaknesses": "- The sequence generation abilities of the model are not really explored (besides one example for 1 protein family in appendix)\n- This claim at the end of the introduction (lines 69-70) appears to be wrong: \"improves prediction of variant effect in sequences with large numbers of mutations.\" Performance on deep mutants (5+ in Table S2) is lower than baselines (Tranception & TranceptEVE). Would suggest removing this sentence and adjust the second to last sentence of conclusion, as well as fixing the bolding in the appendix table.\n- Some other claims are not properly substantiated. For example: \"One advantage of PoET is that is able to not only score indels, but also generate sequences with indels\" (lines 257-258). There does not seem any evidence in the paper that the quality of such \"generated sequences with indels\" would be any good, especially given the specifics of the position encoding chosen which relies on the assumption that \"amino acids at similar absolute positions in homologous proteins are more likely to be drawn from the same distribution\" (168-169). See questions for other claims where evidence seemed light.\n- The performance seems to plateau as the number of parameters grow (the ~600M-param model performs on par with the ~200M-param version) and in terms of context length (12k better than 24k at inference) suggesting this architecture has perhaps already reached its limits", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive comments and suggestions. In the manuscript, we strove to provide a comprehensive analysis of the prompt engineering aspects, which we hope will be relevant for using PoET and also for future work in retrieval-/homologue-augmented protein ML models. We’re glad the reviewer found this part interesting. We address the reviewer’s questions and concerns below. We will also expand the discussion of limitations to include points raised in this discussion.\n\nOn PoET as a generative model, especially relating to the reviewer’s concerns about the quality of sampled indels: we focused our analysis on the variant effect prediction problem, because we view it as a good surrogate for generation in the sense that assigning high likelihood to high functioning variants means that samples or other decodings from the model, which are high likelihood, are also likely to be functional. In addition to our chorismate mutase example in the appendix, we have expanded our analysis of indel generation and will show some additional sampling examples in the revised manuscript. Please see our global response for more details. As suggested by other reviewers, we will also revise the introduction and other sections to better emphasize the variant effect prediction task.\n\n- This claim at the end of the introduction...\n\nWe will revise the manuscript as suggested.\n- The performance seems ...\n\nThis was one of the most surprising findings to us in this work. We also expected larger models and longer context lengths to perform better. We suspect this is related to the fundamental mis-specification of the variant effect prediction problem as density modeling, as recently discussed by [1] and in our Appendix (lines 399-412). We think this is also supported by the fact that holdout perplexity does improve at longer contexts lengths, but variant effect prediction does not. \n\n- “Deep mutational scans ...\n\nFundamentally, deep mutational scanning can only be applied in settings where a high throughput functional assay is available, which usually means that variants can be produced in a single pot and the function of those variants can be read out via a selection and sequencing assay. Low throughput, expensive, and time consuming assays are not amenable to deep mutational scanning due to cost and time limitations. These references [2,3] discuss some of the challenges and considerations for developing functional assays. We will add these references to the text.\n\n- “PoET is a...\n\nWe do not retain the alignment, passing only the unaligned sequences into PoET as the prompt. The challenge with operating on the MSA is not retrieval, but rather that large gaps disrupt the continuity of the sequences as viewed by axial transformer style models. We also do not assume a specific correct alignment.\n- Any particular reason...\n\nWe used Diamond because it was the only tool available to perform the all versus all homology search in a reasonable amount of time. It is substantially faster than alternative homology search programs while attaining highly competitive performance. \n- Given your position...\n\nThe relative positional encodings only provide a prior on the alignment. The model is able to, and does, learn how to properly correspond sequence elements within the sequences. In fact, in newer experiments, we removed the between-sequence relative positional encoding altogether and the model performs nearly identically, suggesting that this does not hinder model performance on long sequences and indels. We also note that our validation uniref50 clusters are extremely gappy and have high variability in sequence length within each cluster (see global response). Thus it directly measures our ability to model such indels.\n- What are the...\n\nWe examine the generative performance of PoET, in terms of perplexity on heldout sequence clusters, using different numbers of homologous sequences, where we see good performance. In fact, PoET performs dramatically better than profileHMMs or other viable protein sequence models in this very few homologs regime (See Appendix Figure 10 for evaluation in terms of number of tokens in MSA and Review Supplement Table 2 in the global response for evaluation in terms of number of sequences), so we view this as a major advantage of PoET. Given our current understanding, we generally agree with your other comments. Clearly, GEMME is much faster than any neural network-based model, but we don’t think this is a significant limitation, especially given that we can efficiently sample from PoET as an alternative. Any kind of exhaustive evaluation of high order variants is likely to be intractable regardless of method chosen simply due to the huge size of those spaces.\n- Figure S1/S2:...\n\nProts2prot was an old working name for PoET. We will correct this in the manuscript.\n- Given the last...\n\nGiven that we think this is a problem related to misspecification of the unsupervised variant effect prediction problem, as well as the challenge in directly linking DMS fitness measurements with specific protein properties (often, many factors contribute to fitness and it may not be clearly associated with one property), this is a challenging question to answer. We think the ability to include additional conditional information in the prompt, such as structures, or specific property specifications could help, and would provide additional mechanisms for controlling the generative distribution and sequence generation.\n\n[1] Weinstein, Eli, et al. \"Non-identifiability and the Blessings of Misspecification in Models of Molecular Fitness.\"\n\n[2] Fowler, D., Fields, S. Deep mutational scanning: a new style of protein science. Nat Methods 11, 801–807 (2014). https://doi.org/10.1038/nmeth.3027\n\n[3] Tiefenauer, L., & Demarche, S. (2012). Challenges in the Development of Functional Assays of Membrane Proteins. Materials, 5(11), 2205–2242. https://doi.org/10.3390/ma5112205"}, {"review_id": "yHcDDY1XTM", "reviewer": "Reviewer_jcwZ", "summary": "Current generative protein language models focus on generating individual protein sequences and are not specifically trained to generate sequences for an entire protein family. The paper introduces a new autoregressive generative model that tackles this limitation by generating Multiple Sequence Alignments (MSAs) for the entire family. The model utilizes a TieredTransformerDecoder to model the within-sequence interactions in an order-dependent manner, as well as the between-sequence interactions in an order-independent manner. The authors evaluate their models on protein variant fitness tasks by measuring the probability of variants conditioned on searched MSA of the parent sequence. The proposed method achieves large improvements on the ProteinGym benchmark. Importantly , the method addresses the problems faced by previous approaches, such as handling proteins with low-depth MSAs and accommodating indel mutations in addition to substitutions.", "questions": "Major points:\n1. The author proposes to encode the MSA independently of the order between different sequences. But they still use an autoregressive decoding method that introduces order dependence in the output probabilities. For example, the probability equation $p(s_1)p(s_2|s_1)p(s_3|s_1,s_2)=p(s_1,s_2,s_3)=p(s_1,s_3,s_2)=p(s_1)p(s_3|s_1)p(s_2|s_1,s_3)$ may not hold when using the proposed autoregressive decoding. In the former case, generating tokens of $s_3$ is conditional on $s_1$ and $s_2$, whereas in the latter case, generating tokens of $s_3$ is conditional on only $s_1$. \nTherefore, when measuring the fitness with $p(v,s_1,…,s_m)$, the results will differ using $p(v|s_1,…,s_m)p(s_1,…,s_m)$ and $p(s_{j+1},…,s_m|v,s_1,…,s_j)p(v|s_1,…,s_j)p(s_1,…,s_j)$. How is this problem addressed during inference? How is the order of MSA sequences determined during training?\n2. A trivial baseline for generating MSA is to use an autoregressive language model, such as ProGen, conditioning the model on the parent sequence and continuing the generation with a start token afterward. If limited by the context window, we can only generate one sequence at each time. Although the model is not trained on MSA datasets, it is now well-known that language models trained with next-token prediction loss can excel at generating next sequences. How would this perform compared with the proposed approach?\n3. The organization of the paper is problematic. It appears that the most significant contribution of this paper lies in the improvement of protein variant fitness prediction tasks. However, the title and introduction sections place more emphasis on the generation of MSAs, which lacks substantial support in their experiments (the main results and ablation study are all based on fitness prediction). I suggest the authors reconsider the introduction section to highlight the \"proposal of a new approach for protein fitness prediction\" rather than \"proposal of a new approach for generating MSA.\"\n\nMinor points:\n1. In Sec. 5, it would be helpful to introduce the included baseline methods and provide the corresponding references before discussing the results.\n2. In Table 1, the performance of TranceptEVE M on Indels is shown to be lower than that of TranceptEVE L. Could you provide an explanation for this observation?\n3. Line 154 in Sec.3.1.2: while it may be well-known, it would be better to explain the abbreviation \"RoPE\" when it is first mentioned and include a reference at that point.\n\nOverall, though with flaws, I think the paper makes a big contribution to the protein design community considering the amount of work in the paper. I’m willing to raise my score if the authors can address my concerns.\n", "limitations": "The authors should add a paragraph to discuss the potential limitations and negative societal impacts of their work.", "rating": 7, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. Protein variant fitness prediction tasks are very important for protein design. The approach, which achieves large improvements on this task, will definitely have large potential impact.\n2. The related work section is well-documented, and it is relatively clear where the authors' contributions lie.\n3. According to Table 1, the proposed method performs exceptionally well on proteins with low-depth MSAs, which supports the motivation behind generating MSA instead of individual sequences.\n4. While the technique used in the paper is not entirely new, the idea of generating MSA instead of individual sequences is novel enough for an application-oriented paper.", "weaknesses": "1. The soundness of autoregressive decoding of MSA is questionable. Please see the questions below.\n2. The authors claim their contribution as proposing a model for generating sequences for the entire family, yet few experiments are conducted to evaluate the generated MSAs.\n3. Considering the significance and volume of this work, it would be disappointing if the code were not provided. The authors only promise to offer an accessible API, but it’s not sure whether they will provide the training code and trained models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive comments and excitement about the potential impact and performance of our method. In particular, we appreciate that the reviewer feels that our “...paper makes a big contribution to the protein design community considering the amount of work in the paper.” We respond to the reviewers specific comments and questions below.\n\n- The soundness of...\n\nTo be clear, we do not decode MSAs. We decode sets of related protein sequences, which are unaligned.\n\n- The authors claim their contribution...\n\nThe model is a generative model of whole protein families as sequences-of-sequences, but this full family generative task is not our primary application of interest. In the same way that GPT is a generative model of text, but is only used conditioned on prompts in practice, we are primarily interested in using our model as a conditional generative model, which is permitted by the structure of the model. The conditional generative capabilities enables both scoring of variants for fitness prediction and efficient generation of novel sequences via sampling. We have expanded our evaluation of the generative capabilities of the model (see global response), and will update the manuscript to include these in the main text.\n\nWe have also explored unconditional generation of whole families with PoET, where we find that sampled families produce reasonable looking multiple sequence alignments, when aligned. We have also found that, interestingly, despite the lack of homology of these families to any natural proteins, conditioning AlphaFold2 on these multiple sequence alignments leads to feasible looking structure predictions with higher pDLLTs than are given to structure predictions using a single sequence from the sampled family alone. Unfortunately, we do not have enough space in the review supplement to include these results, and we don’t want to read too much into what this means, aside from the fact that PoET generates protein families that follow reasonable family-level statistical constraints. We will add some discussion of this to the manuscript.\n\n- Considering the significance...\n\nWe will release code and the trained models with the camera ready version of the manuscript.\n\n- The author proposes...\n\nThis is true, only the individual decoder layers are invariant to the order of the sequences (discussed at lines 131-132); the full model, composed of multiple such layers, is not. When considering variant fitness prediction, we only ever consider the conditional likelihood of the variant given the homologues, p(v | s1 … sn). In practice, we do not find ordering of the sequences to be an issue when considering the full joint likelihood p(s1 … sn), because the model is trained on random orderings of sequences, which requires it to learn to generate sequences given any ordering of the prior sequences. We will emphasize this fact in Section 3.2 for the revised manuscript.\n\n- A trivial baseline for generating ...\n\nProGen is not designed to condition on sequences or to generate multiple sequences. It only conditions on control tags, which are generally annotations from Uniprot. Please see Review Supplement Table 1 in the global response to see how ProGen performs in terms of perplexity when conditioning on a parent sequence. We have also compared PoET with a baseline model that simply autoregressively generates the whole sequence-of-sequences in (Section 5.2.2 Figure 4, Appendix K Figure 10), which shows that general language models do not perform well at this task as they are unable to generalize beyond their training context length. In contrast, our specialized transformer layer allows PoET to generalize well to much longer context lengths.\n\n- The organization of the paper...\n\nWe will revise the manuscript to better emphasize the variant effect prediction task. Also, please see the global response for additional analysis of the generative capabilities of PoET.\n\n- In Sec. 5, it ...\n\nWe will revise the manuscript as suggested.\n\n- In Table 1, the performance of TranceptEVE M on Indels is shown to be lower than that of TranceptEVE L. Could you provide an explanation for this observation?\n\nThese results are reported directly from the TranceptEVE paper. We also aren’t sure why TranceptEVE M performs better than L on indels, but it could be related to the mis-specification of the variant effect prediction task, or some other source we aren’t aware of. The number of datasets with indels is also significantly smaller than the substitution-only DMS datasets, which could contribute (the difference between M and L is not statistically significant, see lines 256-257).\n\n- Line 154 in Sec.3.1.2: ...\n\nWe will revise the manuscript accordingly.\n\n- The authors should...\n\nWe believe an explicit “Limitations” section is optional at NeurIPS this year and sought to discuss limitations throughout the manuscript where appropriate. We will collate these limitations and others raised in this discussion in a clear “Limitations” section in the revised manuscript.\n"}, {"review_id": "NQF3sTEcUN", "reviewer": "Reviewer_sUvw", "summary": "The authors of the paper present Protein Evolutionary Transformer (PoET), an autoregressive transformer-based model that is able to generate sets of related protein sequences, enabled by the proposed novel Transformer layer. This sequences-of-sequences generating method benefits from transfer learning, is able to handle novel indels, and does not need a multiple sequence alignment (MSA) as its input. The model can be conditioned on sequences of a protein family of interest for generation and scoring. The proposed method shows similar or superior performance compared to other models in fitness prediction tasks, while being faster without requiring retraining.", "questions": "1. As mentioned in \"Weaknesses\": Limitations section missing.\n2. As mentioned in \"Weaknesses\": perhaps include some discussion on the (mis)match between density estimation and fitness prediction.\n3. As mentioned in \"Weaknesses\": if at all possible, it would be very informative to report averages over multiple runs, especially for the ablation results (Figure 3) but also for Table 1, to make the observations more convincing.\n\nMinor comments:\n\n4. The abstract states that PoET *outperforms* other models. However, from the number of boxed values in Table 1, this claim might be a bit too strong.\n5. Do you report the dimensionality of embedding size $d$ somewhere? I might have missed it.\n6. If I understand correctly, the relative positional encoding scheme would probably not be beneficial when there are big differences in sequence length amongst homologous proteins. In that case, there must be a better encoding possible (some MSA-like). Could you discuss this or perhaps mention it as a limitation?\n7. The results on ensembling PoET with other models (Table 1 + Table 6 in the appendix) are interesting. Did you also try ensembling PoET with one or multiple other PoET model(s)?\n8. Section 5.2.1: \n    * In general, it could be considered \"cheating\" to monitor the correlation during training since this is essentially your test data. If it's just for these experiments then it's fine, I'm just checking if it's not something you used as a stopping criterion in general?\n    * Context Length: apart from the earlier suggestion to do multiple runs here, it's also worth pointing out that the drawn conclusions depend quite strongly on when it was decided to stop training. \n    * Model Size: the stopping criterion of \"when the performance seemed to plateau\" looks a bit arbitrary. If we compare the lines in the right graph in Figure 3, the blue line was cut off quite early while similar \"plateaus\" can be found in the purple and brown lines, which were allowed to train longer.\n9. Missing references:\n    * Figure 1a is never referenced.\n    * Paper reference missing for RoPE?\n    * Not all appendices are referenced in the main text (B, E, F, I, and O are missing).", "limitations": "Even though some limitations are touched upon in the main text, a thorough discussion of limitations is lacking.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "This paper is very well-written, with clear explanations, great examples and figures, and a comprehensive Related Work section. The experiments are explained well, they align with the claims made in the paper, and the results are discussed adequately. Moreover, there's an abundance of extra information in the appendix.", "weaknesses": "I enjoyed reading this paper, I do not have any major concerns. However, I did miss a \"Limitations\" section or something equivalent, either in the main text or in the supplementary. Moreover, I think some results (mainly those in the \"Ablation\" section 5.2) could be made stronger, for example by averaging performance over multiple runs and showing error bars, because the observed trends are not always that convincing. Finally, and this is more of a general issue, even though the results show improved/competitive performance, the average correlation values to experimental data is relatively low, i.e. around 0.5, which is still quite a weak correlation. The mismatch between density estimation and fitness prediction has been discussed before, for example in [1], and it might be worth discussing to some extent in this paper as well.\n\n[1] Weinstein, Eli, et al. \"Non-identifiability and the Blessings of Misspecification in Models of Molecular Fitness.\" Advances in Neural Information Processing Systems 35 (2022): 5484-5497.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their overwhelmingly positive response and comments. In answer to the reviewer’s questions:\n\n- As mentioned in \"Weaknesses\": Limitations section missing.\n\nWe sought to discuss limitations throughout the paper as appropriate, but we will revise the paper to include an explicit “Limitations” section for the camera ready version.\n\n- As mentioned in \"Weaknesses\": perhaps include some discussion on the (mis)match between density estimation and fitness prediction.\n\nYes, this is an extremely interesting point, which we mention briefly in Appendix D.2 Lines 399-412, but we will expand on this discussion in the revised manuscript.\n\n- As mentioned in \"Weaknesses\": if at all possible, it would be very informative to report averages over multiple runs, especially for the ablation results (Figure 3) but also for Table 1, to make the observations more convincing.\n\nWe did not feel that multiple runs were necessary given the high cost of training these models and the fact that results are already averaged over a large number of datasets, which we think represents a more interesting source of variation, anyway, regarding expected performance on new variant effect prediction problems. We also do not have access to multi-run averages from other methods. Statistically significant claims relating to this source of variation are indicated with a box in Table 1. \n\n- The abstract states that PoET outperforms other models. However, from the number of boxed values in Table 1, this claim might be a bit too strong.\n\nWe will adjust the text as suggested.\n\n- Do you report the dimensionality of embedding size d somewhere? I might have missed it.\n\nThe dimensionality of the embedding is 1024. We report this in Appendix B.\n\n- If I understand correctly, the relative positional encoding scheme would probably not be beneficial when there are big differences in sequence length amongst homologous proteins. In that case, there must be a better encoding possible (some MSA-like). Could you discuss this or perhaps mention it as a limitation?\n\nThe relative positional encoding scheme basically acts as a weak prior on the alignments between the sequences. We thought the same and have also trained a version of PoET without any positional information between sequences, which performs essentially identically to the version of the model we presented here. This suggests the model is able to learn the correct alignments and is not sensitive to this prior. We agree that there are other encodings which may be better, although, it is not-so-obvious how they can be applied correctly with an autoregressive decoder.\n\n- The results on ensembling PoET with other models (Table 1 + Table 6 in the appendix) are interesting. Did you also try ensembling PoET with one or multiple other PoET model(s)?\n\nWe have not, but this would be interesting to try!\n\n- Section 5.2.1:\n- In general, it could be considered \"cheating\" to monitor the correlation during training since this is essentially your test data. If it's just for these experiments then it's fine, I'm just checking if it's not something you used as a stopping criterion in general?\n\nAll hyperparameters, including training time and model size (Figure 3 caption) are tuned based on validation set performance only. This is consistent with other methods evaluated on ProteinGym.\n\n- Context Length: apart from the earlier suggestion to do multiple runs here, it's also worth pointing out that the drawn conclusions depend quite strongly on when it was decided to stop training.\n- Model Size: the stopping criterion of \"when the performance seemed to plateau\" looks a bit arbitrary. If we compare the lines in the right graph in Figure 3, the blue line was cut off quite early while similar \"plateaus\" can be found in the purple and brown lines, which were allowed to train longer.\n\nWe agree that the stopping criterion can have a strong effect on results, and will note this as a limitation of the ablation. While we try to be as objective as possible, our ability to conduct ablations is ultimately constrained by compute resources.\n\n- Missing references:\n- Figure 1a is never referenced.\n- Paper reference missing for RoPE?\n- Not all appendices are referenced in the main text (B, E, F, I, and O are missing).\n\nWe will correct these in the revised manuscript.\n"}, {"review_id": "a6TUDIWOlm", "reviewer": "Reviewer_Uh6C", "summary": "This paper proposes an autoregressive generative model (protein evolutionary transformer, PoET) of whole protein families. Current generative protein language 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. This model can incorporate new sequence information without retraining, generalize to large context lengths, and avoid issues related to conditioning on MSAs. They propose a novel Transformer layer that models order-dependence between tokens within sequences and order-independence between sequences. The advantages include that PoET can be used as a retrieval-augmented protein language model, 2) generate and score novel indels in addition to substitutions, and does not depend on MSAs of the input family, 3) extrapolate from short context lengths allowing it to generalize well even for small protein families.", "questions": "1. Line 64, the contribution 'PoET can be sampled from and can be used to calculate the likelihood of any sequence efficiently.' sampled from which? Moreover, I don't think this is a contribution. \n2. Line 197, why the log-likelihood of the variant can be used as the fitness predictor? What's the intuition behind this? or can you provide the related reference?\n3. Line 227, Comparison to baselines. There are many methods for fitness prediction in directed evolution, such as CLADE and CLADE2.0. Why not compare with them?\n4. In Table 1, What's the result of PoET if no ensemble?\n", "limitations": "This paper presents fitness prediction as the only downstream task of PoET. It is better to show more tasks to verify the effectiveness of PoET.\n", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper proposes a novel Transformer layer that models order-dependence between tokens within sequences and order-independence between sequences.\n- This paper provides a detailed analysis of experiments in the supplementary material.\n- The proposed 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.", "weaknesses": "- As the title stated, the core of this paper is a generative model of protein families, but they only do a downstream task, ie, fitness prediction. In this way, I suggest incorporating 'fitness prediction' into the title.\n- The missing comparison with some fitness prediction baselines.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive assessment of our work and for their constructive comments and questions. These are addressed below.\n\n- As the title stated, the core of this paper is a generative model of protein families, but they only do a downstream task, ie, fitness prediction. In this way, I suggest incorporating 'fitness prediction' into the title.\n\nWe consider generative evaluations of the model in terms of the perplexity of generating heldout sequences (Figure 4, 10) as well as examine samples from the model (Appendix M.2). Furthermore, we have expanded this analysis as discussed in the global response. That said, we do focus primarily on the evaluation of our model in terms of its ability to predict variant fitness. We will revise the manuscript to better emphasize this point.\n\n- Line 64, the contribution 'PoET can be sampled from and can be used to calculate the likelihood of any sequence efficiently.' sampled from which? Moreover, I don't think this is a contribution.\n\nWhat we mean by this is that samples can be efficiently drawn from the distribution over sequences modeled by PoET. PoET can also be used to calculate closed form likelihoods of sequences. This is in contrast to, say, energy based models which are generative models but that can only provide un-normalized likelihoods and are inefficient to sample from, or to GANs which can be sampled from but do not offer any way to estimate likelihoods. Specifically in the variant effect prediction space, most models considered are not proper generative models offering the ability to draw samples and calculate likelihoods efficiently. Thus, it is an important property of our model as we mention in that section.\n\n- Line 197, why the log-likelihood of the variant can be used as the fitness predictor? What's the intuition behind this? or can you provide the related reference?\n\nThe log-likelihood or some un-normalized or approximate version thereof is used by every unsupervised variant effect predictor that we are aware of. The idea is that natural protein sequences must be evolutionarily fit and, therefore, by learning the generative distribution of these sequences we are capturing this density. The probability of a sequence variant is then reflective of fitness, because we assume that more fit sequences are more likely to be observed in nature (examples include refs [3, 4, 9, 10, 11, 13, 23] in main text).\n\n- Line 227, Comparison to baselines. There are many methods for fitness prediction in directed evolution, such as CLADE and CLADE2.0. Why not compare with them?\n\nThese methods appear to be primarily supervised variant effect predictors, whereas we only consider the unsupervised variant effect prediction problem here and compare with other likelihood-based methods. These methods also only report results for four datasets, not for ProteinGym, a far more comprehensive collection of DMS datasets used here. In fact, for its unsupervised component, it appears that CLADE2.0 uses evolutionary scores from DeepSequence VAE, MSA Transformer, profileHMMs, and ESM-1v, all of which we already compare with, or compare with a better version of,and dramatically outperform.\n\n- In Table 1, What's the result of PoET if no ensemble?.   \n\nWe report results without ensembling in Appendix Table 4.\n"}, {"review_id": "P26faBKO2v", "reviewer": "Reviewer_UQRE", "summary": "This paper proposed an autoregressive generation pre-trained model of protein families. The models are trained over the sequences-of-sequences being organized by a set of specific protein sequences. It utilized a shared in-sequence position encoder to capture conditioning among sequences in an order independent manner, and thus is able to generalize to large context lengths. This paradigm can help to explore the correlation among sub-sequences to improve the performance of generation, especially when lack of sufficient multiple sequence alignment supervised training data.\n\nExperiments were conducted over DMS data sets to show the effectiveness of the proposed method.", "questions": "1.Transformer’s attention mechanism still suffer from high time and memory consumption, please give the time and space complexity of the proposed transformer structure.\n2.How to grantee the diversity of generated sequence? In autoregressive models, the maximum probability guided generation usually suffer from simplex result. \n3.Generally, the hyper-parameters in ensembling will affect the evaluation performance. Does this mean a fair comparison in Table 1.\n4.In ablation study, increasing the context length even results in negative performance. This is anti-intuitive. Is this phenomenon due to the decoding strategy (See question 2) or longer sequence weaken the effect of invert count sampling?\n5.How to perform ensemble to fuse the PoET and the other baselines, what is the motivation of ensemble? It seems unfair to compare with other baselines w/o ensemble.\n\nTypos:\n213: All combinations of these are parameters are used", "limitations": "Please refer to weakness.\n\nThe organization and representation (typos) can be further improved. For example, it would be better to give a brief organization/introduction of experimental study at the beginning of Section 5. ", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1.Proposed an autoregressive generation model which is trained in the manner of sequences-of-sequences, which is easy to append more protein sequence and explore the correction among protein to guide the generation.\n\n2.Shared absolute position encoder is employed in self-attention across different sub sequences to relieve the impact of sub-sequence ordering on next amino acid generation.", "weaknesses": "1.Although shared inter sequence position encoder can relive the impact of input sub-sequence order, the generation procedure is still order-dependent since Pr(next amino acid|s1, s2, ...). For example, if we arrange all short sequences at the begging of the combined sequence, it would be difficult to generated longer sequence such as insertion-based variant. It is better to explore such impact or validate whether the generation is robust to the input order by careful experimental study.\n\n2.Lack of deep insight or analysis of the results, even when some results might be anti-intuitive, such as longer input sequence results in negative performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their consideration of our paper, and for their questions and comments. We note that there seems to be some misunderstanding regarding certain aspects of our method and experiments, which we hope to address below in addition to the reviewer’s questions. We will also clarify these points in the revised manuscript.\n\n- Although shared inter sequence position encoder can relive the impact of input sub-sequence order, the generation procedure is still order-dependent since Pr(next amino acid|s1, s2, ...). For example, if we arrange all short sequences at the begging of the combined sequence, it would be difficult to generated longer sequence such as insertion-based variant. It is better to explore such impact or validate whether the generation is robust to the input order by careful experimental study.\n\nThe model is not sensitive to the lengths and ordering of sequences in the way suggested. The model is able to generate novel insertions and the relative positional encodings only act as a prior on the alignment between sequences in multihead attention. The sequence clusters the model is trained and evaluated on have large numbers of indels between sequences and have high length variability. Please see our global response comment for more details.\n\n- Lack of deep insight or analysis of the results, even when some results might be anti-intuitive, such as longer input sequence results in negative performance.\n\nAs pointed out by the other reviewers, we provide extensive analysis of the model and prompt construction methodology and discussion of these results in the appendix. That being said, we cannot explain every phenomena found in these experiments. In particular, we suspect that the model capacity and context length results are both related to the misspecification problem of variant function prediction and density estimation as discussed in [1] and Appendix D.2 Lines 399-412. Hopefully these will prove fruitful for us or others to explore in future work.\n\n- Transformer’s attention mechanism still suffer from high time and memory consumption, please give the time and space complexity of the proposed transformer structure.\n\nWe utilize flash attention [2], so the model requires O(N^2) time and O(N) memory where N is the length of the sequence. \n\n- How to grantee the diversity of generated sequence? In autoregressive models, the maximum probability guided generation usually suffer from simplex result.\n\nSampling from the model yields high diversity results and this diversity can be controlled using well known methods like nucleus sampling [3].\n\n- Generally, the hyper-parameters in ensembling will affect the evaluation performance. Does this mean a fair comparison in Table 1.\n\nThere are no additional hyper-parameters in our ensemble. It is a simple average over each prompt. Other settings such as homolog retrieval method, sampling homologs, context length, etc, are tuned based on validation set performance (lines 204-208) and are not specific to the ensemble. Furthermore, several other methods use ensembles over large numbers of language models, whereas we only use one.\n\n- In ablation study, increasing the context length even results in negative performance. This is anti-intuitive. Is this phenomenon due to the decoding strategy (See question 2) or longer sequence weaken the effect of invert count sampling?\n\nWe think this is probably due to the misspecification problem between variant fitness prediction and density estimation. Increasing the context length does generally improve the generative performance of the model (as measured by perplexity on heldout sequence clusters), but hurts performance on variant effect prediction.\n\n- How to perform ensemble to fuse the PoET and the other baselines, what is the motivation of ensemble? It seems unfair to compare with other baselines w/o ensemble.\n\nWe simply average them together (Appendix H). This is why we report the PoET + TranceptEVE ensemble in a separate section, consistent with the TranceptEVE hybrid model, and separate from the other non-multi-method ensembles.\n\n- Typos: 213: All combinations of these are parameters are used\n\nWe will correct this in the revised manuscript.\n\n[1] Weinstein, Eli, et al. \"Non-identifiability and the Blessings of Misspecification in Models of Molecular Fitness.\" Advances in Neural Information Processing Systems 35 (2022): 5484-5497.\n\n[2] Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv [Cs.LG]. Retrieved from http://arxiv.org/abs/2205.14135\n\n[3] 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.\n"}]

Author rebuttal:
Dear Reviewers,

We thank all of you again for your positive comments and constructive suggestions and questions. Here, we present additional analyses addressing some common themes in the reviews. This discussion is intended to be viewed with the figures in the attached review supplement PDF. We will include these results in the revised manuscript.

Additional analysis shows that PoET generates high quality, high diversity sequences.
- We have performed additional analysis of the sequences generated in the chorismate example and include an additional example on lysozyme. We have calculated the maximum sequence identity between our generated sequences and any natural protein and also folded the generated sequences using AlphaFold2 to show predicted structural conservation and pLDDT. PoET generates high diversity sequences that are predicted to be structurally similar, supporting the quality of sequences generated by PoET (Figure 1a for chorismate, Figure 1b for lysozyme).
- We have added an additional analysis of lysozyme sequences generated by PoET, including a comparison with lysozyme sequences generated by the fine-tuned ProGen model from Madani et al [0]. PoET generates higher diversity sequences than ProGen and these sequences are predicted to be more structurally conserved than sequences sampled from ProGen at similar levels of diversity. These structures are also predicted with high confidence by AF2, having high pDLLTs (>90% for almost all PoET generated sequences).

PoET was trained and evaluated on indel rich protein sequence clusters, has superior perplexities for these heldout sequences, and generates plausible indels in our chorismate mutase example.
- 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 2a). A histogram of the columns by percent gap shows that highly gappy columns are the most common (Figure 2b, left). Across all of the validation clusters (Figure 2b, 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 these indel rich families. It outperforms our baseline transformer, ProGen, and profileHMM, achieving better perplexities at each number of provided homologues (see Tables below and Figures 4 and 10 in the manuscript).
- In our chorismate mutase generative example, PoET generates sequences with novel indels and low sequence identity to natural chorismate mutases that are predicted to fold into the conserved chorismate structure with high pLDDT according to AlphaFold2 (Figure 3).

PoET achieves low perplexities on heldout Uniref50 clusters, outperforming ProGen and profileHMMs, even when conditioning on no or a very small number of homologs.
- We also evaluate heldout perplexity of ProGen continuing from a prompt sequence (as suggested by Reviewer jcwZ), where we find that PoET achieves better perplexity than ProGen without any conditioning, and when conditioned on the cluster seed sequences (Review Supplement Table 1). We note that our holdout set and ProGen’s holdout set are not the same, so sequences from our heldout Uniref50 clusters may occur in the ProGen training set and PoET still outperforms ProGen in this analysis. 

| Model  | # Sequences Conditioned On | Mean Perplexity | Std. Dev |
|--------|----------------------------|-----------------|----------|
| ProGen |                          0 |       15.591333 | 5.524855 |
| ProGen |                          1 |       14.788089 | 4.022846 |
| PoET   |                          0 |       14.545277 | 2.941460 |
| PoET   |                          1 |       10.907637 | 2.817105 |
Review Supplement Table 1: Perplexity evaluation on our heldout Uniref50 clusters. ProGen and PoET are either conditioned on no homologs (0 sequences conditioned on) or the seed sequence from the cluster (1 sequence conditioned on) and perplexities are calculated on the remaining sequences. This evaluation only includes sequences with length <512, due to length limits for ProGen.

- PoET generalizes well from a very small number of homologs. Conditioning on only 5 homologs, PoET already achieves a perplexity of 7.9 on unseen family members, in contrast to 17.0 for the PSSM and 13.7 for the profileHMM (Review Supplement Table 2). This demonstrates that PoET is a significantly better generative modeling option for families with small numbers of homologs and generalizes remarkably well from small numbers of sequences.
| Model | # Sequences Conditioned On | Perplexity |
|-------|----------------------------|------------|
| PSSM  |                          1 |  18.541033 |
| PSSM  |                          5 |  16.982748 |
| PSSM  |                         10 |  15.912713 |
| HMM   |                          1 |  20.137338 |
| HMM   |                          5 |  13.716971 |
| HMM   |                         10 |  11.888966 |
| PoET  |                          1 |  10.173127 |
| PoET  |                          5 |   7.922363 |
| PoET  |                         10 |   7.252593 |
Review Supplement Table 2: Median perplexities of sequences from heldout Uniref50 clusters after conditioning on 1, 5, or 10 cluster members with a PSSM, HMM, or PoET. All sequences from all heldout clusters are included in this analysis. 

We will incorporate these additional analyses into the revised manuscript to better support the generative modeling capabilities of PoET.

[0] Madani, A., Krause, B., Greene, E.R. et al. Large language models generate functional protein sequences across diverse families. Nat Biotechnol (2023). https://doi.org/10.1038/s41587-022-01618-2


Meta-review: This submission draws unanimous praise from the reviewers for its novelty and clarity. Here are the key takeaways based on the collective feedback:

* One of the standout features is strong performance achieved by the proposed model architecture on the zero-shot fitness prediction tasks from ProteinGym. Notably, this was accomplished using a simple yet powerful model architecture. The authors present various ablations, shedding light on the ideal prompt engineering strategy. 

* Novel Contributions: While building upon existing techniques, the paper introduces the inventive approach of generating MSA over individual sequences. This novelty accentuates the paper's credentials, especially for an application-driven submission. The introduction of a new Transformer layer that captures both order-dependence within sequences and order-independence between sequences. The results presented, particularly the superior performance of the proposed PoET over existing models on deep mutational scanning datasets in ProteinGym

Given the positive consensus among reviewers, the submission is suitable for acceptance at NeurIPS.
Decision: Accept (poster)