Training Compute-Optimal Protein Language Models

About the Weakness-1. Thank you for your valuable feedback. Based on your suggestion, we will adjust the x-axis range and include the trained maximum models for 10**22 FLOPs, namely MLM 10.7B and CLM 7.2B in section 5.

About the Weakness-3: the evaluation of protein generation. We explain more details about Protein Generation Comparison as follows:

• OOD Dataset PPL Analysis: PPL represents the probability of the test sequences in the model distribution. The lower the PPL, the closer the distribution of the model and the test set is. In order to test the generalization ability of the model on new data, we use different sequence identity (0.0, 0.3, 0.5) as thresholds to select the test set;
• pLDDT scores from ESMFold predicted Local Distance Difference Test is the confidence level of ESMFold protein structure prediction. This metric is widely used in methods such as AlphaFold2, ESMFold, and OpenFold. pLDDT filters are often used in protein design (such as RFDiffusion), which can significantly improve the success rate of protein design;
• Natural Sequences Comparisons with Foldseek Foldseek takes protein structure as input and searches for proteins with similar structure in the database. We use the experimentally-resolved protein structure as the database (PDB database) to explore how the structure of the generated sequences close to PDB (a higher TM-score indicates higher structural similarity). This method has been used to evaluate other methods for protein sequence generation (ProGen2, ProtGPT2);
• Diversity Analysis: We cluster the two sets of sequences (ProGen2-xlarge and CLM) according to sequence similarity. Sequences with a identity higher than 50% will stay in one cluster. Since the number of input sequences is similar (8,466 vs 8,263), we can measure the diversity of the generated sequences by comparing the number of clusters.

About Question-1: transfer between mlm and clm. The effectiveness of pre-training models under a given training budget fundamentally determines their capacity. In Causal Language Modeling (CLM), the loss gradient is calculated for all tokens in the sequence, meaning that each token contributes to the model's learning process. This comprehensive contribution makes CLM training more efficient. Conversely, in Masked Language Modeling (MLM), the loss is calculated only for a subset of tokens—typically around 20% that are masked. Consequently, only these masked tokens contribute to the model's learning, resulting in lower training efficiency compared to CLM. Therefore, when MLM is transferred from a pre-trained CLM, it benefits from the higher overall training efficiency of CLM, leading to improved performance of the finalized MLM model. However, when CLM is transferred from a pre-trained MLM, the lower training efficiency of MLM results in a sub-optimal CLM model.

About the Question-1: the re-organization of the paper. We appreciate your insights and have taken them into consideration for our revisions. Although we cannot directly modify the paper PDF during the rebuttal period, we demonstrate how we will reorganize and improve the writing style of our manuscript in the next version of our paper.

Add Discussion/Conclusion Sections: We will add conclusions or takeaway findings at the end of each main section. Specifically:

For Section 2: Studying the Data-hungry Observation & Scaling Up Data, we will add discussions to support the two key findings:

• Scaling the model from 150M to 3B, we noted diminishing returns for CLM and an overfitting tendency for MLM when repeating the UR50/S dataset.
• Our proposed Expanding Diversified Metagenomic Data (UniMeta200B) addresses these problems.

For Section 3: Scaling Law for CLM/MLM Protein Language Models, besides the detailed scaling law equations, we will add discussions for the two key findings:

• In CLM, training data scales sublinearly with model size but follows a distinct power-law with a compute exponent of 0.57. A 10× increase in compute leads to a 4× increase in model size and a 3× increase in training tokens.
• In MLM, training data scales sublinearly with model size, approximately following a power-law with a compute exponent of 0.77. A 10× increase in compute results in a 6× increase in model size and a 70% increase in training data.

These simplified takeaway scaling laws will benefit AI researchers who want to train their own optimal PLM under a computational budget.

For Section 4: Transfer Scaling Scenarios, besides the specific transfer scaling law equations obtained by empirical experiments, we will add discussions about:

