Training Compute-Optimal Protein Language Models

Summary

The 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. The study contains a comprehensive number of experiments and its main findings are: * 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. * 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. * 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. * 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. * 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). * 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).

Strengths

* The paper is clearly written and despite the diversity of analysis conducted the main findings and their importance are well conveyed in most cases. * The paper provides the most comprehensive analysis to date on protein language model scaling and how to optimise the computational resources available. * 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.

Weaknesses

* 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). * 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. * 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.

Questions

- Do you have any intuition about why the CLM transferability to MLM scales with model size and the inverse is not true?

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

4

Limitations

I 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.

Summary

The 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.

Strengths

I 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.

Weaknesses

Personally, 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. - 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. - 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…”. - 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. I think if these kinds of concerns are addressed I'd be happy to increase my score. More minor questions/concerns: - 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. - What are the confidence intervals for the scaling exponents for model size and data? - 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?

Questions

- 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? [copied from the weaknesses section] - What are the confidence intervals for the scaling exponents for model size and data? - 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?

Rating

6

Confidence

2

Soundness

3

Presentation

2

Contribution

3

Limitations

I thought the discussion of limitations in appendix C was well done.

Summary

In 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.

Strengths

**Clarity** - The paper is clearly motivated and finding domain-specific scaling laws is valuable for practitioners. - The paper is very easy to follow and experiments are well described. **Experiments** - Authors conduct extensive experiments and the results on scaling laws are solid. - Notably they also derive scaling laws for transfer between CLM and MLM and mixture-of-experts.

Weaknesses

**Lack of discussion of results** - 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. **Novelty** - This work closely follows previous work on English LMs such as Chinchilla and presents very limited conceptual novelty. **Experimental Validation** - 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] [1] Elnaggar, Ahmed, et al. "Ankh: Optimized protein language model unlocks general-purpose modelling." arXiv preprint arXiv:2301.06568 (2023).

Questions

- 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. - 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). [2] Touvron, Hugo, et al. "Llama: Open and efficient foundation language models." arXiv preprint arXiv:2302.13971 (2023).

Rating

7

Confidence

2

Soundness

3

Presentation

2

Contribution

3

Limitations

I believe limitations have been properly addressed in Appendix.

Summary

This 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.

Strengths

Overall, this paper presents a well-motivated study with solid and extensive experimental work, * The paper addresses scaling laws in the context of AI4Science, an area that has been rarely studied but is highly necessary. * The experimental studies are comprehensive, robust and sound. * The paper is very well-written and well-structured.

Weaknesses

* Evaluations in Section 5 need more rigor: 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. 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. * 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. * 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. However, these weaknesses do not negate the paper's contributions

Questions

See above

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

NA

Response to the authors' rebuttal

I appreciate the authors' rebuttal to the weaknesses I've raised. The rebuttal satisfactorily addresses the concerns I have raised and I am confident that the paper should be accepted.

Response to the authors' rebuttal

I appreciate the authors' answer to the limitations I raised and clarifications. It addresses my concerns and I believe the paper should be accepted. I increased my score accordingly.