Efficient Construction of Model Family through Progressive Training Using Model Expansion

Thank you for your feedback. We would like to address the weaknesses as follows:

Weakness 1.

The idea is simply to apply the bert2BERT method to LLMs for model expansion. Although this new experimental setting is good and interesting, the novelty is relatively limited.

Our novelty lies in (1) framing a practical but previously overlooked application—model family construction—as an evaluable problem and (2) empirically demonstrating that tailoring model expansion techniques for this purpose is highly effective.

Weakness 2.

The intuition and insight are not well illustrated. Why can progressive training lead to better results? How does this intuition guide the design choices of each component? This can be made more clear.

The core intuition is that model expansion provides superior initialization by transferring knowledge from smaller models, placing larger models in better regions of the loss landscape compared to random initialization [1]. Our primary contribution is computational efficiency. The performance improvements, while encouraging, represent a bonus rather than the main objective. Better initialization leading to better performance is well-established in neural network optimization [2]. Understanding the precise mechanisms in our setting presents interesting future work.

Weakness 3.

Missing reference. There are many model expansion works in the literature. This paper cites only very few of them. The missing reference can be supplemented. For example, Knowledge Inheritance for Pre-trained Language Models. NAACL 2022.

Our Related Work section currently cites representative and foundational works in model expansion that are directly relevant to our progressive training approach, providing sufficient coverage of the core techniques and recent developments in the field. Nevertheless, we recognize the value of providing a more comprehensive literature review and specifically appreciate you pointing out "Knowledge Inheritance for Pre-trained Language Models" (Qin et al., 2022). While Qin et al. (2022) explore the related and important area of leveraging smaller models to efficiently train larger ones (primarily via knowledge distillation ), our work on progressive training using model expansion offers distinct contributions. Our primary focus is the efficient construction of an entire model family by iteratively applying model expansion. This approach aims to reduce the total computational cost for the family significantly, ideally to that of training only the largest model, while also demonstrating benefits in performance and cross-model behavioral consistency.

References.

[1] Chen et al. NAACL 2022. bert2BERT: Towards Reusable Pretrained Language Models.

[2] Qi et al. NAACL 2021. When and Why Are Pre-Trained Word Embeddings Useful for Neural Machine Translation?

Thank you for your positive feedback. We would like to answer your questions as follows:

Answer to Does the method remain stable at even larger models?

Thank you for your question. If 'stable' refers to maintaining computational efficiency gains, we believe our progressive training method should indeed hold its efficiency for even larger models. Consider a general model family where each model doubles in size (e.g., [X,2X,4X,8X]). Our paper defines the computational cost (FLOPs) as $FLOPs = 6 \times X \times T$. Adopting the Chinchilla law of using 20 tokens per model parameter for training from scratch ($T_{i}^{\mathrm{scratch​}}=20 \times X_i​$), the costs for independent training would be [120$X^2$, 480$X^2$, 1920$X^2$, 7680$X^2$]. With our progressive training approach, the total computational cost to build the entire family is equivalent to the cost of training the largest model (8X in this case) from scratch. So, the percentage reduction in computational cost is therefore ($2520X^2/10200X^2) \times 100 \approx 24.7$%. As long as similar scaling proportions are used for the family, this percentage saving should be largely preserved even as the absolute size X increases.

Answer to How would depth‑only or LEMON expansion change results?

Thank you for this insightful question. While our study demonstrates that bert2BERT's capability to expand both width and depth dimensions is sufficient for effective progressive training, exploring alternative model expansion methods is indeed valuable.

Experiment:

To address this question, we conducted additional experiments using depth-only expansion (specifically, layer stacking as employed in recent works [1] ) on a smaller model family: 500M → 1B → 2B parameters under the Chinchilla setting.

Results:

The above table shows that depth-only stacking also enables progressive training to achieve comparable or better performance than independent training, while constructing both 500M and 1B models using only the computational budget required for training the 2B model independently. These results are consistent with our main findings using bert2BERT. Future work should explore how different model expansion methods perform within the progressive training framework to further optimize the approach.

