PoET: A generative model of protein families as sequences-of-sequences

Summary

This paper proposed an autoregressive generation pre-trained model of protein families. The models are trained over the sequences-of-sequences being organized by a set of specific protein sequences. It utilized a shared in-sequence position encoder to capture conditioning among sequences in an order independent manner, and thus is able to generalize to large context lengths. This paradigm can help to explore the correlation among sub-sequences to improve the performance of generation, especially when lack of sufficient multiple sequence alignment supervised training data. Experiments were conducted over DMS data sets to show the effectiveness of the proposed method.

Strengths

1.Proposed an autoregressive generation model which is trained in the manner of sequences-of-sequences, which is easy to append more protein sequence and explore the correction among protein to guide the generation. 2.Shared absolute position encoder is employed in self-attention across different sub sequences to relieve the impact of sub-sequence ordering on next amino acid generation.

Weaknesses

1.Although shared inter sequence position encoder can relive the impact of input sub-sequence order, the generation procedure is still order-dependent since Pr(next amino acid|s1, s2, ...). For example, if we arrange all short sequences at the begging of the combined sequence, it would be difficult to generated longer sequence such as insertion-based variant. It is better to explore such impact or validate whether the generation is robust to the input order by careful experimental study. 2.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.

Questions

1.Transformer’s attention mechanism still suffer from high time and memory consumption, please give the time and space complexity of the proposed transformer structure. 2.How to grantee the diversity of generated sequence? In autoregressive models, the maximum probability guided generation usually suffer from simplex result. 3.Generally, the hyper-parameters in ensembling will affect the evaluation performance. Does this mean a fair comparison in Table 1. 4.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? 5.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. Typos: 213: All combinations of these are parameters are used

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

Please refer to weakness. The 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.

Summary

This paper proposes an autoregressive generative model (protein evolutionary transformer, PoET) of whole protein families. Current generative protein language models are either difficult to direct to produce a protein from a specific family of interest or must be trained on a large multiple sequence alignment (MSA) from the specific family of interest, making them unable to benefit from transfer learning across families. This model can incorporate new sequence information without retraining, generalize to large context lengths, and avoid issues related to conditioning on MSAs. They propose a novel Transformer layer that models order-dependence between tokens within sequences and order-independence between sequences. The advantages include that PoET can be used as a retrieval-augmented protein language model, 2) generate and score novel indels in addition to substitutions, and does not depend on MSAs of the input family, 3) extrapolate from short context lengths allowing it to generalize well even for small protein families.

Strengths

- This paper proposes a novel Transformer layer that models order-dependence between tokens within sequences and order-independence between sequences. - This paper provides a detailed analysis of experiments in the supplementary material. - The proposed PoET outperforms existing protein language models and evolutionary sequence models for variant effect prediction in extensive experiments on the 94 deep mutational scanning datasets in ProteinGym.

Weaknesses

- As the title stated, the core of this paper is a generative model of protein families, but they only do a downstream task, ie, fitness prediction. In this way, I suggest incorporating 'fitness prediction' into the title. - The missing comparison with some fitness prediction baselines.

Questions

1. Line 64, the contribution 'PoET can be sampled from and can be used to calculate the likelihood of any sequence efficiently.' sampled from which? Moreover, I don't think this is a contribution. 2. 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? 3. 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? 4. In Table 1, What's the result of PoET if no ensemble?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

This paper presents fitness prediction as the only downstream task of PoET. It is better to show more tasks to verify the effectiveness of PoET.

Summary

The authors of the paper present Protein Evolutionary Transformer (PoET), an autoregressive transformer-based model that is able to generate sets of related protein sequences, enabled by the proposed novel Transformer layer. This sequences-of-sequences generating method benefits from transfer learning, is able to handle novel indels, and does not need a multiple sequence alignment (MSA) as its input. The model can be conditioned on sequences of a protein family of interest for generation and scoring. The proposed method shows similar or superior performance compared to other models in fitness prediction tasks, while being faster without requiring retraining.

