Efficient stagewise pretraining via progressive subnetworks
Authors: Abhishek Panigrahi
Nikunj Saunshi
Kaifeng Lyu
Sobhan Miryoosefi
Sashank J. Reddi
Satyen Kale
Sanjiv Kumar
Venue: ICLR 2025 Poster (2025)
Keywords: Efficient stagewise training
modular training
language model pretraining
implicit bias
simple-to-complex learning
Primary area: optimization
TL;DR: We propose progressive subnetwork training for efficient pre-training of large language models, which improves downstream performance over existing stagewise pre-training methods.

Abstract: Recent developments in large language models have sparked interest in efficient
pretraining methods. Stagewise training approaches to improve efficiency, like
gradual stacking and layer dropping (Reddi et al., 2023; Zhang & He, 2020), have
recently garnered attention. The prevailing view suggests that stagewise dropping
strategies, such as layer dropping, are ineffective, especially when compared to
stacking-based approaches. This paper challenges this notion by demonstrating
that, with proper design, dropping strategies can be competitive, if not better, than
stacking methods. Specifically, we develop a principled stagewise training framework, progressive subnetwork training, which only trains subnetworks within the
model and progressively increases the size of subnetworks during training, until it
trains the full network. We propose an instantiation of this framework — Random
Part Training (RAPTR) — that selects and trains only a random subnetwork (e.g.
depth-wise, width-wise) of the network at each step, progressively increasing the
size in stages. We show that this approach not only generalizes prior works like
layer dropping but also fixes their key issues. Furthermore, we establish a theoretical basis for such approaches and provide justification for (a) increasing complexity of subnetworks in stages, conceptually diverging from prior works on layer
dropping, and (b) stability in loss across stage transitions in presence of key modern architecture components like residual connections and layer norms. Through
comprehensive experiments, we demonstrate that RAPTR can significantly speed
up training of standard benchmarks like BERT and UL2, up to 33% compared to
standard training and, surprisingly, also shows better downstream performance on
UL2, improving QA tasks and SuperGLUE by 1.5%; thereby, providing evidence
of better inductive bias.

Published as a conference paper at ICLR 2025
EFFICIENT STAGEWISE PRETRAINING VIA PROGRES- SIVE SUBNETWORKS
Abhishek Panigrahi
Nikunj Saunshi Kaifeng Lyu
Sobhan Miryoosefi
Sashank Reddi
Satyen Kale Sanjiv K
i Princeton University
a GGoogle New York
8 Apple MLR, New York
ABSTRACT
Recent developments in large language models have sparked interest in efficient pretraining methods. Stagewise training approaches to improve efficiency, like gradual stacking and layer dropping (Reddi al., 2023; Zhang He, 2020), have recently garnered attention. The prevailing view suggests that stagewise dropping strategics, such as layer dropping, are ineffective, especially when compared to stacking-basek approaches. This paper challenges this notion by demonstrating that, with proper design, dropping strategies can be competitive, if not better, than stacking methods. Specifically, develop a principled stagewise training frame- work, progressive subnetwork training, which only trains subnetworks within the model and progressively increases the size of subnetworks during training, until it trains the full network. We propose an instantiation of this framework - Random Part Training (RAPTR) that selects and trains only a random subnetwork (e.g. depth-wise, width-wise) of the network at each step, progressively increasing the size in stages. We show that this approach not only generalizes prior works like layer dropping but also fixes their key issues. Furthermore, we establish theo- retical basis for such approaches and provide justification for (a) increasing com- plexity of subnetworks in stages, conceptually diverging from prior works on layer dropping, and (b) stability loss across stage transitions in presence key mod- em architecture components like residual connections and layer norms. Through comprehensive experiments, we demonstrate that RAPTR can significantly speed up training of standard benchmarks like BERT and UI.2, up 33% compared standard training and, surprisingly. also shows better downstream performance on UL2, improving QA tasks and SuperGLUE by 1.5%; thereby, providing evidence of better inductive bias.
INTRODUCTION
Large network based language models (e.g. Transformers) have revolutionized the field of NLP. Intriguingly, these language models have demonstrated remarkable emergent abilities that only begin to manifest at large scales (Wei et al., 2022; Schaeffer et al., 2023). However, training such large models is usually very slow and resource intensive (Brown et al., 2020; Touvron et al., 2023; Chowdhery et al., 2022). This has sparked interest in efficient training of large models, necessitating the development of new algorithms and paradigms for efficient pretraining. Traditionally, this was accomplished by designing better optimization algorithms (e.g., (Chen et al., 2023; Gupta et al., 2018; Shazeer & Stern, 2018; Liu et al., 2023a)) that require fewer training steps to reduce the loss. Recently, other paradigms based on slagewise training, especially on depth, have garnered interest. Two such contrasting paradigms are: (a) layer stacking and (b) layer dropping.
Stacking based approaches have been studied extensively since Chen al. (2021) applied it for BERT. Progressive (Gong et al., 2019) and gradual stacking (Reddi et al., 2023) are layer stacking methods that train large models in stages, starting with small network and gradually growing the
Equal Contribution, Work done during internships at Google, Work done when author Google, Correspondence ap34 2cs.ccc.. prri princceon. edu, nsaunshi @goog] com.

---

Published as a conference paper at ICLR 2025
network size by stacking subset of layers onto itself from previous stages. Although effective in reducing FLOPs and training time, their performance is sensitive to stacking schedules and require careful tuning. Furthermore, since the model grows gradually. it is not possible to assess the full model performance during earlier stages (i.e., it is not an anytime algorithm). Additionally, using small model with way fewer parameters for part of the training can possibly hamper model quality (e.g. long term memory (Geva et al., 2021)), especially for single epoch training where each data sample seen once during training. In our experiments($44.. and downstream performance may even be worse than baseline (full-model) training.
we observe that the pretraining loss
Layer dropping, on the other hand, maintains the model's identity but drops layers random during training. This area of study remains largely unexplored. The closest work for training efficiency progressive layer dropping (PI.D) (Zhang He, 2020), a heuristic approach where FL.OPs are saved by increasingly dropping more layers as training proceeds, thus, decreasing its overall capacity over time. Such a strategy to improve training efficiency typically comes at the expense of quality, and is generally worse than stacking approaches (Kaddour et al., 2023). Through theoretical analysis and experiments, we first show that there a fundamental problem with existing dropping strategies like PLD since dropping more layers later during training be detrimental to learning complex features. Other instantiations (Fan et al., 2019; Zhang et al., 2019; Liu et al., 2023c) do not achieve training efficiency, they were proposed regularization strategy or for improving inference efficiency. This naturally prompts the central question of the paper:
Is it possible to design principled and robust stagewise dropping based techniques that are competitive or better than stacking techniques?
In this paper, we propose a novel generalization of stagewise dropping approaches called pro- gressive subnetwork training that fundamentally addresses the above concers, and provides an affirmative answer to this question. The key components of this framework are:
(P1) Maintain common base model of interest and train subnetworks of varying complexities within the model.
(P2) Progressively increase the complexity of subnetworks in each stage to explicitly impose this simple-to-complex behavior.
For (PI), one can use very general subnetworks in this framework (e.g. subset of layers or layers with smaller width). This strictly generalizes layer dropping techniques where just layers are dropped to derive efficiency. Furthermore, (P2) is motivated by the theoretically and empirically studied phe- nomenon that gradient-based training learns functions of increasing complexity over time (Kalimeris et al., 2019; Abbe et al., 2022) and starkly diverges from earlier layer dropping ideas where more layers are dropped as training proceeds. As we shall see soon, even a simple instantiation of this idea can already be competitive or better than stacking approaches that have been extensively studied.
Given the above framework, a natural question arises: how do we select the subnetworks? As a concrete instantiation, we choose subnetworks to be random parts of the network, and in each stage we perform forward and backward passes over them. The size of the subnetworks is gradually increased over time with a fixed schedule. This approach, which we call Random Part Training (RAPTR), reduces the total number of FLOPs and wall-clock time in the initial stages, but unlike stacking, every parameter of the full model contributes the training loss at all stages, and it allows us to track the full model throughout training. In this paper, we demonstrate that this simple, yet powerful, strategy not only has a solid theoretical foundation, but also shows strong experimental results for BERT and U1.2 pretraining, both in terms of training efficiency and improving model quality. We summarize the main contributions of the paper below:
We introduce novel stagewise training called progressive subnetworks that generalizes prior dropping strategies in $2. Specifically we explore Random Part training (RAPTR) that trains a part comprised of random subnetwork, with the size of the subnetwork progressively increasing in stages. We specifically study variants of RAPTR where the random subnetwork selection is restricted to single model axis e.g. depth (section 2.1) or intermediate MLP width (section 4.3).
We leave the combination of multiple axes for future investigation.
Through analysis on polynomial data in $3, we demonstrate a fundamental problem with prior dropping techniques like progressive layer dropping dropping more layers towards the end can
2

---

Published as a conference paper at ICLR 2025
hurt the models' ability to capture complex correlations in the data. Furthermore in this polyno- mial setting, we show, both theoretically and empirically, that RAPTR can effectively learn higher order components of the underlying ground-truth function much better than progressive layer dropping. This provides a strong justification for progressively increasing complexity of subnet- works in stages. The polynomial setting can potentially aid future research in efficient training.
We conduct extensive experiments on BERT (Devlin et al., 2018) and UL2 language models (Tay et al., 2022) to compare depth instantiation of RAPTR with dropping techniques (PLD), and gradual stacking and its improved variants. On BERT-Base ($4.1), RAPTR demonstrates competitive performance
state-of-the-arr stacking strategies similar FLOOP,, while being better than baseline' training with 1.33x fewer FLOPs. For UL2-1.6B language model ($4.2), RAPTR matches the pretraining perplexity of baseline training by requiring 1.2x fewer FLOPs. suite
Furthermore, despite having the same perplexity, RAPTR has much better performance of 12 downstream benchmarks, evaluated in 1-shot and 5-shot in-context settings. This suggests desirable inductive bias of RAPTR that improves its generalization abilities beyond perplexity.
We discuss novel implementation strategy for subnetwork training translates FLOPs improve- ments to wall-clock speedups. This is particularly valuable for distributed training where naive conditional dropping implementations either fail or do not speedup training (appendix G).
We establish the firyt theoretical basis for stagewise training based on dropping of layers that studies the behavior at stage transitions. In particular, we characterize stability conditions under which RAPTR yields smooth transition in loss across stage transitions ($5). We connect this stronger
to the idea of loss preservation (Chen et al., 2021) and show that RAPTR has an property of loss improvement in each stage. Through an illustrative setting with linear networks, theoretically demonstrate the importance of modern architecture components, like layernorms and residual connections, for such stability.
2 PROGRESSIVE SUBNETWORK TRAINING
Algorithm I: k-stage RAPTR
Require: Blocks fa,.... for Schedule Tia x II.), total steps T. data distribution D.
Initialize 01.... 62. Ta
for 1 1- l: do for T: - Tot1 do
Sample hatch (X, Y) D. Sample for I..
B(p.) for & Ie, set G
Set Y + X
for j 1 + do
( 1
if G: v y +
to (Y--1)
Figure I: Pictorial description of stagewise RAPTR where the number of layers being skipped progressively decreases over stages.
end for Compute loss with ( and take optimizer step.
end for end for
We setup notation for the rest of the paper and describe the proposed framework of progressive subnetwork training and the specific method RAPTR that we study.
Notation. [n] denotes the set 21,2,..., n).. 01:k is used to denote a sequence of scalars al....,, E R.. denotes vector in Rd æi.A denotes a sequence of k vectors æI...., 2s E Rd. We use X E xd to denote a matrix containing a sequence Di:n as rows. We consider the setting where the goal is to learn L-layer sequence-tee -oteenne------eeecee neural network. (e.g., residual networks (with = 1) and standard Transformer network).
model). Let fi...., , R-xxd - R denote set of sequence-tu cectenceett-sqqquee to-seeunnce functions. Then, F R.Nxd x R. + R N takes an input sequence X E R. Nxd scalars a1:L, and outputs sequence Y defined recursively with intermediate outputs Y()) (L-1)) (L) as
Definition 2.1 (L-layer lequence-to-sequeee
( (f-1) + aefe aefe(y ((-1)) for all E L,
F(X;a) = Y
We refer full-model training baseline training in our paper
3

---

Published as a conference paper at ICLR 2025
where Y denotes X, Y is the output of f-th layer.
Standard model training uses scalars a; = 1 for E [L]. For simplicity, we use F(X) to denote the output and y ( ...., (()) as the intermediate layer outputs of F under standard model training, unless specified otherwise.
Progressive subnetwork training is stagewise training method with two crucial components:
Subnetwork selection: At each iteration, we select a (possibly random) subnetwork of the L- layer neural network. The forward and backward passes are computed based on this subnetwork. Progressive subnetwork sizes: The size of the subnetworks is progressively increased in stage- wise manner, starting from small subnetworks and ending with full network in final stage.
This framework is quite general since subnetworks can be selected along any model dimension, e.g., attention heads, hidden neurons in MLP layers. In this work, primarily focus depth.
2.1 RANDOM PART TRAINING: RAPTR
One strategy select subnetwork is choosing a random part within the model. To better isolate the selection mechanism, we consider variants where selection happens only along one model axis (e.g. depth or width or heads in Transformers). The systematic exploration of combination of multiple axes, which is a combinatorially large space, is left as future work. For the purpose of our discussion, focus on the depth axis i.e., random layers of layers selected each step. Concretely, given network F to be trained, we pick random subset of layers for the forward pass and bypass the rest of the layers by residual connections. With slight terminology abuse, we call this approach Random Part Training (RAPTR). We will revisit RAPTR variants on other axes in $4.3.
More concretely, let p denote the probability of randomly selecting layer and I denote subset of layers that are always included, i.e. they are never bypassed. We define a random subnewtork before describing RAPTR in more detail.
Definition 2.2 (p, --subnnetookk. Let 7I.7 be Bernoulli samples, with G 1 for all i I, while G (p) for all & I. The set of all layers for which S = 1 represents a random subnetwork while layers with Ç are bypassed. The output of the selected subnetwork on sequence input X is equivalent to the output of the model given by F(X, SI..) (1).
Stagewise progressive RAPTR. The pseudo-code is provided in Alg. 1. At high level, the total training is split into k stages. Each stage uses random (p.,I.) J-subnetworks between steps T, and Ts+1- We denote this schedule by Titk x "Pi:k.II:A' We progressively increase the random subnetwork selection pattern across stages, moving towards full model training by either increasing the probability of including each layer or fixing more layers in I, or both. Such schedules satisfy: (a) p. S ps. and (b) I, C I3 for all S s S 3 S k."
Training efficiency. Training step in RAPTR involves computing forward and backward passes only on random subnetwork. Training (p, Z) random subnetwork will require (L. / FLOPs relative to standard training that uses all L layers. By using p < 1 for majority of training. RAPTR can improve training efficiency.
General behavior of layer probabilities. RAPTR provides general framework, where we can make discrete shifts in the average subnetwork lengths across stages. One potential concern is training instability due to discrete transitions. In $5, we show the stability of training loss of standard Transformers under R RAPTR, and showcase the importance of different architectural components.
3 ILLUSTRATION: PROBLEMS WITH EARLIER SUBNETWORK TRAINING INSTANTIATIONS
In this section, we discuss problems with existing subnetwork training instantiations like progressive layer dropping (PLD). Recall that PLD drops more layers as training proceeds to enable FLOPS efficiency (see $D.1 for details). This is in contrast to the depth variant of RAPTR, that drops
IIf N 1, represent the input and the intermediate outputs in their vector representation, æ, ( (1) M
4

---

Published as a conference paper at ICLR 2025
Evaluation Loss
PLD
RaPTY
Baseline
FED Rareltee
Degann
Degtes
Doignos
Component
10
20
Steps
Steps (x5000)
Stops (K500))
Stops 045000)
Figure 2: Evaluation loss (left) and component error (3 on the right) on basis polynomials for a 20 layer residual network. Labels are generated from polynomials of degrees
10 (Eq.2). Both RAPTR and PLD schedules use 20% fewer FLOPs than the baseline (dark vertical lines indicate RAPTR'S phase transitions: details in SD.3). (a) RAPTR matches the baseline's 's evaluation loss, while PLD is worse. (b) RAPTR learns lower order terms faster and picks up higher order terms in the later stages. PLD is worse at capturing higher degree terms for its reduced expressivity in the end.
fewer layers as training proceeds. While this difference may seem superficial at first, it has deep implications from both conceptual and performance standpoints. Here, we argue through theoretical and empirical analysis on polynomial data, that this can have very significant impact on learning complex correlations in data. In the BERT and UL2 experiments in $4, we similarly observe that progressive layer dropping performs significantly worse than RAPTR.
We demonstrate the importance of progressively increasing the complexity of subnetworks in RAPTR (as opposed to progressively decreasing it in PLD) through simple, yet instructive, setting: learning neural networks on polynomial data. For simplicity, we restrict the discussion to sequences of length N 1.
Polynomial learning setting. Suppose x are sampled from a uniform distribution over (+1)d. The ground-truth label is polynomial f* of degree k defined
# m
F*(x) = ΣΣ CI 1=
II XxP
(2)
&ESU
where Ci2 (0, and St, random subsets of [d] of cardinality t. Such polynomials have been studied in great detail (Abbe et al., 2022; 2023) where the higher degree terms capture more complex correlations and are harder learn. We consider training of 20-layer residual network He et al. (2016) F on such data using square loss: B(F(()) (2))²² Here, each residual block consists of a single 4d-hidden dimension MLP layer with ReLU activation.
Empirical Observations, We are interested in measuring how well each method captures the higher degree terms. To this end, we estimate the component of learned function F onto each basis polynomials. This can be done simply using Ctj = - F(x)) IesSs, z4 since basis polynomials are orthogonal under the uniform distribution of the boolean data (O' Donnell, 2014). For each degree I < k; we define the error as
E (e..)-..)) In Figure 2, we observe the following:
2/21 4
PLD, that drops more layers as training proceeds, fails effectively learn higher order complex correlations and, ultimately, hurts performance.
RAPTR is competitive to baseline in terms of evaluation loss; in fact, RAPTR learns every compo- nent effectively. Furthermore, RAPTR quickly learns lower order terms and then picks up higher order terms during later stages.
The learning pattern of RAPTR is reminiscent of the implicit bias of standard SGD training to learn functions of increasing complexity (Kalimeris et al., 2019; Xu et al., 2019; Rahaman et al., 2019; Cao et al., 2019; Abbe et al., 2022). RAPTR naturally imposes this simple-to-comple behavior through size of subnetworks and also provides training efficiency in the process.
Theoretical Analysis. To further illustrate our point, we characterize the behavior of RAPTR for a simple 2-layer residual network, where each block consists of single neuron with non-linear sine activation. We consider simple label function F+) x3 v3 Ø 21 for this analysis.
Lemma 3.1 (Informal, of theorem LI). For a small enough learning rate, 2-phase RAPTR first learns lower degree component and then the higher degree component.
5

---

Published as a conference paper at ICLR 2025
Table 1: Losss and fine-tuning results for BERT-Base after 675K steps and using equal peak learning rate (10-4). Key observations: (a) RAPTR achieves similar loss as baseline at 1.33x reduced
FLOPs. (b) RAPTR achieves 0.02 better evaluation loss than gradual stacking for same schedules. (c) RAPTR has slight edge on downstream fine-tuning task.
Rel.
Eval Loss
Method
Schedule
FLOPs
MNLI
QNLI
SST-2
Avg.
Baseline PLD
1.33
1.76
81.5
90.5
91.4
87.8
1.
1.8
81.7
89.0
90.8
87.2
Stacking
6.8 10-10-12
1
1.78
6-9-12
1.77
80.9
89.8
87.3
RAPTR
6-8-10-12
1.75
82.1
89.8
92.4
88.1
6-9-12
1,75
82.3
89.2
91.0
87.5
In contrast, one layer is not expressive enough to represent the true label function (I.2). Thus, pro- gressively dropping more layers (e.g., PLD) will reduce its expressivity. Remark: Stacking (Reddi et al., 2023) approaches, that grow the model size, can be shown to have similar simple-to-compler complex inductive bias in the above setting. We leave a detailed study of this connection to future work.
4 EXPERIMENTS
We present comprehensive experiments that compare RAPTR with baseline (full-model) training, gradual stacking and progressive layer dropping (PLD) for BERT and UL2. Since our primary goal was to make subnetwork based strategies competitive to stacking, we primarily focus on compar- isons to gradual stacking. However, we also show comparisons on BERT with improved versions of gradual stacking, used MSG (Yao al., 2023), LEMON (Wang et al., 2023b), and bert2BERT (Chen et al., 2021). The details of relation between gradual stacking and these methods are in $D.2.
RAPTR setup: We describe the experimental setting of RAPTR from 12.1. We keep equal to the first and last layer (I= (1, L.) across all stages". Additionally, scaling layer outputs during initial stages helps track validation loss better, though its final effect minimal (details $J). Finally, we use constant learning rate in all phases, except in the final phase where we decay the learning rate, following gradual stacking (Reddi et al., 2023).
Notations for RAPTR and stacking schedules: We succinetly denote the stage schedules by set of subnetwork lengths separated by hyphens. For stacking (RAPTR) training for a 24-layer model, 3-6-9-12 refers to 4 stages with 3. 6,9, and 12 layer (average subnetwork) training respectively. The length of each stage is chosen based on the reduction in FL.OPs to achieve during training (see }E). Unless specified otherwise, we split the training steps equally across the stages.
Wall clock speedups: In In SF, we report wall-clock times for UL2 and BERT experiments. For UL2, a 1.2x FLOPs reduction leads to a 1.19x speedup, closely matching FLOPs improvements. For BERT, 1.33x FLOPs reduction gives a 1.26x speedup. We also introduce new implementa- tion strategy for RAPTR in $G, essential for achieving speedups in distributed training where naive conditional dropping fails or doesn't 't give any speedup. As runtime depends on architecture and hardware, we focus on FLOPs reductions in the following discussion.
4.1 EXPERIMENTS WITH BERT
Experiment setup. We pretrain BERT models (Devlin et al., 2018) on Wikipedia- Books dataset with AdamW optimizer (Loshchilov & Hutter, 2019). We report the evaluation loss and fine-tuning performance of the trained models on GLUE tasks (Wang et al., 2018) (refer to $D for details).
We observe that RAPTR achieves similar or better evaluation loss to baseline training, despite baseline training using 1.33x more FLOPs than RAPTR for SERT-BASE Compared to gradual stacking at similar FLOPs, RAPTR again has better validation loss. Additionally, all three methods exhibit similar performance in downstream fine- tuning tasks (see Table 1). However, PL. performs significantly worse compared others.
Comparisons to gradual stacking and baseline:
Comparisons with variants of stacking: We further compare RAPTR to variants of gradual stack- ing used in MSG (Yao al., 2023), LEMON (Wang al., 2023b), and bert2BERT (Chen et al.,
Please see Appendix Table 10 and the discussion later for the relevant ablation study

---

Published as a conference paper at ICLR 2025
Table 2: We compare the performance of RAPTR against depth-wise stacking operations used in ben2BERT, LEMON, and MSG, which implement slight variations of gradual stacking (see ap- pendix D.2) after 300K steps of training. RAPTR demonstrates competitive or superior performance across all variants, achieving comparable loss to the best-performing version of gradual stacking.
Schedule
Loss
SST-2
QNLI
MNLI
SNLI
RTE
MRPC
CoLA
QQP
Avg.
bert2bert MSG
3-6-12
89.3.u.)
84.5 o.1
aus
015
56.1 223)
75.2 11.3)
(34)
(84)
3-6-12
1.87
88.4 013)
83.9 01.)
76.4 (13)
87.0 (03)
56.0 (1:3)
74.8 (1.0)
76.6 (1.1)
89.4 (o.)
79.1
MSG
o.3)
85.6 (12)
sm
(1))
(1:))
77.2 (1:3)
(22)
(8.)
LEMON
6.9-12
1.83
89.8 (1))
87.2 (14)
77.3 (03)
86.4 (0:1)
58.4 (1.
79.8 (9:6)
78.0 (9:))
89.7 (0:))
81.1
MSG
[+ loyr lngrisarpdihicc)
6-9-12
1.84
89.0 (07)
85.1 002)
77.5 00.1
86.9 (01)
$7.7 (U5
79.9 (3:8)
77.5 (2:2
89.5 (80)
80.4
Stacking
6-9-12
1.84
89.0
85.3 (0.1
77.0
86.2
56.6 (.3))
76.4
78.0
89.3
79.7
RAPTR
6-9-12
1.83
89.3 00.)
85.8 (1.3)
78.6 (1.1
87.0 (0.))
57.7 (0.
80.8 (3:4)
79.1 (3:7)
89.6 (0.1)
81.0
Table 3: Validation loss and 1-shot downstream evals on UL2-1166. Key observations: (a) RAPTR with 30k initial training improves performance by atleast 2 4% on Trivia QA and SquADv2.
RAPTR is atleast 5% better than gradual stacking various downstream tasks, (b) Compared to baseline, RAPTR is 1 - 2% better on all downstream tasks on average. (c) Stagewise training achieves (1 2%) lower variance on Tydi QA and SquADv2, implying implicit benefits of stage- wise training. On an extensive suite of 12 tasks and multi-shot settings, we observe a 1.5%
improvement in performance on average. See table 5 in the appendix for extensive evaluations.
Method
Rel. FLOPs
Eval Loss
Trivia QA
Tydi QA
SQuADv2
SuperGLUE
Avg.
Baseline
1.2
2.06 man
25.0 (2))
34.4 a.
42.1 a9)
60.0 ma
40.4
PLD
2.09 m0)))
21.3 (1))
32.4 an
40.2 (s))
59.9 m2))
38.5
12-16-20-24 Stacking
I
2.08 (.00)
20.1
28.6
36.0
60.4
36.3
12-16-20-24 RAPTR
2.06
25.8
36.7 (m.
44.1 (0.5)
60.9
41.9
2021) (details are in appendix D.2). We focus on depth-only stacking operations in these frame- works in table 2. We show that RAPTR performs competitive even better than all variants of stacking, achieving same evaluation loss as the best tuned variant of stacking.
Loss behavior during phase transitions: Perhaps surprisingly, observe that the training loss decreases quite smoothly when transitioning between two stages in RAPTR (see Fig. 3 (a) for BERT training under RAPTR). This is in contrast stacking, where variants have been devel- oped to explicitly impose loss preservation across stage transitions (Shen et al., 2022; Wang et al., 2023b). We delve deeper into this favorable behavior of RAPTR in $5.
Ablations with RAPTR parameters: We observe robustness of the final trained model to different RAPTR schedules (Appendix Table 9). Furthermore, we observe that fixing the first and last layers at all steps during training helps RAPTR (Appendix Table 10). Since RAPTR is robust to schedule selections (Table 9), we recommend the schedule in the box below for all practical purposes.
Recommended RAPTR schedule selection
Schedule contains phases. Initially, random subnetworks of size L/2 are selected, and we increase by T./6 layers in each transition until the full model is trained in the final phase.
4.2 EXPERIMENTS WITH UL2-1.6B
Experiments on BERT ($4.1) show that RAPTR can be competitive or better than gradual stacking and can outperform baseline training. We ask whether the same observations hold when training billion-parameter language models large text corpora. We pretrain 1.6B decoder-only UL2 model (Tay et al., 2022) with 24 layers for 400k steps with batch size 512. Table reports the validation loss and downstream 1-shot performance without fine-tuning. Please see $D.5 for exper- imental details and table for more extensive evaluations.
Schedules for RAPTR and gradual stacking. In table 3. report for schedule with 4 stages, denoted by 12-16-20-24 The length of each stage has been adjusted appropriately to use an average subnetwork length of 20 out of 24, giving 1.2x FLOP reductions compared to the baseline training
7

