Title: PoET: A generative model of protein families as sequences-of-sequences
Venue: NeurIPS 2023
Accepted: yes
Keywords: protein fitness prediction

transformer

retrieval

language model

MSA

generative model

protein engineering
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-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. 1

1 Introduction

Proteins carry out most of the biological functions at the molecular level of life. The function of a protein is encoded by its specific amino acid sequence and the three-dimensional structure that the sequence folds into. Engineering proteins for novel and enhanced function is a key problem in pharmaceuticals and biotechnology and involves designing novel sequences or modifying existing natural proteins for these purposes. Deep mutational scans and directed evolution experiments have been used to successfully design novel proteins [1-3], but can be costly and difficult to implement, which makes these experimental methods inapplicable for many proteins and functions of interest. Accurate computational models of sequence-function relationships can narrow down the protein sequence search space, reduce the need for expensive experiments, and enable the design of more novel proteins and functions.

Protein language models have emerged as promising methods for understanding and designing protein sequences [4-7]. In particular, generative models offer a natural way to produce new protein designs. By training on large corpuses of natural protein sequences, these models learn evolutionary constraints on sequence space. They can then be used either to generate realistic sequences directly by sampling [8, 9], or to identify promising protein sequence variants by predicting the relative fitness of the variants of interest using the sequence likelihoods as a proxy [4, 10, 11].

Traditionally, family-specific models learn evolutionary constraints specific to the protein family of interest by training on a multiple sequence alignment (MSA) of homologous sequences [12, 13]. However, this is ineffective for protein families with few sequences due to the lack of sufficient training data and inability to exploit information across families. These models also assume that MSAs are accurate, and cannot model novel insertions or deletions (indels) not present in the training MSA. More recently, unconditional large protein language models [5, 10, 11] have been developed. Trained on all known natural protein sequences, unconditional protein language models generalize across protein families. However, these unconditional single-sequence models cannot be easily directed to generate a protein from a specific family of interest, and underperform family-specific models for relative fitness prediction. Hybrid models such as Tranception [11] and TranceptEVE [14] combine unconditional language models with family-specific models to enable specialization to protein families. Nonetheless, it is unclear how to use these models to generate sequences with novel indels, and predictions from the family-specific models do not directly benefit from transfer learning across families.

Here, we propose the P r o tein E volutionary T ransformer (PoET), an autoregressive generative model of whole protein families that addresses these limitations. By learning to generate sets of related proteins as sequences-of-sequences across tens of millions of natural protein sequence clusters, PoET is able to generalize about evolutionary processes across protein families, and avoids issues related to conditioning on MSAs. In order to capture conditioning between sequences in an order independent manner (the order of sequences within a family is arbitrary) and to generalize to large context lengths, we propose a novel Transformer layer (§3.1.2) that models order-dependence between tokens within sequences and order-independence between sequences. PoET has the following properties:

PoET can be used as a retrieval-augmented protein language model by conditioning the model on sequences from any family of interest. This also allows PoET to be used with any sequence database and to incorporate new sequence information without retraining.

PoET is a fully autoregressive generative model, able to generate and score novel indels in addition to substitutions, and does not depend on MSAs of the input family, removing problems caused by long insertions, gappy regions, and alignment errors.

By learning across protein families, PoET is able to extrapolate from short context lengths allowing it to generalize well even for small protein families.

PoET can be sampled from and can be used to calculate the likelihood of any sequence efficiently.

We demonstrate these properties and show that PoET outperforms existing protein language models and evolutionary sequence models for variant effect prediction in extensive experiments on the 94 deep mutational scanning datasets in ProteinGym. PoET improves substitution effect prediction across proteins of all MSA depths, and also improves effect prediction of indels. A simple weighted average ensemble of PoET with existing methods further improves performance both across MSA depths and in sequences with large numbers of mutations. Furthermore, when used for generation, PoET controllably produces diverse, novel sequences that more closely match natural statistical constraints, resulting in better folding pLDDTs, than other generative protein language models. We expect PoET to become integral to future protein mutagenesis efforts.

2 Related Work

Evolutionary sequence models are well established methods in biological sequences analysis. To model protein families, these models search large protein sequence databases for homologs, align the positions of these homologs in an MSA, and then fit statistical sequence models to the MSA. Common models include site independent models [15], profile HMMs [16], and coupling models [17]. Newer variants incorporate higher order correlations between sequence positions by training a V AE [13] or by building phylogentic trees [12]. These approaches are often referred to as "alignment-based" and must be fit on a family-by-family basis, requiring large numbers of members to generalize. A significant limitation of these models is that they assume the MSA is an accurate model of the evolutionary process generating the sequences, when in fact, MSA algorithms inevitably make alignment errors; regions with long insertions or lots of gaps can be particularly problematic. Furthermore, they cannot model novel insertions or deletions (indels) that are not present in the training MSA.

Unconditional protein language models that do not condition on homologs at inference have emerged as powerful methods for understanding and generating protein sequences. Both bidirectional models [4, 6, 18-20] and autoregressive generative models [5, 10, 11] have demonstrated competitive performance for variant function prediction. The latter type of model has the advantage of being able to score indels, but both cannot integrate evolutionary context not present in the trained model parameters. In contrast to family-specific evolutionary sequence models trained on sequences derived from a specific protein family [13, 17, 21], these models are trained on large protein databases [22, 23] that span all known sequences. This enables them to learn evolutionary constraints that generalize across families to improve predictions for small families with few homologs, but they generally underperform family-specific models for larger families.

Conditional protein language models fit between the unconditional and family-specific paradigms. Only a few works have explored this direction to date. Masked language models of whole MSAs are able to integrate evolutionary context directly for conditioning site predictions [24], but are unable to model insertions in the alignment. Ram and Bepler [25] use an encoder-decoder framework to generate new sequences conditioned on an MSA, which removes the insertion limitation of Rao et al. [24], but still requires conditioning on aligned input sequences. Notin et al. [11] combine predictions from an unconditional language model and an alignment-based model and show that integrating retrieval-based methods with protein language models can improve variant function prediction performance. However, the reliance on an alignment-based model means that the combined model is still limited by the constraints of MSAs.

Retrieval-augmented language models have shown impressive results in natural language processing, especially on Question Answering tasks. These models incorporate a database search as part of the text generation process in order to generate new text conditioned on prototypes found in the database [26-32]. In this way, they are conceptually similar to the conditional protein language models above. Retrieval-augmented approaches have the advantage of not requiring the entire training corpus to be encoded within the model parameters and the ability to easily integrate new data without retraining by simply adding it to the retrieval database. We adopt a similar approach with PoET in order to realize similar benefits. However, we build on well-established, fast, and accurate protein search tools for retrieval and frame the protein sequence generation problem as a sequence-of-sequences problem to incorporate retrieved-sequence conditioning, a fundamentally different paradigm than that employed by current retrieval-augmented models in natural language processing.

3 The Protein Evolutionary Transformer (PoET)

PoET is an autoregressive generative model of the distribution over protein families, where each family is generated as a sequence-of-sequences. Specifically, it models the distribution P ( X = x ) , where x = s 1 , s 2 , ..., s n is the concatenation of n sequences s i from the same family, and each sequence s i = s i, 1 , s i, 2 , ..., s i,L i is a sequence of L i amino acids padded by a start and end token. For example, below is a sequence of three protein sequences of lengths 4 , 6 , and 5 with start token denoted by $ and stop token denoted by * :

When referring to a sequence-of-sequences, s i , which has one index i , refers to a sequence of tokens the i th sequence in the sequence-of-sequences, whereas s i,j , which has two indices i, j , refers to one token, the j th token of the i th sequence. We use x to denote the full sequence-of-sequences.

PoET generates each token in a sequence x one at a time, decomposing the probability of x as

The order of the individual sequences in a sequence-of-sequences is arbitrary, and we propose a novel Transformer-based architecture to exploit this order invariance.

P=p(s1s1,s)=P(s31s1, ,3,<)

j=1

P1,1P1,2P1,3P1,4P2,1

P2,2P2.3P2,4P2,5

P2,6

P3,1P3,2P3.3P3,4P3,5

目

目

目

目

目

目

目

目

目

目

目

目

DecodedProbabilities

FeedForward

gij

Sequence-Of-SequencesCausal Self-Attention

Nx

TieredTransformerDecoderLayer

Per-Sequence

Per-Sequence

Per-Sequence

Embeddings

$

M

K

*

$

M

H

I

P

*

$

P

M

V

S1

S2

S3

Input x

(a) PoET generates sets of homologous proteins as sequences-of-sequences. The equation in red demonstrates how to compute the conditional probability of a sequence given homologous sequences.

Queries

$

M

1

S1

K

2

0

*

3

2

$

7

3

0

M

1

0

L

2

L

0

H

2

1

1

2

1

S2

I

3

2

3

2

1

P

4

3

2

1

4

3

2

1

0

*

5

4

3

2

5

4

3

2

0

$

M

K

*

$

M

H

P

*

S1

S2

S1:

