Paper: Efficient Stagewise Pretraining via Progressive Subnetworks (ICLR-2024)

Reviewer comment (Experiments):
The experimental results do not show a consistent improvement of the method over the other stacking methods. Table 1 and 2 show very limited improvement in some cases, for the same value of FLOPs; for schedules that have been chosen by unclear means.

Author response:
- **“Results do not show a consistent improvement of the method over the other stacking methods”. "Table 1 and Table 2 show very limited improvement”.**

Table 1 aims to compare RAPTR and the baseline models at equal training steps. The aim is to show that RAPTR can achieve similar perplexity and fine-tuning performance as the baseline. However, motivated by previous works like [2], we conducted extensive experiments (with learning rate tuning) to compare baseline and RAPTR at equivalent flops. Even with extensive hyperparameter tuning, RAPTR still performs better than the baseline at shorter horizon training, when the models observe only a few epochs of training over the training data. This observation can be crucial for large language models, which are mostly trained with a single epoch over data. 

- **“Schedules … chosen by unclear means.”**

We followed the schedules from [1] and reported the best that we observed during our experiments. We provide a description of the schedules and an ablation study on the effect of schedules on bert-base in section C in the appendix in the newer version.

Reviewer comment (Experiments):
The results in Table 3 are surprising. RaPTr does show an improvement on Tydi QA and SQuADv2 over Stacking, but not for the other metrics. The alternative schedule is not detailed and merits more introduction. It is impossible to compare Stacking with RaPTr as long as Stacking has not been trained with this alternative schedule, or a similar one. This could be done by training the entire model with 30K steps, then training only the first layers, and restacking the previously trained layers to extend the model, rather than stacking the same layer. Otherwise, the improvement due to RaPTr seems relatively small without this schedule. The small variance claimed by the authors is only proven with the alternative schedule and not RaPTr in general.

Author response:
- **“RAPTR does show an improvement on Tydi QA and SQuADv2 other Stacking, but not for the other metrics.”**

This is not entirely correct. RAPTR without a 30k warmup is also better on TriviaQA and is almost as good on SuperGLUE. We have included more detailed evaluations and find RAPTR to be better than stacking on multiple metrics.

- **“The improvement due to RAPTR seems relatively small without this (initial warmup based) schedule”,
“It is impossible to compare Stacking with RAPTR as long as Stacking has not been trained with this alternative schedule”**

This is not true; RAPTR without this schedule is already ~4% better on evals compared to Stacking, which is a significant gap. With the alternative schedule, the gap is even larger (5.5%). Thus the schedule is not as important when compared to Stacking.

Besides, the schedule proposed by the reviewer for stacking with warmup is RAPTR in disguise, with a different subnetwork schedule (the first stage trains all layers, the second stage only trains the first few layers, and the depth is increased incrementally over stages). At this point, the method is no longer performing “stacking” of layers.

Reviewer comment (Writing):
All along the article, the notations of the schedule of RaPTr are very inconsistent and unclear. The schedule \(T^{1:k} × (p^{1:k}, I^{1:k})\) is never used after its definition if I'm not mistaken. This is logical as it is often useless: \(I\) is denoted to be the entire set during all of the experiments, and the stage times are always equal. Only \(p\) ever changes during the experiments. Similarly, the sense of \(I\) reverses during the paper. At first, it defines the set of fixed layers before seemingly being reversed during the experiments as \(I\) is the entire set \((1,24)\) (which is also not consistent with the claim that the first and last layers are always fixed). \(I_s\) in Algorithm 1 is not defined and just seems to be \(I\).

Author response:
- **“The schedule $T^{1:k} \times (p^{1:k}, I^{1:k})$ is never used”.**

This notation is used for a more general definition of RapTr in Algorithm 1.

- **“$I$ is denoted to be the entire set during all of the experiments, and the stage times are always equal.”, “The definition of $I$ reverses”**

This is not true. $I$ is defined as the set {1, 24} and not the range $(1, 24)$ in our experiments (see “Notations for RAPTR and stacking schedules” section).

As for stage times, they need not always be equal. For BERT we use equal schedules, but for UL2 experiments, as presented in section B.2, we use a different “proportional” schedule, where we spend time proportional to $i$ in stage $i$. We have added more details in section C in the new version. The only quantity that stays fixed during training is the fixed set $I$ that comprises the first and last set. However, this set in itself can be changed over time to give other RAPTR schedules.

Reviewer comment (Novelty):
The technique proposed is very close to other layer-dropping techniques, in particular Zhang et al. 2019, which also aims at reducing the FLOPs needed for the training. The main difference is that the dropping is done at the start of training rather than the end, and the different rescaling.

Author response:
**Novelty: “The technique proposed is very close to other layer-dropping techniques”**

We would argue that RAPTR (and progressive subnetworks) is different from layer-dropping in motivation & concept, implementation, and effectiveness. Please see the common response for details on differences.

Reviewer comment (Presentation):
The Figures either bring very limited information or are very hard to read. Figure 1 could have been substituted by a simple Table. Figure 2 is never referred to anywhere. Figure 3 is extremely hard to read and never really explained well.