• The significance of the transfer scaling law: Researchers can obtain better MLM models transferred from CLM pre-training under the guidance of the transfer scaling law than merely MLM from scratch following the General MLM scaling law obtained in Section 3.
• The simplified takeaway transfer scaling law: Training MLM from scratch with 10× the compute requires approximately 7.7× the compute compared to MLM from CLM pre-training, implying that around 20% of the compute budget should be allocated for CLM pre-training to get better MLM models transferred from CLM pre-training.

For Section 5: Confirming the Scaling Law Effect in Downstream Tasks:

• Using our concluded scaling law, we obtain a training-optimal 7.2B CLM under the total compute of ProGen2-xlarge, which shows better generation capacities as shown in Figure 6.
• We also obtain a training-optimal 10.7B MLM model under the same FLOPs as ESM2-3B, which performs better on 8 protein understanding tasks, as shown in Table 1 in the attached PDF.
• The 470M MLM transferred from CLM outperforms MLM from scratch under the same training budget.

To summarize: we aim to establish a practical pathway for researchers to develop faithful and powerful protein language models optimized by both CLM and MLM objective in an end-to-end manner. This includes everything from pre-training dataset construction to optimal parameter and dataset allocation, as well as knowledge transfer from other pre-training objectives. Using our approach, we are able to obtain better MLM or CLM models compared to other well-known suboptimal PLM models such as ProGEN2 and ESM. Our work holds significant potential for the application of large language models across various scientific domains.

Finally, we will move the discussion and limitations of our work from the Appendix Sections A, B, and C to the main body of the paper to highlight the great potential and future research directions of our work.

We believe that these revisions address the concerns raised and significantly enhance the clarity and impact of our manuscript.

About the Question-2: the confidence interval. We used scipy's curve_fit to fit the parameters, such as in the following example:

params, covariance = curve_fit(exponential_fit, x, y, maxfev=10000)

We did not detail the confidence intervals in our paper, but we will add them using the pseudo code below:

def get_lower_and_upper_bound(params, cov):     
    perr = np.sqrt(np.diag(cov))     
    alpha = 0.05
    z = stats.norm.ppf(1 - alpha/2)  # 1.96 for 95% CI
    lower_bounds = params - z * perr       
    upper_bounds = params + z * perr 
    return lower_bounds, upper_bounds

The calculated z×perr for the scaling exponents of the CLM is 0.028, and for the MLM, z×perr is 0.014. Thus, we include the confidence intervals at a confidence level of 1−α=0.95 with the lower and upper bounds shown in parentheses:

• CLM:

  • Model size scaling exponent: 0.578 (0.554, 0.602)
  • Data size scaling exponent: 0.422 (0.394, 0.449)

• MLM:

  • Model size scaling exponent: 0.776 (0.768, 0.784)
  • Data size scaling exponent: 0.230 (0.216, 0.244)

About the Question-3: the particular functional form of transfer scaling law. Thank you very much for your suggestion. You raise a very good point about the lack of a single standard definition for the scaling law formula. For instance, OpenAI's scaling law differs from Chinchilla's approach. In fact, an additive scaling law could also be used to address your question. The reason we chose the multiplicative form is mainly empirical and was inspired by the following two papers [1,2] ~(in section 4.2). [2] provides a detailed comparison between additive and multiplicative scaling laws, concluding that there is not much difference between the two, although the multiplicative form achieves slightly lower extrapolation error on average. Therefore, we adopt this for follow-up analysis. Due to space limitations, we couldn't enumerate all possible experiments, so we naturally chose this one. We will add this explanation in the paper.

[1] Scaling Law for Transfer

[2] WHEN SCALING MEETS LLM FINETUNING: THE EFFECT OF DATA, MODEL AND FINETUNING METHOD

About the Weakness-1: the lack of discussions. Thank you for your insightful suggestions. We will make the following modifications to significantly enhance the clarity and impact of our manuscript.

• Adding conclusions or takeaway findings at the end of each main section.
• Moving the detailed discussion and limitations of our work from Appendix Sections A, B, and C to the main body of the paper to highlight the great potential and future research directions of our work.