*

TokensAssociated

WithPosition

S2:

$

M

H

I

P

*

1

2

3

45

Position

(b) Inter-sequence relative position encoding scheme used in sequence-of-sequences self-attention that is invariant to sequence ordering. ( Left ) Relative positions between tokens in a sequence-of-sequences composed of two sequences. ( Right ) Tokens associated with absolute positions.

3.1 Model Architecture

We propose a variant of the common Transformer decoder layer (Figure 1a) to capture order invariance between sequences while preservering order-dependence between tokens within sequences. We accomplish this using two attention modules: (1) a within-sequence module in which the representation at each position of each sequence is updated based on attending only to the other tokens within this sequence, and (2) a between-sequence module in which the representation at each position of each sequence is updated based on attending to other sequences within the sequence-of-sequences (§3.1.2). This tiered approach is critical for capturing long-range dependencies between sequences and uniquely allows our model to extrapolate to much longer context lengths than used during training (§5.2.2), improving sequence generation and performance on downstream tasks (Appendix H). Our full PoET model is a stack of these layers with causal self-attention.

3.1.1 Input Embedding

The input sequence x = s i,j , i ∈ 1 ..n, j ∈ 1 ..L i of amino acids and start/stop tokens is first converted into a sequence of continuous embeddings h i,j by mapping each token to its learned embedding:

where AA is the set of 20 standard amino acids, and W is a matrix of learnable embeddings of dimension d .

3.1.2 Tiered Transformer Decoder Layers

Next, the embeddings h i,j are transformed by N layers of TieredTransformerDecoderLayer s (Appendix Algorithm 1), a novel layer for processing a sequence-of-sequences that is invariant to the order of the individual sequences, and extrapolates to context lengths substantially longer than the training context length.

The TieredTransformerDecoderLayer is composed of two phases. In the first phase, causal self-attention is applied independently to each sequence h i of the input sequence-of-sequences, transforming them into new sequences f i = PerSequenceSelfAttn ( h i ) . Relative positional information is encoded by applying Rotary Positional Encodings (RoPE) [33] to the queries and keys before applying self-attention in the standard manner; the absolute position for f i,j is j .

The second phase applies causal self-attention to the entire sequence-of-sequences by concatenating the individual f i from the previous layer into one sequence before applying self-attention: g i,j = SequenceOfSequencesSelfAttn ([ f <i ; f i,<j ]) . In order to make self-attention in this phase invariant to sequence order, we adopt a simple but novel inter-sequence relative positional encoding scheme (Figure 1b): for g i,j the absolute position is j . Just as in the first phase, the absolute position for tokens in the i th sequence g i is independent of the position i of the sequence in the sequence-ofsequences. Thus, the positional information encoded by RoPE in this layer alone does not distinguish between the positions of tokens in different sequences. For example, the relative position between the first token of the first sequence f 1 , 1 and the first token of the second sequence f 2 , 1 is 0 . The fact that these two tokens come from two different sequences is encoded by the first phase, which operates on the two sequences independently. This inter-sequence relative positional encoding scheme has two useful properties:

it encodes the fact that amino acids at similar absolute positions in homologous proteins are more likely to be drawn from the same distribution

it limits the maximum relative position encoding needed to the number of tokens in an individual protein sequence 2 , rather than the number of tokens in a sequence-of-sequences, allowing the model to generalize to longer sequences-of-sequences than seen during training

3.1.3 Decoded Probabilities

Lastly, the output from the last TieredTransformerDecoderLayer , g i,j , is decoded into token probabilities by applying a linear transformation P ( s i,j | s <i , s i,<j ) = p i,j ( s i,j ) = Linear ( g i,j ) . Here, p i,j is a vector of probabilities, one for each distinct token ∈ AA ∪{ START , STOP } , and p i,j ( s i,j ) is the probability of the token s i,j according to p i,j .

3.2 Training Data

Models were trained on 29 million sets of homologous sequences. Each set corresponds to a sequence in UniRef50 Version 2103, and contains all its homologs in UniRef50 found using Diamond [34]. We removed any sets with fewer than 10 homologs. To avoid overfitting on promiscuous sequences which may belong to a large number of sets, we sample each set with weight inversely proportional to the size of the set ("inverse count" sequence weighting). The order of sequences within a sequenceof-sequences is randomized to promote order invariance. See Appendix B for more details.

4 Protein Variant Fitness Prediction

Protein variant fitness prediction is the task of assigning a score to each sequence in a set of variants { v 1 , v 2 , ..., v n } of a target sequence t that accurately reflects the relative fitness of the variants. A protein variant v i can be any sequence with a limited number of substitutions, insertions, and/or deletions relative to the target t that the experimenter believes may have improved fitness. Fitness

Evaluatecorrelation

3

withexperimentaldata

Homologs

Spearman Correlation: 0.60024

MPGFY .. INYAD

10.0

7.5

SequenceDB

MELVG..DRLAR

5.0

2.5

METQT..SSWCH

PoET

0.0

2.5

TargetSequence

MELVG..DRLAR

7.5

MAADG..HSEPR

METQT..SSWCH

10.0

980-970-960-950-940

Log-Likelihood

-930-920

Variants

Likelihoods

Computelikelihood

WT:MAADG..HSEPR

P(WT|Prompt)

ofeachvariantinthe

V,:AAADG..HSEPR

mutational landscape

P(V,IPrompt)

ofinterest

:MALDKG..HSEPA

P(V-1Prompt)

V:MAA-G..-SEPR

P(V1Prompt)

refers to the value of any property of a protein sequence related to function that the experimenter is interested in optimizing e.g. thermostability, enzymatic activity, etc.

Benchmarking using deep mutational scans Deep mutational scanning (DMS) is a method for measuring the fitness properties of thousands to millions of variants of a protein using next-generation sequencing techniques [1]. Data from these experiments have been used to evaluate fitness prediction methods. We 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 et al. [11] for tuning hyperparameters.

Fitness Prediction with PoET PoET predicts the fitness of a variant as the conditional loglikelihood of the variant v i given a set of sequences S homologous to the target t :

The set S is retrieved by searching a large database of proteins such as UniRef100 [23] for sequences homologous to t . The homologs form a diverse set of sequences that define the protein family of interest. By conditioning on these sequences, PoET can infer evolutionary constraints on the protein family to improve fitness prediction. The sequences are conditioned on in an arbitrary order. Figure 2 illustrates the evaluation of PoET for variant fitness prediction on one DMS dataset.

Based on validation set performance, we selected the best method for (1) retrieving homologous sequences (Appendix C), (2) subsampling and filtering the homologous sequences to a reasonable context length for efficient inference (Appendix D), and (3) 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 homologs, and the final fitness prediction score is obtained by averaging the conditional log-likelihoods across subsamples of the full set of retreived homologous sequences:

The subsets S j are drawn by sampling with sequence weights [17], maximum sequence identity to the target t in { 1 . 0 , 0 . 95 , 0 . 90 , 0 . 70 , 0 . 50 } , and total tokens (i.e. context length) in { 6144 , 12288 , 24576 } . All combinations of these parameters are used, for a total of N ensemble = 15 .

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 in ProteinGym. We compare PoET to the best existing alignment-based models, unconditional and conditional protein language models, and hybrid models that combine multiple approaches (Results Summary: Table 1, Per Dataset: Figures S1, S2 ). A brief overview of the baseline methods is provided in Appendix F. Statistical significance was assessed using a paired t-test. For a fair comparison, we ran the relevant baselines with all homologous sequence retrieval methods considered for use with PoET (Appendix C), and present the results for each baseline method paired with its best retrieval method. We do not tune parameters for 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 to TranceptEVE L, the best performing baseline method (Table 1). PoET improves average Spearman correlation ∆¯ ρ = 0 . 013 , but the difference falls just short of statistical significance p = 0 . 05029 . Although the methods perform similarly, PoET offers a number of advantages. Inference with PoET is substantially faster when considering on the order of tens to hundreds of thousands of variants (Appendix J), which allows users to more quickly accomplish common tasks such as computing the fitness of all single site substitutions. By avoiding the costly step of training a new model for each protein family, users can more easily experiment with different ways of defining the protein family. For example, users can retrieve the homologous sequences using different sequence search programs and settings, and prior knowledge can be incorporated by selecting the subset of most relevant homologous sequences e.g. sequences from the taxon of interest, or that are known to be functional (Appendix K). Since PoET is able to generalize across protein families, whereas the EVE component of TranceptEVE L cannot, it is less critical that the protein family of interest is defined by a large number of homologous sequences (Appendix C).

On the other hand, if one is willing to spend more time on inference and the protein family is well defined a priori , better fitness predictions can be obtained by ensembling PoET with TranceptEVE L (Results: Table 1, Methods: Appendix G). An ensemble of PoET with TranceptEVE L performs significantly better than TranceptEVE L alone ( ∆¯ ρ = 0 . 021 , p < 7 e6 ). 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; GEMMEisthousands of times faster than either PoET or TranceptEVE L [12]. The main disadvantage is that GEMME is unable to score indels.

