Efficient stagewise pretraining via progressive subnetworks

We thank the reviewer for their helpful comments and suggestions regarding our work. Please find our responses to your questions below.

Q1: Although the polynomial example illustrates how RAPTR learns lower and higher degree components, it is overly specific and lacks sufficient theoretical support to convince me that RAPTR is better than PLD.

Response: The theoretical setting on polynomials was meant to be an illustrative model on how RaPTr could perform better than PLD. The primary purpose is to provide useful insights for the success of RaPTr that can guide future work. While the difference in curriculum on the subnetwork size between RaPTr and PLD might seem minor at first, our theory highlights that such a change in curriculum could lead to substantial improvements in performance (which is clearly corroborated by our empirical results). The key insight is that dropping more layers later in training reduces its capacity to capture complex correlations in data. Our experiments on the polynomial setting clearly display this phenomenon. The experiments on BERT and UL2 also show that RaPTr performs substantially better than PLD. Modeling the real language distribution and providing theoretical results for it has proven very difficult for deep learning, and thus such intuitive theoretical results are the best we can hope for. We view the simplicity of the setting as a strength rather than weakness.

Q2: Theorem 5.3 shows that the loss gap $|L_2(F)−L_1(F)|$ is upper bounded by the stability of the network and that is used to show that the losses between stages can be small for linear models in RAPTR. However, it is not clear if the same may hold for other stacking or dropping methods.

Response: In section 5, we show that models trained by RaPTr are stable in training loss behavior across stage transitions. We create the theoretical setting to bound loss gap $|L_2(F) - L_1(F)|$ in order to explain the stability. Our framework shows that the change in loss can be bounded across the RaPTr stages. This, combined with the order in which subnetwork lengths are increased, our framework presents an effective order of minimizing loss over training. On the other hand, for PLD, one could potentially extend the stability framework to show small loss gaps across stages, however the reverse order of dropping of layers will limit the capacity of the model to minimize loss over training. Moreover, for stacking, previous works have already observed instabilities in training loss behavior [1] . In general, it is not possible to show stability results for stacking. This is because, new layers are introduced at each stage of training, which can arbitrarily impact optimization.

Q3: The empirical verification with BERT in section 5 does not seem to support the theory as in Fig 3, there is no linear decrease in loss gap as depth of model increases.

Response: The theoretical bound is an upper bound on the loss gap. The theoretical setting in section 5 was meant to give a worse case guarantee on how the loss could possibly be stable at stage transitions and how we could attribute the stability phenomenon to architecture components. In fact, empirically, we observe that the loss gap between a ($L-1$)-subnetwork and $L$-layer subnetwork is negative, which is an even stronger and desirable finding than the theoretical result, and understanding this is left for future work.

Q4: Can all GLUE tasks be reported together with results on BERT-large as a base model is quite small and easy to pre-train?

Response: Based on the reviewer’s suggestion, we ran evaluations on all GLUE tasks and reported the results below. Overall, we continue to observe that RaPTr trained models are competitive or better than the best tuned stacking versions. We used BERT experiments mainly for ablations and extensive experiments on design choices. These were mainly meant to show that dropping based strategies can be made competitive to stacking (and even their best versions) which PLD had failed to show before.

For a larger scale setting, we have the UL2 experiments using a decoder-only 1.5B parameter model, a setting that is more prominently used these days and is also much larger than BERT-large. Extensive pre-training and evaluations on a suite of 19 benchmarks for UL2 show that RaPTr can show a large improvement in few-shot evaluations (please see table 5 in the appendix for details results).

Revised table 2:

Revised table 4:

Q5: Are we able to use RAPTR for fine-tuning on downstream tasks instead of pre-training?

Response: Thank you for the interesting question. Indeed RaPTr can in principle be used for finetuning from a pre-trained checkpoint, by varying the layer dropping budget in stages. This is another potential advantage over stacking-based method because RaPTr maintains the model identity . However our primary focus in this paper was pre-training which is significantly more resource expensive than fine-tuning. We believe extending RaPTr to fine-tuning is an important future work, which requires understanding the structure of existing pre-trained models, and is beyond the scope of our current work.

References.

1. Efficient Training of Language Models using Few-Shot Learning. Reddi et al.’2023