About the Weakness-2: the clarification of novelty. While Chinchilla focuses on optimizing the model size and number of tokens for training a causal language model within a given compute budget and provides open-source models, our work takes a novel approach. We aim to establish a practical pathway for researchers to develop faithful and powerful protein language models optimized by both CLM and MLM objective in an end-to-end manner. This includes everything from pre-training dataset construction to optimal parameter and dataset allocation, as well as knowledge transfer from other pre-training objectives. Our work holds significant potential for the application of large language models across various scientific domains.

About the Weakness-3: About the experimental validation. We have provided a comprehensive task performance comparison across 8 protein understanding benchmarks, as adopted from the protein understanding benchmarks available at biomap-research on Hugging Face. These benchmarks span four distinct categories of protein-related challenges:

• Protein Structure: Contact Prediction, Fold Prediction
• Protein Function: Fluorescence, Fitness-GB1, Localization
• Protein Interaction: Metal Ion Binding
• Protein Developability: Solubility, Stability

The results, shown in Table 1 in the attached PDF, demonstrate that our 10.7B model outperforms ESM-3B on x out of x tasks. This confirms the rationale behind the observed scaling law and addresses concerns about the scope and rigor of our evaluation tasks.

About the Question-1: the scaling law. It is indeed common practice to use smaller models in many practical applications, even though these models may not follow Compute-Optimal Training principles. Smaller models are more deployment-friendly, especially for models with high traffic and for downstream task fine-tuning. We did not perform this experiment because it falls beyond the scope of this paper, which primarily discusses how to train compute-optimal models. If we were to discuss the scenario of small models with high computational power, it would require further examination of data repetition issues, introducing a new variable into the analysis. We addressed this in the limitations section (Appendix C), noting that for MLM models, smaller models trained with relatively high computational power can have comparable performance to optimally trained models while being more inference and deployment-friendly. For example, a 3B model trained with 5 epochs (1T tokens) uses the same computational budget as an optimally trained 10.7B model with 265B tokens, and the performance gap between the 3B and 10.7B models is minimal. However, we know that continuously repeating data can lead to diminishing returns, and some models may even suffer from overfitting. Different model sizes correspond to an optimal data repetition variable. Without accurately determining this variable, it's challenging to predict how much loss smaller models might incur due to data repetition, unless we have infinite data. This area requires further exploration and generalization of current scaling laws, such as data-constrained scaling laws [1]. We anticipate more work in this direction soon.

[1] Scaling Data-Constrained Language Models

About the Question-2: the downstream task scaling law. Thank you for your suggestion. Research on the downstream performance of scaling has already emerged in the field of LLMs [1]. This research is indeed fascinating, but it falls beyond the scope of this paper. Due to limitations of the paper space and computational resources, a deeper discussion would introduce other scaling factors such as downstream fine-tuning data size. We briefly compared CLM and MLM under the same data and two different model sizes in Appendix F, using the same computational cost (same FLOPs) for the Contact Prediction downstream task. The conclusion is that MLM consistently outperforms CLM. For a more detailed exploration of scaling laws in downstream tasks, including CLM generation and MLM, you can refer to the xTrimoPGLM paper [2].

[1] When Scaling Meets LLM Finetuning: The Effect of Data, Model, and Finetuning Method

[2] xTrimoPGLM: Unified 100B-Scale Pre-trained Transformer for Deciphering the Language of Protein

About Weakness-1: the protein understanding tasks and missing experimental details. We have provided a comprehensive task performance comparison across 8 protein understanding benchmarks, as adopted from the protein understanding benchmarks available at biomap-research on Hugging Face. These benchmarks span four distinct categories of protein-related challenges:

• Protein Structure: Contact Prediction, Fold Prediction
• Protein Function: Fluorescence, Fitness-GB1, Localization
• Protein Interaction: Metal Ion Binding
• Protein Developability: Solubility, Stability

The results, shown in Table 1 in the attached PDF, demonstrate that our 10.7B model outperforms ESM-3B on 7 out of 8 tasks. This confirms the rationale behind the observed scaling law and addresses concerns about the scope and rigor of our evaluation tasks.