Strengths

This paper is very well-written, with clear explanations, great examples and figures, and a comprehensive Related Work section. The experiments are explained well, they align with the claims made in the paper, and the results are discussed adequately. Moreover, there's an abundance of extra information in the appendix.

Weaknesses

I enjoyed reading this paper, I do not have any major concerns. However, I did miss a "Limitations" section or something equivalent, either in the main text or in the supplementary. Moreover, I think some results (mainly those in the "Ablation" section 5.2) could be made stronger, for example by averaging performance over multiple runs and showing error bars, because the observed trends are not always that convincing. Finally, and this is more of a general issue, even though the results show improved/competitive performance, the average correlation values to experimental data is relatively low, i.e. around 0.5, which is still quite a weak correlation. The mismatch between density estimation and fitness prediction has been discussed before, for example in [1], and it might be worth discussing to some extent in this paper as well. [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.

Questions

1. As mentioned in "Weaknesses": Limitations section missing. 2. As mentioned in "Weaknesses": perhaps include some discussion on the (mis)match between density estimation and fitness prediction. 3. 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. Minor comments: 4. 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. 5. Do you report the dimensionality of embedding size $d$ somewhere? I might have missed it. 6. 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? 7. 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)? 8. Section 5.2.1: * 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? * 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. * 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. 9. Missing references: * Figure 1a is never referenced. * Paper reference missing for RoPE? * Not all appendices are referenced in the main text (B, E, F, I, and O are missing).

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Even though some limitations are touched upon in the main text, a thorough discussion of limitations is lacking.

Summary

Current generative protein language models focus on generating individual protein sequences and are not specifically trained to generate sequences for an entire protein family. The paper introduces a new autoregressive generative model that tackles this limitation by generating Multiple Sequence Alignments (MSAs) for the entire family. The model utilizes a TieredTransformerDecoder to model the within-sequence interactions in an order-dependent manner, as well as the between-sequence interactions in an order-independent manner. The authors evaluate their models on protein variant fitness tasks by measuring the probability of variants conditioned on searched MSA of the parent sequence. The proposed method achieves large improvements on the ProteinGym benchmark. Importantly , the method addresses the problems faced by previous approaches, such as handling proteins with low-depth MSAs and accommodating indel mutations in addition to substitutions.

Strengths

1. Protein variant fitness prediction tasks are very important for protein design. The approach, which achieves large improvements on this task, will definitely have large potential impact. 2. The related work section is well-documented, and it is relatively clear where the authors' contributions lie. 3. 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. 4. While the technique used in the paper is not entirely new, the idea of generating MSA instead of individual sequences is novel enough for an application-oriented paper.

Weaknesses

1. The soundness of autoregressive decoding of MSA is questionable. Please see the questions below. 2. 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. 3. 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.

Questions

Major points: 1. 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$. Therefore, 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? 2. 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? 3. 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." Minor points: 1. In Sec. 5, it would be helpful to introduce the included baseline methods and provide the corresponding references before discussing the results. 2. 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? 3. 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. Overall, 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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

The authors should add a paragraph to discuss the potential limitations and negative societal impacts of their work.

Summary

- This paper introduces PoET, a novel autoregressive transformer architecture that learns a distribution over protein families. - 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. - 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 - The former (new sequence generation) is partially covered (one interesting yet limited example in appendix for chorismate mutase) - 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

Strengths

- The performance of the architecture on the zero-shot fitness prediction tasks from ProteinGym is quite remarkable for a model of this size - The model architecture is simple and effective - The different ablations to curate the optimal prompt engineering strategy are thorough and packed with many interesting insights

Weaknesses

- The sequence generation abilities of the model are not really explored (besides one example for 1 protein family in appendix) - 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. - 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. - 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

Questions

