Title: Training Compute-Optimal Protein Language Models
Venue: NeurIPS 2024
Accepted: yes
Keywords: Protein Language Model

Scaling Law
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 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.

Training Compute-Optimal Protein Language Models

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.

1 Introduction

Scaling up transformer-based models has become a guiding principle for enhancing model performance across broad domains, particularly in Natural Language Processing (NLP) [4, 12, 24, 64, 79] and Computer Vision (CV) [20, 67, 89]. In recent years, large transformer-based Protein Language Models (PLMs) such as PROGEN familiy [51, 61], ESM familiy [68, 47] and xTrimoPGLM [14] have also been developed, which leads to significant improvements over model performance on complex downstream tasks [27, 45]. Current language models utilize two main training objectives to encode sequence information: the BERT-like [23] Masked Language Model (MLM) and the GPT-like Causal Language Model (CLM) [11], each applied either separately or in a unified fashion. A common understanding is that bi-directionally MLM excels in sample efficiency and shows enhanced performance in downstream task fine-tuning. This is particularly true in tasks that emphasize understanding complex patterns, making it a prevalent learning objective in modeling protein sequences *

[46, 14]. On the other hand, uni-directional CLM, due to its sequential generation ability, is better suited for generating more coherent and realistic sequences compared to MLM [19, 61, 63].

However, training large protein language models (PLMs) are computational-intensive, and strategies for optimally allocating compute budgets for training PLMs are relatively underexplored, with most efforts focusing on scaling model parameters based on a fixed set of training tokens to achieve performance improvements . A key insight [36, 42, 76] is that large models should not be trained to their lowest possible loss to optimize computing; instead, models and data should be scaled proportionally based on available compute budgets. These scaling laws are broadly found in natural language models [42, 36, 33, 2, 58, 78, 17, 90]. But their applicability has not been validated within biological datasets, such as the primary structures of proteins, which are composed of amino acid sequences forming protein chains. Unlike natural languages, protein sequences are scientific data that are precisely represented using a vocabulary of 20 amino acids, with very little redundancy and are not as semantically smooth. Thus, we consider such data as a distinct modality and ask the question: What are the scaling behaviors for MLM and CLM in protein language modeling?

We focus on the best practices, which include revisiting datasets, optimization objectives, and model parameters as key factors. Our goal is to investigate an optimal training scheme for protein language models given predetermined compute budgets. Our core findings are as follows:

We revisited the protein sequence data used for training PLMs and collected a dataset of 194 billion unique tokens on 939M unique sequences from publicly available sources to address the issue of overfitting and perform plateau in protein language modeling.

We find that, in both MLM and CLM, training data scales sublinearly in the model sizes but follow distinct power-laws. In other words, a 10 × increase in compute leads to a 6 × increase in MLM model size and a 70% increase in data, versus a 4 × increase in CLM model size and a 3 × increase in training tokens.

We also find that models trained with CLM can be transferred to MLM. When given a predetermined amount of computation, and one wants to obtain both a CLM and a MLM model, there is a trade-off in allocating the training token to each model to jointly optimize the performance of the two. Interestingly, the allocation for CLM pre-training was determined by the scaling law of CLM and MLM, and the Effectively Transferred Tokens D t from CLM to MLM. Furthermore, we verify this method experimentally using a 470M model and fine-tuning on downstream tasks.

Building on our scaling strategies, we re-allocate of model size and training tokens under the compute budgets of established PROGEN2-xlarge and ESM-2 (3B) setups. Consequently, with the same compute budgets, we trained two corresponding models, one with 7.2B parameters and the other with 10.7B parameters, which exhibited enhanced performance in a diverse range of downstream tasks.

2 Scaling up data

First, we explore the effects of training PLMs across multiple epochs under token scarcity conditions. We then introduce a dataset, UniMeta200B, used throughout this work. This dataset enhancement alleviates the challenge of insufficient training for protein language models.

2.1 A Data-hungry Observation

Using the UniParc database with 250 million protein sequences, research on ESM [68] shows that the datasets UR50/S and UR50/D, with 45M and 65M unique sequences respectively, outperform Uniref100 in perplexity (PPL) on a ~670M parameter MLM model. These datasets contain ~15B and ~20B unique amino acid tokens. The ESM-2 family models, ranging from 150M to 15B parameters, are trained extensively with nearly 1 trillion tokens over 45 epochs on the UR50/D dataset. In observing the scaling of ESM-2 models, it becomes apparent that increasing model size to 15B parameters from 3B shows marginal improvement. On the other hand, contemporary LLMs are often trained for only one or a few epochs [43, 36, 79, 80, 11]. The repetition of data with limited unique tokens has diminishing returns and hinders scaling model size [65, 34, 58, 70]. This underscores the importance of using rich datasets for training large-scale language models to ensure robust performance across applications. We evaluated models with 150M and 3B parameters on the UR50/S dataset, trained on 200B tokens, as shown in Figure 1. We focus on the Independent and Identically

2.0

2.2

2.4

2.6

2.8

Training Loss

UR50/S-150M

UniMeta200B-150M

UR50/S-3B

UniMeta200B-3B

6

8

10

12

14

16

18

Validation Perplexity

14

15

16

17

18

OOD Dataset Perplexity

15B

77B

138B

200B

Tokens

1.6

1.8

2.0

2.2

2.4

2.6

2.8

Masked Language Model

15B

(1)

77B

(5)

138B

(9)

200B

(13)

Tokens

(UR50/S Epochs)

6

8

10

12

14

16

18

15B

(1)

77B

(5)

138B

(9)

200B

(13)

Tokens

(UR50/S Epochs)

14

15

16

17

18

Distributed (IID) validation and Out-Of-Distribution (OOD) test PPL, which measures the model's randomness in amino acid selection. For our OOD dataset, we utilized the MMseqs2 tool [72] to conduct searches within the UniRef90 database for sequences post-training dataset timestamp, retaining those with no detectable identity. From these, a random sample of 3,000 sequences was selected to constitute the OOD dataset. Notably, we do not adopt dropout regularization, a practice that often reduces model capacity and is infrequently used in contemporary LLMs [43]. This choice is consistent with recent LLM configuration findings [38], including ESM-2 [47].

The results show the 150M model lacks good generalization while increasing to a 3B model resulted in diminishing returns for CLM and severe overfitting for MLM. Principally, the bidirectional selfattention mechanisms used in MLM have a higher capacity to overfit compared to the unidirectional self-attention used in CLM. This is because MLM can utilize the entire context surrounding a masked token, leading to faster memorization of the training data.

2.2 Expanding Diversified Metagenomic Data

To tackle the challenge of data scarcity, we leveraged the ColabFoldDB database [56], which focuses on metagenomic data sources such as BFD [1], MGnify [57], and specific eukaryotic and viral datasets including SMAG [22], MetaEuk [44], TOPAZ[3], MGV [59], and GPD [13]. We applied a stringent deduplication process with a maximum similarity threshold of 0.3 to preserve the diversity of the protein universe. Given that the Uniref90 dataset has proven most

effective for pre-training across various Uniref clustering levels per ESM-1v [54], we incorporated Uniref90/50 (Before 2022-12), which includes incremental data relative to Uniref50/S representatives. ColabFoldDB c and ColabFoldDB m play dominant roles within the dataset, corresponding to cluster representatives and members, respectively. To ensure uniformity during training, we allocate weights within each batch to allow each amino acid token to be evenly processed through the model. This dataset, termed UniMeta200B, contains 939 million unique protein sequences and 194 billion amino acids, which is an order of magnitude larger than UR50/D. We observed significant improvements

in the OOD test set and a consistent learning curve on the IID validation subset extracted from the training set (Figure 1). These enhancements not only ensure a controlled diversity to maintain sample efficiency but also significantly increase the quantity and uniformity of data, facilitating model scaling. †

3 Parameters and Datasize Optimal Allocation

10 7

10 8

10 9

10 10

Model size

2.2

2.3

2.4

2.5

2.6

CLM Validation Loss

1e+18 FLOPs

3e+18 FLOPs

1e+19 FLOPs

3e+19 FLOPs

1e+20 FLOPs

3e+20 FLOPs

1e+21 FLOPs

10 7

10 8

10 9

10 10

Model size

2.0

2.1

2.2

2.3

2.4

MLM Validation Loss

1e+18 FLOPs

3e+18 FLOPs

1e+19 FLOPs

3e+19 FLOPs

1e+20 FLOPs

3e+20 FLOPs

1e+21 FLOPs

10 17

10 18

10 19

10 20

10 21

10 22

10 23

10 24

FLOPs

10 6

10 7

10 8

10 9

10 10

10 11

10 12

Model size

N MLM

opt

=6.19

10

8

C

0.776

N CLM

opt

=1.26

10

3

C

0.578

10 17

10 18

10 19

10 20

10 21

10 22

10 23

10 24

FLOPs

10 9

10 10

10 11

10 12

10 13

Data size

D MLM

opt

=2.02

10 6

C

0.230

D CLM

opt

=1.24

10 2

C

0.421

In this section, we propose a scaling law for protein sequences with MLM and CLM objectives, aiming at optimally balancing model size and data size under a fixed compute budget to improve efficiency on expanded resources.

3.1 Scaling laws for CLM and MLM

We first fit our models in the form of a fundamental power-law based on the existing work [42, 36, 33, 69, 2, 17, 90] in the field of LLMs. Specifically, given a fixed FLOPs formula of C = 6 × N × D , where N represents the

number of forward-activated non-embedding parameters, and D is the number of training tokens, how should one navigate the trade-off between model size and the number of training tokens? The

