Learning to Grow Pretrained Models for Efficient Transformer Training
Authors: Peihao Wang
Rameswar Panda
Lucas Torroba Hennigen
Philip Greengard
Leonid Karlinsky
Rogerio Feris
David Daniel Cox
Zhangyang Wang
Yoon Kim
Venue: ICLR 2023 notable top 25% (2023)
Keywords: Transformer
Efficient Training
Model Reuse

Abstract: Scaling transformers has led to significant breakthroughs in many domains, leading to a paradigm in which larger versions of existing models are trained and released on a periodic basis. New instances of such models are typically trained completely from scratch, despite the fact that they are often just scaled-up versions of their smaller counterparts. How can we use the implicit knowledge in the parameters of smaller, extant models to enable faster training of newer, larger models? This paper describes an approach for accelerating transformer training by learning to grow pretrained transformers, where we learn to linearly map  the parameters of the smaller model to initialize the larger model. For tractable learning, we factorize the linear transformation as a composition of  (linear) width- and  depth-growth operators, and further employ a  Kronecker factorization of these growth operators to encode architectural knowledge. Extensive experiments across both language and vision transformers demonstrate that our learned Linear Growth Operator (LiGO)  can save up to 50% computational cost of training from scratch, while also consistently outperforming strong baselines that also reuse smaller pretrained models to initialize larger models.

Published as a conference paper at ICLR 2023
LEARNING TO GROW PRETRAINED MODELS FOR EFFICIENT TRANSFORMER TRAINING
Peihao Wang Rameswar Panda
Lucas Torroba Hennigen" Philip Greengard? Leonid Karlinsky² Rogerio Feris" David D. Cox2 Zhangyang Wang' Yoon Kim
'University of Texas at Austin, "MI--IBM Watson AI Lab, 'Columbia University, "MIT aataooanng,atlssaaan (awpwangggeteeeaaaeeeee edu, {rpanda, leonidka, david, d. d. cox rs rsferis@us.. wiceusssbmmcommmm wtt.miicoom....................e com, pg21 @ecolumbia.edu, {lucastor, yoonkim]@mi @m.t..
com,
ABSTRACT
Scaling transformers has led to significant breakthroughs in many domains, lead- ing to a paradigm in which larger versions of existing models are trained and released on periodic basis. New instances of such models are typically trained completely from scratch, despite the fact that they are often just scaled-up ver- sions of their smaller counterparts. How can we use the implicit knowledge in the parameters of smaller, extant models to enable faster training of newer, larger models? This paper describes an approach for accelerating transformer training by learning to grow pretrained transformers, where we learn to linearly map the parameters of the smaller model initialize the larger model. For tractable learn- ing, we factorize the linear transformation operators, and further employ a Kronecker factorization of these growth operators to encode architectural knowledge. Extensive experiments across both language and vision transformers demonstrate that our learned Lin- ear Growth Operator (LiGO) can save up to 50% computational cost of train- ing from scratch, while also consistently outperforming strong baselines that also reuse smaller pretrained models to initialize larger models [
composition of (linear) width-
and
INTRODUCTION
The transformer architecture (Vaswani al. 2017) has emerged as a general purpose architecture for modeling many structured domains (Devlin ii et al. 2019; Brown et al. 2020 Rives et al., 2021]
Dosovitskiy al.] [02T]] TTuvron et al. 222 Perhaps more so than other architectures, the transformer empirically seems to have inductive biases that make it especially amenable to scaling (Rosenfeld al. 2019 Kaplan et al. 2020), which has led to a paradigm which larger versions of smaller, existing models are trained and released
periodic basis (c.g.. the GPT lineage of models (Radford et al. 2018 2019 Brown et al. 20200). ). New instances of such models are typically trained completely from scratch, despite the fact that they are often scaled-up versions of their smaller counterparts. Given the compute required to train even the smaller models, we argue that training each model from scratch is wasteful, and that prior knowledge implicit in the parameters of smaller pretrained models should be leveraged to enable faster training of larger models.
One approach to this problem is through the lens of model growth, wherein a smaller model's pre- trained parameters are used to initialize a subset of the larger model's parameters. While earlier works generally froze the parameters initialized from the pretrained model and only trained the new (randomly initialized) parameters (Fahlman & Lebiere 1989) Fahlman [1990; Gutstein et al.) 2008), subsequent work has shown that copying a subset of the pretrained parameters to initialize the new parameters and then finetuning the entire network significantly accelerates training and sometimes even leads to better performance (Chen et al. 2015). When applied to modern transformers, these mechanisms roughly translate to depth-expansion operator in which pretrained models are stacked (or combined with identity layers) to initialize deeper transformers (Gong et al., 2019 Yang et al.) 2020), and width-expansion operator in which the smaller model's 's matrices are copied to initialize the larger model's 's matrices (c.g., in block-diagona fashion) (Chen et al. 2021 Gu et al. 2020).
"Work done during an internship at MIT-IBM Watson Al Lab. Projecc page: https: htmp:://vitas htt/:/vviaaaaaeeee group. gi io/Li///

---

Published as a conference paper at ICLR 2023
expand
We med ao Was
Wi WOdT
W,
W;
W
W!TH B.W
Wi B.W.
Depth Expansion
Width Expansion
Wi.1
A₁ B
LiGO voc (e(nnc))
:
@
Au. @ B
Eeepth
Regidth
Figure 1: Our linear growth operator (LiGO) accelerates training by using the weights of a smaller model + initialize the weights of the larger model e (new), LiGO is parameterized a sparse linear map M that can be decomposed into width- and depth--cpansion operators. The width-operator Rucd and depth-operator Leepth are structured matrices obtained from Kroneeker products of smaller matrices which encode architec- tural knowledge by grouping parameters into layers and neurons. While we show the expansion operators for simple multi-layer perceptrons for illustrative purposes, in practice apply LiGO to enable faster training of transformer networks. In our approach, learn the growth matrix M with 100 steps of SGD, this initialize the larger model, and then continue training as usual. Best viewed color.
Noting the empirical effectiveness of such recipes, we observe that existing mechanisms generally do not have a learning component (e.g., randomly copying over neurons for width-expansion or stacking consecutive layers for depth-expansion). This paper instead proposes efficient, data- driven approach for learning to grow transformers. In particular, our approach frames the problem of initializing the larger model's parameters learning linear mapping from the smaller model's M (small) where θ (amalt) and @ (large) are the vectorized parame- ters of the small/large models. Due to the high dimensionality of the parameters, this mapping is completely intractable to learn without any restrictions on M. We thus factorize the linear mapping to be composition of sparse width- and depth-expansion operators, M = Leepth Roiath where both width and depth matrices are further factorized to be a Kronecker product of smaller matri- ces that express architectural knowledge (c.g.. through grouping parameters by layers and neurons). We show that our growth operators can represent existing approaches such as layer-stacking and neuron-copying special cases. We find that with small amount of learning on M (e.g., 100 gradient steps) to initialize the larger model, we can significantly accelerate training of both vision and language transformers. Figure I illustrates our approach.
parameters, i.e., θ (lerge) MO
We apply our learned linear growth operator (LiGO) to popular families of models- BERT (Devlin al. 2019), RoBERTa (Liu et al. (2019), GPT2 (Radford et al. 2019),, and ViT (Dosovitskiy al.) 2021 Touvron et al. 2021ap b) and find that LiGO can consistently improve transformer training efficiency over the traditional way of training from scratch across domains and model sizes. For instance, LiGO saves 44.7% and 22.5% FLOPs for training BERT-Base and GPT2-Medium from scratch by reusing pretrained smaller models that are half big. Similarly, for vision transformers, when using DeiT-S (Touvron et al. 2021a) for initialization, LiGO yields 55% savings in FLOPs with no performance drop on ImageNet (Deng et al., 2009). These FLOPs savings directly trans- late to similar wall clock savings. We further find that models trained using LiGO achieve similar performance to the trained- tuined-from--oe fcoatscratch baselines when transferred to downstream tasks.
2 RELATED WORK
Efficient training. Efficient training of transformers has been studied from multiple perspectives. Some methods that are orthogonal our work include mixed precision training (Shoeybi et al., 2019). large batch optimization You et al. 2019), distributed training (Huang et al. 2019), and dropping layers <Zhang & He. 2020) or tokens (Hou et al., 2022). Knowledge inheritance QQn et al. 2021) explores knowledge distillation during pretraining to efficiently learn larger transform- ers. Progressive training, which first trains a small transformer with few layers and then gradually expands by stacking layers, has also been applied to accelerate transformer training (Gong et al.) 2019 Yang et al. [2020] Li et al.] 2022 Shen et al.) 2022). 3. Net2Net Chen et al. (2015) uses function- preserving transformations
grow width by copying neurons and depth by using identity layers. Recently, bert2BERT (Chen et al. 202T) extends Net2Net to transformers. In contrast to these ap- proaches, approach learns (liearlyy) transform the parameters of smaller model to initialize
2

---

Published as a conference paper at ICLR 2023
larger model. While there is line of work on learning grow neural networks in data-driven way, these methods are in general difficult to apply to modern-scale transformers since they (for example) involve growing single neuron a time employ expensive optimization/search procedures (Wei 2016, Cai et al. 2018: Wu et al., 2019 2021 Evci et al., 2022).
et al.
Network initialization. Our work is also related work on neural network initialization. Exist- ing works include controlling the norm of the parameters (Mishkin & Matas 2015; Kilcher et al.
2018 Dai al. 2019, [Wu et al., 2019 Glorot & Bengio, 2010) or replacing the normalization lay-
ers (Brock ct al. 2hang et al., 2 Huang et al.J 2020p. Metalnit (Dauphin & Schoenholz
2019) proposes an automatic method that optimizes the norms of weight tensors to minimize the gradient quotient quotient on minibatches of random Gaussian samples. GradInit (Zhu et al. (021) learns to initialize larger networks by adjusting norm each layer. Our work focuses on using smaller pre- trained transformers to better initialize larger transformers, which remains an understudied problem.
Structured matrices. Finally, our work is also related to structured matrices which are typically used to replace dense weight matrices for reducing training and inference computation cost. Ex- amples include sparse and low rank matrices (Chiu et al.] [002]] Han et al.) 2013), Chebyshev matrices (Tang al.] 2019), Toeplitz matrices (Sindhwani
al. 2015), Kronecker-produet ma- trices (Zhang et al. 2015), and butterfly matrices (Dao et al. 22019). 9) A unified framework to learn broad family of structured matrices is presented in Sindhwaai al. (2015). [ao et al. (2022) propose Monarch matrices, which inherit the expressiveness of butterfly matrices and achieve rea- sonable accuracy-efficiene tradeoffs in many applications. While our approach is inspired by these works, we propose to grow pretrained models by learning structured sparse linear operators with Kronecker factorization, which to our knowledge has not been explored in the literature.
3 PROPOSED APPROACH
Notation. We denote the parameters of a neural network with layers and D dimensions as Wz] z) E RLDDD where W, E DxD denotes the weights for the I-th With slight abuse of notation, we denote the vectorization of i as ve = vec(W]. L) Our goal is to re-use the parameters θ = OL,D, from pre- trained smaller model to initialize large model θ (new) O Lg,D, through model growth operator M : LI - IR. b2xxxxx D2 that maps the weights of the smaller network to the weights of the larger one, i.e., (new) M()) where L L2 and D1 < D2- After model growth, adopt + (new) as the initialization of the large model and train it using standard recipes.
LLD = W
layer 2
vee W
3.1 EXISTING GROWTH OPERATORS
Existing works have separately established model growth operators for depth (I1 < L2, Dj = D2) and width (I1 = T22 D1 D2). We summarize these methods below.
Depth expansion. StackBERT (Gong et al. 2019) proposes to duplicate the smaller model to double the depth, based on the observation that upper layers share similar functionality with the lower layers. In contrast, interpolation- based depth expansion methods (Chang et al.) 2017 Dong et al., 2020) interleave every layer to form a deeper model, which can be roughly interpreted as simulating giner-grained solution to the original dynamical system from neural ODE perspective (Chen et al. 2018). Letting L2 = kLL, the two methods' growth operators can be formulated as:
(StackBERT) W ['new) W mod L,, (Interpolation) (new) W WWEK
[
(I)
Width expansion. Net2Net (Chen et al. 2015 expands the width of neural networks by randomly copying neurons while preserving output values via normalization. This can be seen as growing matrix associated with a particular layer by duplicating the columns and rows of its weight matrix.
Suppose a layer has weight matrix W, E D
To expand it to matrix W (new) E D2
FFor notational brevity we assume that each hidden layer has same number of dimensions D, but LiGO can be straightforwardly generalized to layers with different dimensions (e.g., FFN layers transformers).
3 We therefore have vec vec(@ (L.D)" E TR L Our approach also agnostic with regard to vectorization order. Wee define single layer as f()) Wa + by, where the row number of W, corresponds to the output dimension, and the column number of W, corresponds the input dimension.
3

---

Published as a conference paper at ICLR 2023
(D2 > Di). Net2Net copies W its upper-left corner of (neww) fills the new columns via random selection matrix S₁, and finally duplicates and normalizes rows according to the selection matrix from the previous layer. Formally, the growth operator of Net2Net can be written as:
(Net2Net) W (new)
I S DJ w Si], D, = diag(S
E [42]
(2)
where S (0,1) is random selection matrix. The diagonal of D; is D1- dimensional histogram, whose i-th entry indicates numher of times i-th column of W was copied.
3.2 LEARNING TO GROW WITH A STRUCTURED LINEAR GROWTH OPERATOR
While existing operators have been empirically successful in accelerating transformer-basee models such BERT (Gong et al. 2019; Chen et al. 202T). observe that generally do not have learning component and perform the depth- and width-expansions separately. In this section we introduce general framework for learning grow with linear growth operator (LiGO), which generalizes existing operators by combining the width- and depth-grown operators in a data-driven way.
We can formulate the problem of initializing the weights of the larger model (new) from the smaller model O through the following optimization problem,
arg min EE L(x; + (new)), subject θ (new) .())). M
where D is the data distribution and L is the loss function. Is of course intractable to optimize over the entire operator space, and thus we further simplify the function M to be a linear transformation, which results in the following formulation,
vec( O (new)) = vec(M(()) = M vec(O), M R
(4)
This simplified objective is still completely infeasible to apply to contemporary neural networks where L1 D can easily be in the hundreds of millions. We therefore propose an efficient parameter- ization of M for tractable learning.
3.2.1 DECOMPOSITION ALONG DEPTH AND WIDTH
Our first step is to decompose the LiGO operator M = Leepth Ruiithh, where Leppth and Ruiidth expand the depth and width of model separately. Concretely, we decompose M as
diag(E1.1)
diag(L1,L.)
R
M
(5)
diag(Llz.1) diag(ELn.te
RL
Lacpen
Rwidth
where RG E Dj and lúj E D2 In the above, Leppth is an array of diagonal matrices and Riidh is block-diagonal matrix, i.e., both matrices are highly structured and sparse. When applying Rwsstth to weights vec(O), the parameters of each layer will be transformed independently via vec(W (new)) (neww))
R₁ vec(W,) and lifted to higher dimension. The (-th row block of Ldepth corresponds to the growth operator of 1-th layer, which amounts to linearly combining all layers ΣLL 1(rr)) vee (Wi) By this factorization, to D?LL) and encode architectural knowledge by grouping parameters by layers. Later Section 3.4 this representation is also shown to preserve high representation power owing to its connection with Monarch matrices (Dao et al. 2022 2019).
of the smaller model via veo(W<(ng (new)) )
can effectively reduce the complexity of the LiGO operator from L,
3.2.2 PARAMETER SHARING VIA KRONECKER FACTORIZATION
The above LiGO operator requires O(D)D3L, parameters for Rwitth and LLL for Ldepth- The width operator Rovidth is thus still prohibitively expensive given that D1 (and Dz) can easily be in the hundreds or thousands. In this section, we propose a Kronecker factorization to further reduce the number of learnable parameters for each growth operator.
4

---

Published as a conference paper at ICLR 2023
Depth. For depth, treat entire layer as single group and construct new layer by combining existing layers, effectively tying parameters for all neurons in same layer. Formally, each block in Eeopth is simplified to be (iag(L..)) w Then the entire matrix can be written Kronecker factorization, deppth
@ I, where E IR 1 L1 is matrix whose entry Wa,j indicates blending weights of j-th layer of the small model to form i-th layer of the large model. This strategy reduces the number of parameters in Eeepth (L) and is shown on left-hand side of Figure !
Width. For width, we decompose each diagonal block of width expansion operator Rwidth using the Kronecker factorization R. A Bi, where A,, Br E RDxxD: Since vec(CAAB)
(B" @ C) (e)(A) Schacke 2004). we then have,
A & B
Rwidth vec( )
(
(6)
An B vec [BIWIA] ... BL, W,,W, B,,W,, A
(7)
Here observe that B W A performs in- and out-dimension expansion by A and B,, respec- tively. Each new column/row is linear combination of columns/rows of small model's weight matrix. This factorization, which can be seen as grouping parameters by neurons, reduces the num- ber of parameters to O D D2). Figure (right) illustrates LiGO's width-eepansion operator.
Altogether, we obtain the final parameterization of LiGO operator M:
W1,1
11,2
W1,4,
A & B
M
X
(8)
WL,11
WL,22
WL Ln
AL & BL,
Depth expansion
Width expansion
We can exploit the factorization implement the LiGO operator (Eq. 8 efficiently.
Training. LiGO expands a model in three steps: (1) for each layer, inserting new rows by lin- early combining existing rows through B;, (2) for each layer, inserting
columns by linearly combining existing columns through A,, and then finally (3) reconstructing each layer by linearly combining the weight matrices with w along the depth. We then run a few steps (c.g.. 100 itera- tions) of SGD to optimize M, which has negligible compute cost relative to regular training. After obtaining M, we initialize large model with M vec(O), (Q), and train parameters θ (new) through SGD as usual. Algorithm [ summarizes a forward pass of LiGO with transformer. Finally, as shown in Appendix A we note that StackBERT (Eq. [p. Interpolation (Eq. [ and Net2Net (Eq. 2 are all special cases of LiGO (Eq. g with a particular setting of Leppt and Rwidth-
3.3 LIGO FOR TRANSFORMERS
While LiGO can be applied to any multi-layer neural network architecture, in this paper we focus on using LiGO to grow transformers which have been shown to be particularly amenable to scaling. Below we briefly describe how LiGO applied to the main transformer embedding/attennion layers and defer further details (e.g., growing bias vectors, layer norm parameters) to Appendix B.
Embedding layer. The embedding layer can regarded as a linear layer whose inputs are one-hot vectors. We learn a matrix B (emb) to extend its output dimension. This embedding layer is also used the final output layer for our transformer language modeling experiments.
Attention and feedforward Layers. An attention layer consists of multi-head attention weights (Wº, k V and a linear projection (W"). Let A and B where & E (Q, K, V, 03 be the th layer's in- and out-dimension expansion matrices (Eq. for the query, key, value, and projection matrices. To make sure new input and output channels are aligned across modules, we tie the LiGO
operator as follows: for all / [1]. (1) A = (B(mm))) for Vk E {Q. V}. (2) A = (BY)T. B (emb)) The last constraint is added to take into account the residual connections (Cen
(3) B
2021 We similarly tie parameters for the feed-forward networks, A [fel) (B(emt)
ct al.
A
1/22
(B(al))] and B(/e2) (emb)) Since transformers make heavy use of residual layers
5

---

Published as a conference paper at ICLR 2023
with skip connections, we found that simply using the same (emb)) to parameterize A and B for many layers/modules worked well in practice. This reduces the number of learnable parameters even further and enables fast learning of M small amount of data (100 gradient steps).
3.4 CONNECTION TO MONARCII MATRICES
As shown in Section 3.2.1 our depth- decomposition factorizes M into a multiplication of two structured sparse matrices. We examine the expressiveness of this factorized representation by relating it to Monarch matrices (Dao et al. 2022). defined below.
Definition 1. Let the space of Monarch matrices be M S Xman, Then matrix M E M if M = P, LP2 R = P1 diag(Ln,".. P diag( Ri,,... ... ,R,)) where L, E RRbibb,, R E
bbbb are dense rectangular matrices, and njb2 = n2b3.. P1 is the permutation (i) (i
b1 b 1)m1 + by
and P2 the permutation π T(]) = (j balbb1 1)m1 1bg +1
It is clear that the block-diagonal matrix R has the identical form to our width growing operator Rwidth- By applying the permutation matrices P1 and P2 to L. is transformed into exactly the same form with our depth-growth operator deepth in Eq. [5] This implies that our depth- decomposition coincides with Monarch sparsification of dense matrices, which generalize butterfly
matrices (Dao et al. 2019) and enjoy rich expressivity properties (Dao et al. 2020 2022).
4 EXPERIMENTS
We conduct experiments to answer three key research questions. Q1: To what extent LiGO improve the training efficiency (FLOPs and wall time) of transformers compared to training from scratch and other growth operators? Q2: Can LiGO be universally effective across transformers from different domains (e.g.. language and vision) and sizes? Q3: Can models trained using LiGO achieve similar performance compared the baselines when transferred to downstream tasks?
4.1 EXPERIMENTAL SETUP
Datasets. We follow Tan & Bansal 2020) and use the English Wikipedia corpus for training BERT (Devlin al. 2019) and RoBERTa (Liu al. 2019). We use the public C4 (Raftel al. 2020) dataset for training GPT2 (Radford et al. 2019). We use ImageNet (Deng et al.] 2009) for training vision transformers. We use GLUE (Wang al.) 2018). SQuADv1.1 (Rajpurkar et al.] 2016), and SQuADv2.0 (Rajpurkar et al. 2018) for evaluating pretrained BERT models. We test downstream performance of vision transformers (DeiT (Touvron et al.) (221aa) by performing trans- fer learning on 5 downstream image classification tasks, including CIFARI0 (Krizhevsky et al.)
2009), CIFAR100 (Krizhevsky et al. 22099) Flowers 102 (Nilsback & Zisserman 2008), Stanford-
Cars Krause et al. 2013), and ChestXRay8 (Wang et al. 2017).
Models. We experiment with growing the following language andd vision transformers: (1) --a---mall-------------------- (2) RoBERTa- Small-RoBBRRT--Bes for RoBERTa; (3) Serr7-Basc-GGTT-MMeeee
BBRRT-Smlll-BERT-PP
(4) DciT-S-DDiTT and (5) CaiT-XS->Caill BERT-Small has 6 layers with 512 hidden dimensions, while other named models are their usual sizes. See Appendix B.2 for full details.
Baselines. We compare our approach with the following baselines: (1) training from scratch baseline where we train the larger transformer without using any smaller pretrained models; (2) progressive training methods designed for growing depth in transformers (StackBERT (Gong al. 2019) and
MSLT Yang et al. 2020)): : (3) ber12BERT (Chen et al. 2021) that extends Net2Net (Chen et al.)
2015 for width expansion and stacking for depth expansion; (4) KI (Qin al. 2021 which uses distillation for transferring knowledge from the smaller model to the larger model.
Implementation details. We always use 100 gradient steps to learn the LiGO for all models, which is negligible in terms of FLOPs/wall time compared to full training after initialization. We train both BERT and RoBERTa models for 400K steps with a warmup of 10K steps. We remove the next- sentence prediction task (Liu et al. 2019) and use fixed sequence length of 128 for pretraining
While the original BERT Devlin et do not include here since is no longer publicly available.
2019) paper also uses the Toronto Book Corpus 7hu et al. 2015

---

Published as a conference paper at ICLR 2023
3.0
3.0
3.0
Scratch StackBert MSLT
Scratch
KI
Scratch
2.0
beIZBER!
USp
2.8
NSLT
Stac thte
UGO
bei2
2.8
usa from Smell
USO from Base
Perplexity.
2.6
Log Perplextty
2,6
Log Perplexity
2.6
2.4
2.4
180
2.4
1.40
2.2
:09
3:00
2.2
2.0
2.0
2.0
1.8
188
1.8
0
0
20
40
60
80
15
FLOPs (1e18)
Well Time (hrs)
FLOPs (1e18)
(a) BERI--mall--Biine -ERRT-Bas
(b) SEaRT-SSall-BBERT-E
(c) BERT- (BERT-(SSall, Buse) - -BBRT-LLrge
Figure 2: Results on BERT. (a-b) shows validation log perplexity FL.PPs and wall time respectively for training BERT-Base by reusing BERT-Small. (c) shows log perplexity vs. FL.OPs in growing BERT- Small and BERT-Base to BERT-Large. The solid line indicates the final perplexity the larger model trained from scratch, while the dotted line represents performance of the smaller model trained from scratch. LiGO offers about 45% savings in FLOPs and 40% savings in wall time over BERT-Bas Base training from scratch. Our approach is also flexible in reusing either BERT-Small BERT-Basc for accelerating BERT-Large training.
Table 1: Downstream transfer learning performance on GLUE and SQuAD. All of the results are based on B-sse models trained using the different baselines. LiGO achieves similar or even hetter performance than the original training from scratch baseling on several downstream tasks, despite improving training efficiency.
Method
Savings (FOOS)
Savings (Walltime)
SST-1
(Acr.)
MINII
(Acc.)
MRPC
(Axc.)
CoLA
[Acc.))
ONLI
(Acc.)
(Acc.)
QOP
STS.8
(Acc.)
SQuADvLI
(FUEN))
SQBADvI0 (FLEM)
GLLE
Avg.
SQUAD
Avg.
Scratch
83.17
75.41
K5.78
62:10
87.35
57.1%
$6.39
86.57 21
113175707
82.35
78.977 72.19
SHEkBERT
54.1%
33.3%
88.99
79.72
85.29
50.09
87.28
89.17
96.97
86.50
77.42
71.32/67.41
82.36
T8.91 72.41
MSLT
34.9%
30.0%
$8.53
76.10
82.60
64.76
83.58
83.54
85.89
86.07
/ 76.73
70.6//7.11
81.72
73:47/ 71.95
KI
-5.7%
13.9%
38.65
78.33
83.50
64.86
86.25
$3,96
57.09
84.93
/ 26
71.99 099//50 57.41
82.59
78.01 / 71.85
NIRRERT
20.09
25.1%
88.30
80.05
85.54
61.73
88.
86.18
87.00
34
/ 77
71
71.52/66.85
82.42
7888/77 71.97
LIGO
$4.7%
20.7%
35.12
79.29
34.31
62009
35.07
57.00
85.25777:45
71 577
82.57
78.76/
both models. For BERT, we use a batch size of 256 and a learning rate of 2e -4, while we use a batch size of 1024 and a learning rate of 8e for training RoBERTa models.
Following Shen et al. 20022), we train GPT2 models with a batch size of 384 and sequence length of 1024. For vision transformers, we build our models based on DeiT (Touvron et al., 2021a) and CaiT (Touvron et al. 2021b), and apply their default hyper-parametert for training on ImageNet dataset. We train all our vision transformers for 300 epochs with a batch size of 1024. For transfer learning with BERT/RoBERTA. we follow Tan & Bansal (2020) and train for 3 epochs with learn- ing rate of 1 and a batch-size of 32 for all tasks in GLUE. On SQuAD vI.1 and SQuAD 2.0, we fine-tune for 2 epochs with a learning rate of 5e and batch size of 12. We run both GLUE and SQuAD evaluations three times with different random seeds and report the mean numbers. For transfer learning experiments on DeiT, finetune the pretrained models with 1000 epochs, batch size 768, learning rate 0.01, and use the same data augmentation in training on ImageNet. We use the same pretraining data and experimental settings for all the baselines (including our approach) for a fair comparison. Note that we include the additional compute required for training LiGO in all our tables and figures. However, since our LiGO is only trained for 100 steps, the influence on visualization and quantitative saving percentages is negligible.
4.2 RESULTS AND ANALYSIS
BERT. Figure shows the comparison between the different baselines for training BERT models. As seen from Figure aa, LiGO saves 44.7% computational cost (FLOPs) of training BERT-Base (12 layers, 768 dimensions) from scratch by reusing BERT-Small layers, 512 dimensions). LiGO offers 40.7% savings in wall time compared to training from scratch (Figure 2b)) Among the com- pared methods, StackBERT is the most competitive in terms of both FLOPPs and wall time, although LiGO obtains 1006% and +7.2% improvements in FLOPS and wall time on top of StackBERT. Similarly, LiGO significantly outperforms the recent bert2BERT method which saves about 30% computational costs. We observe that KI does not provide any real savings in training as it requires additional computation for knowledge distillation. Figure c) shows that our LiGO approach is flexible in growing either BERT-Small or BERT-Base for accelerating BERT-Large training. As ex- pected, reusing BERT-BE Base instead of BERT-Small Small leads more savings in FLOPs (45.2% vs 30.3%) as BERT-Base contains more implicit knowledge in its parameters. Table T shows the per-task per-
7

