PaperlayerView exact revision on arXiv

Source-linked AI summary

RITA: a Study on Scaling Up Generative Protein Sequence Models

Daniel Hesslow, Niccoló Zanichelli, Pascal Notin, Iacopo Poli, Debora Marks

arXiv:2205.05789v2q-bio.QMcs.LG

TL;DR

Protein-domain generative models lacked a systematic study of how capabilities change with model size. RITA introduces autoregressive protein sequence models and examines scaling; its protein modeling-loss exponent is 0.74 versus around 0.5 in NLP.

  • Problem

    The protein domain lacked a systematic study of how autoregressive model capabilities evolve with size, while the largest ProGen model was reported to be underfitting.

    Source: §1 · p. 1

  • Method

    RITA introduces generative protein sequence models up to 1.2B parameters and studies the relationship between model size and downstream task performance.

    Source: §1 · p. 1

  • Results

    Protein modeling loss scaling has an exponent of 0.74, versus around 0.5 in NLP, in the analyzed regime up to 1.2 billion parameters.

    Source: §4.2 · pp. 3–4

  • Takeaways & Limitations

    The authors release RITA models openly for the scientific community.

    Source: §1 · p. 1

    In the analyzed regime, all but the smallest model still appear to be undertrained.

    Source: §4.2 · p. 4

Abstract

from arXiv · show

In this work we introduce RITA: a suite of autoregressive generative models for protein sequences, with up to 1.2 billion parameters, trained on over 280 million protein sequences belonging to the UniRef-100 database. Such generative models hold the promise of greatly accelerating protein design. We conduct the first systematic study of how capabilities evolve with model size for autoregressive transformers in the protein domain: we evaluate RITA models in next amino acid prediction, zero-shot fitness, and enzyme function prediction, showing benefits from increased scale. We release the RITA models openly, to the benefit of the research community.

1. Introduction

Protein design has made progress, but reliable design remains an open challenge; the protein domain also lacked a systematic study of how autoregressive model capabilities change with size. RITA introduces a large generative model family, studies size and downstream performance, and releases the models openly.

1. Introduction

  • The protein domain lacked a systematic study of how autoregressive model capabilities evolve with model size.Prior work had studied size effects for bidirectional transformers, while ProGen's largest model was noted to be underfitting.Source: §1 · p. 1
  • RITA is a family of generative protein sequence models for protein design, with up to 1.2B parameters.Source: §1 · p. 1
  • The study examines the relationship between model size and downstream task performance as a first step toward protein-sequence scaling laws.Source: §1 · p. 1
  • The authors release RITA models on Hugging Face for the scientific community.Source: §1 · p. 1

2. Related Work

The related work situates protein modeling within NLP's scaling-law context and reviews protein sequence models' existing predictive and design uses. It identifies limited scaling of generative protein models and limited study of how scale affects downstream tasks.

2.1. Language models for natural language processing

  • NLP scaling laws guide estimates of language-model loss and choices of model and dataset size for a compute budget.Source: §2.1 · pp. 1–2
  • Generative protein models had seen limited scaling to NLP-sized models and limited study of scaling effects on relevant downstream tasks.Source: §2.1 · p. 2

2.2. Protein Sequence Models

  • Protein sequence-model representations have supported prediction of structure, stability, function, variant effects, and tertiary structure.The cited examples span UniRep, bidirectional transformer models, DeepSequence, ESM-1v, and RGN2.Source: §2.2 · p. 2
  • ProGen can generate proteins with desired characteristics by conditioning on sequence metadata.Source: §2.2 · p. 2
  • Fixed-backbone design pipelines have combined sequence models with structural conditioning, structure-predictor-guided decoding, or iterative finetuning on AlphaFold 2-refined sequences.Source: §2.2 · p. 2

3. Methods

RITA uses unconditional decoder-only transformers and selects RoPE after comparing positional embeddings. The methods describe corpus selection and training, alongside a reported relationship between model scale and performance.

3.1. Model Architecture

  • RITA's decoder-only transformers use no conditioning information, and RoPE was chosen over AliBi after yielding lower language-modeling loss.Four model sizes were trained to study the relationship between size and downstream capability.Source: §3.1 · p. 2
  • Performance scales smoothly with model size, and RITA-XL gives the best results.Source: §3.1 · p. 3

3.2. Data

  • The authors avoided clustering and compared UniRef-100, MGnify, and Metaclust as pretraining corpora; UniRef-100 followed by Metaclust performed best, while MGnify performed worst.The authors note that combining datasets may be beneficial.Source: §3.2 · p. 3

3.3. Training

  • Training used 150 billion amino acids across all models and exceeded 25,000 Nvidia-V100 GPU hours.All experiments used Adam, batch size 512, and context size 1024.Source: §3.3 · p. 3
  • Figure 1 plots protein modeling loss against compute measured in PetaFLOPS-days (PF-days).Source: Figure 1 · p. 3

4. Evaluation

RITA is evaluated across sequence modeling, mutation-effect prediction, enzyme-function prediction, and prompt tuning. Performance improves with model size across several tasks, while scaling-law analysis and prompt tuning probe additional capabilities.

4.1. Perplexity evaluation

  • Per-amino-acid perplexity improves with increasing model size compared with ProtGPT2 across UniRef-100, MGnify, and Metaclust evaluations.The evaluation also includes twenty protein families withheld during training to test generalization to unseen families.Source: §4.1 · p. 3

