[Paper Review] ProGen2: Exploring the Boundaries of Protein Language Models

Erik Nijkamp, Jeffrey A. Ruffolo+3 more|arXiv (Cornell University)|Jun 27, 2022

Machine Learning in Bioinformatics50 citations

TL;DR

ProGen2 scales protein language models up to 6.4B parameters trained on over a billion sequences, achieving state-of-the-art perplexity, generating viable proteins, and enabling zero-shot fitness prediction.

ABSTRACT

Attention-based models trained on protein sequences have demonstrated incredible success at classification and generation tasks relevant for artificial intelligence-driven protein design. However, we lack a sufficient understanding of how very large-scale models and data play a role in effective protein model development. We introduce a suite of protein language models, named ProGen2, that are scaled up to 6.4B parameters and trained on different sequence datasets drawn from over a billion proteins from genomic, metagenomic, and immune repertoire databases. ProGen2 models show state-of-the-art performance in capturing the distribution of observed evolutionary sequences, generating novel viable sequences, and predicting protein fitness without additional finetuning. As large model sizes and raw numbers of protein sequences continue to become more widely accessible, our results suggest that a growing emphasis needs to be placed on the data distribution provided to a protein sequence model. We release the ProGen2 models and code at https://github.com/salesforce/progen.

Motivation & Objective

• Investigate how very large-scale protein language models (PLMs) perform across generation, structure-aware generation, and fitness prediction.
• Assess the impact of training data scale and composition (genomic, metagenomic, immune repertoire) on model quality.
• Evaluate generation quality by structural predictions and functional proxies without task-specific fine-tuning.
• Examine zero-shot fitness prediction capabilities on diverse experimental landscapes and antibody datasets.

Proposed method

• Autoregressive transformer architectures trained with next-token prediction on protein sequences.
• Model sizes: 151M, 764M, 2.7B, and 6.4B parameters with left-to-right causal masking.
• Rotary positional encodings for sequence positions.
• Parallelized training using TPU-v3 with JAX and pjit for SPMD data/model parallelism.
• Training datasets: Uniref90, BFD30, and a large antibody-focused set (OAS) for specialized models.
• Evaluation via perplexity on held-out data, AlphaFold2 structure prediction for generated sequences, Foldseek TMscore comparisons, and zero-shot fitness benchmarks on multiple datasets.

Experimental results

Research questions

• RQ1 How does increasing model size affect the ability to capture the distribution of observed protein sequences (perplexity at different identity thresholds)?
• RQ2 Can large PLMs generate novel, structurally plausible proteins and antibodies without task-specific fine-tuning?
• RQ3 Do pretrained models provide accurate zero-shot fitness predictions across diverse protein landscapes and antibody datasets?
• RQ4 How does finetuning on structure-focused datasets influence generation quality and structural similarity to natural proteins?

Key findings

• Larger ProGen2 models yield substantially lower perplexities, indicating better distribution capture.
• Generated proteins often fold into plausible structures and resemble natural folds, though with notable sequence deviations and occasional novel folds.
• Finetuning on structural-architecture datasets increases similarity to natural sequences and structures, with sampling parameters influencing diversity and similarity.
• Antibody-generation results show prompted sequences align more with training distribution and that development properties (aggregation, solubility) vary with sampling choices.
• Zero-shot fitness prediction ability is demonstrated across narrow, wide, and antibody-specific landscapes, supported by structural and functional assessments.