model parameters N and data size D can be directly fit with a simple power-law:

Weemployed the IsoFLOPs profiling approach [36, 9], setting 7 distinct training FLOP counts ranging from 1 × 10 18 to 1 × 10 21 . For each FLOP count, we selected models from a pool of candidates (see Appendix N). Models were excluded if the estimated data size ( C/ (6 ∗ N ) ) resulted in more than 200B tokens or if the training steps were fewer than 20K. Ultimately, approximately 260 models were used for fitting. We considered the final validation loss for each model to ensure that every model completed a full cosine cycle with 10 × learning rate decay. For each fixed FLOP count, we employ smoothed loss to determine the optimal model size with the smallest loss (Figure 2 (above)). Subsequently, we use Equation 1 and apply the least_squares method to fit the model. Given

the minimal variations in the final loss among a set of ( N,D ) configurations, we classify these configurations as operating under "IsoLoss" conditions (see Appendix K Figure A15), considered optimal for training. In Figure 2 (below), we illustrate an efficient frontier interval that demonstrates permissible fluctuations in model size and dataset size at a specific FLOP count, while still achieving nearly identical losses. The variation in loss is quantified at 0.25 on a logarithmic scale with a base of 10. This indicates that within this FLOP counts, the model size can be adjusted within a range, increasing up to 80% or decreasing up to 40% without repeating data, to maintain a loss variation within 0.01.

We observe distinct growth rates in the proportional relationship between model size and training tokens for the MLM model compared to the CLM, as detailed in Table 2. Both models demonstrate an increase in the growth of model size that surpasses the growth of training tokens. Up to the intersection point around 1 × 10 22 (see Figure 2, left below), the model size of MLM tends to be smaller than the CLM, thereafter, the MLM rapidly exceeds that of the CLM. Notably, the growth of the MLM's training tokens is greatly lower than that for the CLM, possibly due to MLM's higher sample efficiency. For instance, if the compute budget is increased by 10 × , the size of the CLM model should increase by 4 ×

10 18

10 19

10 20

10 21

10 22

FLOPs Sum (

Csum

)

10 7

10 8

10 9

N

(

Csum

) = 1.497

10

6

C

0.703

sum

and the training data by 3 × , aligning more closely with equally proportional scaling. For the MLM, the model size should increase by 6 × and the training data size by 1 . 7 × .

In exploring the scaling relations of loss, we analyzed various model sizes N , compute budgets C , and training dataset tokens D . These can be described by a similar power-law relation defined as:

where α x is the scaling exponent for different variables. For each FLOP count, we aimed to identify the minimal loss as the fitting target along with the corresponding independent variable x . Table 3 presents these fitting coefficients.

Based on the coefficients obtained from the fitting described above, we can establish the relationship between D and N by eliminating L . The relationship is expressed by the following equation:

By substituting the learned coefficients into this formula, we can derive D opt MLM and D opt CLM when given N . The estimation may be affected when the data exceeds 200 billion or when the quality or quantity of the training dataset changes.

3.2 Scaling law for training two models

When our goal is to optimize both CLM and MLM simultaneously, the strategic allocation of compute resources between these two objectives becomes essential. To facilitate this, we equalize model parameters across objectives to assess specific compute budgets for dual-objective training. Specifically, we seek the compute budgets, C MLM and C CLM, for configurations where the optimal model size is the same, i.e., N ( C MLM ) = N ( C CLM ) . These individual computations are then aggregated to formulate the overall compute budget:

These two objectives share the same parameter size, their compute budget C and the number of training tokens D differ. Thus we further introduce a model-to-ratio r ( N ) as D MLM ( N ) /D CLM ( N ) . We then achieve the relationship between N and C sum by a fitted power-law (Figure 3) form:

The ratio r ( N ) informs us about the allocation proportion of training tokens. Specifically, under equal parameters, the data for MLM should exceed that for CLM until a 10B threshold (achieving a 1:1) is reached, after which more training tokens are allocated to CLM.

We further find that the scaling behavior of sparse parameter counts in Mixture of Experts (MoE) protein model [74], configured with eight experts (see Appendix I), as well as a combined power-law formula used to fit our data (see Appendix J), both exhibit a similarity to the scaling behavior we have proposed.

Findings 2. In both CLM and MLM, training data scales sublinearly with model size, following distinct power laws. With an 'infinite' dataset, where samples are not repeated and training for less one epoch, MLM's model size grows faster than CLM's.

4 Transfer Scaling

We have outlined two independent scaling laws and how to allocate FLOPs under a fixed budget for training two optimal models, one with MLM and the other with CLM. However, we have not explored the interaction between these objectives. This raises important questions: Can models trained with one objective transferred to one with another objective? Is there a synergistic effect from training two models? Does training order impact the results?

4.1 Transferability

We conduct transfer learning experiments on MLM and CLM objectives, selecting eight optimal model sizes based on Equation 1. These models correspond to four increasing FLOP counts from 3 × 10 19 to 1 × 10 21 and undergo training from scratch followed by transfer training. Transfer training involves initially training on MLM or CLM, then training on the alternate model for each size.

We find that optimal pre-training on one objective benefits the target objective in transfer learning, though effects vary between methods. Starting with CLM and then training MLM, benefits increase with model scale. In contrast, starting with MLM then training CLM sees diminishing benefits. As shown in Figure 4 (left), for a model size of 230M with 3 × 10 19 FLOPs, CLM from MLM pre-training reduces the loss by 0.02 compared to CLM from scratch, however, benefit that nears zero for the 1.7B model. Conversely, for models from 85M to 1.2B, transfer benefits grow with model size, the compared validation loss gap increasing from 0.025 to 0.045. This likely stems from the higher loss utilization rate in CLM; CLM calculates losses for all tokens in a protein sequence, whereas MLMonly calculates losses for 15% of the tokens. ‡ .

MLM loss curve from scratch

CLM loss curve from scratch

We use a power-law to model the transfer scaling law, initially excluding the pre-training FLOPs. The scaling behavior of transfer learning is modeled by:

where L ( C t ) and L ( C s ) represent the loss for transfer learning and training from scratch.

Figure 4 (right) shows that the efficient frontier for L ( C t ) has shifted relative to L ( C s ) (it can be directly obtained from Table 3, repeated here for convenience.), indicating an improvement. The coefficients from both are shown in Table 4, where we can infer that C t ∝ C αs αt s = C 0 . 89 s ,

suggesting that training MLM from scratch with 10 × the compute requires approximately 7.7 × the compute compared to MLM from CLM pre-training. Another observation is that mixing training objectives in a single batch tends to be detrimental. Detailed results and settings are in Appendix H. The recommended transfer learning schedule involves pre-training CLM before MLM, as mixed training and order swapping show no benefits. We speculate that this primarily occurs because our MLM, which focuses solely on recovering corruption tokens, is not causal. If it predicted a middle segment in a left-to-right manner, it could mutually adapt with the context to accelerate training [86].

4.2 Effectively Transferred Tokens

Although we observe that MLM benefits from transfer learning from CLM, the pre-training compute budget remains unaccounted for. We focus on two aspects: (1) the actual benefit CLM provides to MLMand its predictability, and (2) performance differences between MLM trained from pre-trained

0

10

20

30

40

50

60

% Compute of CLM Pre-training

8.0

8.5

9.0

9.5

Validation PPL | Optimal MLM Size (FLOPs)

From scratch

Optima interval

34M (1e19)

85M (3e19)

200M (1e20)

470M (3e20)

Eff. Trans.

Dt

0B

12B

50B

59B

Tokens

9.0

9.5

10.0

10.5

11.0

11.5

12.0

12.5

13.0

Validation Perplexity

Pre-trained CLM Tokens

Dt

85M MLM from pre-trained CLM

85M MLM from scratch

CLM (MLM-CLM) and MLM from scratch (MLM-S) under identical FLOP constraints. We define Effectively Transferred Tokens D t as the additional data a model of the same size would need to train from scratch on MLM to achieve the same loss as a model pre-trained on CLM. If the token number in the pre-trained CLM model exceeds D t , then the computations for CLM pre-training was excessive. Knowing D t in advance would guide the allocation of tokens for CLM pre-training.

We compare MLM-S and MLM-CLM models ranging from 33M to 1.2B with FLOP counts from 3 × 10 19 to 1 × 10 21 . By calculating the token distance at the same loss level between these models, we establish our fitting target D t , collecting approximately 2800 sample points. Following similar methods in scaling transfer works [35, 90], D t is defined by a simple multiplicative scaling formula:

where D f represents the tokens used for MLM-CLM, and N is the number of parameters, with k , δ , and γ as fitting coefficients. For instance, a 10 × increase in D f would roughly triple the model size and double D t . We validate these findings by evaluating the compute ratio of CLM pre-training under four specified parameters and FLOPs, as shown in Figure 5 (left), finding that MLM-CLM generally outperforms MLM-S. Specifically, D t / ( D t + D f ) ranges from 10% to 20% of the compute budget for CLM pre-training. Figure 5 (right) schematically illustrates the learning curves of two 85M (3e19 FLOPs) models, with MLM-CLM achieving similar or better loss levels with equal or fewer tokens.

Findings 4. 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.

5 Experimental Validation

Based on the scaling laws we observe, we estimate the model size and training tokens for current leading models by analyzing their FLOPs. In our configuration, the PROGEN2-xlarge model, with 6.4B parameters, is estimated to require training with 7.2B parameters and 265B tokens. Similarly, the ESM-2 model, with 3B parameters, should be trained with a model size of 10.7B parameters and 260B tokens. Additionally, we employed two 470M models to test the transfer scaling strategy, one trained from scratch (470M scratch) and the other from CLM pre-training (470M trans.). The model's details are reported in Table 5.