---

Published as a conference paper at ICLR 2025
Table 4: Comparisons at 300K steps of training on BERT-BASE. (Top): Comparison between Attention-RalTTr and MSG's s attention head growth operator (Yao et al., 2023). Both use schedules that use an average of heads per layer, reducing FLOPs in self-attention by 1.33x. (Bottom): Comparison between MLP- MLP-RaPTr and MSG's MLP growth operator. Both use schedules that use an average intermediate MLP dimension of 2048, reducing MIP FLOPs by 1.7x.
Eval Loss
SST-2
QNLI
MNLI
SNLI
RTE
MRPC
CoLA
QQP
Avg
Head Operators (6-9-12)
MSG
2.07
87.5 o.
83.4 a.
73.3 o.
85,5 mz)
$2.3 84
72.5 a.u.
73.5 mz
88.6 1..))
77.1
Atteution-RaPTr
1.84
89.8
86.4
77.5
89.5
60.5
80.4
78.7
89.7
81.6
Intermediate MLP size Operators (768- 1 1536-3072)
MSG
1.88
89.0 o.
88.5 mm
77.2 0m
87.2 am
58.3 ()
80.0 a
77.1 m
89.6 i
80.5
MLP RaPTr
MLP-RaPPr
1.84
88.7
85.3
77.3
88.6
55.8 (1
79.8 (1
78.3 (0:))
89.6 10.1)
80.4
(see ŞE). RAPTR uses initial full-model training phase for 30K steps before transitioning to its schedule. For parity in FLOPs, we reduce 30K steps from the final phase. This is useful for RAPTR- attribute this to the fast drop in loss that large models can achieve initially. The ability to include full model training is another benefit of RAPTR over stacking due to its flexibility.
We summarize the key findings from table below.
Perplexity results. At 1.2x FLOP reduction, RAPTR achieves similar evaluation perplexities to baseline training and also performs better than stacking (lable 3).
Downstream inductive bias. Intriguingly, we find that despite having similar perplexity, RAPTR has much better downstream metries compared to both baseline and stacking. The improvements are particularly large for TydiQA (3.8%) and SQuADv2 (2.0%). This perhaps suggests RAPTR has desirable inductive bias of RAPTR, phenomenon that has been recently studied for other pretraining settings (Saunshi et al., 2022: Liu et al., 2023b).
Lower variance. Another notable observation is that RAPTR has lower variance on all metrics and more stable training runs. These inductive bias benefits deserve further exploration.
4.3 SUBNETWORK TRAINING BEYOND DEPTH
Stacking based approaches have been studied for other dimensions beyond depth (e.g. width) by (Gu et al., 2020; Shen et al., 2022; Wang et al., 2023a; Gesmundo Maile, 2023). In this see- tion demonstrate the flexibility of the subnetwork training framework handle other growth operators. We extend RAPTR to intermediate MLP dimensions and attention heads. For attention heads, we apply stagewise training by randomly selecting subset of heads in each layer, increasing the number of heads across stages. For MI. dimensions, we divide each MLP into sub-MLPs, pick random subset of sub-MLPs in each layer define random subnetwork, and gradually in- crease the number selected at each stage. We compare these variants, named AttentionRAPPT and MLP-RAPTR, to their MSG counterparts. Attention-RAP than their corresponding operators in MSG, providing further evidence of better performance than stacking- basckd approaches (table 4). More comprehensive study to other axes is left future work.
and MLP-RAPTR perform better
LOSS PRESERVING BEHAVIOR OF RAPTR
RAPTR provides a general framework, as we can adjust the average size of random subnetworks in discrete steps between stages ($2.1). Training time instabilities are a possibility, as one might expect an arbitrarily large spike in the loss stage transitions, due to the discrete shift the training of the model. Remarkably, we observe that the training loss decreases quite smoothly when transitioning between two stages (see Fig. 3 (a) for BERT training under RAPTR). Hence, we ask the following:
What are the general conditions under which training loss under RAPTR stays stable at stage transitions?
This question has been been extensively studied for stacking based approaches (Wang et al., 2023b; Chen et al., 2021; 2015; Shen et al., 2022). Here, we introduce, to our knowledge. the first frame- work to analyze models under discrete dropping based training strategies, and understand the ne- cessity of different architecture components for robustness to discrete transitions. Related works
8

---

