Stacking Your Transformers: A Closer Look at Model Growth for Efficient LLM Pre-Training

Summary

The paper tackles the problem of model growth during training. The authors focus on the problem of efficient LLM pretraining and analyze a plethora of proposed growing techniques. The paper focuses on and established from the very beginning three clear and important objectives: 1. Comprehensive evaluation for decoder transformer language model training, 2. Generalization of small-scale results to bigger scales and 3. The establishment of clear guidelines for practitioners.

Strengths

The paper is excellently written, with clear goals and intentions from the very beginning. It successfully conveys the message that stacking layers once with a specific growth factor, is the best possible solution. The evaluation framework is comprehensive and sufficient. The authors analyze 4 dominant growing (although ignoring many others) techniques and present convincing results for "stacking" as the most promising technique. They then present experiments scaling up pretraining tokens and FLOPs and show that results generalize out of the box. Finally, the paper discusses the other two important choices when growing a model, namely when to perform these operations and how much to grow the model. Inspired by scaling laws, the paper fits curves on the "optimum" iso-FLOPs points derived by growing at different points and by different amounts. Overall the paper is very comprehensive in its experiments both in the main text and the appendix.

Weaknesses

The motivation behind the three objectives is clear, yet at times insufficient. O1. Reported results in the literature are indeed focusing mostly on BERT pretraining. Although the evaluation framework proposed here is complete and comprehensive, it is not clear why decoder-transformer language modeling is fundamentally different and model growing results will in this case be different compared to BERT pretraining (apart from post-LN in the architecture). The authors should compare their findings with established results on BERT pretraining in the literature and discuss how their best configurations compare to them [1]. Note that other papers ([2, 19 from your pdf]) are dealing with growing operations during language model pretraining. O2. Viability for scaling is important. One of the issues of model growth not discussed sufficiently in this paper is the problem of "diminishing returns", meaning that one can save compute, but the ratio of compute saved becomes smaller as training progresses (this is hinted in Figure 6a), especially since you are performing a single growing step. Section 4.2 is very interesting and the experiments enlightening, but the analysis is superficial. [2] has a more detailed analysis to determine growing timings. In general, assuming that models are training following a power law loss, growing timings can be determined based on the gradient of the training losses. This was analyzed in detail in [3], where they describe in detail how to determine these timings. Your case is admittedly more complicated since you are growing in a non-functional-preserving manner, but fundamental take-away messages should be the same. Section 5, is interesting, but both of the results seem to disagree with multiple findings in the literature. Although additional insights are presented in the appendix, these results seem preliminary and should be taken with a grain of salt. [1] Wang, Peihao, et al. "Learning to grow pretrained models for efficient transformer training." arXiv preprint arXiv:2303.00980 (2023). [2] Shen, Sheng, et al. "Staged training for transformer language models." International Conference on Machine Learning. PMLR, 2022. [3] Anagnostidis, Sotiris, et al. "Navigating Scaling Laws: Compute Optimality in Adaptive Model Training." Forty-first International Conference on Machine Learning.

Questions

Some additional comments: - Can you comment with more detail on why you think stacking multiple-times does/should not work? - How specific are the results to Transformers and language models? - Since your growing operator is not function-preserving, should there be spikes in the loss curves? If yes why are these not visible in the current figures?

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

2

Limitations

See above.

Summary

The paper introduces a novel method for pre-training large language models (LLMs) efficiently using model growth techniques. The authors tackle three main obstacles: lack of comprehensive evaluation, untested scalability, and absence of empirical guidelines. They propose a depth-wise stacking operator, "Gstack," showing its effectiveness in reducing training time and computational resources across various LLM sizes, verified through extensive experiments and evaluations on multiple NLP benchmarks.

Strengths

The concept of model growth isn't new and stacking the weight depthwisely is not novel, but the paper innovatively applies it to LLMs through a systematic method, introducing the "Gstack" operator for depth-wise expansion. The paper is grounded in rigorous experimentation, presenting reproducible results with publicly available code and detailed experimental setups, ensuring high quality and reliability.

Weaknesses

Although extensive, the experiments mainly focus on model performance from a computational and speed perspective. The impact on final model accuracy or downstream task performance is less explored, which could be crucial for practical applications.

Questions

Will this stacking method influence the robustness or generalizability of the LLMs on downstream tasks?

Rating

6

Confidence

2

Soundness

4

Presentation

3

Contribution

3

Limitations

The paper lacks theoretical insights that could explain why the Gstack operator works well in practice. Including a theoretical analysis or rationale could strengthen the paper and provide a deeper understanding of the method.

Summary

This paper studies the model growth technique for large language models, which expand a smaller pretrained large language model into a bigger one. The authors consider four natural way of expanding the parameters of the smaller pretrained large language models, and find that duplicating layers is the most effective technique. Therefore, this work delve deep into the dupicating/stacking layers technique, including the following aspects: - Scaling model sizes - Scaling training tokens - estimating laws - determining the growth timing - determining the growth factor and a bunch of other ablation studies.

Strengths

- Except in the introduction, the “model growth” concept is not introduced early enough, the paper is in general well-written and easy to follow. - The conducted experiments are very comprehensive and well-designed. - The problem is well-motivated and useful. - The proposed solution is very natural, simple yet effective. - In general, this paper provides very insightful observations for future usages.

Weaknesses

- [Minor] The experiments are conducted on Transformer architecture. Trying different model architecture such as SSM can be more interesting. - [Minor] The concept “model growth” is not introduced clearly until the related works.

Questions

See weakness.

Rating

7

Confidence

3

Soundness

4

Presentation

4

Contribution

4

Limitations

The paper is mostly about empirical observation. However, the reason behind this phenomenon is still unclear and not much discussed.

Summary

The presented work systematically investigated the major obstacles of applying model growth methods to large language models and the corresponding solution. The empirical results reveals that depthwise stacking methods works the best to LLMs. The paper then studied how to practically use the depthwise stacking methods in detail.

Strengths

- The experiments are comprehensive and convincing. - This paper provides detailed and clear empirical guidelines. - The experiment results shows that the proposed model growth practice does accelerate the training of LLM by a lot.

Weaknesses

- This paper provides practical guidelines on the usage of model growth methods on LLM training, but doesn't provide any theoretical analysis or intuition on why certain methods work or not work. - Current study focuses on to making use of an existing model growth method under the scenario of LLM training, instead of improve it or modify it. It is possible that there exists a better model growth method for LLM that doesn't fit into the 4 catagories summerized by this paper.

Questions

See Weaknesses.

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

4

Limitations

Authors discussed the limitations.