---

Published as conference paper at ICLR 2023
3.0
3.0
4.0
Scratch StackBert MSLT
Scratch
Scaach
uG6
2.8
beIGBERT
UÇO
2.8
Stac titer
MSLT
UGO
ber12BKK
3.8
Perplexity
2.6
1-0
Log Perplextty.
2.6
3.9
Log Perplexity
2.4
2.4
180
2.2
22
1.89
$1.3%
9.0%
2.0
2.0
3.2
1.8
18
0
20
40
60
80
3.0
FLOPs (1e18)
Well Time (hrs)
FLOPs (1e17)
(a) LOBERTa-Smale -Roeil--SSmm----BBB--------- Bos
(b) RooLaaaalll----R------
(c) GPT2-Base GGPT---Mdiim
Figure 3: Results on RoBERTa and GPT2. LiGO reduces FI.OPs by 47. 47.2% and 22.5% for RoBERTa-Base
and GPT2-Medium, demonstrating its effectiveness across different training strategies and architectures.
formance of different BERT-Base models on both GLUE and SQuAD benchmarks, where we find that BERT trained with LiGO achieves very similar performance compared the baselines both benchmarks. Finally, in Table 3 of the Appendix C.3 we show that growing BERT-Small to BERT- Base with 100 steps of LiGO and then finetuning on GLUE tasks without additional pretraining outperforms just directly finetuning SERT-Small.
RoBERTa and GPT2. Figure B(a-b) shows the results on RoBERTa, whose training recipe uses larger batch size and learning rate than BERT. LiGO similarly accelerates RoBERTa training, which indicates that our method is robust to optimization hyperparameters. On GPT2, LiGO saves 22.5% computation cost of training GPT2- rameters) (Figure These consistent improvements show that LiGO is effective for accelerating transformer training across different model architectures and sizes.
(345M parameters) by reusing GPT2-Base (117M pa-
Vision Transformers. Fig- ure 4 shows that by grow- ing from DeiT-S, LiGO can save 55.4% FLOPs and 52% GPU wall time to reach the same performance of 81% 31% ImageNet. Interestingly, the model initialized by our data- driven growth operator (w/ only 100 gradient steps of tuning) can already achieve 72% accuracy at the begin- ning of training and leads to final accuracy of 81.7% at the end of the training. Compared to the next best method, bert2BERT, LiGO obtains more than 15% savings, which once again demonstrates the effectiveness of our approach in growing vision transformers well. Table shows that finetuning results on downstream tasks perform on-par with the model trained from scratch, showing that LiGO does not harm the model's generalization capabilities when transferred downstream datasets. We also find similar savings in CaiT-XS-Caile where LiGO saves FLOPs by 52.6% and wall time by 46.1% over training CaiT-S from scratch on ImageNet (see Appendix C for more details).
80
75
Valid Accurary
5.-45
0.0%
53.0%
0.75
70
70
65
Valid
83
60
68
55
55
50
6
10
12
10 15
20
25
30
35 40
FLOPs (1e18)
Wall Time (hrs)
(a) B-DiTBB
(b) DeiT B
Figure Results on DeiT. (a) Accuracy vs. flops (b) accuracy vs. wall time for training DeiT-B. LiGO saves flops and wall time by more than 50% over training from scratch ImageNet.
Combining with other training strate- gies. We also find that LiGO can ef- fectively combined with orthogonal strate- gies such as layer dropping (Zhang Hel 2020), token dropping (Hou et al. 2022). and staged training (Chen et al. 2021 More details included in Ap-
Table 2: Transfer learning performance of DeiT-B. DeiT-B model trained using LiGO performs similarly to the original trai-ffrom-scratee baseline on all downstream tasks.
11A.
Walfim
Imaggaee
CIPANER
Plaurrs
Lam
141.5
10.73
MUT
36.7%
58.92
81 30
06.03
00.31
00:21
08.11
7M
WOTBET
15/04
06.99
18.15
15.00
TRR
15
15774
01.11
160
0671
10.79
19180
pendix B33 Figure shows that LiGO can be combined with other training techniques improve the computational savings by 4.7%, 7.4%,
8

---

Published as a conference paper at ICLR 2023
3.0
3.0
3.0
Scratch
Scratch
2.0
UGO wi Layer Oropping
Lisa
2.8
LGG
L660 w Toxen Dropping
2.8
COMBR LAD
Perplexity
2.6
1.6
Log Perplextty
2.6
100
Log Perplexity
2.6
2.4
2.4
2.4
139
2.2
22
2.2
2.0
2.0
2.0
1.8
1.8
1.8
0
0
1
FLOPs (1e18)
FLOPS (1e18)
FLOPs (1e18)
(a) LIGO w/ Layer Dropping
(b) LiG0 w/ Token Dropping
(c) LiGO w/ Staged Training
Figure 5: LiGO with other efficient training strategies. Our approach can be combined with (a) layer dropping, (b) token dropping, and (c) staged training (ST), for further accelerate BERT training.
and 8.2% with layer dropping, token dropping and staged training, respectively. Following (Chen al., 2,2TT, we also apply staged training strategy to bert2BRT and observe that LiGO still out- performs bert2BERT with staged training by 16. 7% (see Figure ((c).
4.3 ABLATION STUDIES
Seratch Ialeder
MXT
Sojech
Depth-only expansion. We examine the effectiveness of proposed depth pansion operator (Ldepth) by only grow- ing the depth of BERT from 6 layers 12 layers, i.e, (BERT(6, 768) -BERT12, 768). We compare with stacking (Stack- BERT, Gong et al. 2019). Interpolation (In- terBERT, Chang et al.) 2017 Dong et al. 2020 (see Eq. 11. and MSLT Yang et al.
P
2,0
FLOPs (tedil)
TLOPA (1
(a) BERT(6 -6R)-R 768) (b) BFRTI12 12
2020 1. For LiGO, we only apply its Leepth component to the pre-trained model weights. Results in Figure(()a) show that data-driven approach works well even when just growing across the depth dimension.
Figure 6: Results Depth-only and Width-only growth. LiGO saves 51.7% FLOPS when expanding depth-only, and 41.6% FLOPS when expanding width-only.
Width-only expansion. We also verify the effectiveness of Renitth by only extending BERT width from 512 to 768, i.e., BERT(12, 512)-BBRR(122, 768). We compare LiGO based initialization with
Table 3: Effect of number of gra- dient steps. "+FLOPs" stands for additional flops (in 10-.)
direct copy (Wei al.] 2016), function preserving initialization (FPI, Chen et al. 2015). and advanced knowledge initialization AKKinChen&ee al. 2021). LiGO's width expansion component out- performs all other methods, as shown in Figure(6(b)
# of Steps
+FLOPs
Savings
100
3.61
44.7%
500
18.06
44.5%
1000
36.13
44.2%
10000
361.30
38.9%
Number of growing steps. Our main experiments just use 100 gradient steps to grow. We tune our LiGO on the pretraining set for 100, 500, 1000, and 10000 steps and compute the additional FLOPs for training LiGO within 1000 steps results in the identical model convergence (reaching 1.8 PPL at 215K steps). This suggests tuning model weights under the linear constraints of LiGO can achieve faster convergence. Training LiGO for more than 10000 steps can provide a model with slightly faster convergence (214K steps), but results in less saving overall.
training. Table shows that
CONCLUSION
This paper describes an approach for accelerating transformer training by learning to grow pre- trained transformers, where the larger transformer's parameters are initialized as a linear mapping from the smaller pretrained model's parameters, The linear map is factorized to a composition of sparse width- and depth-expansion operators with a Kronecker factorization that groups parameters into layers and neurons. We demonstrate the effectiveness of our proposed approach on both lan- guage and vision transformers of different sizes, outperforming several competing methods. While our compute resources prevented us from applying LiGO to even larger transformers, would be interesting to see if this can be applied top of even larger models.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
PW sincerely thanks Zhen Wang for the insightful discussion and for providing reference reposi- tories for language model pre-training. PW also appreciates Hao Tan's assistance for reproducing fine-tuning results on GLUE datasets. YK and LTH were partially supported an MIT-IBM Watson AI grant and an Amazon award. We also acknowledge support from the IBM Research AI Hard- ware Center, and the Center for Computational Innovation Rensselaer Polytechnic Institute for the computational resources on the AiMOS Supercomputer. The research of ZW is in part supported by the US Army Research Office Young Investigator Award (W911NF2010240),
REFERENCES
Andy Brock, Soham De, Samuel Smith, and Karen Simonyan. perffreeformnce large-scale image recognition without normalization. In International Conference Machine Learning, PD- 1059-1071, 2021.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In Proceedings of NeurlPS, 2020.
Han Cai, Tianyao Chen, Weinan Zhang, Yong Yu, and Jun Wang. Efficient architecture search by network transformation. In Proceedings of AAAI, 2018.
Bo Chang, Lili Meng, Eldad Haber, Frederick Tung, and David Begert. Multi-level residual net- works from dynamical systems view. arXiv preprint arXiv: 1710..0008. 2017.
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 arXiv:2 / 2021.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Tianqi Chen, Ian Goodfellow, and Jonathon Shlens. Net2net: Accelerating learning via knowledge transfer. arXiv preprint 1.Xi:1511..0644 2015.
Justin Chiu, Yuntian Deng, and Alexander Rush. Low-rank constraints for fast inference in struc- tured models. Advances Neural Information Processing Systems, 34:2887-2898, 2021.
Xiaoliang Dai, Hongxu Yin, and Niraj Jha. Nest: neural network synthesis tool based on grow and and-urun- prune paradigm. IEEE Transactions on Computers,
2019.
Tri Dao, Albert Gu, Matthew Eichhorn, Atri Rudra, and Christopher Ré. Learning fast algorithms for linear transforms using butterfly factorizations. In International conference on machine learning, pp. 1517-1527, 2019.
Tri Dao, Nimit S Sohoni, Albert Gu, Matthew Eichhorn, Amit Blonder, Megan Leszczynski, Atri Rudra, and Christopher Ré. Kaleidoscope: An efficient, learnable representation for all structured linear maps. arXiv preprint arXiv:2012. 2020.
Tri Dao, Beidi Chen, Nimit S Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. Monarch: Expressive structured matrices for efficient and accurate training. In International Conference on Machine Learning, pp. 4690- 4721, 2022.
Yann N Dauphin and Samuel Schoenholz. Metainit: Initializing learning by learning to initialize. Advances in Neural Information Processing Systems, 32, 2019.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. 248-255, 2009.
Jacob Devlin, Ming- Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACI, 2019.
10

---

Published as a conference paper at ICLR 2023
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.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold. Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An Image is Worth 16x16 Words: Transformers for Image Recognition Scale. In Proceedings of ICLR, 2021.
Utku Evci, Max Vladymyrov, Thomas Unterthiner, Bart van Merriënboer, and Fabian Pe- dregosa. Gradmax: Growing neural networks using gradient information. arXiv preprint arXiv:2201.05125 2022.
Scott Fahlman. The recurrent cascade-correlation architecture. Advances in Neural Information Processing Systems, 1990.
Scott Fahlman and Christian Lebiere. The cascade-correlation learning architecture. In Advances in Neural Information Processing Systems, 1989.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249-2449222 246, 2010.
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, 2019.
Xiaotao Gu, Liyuan Liu, Hongkun Yu, Jing Li, Chen Chen, and Jiawei Han. On the transformer growth for progressive bert training. arXiv preprint arXiv:2010. /2562, 2020.
Steven Gutstein, Olac Fuentes, and Eric Freudenthal. Knowledge transfer deep convolutional neural nets. In Proceedings of International Journal Artificial Intelligence Tools, 2008.
Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. arXiv preprint arXiv:1510.00145 2015.
Le Hou, Richard Yuanzhe Pang, Tianyi Zhou, Yuexin Wu, Xinying Song, Xiaodan Song, and Denny Zhou, Token dropping for efficient bert pretraining. urXiv preprint arXiv:2203. 13240, 2022.
Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, An- drea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficieet transfer learning for nlp. In International Conference on Machine Learning, 2019.
Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs. Improving transformer opti- mization through better initialization. In International Conference on Machine Learning, pp- 4475 4483, 2020.
Yanping Huang, Youlong Cheng. Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee. Jiquan Ngiam, Quoc V Le, Yonghui Wu, al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019.
Menglin Jia, Luming Tang, Bor-Chun Chen, Claire Cardie, Serge Belongie, Bharath Hariharan, and Ser-Nam Lim. Visual prompt tuning, arXiv preprint arXiv:2203. 2022.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08836 2020.
Yannic Kilcher, Gary Bécigneul, and Thomas Hofmann. Escaping flat areas via function-preservingg structural network modifications. 2018.
Jonathan Krause, Michael Stark, Jia Deng, and Li Fei-Fei. 3d object representations for fine- grained categorization. In Proceedings of the IEEE international conference on computer vision work- shops, pp. 554-561, 2013.
11

---

Published as a conference paper at ICLR 2023
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficien prompt tuning, arXiv preprint arXiv:2104. 08691, 2021.
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 on Computer Vision and Pattern Recognition, pp. 12486-12496, 2022.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint AXiv:1977.1192. 2019.
Dmytro Mishkin and Jiri Matas. All you need is good init. arXiv preprint arXiv:1511.0642 2015.
Maria-Elena Nilsback and Andrew Zisserman. Automated flower classification over a large number of classes. In 2008 Sixth Indian Conference Computer Vision, Graphics & Image Processing, pp. 722-729. IEFE, 2008
Jonas Pfeiffer, Aishwarya Kamath, Andreas Rücklé Kyunghyun Cho, and Iryna Gurevych. Adapter- fusion: Non-destructive task composition for transfer learning. arXiv preprint (.Xiv:2005500247. 2020.
Yujia Qin, Yankai Lin, Jing Yi, Jiajie Zhang, Xu Han, Zhengyan Zhang, Yusheng Su, Zhiyuan Liu, Peng Li, Maosong Sun, et al. Knowledge inheritance for pre-trained language models. arXiv preprint arXiv:2105.13880, 2021.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training. 2018.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J Liu, al. Exploring the limits of transfer learning with unified text-to-text transformer. J. Mach. Learn. Res., 21 22(1111))))) (140):166) 2020.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv:1606.0523 05250, 2016.
Pranav Rajpurkar, Robin Jia, and Percy Liang. Know what you don't 't know: Unanswerable questions for squad. arXiv preprint 1806v:1806. 038222, 2018
Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Jason Liu, Demi Guo, Myle Ott, C. Lawrence Zitnick, Jerry Ma, and Rob Fergus. Biological structure and function emerge from scaling unsupervised learning 250 million protein sequences. Proceedings of the National Academy of Sciences, 118(15), 2021. ISSN 0027-8424. doi: 10.10733/pas:20112239 18.
Jonathan S Rosenfeld, Amir Rosenfeld, Yonatan Belinkov, and Nir Shavit. A constructive prediction of the generalization error across scales. arXiv preprint arXiv: /6673, 2019.
Kathrin Schacke. On the kronecker product. Master's 's thesis, University of Waterloo, 2004.
Sheng Shen, Pete Walsh, Kurt Keutzer, Jesse Dodge, Matthew Peters, and Iz Beltagy. Staged training for transformer language models, arXiv preprint arXiv:2203.06211 2022.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model par- allelism. arXiv preprint arXiv:1909.0800 2019.
Vikas Sindhwani, Tara Sainath, and Sanjiv Kumar. Structured transforms for footrin deep learning. Advances in Neural Information Processing Systems, 28, 2015.
12

---

Published as a conference paper at ICLR 2023
Hao Tan and Mohit Bansal. Vokenization: Improving language understanding with contextualized, visual-grounded supervision. arXiv preprint arXiv:2010.06775, 2020.
Shanshan Tang, Bo Li, and Haijun Yu. Chebnet: Efficient and stable constructions of deep neural networks with rectified power units using chebyshev approximations. arXiv preprint arXiv:1911.05467, 2019.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International Conference Machine Learning, pp. 10347-10357, 2021a.
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 Computer Vision, 2021b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz. Kaiser, and Illia Polosukhin. Attention is All You Need. In Proceedings of NeurlPS, 2017.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Glue: A multi-tast benchmark and analysis platform for natural language understanding. Proceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 353-355, 2018.
Xiaosong Wang. Yifan Peng, Le Lu, Zhiyong Lu, Mohammadhadi Bagheri, and Ronald M Sum- mers. Chestx-ray8: Hospital-scale chest X-ray database and benchmarks on weakly-supervised classification and localization of common thorax diseases. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2097-2106, 2017.
Tao Wei, Changhu Wang, Yong Rui, and Chang Wen Chen. Network morphism. In Maria Florina Balcan and Kilian Q. Weinberger (eds.), Proceedings of The 33rd International Conference on Machine Learning. pp. 564-572. 2016.
Lemeng Wu, Dilin Wang, and Qiang Liu. Splitting steepest descent for growing neural architectures. Advances in neural information processing systems, 32, 2019.
Lemeng Wu, Dilin Wang, Peter Stone, and Qiang Liu. Firefly neural architecture descent: general approach for growing neural networks. Advances in neural information processing systems, 2021.
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 1XXi:20111.3663: 2020.
Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint arXiv: 1904.0996, 2019.
Hongyi Zhang, Yann N Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. arXiv preprint 1xXiv:1901.0932) 2019.
Minjia Zhang and Yuxiong He. Accelerating training of transformer-based language models with progressive layer dropping. Advances in Neural Information Processing Systems, 33:14011- 14023, 2020.
Xu Zhang, Felix X Yu, Ruiqi Guo. Sanjiv Kumar, Shengjin Wang, and Shi-Fu Chang. Fast orthogo- nal projection based on kronecker product. In Proceedings of the IEEE International Conference on Computer Vision, pp. 2929- 2937,,937, 2015.
Chen Zhu, Renkun Ni, Zheng Xu, Kezhi Kong, W Ronny Huang, and Tom Goldstein. Gradinit: Learning to initialize neural networks for stable and efficient training. Advances in Neural Infor- mation Processing Systems, 14:10-11-1 2021.
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.
13

---

Published as a conference paper at ICLR 2023
A UNIVERSALITY OF LIGO OPERATOR
Proposition StackBERT (Eq.
7v. Interpolation (Eq. 1. and Ne12Net (Eq. 2 are all the special
cases of the LiGO operator (Eq.)
Proof. We prove Proposition by constructing parameters in depth and Rwidth-
Stacking. Stacking-based methods (Gong al. 2019 Yang 2020) duplicate the entire lower blocks on top of the small model to the form new layers (Eq. Formally, we show this operation can be done by the following operator:
M
(9)
I
Rwiden
Ldepph
Interpolation.
Interpolation based methods (Chang et al. 2017 Dong et al. 2020 interleave each
layer for twice. We construct the following matrix achieve layer interpolation (Eq.
I
M
I
(10)
I
Rwidth
Lisepth
We remark that any rearrangement layers to construct new layers (mathematically a permutation of existing layers with replacement) can be constructed in similar way.
Since we show in Eq. 6 the Kronecker factorization on R2 amounts to decomposing the general growth operator into in-dimension and out-dimension expansion. We can construct Net2Net (Chen et al. 2015) based growth by simply letting:
Net2Net.
A1 2 B
Edepth
I E RLIDexLaDa Rwidth
(11)
AL. & BL,
1
A
B
(12)
S1-1
S
where S; E (0, (0,1) (DD:-DI) selection matrix to enlarge the out dimension, and S Si-1 diag(1 S copies the selection from S with normalization to guarantee functionality preserving in expansion.
B IMPLEMENTATION DETAILS
B. GROWING TRANSFORMERS WITH LIGO
The transformer architecture consists of an embedding layer, multi-block attention layer, and an output layer. The core ingredient attention block consists of Multi-Har Head Attention (MHA) module followed by a FeedForward Network (FFN), with skip connection across the both blocks. Applying LiGO requires the following considerations:
14

---

Published as a conference paper at ICLR 2023
Embedding layer. For both language and vision transformers, the embedding layer can be garded as a linear layer, whose inputs are one-hot embeddings in language models. We draw a learnable matrix ((emb) to extend its output dimension.
Multi-head attention blocks. An attention layer in transformer consists of multi-head attention weights (WQ, w w and linear projection (W Let A and B with E {Q, K, V, V be the in- and out-dimension expansion matrices (Eq. 6) for query, key, value, and projection in the I-th layer, respectively. Applying B to W (k E Q, K, V) constructs new heads by weighted summation of rows of all existing heads. To make sure the new input and output channels are aligned across modules, we tie our LiGO operator with the following scheme: (1) A = (B(emb))²
for Ak (Q. V. (2) A )T (3) Bi B Blemb) for VI [L1]. Both the bias and layer normalization inherit the associated linear transformations' out-dimension expansion matrices grow the width. For depth expansion, each module independently combines the same module from other layers (Eq. with learnable coefficients w.
Reed-forward networks. Each attention block is followed by two-layer FFN. Let A and B with k E {fcl, Jc2] be the in- and out-dimension expansion matrices (Eq. 66 for the first and second FFN layer in the I-th layer, respectively. We tie the parameters for feed-forward networks: (fel)
B (emb) A (f(2) = -B( and B Blfe2) = (emb))
Output layer. For output head, we have A (out) B (emb)T since the output dimension of tention layers are always aligned with (emb)) by our construction. The output layer does not need out-dimension expansion. Algorithm summarizes LiGO for growing transformers.
B.2 MODEL CONHIGURATIONS
We summarize the settings of different transformer models used for our experiments in Table For BERT and RoBERTa, we re-use the code base provided by Tan & Bansal (2020). For GTP2, we follow the model configuration of OpenAl and use the pre-training code provided by [Shen et al. 2022.. For DeiT, we use their official codebase (Touvron et al. 2021a).
Table 4: Configuration of diffcrent transformers.
BERT-Small
BERT-Base
RoBERTa-Smail Small
RoBERTs Base
GPT2-Base
GPT2-Mediun
DeiT-S
DeTI-B
layers
12
24
layers
12
hidden
512
512
768
768
1024
hidden
384
768
# heads
8
12
8
12
12
16
heads
6
12
vocab
30522
30522
50265
50265
50257
50257
input
224
224
sog. length
128
128
128
128
1024
1024
patch size
16
16
B.3 ORTHOGONAL EFFICIENT TRAINING STRATEGIES
For layer dropping, we follow the same progressive dropping rate schedule with Zhang & He 2022)) and set the maximum dropping rate 0.1 recover the performance. For token dropping, we randomly set 15% tokens aside in the middle layers. In the first 50k steps of staged training, only sub-network activated and trained, and afterwards, we perform full-model training for 350k steps.
C ADDITIONAL EXPERIMENTS
C. REUSING SMALLER MODELS TRAINED FOR ONLY FEW STEPS
LiGO focuses on utilizing the knowledge of smaller models that have already been pretrained and available. In this section, we investigate how LiGO can leverage smaller existing models that are only trained for few steps to accelerate training of a larger model. We perform an experiment on BERT-Base by reusing BERT-Smal trained for only 50k steps instead full training for 220k steps as used in our experiments. Figure / shows that LiGO can still save 35.2% savings FLOPs and 30.2% savings in wall time over the BERT-Base training from scratch.
15

---

Published as a conference paper at ICLR 2023
Algorithm 1 A forward pass of LiGO with transformer.
1: Input: A small transformer with hidden D2 and number of layer L1. Denote the embedding
layer Wrlemb) RRDr DixE, attention layers w I w WY, W DDixD; FFN layers
as w(fe E IRAD,XD, w(fe2) E RDDxXD, LayerNorm layers as W W!ln)) E DRDxx
W (lln)) w((m2) E ERP, VI [L]. the output head W E RCxD1
2: Output: A large transformer with hidden D2 and number of layer L2. Denote the weight
matrices as with the corresponding superseripts as the small model.
3: (embb) f- B (emb)) (emb)
4: for I = 1,... L1 do
Width Expansion
5:
Ω2 (QQ) W
B jemb]
D + B(K) W (emb)T
6:
Ω - B WY B
(emb)
7:
Ω + Blemb) WY (
8:
(m1) f Blemb) w (ln1)
9:
Ω
10:
Ω
Mc1)
B(fel) w B (embiT
11:
Ω
(c2)
t B (eemh)
W (fflyT
12:
Ω
(l2))
Blemb)
W (lm2)
13: end for
14:
for
1, ... L2 do
Depth Expansion
15:
22 - 2 1
52
16:
2 + Σ j= w K Snk
17:
N + Σ w
.2
D + Σ j=1 "
Ω
(m11)
E wi
ii
(l11) (brl)
Ω
Ω (fal)) + E Li
(fal)
( (fel
20:
n (122)
(fe2)
fc2)
21:
22:
(m2))
Eラ Wi.j
(1n2)
(n2)
23: end for
24: Ω <- (out)) (embb)T
25: Train transformer with parameters Ω.
C.2 RESULTS ON CATT
In addition DeiT (Touvron et al. 2021a), we perform additional experiments with CaiT Touvron et al. 202Tb) on ImageNet and find that while reusing CaiT--XS, LiGO offers about 52.6% savings in FLOPs and 46.1% % savings in wall time over the CaiT-S training from scratch (see Figure g
C.3 TASK-SPECIFIC FINETUNING WITH LIGO INITIALIZATION WITHOUT FURTHER PRETRAINING
We perform additional experiments by directly finetuning BERT-Base initialized by LiGO (from BERT-Small) without any further pretraining. We observe in Table the LiGO-initialized model benefit downstream tasks compared to BERT-Small trained from scratch (1st row 2nd row).
C.4 GLUE PERFORMANCE USING ADAPTERFUSION
LiGO is mainly proposed for improving efficiency of the pre-training stage and hence is compatible with various finetuning schemes like full model finetuning, adapters (Houlsby et al. 2019 Pfeiffer
et al.
2020) prompt tuning (Lester al. 2021: Jia et al. 2022 for adaptation downstream tasks. We test BERT-Bas models trained using different baselines by using adapterfusion (Peiffer
al.
2020 instead of full finetuning GL.E benchmark. Table ō shows that LiGO also achieves
16

---

Published as a conference paper at ICLR 2023
Scratch LIGO
LIGO from 50k
Scratch
LiGO from 50k
LIGO
2.6
9.
2.6
1.40
Log Perplexity
Perplexity
2.4
2.4
44.7%
35.2%
0.0%
2.2
2.2
40.7%
30.2%
0.0%
30
40
5D
2.0
2.0
1.8
1.8
0
2
0
20
40
60
80
FLOPs (1e18)
Wall Time (hrs)
(a) BERT-Small Souul-BBERI --BBB-TTBEE Buse
(b) BERT-Smal SSaal->BERT-Baa
Figure 7: Results BERT-Base by reusing BERT-Small trained for 50k steps. Instead of training BERT-Base from fully trained BERT-Small, we run LiGO on BERT-Small trained 50k steps. LiGO offers about 35.2% savings in FLOPs and 30.2% savings in wall time over the BERT-Bass training from scratch.
85
85
Scratch
LiGO
Scratch
LiGO
80
80
75
75
Valid Accuracy
Valid Accuracy
70
70
52.6%
46.1%
65
65
60
60
55
55
50
50
0.0
2.5
7.5 10.0 12.5 15.0 17.5 FLOPs (1e18)
5
30
35
Wall Time (hrs)
(a) CaT-XS--CaiTS
(b) Cai--XS- -Cii-S
Figure 8: Results on CaiT. (a) Accuracy vs. flops and (b) accuracy vs. wall time for training CaiT-S. LiGO saves flops by 52.6% and wall time by 46.1% over training from scratch ImageNet.
Table GLUE performance different LiGO models.
results are based on BERT-Base models
with BERT-Small as the base model for LiGO optimization.
SST-2
MNLI
MRPC
CoLA
QNLI
QQP
STS-B (Acc.)
Average (Acc.)
Method
(Acc.)
(Acc.)
(Acc.)
(Acc.)
(Acc.)
(Acc.)
BERT-Small (Scratch)
87.21
77.56
82.11
59.93
85.06
85.82
84,99
80.38
BERT-Base (LiGO Init)
88.15
77.62
82.53
60.70
85.79
86.65
85,83
81.04
BERT-Basc (LiGO Init + Pretrain)
88.42
79.29
84.31
62.09
88.07
88.81
87.00
82,57
BERT-Basc (Scratch)
88.19
78.43
85.78
62.09
87.06
87.18
86,99
82.25
17

---

Published as a conference paper at ICLR 2023
Table 6: Downstream performance using AdapterFusion (Pleiffer al. the results are based on BERT-Base models trained using different baselines.
2020
GLUE Benchmark. All of
Method
Savings (FLOPs)
Savings (Walltime)
SST-2
MNLI
MRPC
CoLA
QNLI
QQP
STS-B
Average
(Acc.)
(Acc.)
(Acc.)
(Acc.)
(Acc.)
(Acc.)
(Acc.)
(Acc.)
Scratch
88.41
78.60
86.02
62.39
87.62
88.02
86.52
82.51
StackBERT
34.1%
33.3%
88.78
79.80
85.43
59.56
87.71
89.19
86.27
82.39
MSLT
34.9%
30.0%
88.41
78.35
83,15
63.97
86.19
88.20
86.42
82.10
KI
-5.7%
-13.9%
88.94
78.84
84.00
64.61
86.75
88. 19
87.93
82.75
ben2BERT
29.0%
25.1%
88.47
80.53
85,50
62.33
88.57
86.72
87.10
82,75
LIGO
44.7%
40.3%
88.45
80.01
84.67
63.05
88.06
88.92
87.00
82.88
on-par performance with model trained from scratch under adapter-based tuning with 44.7% savings FLOPs abd 40.5% savings in wall time. This shows that LiGO does not harm the model general- ization capability when adapters used as parameter--ffiiient finetuning strategy for transferring a trained model to downstream datasets.
C.5 INITIAL RESULTS ON BILLION+ PARAMETER MODELS
Our extensive experiments on BERT (Devlin et al. 2019, RoBERTa (Liu et al. 2019), GPT2 (Rad-
ford et al.) 2019), DeiT (Tovvron et al. 202 laj and CaiT (Touvron et al. (021b) show that LiGO can consistently improve transformer training efficiency over the traditional way of training from scratch across domains and model sizes. One interesting future direction of our work is scaling LiGO to very large models with parameters more than 100B, such as GPT3 (Brown et al.) 2020). While we currently do not possess the compute resources for this extreme large-scale study, per- form a preliminary experiment on GPT2-1.5D (Radford et al. 2019) by using GPT2-Medium as the initialization. We train for 15k steps C4 dataset (Raffel al. 2020) and find that proposed LiGO saves about 39% computation cost (FLOPs) training GPT2-1.5B from scratch to reach the same log perplexity (which is 3.3). We believe that is imperative to study the extent which the benefits of LiGO remain at the scale on which the modern large language models are trained. We hope cover this in future work.
18


Reviews:
[{"review_id": "xBR2R7kffx9", "reviewer": "Reviewer_52nq", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper investigates the problem of accelerating the large-scale transformer model's training. The idea is to initialize the large-scale model with the transformed parameters of the small-scale model. Compared to existing methods, the proposed method can expand the network in both width and depth and the Kronecker factorization is used to reduce the computational cost. The experimental results in both BERT and Vision Transformer demonstrate that LEGO can make the model converge faster.", "summary_of_the_review": "Please refer to the 'Clarity, Quality, Novelty And Reproducibility' part above. ", "strength_and_weaknesses": "Pros:\n* The first work to combine both width- and depth-growth operators together and thus expand the model in both width and depth.\n* LEPO shows good capacity in both the linear layer and the attention layer.\n* The empirical results show that LEPO can greatly improve training efficiency.\n\nQuestions:\n* The reported FLOPs and Wall Time of LEPO are counted during training. Would it be more reasonable to also consider the calculation of the LEGO operator before the transformer training? \n* As claimed in the experiment part, the authors claim that they investigate the performance of LEGO when transferring the model to other downstream tasks. Do the authors compare LEGO with some adapter-based methods or visual prompt tuning methods?\n* Some typos can be fixed. For example, the reported results of LEGO show it can reduce 40.5% wall time when transferring BERT-Small to BERT-Base, however in Figure 2(b) the number is 40.7%.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "cckfE1lxZlx", "reviewer": "Reviewer_xesb", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposes a method for initializing a larger model by learning a linearly map using the parameters of a small model to obtain the parameters of a larger one, thus reducing the training cost of the larger model by using the existing small model.\nKronecker factorization is used to reduce the computational parameters to an acceptable level, making it possible to learn a linearly map as mentioned above.\nThe proposed method is applicable not only to multilayer neural network structures but also to transformers, and experiments have been conducted on various transformer-based models with good results.", "summary_of_the_review": "The idea and method in this paper is novel and subtle, makes it easy to learn a linearly map to grow model. However, the specific training process of LEGO is not highlighted in the text, and the experimental part is not sufficient, such as the detailed evaluation of the performance of downstream tasks. ", "strength_and_weaknesses": "Strength\n1.This paper uses a subtle factorization which substantially reduces the number of computational parameters, making it feasible to learn such a linearly map. This idea is very good, simple but useful, and also has some explainability.\n2.Clearly illustrates how to extend the parameters in depth and width from a theoretical point of view, and with some justification.\n3.The paper is well written.\nWeakness\n1.The training process of LEGO method is not described in detail in the paper, although illustrated in the appendix, I suppose it should be described in detail in the main paper in Sec.3 or another place. And the optimization process is not mentioned either. As I understand it, the parameters of LEGO are trained for a specific number of steps together with the parameters of the larger model, and then the training of LEGO is stopped to train only the larger model. Is it right? \n2. Compared to training from sketch, I wonder how about the performance in downstream task of model which use LEGO method initialized. The performance of downstream task is what we all concern, and it’s important. I hope you can add some relevant experiments.\n3.Only one Vision Transformer model has experimented, more vision models should be considered.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "QW-J1CwaF5", "reviewer": "Reviewer_nRWA", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposes a method named LEGO to efficiently train a transformer model by initializing it with a pretrained smaller transformer. The method expands the pretrained transformer model on width and depth by learning linear maps on the parameters. Experiments on various language and vision tasks show that LEGO can obviously reduce the training time compared to training from scratch.", "summary_of_the_review": "Overall, I think this paper is a borderline paper. My major concerns of this paper are the limited improvements compared to the existing methods and the generalizability of the method.", "strength_and_weaknesses": "Strengths:\n1. The idea of this paper is easy to follow. By learning the expansion operators with task loss, the pretrained smaller transformer can be naturally mapped to a larger network, which could obtain a significantly higher accuracy at the beginning of training compared to training from scratch.\n\n2. Extensive experiments on various tasks validate the efficacy of the proposed method. For example, in Table 1, LEGO outperforms Scratch in most cases with 40.5% training time reduced.\n\nWeaknesses:\n1. The method requires a trained smaller transformer to initialize the target network, but I'm not sure whether the training cost is still low if we take the training time of the smaller network into account.\n2. It seems that LEGO only supports expanding the width of a linear transformation or growing it into a sequence of linear transformations, while two transformations connected by a non-linear activation function can not be initialized by one transformation. This may be a limitation of this work.\n3. The improvements compared to existing methods such as StackBERT are marginal. In Figure 2, StackBert has lower log perplexity curves when FLOPs > 4e18 or wall time > 40 hrs.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "JYoOvLg8VA", "reviewer": "Reviewer_bupJ", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper presents a learning strategy to grow smaller pre-trained networks into larger networks with the hope that model growth can save compute compared to a network that is trained from scratch. While mostly in the literature, growth strategies are based on heuristics, e.g., copying the layers of the network for expanding depth, or block diagonal expansion for increasing width, learning-based model growth is not extensively studied. In this work, the authors initialize the parameters of the larger network as a linear transformation of the parameters of the smaller network. This linear transformation is learned through data using a small number of gradient steps (typically ~100) – a cost which is marginal compared to the training of the larger network. The authors propose to decompose the linear transformation into depth- and width-operators, and use structured sparse matrices for these operators. For enforcing structures and sparsity, layer-wise transformations are proposed using kronecker factorization. Through experiments on BERT, RoBERTA, GPT2 and ViT, the authors show that their proposed growth strategy saves up to ~40% of compute without sacrificing the upstream and downstream performance of Transformer models across the domains. ", "summary_of_the_review": "Please refer to the strengths and weaknesses section above for details. ", "strength_and_weaknesses": "**Strengths**\n\n1. The paper is written well and easy to follow. \n2. Given the prevalence of large (foundation) models, and the fact that compute is perhaps the most important constraint while training these models, the paper tackles a very good problem in a timely manner. \n3. The experiments aptly make the points of the main paper. \n\n**Weaknesses/ Questions/ Suggestions**\n\n1. **Performance of the base model**: I would encourage the authors to report the performance of the base (small) models, from where the larger models are trained, in the plots. It would help us clarify a) how quickly the larger model recovers the performance of the base model, b) how much it improves upon the base model. \n\n2. I would encourage the authors to modify the introduction/ story and emphasize the fact that their proposed method allows one to train a family of Transformer models, where new and larger models can utilize the information in the previous and smaller models to save compute, which is not possible when training from scratch. In the absence of this messaging, one could fairly ask, when comparing model growth with training from scratch, the compute budget of training the smaller models should also be included in the plots. \n\n3. **Scalability to larger models**: The authors have reported performance on relatively smaller architectures (up to 774 million parameters) compared to modern Transformer models. I wonder if the proposed method would also result in equivalent gains on these larger models. I see that in the conclusion the authors already concede this point but, nevertheless, I would like to know the authors' take on that. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%