Published as a conference paper at ICLR 2025
look into the generalization theory for dropout (Srivastava et al., 2014) regularization (Gal & Ghahramani, 2016; Arora et al., 2021; Wager et al., 2013; Hu et al., 2019; Helmbold & Long, 2015). Building on these frameworks, The stability analysis, when instantiated for linear networks highlights the importance of two popu- lar architectural choices in Transformer models residual connections and layer normalization.
analysis relies on crucial notion of stability (Arora et al., 2018).
5.1 LAYER STABILITY
We focus on sequences of length N = 1; generalization to N > 1 is fairly easy. Let C R - R represent the training loss of the backbone F. For transformers, F's output passes through a Layer normalization layer before projecting to logits. This makes L scale-invariant (loffe & Szegedy, 2015; Ba et al., 2016), which gives favourable stability condition as any perturbation in F's output results in proportional loss perturbation relative to the norm of F's output. Thus, we make the following assumption.
Assumption 5.1 Reelatie-Lipsscitzzer label y and perturbation 1) E RR, satisfies L(x + 7) y) - L(x, v) S u.c. (Ilvll2 / [|
There exists a constant μlg. > 0 such that for any input
To understand loss changes across stage boundaries of RAPTR, we consider two stage RAPTR: the by dropping single layer at random, and the
first stage trains with subnetwork of length second stage trains the entire model. Suppose L1 and L2 denote the effective loss functions being minimized the two stages. The aim show that F learned in the first stage by minimizing L1 should also result in small value for L2.
Definition 5.2. Let F-g denote the subnetwork after skipping layer É. The stability of network output with respect to dropping layer l is defined (x(x)) [[=+(())) - F())ll2-
Next theorem shows that (F) (F(F) is small if stability scales slower relative output norm. Theorem 5.3 (Informal, ef theorem H.1). Under assumption 5.1 of loss C, (F2(F) (I(F) is upper bounded by (uc/ L)D Ei=1 (x(x))/ [IF(z))l2
The proof follows by observing that (F) = t, EeE(L) C(F-s) and C2(F) = L(F). Thus, the losses are close if F-c is close enough to F for a random &. So, the relative stability of the network determines the success of the approach.
In fig. 3 (b, c, d), we report the stability of BERT-LARGE during training. We observe that the norm of the layers grow linearly with depth, while the stability i grows sub-linearly with depth. This implies that our theoretical upper bound on the loss gap (F) - (FF) will show rough linear decrease with the depth of the model."
Empirical verification with BERT:
We find that (F)) is lower than (F)), even when training with subnetworks of length L 1. This shows that RTTR can reduce training losses at stage transitions without requiring explicit patterns, unlike stacking approaches that need carefully designed growth functions to avoid loss spikes (Wang et al., 2023b; Shen et al., 2022). Moreover, this loss reduction surpasses the concept of loss preservation (Chen et al., 2021), which opens up interesting future directions.
Instead, we focus on the following question, "when is a network stable (by def. 5.2) across the stages of RAPTR?" For simplicity, we consider the special case of linear residual networks. We show that layer normalization layers and the residual connection helps maintain the stability of RAPTR.
Illustrative example: linear networks We present more concrete instantiation of Theorem 5.3 for residual network where the layers fi.z. are linear with parameters WT... The layer output is y (f- + (f-( or y = ((e-1) + (ty(t-1)/ (f=1") depending on whether layernorm is enabled or not respectively We also study another setting with layernorm but no (|ggg]])
residual connection; so y
Lemma 5.4. For 6 > O. a finite input set and dimension d 2 ([(SS)Loo(LL/)) w.p. atleast 8 w.r.. random initialization W ~ (N, 1/2 the following holds true for all input x E S.
'Note here that are talking about our theoretical upper hound, from which can only infer C2(F) < (P))) + (1)L.. However, empirically, we observe (F(F) < Li (F), explaining this is kept for future work.
9

---

Published as a conference paper at ICLR 2025
Tuait
1A
|
CIF
Ton
Best Fit, 0.34141110
-1.
Mecn g00
1903 10.8 B000 30000 Training steps
Layer
Layer
Layer
Figure 3: (left to right) (a): Training trajectory of BERT under 4-stage 6-8- 10-1012 RAPTR, with the stage transitions (denoted by arrows); (b), (c), (d): Stability study on BERT-LARGE trained for 50k steps by RAPTR with subnetworks of length I, - 1. Behavior of (b) norms of intermediate activa- tions y (c) Ve 1 (definition 5.2), and (d) Loss gap between different random subnetwork F- and model F, given by L(F) E ((F-t). Key observations: (b) Norms of the intermediate activations grow linearly with f, (c) V changes slowly with l as (L/E)0.12, suggesting worse-case bound of (L on loss gap based on Theorem 5.3 (d) Interestingly, L(F) S E ((F-e), even when model is trained with I. random subnetworks.
(a) With residual connection & layernorm, Vy(x:) O & [|F(x))l2 = 2(VL). Then the gap in losses between stages scales as O(I/vI).
(b) Without residual connection, V2(x)) = [[1 & [|(()))l2 = O(1). Thus the gap in losses
between stages can be (?(1).
(c) Without layernorm, we have 4 (x(x)) = 12(2(2-1)22) and [F())]l2 O(2 Thus the gap in losses between stages can be (1(1).
In the Appendix, we show similar results for perfectly aligned layers (Lemma H.5). We can consider even more general scenarios, where the layers parameters are expressed as r-combinations of Gaussian and shared matrix (Appendix Fig. 4). We run simulations and observe that for each т, the loss gap between L - 1 random subnetwork and the full model scales (D(L-.44)
RELATED WORKS
Here, we focus on the most relevant works. Refer to appendix C for more related works.
Stochastic depth. RAPTR is similar to stochastic depth, which drops layers with a fixed probability during training to reduce the cost of training deep networks (Huang et al., 2016). A key distinction that the probability of dropping layers is fixed during training and can thus be viewed as a regu- larization (Pham et al., 2019; Steiner et al., 2021; Tolstikhin et al., 2021; Liu et al., 2023c). Fixed stochastic depth has also been used for inference efficiency (Fan et al., 2019). (Devvrit et al., 2023; Valipour al., 2023) explore nested training for dynamic inference with multiple subnetworks.
Subnetwork training. Training random paths and subnetworks has been used in other contexts, such parameter-efficient fine-tuning (Houlsby al., 2019; Pfeiffer al., 2021; Hu et al., 2022; Liu et al., 2022) to reduce memory footprint, distributed and federated training (Dun et al., 2022), and incremental learning (Jathushan al., 2019) avoid forgetting in continual learning. However, these approaches are not aimed at reducing FLOPs during pretraining.
7 LIMITATIONS, CONCLUSION AND FUTURE WORK
We propose a stagewise training framework of progressive subnetworks for efficient pretraining, and evaluate natural instantiation of this framework (RAPTR) based on training random paths/subsets of layers. Overall RAPTR yields better quality language models than baseline training, while showing a reduction of 1.2 1.33x the total FL.OPs. the same speedup, it is also better than prior layer dropping and competitive to stacking based approaches.
The current analysis provides insights into loss stability at stage boundaries but does not explain the observed decrease in loss during transitions. Furthermore, schedule selection is not well understood due to expensive compute necessity and deserves more exploration. Our analysis does not explain the role of initial full-model warmup for UL2 and the desirable inductive biases of RAPTR style training, such as better downstream performance and smaller variance. A deeper understanding of these phenomena could lead to the development of more efficient training algorithms.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staircase property: necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, pp. 4782 4887. PMLR, 2022.
Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. Sgd learning on neural networks: leap complexity and saadll-tt-ssddee dynamics. In The Thirty Sixth Annual Conference on Learn- ing Theory, pp. 2552-2623, PMLR, 2023.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. convergence theory for deep learning via over- parameterization. In International conference on machine learning, pp. 242-232222 242-252. PMLR, 2019.
Raman Arora, Peter Bartlett, Poorya Mianjy, and Nathan Srebro. Dropout: Explicit forms and capacity control. In International Conference Machine Learning, pp. 351-361. PMLR, 2021.
Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Zhang. Stronger generalization bounds for deep nets via compression approach. In International Conference on Machine Learning, pp- PMLR, 2018.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint weXivv1607.06450 2016.
Thomas Bachlechner, Bodhisattwa Prasad Majumder, Henry Mao, Gary Coutrell, and Julian McAuley. Rezero is all you need: Fast convergence large depth. In Uncertainty in Artificial 1352-1361. PMLR, 2021.
Intelligence, pp.
Jonathan Berant, Andrew Chou, Roy Frostig, and Percy Liang. Semantic parsing on freebase from question-answer answer pairs. In Proceedings of the 2013 conference on empirical methods in natural language processing, pp. 1533-1544 2013.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal,
Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 13:1877-1901, 2020.
Yuan Cao, Zhiying Fang, Yue Wu, Ding- Xuan Zhou, and Quanquan Gu. Towards understanding the spectral bias of deep learning. arXiv preprint arXiv: 1912.01 12220/1998, 198, 2019.
Shuning Chang, Pichao Wang, Hao Luo, Fan Wang, and Mike Zheng Shou. Revisiting vision transformer from the view of path ensemble, arXiv preprint arXiv:2308..6655 065448 2023.
Cheng Chen, Yichun Yin, Lifeng Shang. Xin Jiang, Yujia Qin, Fengyu Wang, Zhi Wang, Xiao Chen, Zhiyuan Liu, and Qun Liu. bert2bert: Towards reusable pretrained language models. arXiv preprint 2 2021.
Tianqi Chen, Ian Goodfellow, and Jonathon Shlens. Net2net: Accelerating learning via knowledge 2015.
transfer. arXiv preprint
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Yao Liu, Hieu Pham, Xu- anyi Dong, Thang Luong, Cho-Jui Hsieh, et al. Symbolic discovery optimization algorithms. arXiv preprint arXi::2302.06675 2023.
Eunsol Choi, He He, Mohit Iyyer, Mark Yatskar, Wen-tau Yih, Yejin Choi, Percy Liang, and Luke Zettlemoyer. Quae: Question answering in context. arXiv preprint 18Xiv: 07036, 2018.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint S0MMi::20400231111 2022.
Jonathan H Clark, Eunsol Choi, Michael Collins, Dan Garrette, Tom Kwiatkowski. Vitaly Nikolaev, and Jennimaria Palomaki. Tydi qa: A benchmark for information-seekine question answering in ty pologically di verse languages. Transactions of the Association for Computational Linguistics, 8:454-470 2020.
11

---

Published as a conference paper at ICLR 2025
Luciano Del Coroo, Allie Del Giorno, Sahaj Agarwal, Bin Yu, Ahmed Awadallah, and Subhabrata Mukherjee. Skipdecode: Autoregressive skip decoding with batching and caching for efficient Ilm inference. arXiv preprint arXiv:2307.02626 2023.
Jacob Devlin, Ming- Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint / 180.00400 2018.
Fnu Devvrit, Sneha Kudugunta, Aditya Kusupati, Tim Dettmers, Kaifeng Chen, Inderjit S Dhillon, Yulia Tsvetkov, Hannaneh Hajishirzi, Sham M Kakade, Ali Farhadi, et al. Matformer: Nested transformer for elastic inference. In Workshop on Advancing Neural Network Training: Compu- tational Efficiency, Scalability, and Resource Optimization (WANT@ NeurlPS 2023), 2023.
Yihe Dong, Jean-Baptiste Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine PMLR, 2021.
Learning. pp.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x 16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 1999, 2020.
Dheeru Dua, Yizhong Wang, Pradeep Dasigi, Gabriel Stanovsky, Sameer Singh, and Matt Gardner. Drop: A reading comprehension benchmark requiring discrete reasoning over paragraphs. arXiv preprint arXiv:1903. .0161101611011 2019.
Chen Dun, Cameron Wolfe, Christopher Jermaine, and Anastasios Kyrillidis. Resist: Layer- wise decomposition of resnets for distributed training. In Uncertainty in Artificial Intelligence, pp. 610-620. PMLR, 2022.
Angela Fan, Edouard Grave, and Armand Joulin. Reducing transformer depth on demand with structured dropout. arXiv preprint 1rXiv:1999.11556 2019.
Wikimedia Foundation. Wikimedia downloads. URL https: /ummpsuuuppppp////// . wi welkiimeii...m org.
Yarin Gal and Zoubin Ghahramani. Dropout as bayesian approximation: Representing model
uncertainty in deep learning. In international conference on machine learning, pp. 1050-1059.
PMLR, 2016.
Andrea Gesmundo and Kaitlin Maile. Composable function-presern expansions for transformer architectures. arXiv preprint arXiv:2308.066 06000033, 2023.
Mor Geva, Roci Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers key-value memories. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 5484-5495, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10. 001553/11202..emnlppp m46..446......... URL https: http::::///claaahhooogggo org 2021 emnlp-ma 446.
Linyuan Gong, Di He, Zhuohan Li, Tao Qin, Liwei Wang, and Tieyan Liu. Efficient training of bert by progressively stacking. In International conference on machine learning, pp. 2337-2346. PMLR, 2019.
Xiaotao Gu, Liyuan Liu, Hongkun Yu, Jing Li, Chen Chen, and Jiawei Han. On the transformer growth for progressive bert training. arXiv preprint vrXiv:2010.12562. 2020.
Vinect Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor opti- mization. In International Conference Machine Learning, pp. 1842-1850. PMLR, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
David P Helmbold and Philip M Long. On the inductive bias of dropout. The Journal of Machine Learning Research, 16(1):340334544 2015.
12

---

Published as a conference paper at ICLR 2025
Lu Hou, Zhiqi Huang. Lifeng Shang, Xin Jiang, Xiao Chen, and Qun Liu. Dynabert: Dynamic bert with adaptive width and depth. Advances in Neural Information Processing Systems, 33: 9782-9793, 2020.
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, An- drea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficient transfer learning for nlp. In Interaational Conference on Machine Learning, pp. 2790-2799. PMLR, 2019.
Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation large language models. In International Con- ference Learning Representations, 2022. URL https: wpwp/eeeeeeeeeeeee///ee/e. net forum?
id-nzeVKeeFYE9
Wei Hu, Zhiyuan Li, and Dingli Yu. Simple and effective regularization methods for training on noisily labeled data with generalization guarantee. arXiv preprint arXiv:1905.11368, 2019.
Gao Huang, Yu Sun. Zhuang Liu, Daniel Sedra, and Kilian Q Weinberger. Deep networks with stochastic depth. In Computer Vision-ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 1-14, 2016, Proceedings, Part IV 14, pp. 646-661. Springer, 2016.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Ma- chine Learning Research, pp. 448-456, Lille, France, 07-09 Jul 2015, PMLR. URL https: / //proceedings www.mrrss.ss...... /37/ioffee5.... html.
Rajasegaran Jathushan, Hayat Munawar, H Salman, Khan Fahad Shahbaz, and Shao Ling. Random path selection for incremental learning. arXiv preprint, 2019.
Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. arXiv preprint /70X.v::050000 2017.
Jean Kaddour, Oscar Key, Piotr Nawrot, Pasquale Minervini, and Matt Kusner. No train no gain: Revisiting efficient training algorithms for transformer-based language models. arXiv preprint arXiv:2307.06440. 2023.
Dimitris Kalimeris, Gal Kaplun, Preetum Nakkiran, Benjamin Edelman, Tristan Yang, Boaz Barak, and Haofeng Zhang. Sgd on neural networks learns functions of increasing complexity. Advances in neural information processing systems, 32, 2019.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming-Wes Wei Chang, Andrew M. Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. Natural questions: A benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:452-466, 2019, doi: 10. 10.1162/taclaa URL https: http://aaaaathhooo....... org/ (tgg/01911026
Tao Lei, Junwen Bai, Siddhartha Brahma, Joshua Ainslie, Kenton Lee, Yanqi Zhou, Nan Du, Vin- cent Y Zhao, Yuexin Wu, Bo Li, et al. Conditional adapters: Parameter-efficient transfer learning with fast inference. arXiv preprint arXiv:2304.04947 2023.
Haokun Liu, Derek Tam, Mohammed Muqeeth, Jay Mohta, Tenghao Huang. Mohit Bansal, and Colin A Raffel. Few-shot parameter-efficient fine-tuning is better and cheaper than in-context learning. Advances Neural Information Processing Systems, 35: 15:1590-1965. 2022.
Hong Liu, Zhiyuan Li, David Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochas- tic second-order optimizer for language model pre-training. arXiv preprint arXiv:2305.1434. 2023a.
Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same pre-training loss, better down- stream: Implicit bias matters for language models. In International Conference on Machine Learning, pp. 28-22114, PMLR, 2023b.
13

---

Published as a conference paper at ICLR 2025
Zhuang Liu, Zhiqiu Xu, Joseph Jin, Zhiqiang Shen, and Trevor Darrell. Dropout reduces underfit- ting. arXiv preprint arXiv:2303.01500 2023c.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019. URL https:
https://openreviaw http///pppnreeeeeeeeeeeeeeeeeeep net wwt.:/neeeeeoeeeeeeeeeee forum?id=
Bkg6RicqY7.
Nasrin Mostafazadch, Michael Roth, Annie Louis, Nathanael Chambers, and James Allen. Lsdsem 2017 shared task: The story cloze test. In Proceedings of the 2nd Workshop Linking Models of Lexical, Sentential and Discourse-level Semantics, pp. 46-51, 2017.
Ryan O' O'Donell. Analysis of boolean functions. Cambridge University Press, 2014.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring a broad discourse context. arXiv preprint
2016.
Jonas Pfeiffer, Aishwarya Kamath, Andreas Rücklé, Kyunghyun Cho, and Iryna Gurevych. Adapter- Fusion: Non-destructive task composition for transfer learning. In Proceedings of the 16th Con- ference of the European Chapter of the Association for Computational Linguistics: Main Volume, pp. 487-503, Online, April 2021. Association for Computational Linguistics. doi: 18.18653/v1/ 2021.eacl-main.. URL https: / cellannhoooggo org/
2rrg/202/ 202100112222 1021.eace seacl-mainn 39.
Ngoc-Quan Pham, Thai-Son Nguyen, Jan Nichues, Markus Müller, Sebastian Stüker, and Alexander Waibel. Very deep self-attention networks for end-to-end speech recognition. arXiv preprint 1rXiv:1904.13377 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with unified text-to- text transformer. Journal of Machine Learning Research, 21(140):1667, 2020. URL http:
//jmlr. wttp:/ppppppp//ooooooatoe.ss 2011222000444 html.
Nasim Rahaman, Aristide Baratin, Devansh Arpit, Felix Draxler, Min Lin, Fred Hamprecht, Yoshua Bengio, and Aaron Courville. On the spectral bias of neural networks. In International Confer- ence on Machine Learning. pp. 5301-5310 PMLR, 2019.
Pranav Rajpurkar, Robin Jía, and Percy Liang. Know what you don't 't know: Unanswerable questions for squad. arXiv preprint arXiv: 1806 03822, 2018.
Sashank Reddi, Sobhan Miryoosefi, Stefani Karp. Shankar Krishnan, Satyen Kale, Seungyeon Kim, and Sanjiv Kumar. Efficient training of language models using few-shot learning. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scar- lett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 14553-14568. PMLR, 23-29 Jul 2023. URL ht https: /ttp/o/eecceigggsece/e/eeeees..e press / v2222 reddieddddeee html.
Siva Reddy, Danqi Chen, and Christopher Manning. Coqa: A conversational question answering challenge. Transactions of the Association for Computational Linguistics, 7:249 266,6666666 2019.
Nikunj Saunshi, Jordan Ash, Surbhi Goel, Dipendra Misra, Cyril Zhang, Sanjeev Arora, Sham Kakade, and Akshay Krishnamurthy. Understanding contrastive learning requires incorporating inductive biases. In Proceedings of the 39th International Conference Machine Learning, 2022.
Rylan Schaeffer, Brando Miranda, and Sanmi Koyejo. Are emergent abilities of large language models a mirage? In seirth- Conference Neural Information Processing Systems, 2023. URL https:/ /
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 459 4596-4604. PMLR, 2018
Sheng Shen, Pete Walsh, Kurt Keutzer, Jesse Dodge, Matthew Peters, and Iz Beltagy. Staged training for transformer language models. In International Conference on Machine Learning, pp. 19893- 19908. PMLR, 2022.
14

---

Published as a conference paper at ICLR 2025
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, 1lya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1):1929-1958, 2014.
Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer. How to train your vit? data, augmentation, and regularization in vision transformers. arXiv preprint arXiv:2 100.202000.00.7. 10270, 2021.
Yi Tay, Mostafa Dehghani, Vinh Q Tran, Xavier Garcia, Jason Wei, Xuezhi Wang, Hyung Won Chung, Dara Bahri, Tal Schuster, Steven Zheng, et al. U12: Unifying language learning paradigms. In The Eleventh International Conference on Learning Representations, 2022.
Ilya O Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Un- terthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al. MIp-mixer: An all-mlp architecture for vision. Advances in neural information processing systems, 34:24261- 24272, 2021.
Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou. Going deeper with image transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 32 42, 2021.
Hugo Touvron, Louis Martin, Kevin Stone. Peter Albert, Amjad Almahairi, Yasmine Babaci, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2377.09288 2023.
Mojtaba Valipour, Mehdi Rezagholizadeh, Hossein Rajabzadeh, Marzieh Tahaei, Boxing Chen, and Ali Ghodsi. Sortednet, a place for every network and every network in its place: Towards a generalized solution for training many-in-one neural networks. arXiv preprint arXiv:2309.00255: 2023.
Andreas Veit, Michael J Wilber, and Serge Belongie. Residual networks behave like ensembles of relatively shallow networks. Advances in neural information processing systems, 29, 2016.
Roman Vershynin. High-dimensional probability. University of California, Irvine, 10:11, 2020.
Stefan Wager, Sida Wang, and Percy Liang. Dropout training as adaptive regularization. Advances in neural information processing systems, 26, 2013.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv:1804.07445 2018
Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Superglue: A stickier benchmark for general-purpose language understanding systems. Advances in neural information processing systems, 32, 2019.
Peihao Wang, Rameswar Panda, Lucas Torroba Hennigen, Philip Greengard, Leonid Karlinsky, Rogerio Feris, David Daniel Cox, Zhangyang Wang, and Yoon Kim. Learning to grow pretrained models for efficient transformer training. arXiv preprint arXiv:2303 000900 2023a.
Yite Wang, Jiahao Su, Hanlin Lu, Cong Xie, Tianyi Liu, Jianbo Yuan, Haibin Lin, Ruoyu Sun, and Hongxia Yang. Lemon: Lossless model expansion. arXiv preprint arXiv:2310.07999 2023b.
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, Ed H. Chi, Tatsunori Hashimoto, Oriol Vinyals, Percy Liang, Jeff Dean, and William Fedus. Emergent abilities of large lan- guage models. Transactions on Machine Learning Research, 2022. ISSN 2835-8856. URL https: wwtp/eeeeeeewecc////p/ee net wwtp:/eeeeettttteeeeeeeeeeeeeeet
Survey Certification.
Ji Xin, Raphael Tang. Jaejun Lee, Yaoliang Yu, and Jimmy Lin. Deebert: Dynamic early exiting for accelerating bert inference. arXiv preprint arXiv:2004. 2020.
Zhi-Qin John Xu, Yaoyu Zhang, Tao Luo, Yanyang Xiao, and Zheng Ma. Frequency principle: Fourier analysis sheds light on deep neural networks. arXiv preprint arXiv: 0901.0653, 2019.
15

---

Published as a conference paper at ICLR 2025
Yiqun Yao, Zheng Zhang, Jing Li, and Yequan Wang. Masked structural growth for 2x faster lan- guage model pre-training. arXiv preprint arXiv:2305.02869 2023.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence? arXiv preprint arXiv:1905.07830, 2019.
Hongyi Zhang, Yann N Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. arXiv preprint 19X.v:/901.093211/ 2019.
Minjia Zhang and Yuxiong Hc. Accelerating training of transformer-basee language models with progressive layer dropping. Advances Neural Information Processing Systems, 33:14011- 14023, 2020.
Wangchunshu Zhou, Canwen Xu, Tao Ge, Julian McAuley, Ke Xu, and Furu Wei. Bert loses pa- tience: Fast and robust inference with early exit. Advances in Neural Information Processing Systems, 33:18330-18344 2020.
16

---

Published as a conference paper at ICLR 2025
Appendix
Table of Contents
A Broader Impact
17
B More extensive downstream tests on UL2 models
17
C Additional related works
18
Additional Experiment details
Details PLD
D.2 Details variants of gradual stacking
20
D.3 Details for polynomial training
20
D.4 Details for BERT
20
D.5 Details for UL2
22
E Schedule selection
22
Ablation on different schedules for BERT-BASE
................................
23
Ablation fixed layer set for BERT-BASE
................................
23
F System speed of different subnetworks
23
G Efficient implementation of RAPTR in distributed training
24
H Noise stability for transformers
27
H.1 Noise stability of lincar residual networks
28
I RAPTR motivation with boolean data: Theory
31
I.1
Analysis of the different phases
34
1.2
Useful Lemmas
43
J Scaling
45
J.1 Proofs
45
A BROADER IMPACT
Training LLMs demands extensive computational resources and infrastructure, Our paper proposes an algorithm that aims to accelerate their pre-training. Accelerating the training of large language models (LLMs) can significantly reduce their environmental impact by lowering energy consump- tion and minimizing carbon footprints. Optimizations in algorithms and infrastructure will lead to more efficient use of computational resources and will promote sustainability of the field of AI.
Furthermore, we observe benefits of structured pre-training downstream tasks. Hopefully, these ideas can lead to better pre-training strategies for LLMs.
of these models on
B MORE EXTENSIVE DOWNSTREAM TESTS ON UL2 MODELS
We conduct extensive downstream evaluation of the trained models from table 3 additional downstream tasks and under higher shot in-context setting. Additional tasks include QA tasks like natural QA (Kwiatkowski et al., 2019), web QA (Berant et al., 2013), DROP (Dua et al., 2019), CoQA (Reddy et al., 2019), and QuAC (Choi et al., 2018); and completion tasks like Storyclose (Mostafazadeh al., 2017), Hellaswag (Zellers et al., 2019), and LA (Paperno al., 2016).
17

---

Published as a conference paper at ICLR 2025
Table 5: We extensively compare all trained models from table 3 on multiple downstream tasks and few-shot in-context settings. We follow prompt design from Chowdhery et al. (2022) in each setting. For tasks marked with 1 and we report Exact match and Fl scores respectively. For the rest, we use accuracy. The tasks have been sub-divided into major groups, memorization QA, QA with context, completion, and SuperGLUE. On average, RAPTR is 1-2% better than baseline and stacking.
1-shot
5-shot
Baseline
Stacking
RAPTR
Baseline
Stacking
RAPTR
Trivia QA
25.0 (0.2)
20.1 (1.3)
25.8 (02)
26.5 (I.)
21.1 (1.3)
25.1 (0.5)
Web QA!
6.4 00.0
5.8 (0.6)
7.6 (0.5)
10.6 0..)
9.2 (0.5)
11.2 (0.2)
Natural QAt
4.2 (.5)
3.4 (0.4)
4.4 (0.))
5.7 o.
4.6 (0.)
6.0 (0.3)
Tydi QA!
34.4
28.6 (2.7)
36.7 (1.0)
SQuaDv2
42.1 (2.0)
36.0 (0.9)
44.1 (0.5)
43.2 (30)
36.2 (2.)
44.5 (1.3)
DROP
21.4 (08)
19.5 (0.6)
23.0 (0.4)
CoQA"
49.2 (0.7)
43.9 (0.)
52.4 (0.7)
QuAC
18.1 (0.5)
16.8 (0.6)
18.1 (0.4)
LAMBADA
13.7 (2.9
12.0 (1.1)
18.7 (3.0
30.5 (3.4)
29.5 2.1)
38.7 (2.2)
Storycloze
72.9 (0.4)
71.0 (0.4)
73.3 (0.4)
75.1 (0.))
72.6 (0.8
75.3 (0.6)
Hellaswag
58.3 (0.2)
56.1 (0.1)
58.5 (0.3)
58.3 (0))
56.1 (0.1)
58.4 (0.3)
SuperGLUE
60.0 (0.0)
60.4 (0.9)
60.9 (0.2)
60.7 (0.3)
58.8 (0.5)
62.1 (0.2)
Average
33.8
31.1
35.3
38.8
36.0
40.2
Intermediate norms ||y
WXO/|FoO0
E.A.A, 123
2722 ft:
r-023. Ft 0
eiu 0.0. 1220--1 n- 0.25 Ht - D.S. 1-1
T= 05 Ftt assppt
6.75 Ft 1.00-
e- 0.75, Ft 2
1.0 1146
Ea 1.0.0.00 Ftt
3 3.99//0-1
075, Fi: -22
2.0.1 1322
30
19 Number layers (L)
Layer Index E
Layer Index E
Figure Behavior on a linear residual network with normalization layers with 100 random samples from gd-1, and dimension d = 100. The parameters of each layer E is represented as VFA + vI - ( for shared matrix A E Rddd with ||Al|] S 1 and G ~ N (/2)- Left right: Behavior of (a) the norms of intermediate activation (()) with index f, (b) V. (definition 5.2) for each stack of layers FE. and (c) 1 El= (x(z)/ [[())] that appears in our bounds in theorem 5.3.
C ADDITIONAL RELATED WORKS
Residual networks ensembles. Training with random paths can be weakly viewed ensem- bles over paths, and the connection between ResNets (He et al., 2016) and ensembles of shallow networks was first point out in (Veit al., 2016) for vision models. (Dong et al., 2021) showed the same phenomenon for self-attention models, where longer paths lead to rank-collapse in the self-attention module. Chang et al. (2023) study vision transformers (Dosovitskiy et al., 2020) as cascade of multiple paths and propose pruning and self-distillation to remove long paths and im- prove performance. All these works mainly provide a novel perspective or inference efficiency but do not focus training efficiency.
Learnable scaling of residual blocks. Bachlechner et al. (2021), Zhang et al. (2019). Touvron et al. (2021) consider learnable scales the output of the residual blocks. These works aim understand favorable initialization of the learnalle scales under various constraints for faster training, which is very different from the role of scaling in our algorithm.
18

---

Published as a conference paper at ICLR 2025
Table Additional rows from table compare baseline runs two learning rate schedules, and run compare RAPTR and Gradual stacking with new 6-12-18-24 schedule. Key Observations: (a) Cosine decay learning works best for the baseline, compared to Square-root schedule proposed by Tay et al. (2022), (b) The performance of gradual stacking improves with the new schedule, implying a performance dependence with appropriate scheduling, and (c) RAPTR's performance doesn't I change much with the new schedule.
Rel. FLOPs
Eval Loss
Trivia QA
Tydi QA
SQuADv2
SGLUE
Avg.
Baseline (Squae noot LR decay)
1.2
2,07
23.4
31.9
44.3
60.0
39,9
Baseline (Cosine LR decay)
1.2
2.06
25.0
34.4
42.1
60.0
40.4
6-12-18 18-28 Stacking
1
2,06
22.1
34.6
38.0
60.5
38,8
12-16-20-24 Stacking
1
2.08
20.1
28.6
36.0
60.4
36.3
12-16-20-24 RAPTR
I
2.08
22.2
38.2
40.6
60.1
40.3
(+30k imttal full model trin)
1
2.06
25.8
36.7
44.1
60.9
41.9
6-12-18-24 RAPTR
(+30k ininal full-model tmin)
2.06
24.2
42.3
61.1
41.2
Early exit for efficient inference ficiency for large language models. (Lei et al., 2023; Tay et al., 2022; Del Corro al., 2023; Xin et al., 2020; Zhou et al., 2020; Hou et al., 2020). However, none of these works focus efficient fine-tuning. Lei et al. (2023) modified pre-training by substituting resource-intensive MLP compu- tations with straightforward classifiers for predetermined fraction of tokens within the sequence. It's worth noting that their primary focus was on accelerating inference time rather than enhancing the efficiency of the pre- pee-traiiing process.
A lot of recent works have focused on improving inference ef-
Rel. FLOPs
Eval Loss
Trivia QA
Tydi QA
SQuADv2
SGLUE
Avg.
Baseline
1.2
2.06
25.0
34.4
60.0
40.4
Equiflop Baseline
1
2.07
24.3
35.7
42.4
60.6
40.7
RaPTr
1
2.06
25.8
36.7
44.1
60.9
41.9
Table 7: Performance comparison of models with relative FLOPs and various evaluation metrics.
Train Loss
Eval Loss
10a
Train
Exx10
Eval
4x
4x
3x
3x127
2x
2x19
20000
40000
80000
100000
20000
60000 80000 100000
Training steps
Training steps
Figure 5: Train and Evaluation Loss behavior for BERT-BASE model trained with RAPTR for 100k steps. We have stages with 6-8-10-12 schedule (see section 4 for details). The boundaries 25k, 50k, and 75k. Key observation: the model's train and evaluation loss change smoothly across stage transitions, indicating stability of the model to subnetwork training.
D ADDITIONAL EXPERIMENT DETAILS
D.1 DETAILS ON PLD
The progressive layer dropping algorithm Zhang He (2020) is given in algorithm PLD employs time and depth schedule that determines the probability of dropping each block. The time schedule
19

---

Published as a conference paper at ICLR 2025
begins with zero probability of dropping each block. Then it increases this probability throughout the training process until it reaches a maximum of (1 - a), where the hyperparameter ā = 0.5. We vary ä depending the average FLOP perplexity want achieve in the end. With p. the FL.PPs reduce by - (p/2).
The depth schedule ensures that blocks located earlier in the model are dropped with lower prob- ability than those located later/deeper. An important hyperparameter of the depth schedule in PLD 7,s which controls the rate at which the probability of dropping layers increases. A higher value of 7yf leads to a quicker increase in the probability. (Zhang & He, 2020) set 7/ to 100 in their experiments.
D.2 DETAILS ON VARIANTS OF GRADUAL STACKING
Here, we give the details of the variants of gradual stacking, which we use to compare RAPTR against in table 2. We simply initiate the depth-wise growth operator of each method to compare to RAPTR.
bert2BERT (Chen et al., 2021) bert2BERT is stacking based strategy that aims to progressively grow BERT models along width and depth different stages of training. bert2BERT used progres- sive stacking (Gong et al., 2019) to grow the depth of bert models during the course of training. Here, the model's depth is doubled each transition phase (e.g. 3 layer BERT - layer BERT > 12 layer BERT).
MSG (Yao et al., 2023) MSG a stacking based strategy that aims to progressively grow BERT models along 4 dimensions: width, depth, number of attention heads, and MLP intermediate dimen- sion at different stages of training. The main recipe is a small warm-up during phase transitions, where the model is slowly allowed to grow from its previous configuration to its new one. This is done to preserve the loss during phase transitions. MSG uses gradual stacking to increase the depth of the model, where top few layers are repeated initialize new models at phase transitions.
LEMON (Wang et al., 2023b) LEMON modified gradual stacking to instead use consecutive stacking. Here, the new layers are put consecutive to the source layers, whose parameters are copied to initialize the new ones. For example. for schedule that grows model from 6 layers to 9 layers, we choose the top layers as source layers and insert new layers between layers (4, 5), (5, 6). (6.) repeating the parameters of layers 4, 5, 6 respectively.
MSG with layer interpolation We modify MSG's depth growth operator to LEMON's growth opertor. Here, instead of using the top few layers to stack on top of the existing model, we use consecutive stacking to put the newly initialized layers close to their source layers.
Remark: In order keep the comparisons uniform across the methods, we use the same learning rate schedule as RAPTR across all methods.
D.3 DETAILS FOR POLYNOMIAL TRAINING
We present in fig. the results from a 20-layer residual network trained on 100-dimensional data. Our f has max degree A: 10 restricted to the first f 20 variables (æ1:20) in the input. We pick m = 20 random basis polynomials for each degree f < 10. We set the average drop rate of PLD
0.6, such that the average FLOPs is 80% lower a baseline run. For RAPTR, we use 4-stage training, with the stages represented by its Z) as (8/20, {H), (12/20, (), (16/20, ()) and(20/20 220 ()) respectively. The lengths of each stage have been such that the average FLOPs is 80% lower to a baseline run (please see appendix E).
D.4 DETAILS FOR BERT
The batch-size is set to 256 and 512 for BERT- BERT-BASE and BERT-LARGE respectively. Following (Reddi et al., 2023), we use learning rate warmup of 10k steps and keep learning rate constant in all phases, followed by a linear learning rate decay to in the final full model training phase. We set the peak learning rate 10-1 in all our experiments.
Equal step Pretraining
20

---

Published as a conference paper at ICLR 2025
Table 8: Equal FLOPs comparisons for BERT-BASE and BERT-LARGE with extensive peak LR tuning. For BERT-BAS we use the best performing schedules from table 1. For BERT-LARGE, we use 6-12-18-24 for both RAPTR and stacking. FLOPs denotes the number of steps involved in baseline training. Key observations: for BERT-BASE (a) RAPTR achieves better loss compared to baseline at FLOP measures, with larger differences for fewer FLOPs. (b) RAPTR has 0.02 better loss than gradual stacking all FLOPs. (c) For BERT-LARGH stacking and RAPTR are competitive to each other at all FLOPs. Both methods have better loss than baseline at lower FLOPs.
Model
FLOPs
Baseline
Stacking
RAPTR
75k
2.09
2.02
2.01
BERT BASE
170k
1.90
1.88
1.86
510k
1.74
1.75
1.73
62.5k
1.84
1.78
1.80
BERT-LARGS
140k
1.63
1.60
1.61
625A:
1.40
1.41
1.41
Equivalent flop Pretraining To make fair comparison across all methods, we tune the peak the learning rate in the grid {1, 1.6, 2.5, 4} x 10-4 for BERT-BAR -ASE and grid {1. 1.6} x 10 for BERT-LARGE.
Finetuning We fine-tune BERT models for 3 epochs with batch-size 32, using a learning rate grid search over 2,, 3, 5} x 10-5, We use a linear warmup schedule over initial 6% training steps, followed by constant learning rate schedule for the rest of the training.
Algorithm 2: Progressive layer dropping (PLD) Zhang & He (2020)
1: Input: iterations T', layer keep probability ā, temperature budget y > O, layers L, functions (self-attention, layer-norm, feed-forward) ferTN,INN, frEN, loss function L, data (Xo. y), output layer fo
2: 7 7-
3: for ( f- to T do
4:
p (Keep probability.)
5:
a₁ + (1 - â) exp(-7 t) + ā
6:
L-a; Pa _ [Layer decay.)
7:
for - 0 to L 1- do
8:
Bernoullip)) (Keep drop.) if = then
9:
10:
X++1 x; (Drop.)
11:
else
x/ + x; (fTTN(fEN(x1))
12:
X;+1 + x + (PEN(SEN(s
13:
14:
end if
15:
p p p- pa (Decay prob.}
16:
end for
( (()o(x)))))
18:
fATTN, fLN, fEEN. fo + (pdate({)
end for
D.4.1 ADDITIONAL PRETRAINING EXPERIMENTS
Results at equal FLOPs. Inspired by (Kaddour et al., 2023), we further compare RAPTR and gradual stacking to baseline training by adjusting the number of training steps of baseline training to match its FLOPs to RAPTR (Table 8). For BERT-BASE, we observe that at shorter FLOP experiments, RAPTR achieves better validation loss compared to baseline training and gradual stacking. This difference gets smaller as we move to larger horizon settings. Similar findings for BERT-LARGE.
21

---

Published as a conference paper at ICLR 2025
D.5 DETAILS FOR UL2
Pretraining Similar to our experiments on BERT, we use linear warmup of 10K steps, followed by constant learning rate in all stages, and a cosine decay to 0.1x x the peak learning rate in the final full model training stage. The peak learning rate was searched over {10- 10-2) and was fixed to 10-2
3,5 x 10 2 x 10 10-22.150 2,3
We use Adafactor (Shazeer & Stern, 2018) optimizer and train with a batch size 512 for 400k: steps on mixture of Arxiv, C4 (Raffel et al., 2020), Github, and Wikipedia (Foundation) datasets, with mixing ratios 9%, 57%, 17%, 17% respectively (for total of 100B tokens). This roughly corresponds to 0.8 epochs of C4.
Reported datasets in table 3 include Trivia QA (Joshi et al., 2017). Tydi QA (Clark et al., 2020), SQuADV2 (Rajpurkar et al., 2018), and SuperGLUE (Wang et al., 2019). For QA tasks, we report Exact match scores, and report average accuracy for SupreGLUE tasks. To reduce variance, we report average performance across runs for the most representative setting for each method.
Boundaries for 12- 16-2--2 RAPTR and Gradual Stacking 40k, 120k, and 240k (decided using appendix E). The average subnetwork length used during training turns out to be 20 with this sched- ule. When we introduce initial 30k full-model training for RAPTR, we shift the boundaries by 30k so that the average subnetwork length stays 20.
D.5.1 ADDITIONAL PRETRAINING EXPERIMENTS
In table 6, we have the following additional runs compared to table 3 in the main paper.
We first compare Baseline model with cosine decay learning rate to Baseline model with Square- root learning rate decay that was used in the original paper (Tay et al., 2022). We observe that the model with square-root learning rate performs significantly worse compared to the model with cosine decay learning rate. Hence,
model with cosine decay learning rate
baseline.
Secondly, we try another schedule for Gradual Stacking, which contains 4 stages and is represented by 6-12-18-24. The stage lengths have been adjusted such that the average number of layers used over training is 20. We set the stage boundaries as 17k, 74k, and 172k respectively. We observe that 6-12-18-24 Gradual Stacking performs much better than 12- 12-16-200 16-20-244 Gradual Stacking various downstream tasks. This signals towards the dependency of Gradual Stacking on proper schedule selection.
On the other hand RAPTR doesn't 't show a big difference with 6-12-18-24 schedule, when compared 12-16-20-24 schedule. proper analysis of both the methods on schedule dependencies is kept for future work.
E SCHEDULE SELECTION
We follow two schedules from Reddi al. (2023), Equal and Proportional, to build the lengths of each stage in RAPTR and Stacking. For Equal scheduling, used in Tables and 8, we split training into k equal stages. For Proportional scheduling Table 3. we increase the length of a stage in proportion to index of the stage. For example, for 12-16-20-24 RAPTR and Stacking with proportional scheduling for 400% training steps, the stage lengths are set as 40k, 80k, and 120k: respectively.
The flop counts are decided on the basis of average subnetwork length during training. For example, if the average subnetwork length is 18 for model with 24 layers, we consider relative flops of RAPTR compared to baseline as 0.75.
For 12-16-20-24 RAPTR with Equal schedule, the average subnetwork length during training is given by 18, while for Proportional schedule, the average subnetwork length during training is given by 20. However, when either Equal and Proportional schedules don't t return the target average subnetwork length, we re-consider the stage lengths by taking the closest candidate among Equal and Proportional schedules and remove equal steps from each of the first le 1 stages and add to the final full model training stage. For example, in table 6, with 6-12-18-22 RAPTR and target average subnetwork length 20, Proportional schedule returns average subnetwork length of 18 and stage
22

---

Published as a conference paper at ICLR 2025
lengths 40k, 80k, 120k and 160k respectively. We then aim to find such that stage lengths of 40k I, 80k 2, 120k I and 160k + 3z returns the desired average subnetwork length. On solving for get = 22k.
E. ABLATION ON DIFFERENT SCHEDULES FOR BERT-BASE
We conduct experiments on RAPTR with different stage schedules that have same relative flop counts in table We observe on BERT- BASE that different schedules differ by atmost 0,01 in pre-training loss and 0.1 0.5% in downstream performance average. Thus, we conclude that RAPTR robust to the stage schedule selection.
the following stage schedules minimize computational overhead. We restrict the number of phases in RAPTR to 4. The first stage selects a sub-network with half of the layers at random each step. We increase the sub-network size by equal quantity in each stage transition, such that we train the full model in the final phase.
For UL2,
Table 9: Performance of RAPTR on BERT-BASE with different schedules. The average subnet- work size for each schedule is 9, which makes each schedule 1. 1.33x faster than the baseline model (provided for reference). We use the same hyperparameters as used in table 1. We observe minor differences between different schedules, indicating robustness of RAPTR.
RAPTR schedule
Rel. flops
Eval loss
MNLI
QNLI
SST-2
Avg.
3-6-9-12
1
1.76
82.0
89.6
92.0
87.9
4-8-12
1
1.76
81.9
89.3
91.5
87.6
6-9-12
1
1.75
82.3
89.2
91.0
88.0
6-8-10-12
1
1.75
82.1
89.8
92.4
88.1
Baseline
1.33
1.76
81.5
90.5
91.4
87.8
E.2 ABLATION ON FIXED LAYER SET FOR BERT-BASH
We conduct experiments 6-8- 10- 12 RAPTR BERT-BASH with different fixed layer set I (ta- ble 10). We observe that fixed set selection of {1, 12) performs the best among other candidates. This selection however can be expensive verify for other settings like UL2. Hence, in all subse- quent experiments, we continue with the first and last layer fixed in all stages of RAPTR.
Table 10: Performance of RAPTR on BERT-BASE with different I set for 6-8- 10-12 RAPTR run for 100k steps. We use the same hyperparameters used in table 1, except the training steps reduced to 100k steps. We observe that fixing the first and the last layer at all times lead to slightly better performance, compared to other candidates.
I set
Eval loss
MNLI
QNLI
SST-2
Avg.
2.12
(1)
2.13
[1,
2.11
78.6
86.77
88.69
84.7
{1,2, 12]
2.13
11,
2.12
77.4
87.0
89.4
84.6
11,
2.14
78.0
86.1
88.8
84.3
1,2,3,11,18
2.18
77.4
86.5
88.4
84.1
21, 10, 12)
2.16
F SYSTEM SPEED OF DIFFERENT SURNETWORKS
We report the speed of training in steps/sec for UL2 models with different subnetwork sizes a TPU 16x16 chips with percore batch size of table 11. Suppose is the total time for 12-16-20-24 RAPTR that has
400k steps with 24-la model in UL.2 training. In table
23

---

Published as a conference paper at ICLR 2025
20% "theoretical" flop count reduction, with the stage lengths given by 40k, 80k, 120k, and 240%: respectively (appendix E). Theoretically, the necessary time required be 0.83t. Using the training speed numbers of table 11, it takes approximately 0.84t complete 400k steps with RAPTR machine, close to the theoretical training time.
Table 11: Training Steps/sec at different stages on 24-layer UL2
subnetwork size
steps/sec
12
2.0
16
1.6
20
1.4
24
1.1
We trained our BERT-base models on 4 x TPU cores with percore batch size of We report the speed of training in steps/sec with different subnetwork sizes.
Table 12: Training Steps/sec different stages BERT-base
subnetwork size
steps/sec
6
17.9
8
14.3
10
11.8
12
10.2
So. looking at table 8, let be the total time for completing 675k steps in BERT-base training. In the case of a 66-8-10-12 RAPTR with Equal stage lengths, it takes approximately 0.79t to complete 675k steps
machine. For ideal machine that perfectly realizes the theoretical speedups based on the desired flop count reduction, the total time required will be 0.75t.
G EFFICIENT IMPLEMENTATION OF RAPTR IN DISTRIBUTED TRAINING
Issues with distributed training: RAPTR can be naively implemented with if-else statements that bypass layer if layer isn't 't part of a random sub-network specific time step. This will work with pytorch if we are training with a single GPU. However, if we train a model on multiple GPUs with pytorch, we get the following error.
Expected have finished reduction in the prior iteration before starting a new This error indicates that your module has parameters that were not used in producing loss. You can enable unused parameter detection by passing the keyword argument 'find unused parameters=True to 'torch.nn nn. waw.laalllll.aiiiii... Dww..iiiiiiibbbttaaaaaarrea participate in calculating loss. Parameter indices which did not receive grad for rank I: 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 72 73 74 75 76 78 80 82
RuntimeEEroor:
and by making sure all 'forward' function outputs
85 87 88 89 90 92 93 94 95 96 97 98 99 100
101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123
124 125 126 127 128 129 130 131 132...
We get this error, because the optimizer expects gradients from each trainable model parameter from each GPU at each step of training. We get similar error, when we use jax framework to implement RAPTR.
Solution Layer Selection as Linear Algebraic formulation on Parameters: Instead of view- ing RAPTR as selecting random sub-network composed of a random selection of layers, we should instead view sub-networks as a set of selected parameters in the Transformer model. All Transformer layers are identical in structure but with different parameters.
(3) ((T)) represent the L Transformer layers with corresponding parameters ,.. ((1), ((2) ... O(2)) Each of these layers be expressed as common Transformer layer where
Let f(1), f(2),
24

---

Published as a conference paper at ICLR 2025
Sh. 1
Sh. 2
O
0(1)
((2)
sh
0 = c1,10(1) + 2122 + c1,30(3)
f(1)
f(2)
f(3)
00) ah
c2,10 + 22 + c2
g())
g(2)
g(3)
(b) Shared base sub-network across possible 2-layer sub-networks
(a) A representative 3-layer network
(c1,1, C1,,, C1,3, (c2,1, C2,2, C2,3)
f(1)
f(2)
O (1,0,))
f(1)
f(3)
O
(1,0,0)((0,0,1))
f(2)
f(3
(0, (0,1,0), (0,1,0),(0,0,1)
(c) Representing each 2-layer sub-network in the shared base sub-network
Figure 6: Representative overview of efficient implementation of RAPTR in distributed training setting. A simple 'if-else' implementation of dropping layers doesn't work in distributed setting, as the optimizer expects gradients from each parameter from each GPU at each step of training. We instead represent all possible random sub-networks as shared sub-network, whose parameters are represented as linear combination of all parameters. Then, we can back-propagate through random sub-network by setting the linear coefficients appropriately.
25

---

Published as a conference paper at ICLR 2025
the parameters (()) ())))))....
((()) are used to represent f()) f(2), f()). f())
f(()) respectively. Fur-
thermore, each layer f* can be represented with parameters Ej=1 9906), where Cu = 1 and Cij for all j i.
We can extend this idea to a sub-network that selects k layers at random from the Transformer layers. We can maintain k-layer shared sub-network, where shared layer 1 S l S k uses parameter g(t)
sh = E Ct Ct,j are fixed depending on the random sub-network of layers that we pick at each step of training. For example, for a sub-network that picks layers 1, 2, 3, ..... . k, Cg is set as for i é for all 1 S & S k and 0 otherwise. In doing all the parameters get gradients in the loss computation at each step of training.
Additional FLOPs involved in sub-network selection: Because have additional parameter additions to represent sub-networks at each step of training, one might wonder whether the addi- tional FLOPs incurred will lead to drastic slowdown in model training. To do so, will compare the number of FLOPs necessary for backpropagation for a k-random sub-network to the number of FLOPs necessary to compute gradients for k-layer static sub-network used in stacking. The backpropagation step includes a forward pass, backward pass, and gradient descent step.
To further simplify our explanation, we will focus on an L-layer network that uses linear layers as its Transformer layers with parameters W W(L) Raxad We will foeus Gradient Descent algorithm for simplicity. Batch size is given by B.
Back- preppapopaggtion FLOP Computation for k-layer static sub-network: For simplicity, we will use the first l layer parameters to define the static sub-network. We will refer to the output of each layer l as ( E Rdxd
Forward pass: For input (() E IR Bxd, the forward pass involves computing W (
for each y E ( for each layer 8 E [k:]. Thus, the number of FL.OPs involved is Bdd.
Backward pass: On gradient G E Bxd bao--proggg to layer f, the backward pass computes the gradient w.r.t. the parameter W and the gradient that flows to the next layer.
B The gradient w.r.t. the weight is given as 1 E g (o= This involves Bdd compu- tation, and overall a kBd computation for all weights.
Furthermore, the gradient to be back propapropag beyond layer i given as (E) (W(1) which incurs additional Bdd com- g (2) (W(O)), g; (w((O) .... 9B
putation. Overall, BBd computation across all layers.
* Gradient Descent step: Finally, the descent step is simple addition of negative gradients to each weight, which requires kd for all weights.
Back-propagation FLOP Computation for k-layer random sub-network: For randomly picked sub-network, we will use Σイ ce, W to denote the parameters of the l-layer shared sub-network that we use to perform back propaaappoaa at particular step.
Forward pass: Our proposed method requires an additional sum of weights in each layer, with the coefficients ce,j for S k 1 S S L in the forward pass. Hence, the total FLOOP computation involved is l: Bd + kLd.
Backward pass: Here, there are no changes for backpropagation through the k-layer shared network, as the gradients are computed and backpropagated w.r.t. (
Gradient Descent step: Here, the gradients need to be distributed to the weight parameters ((O) before taking the descent step.
VW
Σ CiE 1=1
(
This requires kLd computation, in addition to the necessary gradient descent steps.
26

---

Published as a conference paper at ICLR 2025
Thus, the relative FLOPs of backpropagating through k-layer random sub-network to k-layer static sub-network (used in stacking) is
2kBd kdl 2kLd
2L
+ 2B +
2k Bdd + kd?
With B set high, this relative FLOPs gap is close
Discussion for language models: For language models, the batch size represents the number of tokens that pass through the layer during hackpropagation. our UL.2 experiments, sequences were of length 512 and the number of sequences in each minibatch was set as 2. Hence, the effective batch size 512 x 2 1024. The number layers FLOPs of backpropagating through k-layer random network compared to k-layer static network (used stacking) would be = (48/(2048 + linear network setting.
our UL2 setting is 24. Hence, the relative
1) + 1) = 1.02 if we had used the same minibatch for
Remark: When training large models, we can structure the selection of random sub- sub-networks to minimize the performance overhead associated with random sampling. For instance, if we train- ing a 100-layer network and wish to train random sub-network of 50 layers at given step, rather than selecting 50 layers arbitrarily from the entire model, each containing 10 layers. Then, we create a 50-layer shared sub-network, where each shared layer linear combination of 10 layer parameters from specific group. This approach reduces the computational overhead associated with combining parameters from 50 10 x d to x 10 x d'
organize the layers into 10 groups,
H NOISE STABILITY FOR TRANSFORMERS
Theorem H.1. Suppose L satisfies the following condition L satisfies the following condition for any input ă. I.' and label v 1), for some constant Ploss
C(E + 7) || v) C(ă, v) S Pe ||
(3)
The difference between expected loss of (L R)-RR and L-RAPTR. i.e. - - (F)),. is upper bounded by DE properties of the head H on top of the backbone F.
(=())) (E
2 some constant C that depends on regularity
Proof of theorem 5.3. Denote by F-t the candidates that we can randomly choose from during L L- RAPTR training. The output of F_g on any input is equivalent to the output of the model F on æ and a1.L, with aę = and the rest set to 1.
Then, the average loss of L random training is given by
(2(F) i=1
Eny-nD pf (F-r(x), (
On the other hand, the loss of the full model is given by
(F) EL (F(m))
Hence, the expected difference between loss of L random training and full model given by
L
-2(F)-C((F
Σ
E BRDDDD (C (CFF)((z))) -
1.
(x) F(x)ll2
<
En
BR
Pc
F(x)
the PL
E
R~D
1=1
Here the pinafinal step uses the definition of fg from eq. (3), and the final step uses the definition of Ve from definition 5.2.
27

---

Published as a conference paper at ICLR 2025
Discussion An example of loss function with bounded mag is transformer model, that uses layer normalization layer before using an embedding matrix & to compute the logits. The logits are then compared with the true label with cross-entropy loss with softmax function on logits. (CC hence depends on the l norm & and the weight parameters of the layer normalization layer.
H.1 NOISE STABILITY OF LINEAR RESIDUAL NETWORKS
H.1.1 USEFUL LEMMAS
To prove, we will require the following theorem (and variaions). For random initialization, we use W1, WL ~0, ).
Lemma H.2 (Norm of the output of linear layers at initialization). Consider the case of a linear network with no residual connections and layernorm, i.e.
y () = (t 1), for all & 2 1,
with (0)) a.
For impul any E O, with probability at least Le -o(E6(1)) over the randomness of ( E e1-
W1. WL. we have
for alll 2
Proof. Because W, is gaussian, (+)) Wyy follows (0 (1 /d). Thus, {V) (0) )
follow Chi-squared distribution, with E || = E E'= [y] ))2 given by (| The proof then follows from applying concentration bounds on the norm of gaussian vectors (Formula 3.7, Vershynin (2020)), i.e. for any & 2 O,
Pr y 03
11 2 1eeeee/01
se -l(dt))
We apply a union bound over all layers, to get the final bound.
Corollary H.3 (Norm of the output of linear layers with residual connections initialization). Con- sider the case of a linear network with residual connections but no layernorm, i.e.
y(t) ((- Way (6-1), for all E 2 1,
with (0) x.
For an input x, any E 2 0, with probability at least Le -ode2) over the randomness of Wi.... WL we have
/ E
e)] for all & 2 1.
Proof. The proof follows similar to H.2. W.h.p., can show that for any 2 0.
Pr (W (e-
- 111 2 8) s -
The only remaining piece show that (f-11) and Why (f-1) roughly orthogonal. First of all, (t)) (t-(t-1)
because We is concentrated around 0. Furthermore, = Eig w L (2 (f) = Y ((-)) Thus, using Hanson wright inequality (theorem 6.2.1. Vershynin (2020),, we have w.p. atleast
W (ty ((()(
Pr (e-1) T
Wey (e-1) 2 2ell
(- (o(ds))
We apply a union bound over all layers, to get the final bound.
28

---

Published as a conference paper at ICLR 2025
Corollary H.4 (Norm of the output of linear layers with residual connections initialization). Con- sider the case of a linear network with residual connections and nnrer-noomaiizztionn i.e.
y (e- Wey! //
//00>////>>>>/>>>/ (t-))) for alll 2 1,
with y0) a.
For an input æ, any E 2 O, with probability at least Le -12)) over the randomness of WI. WL. la we have
(2) 2
E ()
for all 2 1.
Proof. The proof stays extremely similar to H.3. For simplicity, we simply refer to H.3 for any such variations.
H.1.2 PROOFS OF LEMMA 5.4
Proof of lemma 5.4. We outline the proof for each case.
1. With residual connection and layer normalization, the function F computes the intermedi- ate activations (1), (L) on input follows.
y0) = (
W.
(e--))
By corollary H.4, we can show w.h.p. with respect to the randomness of Wi,"! WL, for all & < L, v((1 O(1)) < ( S v((1 + +1)).
Ignoring the error terms, can simply write
y (
We
y (e-1)
Pick general layer e to be dropped. Then,
F-((z)- - F(x)= = II 6=t+1
L
1
We
We
(e-1)) +er,, Ve
where the err term appears because of the change in scales of activations ((++)), (() with the dropping of layer L. This error can be bounded O(VL/() using the same procedure followed below to bound the first term on R.H.S..
Similar bounding the norms of y (f), w.h.p. we can show that
I
L
We Ve
We
(t-1)
S/2
L
1
ve
#'=++1
2
Hence,
14 = [F Fe(x))l2 S SO(VT.
This implies
L Σ we(x) 1 -
We then apply union over all E S such that the above condition holds true for all inputs in set S. Since the gap in L2 and L1 bounded by O Σe (F(=)/|[((( from
theorem 5.3, we have the gap as O
29

---

Published as a conference paper at ICLR 2025
2. With normalization, the function F looks follows.
(I Wz) y ((-))
By corollary H.3, we can show w.h.p. with respect to the randomness of WT, ". Wr, for all & L, 2 (()) & | 21+
With drop in layer, we get
L II (1+ 1=
[-
We II I we) x.
F-e(x)) F(x:)
Similar to corollary H.3, we can show w.h.p. with respect to the randomness of
W1,' WL. Te=411 [I+ w We We 8-1 I W
=|| > 0(2(--1/2).
This implies [|F-((x)- F(x))12 (22(-11)/2
3. Without residual connection, the function looks as follows.
y( = W (e-1)
Similar corollary H.4, show w.h.p. with respect to the randomness of
W1.... WL, for all & < L, (1 O(1)) I ( ll s (1 + O(1)). Thus, ignoring
error terms, the network F roughly looks like
( Wey (e-1)
With a drop in layer, we get
1-1
-f(x))
II w (W 1) II
Wea.
8=t-1
d=1
Using randomness of ( this can be shown to be of norm !(11).
Lemma H.5.
When the layers are perfectly aligned, i.e. all weights W, A for some matrix A
with || = 1 and for simplicity, assume second eigenvalue Xg (A) = - 8 for some 8 >> 0, we have
(a) With residual connection & layernorm, have Ve(x:) = O(1) and [[(x)] ((L)). Thus, the gap losses between stages is O(1LL).
(b) Without residual connection, we have both (2(()) 0 and || [())]] = O(1). Thus the gap in losses between stages O2+))
(c) Without layernorm, we have (x(x)) 0((2-11 and [[()]]] 0 Thus the gap in losses between stages can be ((1).
Proof of lemma H.5. We outline the proof for each case.
1. With residual connection & layernorm, the function F computes the intermediate activa- tions (1),.... (L() on input follows.
(
I
y ((1-1)
A
30

---

Published as a conference paper at ICLR 2025
We show that the above method will behave like power method, i... y will be of mag- nitude (2()) and will be E-close in angle to the top eigenvector of A, denoted as 82 (A)), provided e 2 (((1/8) log(1/s)).
We give an inductive argument as follows. Consider any layer É. Say 0g denotes the angle between ( and di (A). Also, say I1 denotes an orthogonal projection to subspace spanned by the rest of the eigenvectors of A. Then,
(() y (c+)))
|1 ( (c(++)))] Lan 0++1|
(A), y (0) (v) Ay | (v)()) + 11 Ay
2 (A),, y
|1 (1) 11 Ay
2(01 (A), (+)
[T b1 (y)
A2 (A) y
2
ttan 0el
1 (A)
This implies, under (A) 1, ttan 0g] decreases with é. Under the assumption that isn't torthogonal to 0ı (4), the above condition simply shows that if 2 log(1/))) than [tan Oel <
Furthermore, once aligned (or closely aligned) 01 (A), the norm of (() grows linearly. Hence, (0/ S2(L). Furthermore, for any L, the gap between F and Fg simply becomes equal to the gap between L and L - steps of the modified power method, which can be bounded as (1).
2. With no residual connection, the function F computes the intermediate activations (1), (I),.... .... .y(() on an input æ follows.
A
(0)
From the update, it's trivial to come the conclusion that (8) will stay unit norm for any layer L.
This update is exactly power-method, where ( gets s-close to the top eigenvector di- rection of A in ((1/6) log(1/e)) steps. The result then follows from bounding the gap between L and L 1 steps of power-method.
3. The proof is very similar to the proof of lemma 5.4 (3), where the major difference comes from the blowup the norms the intermediate activations.
I RAPTR MOTIVATION WITH BOOLEAN DATA: THEORY
Data and labels: We use uniformly sampled boolean data (({+1))) The true label function ft Rdd + for each data is given as
V3 2
v3 I1
(x))
T1
2
Network to train: We train a 2-layer residual network fno. ur..h. fn Rd - R comprised of two single neuron residual blocks, where Po behaves position bias that we add to the output
31

---

Published as a conference paper at ICLR 2025
of the first block. On an imput æ, the network computes output y E R with intermediate output y( 2y2 E IR as follows.
(a) po + (un, x² + b2),
y(2) = yا) + sin( (w2, x) + (1)) + 62).
(4)
Here = y(8) is returned as the final output. The network is trained with mean squared error, given by
-= (y We consider population gradients for the sake of clean exposition of the important ideas. Interested readers can modify the proof to work for finite sample regimes.
(5)
Initialization of weights and biases: The elements of the weights and biases wi, W22 Po, bi, b2 have been initialized from all Os.
RAPTR Initial position bias training: In the initial bias phase training, simply train the bias Po. i.e. at each step, we train with the following mean squared loss,
L= (=))=((()))))( - f
(6)
RAPTR First phase training: In the first phase, both the layers are trained independently. That is, at each step, we pick random layer 8 E [1,2} and train with the mean squared loss, given by
C = Eg ue (po + sin((we, a) + be) - (x:))²
RAPTR Second phase training: In the second phase of training, we train the full model (eq. (4)) with the mean squared loss (eq. (5)). For simplicity, fix the parameters of the first layer and simply train the parameters W2 b2.
We first recall the theorem corresponding to the stage-wise training of RAPTR. Theorem L.I. For a small enough learning rate 1 < polytdy under parameter initialization from all Os, an initial position bias only training, followed by 2-stage RAPTR on the network fo wn, w. b2 (eq. (4)) with mean squared error shows the following behavior.
After (1//r) steps of position bias only training and ((1/n) steps of first phase, for both layers E (1,2),
v3
v3
po sin((we, x) + bb) 2
2 T1
< 0())
After ((1/n) steps of second phase, the ouput of the second layer is given as
120) - (2) (-2122) : (in(( x) + y + Po + b2) (- (-z102) < O(1).
Hence, after ((1/n) steps of initial position bias and 2-phase RAPTR training, loss L 0()) (5).
Proof. The proof proceeds by analyzing the behavior of the weights and the biases in the different phases.
. In the initial position bias only training, we show that in @(1/n) steps.
po
v3/2 < O(n).
In the first phase, after On log(d/m)) steps, show in lemma I.4 that for any layer e {1, 2}. the weights and the biases satisfy the following conditions.
WE11
S 0())
3
Jee,
S On)). for all 2 2 s d,
lbel < O(n)
32

---

Published as a conference paper at ICLR 2025
Hence, after first phase of training, the network's intermediate outputs y() and yl (from eq. (4)) are given as
v3
v3
y (1) po + sin (1F) a) + b1)
+
2 21 + 0()),
2
y12) = - sin( s1, x.) +1 + + b2)
v3 v3 x1 + sin (w2, x) v3/2n1 + v3/2+ + b2) O()).
2 +
2
In the second phase, after (11/)) steps, we show in lemma 1.5 that the weights and the biases W2, by satisfy the following conditions.
22.1 + v3/2-
- O(n).
2
W2,2 2 O()),
62 po 2 s O()).
|22 O(i), for 3 $ d.
Thus, at the end of the second phase, the network's output y y(2) (from eq. (4)) given
as
y y
sin( (n,2 x) + (1) + po + bz)
v3
v3
D1 + sin (w x) + v3/2n1 + v3/2+ b2) O(1)
2
2
v3
v3
π
π + ~0 + 0
v3
v3
+ sin(
(
+
21 + O(r).
2
2
2
2
Furthermore, we can show that each layer in itself isn't expressive enough to represent the true label function.
Theorem 1.2. For each luyer when trained independently the best we be parameters will have @(1) error.
Proof. The proof the theorem is simple. We argue for the first layer (same argument holds for the second layer). We simply caleulate the projection of the function learned by wi d into the components of the true label function. Suppose the function define dby them is given by aæ + 8x1Z2
Then, using orthogonality of the basis polynomials x and æ1æ2, we have
= En (+ 1 ((wi, a) + b)æ1
E= sin (T/2 (wi, .²) b)r1
( cos w,) sin(w1).
Using corollary I.9
8 = Exe(+1)a sin( (in x) + b)
E Ex sin(T//2 (w) x)) b) 22
cos w.) sin(wi) sin(wz). - cos(b)( II
Using corollary I.9
112
33

---

Published as a conference paper at ICLR 2025
Adding/Subrracting these terms, get
sin (b+ + w2)( II cos w;),
1
a + sin(b w2) II
cos w:),
1,2
+ 3 = sin(-b+ + w2)(
II 1
COs w/);
sin(-b- w2)( (
II cos wt)
1
Thus, we must be restricted by the set |a|. < 1, |o|| label function don't lie in this set, the proof follows.
|8]
<
As the coefficients of the true
1.1 ANALYSIS OF THE DIFFERENT PHASES
Initial position bias only training linearly increases Po to v3/2. Lemma L3. [Position bias only training] For any learning rate nj > 0, after @(1/n) steps, Po satisfies
po v3/2 & O(n)-
The first lemma analyses the behavior of the weights during phase 1.
Lemma L.4. [Phase of RAPTR] For a small enough learning rate n < polytay with position bias Po v3/2+ + O()). and under parameter from all Os, after ((L/n) steps. for any layer é {1,2), the weights and the biases satisfy the following conditions.
We.1
S O(()).
3
[e,l S O(n), for all 2 2sisd,
lbe S O()).
The second lemma analyses the behavior of the second layer weights in phase 2.
Lemma I.5. [Phase of RAPTR]
small enough learning rate 1j < polvtay starting from
1
the weights and biases reached by @(1/n) steps of position bias only and first phase of RAPTR training (lemmas 1.3 and 1.4), after ((1/n) steps of second phase of RAPTR training, the weights and biases w2 and b2 satisfy the following conditions.
W2,1 v3/2
S O(1),
W2.2
7 2
$ <O)))
62 + v3/2
π
< O()),
2
[22.1 S {(n)), for 3 d.dd.
I.1.1 PROOF LEMMA I.3
Lemma L.3, [Position bias only training] For any learning rate 17 > 0, after O(1/n) steps, po satisfies
Po v3/2 < O()).
Proof. Under all Os initialization, the output of both of the layers is 0. Thus, the gradient of is given by
V3
v3
v3
(-1,114
Po
I1 +
11
Po
2
2
2
Hence, Po increases up until reaches v3/2+ O(n)-
34

---

Published as a conference paper at ICLR 2025
1.1.2 PROOF OF LEMMA I.4
Lemma 14. [Phase of RAPTR] For a small enough learning rate nj polyld) with position bias
Po v3/2- + O(n), and under parameter from all Os, after ((1/n) steps, for any layer é e {1,2}, the weights and the biases satisfy the following conditions.
11 3 5
[e, - O(1), for all 2 223ss0 byl $ On))
Proof. Similar the proof of lemma I.3, we can show that the position bias Po stays O()) close to v3/2. Thus, for simplicity exposition, we simply assume it without formally proving so.
With the first stage of RAPTR, we only keep one of the two layers during training. Hence, both the layer parameters train similarly during training. Without loss of generality, we discuss for a single layer weights É E {1,2}.
General formulation for gradient w.r.t weights:
The population gradients of wy is given by
(=- (po + sin sin((we, + bb)
v3
v3
2 + 21
cos((we, + bb)
2
2
V33 2122
(B sin((ws, x)) be) cos((wo x)) + be) Banl-1,41d
2 21
cos((wer, x) be) x
v3 cos cos(((u,, x)) by) w.
E Po
We will argue about the contributions of the three terms separately.
Lemma 1.6 (Contribution of Term 1). For any coordinate j E [d], sin( .r) + be) cos((wee xx) be)ar
E
2 cos()bb)
(20(2 sin(2wes))
Proof. Term in the gradients is given as sin((w: (w,, 2²) + be) cos((wos x) + bb) x sin(2(w,, x) + 2br) xc
Bend-11+tn sin
1 Ex
2
Consider coordinate j E [d]. The gradient concerning the coordinate is given 2be )r,
E Ex-1-1,411 sin(2(w;; a)) 2
E Ban
1-1,111
sin 2
Σ
We,s + 22re, 2b xz
2
ii--disj
1
c2s())))
II
cos cos (2wes))) sin(2we,s):
2
1=1 --1-
using lemma 1.7.
35

---

Published as a conference paper at ICLR 2025
Contribution of Term 2: Term 2 in the gradients is given
V3 21 21
E 1 11,+12
cos( x²) + be)
Consider the coordinate i = 1. The gradient concerning the coordinate is given
V33 11 3132 cos((we,; v²) + be)x
2
v3 ( x) + bς) 1o-([-ittjo22)) cos( cos(( x)) + be)
Replace of by
2
v3
B (1 cos( (os x) + be) + sin(be)
II
cos(we,)) sin(w))2)
2
Simplified second term with lemma I.8
v3 (bs(
II cos(we,)) + sin(be)
cos we,))
sin (we,z)
d
Simplified first term with corollary 1.9
Consider the coordinate i 2. The gradient concerning the coordinate is given
V33 31 31 cos(( .²) be)ar
Ba (-1,t13
2
v3
Paw(-11t11 cos(( x)) + be) P= cos( (wre, x:) + be)
2
Replace 22 by
v3 cos(be)
II
cos(wr,))
II sin(wes) + sin(be) (1.2)
cos(we,))
sin(we,
2
im1->d,
Using corollary
For any other coordinate & {1,2), we have
(v3
Bx-(-1,t1]
21 cos((20,, x) + be)
2
v3
E
B (- cos(( (00s((e,,,)) æ)) be)xx1. 1-1 2222 cos( cos((wr,,,)) be)
cos(wee,) )
(
v3
2 cos(be)
II sin (wees)
JELL.
sin(be)
cos(we,))
II sin(we.s)
Using corollary I.9
(12,0)
(E(1,2.)
Contribution of Term 3:
For any coordinate [d].
v3 Po
E 11,+11
cos((uue, az) + be)ar
2
v3
Po
sin (b)
cos(we,)) sin(we,)
di
O(1),
36

---

Published as a conference paper at ICLR 2025
where the pre-final step follows from using corollary 1.9 and the final step follows since po
1/3
2
(1)) from the initial bias training phase.
Thus, the combination of the 3 terms gives
v3
Vwe,1 cbs(
II I-
cos cos(wee))
sin(be)
II
cos(w)r,)
sin(we,2)
2
1
cos (2be)
II i= i1,ii1
cos((
sin(2we,1) + (1))
(8)
2
v3
Vwi.2
cos (b;)
II
cos(wp, (
II sin(wes) -si((b))
cos(we,)) sin(wa,))
i=1 -nl- (1.2)
(E(1,2)
i=1-ad
1
cos (2
sin(2 + O())
1=1 12
Fort $ {1,2}.
v3
Vwer cos(b:)
cos(wes)
II sin(wes)
2
jE(1,))
+ sin(be)
cos(we,)
sin(we,))
(1.2))
JE(1,2,1)
+ (bee)
cos (2
sin (2we, + O(n)
(10)
General formulation for gradient w.r.t bias: population gradients of be given by
Following a similar approach as the weights, the
v3
v3
Ex-1-1,41 sin(( x) + bb)
I1 + x1
cos((we, x) + bt)
2 2
-Eas--1 sin((we, x)) be) cos( (ew(((w))))) x) + be)
v3
E
Tı - x112 (os((we,, x)) + be)
2
v3
+ B
1,+11
Do
cos( (oouuuu xr) + be)
2
Following similar approach as above, the terms can be simplified as
1
sin()be)
II V=Iー
cos(2we.))
(1I)
2
V3 sin(be)
+
II i=
cos(we,)
sin(we,)
2
- cos((be)
cos(we,+)
II sinwwes) + O(n).
(12)
1=1
(=1-d,ig(122)
(E(1.2)
37

---

Published as a conference paper at ICLR 2025
Behavior of gradients at initialization: from Os, we have for all j,
Since the coordinates of wy and biases by initialized
sin (bbe) O, sin(wer) 0, cos(we,)) 1, cos(be) 1.
(13)
Furthermore, recall that po = v /3/2 (1) after the initial bias training phase.
Using the above, simplify the gradients from eqs. (8) to (10) and (12)
v/
Vung,1 2 O()),
(14)
Vwg. O(), t+1 1 Vb = O(n),
(15)
(16)
Hence, we observe three kinds of behavior at initialization.
First coordinate the weight grows by atleast
W1,1 < W1 + 753
2 + O(1²).
2. Other coordinates of the weight get only O17²²) updates.
Wee + wey + (11²)
3. The magnitude of the bias drops as well.
be by + +
Beyond Initialization: Due increasing magnitude of w2,1, the behavior of the gradients changes slightly from Equations (14) to (16). However, assuming that the weight coordinates [2, d and the biases are smaller than O(r), we can still give similar formulations as before using eq. (13) using the following inequalities for all j & 1.
sin(b,) = On), sin(we,)) = -(m).
We state them directly.
v3
V (br(
cos(we,)) sin(be)
cos(we,)) sin (
2
i=1-d
i=1 122
+ cos(2be)
II
cos(2u)),))
sin(2we,.) +(n)
i=1-d,
sin(we.1) v3 +0(n) S 2 (sin(wz,1) v3 O (17)
1
(1 On)))
for 1) S 1/poly(d).
)
Vwi22 = v3
2 cos(br)
II
cos wee))
II sin(wes) -in(br)
cos(we,) sin(we.)
i=1=d (1,2)
(E(1,2)
1
+ 2
cos(2be)
cos( cos(2))).)
sin(2 + O(r)
+2
O())-
(18)
38

---

Published as a conference paper at ICLR 2025
For any & (1,
Vwt,5 = v3
cos(be)
cos(we,))
II sin(we)))
2
1=1 i
JE(L,0)
sin(be)
cos(we,s)
sin(we.)
i=l idl
(E(11.2))
cos((br)
II
(2w()
sin() + O().
2
dd
O()).
(19)
And finally for the bias,
v3 + sin(be)
bb = sin (26x)
II i=-
cos (2
II cos(wee)) sin(we,,)
2
2
- coos
II i=1 (1,2)
cos(wee.))
II sin(wes)
+ +0)))
(20)
(E(122
= O())
(21)
Thus, we observe three properties.
First coordinate of the weight grows until sin(we,1)
33 or 11,1 reaches n 7/3.
v3 sin(wee.) + O(7²).
WE1 + wes1 +
2 2
Thus, in ((Lnn) steps, w/,1 can reach arbitrarily close to x/3.
Any other coordinate + still receives O(1²) updates as We,t f- Wigt + O(1)). Hence, in @(1/n) steps, Wh,, can only reach (n) magnitude.
by also receives O(1)²) updates
by + by + O17?) Hence, in (1/n)) steps, be only reach O magnitude.
I.1.3 PROOF OF LEMMA 1.5
Lemma 1.5. [Phase of RAPTR] For small enough learning rate 1) -1 "pllyaay starting from the weights and biases reached by @(1/n) steps of position bias only and first phase of RAPTR training (lemmas 1.3 and 1.4), after ((1nn) steps of second phase of RAPTR training, the weights and biases wy and by satisfy the following conditions.
W2,1 v3/2
< 0(n).
W2.2
π 5
< 0()))
62 + v3/2
π
< 0()),
W22, s O(n), for 3 < j $ d.
Proof. The proof follows along similar lines as lemma I.4.
39

