Masked Structural Growth for 2x Faster Language Model Pre-training
Authors: Yiqun Yao
Zheng Zhang
Jing Li
Yequan Wang
Venue: ICLR 2024 poster (2024)
Keywords: Progressive Learning
Large Language Models
Model Growth
Primary area: general machine learning (i.e., none of the above)
TL;DR: We grow up language models in pre-training with efficient schedules and function-preserving operators that yields 2x speedup.

Abstract: Accelerating large language model pre-training is a critical issue in present research. In this paper, we focus on speeding up pre-training by progressively growing from a small Transformer structure to a large one. There are two main research problems associated with progressive growth: determining the optimal growth schedule, and designing efficient growth operators. In terms of growth schedule, the impact of each single dimension on a schedule’s efficiency is underexplored by existing work. Regarding the growth operators, existing methods rely on the initialization of new weights to inherit knowledge, and achieve only non-strict function preservation, limiting further improvements on training dynamics. To address these issues, we propose Masked Structural Growth (MSG), including (i) growth schedules involving all possible dimensions and (ii) strictly function-preserving growth operators that is independent of the initialization of new weights. Experiments show that MSG is significantly faster than related work: we achieve up to 2.2x speedup in pre-training different types of language models while maintaining comparable or better downstream performances. Code is publicly available at https://github.com/cofe-ai/MSG.

Published as a conference paper at ICLR 2024
MASKED STRUCTURAL GROWTH FOR 2X FASTER LANGUAGE MODEL PRE-TNAINING
Yiqun Yao!, Zheng Zhang', Jing Lif, and Yequan Wang"1
Beijing Academy of Artificial Intelligence, Beijing, China 2 Harbin Institute of Technology, Shenzhen, China
fygyao, zhangzheng @ baaiiacc.n, jingli phd hotmail.com t tswwangyequan @ gmail.com
ABSTRACT
Accelerating large language model pre-training is a critical issue in present re- search. In this paper, we focus on speeding up pre-training by progressively growing from a small Transformer structure to a large one. There are two main research problems associated with progressive growth: determining the optimal growth schedule, and designing efficient growth operators. In terms of growth schedule, the impact of each single dimension on schedule's efficiency under- explored by existing work. Regarding the growth operators, existing methods rely on the initialization of new weights to inherit knowledge, and achieve only non-strict function preservation, limiting further improvements training dy- namics. To address these issues, we propose Masked Structural Growth (MSG), including (i) growth schedules involving all possible dimensions and (ii) strictly function-preseet
growth operators that is independent of the initialization new weights. Experiments show that MSG is significantly faster than related work: we achieve up to 2.2x speedup in pre-training different types of language models while maintaining comparable better downstream performances
INTRODUCTION
Pre-trained language models (PLMs) (Devlin et 2019 Brown et al., 2020 OpenAl 2023) have
brought tremendous potential to research and applications. However, the excessive computational cost of pre-training is still a bottleneck. In addition the rents on computational platforms, the delay in research cycles and increased carbon footprints (Schwartz et al. 2020) are also undeniable issues. While structured pruning (Wang et al. 2020) Xia et al. 2022; Chen et al.] [022]) can ex- pedite pre-training, the resultant models are smaller in size and can be inferior in performance and knowledge capacity (Li et al. 2020), limiting their adaptability. Instead, we study the problem ac- celerating pre- re-training while still producing large-scale models. We focus progressively growing from smaller models to larger ones, which is an intuitive idea inspired by the neurogenesis in human brain (Eriksson al. 1998: Van Praag et al. 2002 Deng al., 2010). Although research like the Chinchilla laws (Hoffmann et al.) 2022) provides insights on the optimal model size given a FLOPSs budget, it only considers fixed model sizes in training. Thus, the best practices for progressive growth remain a largely uncharted territory.
We consider two main problems: the growth schedule and the growth operator. Growth schedule determines when and where to grow the model structure. Existing work (Gong et al. 2019 Gu et al., 2021) has studied growth schedules for Transformers (Vaswani al., 2017) invol viig layer number and the width of FeedForward Networks (FFNs). However, finding an efficient multi-staged sched- ule involving all the possible growth dimensions remains difficult. Growth operator stands for the operations applied during growth to inherit knowledge from the preceding model, for which function preservation (Chen et al. 2016 2022) is a thcoretically important property because insures that
"Corresponding author
CCode is publicly available wttttps:///"" /ttttss://gtthub.ccmm www..tth.cc........../
/MSG

---