GLYPH<c=18,font=/AAAAAB+DejaVuSans> GLYPH<c=24,font=/AAAAAB+DejaVuSans>

GLYPH<c=23,font=/AAAAAB+DejaVuSans>GLYPH<c=15,font=/AAAAAB+DejaVuSans>GLYPH<c=20,font=/AAAAAB+DejaVuSans>GLYPH<c=22,font=/AAAAAB+DejaVuSans>GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=4,font=/AAAAAB+DejaVuSans>GLYPH<c=25,font=/AAAAAB+DejaVuSans>GLYPH<c=16,font=/AAAAAB+DejaVuSans>GLYPH<c=19,font=/AAAAAB+DejaVuSans>GLYPH<c=24,font=/AAAAAB+DejaVuSans> GLYPH<c=21,font=/AAAAAB+DejaVuSans>GLYPH<c=20,font=/AAAAAB+DejaVuSans>

GLYPH<c=11,font=/AAAAAB+DejaVuSans>GLYPH<c=5,font=/AAAAAB+DejaVuSans>GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=13,font=/AAAAAB+DejaVuSans>GLYPH<c=3,font=/AAAAAB+DejaVuSans>GLYPH<c=14,font=/AAAAAB+DejaVuSans>GLYPH<c=16,font=/AAAAAB+DejaVuSans>GLYPH<c=17,font=/AAAAAB+DejaVuSans>

GLYPH<c=8,font=/AAAAAB+DejaVuSans>GLYPH<c=6,font=/AAAAAB+DejaVuSans>

GLYPH<c=9,font=/AAAAAB+DejaVuSans>GLYPH<c=6,font=/AAAAAB+DejaVuSans>

GLYPH<c=10,font=/AAAAAB+DejaVuSans>GLYPH<c=6,font=/AAAAAB+DejaVuSans>

GLYPH<c=12,font=/AAAAAB+DejaVuSans>GLYPH<c=6,font=/AAAAAB+DejaVuSans>

GLYPH<c=7,font=/AAAAAB+DejaVuSans>GLYPH<c=6,font=/AAAAAB+DejaVuSans>GLYPH<c=6,font=/AAAAAB+DejaVuSans>

20

80

100

ESMFold pLDDT

PROGEN2

7.2B CLM

(N=8466)

(N=8263)

B

C

FoldSeek TM-score

PROGEN2-xlarge

7.2B CLM

Clusters

Cluster size

D

8,466 seqs

4,818 clusters

8,263 seqs

7,097 clusters

Sequence clustering at 50% ID cutoff

A

(N=514)

(N=769)

(N=1756)

MaxID

Perplexity

SeqID

Figure 6: Comparative Analysis of CLM Models. A. Perplexity analysis for PROGEN2-xlarge and our 7.2B CLM shows lower values for our model across various MaxID levels, suggesting better sequence handling. B. Box plots of pLDDT scores for protein structures by PROGEN2-xlarge and our 7.2B CLM. C. Contour and line plots show our 7.2B CLM sequences mimic natural sequences more closely than PROGEN2-xlarge, assessed using Foldseek with the PDB database. D. Clustering at 50% sequence identity reveals our 7.2B CLM generates more clusters, indicating higher diversity.

5.1 Protein Generation Comparison: 7.2B CLM vs. 6.4B PROGEN-xlarge

We first evaluate the perplexity on OOD data and then compare the protein generation capabilities of the 7.2B CLM and PROGEN2-xlarge models. Each model generated 2,000 sequences for each parameter combination of topp { 0 . 5 , 0 . 7 , 0 . 9 , 1 . 0 } and temperature t { 0 . 2 , 0 . 4 , 0 . 6 , 0 . 8 , 1 . 0 } , totaling 40,000 sequences per model. Sequences with a perplexity greater than 10 and duplicates were removed, leaving 8,263 and 8,466 sequences for the 7.2B CLM and PROGEN-xlarge, respectively. We used four metrics to assess the quality of the models and the generated sequences (See Appendix F for details).

OODDataset PPL Analysis We randomly sampled 5,000 sequences from UniProt released after 2023-01-01 and aligned them to our and PROGEN2's training data (Uniref90 and BFD) using HHblits [66] or Jackhmmer [29]. Sequences below a maximum identity cutoff were used to assess the models' PPL, as shown in Figure 6A. Our 7.2B CLM exhibited lower PPL on three subsets.

pLDDT scores from ESMFold Atomic structures of 8,263 and 8,466 generated sequences were predicted using ESMFold, and compared based on pLDDT scores, displayed in Figure 6B. The 7.2B model's average pLDDT score was 78.69, higher than PROGEN2-xlarge's 74.33.

Natural Sequences Comparisons with Foldseek Using Foldseek [81], we searched the PDB database for sequences similar to those generated by our 7.2B CLM model, which showed better mimicry of natural sequence properties with higher average TM-scores (0.655 vs 0.522) and SeqID (0.194 vs 0.165), as shown in Figure 6C.

Diversity Analysis Generated sequences were clustered using MMseqs2 [72] with a 50% similarity cutoff. The 7.2B CLM model resulted in higher diversity with 7,097 clusters compared to 4,818 clusters for PROGEN2-xlarge, detailed in Figure 6D.

5.2 Protein understanding tasks: 10.7B MLM vs. 3B ESM2

We evaluate different task types from the protein benchmark [14]: Contact prediction as binary classification at the amino acid pair level; fold classification into 1195 classes at the sequence level; and fluorescence as regression tasks. Following [14], we add a Multi-Layer Perceptron (MLP) head to each pre-trained model and apply Low-Rank Adaptation (LoRA) [37] (r=8, α =16) for fine-tuning (see Appendix G for convergence details).

The results, shown in Table 6 and A7, 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. Additionally, the 470M model transferred from CLM pre-training continues to perform effectively in this task, showing the efficacy of the observed transfer scaling law.

6 Discussion and Limitations

Data Repeat Scaling Law. Our scaling law is derived within a single-epoch training setting. It is well known that MLM exhibits higher sample efficiency than CLM due to its dynamic masking strategies across multiple epochs. However, this advantage diminishes when training is restricted to just one epoch. We present an empirical study comparing a 2.8B model trained on 1T tokens (approximately five epochs) with a 10.7B model trained on 265B tokens (roughly 1.4 epochs). The two models achieve similar performance in terms of OOD PPL (10.33 vs. 10.21) while utilizing the same amount of FLOPs. This finding suggests that repeating multiple rounds of MLM training has minimal impact on reducing loss. Notably, the smaller models are more user-friendly during inference and fine-tuning. Therefore, we suggest an alternative approach that adjusts the optimal training token count and model size within the data-repeat scaling law.

Multi-modality Scaling. The multi-modal auto-regressive work [33] suggests the existence of a nearly universal scaling law across various modalities, including images, videos, math, code, and languages. Our results appear in this trend as well, such as, the scaling laws for CLM exhibit similarities to those in natural languages. The same situation may apply to other modalities of biological data, such as RNA and DNA [60].

Various Pre-train Datasets and Strategies. Our datasets cover a substantial portion of the protein universe, yet they might not be entirely representative. Combining BFD [8], Uniref [75], MetaClust [44], and IMG/JGI [52] with 90% clustering results in at least 600 billion unique tokens. However, variations in datasets could affect the power-law behavior. Future work could explore applying our findings to different model architectures. There is ongoing research on scaling LLMs for long sequences [7, 15, 16, 18, 39, 48, 87], and MSA augmentation could significantly improve protein representation regarding contacts and structure. Investigating scaling laws in this context could be a promising direction for future research.

7 Conclusion

In this work, we are the first 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, expanded metagenomic databases such as ColabFoldDB, emphasizing the critical importance of data quality and quantity for scaling language models, to optimal parameter and dataset allocation along with the potential loss prediction, 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.

Acknowledgments. This work has been supported by the National Key R&D Program of China 2021ZD0113304, NSFC for Distinguished Young Scholar 62425601, New Cornerstone Science Foundation through the XPLORER PRIZE.

References

Bfd - big fantastic database. https://bfd.mmseqs.com .

Armen Aghajanyan, Lili Yu, Alexis Conneau, Wei-Ning Hsu, Karen Hambardzumyan, Susan Zhang, Stephen Roller, Naman Goyal, Omer Levy, and Luke Zettlemoyer. Scaling laws for generative mixed-modal language models. In International Conference on Machine Learning , pages 265-279. PMLR, 2023.

Harriet Alexander, Sarah K Hu, Arianna I Krinos, Maria Pachiadaki, Benjamin J Tully, Christopher J Neely, and Taylor Reiter. Eukaryotic genomes from a global metagenomic dataset illuminate trophic modes and biogeography of ocean plankton. bioRxiv , pages 2021-07, 2021.

Rohan Anil, Andrew M Dai, Orhan Firat, Melvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, et al. Palm 2 technical report. arXiv preprint arXiv:2305.10403 , 2023.

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016.

Mohammad Bavarian, Heewoo Jun, Nikolas Tezak, John Schulman, Christine McLeavey, Jerry Tworek, and Mark Chen. Efficient training of language models to fill in the middle. arXiv preprint arXiv:2207.14255 , 2022.

Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150 , 2020.

BFD Team. Big fantastic database. BFD Official Website, n.d.

Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, et al. Deepseek llm: Scaling open-source language models with longtermism. arXiv preprint arXiv:2401.02954 , 2024.

Nadav Brandes, Dan Ofer, Yam Peleg, Nadav Rappoport, and Michal Linial. Proteinbert: a universal deep-learning model of protein sequence and function. Bioinformatics , 38(8):21022110, 2022.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020.

Luis F Camarillo-Guerrero, Alexandre Almeida, Guillermo Rangel-Pineros, Robert D Finn, and Trevor D Lawley. Massive expansion of human gut bacteriophage diversity. Cell , 184(4):10981109, 2021.

Bo Chen, Xingyi Cheng, Pan Li, Yangli-ao Geng, Jing Gong, Shen Li, Zhilei Bei, Xu Tan, Boyan Wang, Xin Zeng, et al. xtrimopglm: unified 100b-scale pre-trained transformer for deciphering the language of protein. arXiv preprint arXiv:2401.06199 , 2024.

Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509 , 2019.

Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009.14794 , 2020.

Aidan Clark, Diego de Las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, et al. Unified scaling laws for routed language models. In International conference on machine learning , pages 4057-4086. PMLR, 2022.

Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems , 35:16344-16359, 2022.

Justas Dauparas, Ivan Anishchenko, Nathaniel Bennett, Hua Bai, Robert J Ragotte, Lukas F Milles, Basile IM Wicky, Alexis Courbet, Rob J de Haas, Neville Bethel, et al. Robust deep learning-based protein sequence design using proteinmpnn. Science , 378(6615):49-56, 2022.

Mostafa Dehghani, Josip Djolonga, Basil Mustafa, Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. In International Conference on Machine Learning , pages 7480-7512. PMLR, 2023.

Grégoire Delétang, Anian Ruoss, Paul-Ambroise Duquenne, Elliot Catt, Tim Genewein, Christopher Mattern, Jordi Grau-Moya, Li Kevin Wenliang, Matthew Aitchison, Laurent Orseau, et al. Language modeling is compression. arXiv preprint arXiv:2309.10668 , 2023.

Tom O Delmont, Morgan Gaia, Damien D Hinsinger, Paul Frémont, Chiara Vanni, Antonio Fernandez-Guerra, A Murat Eren, Artem Kourlaiev, Leo d'Agata, Quentin Clayssen, et al. Functional repertoire convergence of distantly related eukaryotic plankton lineages abundant in the sunlit ocean. Cell Genomics , 2(5):100123, 2022.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, 2018.

Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with mixture-of-experts. In International Conference on Machine Learning , pages 5547-5569. PMLR, 2022.

Zhengxiao Du, Yujie Qian, Xiao Liu, Ming Ding, Jiezhong Qiu, Zhilin Yang, and Jie Tang. Glm: General language model pretraining with autoregressive blank infilling. arXiv preprint arXiv:2103.10360 , 2021.

Zhengxiao Du, Yujie Qian, Xiao Liu, Ming Ding, Jiezhong Qiu, Zhilin Yang, and Jie Tang. Glm: General language model pretraining with autoregressive blank infilling. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 320-335, 2022.

Ahmed Elnaggar, Hazem Essam, Wafaa Salah-Eldin, Walid Moustafa, Mohamed Elkerdawy, Charlotte Rochereau, and Burkhard Rost. Ankh: Optimized protein language model unlocks general-purpose modelling. arXiv preprint arXiv:2301.06568 , 2023.

Ahmed Elnaggar, Michael Heinzinger, Christian Dallago, Ghalia Rehawi, Yu Wang, Llion Jones, Tom Gibbs, Tamas Feher, Christoph Angerer, Martin Steinegger, et al. Prottrans: Toward understanding the language of life through self-supervised learning. IEEE transactions on pattern analysis and machine intelligence , 44(10):7112-7127, 2021.

European Bioinformatics Institute. Jackhmmer tool. EBI Tools Documentation, n.d.

fast.ai. How could the memorization hypothesis be true. fast.ai Blog, 2023. Retrieved May 21, 2024, from https://www.fast.ai/posts/2023-09-04-learning-jumps/#how-could-thememorization-hypothesis-be-true.

Noelia Ferruz, Steffen Schmidt, and Birte Höcker. Protgpt2 is a deep unsupervised language model for protein design. Nature communications , 13(1):4348, 2022.

Michael Heinzinger, Konstantin Weissenow, Joaquin Gomez Sanchez, Adrian Henkel, Martin Steinegger, and Burkhard Rost. Prostt5: Bilingual language model for protein sequence and structure. bioRxiv , pages 2023-07, 2023.

Tom Henighan, Jared Kaplan, Mor Katz, Mark Chen, Christopher Hesse, Jacob Jackson, Heewoo Jun, Tom B Brown, Prafulla Dhariwal, Scott Gray, et al. Scaling laws for autoregressive generative modeling. arXiv preprint arXiv:2010.14701 , 2020.

Danny Hernandez, Tom Brown, Tom Conerly, Nova DasSarma, Dawn Drain, Sheer El-Showk, Nelson Elhage, Zac Hatfield-Dodds, Tom Henighan, Tristan Hume, et al. Scaling laws and interpretability of learning from repeated data. arXiv preprint arXiv:2205.10487 , 2022.

Danny Hernandez, Jared Kaplan, Tom Henighan, and Sam McCandlish. Scaling laws for transfer. arXiv preprint arXiv:2102.01293 , 2021.

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556 , 2022.

Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685 , 2021.

Hugging Face. Llama 2 model documentation, n.d.

Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Leon Song, Samyam Rajbhandari, and Yuxiong He. Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models. arXiv preprint arXiv:2309.14509 , 2023.

Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401.04088 , 2024.

John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature , 596(7873):583-589, 2021.

Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 , 2020.

Aran Komatsuzaki. One epoch is all you need. arXiv preprint arXiv:1906.06669 , 2019.

Eli Levy Karin, Milot Mirdita, and Johannes Söding. Metaeuk-sensitive, high-throughput gene discovery, and annotation for large-scale eukaryotic metagenomics. Microbiome , 8:1-15, 2020.

Francesca-Zhoufan Li, Ava P Amini, Yisong Yue, Kevin K Yang, and Alex X Lu. Feature reuse and scaling: Understanding transfer learning with protein language models. bioRxiv , pages 2024-02, 2024.

Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, Allan dos Santos Costa, Maryam Fazel-Zarandi, Tom Sercu, Salvatore Candido, and Alexander Rives. Evolutionary-scale prediction of atomic level protein structure with a language model. bioRxiv , 2022.

Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, et al. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science , 379(6637):1123-1130, 2023.

Hao Liu, Matei Zaharia, and Pieter Abbeel. Ring attention with blockwise transformers for near-infinite context. arXiv preprint arXiv:2310.01889 , 2023.

Xiaoran Liu, Hang Yan, Shuo Zhang, Chenxin An, Xipeng Qiu, and Dahua Lin. Scaling laws of rope-based extrapolation. arXiv preprint arXiv:2310.05209 , 2023.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101 , 2017.

Ali Madani, Bryan McCann, Nikhil Naik, Nitish Shirish Keskar, Namrata Anand, Raphael R Eguchi, Po-Ssu Huang, and Richard Socher. Progen: Language modeling for protein generation. arXiv preprint arXiv:2004.03497 , 2020.

Victor M Markowitz, Frank Korzeniewski, Krishna Palaniappan, Ernest Szeto, Greg Werner, Anu Padki, Xueling Zhao, Inna Dubchak, Philip Hugenholtz, Iain Anderson, et al. The integrated microbial genomes (img) system. Nucleic acids research , 34(suppl_1):D344-D348, 2006.

Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAI Dota Team. An empirical model of large-batch training. arXiv preprint arXiv:1812.06162 , 2018.

Joshua Meier, Roshan Rao, Robert Verkuil, Jason Liu, Tom Sercu, and Alex Rives. Language models enable zero-shot prediction of the effects of mutations on protein function. Advances in neural information processing systems , 34:29287-29303, 2021.

William Merrill, Vivek Ramanujan, Yoav Goldberg, Roy Schwartz, and Noah Smith. Effects of parameter norm growth during transformer training: Inductive bias from gradient descent. arXiv preprint arXiv:2010.09697 , 2020.

Milot Mirdita, Konstantin Schütze, Yoshitaka Moriwaki, Lim Heo, Sergey Ovchinnikov, and Martin Steinegger. Colabfold: making protein folding accessible to all. Nature methods , 19(6):679-682, 2022.

Alex L Mitchell, Alexandre Almeida, Martin Beracochea, Miguel Boland, Josephine Burgin, Guy Cochrane, Michael R Crusoe, Varsha Kale, Simon C Potter, Lorna J Richardson, et al. Mgnify: the microbiome analysis resource in 2020. Nucleic acids research , 48(D1):D570-D578, 2020.

Niklas Muennighoff, Alexander Rush, Boaz Barak, Teven Le Scao, Nouamane Tazi, Aleksandra Piktus, Sampo Pyysalo, Thomas Wolf, and Colin A Raffel. Scaling data-constrained language models. Advances in Neural Information Processing Systems , 36, 2024.

Stephen Nayfach, David Páez-Espino, Lee Call, Soo Jen Low, Hila Sberro, Natalia N Ivanova, Amy D Proal, Michael A Fischbach, Ami S Bhatt, Philip Hugenholtz, et al. Metagenomic compendium of 189,680 dna viruses from the human gut microbiome. Nature microbiology , 6(7):960-970, 2021.