Author response:
**Figures carry minimal information as well.**

Thank you for the suggestion. We will add additional details to the plots in the new version.

Reviewer comment (Experiments):
The statement “The intuition is that since the first and last layers play different roles” leads to the question of whether not only the first and last but \(k\) first and/or \(k\) last layers may need to be fixed. Are the early or later layers more important to be fixed?

Author response:
**Q1: Further ablation studies are necessary on the layers that are necessary to be fixed during training.**

 Thank you for the suggestion. We conducted more experiments for the BERT-base, which shows that fixing more layers during training can hurt performance. Fixing the first and last layers seems to be the best. Please refer to Table 6 in the revision.

Reviewer comment (Evaluation):
Reviewer does not understand why the experiments on BERT are qualified as “short horizon settings”.

Author response:
**Q2: Why are experiments on BERT qualified as “short horizon settings”?**

 We conducted experiments on BERT-base in Table 2 under 3 different settings, short horizon (75k training steps on baseline and 100k training steps on RAPTR and stacking), medium horizon (170k training steps on baseline and 225k training steps on RAPTR and stacking), and longer horizon setting (510k training steps on baseline and 675k training steps on RAPTR and stacking). By short horizon, we mean training for fewer steps (or epochs), motivated by experiments in Kaddour et al.

Reviewer comment (Writing):
In the Introduction, “Gradual stacking (Gong et al., 2019) and Progressive stacking (Reddi et al., 2023)” are reversed; it should be Progressive for Gong et al. and Gradual for Reddi et al.

Author response:
***Errors and typos***

Thank you for the suggestions. We have fixed the typos and the errors pointed out by the reviewer in the revision.

Reviewer comment (Writing):
In Section 2, “the computed based” is incorrect.

Author response:
***Errors and typos***

Thank you for the suggestions. We have fixed the typos and the errors pointed out by the reviewer in the revision.

Reviewer comment (Writing):
In Figure 1.b, “for fixed set” should be “for the fixed set”.

Author response:
***Errors and typos***

Thank you for the suggestions. We have fixed the typos and the errors pointed out by the reviewer in the revision.

Reviewer comment (Writing):
“table” and “appendix” should be capitalized. Replace “fig” by “Figure”.

Author response:
***Errors and typos***

Thank you for the suggestions. We have fixed the typos and the errors pointed out by the reviewer in the revision.

Reviewer comment (Novelty):
The idea of training sub-layers progressively is not novel, which is similar to [1][2].

Author response:
**W1: “The idea of training sub-layers progressively is not novel [1,2]”**

We would argue that RAPTR (and progressive subnetworks) is different from layer-dropping [1] in motivation & concept, implementation and effectiveness. Please see the common response for details on differences. The motivation of [2] is very different, it is to provide efficiency during inference, not training. The aim of [2] is to train a model on which one can perform early exit for inference speeds. However, such training can possibly degrade the performance of the baseline model.

Reviewer comment (Experiments):
RAPTR introduces many hyper-parameters and it would be difficult to tune these hyper-parameters in the pre-training setup, which would hinder the application of this type of work.

Author response:
**W2: “RAPTR introduces many hyper-parameters and would be difficult to tune.”**

The important hyperparameters introduced in RAPTR are the number of stages, path lengths, and the number of training steps in each stage. We concur that the number of stages and path lengths in each stage is a heuristic that needs to be pre-decided before training.

We ran some more ablation studies (table 5 in the revision) on BERT-base that show that RAPTR is robust to stage schedule selection. Since it is computationally expensive to verify in UL2 settings, we continue with the following design principle. train an L layer network with L/2 - 2L/3 - 5L/6 - L schedule for RAPTR with either uniform schedule (equal time in each stage) with ~1.33x speedup or proportional schedule (time proportional to $i$ in stage $i$) giving ~1.2x speedup. For other target speedups, please see section C in the revised version.

Reviewer comment (Experiments):
How to determine the hyper-parameters of RAPTR, e.g. the number of stages and the training steps in each stage? Does it affect the learning rate schedule? Is it affected by the model scale?

Author response:
**Q1: “How to determine the hyper-parameters of RAPTR?”**
 
In our explorations, we tried simple and intuitive schedules adapted from stacking approaches, and that already provided good results. It is very likely that tuning for the schedule could lead to larger gains. We also explored the effect of different schedules with similar FLOPs in Table 3, and found most of them to be reasonably effective and better than baseline. The schedule we used for most results, and that we prescribe, is: train an L layer network with L/2 - 2L/3 - 5L/6 - L schedule for RAPTR with either uniform schedule (equal time in each stage) with ~1.33x speedup or proportional schedule (time proportional to $i$ in stage $i$) giving ~1.2x speedup. This in addition to other design choices (like scaling, fixing layers, and full training warmup) has been effective for both BERT and UL2.

Reviewer comment (Reproducibility):
What is the training setup of baselines and the proposed method in Table 2?