Published as a conference paper at ICLR 2024
Table I: Summary of related work. stands for "supported but not function- function-preservingg preserving": X stands for "not supported"; 0 stands for "supported, sometimes function- function-preserving but can be problematic in
other cases (Section 2.3.1
2.3.2 ; v stands for "supported and function-preserriing all cases".
Method
lover_nwe hidden fi head
Any Initialization
Schedule
Speed-up Ratio"
Stankong [Gong et a 2019)
x
x
X
stage
1.65%
Compounatoon Come
2021
X
stali-stage
1.82x
Staged Shen Ai 20223
X
X
onehwo-stage
Ben2 (Chen ot mi 2022
X
one stage
1.82x
L:GO [Wang el
one-stage
MSG (ours)
malli-staee
2.2x
are not direcily compurable because the vasily different experimental settings umong work. We just summanize the highest results reporied under cemtiin setting, and discuss the fair comparison topic in Section 4.22
* Most resalts in this
the initialized large model behaves the same as the small model. There are two potential drawbacks associated with existing operators: first, the function preservation is not strict, leading to function disparities in certain cases; second, they rely fully on the initialization of new weights, constraining improvements in training dynamics (Section 2.3.3).
To help address these issues, we propose Masked Structural Growth (MSG), a novel progressive learning framework for language model pre-training with Transformers. We leverage a masking mechanism to ensure function preservation by first eliminating the effects of the new neurons, and then gradually enhancing their roles in subsequent training. MSG offers growth operators for all possible dimensions (Section (.2) with decent flexibility in schedule design, resulting in growth schedules with state-of-the-art speed-up ratios of 2.2x for Bert (Devlin et al., 2019) and 1.4x for GPT-2 (Radford et al., (2019) while maintaining comparable or better downstream performances. By solving a dilemma caused by Layer Normalization (Ba et al. 2016), MSG achieves strict function preservation in arbitrary expansion, for the first time. Moreover, MSG operators are independent of the specific initialization strategy of new weights, making friendly to future study.
Progressive growth also proves feasible in training large language models (LLMs), but it is difficult to conduct systematic and rigorous studies due to limited resource and variety in settings. For related topics, please refer to the FLM-101B technical report (Li et al. 2023) which incorporates operators originated from MSG.
To summarize, our contributions include (i) propose MSG: novel framework for progressive pre-training, and develop growth schedules with ofate-of-the-art speed-up: (ii) we demonstrate that MSG is a solid backbone for future research due its strict function preservation and independence of new weight initialization; (iii) experiments show that MSG outperforms existing methods across diverse model configurations, and provide insights on the role of function preservation.
PRELIMINARIES
We define the task in Section 2.1
describe the growth dimensions in Section 2.2 and discuss func- tion preservation in Section 2.3 to explain the motivation of MSG.
2.1 TASK FORMULATION
Our objective is to find the optimal growth schedule (sg: ..., sT, ST) E S and operator g (90....,TT) E & to minimize the total wall (real-world) time of pre-training a model f. keeping its downstream performance on dataset DD intact. Each s, in s is the intermediate structure after the t-th growth, represented by a hyperparameter set (Section 2.22: operator 9r controls how the weights of s, initialized to form a new model f. based on the previously trained model fi
ft 91 Ui-1 $c-
Let e(f(t) be the wall time cost for training each Se to be fī, and f be model with the same structure as the final fr. but trained from scratch without growing- Then, for each intermediate model, the ratio c((f) is measurable with analytical expressions like (Hoffmann et al. 2022). Finally, for some evaluation metric L (the higher the better), our task is formulated as:
mingeg ses Els cffe), s.t. L(J L(F.D).
2.2 GROWTH DIMENSIONS OF TRANSFORMERS
There are 4 expandable dimensions in Transformers (Vaswani et al. 2017). We name them by the corresponding hyperparameters: s, = (hidden dim, fn dim, head num. layer num). hidden dim is the size of feature embedding; ffn dim is the intermediate layer size of the FeedForward Networks
2

---

Published as a conference paper at ICLR 2024
e- G ?
NetNet
MSG
Figure MSG (right) vs. Net2Net (middle) in the expansion of fully- fully-connected layers.
(FFN); head num is the number of heads each Multi-Head Attention (MHA) module; layer_num num is the total number of Transformer blocks. For instance, Bert-base is represented (768, 3072, 12, 12) and Bert-large is (1024, 4096, 16, 24) (Devlin et al. 2019)
2.3 FUNCTION-PRESERVATUN
In equation 1 if gr insures that
Vz,
(3)
we name a functioon-preserving operator. In this section, we first prove that most existing op- erators do not achieve strict function preservation in Layer Normalization (2.3.I) and the depth dimension (2.3.2). Next, we study the initialization for new weights and discuss in which cases could strict function preservation critical (.3.3). Table
summary of the mentioned methods in
2.3.1 THE LAYER NORMALIZATION DILEMMA
Layer Normalization (LN) (Ba et 2016 is formulated as:
(4) where E Rdd is feature vector, w Rd and 8 Rd are trainable weight and bias. In most implementations, E IR and E IR are the mean and standard deviation of the d elements in æ, respectively, O denotes element-wise product. LN exists in a wide range of models. We prove that it raises a function-preserving dilemma for mainstream growth operators as follows:
LN(x) @ (x 4)/a +.
Net2Net widely-used operator for width expansion of fully-connected networks is Net2Net (Chen al. 2016). We demonstrate it in Figure I (middle). While layer æ is growing from q neurons, Net2Net maps the id each new neuron to an existing neuron m(i) following:
m()) map((1,.....)))) n n<isq"
(5)
Net2Net initializes weights ensure that neuron i is replica of m(i) by copying from the input weights U, and equally splitting the output weights W:
U ki U m))) W ".j [=]m(a)==([]]
W mi()
(6)
where | denotes cardinality. As result, the output values of the consequent layers preserved.
In Transformers, previous state-of-the-ar methods (Chen et al., 2022 Qin et al. 2022b Gu et al. 2 are essentially different implementations of the mapping function (Eq. [5) and (heen al, 2022 holds a similar spirit. Unfortunately, Net2Net is not strictly function- pro.erreeeeengggeeo if LN is hidden dim growth from
applied after æ. We take a = (x1, 12, Z2, .x3) for example:
for the 2 new neurons, Net2Net selects existing neurons to copy. For m(4) = and m15)=1, this yields:
#2, 21) /(xx) p(x') > LN(x) / LN(x'),
even if the LN w and 8's are copied accordingly. Thus, z w (LN(=')) + y, and the function preservation is broken. This holds regardless of the position of the activation functions (e.g. ReLU)..
3

---

Published as a conference paper at ICLR 2024
For weight-dependent methods like Net2Net, the only solution to this dilemma is copying each neuron for exactly the same times to preserve the mean and variance. However, this requires to at least double the width each time, resulting in a fixed exponential schedule that severely harms the flexibility. For example, a small model may have better training dynamics with hidden dim = 512 than 384 (Figure 2. but if our target model has hidden_dinn 768, Net2Net must grow from 768/2 dimensional model to be function-presere
Other Operators Similarly, operators that zeroize the new input weights (Evei et al. 2022) yield x' (I1, 222 23. 0, 0) and do not solve this dilemma. Moreover, they intrinsically require the activation function a to satisfy (0)) = O, which limits their usage. Thus, to our knowledge, there is no strictly function-preserving growth operator for arbitrary hidden_dim growth before MSG.
2.3.2 THE DEPTH DIMENSION
For depth (layer_num) growth, existing work either stacks GGng et al. 2019 Yang et. 2020 Chen
ctt al. 2022; Gu et al., 202I) or inserts (Qin et al., 2022b Li et al. 2022) new Transformer layers with weights copied from existing ones. Neither of them is function preserving.
Shen et al. (2022)
zeroized all the LN weights in the new layer for GPT-like (Radford et al. 2019) Transformers:
x (tt)
x' FFN(LN(z')),
which achieved strict function preservation y = x. However, for Bert-like implementations:
IN + (=:]), = IN(('' + FFN(x"))
this method yields all-zero output. Thus, structure-agnostie function-preesrrine operator for depth growth is also underexplored.
2.3.3 DEPENDENCY ON INITIALIZATION
Most existing growth operators aiming at function preservation determine a fixed initialization for new weights. For example, Net2Net initializes new weights fraction/copy of existing weights, while other methods simply zeroize them. One of the drawbacks of these methods (especially the zeroizing ones) is the Symmetry issue:
neurons always receive similar gradients other new/existing neurons, which brings unnecessary restriction that slows down model convergence D Furthermore, (Evei al.) 2022) and (Wang al. (022) show evidences that weight initialization can even be treated as an optimization problem itself. However, growth operators that depend on initialization deny this chance of further improving the training dynamics. Thus, although most work agree on the importance of function preservation, it seems necessary to answer "ir is worth sacrificing a better initialization."" By MSG, we decouple these two factors and explore this topic by experimenting w/ and w/o function preservation in different conditions.
3 MASKED STRUCTURAL GROWTH
We propose a novel framework for progressive training, namely Masked Structural Growth (MSG). In Section 3.1 we introduce MSG operators for all growth dimensions. In Section 3.2 discuss methodologies and results for schedule design.
3.1 THE MSG OPERATORS
The main idea of MSG is to derive external masks that completely eliminate the effects of new structures on the model's function. Immediately post-growth, with mask-0, we make the function strictly preserved. Next, we gradually increase the masks to raise the influence of new structures, and finally achieve the target structure with mask 1. The underlying rationales are (1) right after growth, the current loss precisely mirrors that of the smaller model, and it tends to reduce in the next step via gradient-descend; descend: (2) when the mask value close to 0, the outputs of existing neurons effectively provide initial guidance for the update of new weights.
3.1.1 FULLY- CONNE-CON LAYERS
Colle-Coneected (FC) layer is linear projection between feature vectors, followed by (optional) activation and LN. It is involved in all the growth dimensions except layer num. Let æ E dh and
Existing solutions symmetry hreak the function preservation Gong al. 2019 Chen al.
2022
4

---

Published as a conference paper at ICLR 2024
R be the input and output of FC layer W, respectively. As depicted in Figure (right), while expanding the size of y from dg to d22, we introduce a mask vector e with de elements value 1 and d2 d2 elements of value O:
c= [1ai 02--22l 442b
(10)
where [-] stands for concatenation along the last dimension. We first copy the existing weights from W and initialize the new weights to be arbilrary value:
W
1 1 d2
W
(11)
any value d2 < 1isd
Then, the new output y Rdds of the grown FC layer computed by:
y' c (a(W' x)) = y: 0 O2g-dal
(12)
where stands for matrix (inner) product and a is any activation function. For existing neurons, the network yields exactly the same value as original; for new neurons, their outputs are masked to
On the other hand, if the input æ: is expanding to x' E RR it should have been in the form of:
x' [c; 025 -d,
(13)
because it is the output of the previous layers and already processed by Equation the new weights following:
12
We initialize
W
isjsd
W
(14)
any value
d1 <jSd
Thus, the output y' E gdd the network is: y' a(W' x²) ( x) y.
(15)
We do not need additional mask here if the output size da is expanded. If both di and d2 are expanded, we first apply (14), then (1.), and finally mask the output by (12). After solving LN (Section 3.1 1.2... the function of the whole FC module will be preserved, illustrated in Figure
3.1.2 LN SOLUTION
We solve the IN dilemma (Section 2.3.1 by incorporating a similar mask inside of LN. While expanding the size of a feature æ and its LN from d to d', we introduce an external mask e E Rd::
LL; Ozr-a)+
(16)
and derive new formula for the mean and variance of the expanded feature x' E Rd':
pa (a'' c)/sum(c).
(17)
(a')² 21)² c)/sum(c).
where sum(-) stands for the sum all elements and is inner product. With zero mask value, the new neurons has no contribution to the mean and variance. As result, if retain the first d values of the LN weight w and bias B, and initialize the d' -d new LN weights to arbitrary value, we have:
IN* [IN(()): a),
(18)
which the values of Ear are decided by the input weights connected to the new neurons in x'. Finally, we multiply mask again to the result:
C LN" [LN(x)) -v
(19)
Thus, for arbitrary d and d', MSG is strictly function presoorinne with LN. As the training goes on, gradually increase the last d' d values in
1.
3.1.3 GROWTH OF SELF-ATTENTION HEADS
The self-attention module (Vaswani et al. 2017) contains multiple heads. Strict function- preservation be achieved with MSG by treating each head "neuron" and adapting a similar process as equations 10 to 15
We leave the proof to Appendix A due to limited space.
5

---

Published as a conference paper at ICLR 2024
layer
hidden dim
(Tr. dim
head
0861126 6.)
4122440024122 51
3.224224558 12 21
Train Loss (Bazch)
Train Loss (Batch)
(1832336.8 8))
013433261
(13244/05611
Train 100 (Detch)
Train
17883324.20 10.0
Wall Time (seconds)
Wall
43
Wall Time (scconds)
Wall
Figure 2: Training loss curves with different structural hyperparameters. We study the impact of each growth dimension on the model's "pre-training rate" y in early stages.
3.1.4 DEPTH GROWTH
We propose depth growth operator inspired by residual connection (He et al. 2016). Let Layer, be a new Transformer layer, - be its input, and c be an all-zero mask in the same size as 2n-1, we compute the layer output as:
2, co (x) c) ---
(20)
With zero mask value, the whole Layer, is skipped to preserve function; while the mask value increases to 1, the residual connection vanishes, yielding a same structure as vanilla Transformers.
3.2 GROWTH SCHEDULE
In this sub-section, study the construction of schedules. Although multi-staged growth inher- ently offers more flexibility than its single-staged counterpart, it is difficult to find an optimal growth schedule without being h 3 Existing work jump to the final schedules without detailed ex- planation Gong et al. (2019); Gu et al. (022). In contrast, we introduce a methodology of grid- searching, which holds the potential to generalize to other Transformer variants. We do not claim to completely solve the problem, but provide empirical insights to inspire further investigations.
3.2.1 SCHEDULE CONSTRUCTION BY GRID SEARCH
An ideal growth schedule entails a series of structures that learn fast in their own stages. We find that a "pre-training rate" yy is a reasonable indicator of learning speed in carly/medium stages |
A
Aloss
Aloss Astep
Astep A a..T.
(21)
Ai
A
Intermediate Structures. Growth scheduling a trade-off: a small model always has better con- stant 7, but its a drops rapidly during training, making it necessary to grow in proper time. Given the challenge of quantifying a mathematically, we opt for grid search with only 10k data to study how each of the 4 growth dimensions influences 7, respectively. We examine possible intermediate structures varying in different dimensions, and compare the training curves. The other hyperparam- eters are kept the same in the final schedules. Models typically converge in several minutes with 8 GPUs, yielding minor time cost. As demonstrated in Figure 2] we find the following heuristic rules for Bert: (i) Large layer_num severely harms with limited gain in a. It's preferable grow it in late stages. (ii) Larger hidden_dim brings significantly better for shallow models. It should start from a large value. (iii) The ffn balances the a and r itself with little impact on y, while slightly smaller head num brings better for deep models. We construct different intermediate structure sequences based on these heuristics. The same methodology transfers to GPT.
Like normal pre-training, the total training steps are pre-decided. Given the total steps and an intermediate structure sequence, we split the training process for each stage to have a roughly equal number of steps, while slightly emphasizing the carly stages. We observe that this simple methodology is efficient in practice despite being theoretically sub-optimal. An optimal strategy may be growing to the next structure once the current 7 goes lower than the full target structure given the same wall time, which is obviously hindsight and not practical.
Stage Duration.
Shen et al. (2022) studied this topic a different setting, see Appendix B for more discussions.
Note that y is longer valid when the loss curve almost levels off. For example, in late stages, we observe no consistent correlation between training/validation loss and downstream performances for Bert.

---

Published as a conference paper at ICLR 2024
Table 2: Growth schedules expanding one dimension in each stage. Sch1-B and sch1-L. used for Bert-base and Bert-large, and schl-G for GPT-2, respectively. Rapid-L is rapid schedule with multiple expansions within 100k steps.
Schedule
2000KK
200K 400k
400K 500%
SOUK 6OOK
GOUK 700k
700k YOUK
Sch1-B
(512,768,8.3)
(512.3072.33)
(5 1123072....6
8072.66)
(768, 3072, 12.61
(768,3072, 22.12)
Schl-L
(768, 1024,12.6)
1768.499612..6
7768,4096.1..1
(1024, 1224,
(1024, 14966,6666111
1024,4096.16.2
Schedule
0'50k
TOOK
100K"1300 130K
130K 160k
160k 190K
end
Schl-G
8512,768,833
3072,3))
(512, (072,0)7)
7768,072,8.88
M182,
(768, 3072,
Schedule
0'400K
400K 420K
420K-HHUK
440K SUOK
S0UK S00K
Rapid-L
(512,10248.6)
8..6409%
(1024,40%,8
[1024,4096) 16,6)
(1024,409,.66,
Table 3: Experiments on the mask strength Ø.
100
250
500
800
1000
loss @ 10k
3.36
3.34
3.25
3.23
3.23
3.23
Schedules Constructed. Following the same methodologies mentioned above, we construct the schedules for both Bert and GPT, as detailed in Table [ Sch1-B, Schl-L, and Schl-G are used in our main experiments for Bert-base, Bert-large, and GPT-2, respectively. These schedules are multi- staged with one dimension expanded per stage, leading to 1.8x, 2.2x, and 1.4x speed-up compared to training from scratch with the same step numbers. This significantly faster than (Gong et al.) 2019 Gu et al. 2021) the same target models. Rapid-L is rapid schedule with several successive growth stages arranged within a span of 100k steps to study the influence of stage duration. Other schedules for ablation study are explained in Appendix G
Mask Strength. We define & as the number of steps used to linearly increase the mask values from 0 to I. Initially, based on intuition, we set 5000 for all our experiments, which worked well. To further study its impact, we experiment with different & values on Bert-large with only 100k data. We train for more than 10 epochs simulate the middle and late phases of the full-data training. The schedule we use here is a compressed version of Rapid-L. (Table 2). We compare the training loss at 10k steps (Table(3) and observe that all é values exceeding 500 can ensure stability in subsequent training and yield similar results.
3.2.2 INITIALIZATION
MSG supports arbitrary initialization of new weights. We find that sampling from normal distribu- tion N(, 0.02) for width growth and following the stacking strategy for depth (complemented with eq. 20 already outperforms state-of-the-art methods. Given . 0.02) is prevalent choice for initializing Transformers from scratch, we believe that it has been optimized through community- driven refinements. We leave the incorporation of other "learning initialize" methods for future. For the optimizer, we copy the stored moving averages of Adam (Loshchilov & Hutter 2018) for existing weights, and start from scratch for new weights.
4 EXPERIMENTS
4.1 SETTINGS
We train both Bert-base and Bert-large (Devlin et al.) 2019) using a pre-processed combination of English Wikipedia and Book Corpus (Zhu et al. 2015). and GPT-2 (Radford et al.) 2019) using
OpenWebText Web'eett (Gokaslan & Cohen] 2019). The growth schedules are presented in Table 2 For
other details, please see Appendix D As for evaluation, we fine-tune our Bert models on the GLUE WWang et al.) 2018) and IQuADvi.1 (Rajpurkar et al. 2018) tasks, and GPT models on Wikitex12 Merity et al. 2016). We report the mean and standard deviation of the metrics across runs on the dev set. Detailed fine-tuning settings provided in Appendix E
4.2 COMPARED METHODS
In complement to Table I
we answer the following two questions for more fair comparison:
On which schedule should we compare? If the growth operators ensure no performance drop. the speed-up ratio is fully dependent the schedule. Since MSG succeeds is faster than the previous state-of-the of-the-art schedules (Section 3.2.I) for the same target models, we compare different operators our own schedules.
own schedules which

---

Published as a conference paper at ICLR 2024
Table 4: Main results and ablation studies on Bert-Large. We report the average GLUE score and
the accuracy/FI for SQuaD. The numbers are mean (standard deviation) computed across 3 runs.
Target
Bert-Large
Schodule
Full
Schl-L
Rapid-L 29h. 20min
Wall time
705, 48min
32b, 10min
Operater
BerLBERT
MSO
No Musk
MSG
No Mask
Glue Avg. 82.20022
82.1(0.3)
83,2022
82.210.2 79.700.3888.0000
79.10.4)
77.90.3)
Squal>
80. 20.2(0.2)87.90)
80.010. 10.00.1187.660 (87.66001) 1)
81.00.2088..10.1
76.500.61844444055)
75.70..184.000
Table 5: Main results and ablation studies on Bert-Base and GPT-2. For GPT-2, we evaluate on the
validation set of Wikitext2 (Merity al. 2016). "PPL-zs" stands for zero-shot perplexity (standard
practice at the moment), while "PPL_-ft" is fine-tuned with the Wikitext2 training data.
Target
Bert Base
Target
GPT-2
Schedule
Full
Sch1-B 14h, 36min
Schedule
Full
SchI-G
Wall time
26h, 10min
Wall time
53h, Imin
37h, 59min
Operator
BentBERT
MSG
Operator
Ben2BERT
MSG
Glue Avg.
80.70022 19.100.2V86.900 (00.))
80.5(0.2)
81.00.22
PPL-2s
41.31
42.24
41.20
SquaD
19.000.1
PPL-ft
23,90
25.13
24.14
What operators should we compare with? Mainstream growth operators include zeroizing and Net2Net. Since zeroizing limits the usage of activation functions (Section 2.3.1
(S.ction[..311 and sometimes lead
to problematic training (Gong et al. 2019), most related work (Ging et [ Gu et al.) 2021 Chen al.) 2022 Qin et al., 2022b) apply combination of Net2Net on width and some "direct copy" strategy on depth (Table Thus, for the main experiments, we compare with our reproduc- tion the latest version of the BeRt2BERT (Qin et al. (022b) operators that achieved state-of-the-art performances. It features the same Net2Net operators for width growth as (Chen et al. [2022), and an interpolation strategy for depth. Other configurations are kept the same as MSG. We compare with the data-driven LiGO method (Wang et al. 2022) in Appendix H due to limited space
4.3 MAIN RESULTS
We present Lert-Large results in Table Bert-base and GPT-2 results are in Table The results on each sub-task is left to Appendix F due to limited space. We have the following observations:
MSG schedules achieve higher speed-up. Comparing with the no-grow baselines (Full-*), MSG reaches equal downstream performances on Bert-base with 1.8x speed-up, and significantly higher performances on Bert-large with 2.2x speed-up. Comparing with the previous best multi-staged strategy (Gu et al., 202T) with 1.74x speed-up on Bert-base and 1.82x on Bert-large, our advantage because MSG supports growth on all the possible dimensions, which enables more flexible schedule design, while CompoundGrow (Gu et al.) (2021) explored only 2 of them. The same holds for GPT-2 with 1.4x speed-up and comparable performances with baselines 6
MSG operators have better potential than Ber12BERT. Comparing with Bert2BERT on the same schedules, MSG achieves significantly better performances on Bert-large and GPT-2. This because (1) MSG is strictly function- random initialization, which naturally solves the Symmetry problem and leads to better dynamics.
function-preservingg preserving, saving the time for recovering function; (2) MSG supports
4.4 ABLATION STUDY
We try to answer the following questions with ablation studies:
Factors involved in schedule design. Comparing the results of Rapid-L and Sch1-L (Table 4 even though the time cost is close, we observe significant performance gap. This indicates that rapid growth (expanding the model for too many times in a short period) is not friendly to Transformers. This supports the effectiveness of our intuitive strategy to equally split the stage duration (Section 3.2.1. Another factor may be the selection of intermediate structures. To study this, as an alternative schedule, we further develop the SchO (Appendix G and Table 8 following the same heuristics as
'The Gradmax Evci al. 2022 operators have proven to be directly applicable Transformers. 'For GPT-2, Beri2BERT starts from a large checkpoint with a hidden of 512 and 12 layers without taking its training cost into account; LiGO reports 1.22x speed-up up with a relatively small starting checkpoint. We observe that most methods have lower speed-up GPTs than Bert and we guess this is because the causal language model loss decrease relatively slow for shallow models (3 layers, etc.).
8

---

Published as a conference paper at ICLR 2024
Schl but with different intermediate structures. The GLUE results with SchO are presented in Table The scores are close to Schl (82.8 vs. 83.2 on Bert-Large, 80.9 vs. 81.0 on Bert-base), supporting the robustness of our schecule-consttruttion methodologies.
9
Importance function preservation. To study the effect of function preservation. we remove all the masks right after growth (No Mask in Table 4, resulting in strategy of randomly initializing new weights without function preservation. Although (Chen et al.] 2022) showed that a similar strat- egy is inferior in one-staged settings, we observe that performs reasonably well on our Schl-L, comparable to both the Full baseline and Bert2BERT (Table D This indicates that random initial- ization can indeed result in better training dynamics than the weight-dependent Bert2BERT. On the other hand, No Mask underperforms MSG on both Schl-L and Rapid-L with the same initialization strategy, which supports that function preservation is beneficial if it does not harm the initialization. Moreover, looking at the training loss curve of No Mask vs. MSG on Rapid-L. (Figure we observe that MSG will be more stable than No Mask if the training continues. This indicates that for rapid structural changes and limited time budget, function preservation improves the training stability.
5 RELATED WORK
5.1 ACCELERATING PRE-TRAINING
Progressive training Wen et al. celerate pre- pre-training for both Computer Vision (Li et al. 2022) and Natural Language Processing Gong et al.) 2019 Gu et al. 2021). Gong et al. (2019) proposed a stacking method which dou- bles the model depth in each stage. CompoundGrow (Gu al. 202T) extends stacking by adding FeedForward Network (FFN) expansion into schedule. Shen et al. (2022) proposed a staged method that further supports expanding the hidden size of features, but limited to integer multiplications in which the feature size, FFN size, and head number are bounded. BerrBBRT (Chen et al.) [022] Qin et al., 20226) and LiGO (Wang et al., 2022) support all possible growth dimensions, but focus on one-stage growth where all the dimensions dimension has different impact on the training dynamics (Section 3.2). which is underexplored. In contrast, we study multi-staged growth with only one dimension expanded each time, resulting in higher degree of flexibility and speed-up. A detailed summary is in Table T Other work that speeds up pre- pre-training include partial update (Ynn et al.) (020), weight sharing (Lan et al. 2019), adversar- ial training (Clark et al. 2019), compute-optimal scaling laws (Hoffmann et al.] and Mixture of Experts (MoHs) (Fedus et al. (022), which orthogonal to our topic. These approaches systematically surveyed by (Bartoldson et al.) 2023)
2020 Dong et al.
2020 has raised attention for its ability to ac-
expanded simultaneously. However, each growth
5.2 FUNCTION-PRESERVIIO GROWTH
Function preserving growth ensures that post-grow post-growth models mirror the outputs of their precursors for any input, which is intuitively beneficial to knowledge inheritance. It typically tackled math- ematically to avoid conducting the time-consuming data-driven distillation (Jiao et al.) 2020) Qin jet al., 2022a). [hen et al. (2016) proposed the Net2Net method that preserves function in width expansion by splitting existing weights to form new neurons. Other methods (Evci et al., 2022] Wei et al. 2016) preserve function by setting some weights to zero. While applied to Transformers structure, existing operators can only achieve non-strict function preservation in most cases, lim- iting their efficacy (Section 2.3 Besides, they fully rely on the initialization of new weights. In contrast, method is strictly function- function-preserving in all growth dimensions and is agnostic to new weight initialization. LiGO (Wang et af. 2022) designed an efficient data-driven method to map the weights between small and large models, which improves the training dynamics by learning to initialize. Due to MSG's unique advantage of supporting arbitrary initialization, it can be combined with LiGO-like methods on the new weights. We leave this interesting topic for future.
CONCLUSION
We proposed a novel framework: Masked Structural Growth (MSG) for progressive pre-training of language models. MSG supports growth operations on all the possible dimensions for efficient schedules. Experimental results show that MSG achieves state-of-the-ar speed-up on Best- Bert-large, and GPT-2 with equal improved performances on downstream tasks. We also prove that MSG is strictly function cenctioo-preserving as well as independent of new weight initialization. These properties make MSG outperform other operators and benefit future research.
9

---

Published as a conference paper at ICLR 2024
LIMITATIONS
In this work, we have proposed the heuristic rules for growth schedules involving one dimension at a time. However, there is still unexplored space because subsets of several growth dimensions can be combined in the same stage. We believe that there is still long way to go for self-adaptive growth schedules that actively search for the next target structures base training dynamics. Another limitation of our work is that the best initialization strategy for Transformer growth remains an open question. Although we have showed that random initialization can be better than Net2Net, we leave more in in-hepth study (i.e. a generalization of (Evci et al. 2022 to Transformers) for future work.
ACKNOWLEDGMENTS
This work was supported by the National Science and Technology Major Project (2022DD0113000) and the National Science Foundation of China (No. 62 62106249).
REFERENCES
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: vrXiv:1607.064500 (607.06450, 2016.
Brian R Bartoldson, Bhavya Kailkhura, and Davis Blalock. Compute-efficient deep learning: Algo- rithmic trends and opportunities. Journal of Machine Learning Research, 24:1-77, 2023.
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 are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
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. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 2134-2148, 2022.
Tianqi Chen, Ian J. Goodfellow, and Jonathon Shlens. Net2net: Accelcrating learning via knowledge transfer. In Yoshua Bengio and Yann LeCun (eds.), 4th International Conference Learning Representations, ICLR 20/6, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceed- ings, 2016. URL
httttp://arxiv..oo/ http:////wwwww..//aaa// abs 151111111111 1551.
Xiaohan Chen, Yu Cheng, Shuohang Wang. Zhe Gan, Zhangyang Wang, and Jingjing Liu. Earlybert: Efficient bert training via early-bird lottery tickets. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the I 11th International Joint Conference on Natural Language Processing (Volume I: Long Papers), pp. 2195-2207, 2021.
Kevin Clark, Minh-Thang Luong, Quoe V Le, and Christopher 1) Manning. Electra: Pre-training text encoders as discriminators rather than generators. In International Conference on Learning Representations, 2019.
Wei Deng. James B Aimone, and Fred H Gage. New neurons and new memories: how does adult hippocampal neurogenesis affect learning and memory? Nature reviews neuroscience, 11(5): 339-350, 2010.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. 4171-4186, 2019.
Chengyu Dong. Liyuan Liu, Zichao Li, and Jingbo Shang. Towards adaptive residual network training: A neural-ode perspective. In International conference on machine learning, pp. 2616 2626. PMLR, 2020.
Peter S Eriksson, Ekaterina Perfilieva, Thomas Björk- BrinkkEnisssn, Ann-Marie Albom,, Claes Nordborg, Daniel A Peterson, and Fred Gage. Neurogenesis in the adult human hippocampus. Nature medicine, 4(11):33331111, 1998.
10

---

Published as a conference paper at ICLR 2024
Utku Evci, Bart van Merrienboer, Thomas Unterthiner, Fabian Pedregosa, and Max Vladymyrov. Gradmax: Growing neural networks using gradient information. In International Conference on Learning Representations, 2022.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research, 23(1): 5232-5270, 2022.
Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. 2019.
Linyuan Gong, Di He, Zhuohan Li, Tao Qin, Liwei Wang, and Tieyan Liu. Efficient training of bert by progressively stacking. 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. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp-
5174-5180, 2021.
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.
Jordan Hoffmmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katherine Millican, George van den Driessche, Bogdan Damoc, Au- relia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Oriol Vinyals, Jack W. Rae, and Laurent Sifre. An empirical analysis of compute-optimal large language model training. In NeurlPS, 2022. URL http:: n nips. http:p/er..iip..c.pppppeeeeeeeee ww/papp:ppppcccccccccccctcce -bst
c 032
html
Xiaoqi Jiao, Yichun Yin, Lifeng Shang. Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. Tinybert: Distilling bert for natural language understanding. Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 4163-4174, 2020.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Sori- cut. Albert: A lite bert for self-supervised learning of language representations. In International Conference Learning Representations, 2019.
Changlin Li, Bohan Zhuang, Guangrun Wang, Xiaodan Liang, Xiaojun Chang, and Yi Yang. Au- tomated progressive learning for efficient training of vision transformers. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 12486-12496, 2022.
Xiang Li, Yiqun Yao, Xin Jiang, Xuezhi Fang. Xuying Meng, Siqi Fan, Peng Han, Jing Li, Li Du, Bowen Qin, et al. FIm- 101b: An open I1m and how to train with $100 k budget. arXiv preprint M.Xiv:3309008555 2023.
Zhuohan Li, Eric Wallace, Sheng Shen, Kevin Lin, Kurt Keutzer, Dan Klein, and Joey Gonzalez. Train big, then compress: Rethinking model size for efficient training and inference of transform- ers. In International Conference on Machine Learning, pp. 5958-5968. PMLR, 2020.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2018.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. 2016.
OpenAl. Gpt-4 technical report. arXiv preprint arXiv: 2023.
Yujia Qin, Yankai Lin, Jing Yi, Jiajie Zhang, Xu Han, Zhengyan Zhang, Yusheng Su, Zhiyuan Liu, Peng Li, Maosong Sun, and Jie Zhou. Knowledge inheritance for pre-trained language models. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Techmologies, pp. 3921-3937. Association for Computational Linguistics, 2022a.
11

---

Published as a conference paper at ICLR 2024
Yujia Qin, Jiajie Zhang, Yankai Lin, Zhiyuan Liu, Peng Li, Maosong Sun, and Jie Zhou. Elle: Efficient lifelong pre-training for emerging data. In Findings of the Association for Computational Linguistics: ACL 2022, pp. 2789-281111 2810, 2022b.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, lly Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Pranav Rajpurkar, Robin Jia, and Percy Liang. Know what you don't know: Unanswerable questions for squad. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), pp. 784-789, 2018.
Roy Schwartz, Jesse Dodge, Noah A Smith, and Oren Etzioni. Green ai. Communications of the ACM, 63(12):54-66 63, 2020.
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. PMLR, 2022.
Henriette Van Praag, Alejandro F Schinder, Brian R Christie, Nicolas Toni, Theo D Palmer, and Fred Gage. Functional neurogenesis the adult hippocampus. Nature, 41566875):1030--0033 2002.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- fion processing systems, 30, 2017.
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. 07761,, 2018.
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. In The Eleventh International Conference on Learning Representations, 2022.
Ziheng Wang, Jeremy Wohlwend, and Tao Lei. Structured pruning of large language models. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 6151-6162, 2020.
Tao Wei, Changhu Wang, Yong Rui, and Chang Wen Chen. Network morphism. In International conference on machine learning, pp. 564-572. PMLR, 2016.
Wei Wen, Feng Yan, Yiran Chen, and Hai Li. Autogrow: Automatic layer growing in deep volutional networks. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 833-841, 2020.
Mengzhou Xia, Zexuan Zhong, and Danqi Chen. Structured pruning learns compact and accurate models. In Proceedings of the 60th Annual Meeting of the Association for Computational Lin- guistics (Volume I: Long Papers), pp. 1513-1528, 2022.
Cheng Yang. Shengnan Wang, Chao Yang, Yuechuan Li, Ru He, and Jingqiao Zhang. Progressively stacking 2.0: A multi-stage layerwise training method for bert training speedup. arXiv preprint 0rXiv:2011.13635 2020.
Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In Proceedings of the IEEE international conference on computer vision, pp. 19-27, 2015.
12

---

Published as a conference paper at ICLR 2024
A PROOF OF MSG FUNCTION-PRESERVATION ON ATTENTION HEADS
In previous work, the size d of query, key and value vectors in each head is decided by hidden_dim / head num. In this work, we break this binding by fixing d = 64 and expanding the self-attention module via adding new heads. Weight matrices Q, K, V are all in size (hidden dim, d). While growing from 11 to n2 heads, we keep the weights for existing heads and initialize the new heads arbitrarily. For head i:
Q K V (Q.,k, V. is
(22)
any value iE(m). n2]
For all the new heads n <is n2, we multiply an all-zero mask e its value vector:
b. @ (V.' =x)= = 0, <is nz;
(23)
and get all-O attention outputs for < i S n12
o, softmax ( (K4) /vd) of = 0,
(24)
where of and Ky contains all the queries and keys for input sequence.
For the FC layer W. E E Rhidden (dm) that collects head outputs, we initialize it as:
(W.Y., (W.)-s j d- any value je (ddmi,ddnz)
(25)
Finally, the self-attention output is function-preserni for any nı and ng:
y W (ofi ony W., + o..
(26)
B DIFFERENT RESEARCH PROBLEMS IN SCHEDULES
Research on staged training (Sen et al, 2022) provided insightful analysis on the problem of auto- matically deciding the time (step) grow. They perform similar experiments as our Section 3.2.1 to find constants related to training curves that help this decision, although still relying on hindsights.
The main difference is that they actually answer the question of "given the structure sequence, on which step should we grow?", while our work in Section 3.2. focuses on answering "how to design a good structure sequence, which dimension should grow first?"
In their work, they solve our question by manually design the schedules, as most existing work did; in our work, we solve their question by assigning roughly equal steps for each stage. We believe that both are important topics for future research.
C AN INTRODUCTION TO BERT2BERT
Bert2BERT (Chen et al. 2022 is the most typical related work that incorporates Net2Net for width growth and a weight-copying strategy for depth growth. Specifically, the expansion of hid- den_dim and ffr _im standard Net2Net applications fully connected layers (Eq. [)). For head num, they also considered each head as a neuron-like unit. The Net2Net mapping function for attention heads is written as:
i of
(27)
(map(().....)))))
where a² and a' are the head numbers of the source and target models, respectively. The parameters in head m(i) are direct copied from head i. the depth dimension, they follow (Gong et al. 2019) to stack the whole model (except embedding layers) on itself repeatedly until the target layer Jumm reached. In a following implementation (Qin et al. 2022b). an interpolation strategy is used instead, which copies each layer separately and puts the duplicate above the original layer. In Bert2BERT, a warmup stage is introduced before the main pre- prainiiin stage, in which only the top layers are tuned. Note that MSG does not have this warmup stage due to its multi-staged nature.
13

---

Published as a conference paper at ICLR 2024
Table 6: GLUE task results for Bert-Large.
Schedule
Operator
GLUE Avg.
CoLA
SST-2
MRPC
STS-B
Full
82.2(02))
57.9(1.3)
91.9(0.1)
5y88 7(.2)
89.000 y88 70.1)
Sch1-I.
Ner2Net
82.1(033)
57.6(1.3)
91.9(0.2)
85.1008)88.6((0.)
09.500.289.000.2
Sch1-L
MSG
83.2(0.2)
61.9(0.4)
91.6(0.3)
87.3(0.7)/91000) ..0.0))))91000004))))
89.700.0)89.20..0
Schl-L
No Mask
82.2(0.2)
56.8(1.0)
91.0(0.4)
86.2(0.6)99)
4(0.3)
89. 1(0. V88 6(0.2)
Rapid-L
MSG
79.1(0.4)
48)
90.7(0.4)
81.3(0.6188.3400.
87.40.2V88.000..
Rapid-L
No Mask
77.9(0.3)
45.0(2.0)
88.6(0.2)
77.3(L0/885.00.6)
1..70..186.3(0.1
Schedule
Operator
QQP
MNLI(m/mm)
QNLI
RTE
SQuADvIJ
Full
30.7(0.3)87.7700)
83.3(0.5 66..366003)
90.4(0.2)
68. ((0.))
80.2(0 2y87 9(0.3)
Schl-L
Net2Net
90.8(0.2)877.8 8(0.3)
90.4(0.2)
68.0(1.1)
80.000. V887 6(0.1)
Sch1-I.
MSG
90.900.11/87.900.
33.7(0.5)8445(00.3
91.2((0.)
69.0(1.6)
81.3002 ((0.))
Schl-L
No Mask
90.8(0.2)877 7(0.2)
83,7(0 33.7 ((8.
91.3(0.3)
68.5(1.1)
79.7(0.3y88 0(0.2)
Rapid-L.
MSG
89.8 39.8(0.4)86.200666
80.1(0.1 (...)
88.2(0.2)
66.1(0.7)
76.500.6//44440.55
Rapid-I.
No Mask
89.8(0.1)/86.4(0.2)
19.7(0.2)/7977(0.2)
87.3(0.1)
66.7(0.3)
75.7(0.4884.00.3
D DETAILS FOR PRE-TRAINING
Bert Bert-base has 110M parameters, while Bert-large has 336M parameters. For the baseline mod- els without growth (Full-"), we train for 1M steps, since the downstream performances continued increase between 900k to 1M steps. We learning rate of (1e-4,le-5) and warm-up step of (10k,30k) for (Bert-base, Bert-large), respectively, with linear learning rate scheduler for both. The learning rate is reset to maximum value the 1st and 2nd growth stage, following (Gu al. 2021). The batch size is set to 256 and maximum sequence length to 128. We clip the gradient norm 1.0 for Bert-large. All the experiments are conducted on single pod with Nvidia A A100 GPUs. The pre-training dataset includes 67 million sentence pairs with 15% tokens masked.
GPT GPT-2 has 110M parameters. We train with a sequence length of 1024 and a batch size of 140 for 240k steps, which is 2 epochs on Open WebTee. For schedule Sch1-G, we continue training with the full model size until reaching the time cost reported Table We use a learning rate of le-4 and learning rate warmup ratio of 0.01. The experiments are conducted on 2 pods with 14 Nvidia A100 GPUs.
E DETAILS FOR FINE-TUNING
For all the GLUE tasks, we use a batch size of 32, sequence length of 128 and learning rate of 2e-5. We fine-tune for 5 epochs for small datasets including CoLA, MRPC, STS-B B, and RTE, (we exclude
WNLI following most related work (Devlin et al. 2019) Chen al. 2022: Gong et al. 2019 and 3 epochs for other tasks.
For SQuAD, we fine-tune with batch size of and learning rate of 3e-5 for 2 epochs. SQuAD metrics are very sensitive to sequence length, and most models work well only with more than 384 sequence length. Thus, we continue pre-training after the whole schedule with a sequence length of 512 for 100k steps for all the methods compared. This yields slight drop in actual speed-up ratios from 2.2x to 2.0x on Bert-large. Obviously, this drop can be alleviated a sequence length of 512 is used through the whole pre-training process. Moreover, the SQuAD results of MSG significantly higher compared to the baselines in all cases (Table additional time savings.
which can be converted to
For Wikitext2 fine-tuning, we use a learning rate le-4 for 3 epochs with a sequence length of 1024 and batch size of We report the zero-shot and fine-tuned perplexities on the validation set.
F RESULTS ON SPECIFIC GLUE TASKS
We present the scores each GLUE task in Table and 7 For metrics, Matthews correlation for CoLA, Pearson/Spearman correlation for STS-E, E, accuracy/fl for MRPC, QQP, and SQuAD, and accuracy for all the others. The numbers mean (standard deviation) computed across runs.
14

---

Published as a conference paper at ICLR 2024
Table 7: GLUE task results for Bert-Base.
Schedule
Operator
GLUE Avg.
CoLA
SST-2
MRRC
SISB
Full
80.7(02))
52.2(0.8)
90.4(0.2)
85.900.999.11033)
58.8(0.1V88.
Sch1-B
Net2Net
80.5(0.2)
52,4(1.5)
91.1(0.4)
84,30 84,300 8(0.6)
88,3 1188.000. 000.1)
Sch1-B
MSG
81.010.2)
58.2(1.6)
91.0(0.2)
85.00.55899440.3
18.7 187.11877.600111111
Schedule
Operator
QQP
(N/mm//m))
QNLI
RTE
SQuADv1.1
Full
90.610.11/87.330.11
32.590.3W92.((011
89.910.1)
65.1(0.7)
Sch1-B
NetNNet
90.1 1)
81. 81.100.2882.100.
89.2(0.))
66.3(0.5)
79.070. (86.7(0)
Sch1-B
MSG
90.0(0.1V87.000.1
81.80.3)62.40.2)
89.9(0.1)
63.1(1.6
79.6(0.5V88.210.44
No-mask
60x
(Batch)
10"
100
20000
40000
60000
60000
100000
Wall Time (seconds)
Figure Training loss curves of MSG with ans without mask on the Rapid-L schedule.
G ALTERNATIVE SCHEDULES AND RESULTS
According to our heuristic rules in Section 3.2.1 starting from maximum head num and liiddenddin and growing only the other dimensions can be an alternative schedule. For ablation study, we struct such schedules, namely SchO-L/B for Bert-large/base, as presented in Table g The GLUE results with Sch0 shown in Table 9
Table 8: Additional growth schedules. "B" and "L" stand for base and large model size, respectively.
Schedule
0-200k
200k-4004 400k
400k-600k (768,4096.16.12))
600k~900k
sch0-L
1768,1024.16.6)
(768,409,,16.6)
(768,4096,16,24))
Schedule
0'200k
400k
400k 600k
600k-700k (768,768,12.12)
700k 900k
sch0-B
1768,768,12.3)
(768,768,1266)
(768,768,1299)
(768,3072,12,12))
H COMPARING WITH LIGO
We take the reported wall-time speedup numbers of LiGO from the Figure in (Wang et 2022) for comparison because the training settings are very close to ours. The speedup of MSG vs. L.GO
w.r.t. baseline models trained from scratch 1.8x vs. 1.69x on Bert-base, and 2.2x vs. 1.82x
on Bert-large. This indicates that both our mathematically strict operators and their data-driven operators can be effective for growth. Again, note that they are partially orthogonal (Section 5.2) because MSG is potentially compatible with arbitrary new weight initialization.
I SANITY CHECK: SCALING TO 100 BILLION PARAMETERS
MSG-oriented growth operators have been applied train LLM by growing from 16B to 51B, and then to 101B parameters (Li et al. 2023). Since a head-to-heer to-d comparison to full-size training is impossible given the computational budgets, we hereby summarize some sanity-check studies in
15

---

Published as a conference paper at ICLR 2024
Table 9: Additional results on GLUE with MSG and SchO-L//B..
Schedule
Operator
GLUE Avg.
CoLAA
SST-2
MRPC
Scho-L (5Gh, 8min)
MSG
12
60.000.8
91 910
80.8007999000033
Seb0-B (17h, 46min)
MSG
80.910.1)
55.0(0.2)
90.7(0.1)
84.30.77888.90..
Schedule
Operator
ST8-B
QQP
MNL(mu/mm)
QNLI
RTF
Sch0-L (36h, Smin)
MSG
589.0.5888.6000
90990.3387.90023
81.700.3y84:60066
91.2(0.3
68.100 18.16 65.5(1 2)
Scb0-B (17h, 46min)
MSG
88.610.2888.1022 90.3001188.110.00 82.4(0.2y82.:(011)
90.2(0.0)
that technical report to show evidences that the training dynamics of MSG are stable for models at 100B+ scales, and the growth strategy successfully leverages the capacities of large models.
Settings. A GPT-like LLM is trained with 24 DGX-A800 GPU (8x80G) servers. The training data is composed of around 300B Chinese-Englishi mixed tokens. Due to limited budgets, (Li et al.) 2023) does not follow the methodologies in Section 3.2.1] Instead, the 16B, 51B, and 101B models are trained with 245.37B, 39.64B, and 26.54B tokens, respectively. This is a very "rapid" growth schedule. The target model hyperparameter is (10240, 40960, 80, 80). The training is completed in 21.54 days with budget of S100K.
Results. As depicted in Figure 4 (picked from Figure (Li al. 2023 ). the training loss curve of MSG is stable in general after each growth, and the loss reduction rate w.r.t. number of tokens gets larger after growth. Besides, according to the evaluation results in their Table 5, the 101B model yields greater improvements than 50B with less data consumed. This indicates that the new weights are successfully incorporated in the training after MSG growth, which consistent with the loss curve. The benchmark evaluations show that comparing with 100B-scale models trained with a similar amount of data (i.e. GLM-130B), the model achieves at least 80% of the baselines' performances with 10% of their training cost. Note that this is not strict comparison since both the models and pre-training data are vastly different. To summarize, MSG-oriented operators show abilities stabilize the post-growth training and inherit knowledge from small models in progressive training of 100B-scale models.
518
Provsse Taneo (Milnn))
Figure 4: Training curve of the 101B model.
J ADDITIONAL GPT RESULTS
As our GPT ppls on Wikitext2 in Taree are larger than those reported in the original GPT-2 work (Radford al. 2019), we doubt that our settings in Appendix D led to insufficient pre-training. Thus, we tried à new setting: we increase the batch size to 480 and train for 100k steps in total, which is around 6 epochs of Open WebTee. The data processing is kept the same as public imple- mentation (NanoGP17). This new setting results in a training loss of 10, which is also close to the value reported by NanoGPT.
We develop a new schedule (Table 10p by our grid-searching methodology (Section 3.2. This schedule reduces the training time 39h, 30min, comparing to 54h, 33min for full-size training. which is 38% speedup. The training loss curves w.r.t steps are presented in Figure 5 The Wikitext2 ppls are 37.65 (MSG) vs 3 (Full), which are comparable. However, there is still significant gap compared to (Radford et al.) 2019) ("30). A possible explanation is that the open-sourced OpenWeb- Web- Text data is different from the close-sourced ones used by GPT-2. Note that this does not influence our conclusions since all our methods are compared with the same data.
hhttps:////t http://gtht....ooo///aaatataaaaa
16

---

Published as a conference paper at ICLR 2024
Table 10: Schedule for additional GPT results.
Schedule
0330k
30k 60k
60k 80k
80k 100k
Sch-GPT-New
(512.768,8.12)
(768,768,12.12))
(768.3072.12.12)
K MSG LOsS CURVES
The MSG training loss curves vs. steps are presented in Figure 5 on the settings in Appendix [
The curves for GPT-2 are based while those for Bert from the main experiments (Section 4.3). We observe that MSG gradually catches up in training loss with progressive growth. We also observe that lower loss in late training does necessarily indicate superior performances (for the full corpus is process for more than epoch), and thus the loss curves are just qualitative indicator of the models' capabilities.
17

---

Published as a conference paper at ICLR 2024
6'"
TTH LOS (80000
3x300
The Steps
105
(a) MSG vs. Full on Bert-base
Sch34
Full
Tha Loss (810h))
8x300
3x300
2x300
Traning Stops
(b) MSG vs. Full on Bert-large
7x300
Sch-
6x300
Pain 005 BRCC7)
4X300
20000
40000
80000
100000
Stops
(c) MSG vs. Full on GPT-2
Figure 5: MSG loss curves. The pink vertical lines mark the beginning of each growth stage.
18


Reviews:
[{"review_id": "PTYVxtxkaH", "reviewer": "Reviewer_Bd3L", "summary": "This paper proposes an interesting mechanism (masked structural growth) to accelerate the pre-training of foundation models. Technically, two main research problems are studied: i) how to determine the optimal growth schedule; and ii) how to design efficient growth operators. Some experiments are conducted to show the advances of the proposed method with a 2.2 times speedup.", "questions": "Please address my concerns listed in the Weaknesses section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper studies a very important research problem in foundation model training. Such a technique can effectively reduce the end-to-end training time of the foundation models. \n\n- Some theoretical analysis was presented to show the guidance of the design in the proposed method. \n\n- The organization and presentation of the evaluation section are clear, with the hypothesis explicitly stated in each subsection.", "weaknesses": "- I am a little bit confused by the introduced method; intuitively, it seems to me that the proposed MSG operator would extend the parameter space to a more rugged space when compared with the average-based methods. Thus I cannot understand why it works well intuitively. \n\n- The experimental results can be further improved:\n  - The scale of the benchmarked model is small. Would the proposed method be deployed for a practice scale model, e.g., 7B to 13B or 13B to 30B?\n  - The reported results can not demonstrate the real effectiveness of the proposed method in terms of ability to generalize, in fact, it seems that the results are mainly based on training loss while the most important metric should be a set of metric to measure the generalization ability of the model.\n  - There should be some plot to show the training loss of the proposed method w.r.t different training iterations, this is essential to understand the performance.", "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": "HEy5jocClz", "reviewer": "Reviewer_vhLH", "summary": "The manuscript focuses on training transformers faster by progressively growing the architecture during training. Problems with previous methods in this area include limited growth \"dimensions\" (e.g. inability to progressively grow the hidden dimension), absence of function preservation in networks with LayerNorm, and a forced choice between function preservation and initialization flexibility. The proposed method (MSG) addresses these problems, and experiments support its effectiveness (i.e., it accelerates training). Ablation and hyperparameter studies further justify the chosen methodology and offer insights that are helpful for understanding the progressive growth problem more broadly.", "questions": "Score-affecting:\n\n1. Please ensure all baseline approaches are faithful so that readers know MSG is effective in high-performing training runs. Examples follow.\n   - In Tables 4 and 5, the SQuAD scores are 2-3% lower than what is reported in the BERT paper.\n   - In Table 5, GPT-2 perplexity on Wikitext2 is high (41.3). There are strong public GPT-2 implementations; e.g., https://github.com/karpathy/nanoGPT. \n      - It would be nice to see an MSG version of the LiGO (Wang et al., 2022) paper's Figure 3C: i.e., show perplexity/loss dynamics of GPT trained with and without MSG.\n\nHelpful:\n\n1. Please summarize the MSG operator conclusions/discussion from the referenced technical report (Li et al., 2023). \n2. Consider moving tables and figures to the page they are discussed on, or to the next page. This helps readability.\n3. Clean up figures and tables. Some example suggestions follow:\n   - Table 2 could have uniform line lengths. \n   - Figure 2 could have larger axis text, lines, legend text, etc.\n\n\nMinor:\n\n1. The manuscript would benefit from a proofreading. Some examples follow:\n   - \"Proof\" is not a verb that means \"to mathematically demonstrate\". Say \"we prove\" instead of \"we proof\". A \"proof\" is what proves something.\n   - On page 6, rephrase the second sentence of the \"Stage Duration\" paragraph -- the training steps are what gets split, not the stage steps.\n2. I think that not all the inequalities in the second lines of equations 11 and 14 are strict (i.e., change a \"<\" to a \"<=\"). \n3. All of the methods listed in the Related Work section (progressive growth, weight sharing, MoEs, etc.) are discussed in the survey \"Compute Efficient Deep Learning\" (Bartoldson et al., 2023).", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The proposed progressive growth method speeds up training more than competing approaches, and it maintains or improves accuracy/perplexity performances on test datasets (relative to the model without progressive growth).\n\nThe study of the impact of different growth dimensions is broadly helpful and clear (Figure 2). This analysis provides intuition for the MSG approach, making it seem less tailored to perform well on just one dataset/architecture (i.e., more generally applicable).\n\nCare is taken to ensure that comparisons across progressive growth methods are fair and that the baseline progressive growth methods are relevant. In Section 4.2, salient features of progressive growth algorithms are isolated appropriately to facilitate analyses of the proposed method, MSG.\n\nThe manuscript offers an original and well-designed approach to disentangling the function preservation and initialization decisions of progressive learning. This facilitates the creation of a more flexible growth algorithm and better analysis of the importance of these algorithm components to progressive learning. For example, the final paragraph of Section 4.4 leverages this disentangling to show that both function preservation and flexible initializations matter to performance, supporting the form of the MSG method.", "weaknesses": "While the manuscript makes careful comparisons to competing progressive growth methods, some of the baseline performances (without progressive growth) are a little lower than expected. Please see my comments about this in the \"Questions\" section below.", "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": "MR5Uxf6EcO", "reviewer": "Reviewer_aWRp", "summary": "The authors propose a framework called Masked Structural Growth (MSG) for progressive pre-training of transformer-based language models. The MSG includes the following two parts:\n(1) growth schedules involving all the four possible dimensions;\n(2) strictly function-preserving growth operators that are independent of the initialization of new weights.\n\nTheir experimental results show that the MSG achieves state-of-the-art speed-up on Bert-base, Bert-large, and GPT-2 with equal or improved performances on downstream tasks.", "questions": "I list my questions and some minor issues as follows:\n(1).\nFor on page 4 and page 5:\nAre vectors x, y, c in your section 3.1.1 all column vectors? If yes, then in Eq. (12),\nx^T (the transpose of x) is not necessary, x will be good enough.\n\n(2).\nFor Eq. (11) on page 5:\nFor d_2 < i < d_2', it should be:  d_2 < i <= d_2'\n\n(3). Similarly,\nFor Eq. (14) on page 5:\nFor d_1 < j < d_1', should be: d_1 < j <= d_1'\n\n(4).\nLast line on page 7:\nThe stated (Qin et al., 2022b) is not the reference for the bert2BERT method.\nDo you really mean the bert2BERT or the ELLE here?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The greatest strength of the paper is that it introduces the MSG operators for all growth dimensions by using external masks.\nThese growth operators are strictly function-preserving (even with layer normalization), and independent of the initialization of new weights.\nThis apparently has the advantages over the growth operators proposed in the previous work.\n\nAdditionally, the growth schedule proposed in the MSG uses grid-search to grow only one dimension each time, which provides empirical insights on the growing models and holds the potential for generalization.", "weaknesses": "One weakness of the paper is that the authors have not empirically shown the advantages of the MSG in today's large-scale transformer models (which is much larger than the scale of BERT-large and GPT2). Such a weakness may be due to the lack of the computational resources.\nIn terms of presentation, the authors may want to add more description on the bert2BERT method, since it is one main competitor to the proposed MSG.\n\nSome minor issues will be listed in the \"Questions\" section.", "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)