Indels Benchmark On the indels benchmark, the best performing baseline model is TranceptEVE M, a variation of TranceptEVE L that uses a protein language model with fewer parameters ( ∆¯ ρ = 0 . 05 ). PoET and the ensemble of PoET with TranceptEVE L both outperform TranceptEVE L ( ∆¯ ρ ∈ [0 . 044 , 0 . 055] ), and perform comparably to TranceptEVE M ( ∆¯ ρ ∈ [ -0 . 006 , 0 . 005] ). However, no differences between the aforementioned models are statistically significant. One advantage of PoET is that it is able to not only score indels, but also generate sequences with indels (§6). It is also able to score and generate insertions not present in the MSA.

5.2 Characterizing the PoET architecture

5.2.1 Training Distribution and Model Size

We train variations of PoET with different training distributions and model sizes and investigate their effect on fitness prediction (Figure 3).

Context Length We trained 57M parameter versions of PoET for up to 3 days on 7 x A100 GPUs with three context lengths: 4K, 8K, and 16K. Notably, these context lengths far exceed the 1K or 2K context lengths typically used to train other protein and natural language models [5, 11, 24, 36, 37]. We reason that longer context lengths are needed for the model to observe a sufficient number of homologs; even with a context length of 8K, for a sequence of length 500 , which is the length of the typical human protein in ProteinGym, the model would only observe ~ 16 of possibly thousands of homologs. Interestingly, while a context length of 8K performs better than a context length of 4K, increasing the context length further to 16K had a neutral or negative effect on variant effect prediction. It does improve perplexity of heldout sequences.

Sequence Set Weighting Next, we trained an 8K context length model using the naive strategy of sampling UniRef50 sequence sets uniformly instead of using our 'inverse count' strategy that weights sequence sets based on their size (§3.2). We find that the "inverse count" sampling strategy substantially outperforms the naive strategy. Other methods that also train on sequence sets like MSA Transformer [24] may also benefit from this sampling strategy.

Model Size Lastly, using the optimal parameters found thus far, we trained 57M, 201M, and 604M parameter models to explore the effect of model size. Models were trained until validation performance plateaued. Increasing the model size to 201M parameters signficantly improves performance, but increasing the size further to 604M appears to have a neutral effect. We chose the 201M version as the standard PoET model.

5.2.2 Model Architecture

To determine the importance of using TieredTransformerDecoderLayer s (§3.1.2) instead of regular Transformer layers, we trained, on the same task, a regular 201M parameter RoPE-based Transformer that ignores the sequence-of-sequences structure of the input. We compared the generative capabilities of PoET and the Transformer by measuring the perplexity at context lengths both within and beyond the training the context length (Figure 4).

We trained both models with a context length of 8K for 500K steps. As expected based on previous studies[38], the RoPE-based Transformer does not generalize to context lengths beyond the training context length. In contrast, PoET generalizes to ~8x the training context length (61K), which is the longest we tested. PoET also performs better on context lengths within the training context length,

0.46

#Params

ContextLength

Weighting

0.45

57M

8,192

InverseCount

Spearman

57M

4,096

Inverse Count

57M

16,384

InverseCount

0.44

57M

8,192

Uniform

201M

8,192

Inverse Count

0.43

604M

8,192

InverseCount

0

24

48

72

0

4896144192240288336384432

Training Hours

Training Hours

32

8.0

1

1

1

7.5

1

1

Perplexity

7.0

Architecture

Regular Transformer

6.5

PoET Transformer

Train ContextLength

8K

6.0

16K

32K

4096

16

61440

Inference Context Length(log-scale)

but by a smaller margin. Next, we finetuned both models with 16K and 32K context lengths for 50K and 30K steps respectively. While the perplexity of the Transformer improved significantly at the 16K and 32K context lengths, the Transformer still does not generalize to context lengths beyond the training context length, and underperforms all variants of PoET at all context lengths. The Transformer also underperforms PoET for variant fitness prediction (Appendix H).

6 PoET generates novel sequences that preserve structure within a protein family

PoET's state-of-the-art performance on variant fitness prediction demonstrates that it assigns higher likelihood to regions of sequence space with high fitness variants. This suggests that PoET can be used to directly generate high fitness variants belonging to a protein family by conditioning PoET on sequences from the protein family and sampling from the resulting conditional distribution. Direct generation of variants makes the exploration of higher order mutants computationally tractable as sequence space grows exponentially large with number of mutations, making it impossible to explore this space by scoring all such variants.

To evaluate PoET's ability to generate sequences from specific protein families, we sampled 1,000 sequences each from PoET conditioned on homologs from two protein families, phage lysozymes (PF00959) and chorismate mutases with activity in E. coli (Appendix K.2), and examined the novelty of the generated sequences and their predicted structural conservation. Sequence novelty was measured as the max sequence identity to any natural sequence (of the same protein family). Predicted

A:Functional Chorismate Mutases

B:PhageLysozymes

100

More

95

80

Confident

pLDDT

口

Prediction

60

85

40

08

0.5

0.6

0.7

0.8

0.9

1.0

0.6

0.8

1.0

Max seq id to a natural sequence

Max seq id to a natural sequence

Model

Progen

PoET

1.0

TM-score compared to

0.95

fanaturalsequence

0.9

Higher

PDB:1ECM

0.90

solved structure

0.8

Structural

Similarity

0.85

0.7

0.80

0.6

0.75

0.5

0.5

0.6

0.7

0.8

0.9

1.0

0.6

0.8

1.0

Max seq id to a natural sequence

Max seq id to a natural sequence

More Novel Sequences

structural conservation was measured as the max TM-score between the predicted AlphaFold2 [39] structure of the generated sequence and the solved structure of any natural sequence. For both protein families, PoET generates high diversity sequences that are predicted with high confidence to be structurally similar, supporting the quality of sequences generated by PoET (Figure 5).

Furthermore, when tasked to generate phage lysozymes, we find that sequences generated by PoET are more diverse and better preserve structure at similar levels of sequence novelty than sequences generated by ProGen [8] (Figure 5B), a 1.2B parameter model fine-tuned specifically for generating phage lysozymes. As many sequences generated by ProGen have been experimentally confirmed to be functional phage lysozymes, these results suggest that PoET generates functional sequences with even greater diversity even without fine-tuning. Since fine-tuning ProGen is critical for achieving good variant effect prediction [8], and a version of ProGen fine-tuned on chorismate mutases is not available, we did not compare PoET and ProGen for generating novel chorismate mutases.

7 Conclusion

PoET is a Transformer-based autoregressive generative model of whole protein families. By framing family generation as a sequence-of-sequences generation problem, we are able to train across tens of millions of protein sequence clusters to encode fundamental rules of protein evolution into the PoET model. This enables the model to generalize to protein families unseen during training and extrapolate from small numbers of conditioned-upon sequences. The sequence-of-sequences generative framework allows PoET to be used as a retrieval-augmented language model, generating new sequences conditioned on a set of sequences representing the family or other properties of interest. We demonstrate that PoET improves over other protein language models and evolutionary sequence models for variant fitness prediction across a wide range of deep mutational scanning datasets. PoET also enables efficient sequence generation and the generative distribution can be controlled via conditioning. Phage lysozyme- and chorismate mutase-like sequences sampled from PoET are novel and predicted to fold with high confidence. PoET can be backed by other sequence databases and naturally improves as databases grow without the need for retraining. We anticipate that PoET will become a fundamental component of ML-enabled protein design in the future.

References

Douglas M. Fowler and Stanley Fields. Deep mutational scanning: a new style of protein science. Nature Methods , 11(8):801-807, Aug 2014. ISSN 1548-7105. doi: 10.1038/nmeth.3027. URL https://doi.org/10.1038/nmeth.3027 .

Ryan E Cobb, Ran Chao, and Huimin Zhao. Directed evolution: Past, present and future. AIChE J. , 59(5):1432-1440, May 2013.

Louis Tiefenauer and Sophie Demarche. Challenges in the development of functional assays of membrane proteins. Materials (Basel) , 5(11):2205-2242, November 2012.

Joshua Meier, Roshan Rao, Robert Verkuil, Jason Liu, Tom Sercu, and Alex Rives. Language 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. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper/2021/file/ f51338d736f95dd42427296047067694-Paper.pdf .

Erik Nijkamp, Jeffrey Ruffolo, Eli N. Weinstein, Nikhil Naik, and Ali Madani. Progen2: Exploring the boundaries of protein language models, 2022. URL https://arxiv.org/abs/ 2206.13517 .

Tristan Bepler and Bonnie Berger. Learning the protein language: Evolution, structure, and function. Cell Systems , 12(6):654-669.e3, 2021. ISSN 2405-4712. doi: https://doi.org/ 10.1016/j.cels.2021.05.017. URL https://www.sciencedirect.com/science/article/ pii/S2405471221002039 .