Eric Nguyen, Michael Poli, Matthew G Durrant, Armin W Thomas, Brian Kang, Jeremy Sullivan, Madelena Y Ng, Ashley Lewis, Aman Patel, Aaron Lou, et al. Sequence modeling and design from molecular to genome scale with evo. bioRxiv , pages 2024-02, 2024.

Erik Nijkamp, Jeffrey A Ruffolo, Eli N Weinstein, Nikhil Naik, and Ali Madani. Progen2: exploring the boundaries of protein language models. Cell systems , 14(11):968-978, 2023.

Pascal Notin, Aaron Kollasch, Daniel Ritter, Lood Van Niekerk, Steffanie Paul, Han Spinner, Nathan Rollins, Ada Shaw, Rose Orenbuch, Ruben Weitzman, et al. Proteingym: largescale benchmarks for protein fitness prediction and design. Advances in Neural Information Processing Systems , 36, 2024.

Jiezhong Qiu, Junde Xu, Jie Hu, Hanqun Cao, Liya Hou, Zijun Gao, Xinyi Zhou, Anni Li, Xiujuan Li, Bin Cui, et al. Instructplm: Aligning protein language models to follow protein structure instructions. bioRxiv , pages 2024-04, 2024.

Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11446 , 2021.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research , 21(140):1-67, 2020.

Michael Remmert, Andreas Biegert, Andreas Hauser, and Johannes Söding. Hhblits: lightningfast iterative protein sequence searching by hmm-hmm alignment. Nature methods , 9(2):173175, 2012.

Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. Advances in Neural Information Processing Systems , 34:8583-8595, 2021.

Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Jason Liu, Demi Guo, Myle Ott, C Lawrence Zitnick, Jerry Ma, et al. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. Proceedings of the National Academy of Sciences , 118(15):e2016239118, 2021.

Philippe A Robert, Rahmad Akbar, Robert Frank, Milena Pavlovi´ c, Michael Widrich, Igor Snapkov, Andrei Slabodkin, Maria Chernigovskaya, Lonneke Scheffer, Eva Smorodina, et al. Unconstrained generation of synthetic antibody-antigen structures to guide machine learning methodology for antibody specificity prediction. Nature Computational Science , 2(12):845-865, 2022.

Elke Schaper, Andrey V Kajava, Alain Hauser, and Maria Anisimova. Repeat or not repeat?-statistical validation of tandem repeat prediction in genomic sequences. Nucleic acids research , 40(20):10005-10017, 2012.

Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202 , 2020.

Martin Steinegger and Johannes Söding. Mmseqs2 enables sensitive protein sequence searching for the analysis of massive data sets. Nature biotechnology , 35(11):1026-1028, 2017.

Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing , 568:127063, 2024.

Ning Sun, Shuxian Zou, Tianhua Tao, Sazan Mahbub, Dian Li, Yonghao Zhuang, Hongyi Wang, Xingyi Cheng, Le Song, and Eric P Xing. Mixture of experts enable efficient and effective protein understanding and design. bioRxiv , pages 2024-11, 2024.

Baris E Suzek, Yuqi Wang, Hongzhan Huang, Peter B McGarvey, Cathy H Wu, and UniProt Consortium. Uniref clusters: a comprehensive and scalable alternative for improving sequence similarity searches. Bioinformatics , 31(6):926-932, 2015.

Yi Tay, Mostafa Dehghani, Jinfeng Rao, William Fedus, Samira Abnar, Hyung Won Chung, Sharan Narang, Dani Yogatama, Ashish Vaswani, and Donald Metzler. Scale efficiently: Insights from pre-training and fine-tuning transformers. arXiv preprint arXiv:2109.10686 , 2021.

Yi Tay, Mostafa Dehghani, Vinh Q Tran, Xavier Garcia, Jason Wei, Xuezhi Wang, Hyung Won Chung, Siamak Shakeri, Dara Bahri, Tal Schuster, et al. Ul2: Unifying language learning paradigms. arXiv preprint arXiv:2205.05131 , 2022.

Yi Tay, Jason Wei, Hyung Won Chung, Vinh Q Tran, David R So, Siamak Shakeri, Xavier Garcia, Huaixiu Steven Zheng, Jinfeng Rao, Aakanksha Chowdhery, et al. Transcending scaling laws with 0.1% extra compute. arXiv preprint arXiv:2210.11399 , 2022.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 , 2023.

Michel van Kempen, Stephanie S Kim, Charlotte Tumescheit, Milot Mirdita, Cameron LM Gilchrist, Johannes Söding, and Martin Steinegger. Foldseek: fast and accurate protein structure search. Biorxiv , pages 2022-02, 2022.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Robert Verkuil, Ori Kabeli, Yilun Du, Basile IM Wicky, Lukas F Milles, Justas Dauparas, David Baker, Sergey Ovchinnikov, Tom Sercu, and Alexander Rives. Language models generalize beyond natural proteins. bioRxiv , pages 2022-12, 2022.

Alex Wang and Kyunghyun Cho. Bert has a mouth, and it must speak: Bert as a markov random field language model. arXiv preprint arXiv:1902.04094 , 2019.

Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang, Dongdong Zhang, and Furu Wei. Deepnet: Scaling transformers to 1,000 layers. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2024.

Thomas Wang, Adam Roberts, Daniel Hesslow, Teven Le Scao, Hyung Won Chung, Iz Beltagy, Julien Launay, and Colin Raffel. What language model architecture and pretraining objective works best for zero-shot generalization? In International Conference on Machine Learning , pages 22964-22984. PMLR, 2022.

Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems , 33:17283-17297, 2020.

Aohan Zeng, Xiao Liu, Zhengxiao Du, Zihan Wang, Hanyu Lai, Ming Ding, Zhuoyi Yang, Yifan Xu, Wendi Zheng, Xiao Xia, et al. Glm-130b: An open bilingual pre-trained model. arXiv preprint arXiv:2210.02414 , 2022.

Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 12104-12113, 2022.

Biao Zhang, Zhongtao Liu, Colin Cherry, and Orhan Firat. When scaling meets llm finetuning: The effect of data, model and finetuning method. arXiv preprint arXiv:2402.17193 , 2024.

Zaixiang Zheng, Yifan Deng, Dongyu Xue, Yi Zhou, Fei Ye, and Quanquan Gu. Structureinformed language models are protein designers. In International Conference on Machine Learning , pages 42317-42338. PMLR, 2023.

Appendix

Table of Contents

A Related Work

Protein Language Model Since the advent of AlphaFold2 [41], the masked language model (MLM) has been integrated as a subtask within the Evoformer architecture. In this context, an assumption is that large language models can be considered as a lossless compression method [21]. This was followed by a series of language modeling efforts [31, 10, 32, 28, 27], which aimed to conduct pre-training on single-sequence proteins using larger datasets and model scales. These efforts sought to harness the scale of the models to learn complex co-evolutionary information, although detailed investigations on how to optimally scale these models remain scarce. Our work primarily focuses on these finer aspects, aiming to fill this gap in the research.

Training objectives In natural language processing (NLP), masked language models (MLM) are rarely adopted due to the self-explanatory nature of natural language, which inherently prompts the meta-knowledge of tasks and generates task targets through CLM (Conditional Language Modeling) training models. However, a unified language modeling objective for Protein Language Models has yet to be fully consented. Those based on causal language modeling (CLM) have been primarily explored for protein design. Benchmarks in protein design using MLM [84] have also shown promising results for generation [62], exhibiting variable performance when compared to CLM [91, 83]. Additionally, the potential of the in-filling task objective remains largely unexplored [6, 77, 25]. Our research aims to thoroughly discern the scaling behavior of the two most common optimization objectives in this domain.

Scaling Laws To our knowledge, the concept of scaling laws of language model is first introduced by OpenAI [42]. Subsequently, numerous variants and modifications [36] have been developed around this theme. Recently, an array of new scaling laws has emerged. These include scaling laws related to learning rates and batch sizes [9], data-constrained scaling laws [58], scaling laws for downstream tasks and Transfer [90, 35], as well as scaling laws within the Mixture of Experts (MoE) framework [17], and those concerning long sequences and positional encoding [49]. While these laws are primarily derived using auto-regressive models in resource-rich domains, their application in the biological data sector is less common. Our work seeks to address this gap. Furthermore, scaling laws for Masked Language Models (MLM) are notably scarce. Given that MLMs are currently one of the most effective training methods for biological data, our research on MLMs could also be extended to other non-text domains.

B UR50/S Repeat Experiments

0K

50K

100K

150K

200K

Step

1.6

1.8

2.0

2.2

2.4

2.6

Loss

Training loss curve

UniMeta200B

UR50/S bootstrap

UR50/S every_epoch_shuffle

UR50/S global_shuffle

0K

50K

100K

150K

200K

Step

1.8

1.9

2.0

2.1

2.2

2.3

2.4

2.5

2.6

Loss

Validation loss curve

UniMeta200B

UR50/S bootstrap

UR50/S every_epoch_shuffle

UR50/S global_shuffle

0K

50K

100K

150K

200K

Step

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1.0

Loss

Grad norm curve

UniMeta200B

UR50/S bootstrap

UR50/S every_epoch_shuffle

UR50/S global_shuffle

We employed three different methods to repeat training on the UR50/S dataset, all of which ultimately led to overfitting. The reference for these experiments is shown by the blue curve in Figure A7, which represents UniMeta's loss for approximately one epoch.

Firstly, using bootstrapping, we processed 200 billion tokens from UR50/S with replacement. In each epoch, 65% of the dataset was randomly selected, leading to a diminished proportion of unsampled tokens by the fifth epoch, as depicted by the orange curve.

Secondly, we shuffled the unique data for each epoch to ensure that all UR50/S tokens were used per epoch, resulting in a stair-step pattern [30] in the training loss, illustrated by the green curve. It has simply memorized the dataset but isn't improving at generalizing. Over-confident predictions of the first batch of the next epoch lead to a big step update, and then the model is not adapted to the next batches, resulting in no longer a decrease in loss.

Lastly, we shuffled the entire training dataset less stringently, which did not strictly ensure that all UR50/S tokens were used every epoch, but guaranteed that each token was used an equal number of times over the entire training period. We term it global shuffle, this approach is shown by the red curve.

From the gradient norm curve shown in Figure A7 (right), we observe an uptick in gradient norm for the overfitting curves, indicating that the model is no longer optimizing effectively. In machine learning, such an increase in gradient norm typically suggests that the model is encountering areas of the parameter space where gradients are steeper or more erratic, often occurring when the model starts to memorize the training data rather than generalize from it, approaching a saturated network [55]. This behavior can result from overly complex models, too many training epochs without sufficient regularization, or training on non-representative data.

C Choice of Masking Ratio

model:154M

model:85M

Validationloss

2.6

mask10

SSOl

2.6-

mask10

2.5

mask15

mask15

mask20

Validation

2.5

mask20

2.4

mask25

mask25

mask30

2.4

mask30

2.3

mask40

mask40

2.2

mask60

2.3

mask60

0

10

20

30

40

50

0

5

10

15

20

25

30

Training tokens (billion)

Training tokens (billion)

0.05 0.1 0.15 0.2 0.25 0.3

0.4

0.6

Mask prob.

0.0

0.2

0.4

0.6

Performance

Contact Prediction (154M)

0.05 0.1 0.15 0.2 0.25 0.3

0.4

0.6

Mask prob.

0.0

0.2

0.4

0.6

Performance

Fold Prediction (154M)

0.05 0.1 0.15 0.2 0.25 0.3

0.4

0.6

Mask prob.

0.0

0.1

0.2

0.3

0.4

Performance

Contact Prediction (85M)

0.05 0.1 0.15 0.2 0.25 0.3

0.4

0.6

Mask prob.

0.0

0.2

0.4

0.6

Performance

Fold Prediction (85M)

In the original BERT work [23], the absence of masked tokens in downstream tasks presented a mismatch with the pre-training data distribution. The authors investigated various masking ratios and concluded that a 15% masking rate was most beneficial for downstream tasks. This was implemented alongside an 80-10-10 strategy: 80% of the tokens were replaced with a mask, 10% were randomly substituted, and the remaining 10% were left unchanged.

However, given the significant differences between protein sequences and natural language processing data, we employed two models, sized at 85M and 154M, to explore a range of masking ratios from

5% to 60% (see Figure A8). The best masking ratios for validation loss drop ranged from 10% to 20%; ratios too small (5%) or too large (greater than 25%) degraded the performance.

We further used pre-trained eight different models to perform full fine-tuning on downstream tasks such as Contact Prediction and Fold Classification in Figure A9. Results from the test datasets revealed that, similar to NLP, the optimal performance was achieved within a 10%-20% masking range. Specifically, a 20% masking ratio slightly outperformed 15% in Contact Prediction, while the 15% ratio yielded the best results in Fold Prediction. Consequently, for our Masked Language Model (MLM), we decided to adhere to the 15% masking ratio with the 80-10-10 strategy for training all our models.

D MLM/CLM for Protein Contact Prediction

0

10000

20000

30000

40000

50000

Step

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

P@L/5

3B MLM LoRA

3B CLM LoRA

3B MLM Probing

3B CLM Probing

0

10000

20000

30000

40000

50000

Step

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

P@L/5

880M MLM LoRA

7.2B CLM LoRA

880M MLM Probing

7.2B CLM Probing

We compared the effectiveness of CLM in the downstream task of contact prediction, using two different setups (Figure A10). In the first setup, two 3B models were trained under identical computational resources on 200 billion tokens, 3 . 4 × 10 21 FLOPs . Their performance was evaluated through two training approaches: Probing (freezing the pre-trained model) and LoRA fine-tuning, with an added MLP head for comparison.

In the second setup, we compared the effects of MLM and CLM under similar loss conditions. Here, a 7.2B CLM model and an 880M MLM model were selected, both achieving a loss of 1.98 on our validation set. Despite the MLM model having a simpler loss calculation, involving a 15% mask rather than a one-by-one mask-which would result in a higher loss-the MLM significantly outperformed the CLM. Importantly, the CLM model's computational power was an order of magnitude greater than the MLM model ( 1 . 68 × 10 22 vs 1 . 0 × 10 21 FLOPs). This suggests that despite the lower loss achievable by the CLM model compared to MLM with a one-by-one mask, the unidirectional limitations of CLM do not translate into better downstream task performance.

E Pre-training Dataset Quality

Compared to Uniref90, ColabFoldDB offers a higher diversity and larger numbers of protein sequences, though with generally shorter sequence lengths, likely suggesting potentially lower data quality. To evaluate the efficacy of our expanded dataset, ColabFoldDB, we initially trained two 85M models separately on Uniref90 and ColabFoldDB. Uniref90 in our dataset comprises two subsets: Uniref50/S and the incremental dataset over Uniref50/S, termed Uniref90/50. Similarly, ColabFoldDB consists of representative and member data. We controlled the sampling proportion to ensure uniform sampling across both datasets, with results reported in Table A8. Both models were then trained using identical configurations on a 50B scale.

From the perspective of validation loss in pre-training, the higher loss on ColabFoldDB might be attributed to its lower diversity and shorter sequence lengths compared to Uniref90. However, the performance on downstream tasks, such as contact prediction and fold classification, shows negligible differences between models trained solely on ColabFoldDB and those trained on Uniref90, as illustrated in Figure E. This confirms that ColabFoldDB is an effective expansion of Uniref90 that maintains sample efficiency.

0

20000 40000 60000 80000100000

Step

2.0

2.1

2.2

2.3

2.4

2.5

2.6

2.7

Validation loss

Learning curve

85M_Uniref90

0

1000

2000

3000

4000

5000

6000

Step

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

P@L/5

Contact Prediction

85M_Uniref90

0

1000

2000

3000

4000

Step

0.0

0.1

0.2

0.3

0.4

Fold Classification

85M_Uniref90

F The Evaluation of Protein Generation.

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.

G Convergence Analysis of Downstream Fine-tuning Tasks

Observing the learning curves in Figure A12a, we can assess the effectiveness of different finetuning scenarios. For the contact prediction task, the convergence speed under the LoRA setting is very similar for both models. Our testing reveals closely matching results for ESM-2 models with capacities of 650M, 3B, 15B, consistent with the findings reported by Ankh et al. [27]. This similarity suggests possible saturation of the dataset under single-sequence pre-trained models. Additionally, the convergence rates for tasks such as fold classification and fluorescence are generally faster than those for ESM-2, indicating robust generalization following our data augmentation strategies.

Based on the two 470M models defined in our Table 5, despite using the same computational power, we observe distinct outcomes (Figure A12b) in contact prediction and fold classification tasks. The MLMmodel from CLM pre-training converges slightly faster than MLM from scratch. However, the distinction is less pronounced in the two downstream regression tasks. This suggests that perplexity is more sensitive to protein structure related tasks, i.e., contact prediction and fold classification, but shows less sensitivity to regression tasks, particularly when assessed using the Spearman metric, which is prone to variability.

0

2000

4000

6000

8000

Step

0.0

0.2

0.4

0.6

0.8

P@L/5

Contact Prediction

our 10.7b

ESM-2 (3B)

0

20000

40000

60000

80000

Step

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Fold Classification

our 10.7b

ESM-2 (3B)

0

20000

40000

60000

Step

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Spearman coefficient

Fluoresence Prediction

our 10.7b

ESM-2 (3B)

0

5000 1000015000200002500030000

Step

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

P@L/5

Contact Prediction

470M from scratch

470M from CLM pre-training

0

1000

2000

3000

4000

5000

Step

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Fold Classification

470M from scratch

470M from CLM pre-training

0

2500 5000 7500 100001250015000

Step

0.0

0.1

0.2

0.3

0.4

0.5

0.6

Spearman coefficient

Fluoresence Prediction

470M from scratch

470M from CLM pre-training

H Mixed Objectives Training

We also employed an untied model to simultaneously optimize two objectives:

and

where V represents the protein vocabulary embedding, and W 1 and W 2 are the parameters corresponding to the CLM and MLM tasks, respectively. CE is the cross-entropy operator. The σ is the Tanh activation function.

We compared CLM and MLM under our scaling law of optimal model and data size distributions. One approach involved training from scratch, while the other used mixed training. In the mixed training approach, the actual number of training tokens was higher due to the additional FLOPs consumed by another optimally trained objective, in other words. In other words, mixed training consumes the FLOPs of two optimal allocations; we only extracted the loss curve of one target for comparison. We extracted the loss curve of just one target for comparison with the from-scratch training. Our findings indicate that mixed training of the two targets can lead to detrimental interference, an effect not observable in smaller models, as depicted in Figure A13. As the model size increases to a hundred million or billion parameters, the differences become more pronounced. The possible reason for this situation is that mixed training has reduced the batch size for one of the objectives, making optimization difficult. We did not further investigate the impact of increasing the batch size and only observed based on the training tokens. However, we cannot rule out the possibility that they are mutually detrimental. Therefore, if both objectives are to be optimized concurrently, a sequential training strategy should be employed: first optimizing CLM, followed by MLM training. We consider that CLM is more challenging to predict than MLM, which may allow the model to capture more complex and implicit sequential features initially, thereby enhancing its ability to understand and predict masked words in subsequent MLM training.

