Research map / Architectures and efficiency / RNNs and normalization
Training Compute-Optimal Protein Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Trains over 300 protein language models (3.5M to 10.7B params, 5 to 200B tokens) on a 939M-sequence metagenomic-augmented dataset to derive compute-optimal scaling laws for CLM and MLM, including CLM-to-MLM transfer, and validates by comparing re-scaled 7.2B CLM and 10.7B MLM against PROGEN2-xlarge and ESM-2 3B.
Problem
Compute-optimal training of protein language models is underexplored; it is unclear how model size, training tokens, and objective (MLM vs CLM) should scale under a fixed compute budget, and whether repeated UniRef data limits scaling.
Approach
Builds a 194B-token (939M sequence) dataset UniMeta200B by adding ColabFoldDB metagenomic data to UniRef to avoid overfitting, fits power-law IsoFLOP scaling laws separately for CLM and MLM, and introduces Effectively Transferred Tokens to quantify how CLM pre-training transfers to MLM.
Key results
- Under a 10x compute increase, CLM model size should grow about 4x and tokens 3x, while MLM model size should grow 6x and tokens about 1.7x.
- The 7.2B CLM model achieved average pLDDT 78.69 vs 74.33 for PROGEN2-xlarge, and higher diversity (7,097 vs 4,818 clusters).
- The 10.7B MLM model outperformed ESM-2 (3B) on 7 of 8 downstream tasks (e.g., fold classification 0.72 vs 0.69, fluorescence 0.69 vs 0.65).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use the 7.2B CLM scaling allocation instead of PROGEN2-xlarge | Protein sequence generation · Generated sequences evaluated with ESMFold, Foldseek, and clustering | generation quality and diversity: 74.33 -> 78.69 average pLDDT; 0.522 -> 0.655 average TM-score; 0.165 -> 0.194 SeqID; 4,818 -> 7,097 clusters | Helped |
| Use the 10.7B MLM scaling allocation instead of ESM-2 (3B) | Protein understanding tasks · Contact prediction, fold classification, fluorescence, and five appendix test datasets | downstream task performance: contact 0.91 vs 0.91; fold 0.69 -> 0.72; fluorescence 0.65 -> 0.69; fitness 0.94 -> 0.96; localization 0.81 -> 0.79; met | Mixed |
| Pre-train on CLM, then train on MLM | 470M model · Contact prediction and fold classification | downstream task performance: contact 0.78 -> 0.80; fold 0.65 -> 0.66 | Helped |
| Pre-train on CLM, then train on MLM | 470M model · Fluorescence | regression score: 0.67 -> 0.67 | No clear effect |
| Expand training data from repeated UR50/S to diversified UniMeta200B | Protein language model pre-training · OOD test PPL and IID validation learning curve | OOD test PPL | Helped |
| Repeat UR50/S across epochs instead of using diverse, approximately single-epoch UniMeta200B data (ablation) | 150M and 3B CLM/MLM · OOD generalization and training loss | generalization | Hurt |
| Train with 5% or greater than 25% masking rather than a 10%-20% mask ratio (ablation) | 85M and 154M MLM · Validation loss | validation loss | Hurt |
| Use a 20% rather than 15% masking ratio (ablation) | MLM · Contact prediction | contact prediction performance | Helped |
| Use a 15% rather than 20% masking ratio (ablation) | MLM · Fold prediction | fold prediction performance | Helped |
| Replace Uniref90 training data with ColabFoldDB (ablation) | 85M protein language model · Contact prediction and fold classification | downstream task performance | No clear effect |
| Transfer from CLM pre-training before MLM instead of training MLM from scratch (ablation) | 470M model · MLM fine-tuning on downstream tasks | convergence speed | Helped |
| Mix CLM and MLM objectives during training instead of training from scratch with a single objective (ablation) | Protein language models · Validation loss | validation loss | Hurt |
Surprising findings
- Mixing CLM and MLM objectives in a single training run was detrimental compared with training each objective separately.
- Transferring from MLM to CLM yielded diminishing returns, while CLM-to-MLM transfer stayed effective as scale grew.
Limitations
- Scaling laws are fitted within a single-epoch regime; the authors note the data-repeat setting is not covered by the main law.
- Results depend on the chosen dataset composition and may change with different pre-training datasets or architectures.
- Only dense Transformers were studied in the main analysis; MoE scaling was limited to eight experts.
Experimental setup
Datasets: UniRef50/S UniRef90/50 ColabFoldDB (representatives and members) UniProt OOD test sequences PDB (Foldseek search)
Benchmarks: protein contact prediction fold classification fluorescence fitness prediction localization metal ion binding solubility stability OOD perplexity
Metrics: perplexity validation loss P@L/5 accuracy accuracy Spearman correlation pLDDT TM-score cluster count diversity
Baselines: ESM-2 (3B) PROGEN2-xlarge
Builds on: Chinchilla scaling laws ESM-2 ProGen2 LoRA
Closest papers by meaning
- Feature Reuse and Scaling: Understanding Transfer Learning with Protein Language Models (ICML 2024 · similarity 0.91)
- Concept Bottleneck Language Models For Protein Design (ICLR 2025 · similarity 0.90)
- LEMON-ZEST: Evolution-Informed Tokenization for Efficient Protein Language Modeling (NeurIPS 2026 · similarity 0.90)
- Diffusion on Language Model Encodings for Protein Sequence Generation (ICML 2025 · similarity 0.90)
- Diffusion Language Models Are Versatile Protein Learners (ICML 2024 · similarity 0.89)
- Steering Protein Language Models (ICML 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.