Answer to What is the val‑ppl change on reused tokens in the Fixed‑Data ablation?

The slight performance variations are expected due to different data ordering and sampling during training. This comparable performance is consistent with prior work showing that training with repeated data up to ~4 epochs maintains similar performance without significant degradation [2]. In our Progressive+Fixed Data setting, we reuse 140B tokens (less than 1.5 epochs of repetition), which falls well within the range where data repetition doesn't harm performance.

References.

[1] Du et al. NeurIPS 2024. Stacking Your Transformers: A Closer Look at Model Growth for Efficient LLM Pre-Training.

[2] Muennighoff et al. NeurIPS 2023. Scaling Data-Constrained Language Models.

Thank you for your review. We would like to address your concerns.

Weakness 1.

The 25% reduction in training compute seems rather small in the large picture. Many existing research (data engineering techniques, architectural changes, etc.) works have shown much greater compute savings. Moreover, this only reduces the training cost, but not the inference cost. So the total compute saving throughout the life of a model family is even lower when accounting for inference costs.

Thank you for your feedback. On the impact of 25% savings, we respectfully disagree with the reviewer for the following two reasons:

(1) 25% is not small in the development of LLM While a 25% reduction in training compute might initially seem modest compared to some other optimization techniques, its practical and financial significance in the context of large-scale language model (LLM) development is substantial. For instance, the 25% compute savings achieved in our experiments, equivalent to 3,200 GPU hours, translates to a direct cost reduction of approximately $44K USD (calculated using H100 GCP pricing https://cloud.google.com/compute/gpus-pricing ). This is a considerable sum that can impact project budgets. Furthermore, when considering the application of our progressive training to even larger, industry-standard model families like Llama 2 (e.g., 7B, 13B, 70B parameters) under a Chinchilla-optimal setting, the potential savings could escalate to an estimated$154K USD. This level of cost reduction is not trivial; it can make large-scale model development more accessible and sustainable, and our method has the potential to save large budgets and prevent wasting unnecessary computational resources.

(2) Progressive training is complementary to existing techniques We believe that the value of our progressive training technique should not be judged solely by comparing its cost reduction percentage with that of orthogonal approaches, especially considering its complementarity. Regarding inference costs, while not a direct focus of this study, the procedural construction of a model family might offer benefits in selecting optimally sized models for diverse inference scenarios, potentially leading to indirect savings.

Weakness 3.

Section 5.1 provides an extended discussion on the advantages of consistency across model sizes, such as in speculative decoding. However, these claims are mostly speculative. While KL divergence analysis (Section 5.2) does quantify similarity across adjacent models, the claimed practical benefits are not experimentally verified.

We thank the reviewer for highlighting the need to experimentally validate the claimed benefits of model consistency (Section 5.1), particularly for speculative decoding. To address this, we conducted new experiments.

Speculative Decoding Experiment:

• Setup:

  • 8B generator model, 4B drafter model
  • K=8 speculative tokens
  • 1,000 prompts from FineWeb-Edu validation.

• Metrics

  • $D_{KL}(P_{\mathrm{drafter}} \textbar \textbar P_{\mathrm{generator}})$: KL divergence between the draft and generator model's output distributions, calculated on the FineWeb-Edu validation set (lower indicates higher consistency) as noted in Section 5.
  • Acceptance Rate (%): Average percentage of tokens drafted by the 4B model accepted by the 8B model as evaluated in the original speculative decoding paper [1].
  • Generation Time (seconds): Average generation time for all prompts.

• Comparison: To specifically isolate how the draft model's characteristics, particularly its consistency with the generator, affect speculative decoding performance, we fixed the generator model as the 8B Progressive model. We then compared outcomes when using 4B draft models constructed via Progressive versus Independent training approaches.

• Results: | Draft Model (4B) | Generator Model (8B) | $D_{KL}(P_{\mathrm{drafter}} \textbar \textbar P_{\mathrm{generator}})$ | Acceptance Rate (%) | Generation Time (s) | | :--------------- | :------------------- | :------------------------------------------------- | :------------------ | :-------------------- | | Progressive | Progressive | 0.3378 | 93.20 | 7.02 | | Independent | Progressive | 0.4281 | 87.2 | 8.24 |

• Analysis:

The results clearly demonstrate the practical advantages of employing a draft model constructed through Progressive Training. The pairing of a 4B Progressive draft model with an 8B Progressive generator model (Prog-Prog) exhibited superior performance, underpinned by greater consistency between the two models. This enhanced consistency is quantified by a lower $D_{KL}(P_{\mathrm{drafter}} \textbar \textbar P_{\mathrm{generator}})$ value for the Prog-Prog pair, compared to using an independently trained model as a drafter. This improved alignment in output distributions directly translated to more effective speculative decoding: the Prog-Prog configuration achieved a significantly higher acceptance rate of 93.20%, as opposed to 87.2% for the Inde-Prog setup. Consequently, this led to a reduction in inference time, with the Prog-Prog pair completing generation approximately 14.8% faster (7.02s vs. 8.24s). These findings experimentally substantiate our claim in Section 5.1 that the consistency fostered by Progressive Training offers tangible benefits. The lower KL divergence between progressively trained models leads to more effective draft proposals in speculative decoding, thereby improving acceptance rates and reducing inference latency. This highlights that Progressive Training not only provides an efficient way to construct model families but also yields models that are better suited for downstream applications like speculative decoding due to their enhanced inter-model consistency.

References.

[1] Leviathan et al. ICML 2023. Fast Inference from Transformers via Speculative Decoding.

We thank the reviewer for reading our paper and feedback. We would like to address the weaknesses as follows:

Weakness 1.

The novelty of paper is limited. The method employs existing off-the-shelf model expansion tools. The idea of progressive training is also not new. The paper essentially repackages them for constructing a model family without proposing a fundamentally new algorithm or mechanism.

Our primary contribution lies in proposing a comprehensive learning framework for effectively constructing a model family and empirically demonstrating its effectiveness. While the reviewer appears to have concluded that our work lacks novelty because we employed an existing method for the model expansion part, we believe that the novelty of our work is in line with COLM's recognition of contributions in Empiricism, Data, and Evaluation, Technological Impact, rather than being limited to algorithmic innovation (https://colmweb.org/ReviewGuide.html).

Our paper presents novel empirical findings specific to this procedural approach to model family construction. For example, we demonstrate that our iterative expansion procedure using existing methods leads to greater behavioral consistency across model sizes, as quantified by KL divergence between adjacent models in the family. Crucially, we also show that strategically adjusting learning rates based on model size during this progressive family construction yields more efficient training and better-performing model families. These empirical insights into building model families with existing tools offer new, practical contributions to the field.

Weakness 2.

The most interesting aspect is the problem definition and the constraints on the FLOP computations and the token budgets. However, this is not algorithmically addressed. In the experiments, instead of searching or optimizing for configurations under constraints, the authors use a predefined token budget schedule. This makes the paper pure empirical without any new proposed methods or surprising findings.

Our research objective is to efficiently construct model families within given computational budgets that reflect standard and practical settings. We tackle this objective under widely recognized resource constraints (Chinchilla law and 2x Chinchilla law), which represent standard LLM pre-training research practices rather than arbitrary choices [1][2]. Our experimental design comprehensively addresses the core research question through controlled comparisons that isolate the effects of progressive training. While empirically focused, our methodology is both rigorous and complete for the stated objectives. We would also like to emphasize that pretraining demands substantial computational resources and is time-intensive. As such, we generally refrain from using complex methods during pretraining in order to ensure fast and efficient training. For this reason, we adopted a simple, predefined scheduling strategy. We believe that most researchers and developers with experience in pretraining would agree with this point of view. Furthermore, our findings are not trivial: (1) computational efficiency with maintained performance, (2) behavioral consistency across model families with practical implications for techniques like speculative decoding, and (3) empirical validation that benefits persist when controlling for data exposure. These represent meaningful advances in practical model family construction.

References.

[1] Lialin et.al. ICLR 2024. ReLoRA: High-Rank Training Through Low-Rank Updates.

[2] Zhao et al. ICML 2024. GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection.