Efficient Stagewise Pretraining via Progressive Subnetworks

Clarity: Section 4.2 is extremely unclear.

• “Normalization is dependent on initialization, however, the choice of initialization is never discussed.”

The initialization has been clearly stated in the statement of Lemma 4.3. We show that under Xavier initialization normally used in deep learning, normalization layers, and residual connections help stabilize the training of a linear model under RAPTR.

• “What we are supposed to conclude from Lemma 4.3”. “... the link between Section 4.2. and the experimental results in Table 2 (is unclear)”

Lemma 4.3 aims to understand how stable the loss is at the transition between different stages, which is key to why RAPTR can provide speedups. In general, one could expect the loss to spike and training could spend additional training steps (FLOPs) to first recover after the stage transition. However, Lemma 4.3 shows that because of the presence of normalization layers and residual connections, the blowup in training loss at the transitions is bounded by $\mathcal{O}(1/\sqrt{L})$ for linear networks. We believe a theoretical understanding of the benefit of RAPTR for training is an important future direction to consider as a follow-up to this work.

• Many claims do not seem to make sense. “Suggesting a worst-case lower bound of $\mathcal{O}(L^{-0.88})$”?

The lower bound of $\mathcal{O}(L^{-0.88})$ follows from the observation that $\Psi_{\ell}$ varies with $\ell$ as $(L/\ell)^{0.12}$ (please see figure 2b), while the norm of the output changes linearly with $L$ (please see figure 2a). Hence, the net difference between the two losses under study by theorem 4.2 should be bounded by $\mathcal{O}$ $(1/L$ $\sum_{\ell} \Psi_{\ell} / ||F||)$ $=$ $1/L$ $\mathcal{O}(1/L^{1.88} \cdot \sum_{\ell} 1/\ell^{0.12} )$ $=$ $\mathcal{O}(L^{-0.88})$.

• “A simple calculation shows the gap between L-1 random subnetwork and full model scales as $\mathcal{O}(L^{-c})$” even mean?

We observe that the norm of the output increases as $\Omega(L^{0.88})$ in all settings, while $\psi_{\ell}$ varies as $(L/\ell)^{0.5}$ with $\ell$. Hence, the net difference between the two losses under study by theorem 4.2 should be bounded by $\mathcal{O}(1/L \sum_{\ell} \Psi_{\ell} / ||F||)$ $= \mathcal{O}(1/L^{1.38} \cdot \sum_{\ell} (1/\ell)^{0.5} ) = \mathcal{O}(1/L^{0.88})$.

We will clarify more on this in the main paper in the revision.

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.

Figures carry minimal information as well.

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

Questions

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.

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.

Errors and typos

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

References

1: Efficient training of language models using few-shot learning. Reddi et al'23

2: No train no gain: Revisiting efficient training algorithms for transformer-based language models. Kaddour et al.'23

We thank the reviewer for their remarks and suggestions. Please find our responses to your questions below.

Weaknesses

Improvement

• “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.

• The set of evaluation tasks in Table 3 isn’t comprehensive enough.

We further perform evaluations on other downstream tasks in Table 11 in the appendix. In summary, RAPTR performs better than baseline models in completion tasks like Lambada, Hellaswag, and Storycloze, in addition to QA tasks like web questions and natural questions. We also provide downstream performance at higher shot in-context settings. On average, RAPTR performs 1-2% better than baseline training while being 20% faster.

• “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.

Notation issues

• “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.

• "How were schedules selected", "Changes in schedule from table 2 to table 3?"

The decision of schedules is motivated by the schedules in [1]. The schedule lengths were decided to be either Equal or Proportional, depending on the flop counts we need in the end. Please see the response to the previous question for more details. Please find more details in Appendix C in the new version. As for the difference between RAPTR and Stacking schedules in Table 2, we picked the best schedule for each based on Table 1.

We thank the reviewer for their remarks and suggestions. Please find our responses to your questions below.

Questions

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.

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.

Q3. “...show if the design choice affects both pretraining and downstream” (for fixed layers and scaling)

In Table 5 in the appendix of the revision, we conduct an ablation study on different fixed player candidate sets for BERT-base training with RAPTR. We observe that fixing the first and the last layers helps pre-training loss. Furthermore, differences in pre-training also show small differences in fine-tuning performance on different downstream evaluations.

Weaknesses

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.

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.

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

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.

References

1: No train no gain: Revisiting efficient training algorithms for transformer-based language models. Kaddour et al.'23

2: Understanding contrastive learning requires incorporating inductive biases. Saunshi et al.'22

3: Same pre-training loss, better downstream: Implicit bias matters for language models. Lieu et al'22

We thank the reviewer for their comments on the simplicity and effectiveness of the proposed method. Please find our responses to your questions below.

Questions

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.

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.

Weaknesses.

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.

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.

References

1: Accelerating Training of Transformer-Based Language Models with Progressive Layer Dropping

2: Towards Efficient NLP: A Standard Evaluation and A Strong Baseline

We thank the reviewer for their comments and suggestions. Please find our responses below.

How competitive are UL2 baselines? A reference on hyperparameters will be useful.

We use the default settings from [1] and do not optimize for hyperparameters. The mixture rates for the dataset are used from [2]. We tuned hyperparameters for the learning rate and learning rate schedule (square-root decay used in the original paper, and cosine decay).

How does the compute-efficient frontier change with this pre-training procedure? Model sweep will help quantify the amount of improvement in downstream performance.

We are sorry but we're having difficulty understanding the question. If it pertains to how RAPTR behaves concerning model scale, scaling behaviors for pretraining beyond 1 billion parameters can become quite costly. In the eventual version, we may incorporate additional experiments on a 3 billion scale if the reviewer deems the results worthwhile. Investigating the scaling laws for RAPTR could be a subsequent exploration following this research.

[1] UL2: Unifying Language Learning Paradigms. Tay et al.’22

[2] LLaMA: Open and Efficient Foundation Language Models. Touvron et al.’23

We have added more results in Table 8, where we compare RAPTR to the progressive layer drop algorithm of [1] on UL2 models. RAPTR is $3-4$% better on downstream tasks on average.

1: Accelerating training of transformer-based language models with progressive layer dropping. Zhang and He et al.'20