We thank the reviewer for continuing the discussion in the rebuttal. Please find our responses to your questions below.

Comment 1: Regarding Q1, while I do agree that theoretical results are quite difficult for deep learning, the example given seemed too specific and I am curious if it can be extended to a family of polynomials instead?

Response: We can generalize the polynomial to a higher degree and more general polynomial of the form $\sum_i \prod_{j \in p(i)} x_j,$ where each $p(i) \subseteq [n]$ and $p(i) \subseteq p(k)$ for all $k>=i$. The important property that we need is the hierarchical nature of components in the polynomial. Our proof builds on existing theoretical works [1,2] that show gradient descent analysis for such hierarchical polynomials. However, our primary aim was to formalize the difference between PLD and RaPTr through a short example, which can be generalized to higher degree polynomials. Experimentally, we show the behavior of RaPTr and PLD in figure 2 on a degree 10 polynomial.

The primary message of our theoretical statement and figure 2 experiments is that RaPTr ensures simple-to-complex learning, which is useful when the simple features can help the model to learn difficult features in the later stages of training. On the other hand, due to decreasing capacity of the model, PLD will restrict the model’s ability to learn the difficult features towards the end of training.

References:

1: Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics. Abbe et al’2023

2: The staircase property: How hierarchical structure can guide deep learning. Abbe et al’2021

Comment 2: For Q2, Is it possible to show that the loss gaps between stages for PLD has a higher upper bound than RaPTr?

Response: Our theoretical upper bounds in section 5 won’t differentiate between PLD and RaPTr on loss gap across stages, and PLD will also be stable. The issue of PLD is not that of loss stability, but the increasing order of layer dropping probabilities and its restrictive capability to capture higher order features, which is studied in section 3.

Comment 3: For Q3, I refer to line 468, "This implies that our upper bound on the loss gap $|L_2(F)−L_1(F)|$ will show a rough linear decrease with the depth of the model." which is the part that confused me about the linearity of the loss gap.

Response: Thank you for catching this! We meant to say that our theoretical upper bound on the loss gap, which bounds the absolute difference between $L_2(F)$ and $L_1(F)$, will show a linear decrease with the depth of the model. Our theoretical upper bound is a worst case upper bound, through which we can only infer $L_2(F) < L_1(F) + O(1/L)$. On the other hand, in practice, we observe that $L_2(F)$ is less than $L_1(F)$, which shows that loss can decrease across stages in RaPTr, and is an even stronger and desirable finding than the theoretical result. We will clarify this clearly in the next version.

Comment 4: Can I have an efficiency comparison in terms of flops and run time for the results in table 5? with regards to Baseline, Stacking and RAPTR.

Response: Table 5 uses Baseline, Stacking and RaPTr runs from table 3. The theoretical FLOP benefits for Stacking and RaPTr compared to baseline are 1.2x, while the runtime benefits compared to baseline on our machine are roughly 1.19x. The run-times of training each subnetwork have been outlined in Appendix F and the efficient RaPTr implementation in distributed setting is outlined in Appendix G. We are happy to answer more questions about the UL2 experiments.

Comment 5: I also acknowledge that fine-tuning is not the main focus of the paper and should be deferred to future work. Though, it seems that LLMs will no longer be pre-trained from scratch but will instead rely on continual pre-training and fine-tuning.

Response: Thank you for proposing this! Understanding the structure of existing pre-trained models is important for using RaPTr during continued pre-training, which is something that we are planning to do as future work.

We thank the reviewer for their continued discussion in the rebuttal phase. Please find our responses to your comments below.

Comment 1: For comment 2, I refer to Q2 in the question part of the review, then if the theoretical upper bounds in section 5 can differentiate RaPTr and other stage-wise pre-training methods? Otherwise, it is not convincing that RaPTr is better.

Response: The theoretical statement in section 5 was meant as a study on RaPTr and not comparison to other training methods. We start the section by asking the question “What are the general conditions under which training loss under RAPTR stays stable at stage transitions?” Furthermore as we remarked in our previous response, stacking does not demonstrate this loss stability empirically, whereas our results show that RaPTr can without requiring any additional tricks. This provides a contrast between stacking and dropping based methods. Results in Section 3 already distinguish between RaPTr and PLD, and the experiments in section 4 show performance improvements with RaPTr on BERT and UL2.

