PaperlayerView exact revision on arXiv

Source-linked AI summary

Tranception: protein fitness prediction with autoregressive transformers and inference-time retrieval

Pascal Notin, Mafalda Dias, Jonathan Frazer, Javier Marchena-Hurtado, Aidan Gomez, Debora S. Marks, Yarin Gal

arXiv:2205.13760v1cs.LG

TL;DR

Protein fitness prediction is limited by the depth and alignability of protein-family MSAs, while many mutation effects and protein regions require broader coverage. Tranception combines autoregressive modeling with inference-time retrieval and is evaluated with ProteinGym, achieving strong performance across substitutions, indels, shallow alignments, and multiple mutants.

  • Problem

    MSA-based fitness predictors depend on sufficiently deep alignments, but many proteins are difficult or impossible to align and existing methods have limited mutation and coverage scope.

  • Method

    Tranception uses an autoregressive transformer pretrained on non-aligned sequences and combines its predictions with homologous-sequence retrieval at inference.

  • Results

    Tranception achieves state-of-the-art fitness prediction performance, with stronger multiple-mutant extrapolation, robustness to shallow alignments, and support for substitutions and indels.

  • Takeaways & Limitations

    The method offers high performance with deep, shallow, or no alignments and can score mutations across proteins and variation types beyond the usual substitution setting.

  • Takeaways & Limitations

    ProteinGym remains biased toward single-mutant assays: only 11 of 87 substitution assays included multiple mutants.

Abstract

from arXiv · show

The ability to accurately model the fitness landscape of protein sequences is critical to a wide range of applications, from quantifying the effects of human variants on disease likelihood, to predicting immune-escape mutations in viruses and designing novel biotherapeutic proteins. Deep generative models of protein sequences trained on multiple sequence alignments have been the most successful approaches so far to address these tasks. The performance of these methods is however contingent on the availability of sufficiently deep and diverse alignments for reliable training. Their potential scope is thus limited by the fact many protein families are hard, if not impossible, to align. Large language models trained on massive quantities of non-aligned protein sequences from diverse families address these problems and show potential to eventually bridge the performance gap. We introduce Tranception, a novel transformer architecture leveraging autoregressive predictions and retrieval of homologous sequences at inference to achieve state-of-the-art fitness prediction performance. Given its markedly higher performance on multiple mutants, robustness to shallow alignments and ability to score indels, our approach offers significant gain of scope over existing approaches. To enable more rigorous model testing across a broader range of protein families, we develop ProteinGym -- an extensive set of multiplexed assays of variant effects, substantially increasing both the number and diversity of assays compared to existing benchmarks.

1. Introduction

Protein fitness predictors must handle proteins that are difficult to align and mutations beyond single substitutions. Tranception combines an autoregressive transformer, inference-time homology retrieval, and the ProteinGym benchmark to broaden performance and evaluation scope.

  • MSA-based predictors perform well but depend on sufficiently deep and diverse alignments, which are unavailable for some protein families.
  • Tranception is an autoregressive transformer pretrained on large quantities of non-aligned protein sequences and augmented with retrieval at inference.
  • The approach achieves state-of-the-art fitness prediction performance on substitutions and indels, especially for proteins with shallow alignments.
  • ProteinGym substantially increases the number and diversity of multiplexed variant-effect assays available for benchmarking protein models.

2. Background

Existing protein-variant predictors have leveraged aligned sequences to capture evolutionary constraints, while alternative approaches seek to model non-aligned sequences and retrieval. Benchmark development has expanded from small assay collections toward broader evaluations.

  • Alignment-based models capture information about protein functions and structures, but some proteins are unalignable or lack alignments deep enough for complex modeling.
  • MSAs provide homolog retrieval and a coordinate system for comparing amino acids across aligned protein sequences.
  • Non-aligned sequence models reduce reliance on alignments but some still use MSA routines to recover protein-specific homologous training sets.
  • Fitness benchmarks grew from approximately 20 assays to approximately 40, followed by only modest updates before newer diverse benchmarks.
  • Retrieval has been used in language modeling and bioinformatics to identify relevant information or homologous sequences from reference databases.

3. Tranception

Tranception is an autoregressive protein transformer designed to specialize attention heads and extract contiguous-sequence patterns. It scores protein fitness from sequence likelihoods while supporting long sequences and complementary inference signals.

  • Tranception promotes attention-head specialization and explicitly extracts patterns from contiguous subsequences.
  • Tranception converges faster and to lower training loss than the compared GPT2, Primer, and Tranception learned-position architectures.
  • Its attention mechanism divides heads into groups using distinct convolutional ranges, while grouped ALiBi supplies group-specific distance biases.
  • The final model is trained on UniRef100, producing approximately 250 million protein sequences after preprocessing, with 700M parameters.
  • The model predicts each token from prior context and measures mutation effects using the log-likelihood ratio between mutated and wild-type sequences.
  • Inference averages scores from canonical and reversed sequence directions and combines autoregressive predictions with retrieval-based predictions.

4. Inference-time retrieval

