[Paper Review] PoET: A generative model of protein families as sequences-of-sequences

Tristan Bepler, Timothy F. Truong|arXiv (Cornell University)|Jun 9, 2023

Genomics and Phylogenetic StudiesBiochemistry, Genetics and Molecular Biology26 citations

TL;DR

PoET is an autoregressive Transformer that models whole protein families as sequences-of-sequences, enabling retrieval-augmented conditioning, indel-aware generation, and improved variant fitness prediction across many protein families without relying on MSAs.

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 $ extbf{P}$r$ extbf{o}$tein $ extbf{E}$volutionary $ extbf{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.

Motivation & Objective

• Motivate improving protein design by modeling evolutionary constraints across many families without relying on MSAs.
• Develop a scalable, order-invariant Transformer architecture that generates sets of related proteins as sequences-of-sequences.
• Enable retrieval-augmented conditioning and efficient scoring/generation of sequence variants, including indels.
• Demonstrate improved variant fitness prediction across deep mutational scanning datasets and show capability to generate novel, structurally plausible sequences.

Proposed method

• Introduce PoET, an autoregressive model over a sequence-of-sequences distribution P(X=x) = product over sequences and tokens.
• Propose TieredTransformerDecoderLayer with two attention modules: PerSequenceSelfAttn (within-sequence) and SequenceOfSequencesSelfAttn (between-sequences) to achieve order invariance across sequences and order dependence within sequences.
• Use Rotary Positional Encodings for within-sequence attention and a novel inter-sequence relative positional encoding ensuring order-invariance across sequences while preserving within-sequence structure.
• Train on 29 million sets of homologous sequences drawn from UniRef50, with inverse-count weighting to balance set sizes and randomized sequence order to promote invariance.
• Condition on retrieved homologous sequences S to compute conditional fitness scores and enable retrieval-augmented generation and scoring (e.g., PoET’s fitness prediction uses log P(v|S)).
• Evaluate on ProteinGym deep mutational scanning datasets, comparing against alignment-based, unconditional, conditional, and hybrid protein language models; show ensembles improve performance.

Experimental results

Research questions

• RQ1 Can PoET generalize evolutionary constraints across millions of protein sequence clusters to improve variant fitness prediction for small or large families?
• RQ2 Does a sequence-of-sequences Transformer with order-invariant inter-sequence attention outperform MS-based or unconditional models for variant effect prediction and indel handling?
• RQ3 Can PoET be used as a retrieval-augmented language model to generate and score protein sequences conditioned on a target family without requiring MSAs?
• RQ4 How does PoET perform in generating novel, structurally plausible sequences that preserve family-specific characteristics?

Key findings

• PoET achieves state-of-the-art or competitive performance for variant fitness prediction across ProteinGym datasets, improving substitution predictions across all MSA depths.
• An ensemble of PoET with TranceptEVE L significantly improves substitution predictions over either method alone.
• PoET outperforms baselines on indel variant predictions and can score and generate sequences with indels not present in the training MSA.
• Longer context lengths (up to thousands of tokens) enable PoET to observe more homologs and improve performance; PoET generalizes well beyond training context lengths.
• PoET can generate diverse, novel sequences that preserve structural plausibility (high pLDDT, TM-score clustering around native-like folds) and maintain family-level structural integrity.