Comment 2: Regarding comment 3, it would be good if the authors can revise the paper to clarify the comment on the Empirical verification with BERT. As we can submit revisions during this discussion period, I would encourage the authors to do so with their clarification.

Response: We have uploaded a new version, with the aforementioned clarification. Changes have been highlighted in magenta for easy identification.

Comment 3: Okay, I understand the results on UL2. However, if the experiments on BERT were primarily for empirical evidence, the experiments to show that RaPTr is a more effective pre-training recipe are not very extensive. Do you have other experiments on llms such as Llama-3.2-1B as well?

Response: We respectfully disagree with the reviewer’s claims that experiments are not very extensive. We have covered two popular settings of pre-training masked language models and causal language models, both at BERT scale (with GLUE evaluations) and a much larger 1.6B parameter scale (with evaluations on 19 benchmarks, 8 benchmarks marked as SuperGLUE). This is as, or arguably more extensive compared to previously published works on efficient training (please see a detailed note on this towards the end). Furthermore as an added contribution, we also provide theoretical results and intuitions supporting various design choices like schedule of dropping probability and loss stability. Such analyses are largely missing from the largely empirical work in this field.

Scale of experiments: We note that each run on UL2-1.6B for 400k steps took us 3 days of training. We have reported numbers roughly for 30 UL2 runs (including all random seeds), amounting to 90 days worth of training. These experiments are at a much larger scale, something not considered in previous efficient training works.

Comparison of extensive-ness with prior works: Relevant published works on stacking (as recent as ICLR 2024) primarily report results on Bert-base models and one additional setting at a scale that is much smaller than the results that we have reported on. Here are some references:

• Masked Structural Growth for 2x Faster Language Model Pre-training. Yao et al. ICLR 2024. (Bert-base, BERT-large, and GPT-2)
• LEMON: Lossless model expansion. Wang et al. ICLR 2024: (Bert-base and Vision Transformer.)
• Efficient Training of Language Models using Few-Shot Learning. Reddi et al. ICML 2023. (Bert-base and Bert-large)
• bert2BERT: Towards Reusable Pretrained Language Models, ACL 2021. (Bert-base, GPT-2)
• Efficient Training of BERT by Progressively Stacking. Gong et al. ICML 2019. (Bert-base)

Regarding LLama 3.2, the pre-training data details are not public so those experiments are not possible to run. Even so, LLama 3.2 1B has been pretrained with 9T tokens for 370k GPU hours which is nearly impossible to run under any reasonable resource availability (e.g. even with 100 GPUs it would require 150 days!). Conceptually there is no reason to believe that the results will be very different from the 1.6B UL2 experiments presented in the paper since they are solving similar autoregressive problem.

Link: https://huggingface.co/meta-llama/Llama-3.2-1B

We thank the reviewer for their positive and supportive comments about our work. Please find our responses to your questions below.

Q1:In the numerical evaluations, the selected settings of Rel. FLOPS seemed somewhat arbitrary (e.g., fixing to 1.33x for Table 1). It would be nice to get some intuition why these were chosen and the sensitivity to these experimental settings (e.g., how do the results hold as we sweep the Rel. FLOPS).

Response: Our experiments follow schedule selections from prior work [Reddi et al. 2023] -- Inverse, Equal and Proportional. For Equal, we split training into equal stages. For Proportional, we increase the length of a stage in proportion to index of the stage (more details in Appendix E). For Inverse, we decrease the length of a stage in proportion to the index of the stage. The FLOP computation was then given by the average number of layers used in subnetworks during training. Thus, for a 4-stage schedule with subnetwork lengths 6, 8, 10, 12 (used in table 1); Equal schedule will give an average layer use of 9 layers out of 12, effectively reducing the amount of FLOPs 1.33x. Similarly Proportional schedule will reduce FLOPs by 1.2x and Inverse schedule will reduce FLOPs by 1.5x. We primarily focus on Equal scheduling for BERT-base and Proportional scheduling for UL2.

