[Paper Review] Training Compute-Optimal Protein Language Models

Xingyi Cheng, Bo Chen+4 more|arXiv (Cornell University)|Nov 4, 2024

Machine Learning in Bioinformatics4 citations

TL;DR

This paper proposes compute-optimal training strategies for protein language models by establishing scaling laws tailored to protein sequence data. Using a 194-billion-token dataset and training over 300 models, it demonstrates that sublinear scaling of data and model size under fixed compute budgets yields superior performance, with CLM pre-training enabling effective transfer to MLM tasks through estimated Effectively Transferred Tokens, outperforming baseline models like ESM-2 and PROGEN2 under equivalent compute constraints.

ABSTRACT

We explore optimally training protein language models, an area of significant interest in biological research where guidance on best practices is limited. Most models are trained with extensive compute resources until performance gains plateau, focusing primarily on increasing model sizes rather than optimizing the efficient compute frontier that balances performance and compute budgets. Our investigation is grounded in a massive dataset consisting of 939 million protein sequences. We trained over 300 models ranging from 3.5 million to 10.7 billion parameters on 5 to 200 billion unique tokens, to investigate the relations between model sizes, training token numbers, and objectives. First, we observed the effect of diminishing returns for the Causal Language Model (CLM) and that of overfitting for the Masked Language Model~(MLM) when repeating the commonly used Uniref database. To address this, we included metagenomic protein sequences in the training set to increase the diversity and avoid the plateau or overfitting effects. Second, we obtained the scaling laws of CLM and MLM on Transformer, tailored to the specific characteristics of protein sequence data. Third, we observe a transfer scaling phenomenon from CLM to MLM, further demonstrating the effectiveness of transfer through scaling behaviors based on estimated Effectively Transferred Tokens. Finally, to validate our scaling laws, we compare the large-scale versions of ESM-2 and PROGEN2 on downstream tasks, encompassing evaluations of protein generation as well as structure- and function-related tasks, all within less or equivalent pre-training compute budgets.

Motivation & Objective

• To identify optimal compute allocation strategies for training protein language models under fixed budgets.
• To address overfitting and performance platetuing in CLM and MLM by expanding training data with metagenomic sequences.
• To derive scaling laws specific to protein sequence data for both Causal Language Model (CLM) and Masked Language Model (MLM) objectives.
• To validate the effectiveness of CLM pre-training for downstream MLM fine-tuning via transfer learning.
• To re-optimize existing model setups (e.g., ESM-2, PROGEN2) using derived scaling laws to improve performance within the same compute budget.

Proposed method

• Constructed a 194-billion-token, 939-million-sequence dataset (UniMeta200B) by combining Uniref with metagenomic sequences to enhance diversity and reduce overfitting.
• Trained 300+ models ranging from 3.5M to 10.7B parameters on 5 to 200B unique tokens to study scaling behaviors under controlled compute budgets.
• Fitted power-law scaling functions to model loss and performance across model size and training tokens, estimating scaling exponents for CLM (0.57) and MLM (0.77).
• Introduced the concept of Effectively Transferred Tokens (Dₜ) to quantify transfer efficiency from CLM to MLM, enabling optimal compute allocation between objectives.
• Reallocated model size and training tokens for ESM-2 (3B) and PROGEN2-xlarge using derived scaling laws, training new models at 7.2B and 10.7B parameters.
• Used FlashAttention, bfloat16, and optimized learning rate schedules (cosine decay, 2.5% warmup) across A100 GPUs with data parallelism to accelerate training.

Experimental results

Research questions

• RQ1 How do scaling laws for CLM and MLM differ in protein language modeling, and what are their distinct power-law exponents?
• RQ2 Can overfitting and performance plateau in CLM and MLM be mitigated by increasing training data diversity beyond Uniref?
• RQ3 To what extent can CLM pre-training be effectively transferred to improve MLM performance under a fixed compute budget?
• RQ4 What is the optimal allocation of compute between CLM and MLM pre-training when both are required?
• RQ5 Can re-allocating model size and training tokens based on scaling laws improve downstream performance compared to established models like ESM-2 and PROGEN2?

Key findings

• The scaling law for MLM shows a higher exponent (0.77) than CLM (0.57), indicating that MLM benefits more from increased compute in terms of model size growth.
• A 10× increase in compute leads to a 6× increase in model size and 70% increase in data for MLM, versus a 4× increase in model size and 3× in data for CLM.
• CLM pre-training enables effective transfer to MLM, with the transfer efficiency quantified by Effectively Transferred Tokens (Dₜ), allowing optimal allocation of compute between objectives.
• Reallocating compute under ESM-2 (3B) and PROGEN2-xlarge budgets led to models of 7.2B and 10.7B parameters, respectively, which outperformed the original models on downstream tasks.
• The 10.7B-parameter model trained under the new strategy achieved state-of-the-art performance on protein generation and structure/function prediction tasks within the same compute budget.
• The study confirms that data diversity (via metagenomic sequences) is critical to avoid overfitting and performance saturation in both CLM and MLM objectives.