[Paper Review] Are Protein Language Models Compute Optimal?

Yaiza Serrano, Álvaro Ciudad+1 more|arXiv (Cornell University)|Jun 11, 2024

Machine Learning in Bioinformatics4 citations

TL;DR

This study investigates compute-optimal protein language models (pLMs) by adapting NLP scaling laws to pLMs, revealing that performance plateaus at a fixed loss regardless of model size or data scale. Using single-pass training, a 35M-parameter model achieves perplexity of 11.43—comparable to larger models like ESM-2 (15B) and xTrimoPGLM (100B)—demonstrating that smaller models can be more compute-efficient when trained optimally.

ABSTRACT

While protein language models (pLMs) have transformed biological research, the scaling laws governing their improvement remain underexplored. By adapting methodologies from NLP scaling laws, we investigated the optimal ratio between model parameters and training tokens within a fixed compute budget. Our study reveals that pLM sizes scale sublinearly with compute budget, showing diminishing returns in performance as model size increases, and we identify a performance plateau in training loss comparable to the one found in relevant works in the field. Our findings suggest that widely-used pLMs might not be compute-optimal, indicating that larger models could achieve convergence more efficiently. Training a 35M model on a reduced token set, we attained perplexity results comparable to larger models like ESM-2 (15B) and xTrimoPGLM (100B) with a single dataset pass. This work paves the way towards more compute-efficient pLMs, democratizing their training and practical application in computational biology.

Motivation & Objective

• To investigate whether widely used protein language models (pLMs) are compute-optimal under fixed FLOP budgets.
• To determine the optimal ratio between model parameters and training tokens for minimizing training loss within a fixed compute budget.
• To assess whether current pLMs, such as ESM-2 and xTrimoPGLM, follow compute-optimal scaling laws or exhibit diminishing returns.
• To evaluate the performance of a small model (35M parameters) trained in a single pass against larger models, using perplexity as a key metric.
• To derive scaling laws for optimal model size and dataset size that minimize loss at different compute levels.

Proposed method

• Adapted NLP scaling laws from Hoffmann et al. (2022) and Kaplan et al. (2020) to model the relationship between model size (N), dataset size (D), and training loss under fixed FLOP budgets.
• Trained a range of pLMs (5M to 650M parameters) on randomly sampled subsets of UniRef50, varying token counts from 6.5B to 20.8B to simulate different compute allocations.
• Used spline interpolation of training loss curves to map FLOPs to optimal model size and token count at each compute level, enabling precise identification of the most efficient configurations.
• Fitted a combined scaling law: $ L(N,D) = ig[ (N_c/N)^{\alpha_N/\alpha_D} + (D_c/D) \big]^{\alpha_D} $, to model loss as a function of model and data scale.
• Identified the optimal model size for fixed compute budgets (e.g., $10^{17}$ FLOPs) via polynomial fitting of loss plateaus, predicting optimal sizes of 50.6B and 157B parameters at $10^{17}$ and $10^{18}$ FLOPs, respectively.
• Evaluated model performance on held-out sets using perplexity, comparing a 35M-parameter model trained in one pass to larger models like ESM-2 (15B) and xTrimoPGLM (100B).

Figure 1: FLOPs vs loss with token-based learning rate decay.

Experimental results

Research questions

• RQ1 What is the optimal ratio between model parameters and training tokens for protein language models under a fixed compute budget?
• RQ2 Are widely used pLMs such as ESM-2 and xTrimoPGLM compute-optimal, or do they exhibit diminishing returns?
• RQ3 Can a small pLM achieve performance comparable to large models when trained with optimal compute allocation and early stopping?
• RQ4 How does the optimal model size scale with increasing compute budgets in pLM pretraining?
• RQ5 Does multiple-pass training over the same data significantly improve performance beyond the loss plateau observed in single-pass training?

Key findings

• A 35M-parameter model trained on a 20M-sequence subset of UniRef50 in a single pass achieved a perplexity of 11.43, outperforming larger models like ESM-2 (15B) and xTrimoPGLM (100B) in terms of compute efficiency.
• The study identified a consistent loss plateau across model sizes, indicating that performance gains plateau regardless of model or data scale, suggesting a fundamental limit in current pLM training.
• For a compute budget of $10^{17}$ FLOPs, the optimal model size is approximately 50.6 billion parameters, increasing to 157 billion at $10^{18}$ FLOPs, showing sublinear scaling of optimal model size with compute.
• Larger models reach the loss plateau with fewer FLOPs due to higher sample efficiency, indicating that increased capacity improves compute efficiency despite higher parameter counts.
• The optimal model size increases sublinearly with compute budget, and the optimal token count also scales sublinearly, indicating diminishing returns in performance gains as model size grows.
• The results suggest that current pLMs like ESM-2 and xTrimoPGLM are not compute-optimal, as they require significantly more FLOPs to reach the same performance plateau as smaller, better-optimized models.

Figure 2: FLOPs vs loss without learning rate decay.