Author response:
**Q2: “What is the training setup of baselines and the proposed method in Table 2?”.**

We follow the same training setup as Table 1. To equate the flop counts of the baseline and the proposed method, we train the proposed methods longer. For example, when the baseline is trained for 510k steps, stacking and RAPTR (which are supposed to be 1.33x faster) are run for 675k steps (675k $\approx$ 510k $\times 4/3$ ). To ensure we have a fair comparison, we run each method at different learning rate scales and report the best eval loss among them. 

The results show that under extensive hyperparameter tuning which allows fair comparison, RAPTR and stacking are better than the baseline model, especially at a shorter horizon (fewer epochs) settings.

Reviewer comment (Theory):
Theoretical results are good while limited to simplified linear residual network settings.

Author response:
**W1: “Theoretical results are good while limited to simplified linear residual network settings”**

Thank you for appreciating the theoretical results. The results aim to understand the importance of layer normalization and residual connection layers in the model, and analysis for Transformers is still in very nascent stages. Thus results in the linear setting allow us to clearly isolate the role of these design choices.

Reviewer comment (Evaluation):
The gains in pretraining flops diminish in the asymptotic long training setting.

Author response:
**W2: “Gains in pretraining flops diminish in the asymptotic long training setting.”**

That is correct, this observation aligns with those made in [1] for BERT training. However recent language models (like GPT, UL2, Llama) only make 1 (or few) passes over a huge dataset, and are not in the “asymptotic” setting. In this setting our paper shows that RAPTR (and even stacking) improves training time by 20% or more, which could be a significant saving in such settings where models are trained for a few months.

Reviewer comment (Evaluation):
Downstream task improvements lack sufficient analysis on why RAPTR has implicit biases on different tasks, for example, it seems to hurt the multilingual QA performance when adding the 30k.

Author response:
**W3: “Downstream task improvements lack sufficient analysis on why RAPTR has implicit biases on different tasks”, and “hurt the multilingual QA performance when adding the 30k”**

This is indeed an interesting question, and we could not find a simple explanation for these findings. Perhaps it deserves a separate exploration like some recent papers that are dedicated to this topic of inductive bias [2, 3].

Reviewer comment (Evaluation):
Besides the flops, the real wall-clock time might also be good to provide, given in some cases, flops disagree with wall-clock time for efficiency measurements [1].

Author response:
**W4: “...whether the proposed method scales with model sizes and fits to LM is unknown”**

In this paper we have tried 3 different model sizes, 100M, 300M, and 1B, and also two kinds of language models: encoder-only MLM and the more recent decoder-only causal LM. In both cases, we find RAPTR provides speedup and downstream benefits. This exploration is more diverse than many prior papers on efficient pretraining.

Our 1B baseline model already demonstrates good downstream 1-shot performance, and thus we considered it an appropriate setting. Exploring scaling behaviors for pretraining beyond 1B parameters can get very expensive. However, if the reviewer believes that this will add more value, we can include some experiments on a 3B scale in the final version.

Reviewer comment (Experiments):
Besides different architecture and objectives (BERT/UL2), whether the proposed method scales with model sizes and fit to LM is unknown.

Author response:
**W4: “...whether the proposed method scales with model sizes and fits to LM is unknown”**

In this paper we have tried 3 different model sizes, 100M, 300M, and 1B, and also two kinds of language models: encoder-only MLM and the more recent decoder-only causal LM. In both cases, we find RAPTR provides speedup and downstream benefits. This exploration is more diverse than many prior papers on efficient pretraining.

Our 1B baseline model already demonstrates good downstream 1-shot performance, and thus we considered it an appropriate setting. Exploring scaling behaviors for pretraining beyond 1B parameters can get very expensive. However, if the reviewer believes that this will add more value, we can include some experiments on a 3B scale in the final version.

Reviewer comment (Evaluation):
Could the author elaborate on the real wall-clock time gain with the pretraining experiments.

Author response:
**Q1: "...elaborate on the real wall-clock time gain with the pretraining experiments"**

The reason we reported FLOPs is because wall-clock time can be highly hardware and implementation-dependent. We tested wall-clock times on the machines we used and found that a theoretical FLOPs speedup of 1.33x for BERT-base was equivalent to a wall-clock speedup of 1.26x. For UL2, FLOPs speedup matches wall-clock speedup (due to the increased depth and width of the model), so 1.2x speedup in the paper is also the wall-clock speedup. Please see section D in the appendix in the revision.

Reviewer comment (Evaluation):
Could the authors explain more on the varied performance gain in Table 3.

Author response:
**Q2: "...explain more about the varied performance gain in Table 3"**

We observe that RAPTR improves on some tasks more than others compared to the baseline. The trend is that gain is higher for contextual QA tasks (TydiQA and SQuADv2) and lower for non-contextual QA tasks (TriviaQA). Overall RAPTR improves over the baseline across the board, while requiring ~20% fewer FLOPs. Furthermore, we have added more results on other downstream tasks and higher in-context settings (5-shot) in the revision. RAPTR shows similar improvements over the baseline model in all settings.