J. Dauparas, I. Anishchenko, N. Bennett, H. Bai, R. J. Ragotte, L. F. Milles, B. I. M. Wicky, A. Courbet, R. J. de Haas, N. Bethel, P. J. Y . Leung, T. 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 learning&#x2013;based protein sequence design using proteinmpnn. Science , 378(6615): 49-56, 2022. doi: 10.1126/science.add2187. URL https://www.science.org/doi/abs/ 10.1126/science.add2187 .

Ali Madani, Ben Krause, Eric R Greene, Subu Subramanian, Benjamin P 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.

William P. Russ, Matteo Figliuzzi, Christian Stocker, Pierre Barrat-Charlaix, Michael Socolich, Peter Kast, Donald Hilvert, Remi Monasson, Simona Cocco, Martin Weigt, and Rama Ranganathan. An evolution-based model for designing chorismate mutase enzymes. Science , 369 (6502):440-445, 2020. doi: 10.1126/science.aba3304. URL https://www.science.org/ doi/abs/10.1126/science.aba3304 .

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.05789 , 2022.

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://proceedings.mlr.press/ v162/notin22a.html .

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-2619, 08 2019. ISSN 0737-4038. doi: 10.1093/molbev/msz179. URL https: //doi.org/10.1093/molbev/msz179 .

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.1038/ s41586-021-04043-8. URL https://doi.org/10.1038/s41586-021-04043-8 .

Pascal Notin, Lood Van Niekerk, Aaron W Kollasch, Daniel Ritter, Yarin Gal, and Debora Susan Marks. TranceptEVE: Combining family-specific and family-agnostic models of protein sequences for improved fitness prediction. In NeurIPS 2022 Workshop on Learning Meaningful Representations of Life , 2022. URL https://openreview.net/forum?id=l7Oo9DcLmR1 .

P C Ng and S Henikoff. Predicting deleterious amino acid substitutions. Genome Res. , 11(5): 863-874, May 2001.

S R Eddy. Profile hidden markov models. Bioinformatics , 14(9):755-763, 1998.

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):128-135, Feb 2017. ISSN 1546-1696. doi: 10.1038/nbt.3769. URL https://doi.org/10.1038/nbt.3769 .

Tristan Bepler and Bonnie Berger. Learning protein sequence embeddings using information from structure. In International Conference on Learning Representations , 2019. URL https: //openreview.net/forum?id=SygLehCqtm .

Ethan C. Alley, Grigory Khimulya, Surojit Biswas, Mohammed AlQuraishi, and George M. Church. Unified rational protein engineering with sequence-based deep representation learning. Nature Methods , 16(12):1315-1322, Dec 2019. ISSN 1548-7105. doi: 10.1038/ s41592-019-0598-1. URL https://doi.org/10.1038/s41592-019-0598-1 .

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.

Adam J. Riesselman, John B. Ingraham, and Debora S. 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/s41592-018-0138-4. URL https://doi.org/10.1038/ s41592-018-0138-4 .

Jaina Mistry, Sara Chuguransky, Lowri Williams, Matloob Qureshi, Gustavo A Salazar, Erik L 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 , 49(D1):D412-D419, 10 2020. ISSN 0305-1048. doi: 10.1093/nar/gkaa913. URL https://doi.org/10.1093/nar/gkaa913 .

The UniProt Consortium. UniProt: the Universal Protein Knowledgebase in 2023. Nucleic Acids Research , 51(D1):D523-D531, 11 2022. ISSN 0305-1048. doi: 10.1093/nar/gkac1052. URL https://doi.org/10.1093/nar/gkac1052 .

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/rao21a.html .

Soumya Ram and Tristan Bepler. Few shot protein generation, 2022. URL https://arxiv. org/abs/2204.01168 .

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.1162/tacl_a_00030. URL https://aclanthology.org/Q18-1031 .

Tatsunori B Hashimoto, Kelvin Guu, Yonatan Oren, and Percy S Liang. A retrieve-and-edit framework for predicting structured outputs. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 31. Curran Associates, Inc., 2018. URL https://proceedings.neurips. cc/paper/2018/file/cd17d3ce3b64f227987cd92cd701cc58-Paper.pdf .

Omar Khattab and Matei Zaharia. Colbert: Efficient and effective passage search via contextualized 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 9781450380164. doi: 10.1145/3397271.3401075. URL https://doi.org/10.1145/3397271.3401075 .

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.net/forum?id=Ghk0AJ8XtVx .

Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Ming-Wei Chang. Realm: Retrieval-augmented language model pre-training. In Proceedings of the 37th International Conference on Machine Learning , ICML'20. JMLR.org, 2020.

Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George Bm Van Den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, Diego De Las Casas, Aurelia Guy, Jacob Menick, 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.mlr.press/v162/borgeaud22a.html .

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. URL http://arxiv.org/abs/2208.03299 .

Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, 2021. URL https://arxiv.org/ abs/2104.09864 .

Benjamin Buchfink, Klaus Reuter, and Hajk-Georg Drost. Sensitive protein alignments at treeof-life scale using diamond. Nature Methods , 18(4):366-368, Apr 2021. ISSN 1548-7105. doi: 10.1038/s41592-021-01101-x. URL https://doi.org/10.1038/s41592-021-01101-x .

Milot Mirdita, Konstantin Schütze, Yoshitaka Moriwaki, Lim Heo, Sergey Ovchinnikov, and Martin Steinegger. Colabfold: making protein folding accessible to all. Nature Methods , 19 (6):679-682, Jun 2022. ISSN 1548-7105. doi: 10.1038/s41592-022-01488-1. URL https: //doi.org/10.1038/s41592-022-01488-1 .

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 1877-1901. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper/2020/file/ 1457c0d6bfcb4967418bfb8ac142f64a-Paper.pdf .

Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, 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

Lazaridou, Arthur Mensch, Jean-Baptiste 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, Iason 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/2112.11446, 2021. URL https://arxiv.org/abs/2112.11446 .

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://openreview.net/forum?id=R8sQPpGCv0 .

John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A. A. Kohl, Andrew J. Ballard, Andrew Cowie, Bernardino Romera-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, Koray Kavukcuoglu, Pushmeet Kohli, and Demis Hassabis. Highly accurate protein structure prediction with alphafold. Nature , 596(7873):583-589, Aug 2021. ISSN 1476-4687. doi: 10.1038/s41586-021-03819-2. URL https://doi.org/10.1038/s41586-021-03819-2 .

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 on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pages 4596-4604. PMLR, 10-15 Jul 2018. URL https://proceedings.mlr.press/v80/ shazeer18a.html .

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. doi: 10.1186/1471-2105-11-431. URL https://doi.org/10.1186/1471-2105-11-431 .

Martin Steinegger and Johannes Söding. Mmseqs2 enables sensitive protein sequence searching for the analysis of massive data sets. Nature Biotechnology , 35(11):1026-1028, Nov 2017. ISSN 1546-1696. doi: 10.1038/nbt.3988. URL https://doi.org/10.1038/nbt.3988 .

Vineet Sangar, Daniel J. Blankenberg, Naomi Altman, and Arthur M. Lesk. Quantitative sequence-function relationships in proteins based on gene ontology. BMC Bioinformatics , 8(1): 294, Aug 2007. ISSN 1471-2105. doi: 10.1186/1471-2105-8-294. URL https://doi.org/ 10.1186/1471-2105-8-294 .

Eli N Weinstein, Alan Nawzad Amin, Jonathan Frazer, and Debora Susan Marks. Nonidentifiability and the blessings of misspecification in models of molecular fitness. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https://openreview.net/forum?id= CwG-o0ind6t .

Martin Steinegger and Johannes Söding. Clustering huge protein sequence sets in linear time. Nature Communications , 9(1):2542, Jun 2018. ISSN 2041-1723. doi: 10.1038/ s41467-018-04964-5. URL https://doi.org/10.1038/s41467-018-04964-5 .

Michael Remmert, Andreas Biegert, Andreas Hauser, and Johannes Söding. Hhblits: lightningfast iterative protein sequence searching by hmm-hmm alignment. Nature Methods , 9(2): 173-175, Feb 2012. ISSN 1548-7105. doi: 10.1038/nmeth.1818. URL https://doi.org/ 10.1038/nmeth.1818 .

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, 01 2013. ISSN 0737-4038. doi: 10.1093/molbev/mst010. URL https: //doi.org/10.1093/molbev/mst010 .

Nadav Brandes, Dan Ofer, Yam Peleg, Nadav Rappoport, and Michal Linial. ProteinBERT: a universal deep-learning model of protein sequence and function. Bioinformatics , 38(8): 2102-2110, 02 2022. ISSN 1367-4803. doi: 10.1093/bioinformatics/btac020. URL https: //doi.org/10.1093/bioinformatics/btac020 .

Lin Li, Esther Gupta, John Spaeth, Leslie Shing, Tristan Bepler, and Rajmonda Sulo Caceres. Antibody representation learning for drug discovery, 2022.

Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. In International Conference on Learning Representations , 2020. URL https://openreview.net/forum?id=rygGQyrFvH .