---

Published as a conference paper at ICLR 2025
First layer output: At the end of the first phase training, the weights w1, w2 and the biases look as follows. For E {1,2}.
Wx,1 - On),, |wee I O()), for all 2, |bdl - O()),
and po
/3
(n).
Hence, using a Taylor expansion, the output of the first layer is given as (for simplicity, we denote as (1))
v3
11) sin (wi, x)) + b1)
2 21 O(n)-
(22)
Second layer output: The output of the second layer is now given by (for simplicity, we denote as o (2)
0l²) = sin
(w2, x) + ( + po + bz)
sin ((w2, a)) + sin((w,, x²) b1) po b2)
v33 11 +O()) + +V33/2 b2
sin
(w2, x) +
v3
sin
(w2, x))
#1 + v3/2
+ by +0(1)
With Taylor expansion
2
sin
(W2,1 + v3/2)x1
W2j v3/2 +
+
For brevity, we will use new notations represent the coefficients of æ1, Fd in the above for- mulation.
ú = W2,1 + v3/2, ú W2.j; for 2
b = v3/2 b2.
At initialization, their corresponding values are
W1 3 + v3/2 0(1),
ú (n)) for all 222
v3
+ On). 2
(23)
Then, the above formulation is given as
(2)) sin
(10 x²) + b) O(n).
(24)
General formulation for population gradients:
The gradients w.r.. the weight wy and by are
given
V3
Vwz Ba (11t111 (1) + (2)
cos (10, m²) + 8) x + O(n).
2 .1 01
V3
Vb2 Egeny -1++|]
(11) (2)
01 +
11
(w, + 6) +O())
2
40