Tranception augments autoregressive scoring with retrieval of homologous sequences at inference. It converts aligned homologs into position-wise amino-acid distributions and combines both signals for sequence likelihood estimation.

  • MSAs retrieve and align neighboring proteins, whose position-wise amino-acid distributions reflect evolutionary constraints.
  • Retrieval inference obtains an MSA for the wild-type family and computes smoothed, reweighted amino-acid distributions at aligned positions.
  • For substitutions, one retrieval is shared across wild-type and mutated sequences, whereas indel handling uses tailored retrieval.
  • The normalization constant cancels in the log-likelihood ratio, so it need not be estimated in practice.
  • Position-level scoring allows retrieval to be ignored where needed, such as at insertion positions when scoring indels.

5. ProteinGym

ProteinGym is a broad benchmark of multiplexed variant-effect assays designed to compare predictors across substitutions, indels, functional properties, protein families, and taxa.

  • ProteinGym comprises substitution and indel benchmarks covering approximately 1.5M missense variants across 87 assays and 300k mutants across 7 assays.
  • The benchmark prioritizes assays whose measurements reflect organism fitness and, where available, have experimental replicate quality.
  • Its assays span properties including thermostability, ligand binding, aggregation, viral replication, and drug resistance across diverse protein families and taxa.
  • ProteinGym includes multiple amino-acid variants and indels, addressing gaps in prior benchmarks that focused mainly on single substitutions.
  • Model performance is evaluated primarily with Spearman’s rank correlation, supplemented by AUC and MCC for assays with bimodal measurements.

6. Results

Tranception is evaluated against alignment-based and alignment-free protein models on ProteinGym substitution and indel assays. It performs especially strongly with shallow alignments and multiple mutations.

  • Tranception with retrieval outperforms all other baselines overall on the ProteinGym substitution benchmark.
  • Tranception with retrieval achieves consistently high performance across taxa, particularly on human proteins and other eukaryotes.
  • The largest performance gaps occur for proteins with low-depth MSAs, where Tranception outperforms other baselines overall.
  • Without retrieval, Tranception outperforms ESM-1v, the only baseline that also avoids alignments during inference.
  • On the indel benchmark, Tranception outperforms Wavenet, the only listed baseline able to quantify insertion and deletion effects.

7. Discussion

Tranception extends fitness prediction to shallow or difficult alignments, full-length proteins, indels, and variants far from wild type. Its advantage is particularly pronounced for multiple mutants as mutation depth increases.

  • Gain of scope for proteins with shallow alignments: Tranception maintains consistently high performance as MSA sequences are removed, unlike EVE or MSA Transformer.
  • Gain of scope for proteins with shallow alignments: Tranception can ignore retrieval for indels, non-alignable proteins, or proteins with extremely shallow homologous-sequence sets.
  • Gain of coverage: Tranception scores all potential BRCA1 mutations without the coverage limitations faced by alignment-based models trained on full-protein alignments or separate domains.
  • Extrapolating far away in sequence space: Tranception’s performance gap over other baselines widens with mutation depth on multiple mutants.
  • Extrapolating far away in sequence space: ProteinGym remains biased toward single-mutant assays, with only 11 of 87 substitution assays including multiple mutants.

8. Conclusion

Tranception combines autoregressive modeling with inference-time retrieval to broaden protein fitness prediction beyond proteins with deep alignments. The approach reports strong performance across mutation types and supports applications spanning difficult-to-align proteins, viral mutations, human variants, and protein design.

  • 8. Conclusion: Tranception with retrieval delivers state-of-the-art fitness prediction and stronger extrapolation to multiple mutants.The conclusion attributes this result to combining the Tranception architecture with retrieval at inference.
  • 8. Conclusion: Retrieval inference remains robust across alignment characteristics, allowing Tranception to use deep, small, or no alignments when needed.The authors contrast this scope with EVE, which requires sufficiently deep alignments to capture complex residue relationships.
  • 8. Conclusion: Tranception handles insertions and deletions out-of-the-box and outperforms prior baselines in that regime.
  • 8. Conclusion: The approach is robust across taxa and protein families, including mutation-effect prediction for viruses and disease-causing variants in humans.The paper connects viral prediction to outbreak forecasting and human variant prediction to diagnosis and preventative care.
  • 8. Conclusion: The autoregressive model is naturally suited to sequence generation and therefore has potential for protein design.
  • 8. Conclusion: MSA depth is the number of sequences in an alignment, and gain of scope denotes increasing the number of distinct protein families that can be modeled.Historically difficult cases include proteins with shallow MSAs or difficult-to-align proteins such as disordered proteins.

B.1. Ablation studies