- “Deep mutational scans and directed evolution experiments have been used to successfully design novel proteins [1, 2], but can be costly and difficult to implement, which makes these experimental methods inapplicable for many proteins and functions of interest.” -- please elaborate on which proteins / functions DMS and directed evolutions are "inapplicable" (reference paper would be great to add here). - “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? - 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) - 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? - 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)? - Figure S1/S2: there is a reference to Prots2prot? Please also provide the performance at the DMS level for PoET alone - Given the last point in the "weaknesses" box above, which avenue(s) do you see to further improve the performance of your model?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Not particularly discussed, would suggest adding a couple sentences about it, as per the questions above.

Addendum: Sample Size for Statistics in Review Supplement Table 1

Dear Reviewers, The sample size to calculate the perplexities in Review Supplement Table 1 is 7228. Below, we have reproduced Review Supplement Table 1 with an additional column indicating the standard error of the mean perplexity. Using a paired t-test, the differences in mean perplexity between ProGen and PoET when both conditioned on 0 sequences and both conditioned on 1 sequence are statistically significant, with $p\ll1e-50$. We hope that this addendum helps to better contextualize the variance and statistical significance of the data in Review Supplement Table 1. | Model | # Sequences Conditioned On | Mean Perplexity | Std. Dev | Std. Error | |--------|----------------------------|-----------------|----------|----------| | ProGen | 0 | 15.591333 | 5.524855 | 0.06498 | | ProGen | 1 | 14.788089 | 4.022846 | 0.04732 | | PoET | 0 | 14.545277 | 2.941460 | 0.03460 | | PoET | 1 | 10.907637 | 2.817105 | 0.03314 | Review Supplement Table 1: Perplexity evaluation on our heldout Uniref50 clusters. ProGen and PoET are either conditioned on no homologs (0 sequences conditioned on) or the seed sequence from the cluster (1 sequence conditioned on) and perplexities are calculated on the remaining sequences. This evaluation only includes sequences with length <512, due to length limits for ProGen.

Addendum 2: on sequence ordering

Dear reviewers, 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, following other experiments in our manuscript, 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 (Review Supplement Table 3). | | Perplexity of Next Sequence | | | Joint Log Likelihood | | | |---------------------|-----------------------------|-----------------------------|--------------------------------------|---------------------------|-----------------------------|--------------------------------------| | Sequence Order | Mean Perplexity | Mean Difference from Random | Mean Relative Difference from Random | Mean Joint Log Likelihood | Mean Difference from Random | Mean Relative Difference from Random | | Random | 7.44075 | - | - | -8099.74973 | - | - | | Shortest-to-Longest | 7.44828 | +0.00753 (0.07686) | +0.00109 (0.01033) | -8123.18167 | -23.43149 (55.76313) | -0.00335 (0.00838) | | Longest-to-Shortest | 7.44410 | +0.00336 (0.07608) | +0.00055 (0.01030) | -8091.89462 | +7.85511 (52.92428) | +0.00105 (0.00746) | **Review Supplement Table 3**: Comparison of next sequence perplexities and joint log likelihoods, $\text{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. We find that sequence ordering has a 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. Contrasting this with the next sequence perplexities of HMMs and PSSMs conditioned on 10 sequences from Review Supplement Table 2, we can see that PoET performs significantly better regardless of sequence ordering.

Response to rebuttal

Dear authors, Thank you for the detailed responses and additional analyses provided during rebuttal. One point that was not addressed in your response is about the DMS-level performance of PoET alone (second-to-last question), ie., adding PoET alone to Fig S1/S2. This figure is quite insightful -- for instance, we see that ESM1v is performing fairly well across the board, but tanks on a handful of assays, bringing the overall average Spearman down. Since PoET is not shown there, I was thus curious about the variance of the performance across assays, relative to other baselines (eg., ESM1v being "high variance" and TranceptEVE "low variance"). Based on ablations in Tables S1-S3, performance seems relatively stable across settings, but perhaps you noticed something insightful at a more granular level?