4.2. Scaling Laws

  • The scaling-law exponent is 0.74, steeper than the approximately 0.5 observed in NLP, although loss deviates from a close power law and most models appear undertrained.The analysis covers models up to 1.2 billion parameters, trained on 150 billion amino acids.Source: §4.2 · pp. 3–4

4.5. Prompt Tuning

  • Prompt tuning significantly reduces perplexity for held-out protein family PF03272, indicating that the model learned to generate proteins from that family.The experiment specializes a pre-trained model using a learned prompt.Source: Table 4 · p. 4, §4.5 · p. 4

4.3. Mutation Effects Prediction

  • RITA’s mutation-prediction performance increases with model size and exceeds ESM-1v for the Large and XLarge variants on the substitution benchmark.RITA models can also score indels out of the box, performing on par with specialized models introduced in Shin et al. (2021).Source: §4.3 · p. 4

4.4. Enzyme Function Prediction

  • Enzyme-function classification performance increases with scale using final-token representations and a one-epoch linear classifier over 4,793 classes.The labels are the finest-level enzyme commission tags from SwissProt.Source: §4.4 · p. 4

5. Conclusion and Future Work

RITA is presented as a family of generative protein sequence models intended to accelerate future protein-design work. The authors identify in-vitro study, further scaling, and target-structure embeddings as directions for future research.

5. Conclusion and Future Work

  • The authors release RITA as a building block for future protein-design work and point to in-vitro evaluation, further scaling, and target-structure embeddings as next steps.The proposed structure embeddings are for fixed-backbone protein design.Source: §5 · p. 4

A. Full results from our fitness evaluation.

The appendix captions describe ProteinGym substitution and indel fitness evaluations using Spearman’s rank correlation between predicted and measured fitness. The substitution comparison aggregates by UniProt ID and restricts evaluation against evolutionary-data models to mutations with adequate coverage.

A. Full results from our fitness evaluation

  • The substitution benchmark reports Spearman’s rank correlation between predicted and experimentally measured fitness; baseline results use a single seed.Tranception NR and Tranception R are variants without and with retrieval, respectively.Source: Table 5 · p. 8
  • Substitution results are aggregated at the UniProt ID level, and comparisons with multiple-sequence-alignment models use only mutations with sufficient coverage.The aggregation is intended to avoid bias toward proteins with several assays.Source: Table 5 · p. 8
  • The indel benchmark reports Spearman’s rank correlation for experimentally measured versus predicted fitness and includes Wavenet and Tranception variants.Tranception NR and Tranception R are variants without and with retrieval, respectively.Source: Table 6 · p. 9

B. Positional embedding ablation

The positional-embedding ablation reports better results for RoPE than ALiBi in a short, small-model comparison. The authors caution that the results need larger-scale testing and should not be directly compared with the main results.

B. Positional embedding ablation

  • RoPE outperforms ALiBi in a small-model ablation trained for 3 billion amino acids, but short runs and a different training dataset limit interpretation and direct comparison with Table 1.The authors state that larger-scale ablation is needed for reliable results.Source: Table 7 · p. 9

C. Dataset Selection

The dataset ablation used short training runs and evaluated models on a combination of datasets.

C. Dataset Selection

  • Models were trained for approximately 3 GT to evaluate alternative pre-training datasets.Source: Table 8 · p. 9
  • All models were evaluated on a combination of the datasets.Source: Table 8 · p. 9

D. Comparisons with ProtXLNet

The paper removed its ProtXLNet comparisons from the main text after discovering faulty perplexity comparisons, and supplied a corrected evaluation in the appendix.

D. Comparisons with ProtXLNet

  • After finding faulty ProtXLNet perplexity comparisons, the authors removed them from the main text and supplied corrected appendix results.The corrected evaluation also excludes proteins shorter than 100 amino acids; the authors note that XLNet differs from decoder-only autoregressive transformers.Source: D. Comparisons with ProtXLNet · p. 9

E. Comparing perplexity across different vocabularies

Per-token perplexity is not directly comparable across vocabularies, so the paper uses perplexity per amino acid while cautioning that this metric has limitations. Its results also show size-correlated performance across four datasets, with RITA-XL best.

E. Comparing perplexity across different vocabularies

  • Per-token perplexity does not transfer across vocabularies because predicting among tens of thousands of tokens is harder than among dozens.Source: E. Comparing perplexity across… · p. 9
  • Normalizing by untokenized sequence length yields perplexity per byte, or per amino acid for proteins, which remains constant when tokens are merged using their joint probability.Source: E. Comparing perplexity across… · p. 9
  • Perplexity-per-byte performance correlated with model size on all four datasets, and RITA-XL performed best.ProtGPT2 and ProtXLNet had also been trained on the Pfam families held out from RITA's training set.Source: Table 9 · p. 10
  • With redundant vocabularies, perplexity per byte can increase unfairly because the model must guess which tokenization represents the target sequence.Table 10 illustrates changes in token and byte perplexity across vocabularies under uniform token probabilities.Source: E. Comparing perplexity across… · p. 10, Table 10 · p. 10
  • The authors hypothesize that longer tokenized words make deterministic tokenization harder to learn and report preliminary gains from splitting proteins into k-mers before tokenization.Source: E. Comparing perplexity across… · p. 10
  • The authors caution against interpreting exact perplexity-per-byte values and note that they know of no better cross-vocabulary metric.Source: E. Comparing perplexity across… · p. 10

Related papers

Loading 2205.05789v2…