10 8

Model size

2.40

2.45

2.50

2.55

2.60

CLM Validation Loss

3e+18 FLOPs

1e+19 FLOPs

3e+19 FLOPs

1e+20 FLOPs

10 17

10 19

10 21

10 23

FLOPs

10 6

10 7

10 8

10 9

10 10

10 11

10 12

Model size

N CLM

opt

=1.48

10

3

C

0.573

10 17

10 19

10 21

10 23

FLOPs

10 9

10 10

10 11

10 12

10 13

Data size

D CLM

opt

=1.57

10 2

C

0.424

10 8

Model size

2.2

2.3

2.4

2.5

MLM Validation Loss

3e+18 FLOPs

1e+19 FLOPs

3e+19 FLOPs

1e+20 FLOPs

10 17

10 19

10 21

10 23

FLOPs

10 6

10 7

10 8

10 9

10 10

10 11

10 12

Model size

N MLM

opt

=6.36

10

7

C

0.738

10 17

10 19

10 21

10 23

FLOPs

10 9

10 10

10 11

10 12

10 13

Data size

D MLM

opt

=2.04

10 5

C

0.261

I MoE Scaling

We find that the scaling behaviors of sparse parameter counts in Mixture of Expert (MoE) models are remarkably similar to those of dense model sizes, potentially allowing for a reduced compute budget for modeling scaling behaviors due to less activated parameters per token.

In our experiments, we evaluate MoE models ranging from 10M to 500M sparse parameter counts, using a model size of 17 with eight experts, following the settings outlined in Mixtral of experts [40], including its load-balancing scheme. The figure below shows different IsoFLOPs curves. Notably, the FLOPs here are calculated based on sparse parameters rather than actually activated ones. We use the method described in the main text to select optimal loss points and fit these around the sample points, enabling us to project the optimal model size and number of tokens for larger models (center and right). We observe that the power-law coefficients for CLM and MLM are similar to those of dense models, with MoE CLM vs. Dense CLM at approximately 0.57 vs. 0.58, and MoE MLM vs. Dense MLM at 0.74 vs. 0.77.

Our study strictly focuses on models with eight experts, which may not be entirely rigorous. Clark et al. [17] proposed a unified scaling law defining effective training parameters for MoE, aiming to harmonize the scaling laws for Dense and MoE models. Investigation of biological data will be considered as future work.

J Combined Power-law

We applied the fitting function proposed by Chinchilla [36], detailed in Equation 8, to model the effects of various factors on model performance. It can provide a loss prediction where neither the parameters or model size are not optimal allocation. This loss function simultaneously depends on parameters N and D :

where E denotes the irreducible loss. Parameters A , B , α , and β are learned through the fitting process. As N →∞ or D →∞ , the function degenerates to a form similar to Equation 2, which indicates that it models the scenarios under perfect conditions of other variables.

Given that most of our training tokens are used for less than or equal to one epoch, and that the model size is prone to underfitting at fixed FLOPs, the asymptotic behaviors L ( N ) at D →∞ and L ( D ) at N →∞ are enough for determining the parameters in L ( N,D ) .

To enrich data points, we randomly added several FLOP counts into 25% of the model size and trained these models for 0.25, 0.5, 0.75, and 1 epoch. And we adopt the Huber loss to fit these coefficients:

where LSE represents the log-sum-exp operator, and δ = 10 -3 . The terms N i , D i , and L i denote the model size, dataset size, and loss of the i -th run, respectively. We fitted the MLM validation loss from 110 samples and the CLM validation loss from 149 samples using grid search with α ∈ { 0 , 0 . 5 , . . . , 2 } , β ∈ { 0 , 0 . 5 , . . . , 2 } , e ∈ {-1 , -0 . 5 , . . . , 1 } , a ∈ { 0 , 5 , . . . , 25 } , and b ∈ { 0 , 5 , . . . , 25 } . The final initialized parameters of CLM and MLM both are [ e, a, b, α, β ] = [1 , 5 , 10 , 0 . 5 , 0 . 5] . We set the maximum number of iterations to 1000, and the two objectives were essentially achieved after 360 iterations. The exponential powers of learned a and b yielded the coefficients A , B , which were reported in Table A9.

Substituting all learned coefficients into the following Equation from the original Chinchilla paper:

The results closely approximate the trends given in Equations 1 and 2, confirming our overall findings.

K IsoLoss

In addition to using the seven different FLOPs counts reported in the main text to determine the optimal model sizes and fit our scaling law, we also incorporated additional model points into our analysis. We trained using the final loss points of all the CLM and MLM that are run. Figure A15 depicts the contour of the fitted function L and the efficient frontier as a red dashed line, presented in log-log space. The frontier interval of Figure 2 is computed from this observation. From this approach, it revealed the scaling exponents for model size to be 0.77 in MLM and 0.57 in CLM, very similar to the IsoFLOPs profiling method in Section 3.1.

L Training Procedure

We conducted all experiments using Ampere A100 GPUs (80G) equipped with NVLink, utilizing the GLM framework [88, 26] developed based on DeepSpeed and Megatron. We have used a total of around 1 million GPU hours. Our approach predominantly utilized data parallelism, avoiding model parallelism and pipeline parallelism to simplify deployment. Modifications were made to the standard Transformer architecture [82], adopting a DeepNorm [85] strategy and layer normalization [5]. The activation function was set to GLU [71], RoPE [73] was used to encode position, similar to the settings found in the Transformer++ architecture [79]. We further adopt FlashAttention [18] to

10 19

10 20

10 21

FLOPS

10 8

10 9

IsoLoss for CLM

Efficient frontier

N

C

0.56

10 19

10 20

10 21

FLOPS

10 8

10 9

IsoLoss for MLM

Efficient frontier

N

C

0.76

accelerate our training process. The used max LR empirically found to range between 6 × 10 -4 and 1 . 2 × 10 -4 from small to large model size, was used along with a cosine decay strategy to reduce it to 0 . 1 × max LR. Both CLM and MLM were trained under similar settings for model size, with a consistent LR and a minimum warm-up period of 2.5% steps, extending to at least 100K training steps. All sequences were set to a length of 1024, with sequences concatenated using an <EOS> delimiter. Based on findings related to loss magnitude and batch size [53]. The AdamW optimizer [50] was used with β 1 = 0 . 9 , β 2 = 0 . 95 , ϵ = 1 × 10 -8 , and a weight decay of 0 . 01 . All experiments omitted the dropout (it reduced the capacity to hinder model scaling) and trained with bfloat16. Most pre-training experiments were confined to the ≤ 1 epoch, with some models extending up to 30% beyond one epoch. For the transfer learning setting, we load the finished checkpoint of the pre-training model and disregard the pre-trained optimized state, and learn rest tokens with warmup 5% steps the max LR.

M Broader Impact

If the scaling law of the protein language model improves predictions or understanding of protein structure and function, it could potentially have positive impacts on scientific research in fields such as biology, medicine, and drug development. This may facilitate the development of new drugs, accelerate progress in disease diagnosis, or drive advancements in frontier research in the life sciences.

N Model Parameters

Table A10 details the sizes and configurations of all models utilized in this research, training only with data parallel expcept 10B with tensor parallel size 2:

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We clearly made statements to claim the contributions in the abstract and introduction sections.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discussed several limitations of our work in Section 6.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: This work is primarily focused on empirical experiments and extensive validation rather than providing theoretical results. Therefore, the paper does not include a set of assumptions or proofs for theoretical results since it is not the main emphasis of the research.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide the implementation details in the main text and Appendix to ensure the reproduction of all the experimental results.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We provide the implementation details in the Appendix to ensure the reproduction of all the experimental results. Moreover, we release code and script at https://github.com/cxysteven/ScalingProteinLM .

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide the implementation details in the Appendix L to ensure the reproduction of all the experimental results. Moreover, we release code and script at https://github.com/cxysteven/ScalingProteinLM .

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We have presented our results with the error bar in Figure 2.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide the implementation details in the Appendix L to ensure the reproduction of all the experimental results.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have reviewed and obeyed the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discussed the broader impact in Appendix M.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We have mentioned all the source of used data, code, and models.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets

has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We release code and script at https://github.com/cxysteven/ ScalingProteinLM .

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: We don't include human subjects or crowdsourcing in this work.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: We don't include human subjects or crowdsourcing in this work.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe main objective of the paper is to explore the scalability of Protein Language Models (PLMs) and to provide models to determine the optimal number of parameters, pre-training dataset size, and transferability between pre-training tasks as functions of the available computational budget.\n\nThe study contains a comprehensive number of experiments and its main findings are:\n\n* The diversity of sequences of the pre-training database can lead to lower perplexity in out-of-distribution validation dataset. This is reflected by the overfitting observed when training models with masked language modelling (MLM) pre-training objective on the Uniref 50S database and a custom made pre-training database enriched with more diverse metagenomic sequences from a variety of sources.\n\n* It establishes scaling laws for that allow for determining the optimal model size and optimal pre-training data size  given computational budgets in the interval 10e17-10e21 both for the MLM pre-training objective and for the causal language modelling (CLM) pre-training objective.\n\n* It also established what the optimal proportion of computational budget should be allocated to pre-training two models one on CLM and another on MLM finding that in general the computational resources allocated to MLM should exceed those for CLM the training tokens reach the 10B threshold where a 1:1 ratio is reached, after which more training tokens should be allocated to CLM.\n\n* It also explores the benefits of transferring models pre-trained with one objective to another objective (MLM to CLM) or (CLM to MLM). Its findings are that the benefits of transferring from MLM to CLM diminish with model size in the 2.7M to 1.7B parameters range; and that the benefits from CLM to MLM grow with model size in the 85M to 1.2B parameters range.\n\n* They validate the scaling laws by pre-training and evaluating models for CLM and MLM tasks with the optimal sizes and comparing them with state-of-the-art methods: ProGen-2 and ESM-2 3B. In the case of CLM, their proposed method achieves greater sequence diversity between the sequences generated, but less structural diversity. In the case of MLM, their proposed method (10.7B) outperforms ESM2 3B by a small margin in two out of the three tasks (fold classification, improvement of 0.03; fluorescence prediction, improvement of 0.04).\n\n* Regarding transfer from CLM to MLM, they also show that a 470M parameter model pre-trained with the CLM to MLM strategy outperforms a model pre-trained in MLM from scratch in two of the benchmarks by tiny margins (contact, improvement of 0.02; fold classification, improvement of 0.01).\n\nSTRENGTHS:\n* The paper is clearly written and despite the diversity of analysis conducted the main findings and their importance are well conveyed in most cases.\n\n* The paper provides the most comprehensive analysis to date on protein language model scaling and how to optimise the computational resources available.\n\n* The paper demonstrates that the scaling laws induced from their experiments can be used to derive the most optimal pre-training configuration for PLMs both for the generative and predictive purposes.\n\nWEAKNESSES:\n* Figure 2 has quite an extended x axis from 10e17 to 10e24. The interpretation a reader may make of this (and it is the one I'm making) is that the authors are implicitly claiming (or at least suggesting) that the scaling laws are able to extrapolate to that range. However, their data only supports the claim that they are valid in the 10e17 to 10e21. In order, to make the claim (or suggestion) that they extrapolate to the range they show in the figure, there should be at least one data point close to the upper bound of the range (10e24).\n\n* Sections 4.1 and 4.2 were the most hard to follow, in contrast with the rest of the paper that was quite clear despite the complexity and diversity of the analysis. I'm not entirely sure how it could be improved, but I had to read it attentively 3 or 4 times to get the point of the experiment and what the results were. However, it may be a personal issue.\n\n* I think that section 5.1 could benefit of a more detailed discussion on what the results mean regarding protein sequence generation and how the optimal model proposed is better/worse than the state-of-the-art. I am aware of the the constraints on length, so I would suggest including it as an additional section on the appendix. I don't think this is a significant issue, just a minor improvement to make the contributions to the overall field clearer.\n\nQUESTIONS:\n- Do you have any intuition about why the CLM transferability to MLM scales with model size and the inverse is not true?\n\nLIMITATIONS:\nI think the authors have overall addressed properly the limitations of their setup and I don't think there is any significant societal impact on their work that has not being acknowledged. The only limitation that I don't think is properly address is that the results are limited to the range of model sizes and compute power they have experimented on and that the experimental setup does not allow for any claim about whether the laws induced can extrapolate outside that range. I don't think the authors have properly addressed this limitation, as I stated in the weaknesses the layout of Figure 2, seems to argue the opposite.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-06-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors fit scaling laws to determine how to train protein language models compute-optimally, for both causal LM and masked LM tasks. They also consider a setup to consider possible effective data transfer when training on one task and transfer training on the other. Finally, they use their derived scaling laws to train new protein language models compute-optimally, achieving better downstream performance.\n\nSTRENGTHS:\nI think the biggest strength of this paper is its importance – I think that the impact of these kinds of systems will likely be very substantial over the coming years and empirical scaling laws for protein LMs could be very influential. It’s also fairly original because there hasn’t been much work done on this in the past. So I really appreciated that the authors took the time to obtain these results on compute-optimal training.\n\nWEAKNESSES:\nPersonally, I feel the presentation of the writeup could be substantially improved. For example: The main body of the paper doesn’t have a discussion/conclusion, which unfortunately made the paper hard to read. In general, I found the writeup to be largely listing results without much discussion of significance. \n- Personally I think it’d be helpful if the authors made some of their discussion more concise (e.g. in section 2.1), and instead adding more detail about how significant their results are. \n- E.g. this could mean calculations saying something like “using our approach, we were able to reach X performance using a factor of Y less training compute, which suggests…”. \n- In addition, I’d suggest including a discussion section in the main body of the paper that outlines the key results and implications, especially in the context of previous work (in a quantitative fashion). What does their work mean that future AI researchers should do? How far off of compute-optimality was previous work? How important are overtraining considerations? Etc.\n\nI think if these kinds of concerns are addressed I'd be happy to increase my score. \n\nMore minor questions/concerns: \n- Lines 164-166: I feel it’d be best to include more decimal places in the reported numbers, so that the product of the data and model size increase equals 10x.\n- What are the confidence intervals for the scaling exponents for model size and data?\n- Why was this particular functional form for the effectively transferred tokens chosen? (equation 7) Why does it make sense for the data and parameter components to be multiplicative rather than additive?\n\nQUESTIONS:\n- Would it be possible to add a lot more detail on the implications of the work, as well as context on how it improves upon prior attempts? In particular, to what extent have previous models been suboptimally trained, and what are the core recommendations the authors have for future related work?\n\n[copied from the weaknesses section] \n- What are the confidence intervals for the scaling exponents for model size and data?\n- Why was this particular functional form for the effectively transferred tokens chosen? (equation 7) Why does it make sense for the data and parameter components to be multiplicative rather than additive?\n\nLIMITATIONS:\nI thought the discussion of limitations in appendix C was well done.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this work authors propose to explore training compute-optimal protein language models and aim to derive scaling laws. Specifically, they (i) explore new pre-training data for protein language models (ii) derive scaling laws for protein language models using these datasets for both Masked Language Modeling and Causal Language Modeling objectives (iii) explore scaling laws and transferability between those objectives (iv) use the findings to pre-train two pLMs which reach SOTA performance on downstream tasks.\n\nSTRENGTHS:\n**Clarity**\n- The paper is clearly motivated and finding domain-specific scaling laws is valuable for practitioners.\n- The paper is very easy to follow and experiments are well described.\n\n**Experiments**\n- Authors conduct extensive experiments and the results on scaling laws are solid.\n- Notably they also derive scaling laws for transfer between CLM and MLM and mixture-of-experts.\n\nWEAKNESSES:\n**Lack of discussion of results** \n- Conclusion, Related Works and Discussions and Limitations have been moved to the Appendix. While I understand the 9-page constraints, I recommend to include at least a discussion/conclusion part in the main body.  \n\n**Novelty**\n- This work closely follows previous work on English LMs such as Chinchilla and presents very limited conceptual novelty.\n\n**Experimental Validation**\n- While the scaling laws analysis is solid, the experimental validation on Protein understanding tasks is limited to 3 tasks whereas previous works have proposed a more extensive benchmark [1]\n\n[1] Elnaggar, Ahmed, et al. \"Ankh: Optimized protein language model unlocks general-purpose modelling.\" arXiv preprint arXiv:2301.06568 (2023).\n\nQUESTIONS:\n- In Section 5, did you try training the 470M version of your model for the same computing budget as ESM2 ? This would emulate works in NLP [2] where authors find that training smaller models for longer can yield good results with limited inference cost.\n- It would be interesting to see if the derived scaling laws are also reflected in downstream task performance, *e.g.* by fine-tuning models pre-trained for the same computing cost and training data (as I understand, PROGEN2 and ESM2 use different pre-training data).\n\n[2] Touvron, Hugo, et al. \"Llama: Open and efficient foundation language models.\" arXiv preprint arXiv:2302.13971 (2023).\n\nLIMITATIONS:\nI believe limitations have been properly addressed in Appendix.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies scaling law for two types of protein language models: masked language modelling and causal modelling. The paper finds that both types scale sublinearly but the coefficients of the scaling law are very different. Based on the derived scaling law, the paper then reevaluates the allocation of resources for ESM and PROGEN and shows improved results.\n\nSTRENGTHS:\nOverall, this paper presents a well-motivated study with solid and extensive experimental work, \n* The paper addresses scaling laws in the context of AI4Science, an area that has been rarely studied but is highly necessary.\n* The experimental studies are comprehensive, robust and sound. \n* The paper is very well-written and well-structured.\n\nWEAKNESSES:\n* Evaluations in Section 5 need more rigor:\n    a. Narrow Task Selection: The evaluation tasks, particularly for protein understanding, are limited in scope. The absence of important tasks such as fitness prediction and EC number predictions means that the chosen tasks may not provide a complete picture of the model's performance across various aspects of protein understanding.\n   b. Missing Experimental Details: The paper omits crucial experimental details, including: hyperparameter selection process for the evaluation, train/validation/test set division. Also, it's not clear whether the evaluation data is omitted from training. These make it challenging to evaluate the fairness of comparisons.\n\n\n* Lack of Diverse Sampling Strategies: The paper does not explore different sampling methods, which is a common practices in the field. For instance, it doesn't consider the approach used in training models like ESM, where sampling by UniRef is employed to upsample rare protein sequences. This practice might drastically change the scaling behaviour.\n\n\n* Lack of insights: The paper can feel a bit like a lot of experiment findings stacked without insights. Tt would be good to provide some intuitions about the findings, such as why is masked language modelling overfitting while causal modelling does not, and the difference in the scaling behaviour between protein sequences and natural language sequences. \n\nHowever,  these weaknesses do not negate the paper's contributions\n\nQUESTIONS:\nSee above\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}