---

Published as a conference paper at ICLR 2025
We will consider the first two terms in the above gradient formulations and add the enror term later. First all, we observe that
v3 21 2212
(1)) + o
sin ( .²) +
+ b) + 21 + O(n).
We have the gradients w.r.t. u2 as
E Eg (-1, sin (no, æ) + b) cos (10 a) + b) x
2w2-
E (=- cos (τo, x) + b)
O(i).
We treat the two terms separately.
CONTRIBUTION TERM Using similar strategy as lemma 1.6, for any coordinate E [d].
Ex (-1,+11)) sin (tu .²) + b) cos (10 æ>) + b) xa
1 cos(28)
II
cos(2 sin(2:;:)
CONTRIBUTION OF TERM
For coordinate
E Ex [ 11,111122122 cos
(ww, x) + b) 22
P--(-1+1)aa cos
(i .²) b)
Replace 22 by
= sin(b)
cos ( sin(
Using corollary I.9
i= i
For coordinate 1,
Byn (11,+11441111 cos (10 x²) + b) ii
E cos
(w, x) + b) æ2
Replace of by
= sinnn
( sin((z))
Using corollary 1.9
i=1--d,i72
For any other coordinate & (1,2),
Ea-(-1+414 cos ((w) x) + b) 01
sin(6))
II
(wxz)
II sin(wiz)
Using corollary I.9
(E(12200)
Combination of all terms:
For coordinate =
Vw2,1 = cos(2
II d
sin(( c2s()
sin
II
cos(
sin( + O(r).
i=1=d.i42
(25)
41

---

Published as a conference paper at ICLR 2025
For coordinate 2,
Vw2,2 = c2s)
II 1-d
(20))
sin(2 sin(i)
cos(ws) sin(wi) + O(r))
(26)
For any other coordinate
& {1,2),
1
V Wwz,
cOs(28)
II
cos()
sin(2x),) - -in(
cos(wū;
II sin((:))
+ +0l))
2
i=1-d,i
i-1 i---di (1.24)
(1.2..4) (27)
We have similar computation for the gradient of the bias b2, which we report directly.
v3
V62
E E
1,+13a
( t0)
21 + 2122
cos ( x) + +b) + O(n)
2
sin(2b)
II l=l
cos(2u;)
-cos(b)
II (=1-)
cos( sin(w;n sin( + O()).
(28)
AT INITIALIZATION:
With i being initialzied at 2 (n) (from eq. (23)),
cos(b) cos( v3/2) + O(n) > 0
cos cos(2b) = cos( v3) + On) < <0
sin (b) sin( (in(v3/2) (()) > 0.
sin (2b) sin( (2nv(v)) O()) >
The initial value of in is
v3 2
On)) (from eq. (23)) and hence the values of cos(2 ) = cos( (os(v3 + 2n 33) + O(r) <
sin(2 sin(v3 2 2 0(n) < 0
cos (w) = cos( cos(v3/2 x/3) + 0(1) <
sin (1)) sin(v3/2 + m/3) O(n) > 0.
The other coordinates úy are of order (n)). Thus, the behavior of the weights and biases at initial- ization can be summarized as follows.
W2,1 gets a decreasing update as both the terms involved eq. (25) are positive.
1
Vu2,1 = 1/ cos(22)
II
cos()
sin((mn)
17 sin(b)
II cos(iū)) sin(wz)) + +0(1))
2
1=1-1
1=1
= =- 7 cos (V3) (1 ((2))
sinn
(sinnv33
/33 1/3+27/31 2x/3)
+0 S
" (oosv33)sse sin /3 22/3) -(77)
as 17 (1/ (1/poly(dd) small.
W2,2 gets an increasing update as both the terms involved eq. (26) are negative.
V2 17
(26)
II
c25() )
|sin(2 sinn
II
co:(
i=1--dii22
in1-
=nsin (33) sin(v3/2 m/3) ((n)) - +007) > " sin(33))n( 0
7) (11/poly(d) small.
42

---

Published as a conference paper at ICLR 2025
Other coordinates W2,1 for & {1,2} get small gradients as both the involved terms in eq. (27) depend on how large di = W22 = O(n) is.
V2,1 = ? cos(2b) 2
II iddiig
cos(22iz)
sin (2ūr))
+ s)n((())
II im im (1.2.))
cos (wss
sin(w:) (E(1.2.))
II
+ - -
Bias by get increasing update both the involved terms in eq. (28) negative.
(
| + 00000088)
V2 [sin((2))
II (20()) i=1-+d
cos(1ū;) sin(w1) [sin( + O(n)
2
(11-dd
2 sin( sin(v3) cos( /3 + 2n /3) (1 O1))(())))-))) + +()) sin v3) cos( v3 + 2TT33)
" 4
BEYOND INITIALIZATION: In fact, we can extend the above argument beyond initialization, where 12,1 decreases and W2,2, b2 increase by O(1) updates, up until reach the conditions
W1 2
O())
(or)
W2,1
v3/2 S O(n),
2
02 2
S (()). (or)
W2,2 2
$ (n))
i 1 < O(r), (or)
b2
π 2 + v3/2 < O(n).
Since, we get (1) updates in each step, this condition can be reached in (1/n)) steps.
The argument is as follows. For any 11 W2 satisfying the following ranges,
v3
E 2
3
82 (0, 5) be (v (V3/2. ")
the following conditions hold true.
cos(b), sin(b), sin(2b) > 0, cos(2b) < 0 cos(2201) ) sin(2w1), cos (m < 0, sin (wz)), cos( cos(w)), sin(2 >
sin(w1) > 0.
cos (20:) < 0.
We can then show that the terms involved in the update rule of w (eq. (25)) positive, implying -n < (()) up until úi reaches O()) close to म Similar argument can be given for the 2
updates of by and w2,2 respectively.
Furthermore, the coordinates W2,t for & & () get O17²) updates per step and stay O()) small for O(1/n) steps of training.
I.2
USEFUL LEMMAS
Lemma I.7.
Eg-(-1,41p sin ((w,x) + +b) x. cos(b)
cos(w,) sin(w,).
i= -dii
43

---

Published as a conference paper at ICLR 2025
Proof.
(=-(-1,+13 sin (w,xx) + b) xz
-B-(-1,+17 sin Σ
Wizt + Wyaz + bb zjj
ccos(b)
Ea-l-1141jast sin
Σ
Buy cos (wyar,) zj
i
Equals 0 as an odd function
+ cos(b)
B2 cos
Eug sin (wya2z)
i= i
Equals 0 as an odd function
E (- cos + sin(b)
Σ
Wazs
Ez, cos (wyzz) æy
i= diss
R B (-1,t sin sin(b)
Σ
Bary sin (wyary) mr
Equals 0 as an odd function
cos(b)
E Pmm(-1,41]- cos
wiss
Bay sin (wyrz) z.
(=1-d
cos(b)
cos(w:)
sin(w;)
Lemma 1.8.
Em-[- (11(- cos ((w, x) + b)b;; xj bb)
sin(b)
cos(wi) sin(w,).
i=1-d.11p
Proof.
E Ex (- cos ((w, x) + b)
π
(w, x)) b) xz
Em-(--1tt]a cos
2
= =Cos(T/22 12-22-11 - b)
cos(-
sin( sin(-w;))
Using lemma 1.7
sin
cos(w:) sin(w;).
jmm im1-dd,pp
Corollary 1.9. Consider S [d].
(II [I jES
E Eg -1,-1- cos (w) b) I[
II
cos(w:t
sinnww;)
3; Cb
JES
44

---

Published as a conference paper at ICLR 2025
Eval Loss
Eval Loss
5x
No Scale
1=
Linear Scale
I= I=
Sqrt Scale
I= (1,
1x
3x
20000
40000
20000
40000
Training steps
Training steps
Figure 7: Ablation study on choices for RAPTR algorithm. We train BERT-base with RAPTR for 100k steps with 10-12 schedule (see Section 4). Left to right: (a) Square-root scaling (harrt) has better validation loss compared with linear scaling and no scaling at training, especially in the earlier stages. With no scale, scale layer's output by 1/p, during inference (following Huang et al. (2016)). (b) Different candidates for the fixed set are compared for RAPTR. We find that training with first and last layers fixed helps faster training.
where
sin()), cos(b),
S| = 4t + 1 for some t. E N
i/|S 4l for some E N
C6
sin(b),
S = + 3 for some E
cos(b),
S|S = 41 for some 1 N.
SCALING
While working subnetworks, it is important to appropriately rescale the network. In particular, bypassing layer(s) in RAPTR can shift the input distribution for layers compared to the full model. To mitigate this shift, we scale the output of the layers to maintain the norms of the input each layer throughout training. The idea of scaling layer outputs has also been explored in prior work Huang al. (2016); Fan et al. (2019). We use a different square-root scaling mechanism that motivated by analysis randomly initialized Transformer architecture. We show that at initialization, the norm of the intermediate outputs ( scale as vt.
Theorem J.1 (Square-root scaling of input norms). At initialization, given any input sequence i..N and scalars 01:L, w.h.p. the intermediate sequences of (definition 2.1) satisfy
10 y |
2.112 + + azd), for all Iis
Inspired by this, we define the function hegrt that takes in sequence of binary values CI.L E (0, 13 and returns scalars Z1.z where 2, is the scaling for output of layer i. We find a scaling that satisfies the following two conditions hold (a) El=rs Si j for all < (maintaining the norm), and (b) C 0 0 & = 0 for all é < L (maintaining layers in random subnetwork). Formally, for each index j with S 1, it finds the index minimum j > j with Ej and sets ? = v - The RAPTR algorithm with square root scaling is presented in algorithm 1. In Figure 7 compare square root scaling with some other scaling methods and find it to be slightly better.
J.1 PROOFS
For simplicity, we present the results for transformers with single head in the self-attention layer. Furthermore, for simplicity, we use Grela activation in the MIPs.
Definition J.2. [Layer Normalization] Define a normalization function f Rd + Rd that performs f(x) (xc pa/a, where z and are the mean and standard deviation of æ. respectively. Then,
45

---

Published as a conference paper at ICLR 2025
Algorithm 3: Transformer Layer
Require: 2 layer normalization layers fat Impp (J.2), an MLP layer fmlp (J.4), and a softmax
self-attention layer fattn (I.3), input sequence { x RAN
jn=1 with yottnnnn = fLN - (xn) 2: Softmax self-attention; returns (yatt) N fatti (Ly m==)
1: Attention Layer normalization: returns yntttni V.
) for all n S N.
3: Residual connection: returns (y J.
1 n=1- with Nottmblock = In + y for all S N.
4: MLP Layer normalization: returns yn gmuppt
mipp (yn (ynttnbookk) for all N.
5: MLP function: returns {ymmiyy with yn = fmmp S for all n S N.
6: Compute Yn Vn tpl + Un Dittmblock for all N.
7: Return {Un Jn
layer normalization JLN, Rd - gd with parameters 7, b IR takes as input E IRd and outputs y E Rd, which computed as z = =(x), y = & =
Definition J.3 (Softmax self-attention). A self-attention layer fatm : RNxd _ R with pa- rameters (Wo, WK, Wv, Cattn E [Rdxd) takes a sequence (mn)agN and outputs sequence LynJneN such that
yv Cuttt
Σ
on Un; R=1
with Un,7 softmma (q Kqn)m. qn. WQ kn WK Un Wy
for all n S N, and K NxN defined with rows (kn) Definition J.4 (MLP). An MLP layer fmip Rd + R with parameters { E Rmxd Cmlpp E Rmxd] and activation drauus takes input x: E Rd and outputs y Rd, such that
y Cmil (eePoreeaWW)
Definition J.S (Transformer layer). A pre-layernomm transformer layer with three sub-layers; layer normalization layers LNN mmip (J.2), an MLP layer fmipp (J.4), and a softmax self-attention layer fottn (J.3); takes a sequence {# and outputs a sequence (yn in four steps.
1. First computes (yattmin) N n=1 using a layer normalization, i.e, yo gattnin SLN (In)) for all
n < N.
2. Then it runs softmax self-attention to get (y 1
Jn=1 = fattn ((ytm(((a)te (Lyatnin )
3. The net output of the self-attention block is given by attnblokkk Jn=1, with alttnblock
Vn + yo for all n N.
4. Before passing to the MLP function, is passed through another layer normalization, i.c. yo mipin (y atttnbookk) for all N.
5. MLP function then returns m N with ymm = (y for all S N.
6. yn ymmip y atttnblock for all S N.
Definition J.6 (Initialization of the weights in the transformer layer). The weights initialized as follows:
1
Cmlp, Chestn Wo, WK. Wv ~ NO
Vd ))
V2
W ~ N(O,
I).
Vm
The parameters y, b of the functions FLN fnlp have been initialized as 1 and 0 respectively.
Lemma J.7 (Norm of the output of the MLP function, modification of lemma 7.1 Allen-Zhu al. (2019)). For a given input sequence fy InSN, if E (0, 1], with probability at least 1 - O(N) e [(mm?) over the randomness of mmly, W, we have
Vi [N] ||
mipin ||
[1 E, +e).
46

---

Published as a conference paper at ICLR 2025
Lemma J.8 (Norm of the output of the layer normalization layers). Given any input sequence
// under the assumption for all N, In Lisi Ind isn't identically we have
yn = Vd
ofttinn
for all n S.N.
Similar result holds for (y p Jn=1
Lemma J.9 (Norm of the output of the MLP block). For given input sequence (ynttbblokk with probability al least 1 - (1) over the randomness of Crmpp, W, we have
JnSN.
1 11 = || + Yn yn 1nott = Bm atttmblook || +d
log N + (d + g (attnblock] 3/2
+
m
for all n S N.
Proof. Combining lemmas J.7 and J.8, we have for the sequence ayattblook } JnsNs
vlog N vmm
|| milp v|
w.p. atleast (()).
Furthermore, due to the randomness of Cully, we can further show that w.p. at least 1 O(1),
(y (yattmblock ymil s [ attnblock loggN)
for all n S N. Combining the two results, we have
ll 11? // yn w
y atttbblockk +
10g N
||0] log N + (d+ y gattnblock |]
+ o
Vm
Lemma J.10 (Norm of the output of the attention block). For given input sequence {ara (DrannNi if & E (0, 1), with probability at least O(1) over the randomness of Gutnn Wv, we have
Vlog N
Vi E [N]:
|
y
[ z ||
||x,|||
d 0
(d + ||
Vd
Proof. From the definitions of {y and (yattn), }. we have
yn attn
Ciattn Wv
Onsy attnlaa
JEN
Thus, we can use a proof similar to the proof of lemma J.9 to argue that with the randomness of Cottn and Wv,
|| 1y attn
Vlog.N N
50
ESSN On M pittnin
Vd
for all n S N w.p. atleast 1 - O(1). Furthermore, due to the randomness of Cattn we can further show that w.p. atleast O(1),
(y wn)! s O[1 log N),
for all S N.
47

---

Published as a conference paper at ICLR 2025
Using Souchy-Schward inequality, must have
n
gattnin
I max JSN
y || = Vd.
JSN
Thus, combining the results, we have
ym
attbblockk]"
1|²
1 g + xx11
||aen|1 + d
1og²/221 N
+
| log
+ (d + || Il)
Vd
Lemma J.11 (Norm of the output of the transformer layer). For a given input sequence {anInsM if's (0, 1}, with probability least 1 O11) over the randomness of Caltn Wv, Cmilp, we have
Vi E [N]
||> 2 + d O(err) S l1 S || 11 + 2d + O(err)
where err = (1 log N + (d+ + ||
log vm
Theorem J.1 (Square-root scaling input norms). initialization, given any input sequence 21:N and scalars Q1.L w.h.p. the intermediate sequences of F (definition 2.1) satisfy
1 yi
10,1 + + azd), for all 1
Proof. This follows from the fact that the transformer architecture is stack of struturally identical L transformer layers, and each transformer layer's output norm increases by @(d) compared to its input norm, as given by lemma 1.
48


Reviews:
[{"review_id": "eCkARa7PD5", "reviewer": "Reviewer_Dgxw", "summary": "This paper leverages the idea of stochastic depth for the training of language models. The paper proposed a particular training schedule (RAPTR) suitable for training with stochastic depth, where the depth of the network increases over the course of model training, ultimately training the full model. As a significant number of layers are skipped, this translates to direct gains in wall-clock time. The proposed RAPTR training scheme was evaluated on two encoder-decoder language models including BERT and UL2. Results on a small set of benchmarks indicate competitive performance w.r.t. the full model training and even outperforms it in some cases due to this acting as a regularization scheme.", "questions": "- Unclear how the authors decided to take specific mixtures of specific datasets. Was there no off-the-shelf dataset suitable for this analysis (such as FineWeb / FineWeb-edu)?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well written\n- The approach is well motivated by the literature on stochastic depth\n- The obtained results on a small set of benchmarks indicate either competitive or superior performance than the control (i.e., full model training) and directly translate to savings in wall-clock time", "weaknesses": "- The paper only focused on encoder-decoder models. I would like to see similar results on simple decoder-only language models, which are much more prevalent.\n- Evaluation on only a limited number of tasks. I would expect evaluations on a very large number of tasks from eval-harness as evaluation costs are relatively modest in comparison to the training cost.\n- Important implementation details are scattered throughout the paper. I would expect all of the details such as sequence length, and batch sizes to be specified at a single location to make it easier to understand.\n- 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. It can be significantly shorter though by just stating that we need a way to compute gradients for all the layers even if that layer didn't participate in that round. Perhaps just include another trivial implementation that works even in distribution reduction i.e., to multiply block $i$ output with $\\alpha_i$. This ensures that gradient is computed for all parameters, however, without any gain in wall-clock time.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6w8gWskhkH", "reviewer": "Reviewer_Edxi", "summary": "This paper proposes a stagewise pretraining method, progressive subnetwork training. This method randomly selects subnetworks to train in each stage and gradually increases the size of the selected subnetworks in stages. They also discussed why layer dropping can hurt the models’ ability to capture complex correlations in the data. The authors verify the effectiveness of the proposed progressive subnetwork training method by pretraining BERT and UL2.", "questions": "Why can RaPTr achieve better performance than the baseline full-model? The baseline full-model is well-trained and is supposed to achieve the best performance.\nHow about the performance when more FLOPs are reduces? For example, reduce the FLOPs of the baseline model to half or 1/3?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The proposed progressive subnetwork training method is straightforward and easy to understand.\nThe experiments show that the proposed method is simple but effective.\nThe paper is well-structured and written.", "weaknesses": "It is good that try to theoretically explain why layer dropping methods cannot achieve good performance. However, the illustration setting that using Polynomial learning and 2-layer residual network is too naïve, which makes the conclusion less convincing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Q9qUxJfVBt", "reviewer": "Reviewer_zzqn", "summary": "This paper argues that stagewise dropping strategies, such as layer dropping, can be effective for efficiently training large language models. The paper proposes a stagewise training framework of progressively growing a subnetwork and shows that it generalizes layer dropping. They theoretically illustrate the effectiveness of this strategy and moreover empirically demonstrate effectiveness in speeding up training on standard benchmarks.", "questions": "Please see Weaknesses.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The problem of efficient training is of increasing importance as we scale to larger models and data sets. The paper is well-written and comprehensive, covering theoretical justification, detailed numerical analysis, and implementation guidelines.", "weaknesses": "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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GyEZME0aAC", "reviewer": "Reviewer_tJ3v", "summary": "This paper proposes a new approach to stage-wise pre-training of large models to improve the efficiency and computational cost. The main idea is to train random subnetworks within a base network for each stage of training and increase the complexity of the subnetworks as training progresses. Experimental results show some improvement in language models. Additionally, they include theoretical work to support their method, using a simple example of learning a polynomial function and demonstrate stability in stage-wise training.", "questions": "1. Is there a generalisation beyond the polynomial example for RAPTR capturing higher or lower degree components in each stage?\n\n2. Can it be shown that the loss gap for other stacking or dropping methods are either unbounded or have a larger upper bound than RAPTR?\n\n3. Can all GLUE tasks be reported together with results on BERT-large as a base model is quite small and easy to pre-train, while RAPTR is aimed at improving pre-training efficiency. More convincing experimental results using other models are also appreciated as the current results do not persuade me to use RAPTR. \n\n4. Are we able to use RAPTR for fine-tuning on downstream tasks instead of pre-training?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The importance of improving pre-training efficiency is well justified, and it’s clear that addressing this issue is valuable. \n\n2. Efforts have been made to offer theoretical justifications for the method.", "weaknesses": "1. For experiments on BERT, only a base model is used and the improvements are relatively marginal compared to other stacking or dropping methods and only 3 tasks in GLUE are reported. Some improvements for UL2 are also within the variance of the baseline. \n\n2. 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. \n\n3. 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. \n\n4. 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. \n\nFormatting: Many typos, for example line 482: \"layernorm is enabled or not respectively\" should be the other way around, and figures with subplots are not labelled clearly.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)