A Detailed description of TieredTransformerDecoderLayer

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 diamond/uniref50 -f 6 -header -k 200000 -max-hsps 1 -e 0.001 -p 96 -o output.tab

The command returns, for each sequence in UniRef50, a set containing all its putative homologs in UniRef50. We call each such set a 'Diamond-UniRef50 Cluster'. Diamond was used over other homology search tools due to its high performance (>100x speed of BLAST).

To form a training example, sequences are sampled without replacement from a 'Diamond-UniRef50 Cluster' until the total number of tokens reaches a predetermined limit, and then concatenated to form a sequence-of-sequences. Each sampled UniRef50 sequence is replaced with a UniRef100 sequence by sampling a random UniRef100 sequence from the same UniRef50 cluster as the UniRef50 sequence being replaced. The UniRef100 sequences are randomly sampled with weight inversely proportional to the size of the UniRef90 clusters they belong to. Each UniRef100 sequence is sampled at most once in each sequence-of-sequences.

As a final data augmentation, the order of the tokens in a sequence-of-sequences is reversed with probability 50% (i.e. all sequences in a sequence-of-sequences are ordered from either N-terminus to C-terminus only or C-terminus to N-terminus only). This augmentation has been shown to improve the performance of other protein language models [11] .

Following this sampling procedure, the order of sequences in a sequence-of-sequneces is random, which promotes order invariance. To evaluate the sensitivity of PoET to the order of sequences, we compare heldout sequence perplexities after conditioning on 10 sequences ordered either 1) randomly as described above, 2) shortest-to-longest, or 3) longest-to-shortest. We also compare the full joint log likelihoods of these sequences of 10 sequences with the same orderings (Table 2).

We find that sequence ordering has little-to-no effect on next sequence perplexity and joint log likelihoods. Ordering sequences shortest-to-longest changes next sequence perplexity by only 0.00753 on average and has a relative difference of < 0 . 5% from random for both next sequence perplexity and joint log likelihood.

Table 2: Comparison of (a) next sequence perplexities and (b) joint log likelihoods log p ( s 1 , s 2 , . . . , s n ) , of sequences of 10 sequences with different orderings from heldout Uniref50 sequence clusters. We report mean perplexity and joint log likelihood as well as mean difference from random and relative difference from random for shortest-to-longest and longest-to-shortest orderings with the standard deviation across clusters in parenthesis. Relative difference is calculated as ( a -r ) /r where a is the ordered value and r is the value with random ordering.

B.2 Optimizer and learning rate

We used the AdaFactor optimizer [40] with initial learning rate 1 e2 , 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 3.

C Homologous Sequence Retrieval Methods

We experimented with two methods for retrieving homologous sequences from UniRef100, JackHMMer [41] and MMseqs2 [42]. Both programs return an MSA of homologs. For JackHMMer, we used the MSAs provided by ProteinGym [11], which were constructed from UniRef100 Version 2104. For MMseqs2, we used the protocol defined by ColabFold [35], and retrieved homologs from UniRef100 Version 2202. The main advantage of the ColabFold protocol is that it is substantially faster than the ProteinGym protocol (minutes vs up to hours), whereas the ProteinGym protocol may be more sensitive.

We evaluated all relevant baselines and PoET with homologs retrieved by both protocols (Table 4). Note that in the main result table (Table 1), we report the results of each method with the best protocol based on validation set performance. We replicate the findings of Laine et al. [12], who show that GEMME performs similarly with both protocols, and also find that most models perform similarly regardless of the protocol. The notable exception is EVE, which performs significantly worse. This is

RetrievalMethod

ProteinGym

60-

50

40

Count

30

20

10

0

0.0

0.5

1.0

1.5

2.0

2.5

Redundancy log(N/Neff)

most likely because the MSAs returned by the ColabFold protocol are substantially less redundant (Figure 6). Thus they do not contain enough sequences to properly train EVE models, which are trained from scratch for each protein family and can contain millions of parameters. Removing redundancy filters in the protocol could improve results for EVE; we leave this to further research.

Given the poor performance of EVE with the ColabFold protocol, we did not also evaluate the EVE ensemble or TranceptEVE M/L with ColabFold, presuming that the performance of both would also be worse with ColabFold because they both depend on EVE. Furthermore, we did not evaluate the performance of Tranception with ColabFold because Tranception relies on homologs via the site independent model, which performs similarly with both protocols.

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:

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]

the set of sequences retrieved will be biased towards those that have been chosen to be sequenced by humans, or by natural processes that may not reflect relative fitness [17]

To resolve the second issue, we adopt the sequence weighting scheme used by Hopf et al. [17], which has been shown to work well in many previous studies [4, 11, 13, 40]. The weighting scheme is detailed in Appendix D.5.

To address the first issue, existing methods have used sequence identity as a simple heuristic for filtering out irrelevant or misleading sequences. Laine et al. [12] apply a max dissimilarity filter of 0.8 sequence identity to remove sequences that are less likely to be homologous, and a max similarity filter of 0.98 sequence identity to remove sequences that are redundant with the target. Meier et al. [4] explore using a variety of max similarity filter thresholds and find that a threshold of 0.9 or 0.75 tends to be optimal, depending on the target. Both Laine et al. [12] and Meier et al. [4] also limit the total number of sequences in a subsample, which helps to limit the computational cost of inference, and ensemble the results from applying their model to multiple subsamples to obtain better predictions.

To determine the best way to subsample and filter homologous sequences for PoET, we evaluate the performance of the model on the ProteinGym validation set with (1) various max dissimilarity and (2) similarity thresholds for filtering, and (3) various context lengths (i.e. total number of tokens used to compute the conditional log-likelihood (Equation 3), including the tokens of the variant sequence being scored) for limiting the computational cost of inference.

D.1 Maximum Sequence Dissimilarity

First, we explored the impact of the maximum sequence dissmilarity threshold. We evaluate PoET with (1) homologous sequences retrieved using either the ProteinGym or the ColabFold protocol, (2) a context length of 8192 , and (3) various max sequence dissimilarity thresholds (Figure 7). When retrieving homologous sequences using the ProteinGym protocol, a threshold of 0.7 is optimal for both the validation and substitution benchmarks, and is significantly better than no filtering i.e. using a threshold of 1.0 (substitutions benchmark: ∆¯ ρ = 0 . 024 , p < 0 . 01 ). Interestingly, filtering was particularly helpful for prokaryotic datasets, suggesting that the ProteinGym protocol may be prone to finding false positives in this taxon. In contrast, for the ColabFold protocol, no filtering based on max dissimilarity is necessary; there is essentially no difference in performance when filtering with a threshold of 0.7.

D.2 Maximum Sequence Similarity and Context Length

Next, using the optimal maximum sequence dissimilarity thresholds, we evaluated PoET while jointly varying the max similarity threshold ∈ { 1 . 0 , 0 . 95 , 0 . 90 , 0 . 70 , 0 . 50 } and context length ∈ { 6144 , 12288 , 24576 } (Figure 8). When averaged across datasets, we found that PoET generally performed better when some sequences are filtered i.e. max similarity threshold < 1 . 0 , and with smaller context lengths ( 6144 or 12288 ). However, when looking at individual datasets (Figure 9), we find a wide variety of behaviors. On some datasets, increasing the context length consistently improves performance across max similarity thresholds and the optimal max similarity threshold is 1 . 0 (no filtering), whereas for others the opposite is true, and yet for others the behavior lies somewhere in between.

From these observations, we draw two main conclusions. First, although the average performance is lower at higher context lengths, the ability of PoET to improve in performance with increased context length on some datasets suggests that PoET is able to generalize to context lengths much longer than the training context length, and that decreased average performance at higher context lengths may be due to factors other than the model's ability to handle long context lengths. Weinstein et al. [44] suggest that optimal modeling of the sequences in an observed protein family may be detrimental to protein fitness prediction due to biases in the sequences we observe. Consequently, protein fitness prediction based on evolutionary models perform better when they do not or cannot model the family exactly e.g. via model misspecification. In the context of PoET then, limiting the max sequence similarity and context length when sampling sequences from the full set of retrieved homologs may

(a) Performance using ProteinGym protocol for retrieving homologous sequences

Validation

Taxon

0.450

Substitutions

Dissimilarity

0.8

Virus

Prokaryote

0.425

Human

0.6

Eukaryote

rho

0.400

Spearman

0.375

with

0.4

0.350

rho

0.2

Spearman

0.325

0.0

0.300

1.0

0.8

0.6

0.4

0.2

0.0

0.2

0.4

0.6

0.8

SpearmanrhowithMaxDissimilarity1.0

0.475

Validation

Taxon

0.7

0.450

Substitutions

0.8

Human

Virus

0.425

0.6

Prokaryote

Eukaryote

rho

0.400

Spearman

0.4

0.375

with

0.350

0.2

Spearman

0.325

0.0

0.300

1.0

0.8

0.6

0.4

0.2

0.0

0.2

0.4

0.6

0.8

MaxDissimilarity

SpearmanrhowithMaxDissimilarity1.0