About the missing experimental details. The downstream data are collected from corresponding publication and have been well split, we employed the same division for evaluation. We will add the description in the future version. Additionally, we did not implement any specific strategy to filter out downstream data during evaluation. Given the nature of pre-training, the pre-training process is self-supervised, and no label-related information is seen. Although some samples have a high sequence identity with pre-trained sequences, the lack of exposure to label information during pre-training should prevent any substantial data leakage effects.

About Weakness-2: the lack of diverse sampling strategies. Thank you for your question and great thinking. Indeed, different sampling strategies may change the loss curve and potentially alter the scaling laws. Regarding the diversity of data selection, we chose to use i.e., UniRef90, and ColabFoldDB. These is one result of sampling strategies from the original datasets. Our choice is based on balancing diversity and data volume. For example, previous ESM models performed very well on UniRef90, as shown in Figure 4 of the referenced paper [3]. Using the entire Uniref100 would degrade model performance [2], while using UniRef50 would result in insufficient data, necessitating data repetition, which introduces another variable. This would require considering scaling laws under different sampling strategies as well as under repetition [1]. Thus, our initial assumption was to pass the data through only once or just slightly more, following a uniform distribution. This assumption allows us to preliminarily avoid the issue of data repetition. Introducing high diversity sampling strategies would require further research on data-constrained scaling laws [1]. Therefore, we initially considered simpler settings, such as those used in early OpenAI and Chinchilla studies. Regarding ColabFoldDB, its sequence identity is relatively low. In Appendix G, we compared two small models on UniRef90 and ColabFold datasets, both under uniform sampling conditions, and their performances were quite similar. This outlines our rationale for data and sampling strategy selection. The proposed diverse sampling strategy is an excellent point, In the field of LLMs, there has been some work done[4], but there is still no optimal conclusion. Exploring new scaling laws by considering sampling strategies and data repetition variables would be very interesting for future work.

[1] Muennighoff, Niklas, et al. "Scaling data-constrained language models." NIPS'24.

[2] Rives, Alexander, et al. "Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences." PNAS'21

[3] Meier, Joshua, et al. "Language models enable zero-shot prediction of the effects of mutations on protein function." NIPS'21

[4] Xie, Sang Michael, et al. "Doremi: Optimizing data mixtures speeds up language model pertaining." NIPS'24

About Weakness-3: the lack of insights.

Regarding the overfitting issue in masked language modeling: In MLM, a certain percentage of input tokens are masked randomly, and the model is trained to predict these masked tokens using bidirectional context. The bidirectional self-attention mechanisms used in MLM have a higher capacity to overfit compared to the unidirectional self-attentions used in CLM. This is because MLM can utilize the entire context surrounding a masked token, leading to faster memorization of the training data. A more systematic study investigating the training dynamics of MLM and CLM is presented by Tirumala et al. [1]. This study observes that MLM tends to memorize faster than CLM, which leads to overfitting, especially when the pre-training data is repeated over multiple epochs. The bidirectional nature of MLM allows for a more comprehensive understanding of the context, but it also means that the model can become overly reliant on specific patterns in the training data, leading to overfitting.

[1] Tirumala, Kushal, et al. "Memorization without overfitting: Analyzing the training dynamics of large language models." NIPS'22.

Differences in Scaling Behavior: Protein Sequences vs. Natural Language Sequences The scaling behavior between protein sequences and natural language sequences also differs significantly. Protein sequences are governed by the rules of biological systems and have a different structural complexity compared to natural language sequences. Proteins have specific folding patterns and functional constraints that are not present in natural language.

When scaling models for protein sequences, the complexity of the biological rules and constraints means that the models need to capture a wide range of interactions and dependencies. This can lead to different scaling laws compared to natural language models, which primarily focus on syntactic and semantic coherence. The specific coefficient parameters between protein language model (PLM) scaling laws and natural language processing (NLP) scaling laws confirm this difference.

Thanks for your clarifications! My concerns are addressed.