Effect of Increase in Rel. FLOPs Intuitively one would expect that a more aggressive schedule with higher FLOPs reduction will eventually lead to degradation of model quality. To verify this, we ran some experiments increasing the number of Rel. FLOPs by switching between Inverse, Equal, and Proportional. As expected, we see an improvement in the model’s performance with lesser speedup in both loss and downstream evaluations. We report experiments on BERT-base after 300k steps of training below for reference.

References

1. Efficient Training of Language Models using Few-Shot Learning. Reddi et al.’2023

We thank the reviewer for their positive comments about the structure and motivation of the paper. Please find our responses to your questions below.

Q1: Why can RaPTr achieve better performance than the baseline full-model?

Response: Thank you for the question! We only have some preliminary hypotheses, like improved inductive bias of a model when the model is explicitly trained to build in a simple-to-complex fashion. Earlier stages train the model to behave as a mixture of random sub-networks, while later stages merge those random sub-networks to form one single model. On the other hand, baseline model training trains all the layers together, which could potentially result in the layers overfitting to noise. We believe that the study on inductive bias of RaPTr deserves extensive exploration on its own.

Q2: How does RaPTr compare to the performance of a baseline model, when the amount of FLOPs is reduced in the baseline model?

Response: We have compared RaPTr to a baseline model, whose FLOP counts have been reduced to match RaPTr’s, for BERT (table 7 in appendix) and UL2 (below). We reproduce those numbers here for ease of reference. With reduced FLOPs, baseline is worse than RaPTr that uses the same FLOPs.

Performance on UL2:

Performance on BERT-Base:

Here, we report the performance of the model across 3 training horizons, where we match the FLOPs of baseline model, Stacking, and RaPTr.

We thank the reviewer for their positive comments about the structure of the paper and the motivation behind RaPTr. Please find our responses to your question below.

Q1: The paper only focused on encoder-decoder models. How would the results change with simple decoder-only language models, which are much more prevalent?

Response: In fact UL2 experiments are with a decoder-only language model. UL2 training only changes the objective function (60% causal LM, 20% prefix LM and 20% span corruption), but still maintains the model architecture to be decoder-only.

Q2: “limited number of evaluations on UL2 in the main paper?”

Response: Our UL2 evaluations have reported the performance of trained models across a suite of 19 benchmarks (8 benchmarks averaged as SuperGLUE) in single-shot and multi-shot settings (please refer to Table 5). These benchmarks already evaluate a wide range of capabilities like closed-domain question answering, open-domain question answering, completion, and SuperGLUE. We restrict our evaluations to these benchmarks, as our experiments are at small scale with the model size being 1 billion parameters, and so we won’t observe non-trivial performance on math benchmarks like Gsm8k and MATH. However, the evaluation scores from these 19 datasets already shows that RaPTr can perform better than baseline (full-model) and stacking.

Q3: How did the authors decide the mixture composition of pre-training dataset? Why didn’t the authors use off-the-shelf dataset (such as FineWeb / FineWeb-edu)?

Response: We had conducted our experiments on UL2 prior to the release of FineWeb and DCLM. Our dataset mixtures were decided based upon the pre-training data mixtures used in LLaMA-1 [1]. Arxiv, C4 (Raffel et al., 2020), Github, and Wikipedia (Foundation) datasets were mixed with mixing ratios 9%, 57%, 17%, 17%, roughly 1:6:2:2 ratio as used in LLaMA-1.

Suggestions:

1: “Simple implementation details buried in the appendix such as how to provide gradients to all the layers, which I expect to be a part of the main paper. “

Response: We thank the reviewer for pointing this out. As our paper was dense with both empirical results and theoretical claims, we decided to move all our implementation details of RapTr to appendix G. We will include a summary of section G in the final version.

2: ”Important implementation details are scattered throughout the paper. All of the details such as sequence length, and batch sizes should be specified at a single location to make it easier to understand.”

Response: We thank the reviewer for pointing this out. We tried to give all details on our experiments in sections D and E. We will improve these sections in the next version.

References:

1. LLaMA: Open and Efficient Foundation Language Models . Touvron et al.

I want to thank the authors for responding to my comments. ICLR allows updates to the paper, so the comment from the authors that "they will update it in the next version" indicates they would like to retain their current version, which is of course up to the authors as it was just a suggestion from a reader's perspective. I think that this is a useful contribution, and I would increase my score from 6 to 8 to get it over the line.