(b) Performance using ColabFold protocol for retrieving homologous sequences

help avoid some of these biases while still providing the model information about the specific protein family of interest. Second, the large variance in the optimal max similarity threshold and context length on a per dataset level suggests that better sampling approaches, possibly dataset specific, may be an avenue for further improving protein fitness prediction.

D.3 Ensemble over subsamples

Given the variance in best max sequence similarity and context length parameters on the validation datasets, we explored ensembling (Equation 4) over subsamples where the subsamples are drawn with different values for max similarity and context length. Table 5 compares (1) not ensembling, (2) ensembling with N ensemble = 15 subsamples where the max sequence similarity and context length parameters are fixed to the best values on the validation set (Figure 8), and (3) ensembling subsamples with varied parameters (max similarity threshold ∈ { 1 . 0 , 0 . 95 , 0 . 90 , 0 . 70 , 0 . 50 } and context length ∈ { 6144 , 12288 , 24576 } ; 15 combinations in total). On the validation and substitutions benchmarks, we find that both types of ensembling improve performance. On the validation benchmark, the two ensembling methods perform similarly; fixed has a slight edge when using the ProteinGym protocol

(a) Performance using ProteinGym protocol for retrieving homologous sequences

1.0

0.46

0.95

rho

0.9

Spearmanrho

0.45

Spearmanr

0.46

0.7

0.5

0.44

Validation s

0.43

0.45

0.42

0.41

0.44

0.40

6144

12288

24576

6144

12288

24576

ContextLength

Context Length

MaxSimilarity

0.47

1.0

0.95

ValidationSpearmanrho

0.46

SubstitutionsSpearmanrho

0.47

0.9

0.7

0.5

0.45

0.46

0.44

0.43

0.45

6144

12288

24576

6144

12288

24576

ContextLength

ContextLength

(b) Performance using ColabFold protocol for retrieving homologous sequences

for retrieving homologs, and varied has a slight edge when using the ColabFold protocol. On the full substitutions benchmark, the varied method consistently performs better, and by a larger margin; this result may be less biased than the validation benchmark result because the parameters for the fixed method are tuned on the validation benchmark.

D.4 Summary

Based on validation benchmark performance of the PoET variants explored in the above sections, for our final PoET model, we choose the following parameters

Spearman

A: Increased context tends to improve performance

BLAT_ECOLX_Jacquier_2013

Q2N0S5_9HIV1_Haddox_2018

0.52

MaxSimilarity

0.72

0.5

0.51

0.71

0.7

0.70

0.50

0.9

0.95

0.69

0.49

1.0

0.68

0.67

0.48

0.66

B: Increased context

0.47

0.65

has neutral effect

6144

12288

24576

6144

12288

24576

CALM1_HUMAN_Weile_2017

PA_I34A1_Wu_2015

0.230

0.5

FSE'O

0.225

0.4

0.30

0.220

0.3

0.25

0.215

0.2

F

0.20

0.210

0.1

6144

12288

24576

6144

12288

24576

6144

12288

24576

RL401_YEAST_Roscoe_2013

DLG4_RAT_McLaughlin_2012

SPIKE_SARS2_Starr_bind_2020

0.56

0.63

0.34

0.54

0.62

0.32

0.61

0.52

0.60

0.30

0.28

0.50

0.59

0.26

0.58

0.48

0.24

0.57

0.46

0.56

0.22

6144

12288

24576

6144

12288

24576

6144

12288

24576

TPOR_HUMAN_Bridgford_S505N_2020

CCDB_ECOLI_Tripathi_2016

0.45

?

0.52

0.44

0.43

0.50

C: Increased context tends

0.42

to decrease performance

0.48

0.41

0.40

0.46

0.39

6144

12288

24576

6144

12288

24576

ContextLength

Retrieve homologous sequences using the ColabFold protocol

Do not filter homologous sequences by max dissimilarity when subsampling the set of homologous sequences

Ensemble the predictions of PoET across subsamples of the full set of homologous sequences with max similarity thresholds ∈ { 1 . 0 , 0 . 95 , 0 . 90 , 0 . 70 , 0 . 50 } and context lengths ∈ { 6144 , 12288 , 24576 } (15 combinations total)

D.5 Sequence Weighting During Homologous Sequence Sampling

Wesampled homologous sequences using the sequence weighting scheme described in Hopf et al. [17]. Specifically, sequences are sampled from a probability distribution where a sequence's probability is inversely proportional to its number of neighbors in the MSA of the homologous sequences. A target sequence is considered a neighbor of a query sequence if the sequence identity to the query is greater than 0.8. The sequence identity is computed based on the alignment in the MSA, and gaps in the query sequence are ignored.

D.6 Full length sequences vs aligned fragments

In the subsampling and filtering of homologous sequences experiments above, we condition on only the amino acids in the homologous sequences that are aligned to a position of the target sequence in the MSA of the homologous sequences i.e. a sequence in a subset S j (Equation 4) may not be the full length sequence as recorded in UniRef100, rather it is potentially a non-contiguous subsequence of the UniRef100 sequence. We find that using only these 'aligned fragments' performs similarly to using the full length sequences (Table 6), and use the former over the latter as the former can be derived directly from the MSA.

Table 6: Average Spearman's rank correlation on the validation and substitutions benchmarks using full length sequences vs aligned fragments sampled from MSAs of homologous sequences generated by the ColabFold protocol.

E MSA Depth

MSA depth measures the amount of sequence information the MSA contains about the target protein. We use the same definition and low/medium/high categorization as Notin et al. [11]. At a high level, the MSA depth N eff /L is computed as the number of non-redundant sequences N eff (at 0.8 sequence identity for non-viral proteins, and at 0.99 sequence identity for viral proteins) divided by the sequence length L . The MSA depth categories thresholds are - Low: N eff /L < 1 ; Medium: 1 < N eff /L < 100 ; High: N eff /L < 100 .

Although the MSA depth changes based on the method used to retrieve homologous sequences, we categorize proteins/datasets using the ProteinGym MSAs (1) in order to maintain consistency with the categorizations in the ProteinGym benchmark, which allows for easier comparison with existing work, and (2) because the MSA depth is highly correlated between the retrieval methods considered in this paper (Figure 10).

F Protein Variant Fitness Prediction Baselines

F.1 Overview of Baselines

We compare PoET to the following baseline models for protein variant fitness prediction, which includes the best existing alignment-based models, unconditional and conditional protein language models, and hybrid models that combine multiple approaches models:

4.0

3.5

3.0

2.5

2.0

1.5

1

2

3

4

5

F.1.1 Alignment-based Models

Alignment-based models first compute a MSA of homologous sequences of a given protein family, and compute sequence likelihoods based on MSA statistics or models trained on the MSA.

Site Independent Model The simplest alignment-based model; sequence likelihoods are computed from the frequency of amino acids at each position of the MSA, considered independently from other positions of the MSA.

GEMME[12] is a fast, scalable, and simple algorithm for predicting variant effect that models the 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 it cannot score indels.

EVE [13] (ensemble) is an ensemble of five Bayesian variational autoencoder trained on the MSA.

F.1.2 Unconditional PLMs

Unconditional PLMs are large language models trained on protein databases spanning many protein families e.g. UniRef [23], using either the causal or masked language modeling objective.

ESM-1v [4] (ensemble) is an ensemble of five 650M parameter Transformer-based masked language models trained on UniRef90.

ProGen2 [5] (ensemble) is an ensemble of five Transformer-based causal language models ranging between 150M and 6.4B parameters and trained on UniRef90 and BFD [45].

Tranception L (no retrieval) [11] is a 700M parameter Transformer-based causal language model trained on UniRef100.

F.1.3 Conditional PLMs

Conditional PLMs are large language models trained on sets of homologous protein sequences spanning many protein families.

MSA Transformer [24] (ensemble) is a 100M parameter masked language model trained on MSAs of all UniRef50 sequences. MSAs are generated using HHblits [46].

F.1.4 Hybrid

Hybrid models combine predictions from different types of models.

Tranception L [11] combines predictions from Tranception L (no retrieval) and a site independent model in a manner that allows novel indels to be scored by aligning the sequence to be scored to the MSA used by the site independent model.

TranceptEVE [14] combines predictions from Tranception and EVE in a manner that allows novel indels to be scored by using EVE to 'define a family-specific prior distribution over amino acids at each sequence position that is independent from the particular protein sequence'. It is the best performing baseline method.

F.2 Methodology

We run all baselines using the same methodology as Notin et al. [11, 14], with two exceptions:

Homologous Sequence Retrieval Method We explored using two methods, the ProteinGym and ColabFold protocols, for retrieving the homologous sequences used by baselines (Appendix C, Table 4), and compare all methods using the best homologous sequence retrieval method (Table 1).

Site Independent Model We present results for the site independent model using sequence weighting (Appendix D.5) and filtering sequences with more than 0.8 (in terms of sequence identity) dissimilarity to the target sequence. These are common preprocessing steps used to improve the performance of other models for protein variant fitness prediction (Appendix D), and are the same parameters used to define the retrieval-based site independent prior in Tranception [11] and TranceptEVE [14]. We believe that this variation of the site independent model better reflects the model class's ability to predict protein variant fitness.

G Ensembling PoET with Baseline Methods for Protein Variant Fitness Prediction

To ensemble PoET with a baseline method for protein variant fitness prediction, we compute a weighted sum of the score predicted by PoET and by the baseline method i.e.

where ˆ F Baseline ,i is the score predicted by the baseline method, ˆ F PoET ,i is the score predicted by PoET, and ˆ F PoET+Baseline ,i is the ensemble prediction score.

The weight α is determined by optimizing performance over the ProteinGym validation set using ternary search over the range [1 e3 , 1 e 4] .

H Effect of Architecture on Protein Variant Fitness Prediction

Figure 11 compares the performance of the regular RoPe-based Transformer and PoET on the ProteinGym validation benchmark. We use homologous sequences retrieved using the ColabFold protocol, and explore various max similarity thresholds and context lengths as in Appendix D.2. PoET consistently outperforms the Transformer. It also extrapolates to context lengths longer than the training context length (8K), whereas the Transformer does not.

(a) Performance using regular Transformer.

(b) Performance using PoET.

0.48

1.0

0.95

Spearmanrho

0.47

0.9

0.7

0.46

0.5

ValidationS

0.45

0.44

0.43

0.42

6144

8192

12288

24576

Context Length

MaxSimilarity

0.48

1.0

0.95

Validation Spearman rho

0.47

0.9

0.7

0.46

0.5

0.45

0.44

0.43

0.42

6144

12288

24576

Context Length

I Comparison of Models for Generating Sequences From a Protein Family

To benchmark PoET's ability to generate sequences from a protein family, we evaluated the mean perplexity of a protein sequence conditioned on a fixed number of tokens from other sequences in the same protein family across all 'Diamond-UniRef50 clusters' (Appendix B) in the validation set. PoET significantly outperforms two baseline models, a profile HMM and a regular Transformer (Figure 12).

14

Model

HMM

13

Regular Transformer

PoET Transformer

12

Perplexity

11

10

9

8

6

16384

32768

61440

Number of Tokens

PoET effectively models indel rich protein families The UniRef50 clusters in the validation set for perplexity evaluation are highly indel rich. An alignment for an exemplar cluster shows that there are large insertions and high diversity between sequences (Figure 13a). A histogram of the columns by percent gap shows that highly gappy columns are the most common (Figure 13b, left). Across all of the validation clusters (Figure 13b, right), this trend remains true, with alignments becoming

more indel rich when more sequences (longer context lengths) are considered. PoET achieves low perplexities on these heldout sequences, showing that it is a good generative model of indel rich families.

%Gaps

1

UPI00112DDEBD

UPI000D538369

UPI00156C41CF

A0A2K5Q4C2

H2PST6

MPACWSVLQPTHVLQIPEDLAACPTCPCSSLHVRSQAFLQHTVWFHLLAHCLPHLVDVMKTHLLQALLQASOVQEOEQVTFEDVVVDFTOEEWEQLTPAQRM

MPADVNFSQKPQVLG----------

MPA-------

-.....--NKSLKCAGOGSVSFKDVTVDFTRDEWLOLTGTORT

----PEEQDGSCRASVSFEDVTMDFSF

--GASAQVPVTFEDVAVYFSQ

SVODVT

VEFI

EWEGLEEWQKE

WQH

LGLVERT

23

102

50

% Gaps

UPI00112DDEBD

UPI000D538369

A0A2K5Q4C2

H2PST6

MYRDVVLENYSHLI

LYSDVMLDTFELLVSIGEAMS--NP

LYKEVMKENYONLDSLGT

LYQDL---

HLFS

TGSPTITPE】

O--NPDVIFRRLKEKEPCMEEAEFSHQRCQ

--PP

ERGEEPYIRDEPGSEEGGT-

AHCN

GKSSCSEIDDFKKRHKVTLFFLECPVN-

-IFP---DAA

125

UPI00156C41CF

LYRDVMLENYSHLVSVGCHLT--KPDVIFRLEOGEELWPLEGEFPDOGDOVLRDSLWHGLT-

---RLLSKAFGILILCLVSTAHOEWCTLDPAGWDOS

%Gaps

UPI00112DDEBD

UPI000D538369

A0A2K5Q4C2

H2PST6

RVTQORRKR-.....

K-IESR--

...-FLRPGVLQL-..-RPL.....

GLAQL-GPPFPS-

SVAQA-G--

--FLFHFPCICVVR---.

--VQWH

PSSI

PNGF---

--- DLGSLQPPPPR FK---

---.--LQRVNFFF

158

227

129

(a) Typical example of a MSA of a UniRef50 cluster. Bar plot above the MSA shows the percentage of gaps in each corresponding MSA column.

MSAs with<4096 tokens

0.40

0.2

Proportion of MSA Columns

0.35

Proportion ot MSA Columns

0.1

0.30

0.0

0.25

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.20

MSAs with<61440 tokens

0.15

0.6

0.4

0.10

0.2

0.05

0.0+

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

0.00

0.0

0.1

0.2

Proportion of MSA Column Composed of Gaps

0.30.4

0.5

0.6

0.7

0.8

0.9

1.0

Proportion of MSA Column Composed of Gaps

J Inference Speed of PoET vs TranceptEVE L

In the section, we compare the time it takes for PoET and TranceptEVE to predict the fitness of a set of protein variants on A100 GPUs. We begin by outlining the major steps required to perform inference for each method. For PoET, there are two main steps:

Retrieve homologs of the target protein using the ColabFold protocol

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:

Retrieve an MSA of homologs of the target protein using the ProteinGym protocol

Use the MSA to compute the site independent log prior

Train 5 EVE models on the MSA

Use the 5 EVE models to compute the EVE log prior

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 a lower bound on the number of variants that can be scored by PoET before a single variant can be scored by TranceptEVE by computing the number of variants that can be scored by PoET in the time it takes to train one of the EVE models required for TranceptEVE. The homolog retrieval time is irrelevant for this analysis since both methods require a homolog retrieval step, although we note that the ColabFold

protocol used to retrieve homologs by PoET is substantially faster than the ProteinGym protocol used by TranceptEVE.

We find that independent of sequence length, on a single A100 GPU, PoET is able to score approximately 50,000 variants in the time it takes to train one EVE model. Furthermore, since variant scoring with PoET can be parallelized whereas training an EVE model cannot, inference with PoET can be sped up further by a factor of the number of A100 GPUs available. This allows up to hundreds of thousands of variants to be scored by PoET with a reasonable amount of compute before TranceptEVE is able to score a single variant.

K Function-specific variant fitness prediction and sequence generation via prompt engineering

K.1 Function-specific variant fitness prediction

Thus far, we have evaluated the ability of PoET to predict the 'general' fitness of a protein variant, where 'general' fitness can refer to any property that is related to the function of the protein. But in practice, we are generally interested in optimizing specific properties of a protein. Since many properties of interest are correlated and together contribute to the general fitness of a protein (e.g. a more thermostable variant is also more likely to have higher expression or enzyme activity), optimizing general fitness is likely to optimize the true properties of interest, albeit indirectly. PoET is a useful tool for protein engineering because it is able to successfully predict general fitness by conditioning on and inferring evolutionary constraints from a diverse set of homologs of the target protein. This raises a natural question - can PoET be used to optimize specific properties of interest by conditioning on only the subset of relevant homologs that are known to or are predicted to display the specific properties of interest? It is difficult to answer this question in its full generality, since the process of selecting the subset of relevant homologs must be tailored to the specific target protein of interest, and there is an inevitable trade-off between learning from a smaller set of homologs that are more relevant, and a larger set of homologs from which there is more data to learn.

As a proof of concept, we demonstrate that PoET is able to learn function-specific evolutionary constraints for the target protein and property of interest in the chorismate mutase indels dataset [9] from ProteinGym. The dataset contains measurements of the catalytic activity, in E. coli, of 1130 natural chorismate mutase sequences, and 1618 designed chorismate mutase variants. The natural sequences are comprised of the target protein, a chorismate mutase found in E. coli, and homologs of the target protein found using the PSI-BLAST program for sequence search. The designed variants were selected by Monte Carlo sampling from a Potts model trained on an MSA of the natural sequences. This data presents an ideal scenario for selecting the subset of most relevant homologs; we can simply select the subset of natural homologs that are measured to be functional. In the absence of such data, we could instead use predictions from another model, or other relevant known attributes of the sequences e.g. if we would like to optimize for activity at high temperatures, we may select only the homologs from thermophiles.

