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

TL;DR

PoET introduces a Protein Evolutionary Transformer that treats a protein family as a sequence-of-sequences, enabling cross-family transfer learning and conditioning on any family without relying on MSAs. Its tiered Transformer architecture captures within-sequence and order-invariant between-sequence dependencies, enabling extrapolation to very long contexts. Empirically, PoET achieves state-of-the-art or competitive performance on deep mutational scanning fitness prediction across proteins with varying MSA depths and can generate diverse, structure-preserving sequences conditioned on family data. The work demonstrates retrieval-augmented design, faster inference than competing baselines, and the ability to generate novel functional sequences, suggesting PoET as a practical tool for ML-enabled 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 $\textbf{P}$r$\textbf{o}$tein $\textbf{E}$volutionary $\textbf{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.

Figures (19)

• Figure 1: PoET Architecture
• Figure 2: Illustration of evaluating PoET for variant fitness prediction on a DMS dataset
• Figure 3: Performance of PoET on the ProteinGym validation set when trained with (left) various training distributions and (right) model sizes.
• Figure 4: Comparison of the perplexity of a regular Transformer and PoET when generating a protein sequence conditioned on a fixed number of tokens from other sequences in the same protein family. Protein families consist of UniRef50 sequences, and both models perform much better than a profile HMM baseline (Appendix Figure \ref{['fig:appendix_ppl_with_hmm']}).
• Figure 5: Sequence novelty and predicted structural conservation of (A) functional chorismate mutases generated by PoET and (B) phage lysozymes generated by PoET and ProGen. PoET generates diverse sequences (50-100% seq id to a natural sequence) while preserving 3D structure within the protein family (TM-score $>0.8$ to a structure of a natural sequence and pLDDT $>90$).