The ablation studies examine Tranception architecture, training-data granularity, and benchmark design choices. They compare variants under controlled ProteinGym validation settings and report characteristics of the training data and models.

  • B.1. Ablation studies: Tranception is an autoregressive transformer designed to promote attention-head specialization and extract contiguous protein subsequence patterns.The largest model, Tranception L, has 700M parameters and is trained on UniRef100.
  • B.1. Ablation studies: Ten ProteinGym substitution assays form a representative validation subset spanning taxa, mutation depths, and one multiple-mutant assay.The subset includes 3 viral, 4 human or other eukaryotic, and 3 prokaryotic proteins.
  • B.1. Ablation studies: The ablation table measures Spearman’s rank correlation between model scores and experimental measurements, excluding retrieval inference.Model selection uses the validation set described in Appendix B.1.
  • B.1. Ablation studies: About 98% of UniRef100 sequences after preprocessing are shorter than 1,024 amino acids.
  • B.1. Ablation studies: Training uses approximately 249 million sequences, with 99% assigned to training and 1% held out for validation.The preprocessing applies mild filtering while preserving as much sequence diversity as possible.

B.3. Model training

The training and scoring pipeline combines autoregressive sequence scoring with optional retrieval, bidirectional traversal, and alignment construction. Training choices include sequence reversal, context-length handling, UniRef clustering comparisons, and a tuned retrieval weight.

  • B.3. Model training: All model variants train for 150k steps with batches of 1,024 sequences, randomly reversed sequences, and truncation beyond the maximum context size.
  • B.3. Model training: The Tranception S loss curves compare UniRef50, UniRef90, and UniRef100, while the right panel compares Small, Medium, and Large architectures trained on UniRef100.
  • B.3. Model training: Bidirectional scoring averages log-likelihoods from canonical and reverse sequence traversals for Tranception L, with and without retrieval.The comparison is reported on both the validation DMS set and the full DMS set.
  • B.3. Model training: Retrieval inference augments autoregressive predictions only at positions covered by the retrieved MSA, while uncovered positions rely fully on autoregressive predictions.Weighted Laplace smoothing uses a small smoothing parameter of 10^-5.
  • B.3. Model training: A linearly spaced grid search on the validation DMS set selects an optimal retrieval inference weight α of 0.6.
  • B.3. Model training: For ProteinGym assays, MSAs are built using five Jackhmmer search iterations over UniRef100 and selection among nine bit-score thresholds.The selected alignment has the highest number of significant evolutionary couplings.

D.1. ESM-1v and MSA Transformer

ESM-1v and MSA Transformer are extended and evaluated using masked-marginals scoring, with adaptations for multiple mutants, long sequences, and weighted MSA sampling. Their scoring heuristic has explicit limits for epistasis and indels.

  • D.1. ESM-1v and MSA Transformer: The ESM codebase is extended to score multiple mutants, long sequences, and weighted samples from an input MSA.Multiple-mutant scores sum the effects of constituent single mutations, while long sequences use Tranception’s scoring-window routine.
  • D.1. ESM-1v and MSA Transformer: Masked-marginals scoring computes multiple-mutant effects by summing individual mutation effects, thereby ignoring potentially important epistatic effects.
  • D.1. ESM-1v and MSA Transformer: Masked scoring assumes each mutated position exists in the original wild-type coordinate frame, so ESM-1v and MSA Transformer cannot score indels with this heuristic.
  • D.1. ESM-1v and MSA Transformer: The comparison uses the official EVE and DeepSequence implementations with protein-specific MSAs, while alignment-based models share the MSAs constructed for evaluation.
  • D.1. ESM-1v and MSA Transformer: EVE and DeepSequence remove sequences with more than 50% gaps and positions with more than 30% gaps, leaving some mutations without scores.When coverage is difficult, an entire protein domain may be dropped.
  • D.1. ESM-1v and MSA Transformer: Except for Wavenet, alignment-based models cannot score indels because they use the fixed coordinate system of the original MSA.

E.1. Performance reporting methodology

ProteinGym reports standardized model performance across substitution and indel assays using multiple metrics, while additional analyses assess alignment-depth robustness and ensembling. The benchmark aggregates results by taxon and provides detailed assay-level evaluations.

  • Metrics and aggregation: Signed Spearman’s ρ is used because some assays produce very low or negative correlations, with phenotype signs adjusted for consistent directionality.Higher DMS scores indicate better agreement with the standardized direction of experimental effects.
  • Benchmark construction: ProteinGym preprocessing removes silent mutations and missing measurements and averages measurements for duplicate mutants.These steps standardize assay inputs before benchmarking.
  • Metrics and aggregation: Performance is reported using Spearman’s rank correlation, AUC, and Matthew’s correlation coefficient across ProteinGym substitution assays.Spearman’s correlation compares model scores with DMS measurements; results are aggregated by taxon.
  • Robustness analyses: Tranception performance is less sensitive to MSA depth than EVE or MSA Transformer in the alignment-filtering analysis.The analysis progressively excludes aligned sequences based on similarity to the sequence used to construct each alignment.
  • Benchmark construction: The final benchmark contains 94 deep mutational scanning experiments: 87 substitution assays and 7 indel assays.The assays span 77 publications, and raw and processed data are publicly available.
  • Ensembling analyses: Ensembling five EVE models marginally outperforms a single Tranception model with retrieval, while combining Tranception with retrieval and EVE reaches an overall performance of 0.475.The paired-architecture analysis suggests retrieval and alignment-based modeling capture analogous information, while autoregressive inference contributes complementary information.
Loading 2205.13760v1…