On the chorismate mutase dataset, we find that the catalytic activity of designed chorismate mutase variants is better predicted when PoET is conditioned on only the subset of functional natural sequences rather than all natural sequences ( ∆ ρ = 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 a BERT-like protein masked language model ( ∆ ρ = 0 . 06 , Table 8). Such embeddings have been shown to be highly predictive of a variety of protein properties [20, 48, 49] and provide a strong baseline. These fully supervised methods are trained on more data than PoET because they train on the measured catalytic activities of all the natural sequences, whereas PoET is simply conditioned on positively labeled natural sequences and does not have access to the measured activities. This enables PoET to be used with assays that only measure binary endpoints rather than continuous values.

Natural (rho=-0.03533)

Designed (rho=0.40926)

1.00

0.75

Fitness

0.50

Measured

0.25

0.00

0.25

-0.50

0.75

-150-125-100-75

-50

25

0

25

-150-125-100-75

-50

-25

0

25

Predicted Fitness

Predicted Fitness

Natural (rho=0.11743)

Designed (rho=0.36388)

1.00

0.75

Fitness

0.50

Measured

0.25

0.00

-0.25

-0.50

0.75

-250-225-200-175-150-125-100-75

Predicted Fitness

-250-225-200-175-150-125-100-75

Predicted Fitness

Natural(rho=0.40608)

Designed(rho=0.50855)

1.00

0.75

Fitness

0.50

Measurec

0.25

0.00

0.25

-0.50

0.75

-250

-200

-150

-100

-50

-250

-200

-150

-100

-50

Predicted Fitness

Predicted Fitness

Natural(rho=0.61857)

Predicted Fitness

1.00

0.75

Fitness

0.50

Measurec

0.25

0.00

0.25

-0.50

-0.75

-300

-250

-200

-150

-100

-50

-300

-250

-200

-150

-100

-50

Designed (rho=0.56772)

Predicted Fitness

(d) Predicted fitness using PoET conditoned on functional natural sequences. Natural sequences with measured fitness >0.42 are considered functional by Russ et al. [9].

K.2 Function-specific sequence generation

Conditioned on the functional natural sequences, we used PoET to generate 1000 novel putative chorismate mutases using nucleus sampling [50] with p = 0 . 9 .

Figure 17 visualizes the sequence space spanned by natural chorismate mutases, chorismate mutases functional in E. coli, and sequences sampled from PoET when conditioned on the latter sequences. We find that the sequences sampled from PoET are concentrated in clusters around the known functional sequences, suggesting that the generated sequences sample from the functional subspace. The generated sequences are highly diverse and not simply a recapitulation of the natural sequences; the maximum sequence identity to any natural sequence is between 0.4 and 1.0, with a mode around 0.55 (Figure 15a). Nevertheless, the entropy at each position of an MSA of the functional natural sequences and an MSA of the generated sequence are closely matched (Figure 15b), indicating that the generated sequences reflect at least the first order evolutionary constraints of functional chorismate mutases.

High confidence (pLDDT>90) AlphaFold2 [35] predicted structures of generated sequences aligned to the experimentally solved structure of the target chorismate mutase (PDB: 1ECM) show that structure is preserved despite significant divergences in sequence identity; two examples are shown in Figure 16 and the full distributions of pLDDTs and TM-scores are shown in Figure 5A.

0.20

0.15

Proportion

0.10

0.05

0.00

0.4

0.5

0.6

0.7

0.8

0.9

1.0

Max sequence identity to any natural sequence

(a) Distribution of max seq id between generated and natural sequences

4.0

3.5

3.0

2.5

Entropy

2.0

1.5

1.0

0.5

Natural Sequences

0.0

0

20

40

60

80

Position

(b) Per position entropy 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 p ∈ { 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 on some proteins with low MSA depth, it 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 as 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 a couple of mutations. For

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.

Natural sequence space

Functional subspace

N Supplementary Tables

O Supplementary Figures

GFP_AEQVI_Sarkisyan16

P53_HUMAN_Kotler18

DLG4_HUMAN_Faure21

DLG4_RAT_McLaughlin12

BRCA1_HUMAN_Findlay18

CAPSD_AAV2S_Sinai_substitutions21

UBE4B_MOUSE_Starita13

TPOR_HUMAN_Bridgford_S505N20

P53_HUMAN_Giacomelli_NULL_Etoposide18

P53_HUMAN_Giacomelli_WT_Nutlin18

A4_HUMAN_Seuma21

YAP1_HUMAN_Araya12

P53_HUMAN_Giacomelli_NULL_Nutlin18

SPG1_STRSG_Olson14

GCN4_YEAST_Staller_induction18

TADBP_HUMAN_Bolognesi19

POLG_HCVJF_Qi14

PABP_YEAST_Melamed13

MTH3_HAEAE_Rockah-Shmuel15

R1AB_SARS2_Flynn_growth22

TRPC_SACS2_Chan17

PA_I34A1_Wu15

HIS7_YEAST_Pokusaeva19

P84126_THETH_Chan17

GAL4_YEAST_Kitzman15

A0A2Z5U3Z0_9INFA_Doud16

SC6A4_HUMAN_Young21

SRC_HUMAN_Ahler_CD19

GRB2_HUMAN_Faure21

PTEN_HUMAN_Mighell18

A0A2Z5U3Z0_9INFA_Wu14

UBC9_HUMAN_Weile17

TPMT_HUMAN_Matreyek18

ADRB2_HUMAN_Jones20

Q2N0S5_9HIV1_Haddox18

VKOR1_HUMAN_Chiasson_abundance20

B3VI55_LIPST_Klesmith15

A0A192B1T2_9HIV1_Haddox18

HSP82_YEAST_Mishra16

RL401_YEAST_Roscoe13

NCAP_I34A1_Doud15

KCNH2_HUMAN_Kozek20

SPIKE_SARS2_Starr_expr20

DYR_ECOLI_Thompson_plusLon19

POLG_CXB3N_Mattenberger21

PTEN_HUMAN_Matreyek21

TRPC_THEMA_Chan17

RL401_YEAST_Mavor16

A0A140D2T1_ZIKV_Sourisseau_growth19

VKOR1_HUMAN_Chiasson_activity20

I6TAH8_I68A0_Doud15

RL401_YEAST_Roscoe14

A4D664_9INFA_Soh_CCL14119

MSH2_HUMAN_Jia20

C6KNH7_9INFA_Lee18

ENV_HV1B9_DuenasDecamp16

ENV_HV1BR_Haddox16

HSP82_YEAST_Flynn19

SPIKE_SARS2_Starr_bind20

TPK1_HUMAN_Weile17

REV_HV1H2_Fernandes16

MK01_HUMAN_Brenan16

SYUA_HUMAN_Newberry20

SCN5A_HUMAN_Glazer19

A0A1I9GEU1_NEIME_Kennouche19

BLAT_ECOLX_Firnberg14

BLAT_ECOLX_Stiffler15

BLAT_ECOLX_Jacquier13

A4GRB6_PSEAI_Chen20

CP2C9_HUMAN_Amorosi_activity21

NUD15_HUMAN_Suiter20

KKA2_KLEPN_Melnikov14

CP2C9_HUMAN_Amorosi_abundance21

AMIE_PSEAE_Wrenbeck17

CCDB_ECOLI_Tripathi16

IF1_ECOLI_Kelsic16

BLAT_ECOLX_Deng12

CCDB_ECOLI_Adkar12

RASH_HUMAN_Bandaru17

SUMO1_HUMAN_Weile17

F7YBW8_MESOW_Aakre15

ESTA_BACSU_Nutschel20

TAT_HV1BR_Fernandes16

CALM1_HUMAN_Weile17

0.0

0.2

0.4

0.6

0.8

1.0

Spearman rho

Low

Medium

High

ESM-1v

MSA Transformer

PoET

PoET + TranceptEVE L

PTEN_HUMAN_Mighell_deletions18

HIS7_YEAST_Pokusaeva_indels19

CAPSD_AAV2S_Sinai_indels21

BLAT_ECOLX_Gonzalez_indels19

P53_HUMAN_Kotler_deletions18

B1LPA6_ECOSM_Russ20

A0A1J4YT16_9PROT_Davidi20

0.0

0.2

0.4

0.6

0.8

1.0

Spearman rho

Tranception L

TranceptEVE M

PoET

PoET + TranceptEVE L

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1.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.\n\nWEAKNESSES:\n1.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.\n\nQUESTIONS:\n1.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\n\nLIMITATIONS:\nPlease 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\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- The missing comparison with some fitness prediction baselines.\n\nQUESTIONS:\n1. 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\nLIMITATIONS:\nThis paper presents fitness prediction as the only downstream task of PoET. It is better to show more tasks to verify the effectiveness of PoET.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\n1. 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).\n\nLIMITATIONS:\nEven though some limitations are touched upon in the main text, a thorough discussion of limitations is lacking.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nCurrent 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nMajor 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\nLIMITATIONS:\nThe authors should add a paragraph to discuss the potential limitations and negative societal impacts of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n- 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\n\nSTRENGTHS:\n- 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\n\nWEAKNESSES:\n- 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\n\nQUESTIONS:\n- “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?\n\nLIMITATIONS:\nNot particularly discussed, would suggest adding a couple sentences about it, as per the questions above.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-15"}