LEMON: Lossless model expansion
Authors: Yite Wang
Jiahao Su
Hanlin Lu
Cong Xie
Tianyi Liu
Jianbo Yuan
Haibin Lin
Ruoyu Sun
Hongxia Yang
Venue: ICLR 2024 poster (2024)
Keywords: model growth
efficient deep learning
continual learning
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We propose LEMON, a method that initializes large model with pretrained small model to save computational resources.

Abstract: Scaling of deep neural networks, especially Transformers, is pivotal for their surging performance and has further led to the emergence of sophisticated reasoning capabilities in foundation models.
Such scaling generally requires training large models from scratch with random initialization, failing to leverage the knowledge acquired by their smaller counterparts, which are already resource-intensive to obtain.
To tackle this inefficiency, we present $\textbf{L}$ossl$\textbf{E}$ss $\textbf{MO}$del Expansio$\textbf{N}$ (LEMON), a recipe 
to initialize scaled models using the weights of their smaller but pre-trained counterparts. This is followed by model training with an optimized learning rate scheduler tailored explicitly for the scaled models, substantially reducing the training time compared to training from scratch.
Notably, LEMON is versatile, ensuring compatibility with various network structures, including models like Vision Transformers and BERT.
Our empirical results demonstrate that LEMON reduces computational costs by 56.7\% for Vision Transformers and 33.2\% for BERT when compared to training from scratch.

Published as a conference paper at ICLR 2024
LEMON: LOSSLESS MODEL EXPANSION
Yite Wang- Jiahao Su Hanlin Lu², Cong Xie Tianyi Liu?, Jianbo Yuan,, Haibin Lin", Ruoyu Sun Hongxia Yang" USA *ByteDance Inc.
University of Illinois Srbana-Champpar
'hhe Chinese University of Hong Kong, Shenzhen, China *Shenzhen Research Institute of Big Data yi witew2@ill1e wele001lniiiiiii edu {jiahac. su, hanlin. cong.xie, tianyi. liu,
d gianbo. yuan, haibin. lin, hx. @ com sunruoyuecuhk.cm enu.on
ABSTRACT
Scaling of deep neural networks, especially Transformers, is pivotal for their surg- ing performance and has further led to the emergence of sophisticated reasoning capabilities in foundation models. Such scaling generally requires training large models from scratch with random initialization, failing to leverage the knowl- edge acquired by their smaller counterparts, which are already resource-intensive obtain. To tackle this inefficiency, we present LosslEss MOdel ExpansioN (LEMON), a recipe to initialize scaled models using the weights of their smaller but pre-trained counterparts. This followed by model training with an optimized learning rate scheduler tailored explicitly for the scaled models, substantially re- ducing the training time compared to training from scratch. Notably, LEMON is versatile, ensuring compatibility with various network structures, including models like Vision Transformers and BERT. Our empirical results demonstrate that LEMON reduces computational costs by 56.7% for Vision Transformers and 33.2% for BERT when compared to training from scratch.
INTRODUCTION
Deep neural networks (DNNs) have become increasingly popular, showcasing their adaptability across natural lan- guage processing (Liu & Lapata, 2019 Achiam et al. 2023), computer vision (Chen et al. 2023a|b); (b), and code generation Yu et al.) (023)). Recent advances in archi- tectural design, especially Transformers, have further en- hanced the scalability of DNNs. However, it is com- mon practice to train large-scaled models from scratch, discarding the learned knowledge in their smaller coun- terparts. Such an approach can be highly inefficient, es- pecially given the intensive computational resources re- quired to train large language models such as Generative Pre-trained Transformer (GPT) [Brown al. 2020), and the resultant huge carbon footprints. For instance, train- ing GPT-3 incurs costs around $4.6M LLi 2020). Given these challenges, researchers are keenly exploring ways leverage the prior knowledge of smaller models for more efficient scaling.
inkidlizd
Salatantial
faage
1mm ormach
Furgge
Pritzited
Badaced tim
Finsh
flarge
fiarge
Figure 1: Comparison between training from scratch and model expansion. In model expansion, a smaller pre-trained model is expanded without any performance drop, requir- ing significantly less training time than training from scratch.
larger model
Knowledge inheritance and model expansion are two primary methodologies to achieve this goal. Knowledge inheritance (Qin et al., 202T), the reverse knowledge distillation (Hinton al., 2015). allows the large model to learn the predictions of a smaller pre-trained model. However, this method often necessitates additional computational resources and modifications to the training pipeline due to the involvement of 'teacher network.' In contrast, model expansion directly utilizes the weights
from the pre-traines trained small source network, either without training (Chen et al., 2015] 2021a: Yang
et al. 2020 Shen et al. 2022) with negligible training (Wang et al.) 2023a). Hence, work
Wook done during internship at ByteDance.
'Corresonning author.

---

Published as a conference paper at ICLR 2024
mainly focuses on model expansion due to its minimal impact the training pipeline and negligible computational overhead.
A compelling requirement for model expansion is to ensure it is "lossless," meaning information from the source model is lost. Specifically, the goal is for the larger target model to inherit the exact functional mapping as the smaller source model, thus preserving the performance. Net2Net (Chen et al. 2013) represents a foundational study of lossless model expansion for convolutional networks (CNNs) and multi-layer perceptrons (MLPs) where it duplicates neurons and averages their fan- out weights. However, a challenge arises with the 'weight symmetry' issue. This problem occurs when duplicated neurons in expanded layers introduce redundancy, which persists during subsequent training. In this sense, the expanded model will never gain more capacity than the source model. To counteract this problem, previous researchers introduced additional noise into the expansion process, leading to a shift away from a genuine lossless expansion.
Transformers, despite their rising popularity in modern deep learning, introduce additional complex- ities in achieving lossless expansion that goes beyond traditional issues like weight symmetry. One key obstacle arises from the intricacy of the LayerNorm, which was evident when bert2BRT (Chen) et al., 202 tried extending the Net2Net approach to Transformers, leading to lossy outcomes. Staged training (Shen et al. 2022) demonstrated the feasibility of lossless model expansion, but with specific constraint: doubling the width during expansion and only for variant of Transformers known as Pre-Layer Normalization (Pre-I.N) Transformers. However, real-world applications often require width increases in the expanded model that are indivisible by the smaller source model's width, highlighting a limitation in existing methodologies. A typical scenario involves expanding the hidden dimension from 512 to 768.
In exploring the possibilities of lossless model expansion, our research focuses on the ability to break weight symmetry, handle indivisible width and depth increments, and remain compatible with almost all Transformer varieties. We have discovered affirmative answers, revealing that mul- tiple solutions exist, enabling the selection of an optimal candidate break the weight symmetry or find an initialization point with specific properties. Specifically, we break the weight symme- try of replicated neurons by setting their fan-out weights to be unequal, and we introduce average expansion to deal with LayerNorm for indivisible width increment.
In addition to lossless model expansion techniques, our study also delves into training recipes for the expanded models. It is often overlooked whether applying the original training recipe remains optimal or whether the expanded models necessitate tailored approaches. Our empirical studies reveal two key insights: expanded models can benefit from utilizing a default maximum learning rate and, intriguingly, a learning rate scheduler that decays more rapidly.
Our contributions are summarized follows:
1. We propose LEMON, a suite of algorithms designed for lossless model expansion across a variety of architcctures, ensuring compatibility with indivisible width and depth increments.
2. Drawing inspiration from our empirical results, we propose optimized learning rate sched- uler for the expanded models. This scheduler maintains the maximum learning rate used by training from scratch, but features accelerated decay rates.
3. LEMON reduces the computational costs by up to 56.7% for Vision Transformers and 33.2% for BERT compared to training from scratch, thereby setting a new benchmark in performance.
RELATED WORKS
From small models to larger models. There are two main approaches to transferring the knowledge of the smaller models to larger models: knowledge inheritance and model expansion. Knowledge inheritance (Qin et al.] (02T) enables student network to learn the logits provided by teacher network. Net2Net (Chen et al. 2015)) was the first work to explore the idea of model expansion. It involves randomly duplicating neurons while preserving the output values through proper nor- malization and increasing depth by adding identity layers. However, Net2Net resorts to introducing weight perturbations to overcome weight symmetry, resulting in performance deterioration. Follow- up work bert2BERT (Chen et al. 2021a) extends Net2Net to Transformer while others study depth growth (Gong al. 2019 Yang et al. 2020
Chang et al. 2017 Dong et al. 2020). Staged training
2

---

Published as a conference paper at ICLR 2024
Table 1: Overview of model expansion or knowledge inheritance methods. In the first three columns, we use symbols v. X, and N/A to denote whether the method is (1) lossless, (2) non-lossless, or (3) not applicable in the given scenarios. Here, 'Depth' represents the scenario where the large model has more layers than the smaller model, and "Width (divisiblefnddvviiibee)) denotes whether the large model's hidden dimension is multiple of the smaller model's. the subsequent columns, 'Non- unique Expansion' denotes whether the expansion is unique (e.g.. produce target models to break weight symmetry). 'Data-free' specifies whether the algorithm requires training data. LEMON is the most versatile method compared to previous methods.
Method
Depth
Width (divisibbee
Width (indivisible)
Now- Nou-unique Expansion
Datu-dree
SukkoLKE m Stincctal 2019
NIA
NIA
Yes
MSLT Yorg mi 2rzy
NIA
N/A
No
Yes
ben2BB.R
Ceere 2021a
V
X
No
Yes
Staged Training (S5e ot 2022
NIA
No
Yos
L,Gi [Wing die
3
Yos
No
LEMON
Yes
(Shen et al., 2022) made significant progress by proposing a lossless model expansion method for Pre-LN Translormer, but with the constraint of width doubling. LiGO (Wang et al. (023a) sug- gests employing multiple training steps to find an appropriate linear combination of weights from the source networks. Despite these advancements, all existing methods still face the challenge of the performance drop or strict restrictions on the model width. TTable I compares the related methods.
Network initialization. Numerous studies aim to seek optimal initialization methods for neural net- works, primarily focusing on regulating the norm of network parameters (Glorot & Bengio, 2010 He et al. 2015). Theoretical works try to study these methods through dynamical isometry (Saxe ct al. 2013) or mean field theory (Poole et al. 2016). Orthogonal initialization, which supports layer-wise dynamical isometry in fully-connected layers, has been extended to CNNs via Delta or- thogonal initialization (Xiao et al., 2018). However, there has been limited research on initialization methods specifically for Transformers. Most of these works focus theoretical approaches to train Transformers without skip connections or normalization layers (Noci et al.) 2022
He et al. 2023) Mimetic initialization (Trockman & Kolter 2023) seeks initialize attention based on the principles
of pre-trained Transformers.
Continual pre-training Recent research explores adapting pre-trained networks for new or im- proved datasets. While some target datasets from different domains (Scialom et al. [2022] Ke et al.,
2022
Gupta et al., [2023] Qin et al.. (022). others focus on datasets that evolve over time (Han al.
2020 J Jang et al.) [202]] Loureiro et al. 2202).. Model expansion is similar to continual pre- gre-training,
with the distinction being a change in the model size rather than the data distribution.
3 PRELIMINARIES
Model expansion aims to initialize a large model with the weights from its smaller pre-trained counterparts. Concretely, suppose we have pre-trained weights 0g in source network fs(:: grgaine). ), our goal is to design a mapping gespanded = M(ggained), where the expanded weights initialize the target network as (:; + gespanded). Since these expanded weights contain knowledge acquired by the small pre-trained model, it should accelerate the training of Sr compared to random initialization. Moreover, we call a model expansion algorithm lossless if (x; 0 gexpanded,
fs( Osrained) Vx.
An example for model expansion is use pre- trained ResNet-50 (He et al., 2016) or BERT-Small (fs) to facilitate the training of WideResNet-11 10 or BERT-Base (ST), respectively. Instead of training the larger models from scratch, the idea is to initial- ize them with the weights of their smaller pre-trained counterparts, i.c., ResNet-50 or BERT-Small.
LayerSoom
MLF
LaverNoon
MLP
LayorNoem
MLP
LavetNocn
+
MIlA
Liperiloen
MHA
LayerNorm
Transformer architecture, introduced by Vaswani et al. (0177), consists of multiple Transformer blocks g(-), where each block is stack of two modules, helti-head attention (MHA) and a two-layer MLP. Depending on the location of LayerNorm, Trans-
(b)
Figure 2: Varieties of attention blocks. (a) Post-LN block. (b) Pre-LN block. (c) Res- Post-Norm block.
3

---

Published as a conference paper at ICLR 2024
former blocks can be categorized (1) Post-LN block used by the original BERT (Devlin et al. 2019) where LN is applied after the residual block, i.c., (() = [MModuleex)) (2) Pre- LN used by GPT (Brown et al.) 2020). Pre-LN BERT, Vision Transformers (Dosovitskiy et al. 2021
and SWin Transformer (Liu et al. 2021b) where LN is applied inside the residual connection and before all other transformations, 1.e., ((z) + Modu 1 (LL(()))) and (3) Res-Post-Norm used by SWin Transformer V2 (Liu et al. (2022) where LN is applied inside the residual connection and after all other transformations, i.c., ((x)
IN(Module(z)) See Figure 2 for illustration.
Multi-head attention (MHA) uses multiple self-attention heads to attend to information from dif- ferent representation subspaces of the input. Given an input sequence X IR Fxx where E the sequence length, and D is the embedding dimension, each head projects the inputs into different sub- spaces using linear transformations. For the i-th head, its query is defined as Q = XW? its key K, XW, and its values as V; XWV. where W w Dxdxx and WV RDdyy Here, dæ and dy represent the dimensions of the key and value, respectively. Each head then computes the attention with Head, Attention Q,, K,, Vi) = softmax (Q /VK/ V,. The outputs from all H heads are concatenated and linearly transformed to yield the final output:
MHA(X) = Concat [head]....,hadd: W
..
where W E R x is the weight matrix. Please refer to Vaswani et al. (017)) for more details.
Weight symmetry. Consider a two-layer MLP with two hidden neurons in the form of MLP (x) vTo(W 1x) = v1 (1o( + w1,222) + 020 120 (W2,121 W2,242), where is the nonlinear activation, and 01, 12 are the weights associated with the hidden neurons. the weights are initialized such that 01 U2, W1,1 W2.1, W1,2 W2,2 the two neurons will always compute identical values throughout training. This symmetry results from the fact that, at each iteration, the gradients for the corresponding weights are the same, i.e., th1, 1 2//,1,, 11 12z.2. Weight symmetry is detrimental as it implies that the two symmetric neurons do not contribute independently to the model's s learning, potentially harming the model's 's expressive power and learning capability.
4 LOSSLESS MODEL EXPANSION
Projecton
Projection
Cenon
Concet
Kosd
Honiy
Mar
Newi
Ko Ro
(b) Expand the number of heads in MHA from 2 to 3.
(a) Width expansion of MI.P from 2 to 4 (left) 5 (right).
Figure 3: Lossless width expansion with weight symmetry breaking for multi-layer perceptron (MLP) and multi-head attention (MHA). (a) Left: MIP expansion with divisible width. We replicate neurons h1/hg to h /h² and set a + 3 = with a & B. Right: MLP expansion with indivisible width. We further replicate the neuron h, to h1" and set β + I with # B 7. (b) MHA expansion with head dimension unchanged. We duplicate Head; to Head; (i.e., duplicate keylquery/value projections) and expand the projection layer in MLP module.
We decompose the expansion operator M two operators, i.e. the depth expansion operator and the width expansion operator W, each applied to individual layers.
Our expansion method mainly consists of three main components, i.e., (1) general lossless width expansion with symmetry breaking. (2) average width expansion for LayerNorm, and (3) lossless depth expansion. In the expansion process, each layer is independently subjected to these methods, ensuring layer-level lossless expansion. This entails systematic, recursive application of dupli- cating inputs for each layer in a lossless manner, and every layer, in turn, guarantees the lossless duplication of its output.
4

---

Published as a conference paper at ICLR 2024
4.1 GENERAL LOSSLESS WIDTH EXPANSION WITH SYMMETRY BREAKING
We first show how to apply lossless expansion with symmetry breaking for (1) fully-connected layers (FC-layers) and (2) multi-head attention (MHA).
Lossless width expansion for FC-layers. Transformers consist of a set of FC-layers. We first use MLP as an example to show the basic width expansion operator for the FC-layers.
For width expansion, we create copies of neurons similar to Net2Net and ber12ERR,, as this step necessary due to the nonlinear activation used in MLP. However, the essential difference is that we do NOT set the fan-out weights of replicated neurons to be equal. Out of simplicity, we use Lungl-hidden--e
MLP for illustration, and we show on the left half in Figure 3a} We first replicate neurons h1 h2 to h h2 in circular pattern. Consider the same neurons h1 and hf in the plot with the original fan-out weight v1,1: ; we can set the expanded fan-out weights to be 0011,1 and B01,1 where a + B = 1 to ensure lossless expansion.
The selection of (a. (3 corresponds a specific lossless model expansion algorithm, and our method can be considered as generalization of existing model expansion methods. Specifically, Net2Net and bert2BERT perform width expansion by setting a = 3 1/22. However, such a choice causes weight symmetry problems where two neurons learn the exact same representations when it is ini- tialized and for the subsequent training. We introduce a simple modification to fix the issue, i.c., by setting a # | is enough to break weight symmetry for commonly-use nonlinear activation α. This concept extends to cases where neurons are replicated more than twice, illustrated on the right half + 7
of Figure 3a In such cases, set coefficients such that
and # 3 +
MHA expansion. We make sure that we directly copy the entire head in a circular pattern similar to FC-layers as mentioned in the previous section. We then perform width expansion for the cor- responding key, query, and value matrices. Then, it reduces to a case similar to MLP due to the following projection matrix. Symmetry breaking is realized by setting the corresponding fan-out weights in the projection matrix differently. We illustrate the process in Figure 3b
4.2 AVERAGE WIDTH EXPANSION FOR LAYERNORM
When dealing with indivisible width increments, we need to design a specific expansion method for the LayerNorm layer. In this section, we demonstrate that achieving lossless expansion feasible provided that FC-layers are positioned before the LayerNorm layer.
Average width expansion. We first show that it is casy to perform the average expansion method padded
y1
1/2
such that the output of FC-layers with its average. We do so by adding neurons whose weights are the average of existing neu- rons. Specifically, we pad the original weight W X with rows 1 1/Dom DPen Wii and pad bias b E Pou with 1/Dou Σ Don b[i] See Figure for illustration.
LoyenNNorm
LayerNom
hv2
iS
ltz
0.5(%, + Ava)
0.B and
X;
LayerNorm layer. We now show that if the in-
Figure 4: Lossless average expansion. When the fully-connected layer right before Layer- Norm is average expanded, the output of Lay- erNorm is expanded with zeros.
put of LayerNorm is average expanded, lossless width expansion is possible. Specifically, con- sider LayerNorm layer with element-wise affine- transformation in the form of LN(:; p., b) = a Norm(.) + b, where μ1, b E PR and Dr S 2D Define average expanded of
D to be x" Rr. It can be shown that LN(x"; :1xx";/00 11" b²) Concat [x; b). 0) if μ4 = Concat [mm, and b* = Concat [b, 0J, where 0 E R Dr Ds /(sg/Dr is a scalar. See sec-
is zero vector, E I Dr Ds is an arbitrary vector, and tion E. for results and proof with a more generalized case where Dr 2 Ds-
4.3 LOSSLESS DEPTH EXPANSION
In this section, we detail our approach for increasing model depth in lossless manner.
Input dimension should he expanded as well depending on how inputs expanded.
5

---

Published as a conference paper at ICLR 2024
A
A;
Wissl
(a) Arrangement of block stacking. (b) Type-1 depth expansion.
(c) Type-2 depth expansion.
Figure 5: Lossless depth expansion. (a) We place new block next to the block where it originates. (b) For type-1 depth expansion, we set the weights of the last fully-connected layer to zeros. (c) For type-2 depth expansion, contributions from replicated neurons cancel each other. For example, assume / is a duplicate of hi, their fan-out weights to be 0 and -0 to enforce zero output.
specify the weights of the last fully-connected layer that the
Arrangement of added layers. Similar to how Chang et al. 2017); Dong et al. (2020) deal with ResNet, put added layers directly next to the source layer. For example, when expand- ing two-layer network with blocks (91.95 }2). we perform depth expansion with the resulting model (W[n)), D(W
W[gz], D[W[ga]]) See Figure 5a 5a for an illustration.
Lossless depth expansion. We now provide two ways to perform lossless depth expansion. Firstly, can simply the output of each module (MLP or MHA) to be zero, i.e. = B = O. Hence, the residual branch does not contribute to the output. This choice gives great flexibility to the rest of the parameters since we can (I) copy weights from other layers (2) randomly initialize the weights. See Figure 5b) for an illustration. Secondly, we can enforce the output to be zero by setting the summation of fan-out weights for replicated neurons to zero. With the example shown in FFigure 3a we can set the fan-out weights of replicated neurons to be =3 * 0 to ensure all outputs are zeros See Figure 5c for an illustration.
4.4 A SUMMARY OF IMPLEMENTING MODEL EXPANSION
We summarize the procedure of model expansion for Pre-LN Transformer architecture with both depth and width increments. We first average expand the embedding weights. Then, make sure the output of each layer average expanded. Hence, the input to the decoder layer is the original output padded with zeros after the last LayerNorm. We provide detailed description of our expansion method in section C.I. Furthermore, we explain how to use our method for Post-LN and Res-Post Norm architectures in Appendix D.
HOW TO TRAIN THE EXPANDED MODELS
In this section, we delve into the influence of different factors in the training recipe, in particular the maximum learning rate and the learning rate scheduler, when training expanded models.
Experiment setup. Throughout this study, we adopt VIT (Dosovisskiy et al., 2021) as our exem- plary model and train it on the standard ImageNet-1k 1k dataset. In particular, we choose to expand (i, 512) to VIT(12, 768), where 6/12 represent the number attention blocks and 512/700 76668 de- note the hidden dimensions. When training these models from scratch, we apply default maximum learning rate of 1 x 10-3 and run the training for 300 epochs with a batch size of 1024. We use a cosine learning rate scheduler that decays to minimum learning rate of 10-5, However, will modify this training recipe for continual training of the expanded model ViT(12, 768).
5.1 THE EFFECTS OF MAXIMUM LEARNING RATE
Suppose we have an expanded model. fr. that maintains the same accuracy as smaller source model, A(fs). One might naturally opt for smaller learning rate, expecting the validation accu- racy of the expanded model to continue to decrease. If this were the case, we could smooth the
IIf neurons are replicated, then have the fan-out weights be zero.

---

Published as a conference paper at ICLR 2024
x23
Validation Accuracy (%)
10
Validation Accuracy (%)
Thaning temoentee
Learning rate
75-
73-
Training
70-
Eopaniad 2
65
Dipantel Tuxe330
220 Training Epochs
100
302
60
200
200
350
120
202
300
60
120
250
300
Training Epochs
Training Epochs
Training Epochs
(a) Train loss (LR)
Valid Aec. (LR)
(c) Used LR (Sched)
(d) Valid Acc. (Sched)
Figure 6: Influence of maximum learning rate (LR; a,b) and learning rate scheduler (Sched; c,d) for training expanded Vision Transformers. Dashed and solid horizontal lines represent the validation accuracy of small and large models, when trained from scratch. (a) Train loss when changing maximum LR, (b) validation accuracy when changing maximum LR, (c) different LR scheduler used in experiments, (d) validation accuracy when changing LR scheduler. We find that (1) using a smaller maximum I.R results in smaller training loss but yields worse validation accuracy; (2) expanded models require significantly fewer epochs match the performance of the larger model.
transition between the training processes of the small model and the expanded model. However, our investigations reveal that the relationship is more complex than it initially seems.
We conducted experiments with three different maximum learning rates: x 10-3 (default), 2 x 10 and 1 x 10-4, maintaining consistent minimum learning rate of x 10 across all cases. The results are shown Figure 6b) We summarize our findings in the following paragraphs.
Performance drop early at training. An interesting observation is the immediate decrease in validation accuracy experienced by all three expanded models early during the learning rate warm- This This performance drop is correlated with the magnitude of the learning rate; the larger it is, the more pronounced the drop. This aligns with our anticipation as smaller learning rates are critical for model convergence, especially when the source model is already near local optima. Adopting a larger learning rate can displace the weights from this local minimum, leading to an increase in training loss.
Maximum learning rate and model generalization. We observe that maintaining the default max- imum learning rate is pivotal to recovering the performance of the large model. To investigate whether adopting smaller learning rates hinders model learning, we also examine the training loss of all cases, as illustrated in Figure 6al The results show that models trained with reduced learning rates incur smaller training losses compared to training from scratch. Hence, we postulate that the deterioration in performance, induced by a smaller maximum learning rate, is detrimental to the gen- eralization capability of the expanded networks rather than the optimization capability. This concept also theoretically examined by [i et al. (2020), illustrating how the learning rate can influence the sequence of learning varied patterns, therehy affecting generalization capacities.
5.2 HOW FAST THE LEARNING RATE SHOULD DECAY
After settling the maximum learning rate, the next important parameter consider is the total num- ber of epochs. Most works use the default learning rate scheduler (Wang et al.) 2023a Chen et al.) 2021a), maintaining the same number of epochs as if the model were training from scratch. We, however, note that the expanded model, having inherited knowledge from the source model, starts with a small training loss this holds true even when accounting for the significant loss drop dur- ing warm-up up. This indicates the expanded model is closer to the local optimum, requiring smaller learning rate for continued loss reduction. Thus, we should adopt a learning rate scheduler where the learning rate decays faster.
We examine four different epoch totals Tesa 130, 150, 200, and 300, with the corresponding learn- ing rate schedulers illustrated in Figure 6c Experiment results are shown in Figure 6d
We tried change the numher of warm-up steps, hut the results were not greatly affected.
7

---

Published as a conference paper at ICLR 2024
Expanded model necessitates faster learning rate decay. As depicted in Figure 6d notable observation is that employing learning rate scheduler with faster decay enables the expanded model to quickly attain the performance of the corresponding large target model. Remarkably, the expanded model requires only 130 epochs of training to match the performance of the target model that was trained from scratch, translating to computational cost saving of up 56.67%. This corroborates our earlier conjecture that expanded models need a learning rate scheduler that decays faster.
In summary, we recommend employing the same maximum learning rate used for training from scratch but with accelerated decay.
MAIN EXPERIMENTS
Vallidation Accuracy (%6
Accurare 199
THE PEEEEET
Log MLM loss
Log MLM loss
BERTBERT-FR berLBBETAA hard
bortIBENTAK) hersto
Hallidaton
Training Epochs
Training Epochs
Training iterations (x10P)
Training Iterations (x1D=)
(a) From VIT 384) to VIT 768)
(b) From (G 512) (c) From (ERT 384)
(d) From (ERT 512) to BERT 768)
(22 768)
(12 768)
Figure 7: Results of VIT on ImageNet (a,b) and BERT on English Wiki (c,d). Dashed and solid horizontal lines represent the validation accuracy/MLM loss the trained small model and target model. LEMON outperforms baselines, yielding computational savings of 56.7%, 56.7%, 25.5%, and 33.2% in panels (a), (b), (c), and (d) compared to training from scratch, respectively.
In this section, we compare our method with existing model expansion algorithms on Vision Trans- formers and BERT. We name our method LosslEss MOdel ExpansioN (LEMON), which uses the expansion algorithm explained in Section 4 with an optimized learning rate scheduler that decays
faster, suggested in section 5
Baselines. We consider several baselines to compare with our proposed method: (1) training the
target model from scratch. (2) bert2BERT-FP (Chen et al. 2015). generalization of Net2Net, (3) her22BRRT-AKI (Chen et al. 2021a), which uses advanced knowledge initialization (AKI) to break weight symmetry, (3) soft KT (Qin al. 2021) which learns the output of the source model by min- imizing the KL-divergence of the two distributions, and (4) hard KI which learns the predictions of the source model. We do not include StackBERT (Gong et 2019), Yang al. (2020), and Staged training (Shen et al., 2022) as they are not compatible with indivisible width expansion. LiGO (Wang et al. (2023a) is unavailable for direct comparison due to the absence of open-source code; hence, comparisons are made using reported values on ViT(12.512) to ViT(12,768) in sectieionn F.I Experiments of CNN and Post-LN BERT can found in section F.2 and section F3 respectively.
6.1 VISION TRANSFORMERS
Experiment setting, We adopt the default experimental setup described in section 5 Jnleess stated otherwise. For LEMON, the learning rate is decayed to its minimum value over Tiotal = 130 epochs in both experiments. Parameters choices of LEMON are discussed in section C.4
Experiment results. As demonstrated in Figure 7a and Figure 7b LEMON is able to achieve lossless model expansion. For both experiment settings, LEMON is able to recover the performance of the target model in 130 epochs, outperforming other baselines.
Several additional observations were made during the study. First, both bert2BERT-FPI and ber22ERT-AKII exhibited performance inferior to training from scratch. Second, consistent with the observations in Chen et al. (2021a) and Wang al. (2023a), soft KI did not enhance the training speed of the target model, while hard KI did, possibly by functioning akin to curriculum learning and filtering out the challenging training samples for the target model early at training.
8

---

Published as a conference paper at ICLR 2024
Table Downstream performance of BERT [2, 768) on the GLUE dataset: Large model expanded from BERT(6,384) achieves the best downstream performance. A potential reason for outperforming BERT(6,512) may be its longer training duration (165k) compared to the BERT(6,512) (132k).
Total
Dntneet
STS-B
MRPC
CoLA
SST-2
QNLI
MNLI
MNLI-mm
QQF
training steps
(Meeric)
(Corc)
(Aec.)
(Mec.)
(Aec.)
(Apc.)
(Acc.)
(Acc.)
(Are)
2206
Train Erem scniich
0.744
83.33
0.19
58.58
57. 7..00
80.25
$1.17
89.62
132K
LEMON (Ous)) from BERT(6, 512)
0.848
$3.82
0.36
90.14
58.76
80.92
$1.57
8991
LEMON (Curs), from BFRT 354)
0.866
81.81
$1.81
90.40
6.2 LANGUAGE MODELS
Experiment setting. For our experiments, we train Pre-LN BERT (Xiong et al., 2020 on masked language modeling task. The model is trained on the English Wiki corpus as per the methods in Tan Bansal (2020) for 220k iterations with 5k warmup steps and a batch size of 256. We use maximum learning rate of x2100- and a cosine learning rate scheduler which decreases the learning rate to x 10 5 Following Liu et al. (2019), we remove the next sentence prediction task and use fixed sequence length of 128 for model pre-training
We consider the following expansion procedure: (I) BERT(6, 384) to BERT 768), and (2) BERT 512) to B1RT 768). We remove KI our baseline. For LEMON, we decay the learn- ing rate to the minimum values 165k and 132k iterations for (6 384) and BERT 512), respectively. Parameters choices of LEMON are discussed in section C.4 We report the number of iterations needed to achieve a log validation MLM loss of 1.64.
Experiment results. As shown in Figure 7c and Figure 7d LEMON successfully expands smaller models without incurring loss. outperforms baselines and achieve computational cost savings of 25.5% and 33.2% for (ERT 384) and BERT (6, 512), respectively.
Downstream task. We also present downstream performance of BERT trained by LEMON on the GLUE (Wang al. 2018) dataset. We report correlation for the STS-B dataset and Matthews cor- relation coefficient for the CoLA dataset. Accuracy is reported for the remaining datasets. The results reveal that (2,T(12,768) exhibits superior downstream performance when expanded from (,RT(6,384) as opposed to being trained from scratch or being expanded from (6RT(),)12). This likely stems from its more extensive training duration (165k iterations) compared to the model ex- panded from (ERT(,,512) (132k iterations).
6.3 ABLATION STUDIES: THE EFFECTS OF THE TRAINING RECIPE
To study the effects of our proposed train- ing recipe baselines, conduct ab- lation study where we apply our training recipe on them. The results are shown Figure Sa It is shown that expanded mod- els indeed require faster learning rate de- cay. Additionally, LEMON continues outperform other baselines under the same modified training recipe.
Validation Accuracy (%)
Validation Accuracy (%6)
2 3
burtaEar- AAA-- redzo
forelt Ourr recip
Scrtcheu roope
Saxtch-Ou ridas
158
Training Epochs
Training Epochs
(a) ViT(6, 384) (12, 768). (b) VIT(6, 512) - 762.
7 CONCLUSION
Figure 8: LEMON outperforms other baselines even when they employ the same optimized learning rate method that combines lossless model ex- schedulers.
In this paper, we propose LEMON, a
pansion and optimized learning rate scheduler, showing compatibility and significant performance improvements for a variety of Transformer architectures. However, LEMON does have its limita- tions, including the need for tuning the total number of training epochs, and our evaluation scale was constrained by available computational resources. Looking ahead, we are working on extending the application of LEMON to larger models and developing methodologies for selecting optimal free parameters when initializing LEMON.

---

Published as a conference paper at ICLR 2024
REFERENCES
Mohamed S Abdelfattah, Abhinav Mehrotra, Lukasz Dudziak, and Nicholas D Lane. Zero-cost proxies for lightweight nas. arXiv preprint arXiv:2101.08134 2021.
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
Bowen Baker, Otkrist Gupta, Nikhil Naik, and Ramesh Raskar. Designing neural network architec- tures using reinforcement learning. arXiv preprint arXiv:1611.0216/ 2016.
Guillaume Bellec, David Kappel, Wolfgang Maass, and Robert Legenstein. Deep rewiring: Training very sparse deep networks. arXiv preprini arXiv: 2017.
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.
Bo Chang, Lili Meng, Eldad Haber, Frederick Tung, and David Begert. Multi-level residual net- works from dynamical systems view. arXiv preprint arXiv: 10348, 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:2110.07145 2021a.
Tianqi Chen, lan Goodfellow, and Jonathon Shlens. Net2net: Accelerating learning via knowledge transfer. arXiv preprint arXiv: 151 115511.0000 1.05641, 2015.
Tianqi Chen, Yongfei Liu, Zhendong Wang, Jianbo Yuan, Quanzeng You, Hongxia Yang, and Mingyuan Zhou. Improving in-context learning in diffusion models with visual context- modulated prompts. arXiv preprint arXiv:2312.014000 2023a.
Wuyang Chen, Xinyu Gong, and Zhangyang Wang. Neural architecture search on imagenet in four {gpu} hours: A theoretically inspired perspective. In International Conference on Learning Representations, 2021b. URL
Xiaohui Chen, Yongfei Liu, Yingxiang Yang, Jianbo Yuan, Quanzeng You, Li-Ping Liu, and Hongxia Yang. Reason out your layout: Evoking the layout master from large language mod- els for text-to-image synthesis. arXiv preprint arXiv:2311.17126 2023b.
Pau de Jorge, Amartya Sanyal, Harkirat S Behl, Philip HS Torr, Gregory Rogez, and Puneet K Dokania. Progressive skeletonization: Trimming more fat from a network at initialization. arXiv preprint arXiv:2006.0900 2020.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hier- archical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pp. 2009. doi: 10. 1011100/C//.2000 52008488
Jacob Devlin, Ming- Aeig-Wei&o Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, 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.
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 worth 16x words: Transformers for image recognition at scale, 2021.
Utku Evei, Trevor Gale, Jacob Menick, Pablo Samuel Castro, and Erich Elsen. Rigging the lottery: Making all tickets winners. In International Conference on Machine Learning, pp. 2943-2952. PMLR, 2020.
10

---

Published as a conference paper at ICLR 2024
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint arXivv1803.0330 2018.
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-256. JMLR Workshop and Conference Proceedings. 2010.
Chengyue Gong, Dilin Wang, Meng Li, Vikas Chandra, and Qiang Liu. Keepaugment: A simple information-piee preserving data augmentation approach. In Proceedings the IEEE/CVF conference on computer vision and pattern recognition, pp. 1055-1064. 2021.
Linyuan Gong, Di He, Zhuohan Li, Tao Qin, Liwei Wang, and Tieyan Liu. Efficient training of bert by progressively stacking. In International conference on machine learning, pp. 2337-2346. PMLR, 2019.
Kshitij Gupta, Benjamin Thérien, Adam Ibrahim, Mats L Richter, Quentin Anthony, Eugene Belilovsky, Irina Rish, and Timothée Lesort. Continual pre-training of large language models: How to (re) warm your model? arXiv preprint arXiv:2308.04000 2023.
Rujun Han, Xiang Ren, and Nanyun Peng. Econet: Effective continual pretraining of language models for event temporal reasoning. arXiv preprint arXiv:2012. 2020.
Song Han, Jeff Pool, John Tran, and William Dally, Learning both weights and connections for efficient neural network. Advances in neural information processing systems, 28, 2015.
Babak Hassibi, David G Stork, and Gregory J Wolff. Optimal brain surgeon and general network pruning. In JEEE international conference on neural networks, pp. 293-299. IEEE, 1993.
Bobby He, James Martens, Guodong Zhang. Aleksandar Botev, Andrew Brock, Samuel L. Smith, and Yee Whye Teh. Deep transformers without shorteuts: Modifying self-attention for faithful signal propagation. arXiv preprint 2003 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pp. 1026 1026-1034, 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp- 770-778, 2016.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprint urXiv:1503.0253: 2015.
Itay Hubara, Matthieu Courbariaux, Daniel Soudry, Ran El- Yaniv, and Yoshua Bengio. Quantized neural networks: Training neural networks with low precision weights and activations. The Jour- nal of Machine Learning Research, 18(1):6869-688 -6988 2017.
Joel Jang, Seonghyeon Ye, Sohee Yang. Joongbo Shin, Janghoo Han, Gyeonghun Kim, Stan- ley Jungkyu Choi, and Minjoon Sco. Towards continual knowledge learning of language models. arXiv preprint 1000:2210. 2021.
Zixuan Ke, Yijia Shao, Haowei Lin, Tatsuya Konishi, Gyuhak Kim, and Bing Liu. Continual pre- training of language models. In The Eleventh International Conference on Learning Representa- tions, 2022.
Yann LeCun, John Denker, and Sara Solla. Optimal brain damage. Advances in neural information processing systems, 2, 1989.
Namhoon Lee. Thalaiyasingam Ajanthan, and Philip HS Torr. Snip: Single-shot network pruning based on connection sensitivity, arXiv preprint 0rXiv:1810.023400 2018.
Namhoon Lee, Thalaiyasingam Ajanthan, Stephen Gould, and Philip HS Tor. A signal propagation perspective for pruning neural networks at initialization. arXiv preprint arXiv: (ee:::90666000 2019.

---

Published as a conference paper at ICLR 2024
Chuan Li. Openai's 's gpt-3 language model: A technical overview. https:/ hatp::://aambaaabbbb lambdalabs. 2020. Accessed: 2023-09-22.
com/bloq/ (olloo//denystiEiiin----g-o
Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks, 2020.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search. arXiv preprint arXiv: 180.009055 2018.
Shiwei Liu, Lu Yin, Decebal Constantin Mocanu, and Mykola Pechenizkiy. Do we actually need dense over- parameterization? in-time praneeetttttiiiiionnnn in sparse training. In International Conference on Machine Learning, pp. 6989-7000. PMLR, 2021a.
Yang Liu and Mirella Lapata. Text summarization with pretrained encoders. arXiv preprint arXiv:1908. 08345, 2019.
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, 2019.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 10012-10022, 2021b.
Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, et al. Swin transformer v2: Scaling up capacity and resolution. In Proceedings of the JEEE/CVF conference on computer vision and pattern recognition, pp. 12009- 12019, 2022.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. 1rXiv: 2017.
arXiv preprint
Christos Louizos, Max Welling, and Diederik P Kingma. Learning sparse neural networks through 1_0 regularization. arXiv preprint drXiv:171..01312, 2017.
Daniel Lourciro, Francesco Barbieri, Leonardo Neves, Luis Espinosa Anke, and Jose Camacho- Collados. Timelms: Diachronic language models from twitter. arXiv preprint arXiv:2202.03829, 2022.
Joe Mellor, Jack Turner, Amos Storkey, and Elliot J Crowley. Neural architecture search without training. In International Conference on Machine Learning, pp. 7588-7598. PMIRR, 2021.
Decebal Constantin Mocanu, Elena Mocanu, Peter Stone, Phuong H Nguyen, Madeleine Gibescu, and Antonio Liotta. Scalable training of artificial neural networks with adaptive sparse connec- tivity inspired by network science. Nature communications, 91):1-12, 2018.
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. arXiv preprint arXiv:2206.03/22 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances in neural information processing systems, 32, 2019.
Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl- Sohl-Dickstein, and Surya Ganguli. Exponen- tial expressivity in deep neural networks through transient chaos. Advances in neural information processing systems, 29, 2016.
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 axXiv:2105.13880 2021.
Yujia Qin, Jiajie Zhang, Yankai Lin, Zhiyuan Liu, Peng Li, Maosong Sun, and Jie Zhou. Elle: Efficient lifelong pre-training for emerging data. arXiv preprint arXiv:2203.06311 2022.
12

---

Published as a conference paper at ICLR 2024
Mohammad Rastegari, Vicente Ordonez, Joseph Redmon, and Ali Farhadi. Xnor-net: Imagenet classification using binary convolutional neural networks. In European conference on computer vision, pp. 525-542 Springer, 2016.
Esteban Real, Alok Aggarwal, Yanping Huang, and Quoc V Le. Aging evolution for image classifier architecture search. In AAAI Conference on Artificial Intelligence, volume 2, 2019.
Alex Renda, Jonathan Frankle, and Michael Carbin. Comparing rewinding and fine-tuning in neural network pruning. arXiv preprint 2ttt::0000023899 2020.
Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynam- ics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120, 2013.
Thomas Scialom, Tuhin Chakrabarty, and Smaranda Muresan. Fine-tuned language models are continual learners. In Proceedings of the 2022 Conference Empirical Methods in Natural Language Processing, pp. 6107-6122, 2022.
Sheng Shen, Pete Walsh, Kurt Keutzer, Jesse Dodge, Matthew Peters, and Iz Beltagy. Staged training for transformer language models, 2022.
Hao Tan and Mohit Bansal. Vokenization: Improving language understanding with contextualized, visual-grounded supervision, 2020.
Hidenori Tanaka, Daniel Kunin, Daniel Yamins, and Surya Ganguli. Pruning neural networks
without any data by iteratively conserving synaptic flow. Advances in Neural Information Pro- cessing Systems, 33:6377-6389 2020.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention, 2021.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint 13977,, 2023.
Asher Trockman and Zico Kolter. Mimetic initialization of self-attention layers. arXiv preprint arXiv:2305.09888, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion 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
2018
Chaoqi Wang, Guodong Zhang, and Roger Grosse. Picking winning tickets before training by preserving gradient flow. arXiv preprint arXiv:2002.07376 2020a.
Haoxiang Wang, Yite Wang, Ruoyu Sun, and Bo Li. Global convergence of maml and theory- inspired neural architecture search for few-shot learning. In Proceedings of the IEEE/CVF Con- ference Computer Vision and Pattern Recognition, pp. 9797-9808, 2022a.
Peihao Wang, Rameswar Panda, Lucas Torroba Hennigen, Philip Greengard, Leonid Karlinsky, Rogerio Feris, David Daniel Cox, Zhangyang Wang, and Yoon Kim. Learning to grow pretrained models for efficient transformer training. arXiv preprint arXiv:2303.00998. 2023a.
Ruochen Wang, Minhao Cheng, Xiangning Chen, Xiaocheng Tang, and Cho-Jui Hsieh. Rethinking architecture selection in differentiable nas. In International Conference on Learning Representa- tions, 2020b.
Yite Wang, Dawei Li, and Ruoyu Sun. Ntk-sap: Improving neural network pruning by aligning train- ing dynamics. In The Eleventh International Conference on Learning Representations, 2022b.
Yite Wang, Jing Wu, Naira Hovakimyan. and Ruoyu Sun. Double dynamic sparse training for gans, 2023b.
13

---

Published as a conference paper at ICLR 2024
Wei Wen, Chunpeng Wu, Yandan Wang, Yiran Chen, and Hai Li. Learning structured sparsity in deep neural networks. Advances neural information processing systems, 29, 2016.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, al. Huggingface's 's transformers: State-of-the-ard natural language processing. arXiv preprint arXiv:
2019.
Jing Wu, Jennifer Hobbs, and Naira Hovakimyan. Hallucination improves the performance of un- supervised visual representation learning. In Proceedings of the IEEE/CVF International Confer- ence on Computer Vision, pp. 16132-16143, 2023a.
Jing Wu, Naira Hovakimyan, and Jennifer Hobbs. Genco: An auxiliary generator from contrastive learning for enhanced few-shot learning in remote sensing. arXiv preprint arXiv:230.14611, 2023b.
Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Pennington. Dynamical isometry and a mean field theory of cnns: How to train 10,000-layer vanilla con- volutional neural networks. In International Conference on Machine Learning, pp. 5393-5402.
PMLR, 2018
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng. Shuxin Zheng, Chen Xing. Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pp.
PMLR, 2020.
Jingjing Xu, Liang Zhao, Junyang Lin, Rundong Gao, Xu Sun, and Hongxia Yang. Knas: green neural architecture search. In International Conference on Machine Learning, pp. 11613-11625. PMLR, 2021.
Aiyuan Yang, Bin Xiao, Bingning Wang, Borong Zhang, Chao Yin, Chenxu Lv, Da Pan, Dian Wang, Dong Yan, Fan Yang, et al. Baichuan 2: Open large-scale language models. arXiv preprint arXiv:2309.10330 2023.
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 arXiv:201 1. 13635, 2020.
Kaicheng Yu, Christian Sciuto, Martin Jaggi, Claudiu Musat, and Mathieu Salzmann. Evaluating the search phase neural architecture search. In International Conference on Learning Repre- sentations, 2019.
Zishun Yu, Yunzhe Tao, Liyu Chen, Tao Sun, and Hongxia Yang. B-coder: Value-based deep reinforcement learning for program synthesis. arXiv preprint arXiv:2310.03173, 2023.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks, 2017.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Infor- mation Processing Systems, 32, 2019.
Barret Zoph and Quoc V Le. Neural architecture search with reinforcement learning. arXiv preprint 2016.
14

---

Published as a conference paper at ICLR 2024
OVERVIEW OF THE APPENDIX
The Appendix is organized as follows:
Appendix A introduces the general experiment setup.
Appendix provides backgrounds and notations for model expansion.
Appendix shows details for applying LEMON on Pre-LN Transformers.
Appendix D
shows details for applying LEMON on other architectures.
Appendix provides related proofs.
Appendix provides additional experiments.
Appendix G
provides additional related works for efficient deep learning.
A EXPERIMENT SETUP
We conduct all experiments with NVIDIA-VI00 and NVIDIA-A100 GPUs. We use the official code base of DeiT Touvron et al. 2021) for training Vision Transformers and the code base of VLM (Tan & Bansal 2020) for training BERT. For CNN experiments, we adopt the code provided by Pytorch (Paszke et al. 2019
A.1 NETWORK ARCHITECTURE
For Vision Transformers, we the default network architecture adopted in Touvron et al. (2021) We implemented Pre-LN BERT in Huggingface's Transformers package (Wolf et al. 2019) such that:
Within the residual branch each Transformer block, we positioned LayerNorm precede both the multi-head attention (MHA) and multi-layer perception (MLP) modules.
For the MLM classification head, we use only one fully-connected layer (shared with the embedding).
We implemented Post-LN BERT in Huggingface's Transformers package (Wolf et al. 2019 such that:
For the MLM classification head, we use only one fully-connected layer (shared with the embedding).
A.2 DETAILED TRAINING CONHIGURATIONS
Vision Transformers. We train Vision Transformers on the ImageNet-lk Deng et al. 2009) dataset. When training Vision Transformers from scratch, we apply maximum learning rate of 1 x 10 and run the training for 300 epochs with a batch size of 1024. We use Adam W (Lohhchlovv Hutter 2017) as the optimizer. We use cosine learning rate scheduler that decays minimum learning rate of 10-5 with 5 warm-up epochs.
BERT pre-training. We train BERT (Devlin al., 2019 Xiong al.) 2020 on masked language modeling task. The model is trained on the English Wiki corpus as the methods in Tan & Bansal 2020) for 220k iterations with 5k warmup steps and batch size of 256. We use AdamW as the optimizer. We use maximum learning rate of 2 10-4 and a cosine learning rate scheduler which decreases the learning rate to 2 x 10-3, Following Liu et al. (2019), we remove the next sentence prediction task and use fixed sequence length of 128 for model pre-training.
BERT fine-tuning, For mee-tunnng task of BERT on the GLUE (Wang et al. 2018) dataset, we train 3 epochs with learning rate of x - and batch size of 32 for all tasks. We report correlation for
15

---

Published as a conference paper at ICLR 2024
the STS-B dataset and Matthews correlation coefficient for the CoLA dataset. Accuracy is reported
for the remaining datasets.
Convolutional neural networks. We train ResNets (He al. 2016 and WideResNets Zagoruyko & Komodakis) 2017) on the ImageNet- Ik dataset for 90 epochs using SGD with an initial learning rate of 0.1. We set the batch size to be 128. Learning rate is decreased by 10 times at epochs 30 and 60.
A.3 DETAILS OF BASELINES
We provide our implementation details of knowledge inheritance (KI) (Qin et al. 2021 in this section. Given training dataset denoted as DD = (x), yz) we define the total loss Lioal
Tonal fs. D)
Σ (1-a) (f + a..x (X.Y))ED
where a a scalar controlling the strength of KI; The functions fs and ft respectively represent the small source model and the large target model; The loss function Leet computes the standard training loss, such as cross-entropy, between the prediction fr. and the actual label y;. For soft KI, we set Lxi
(x, (x,fs(x))).. For hard KI, we set ΣxI where KL stands for Kullback Leibler divergence, and e is the standard basis vector.
(L(ff(x))) Carz max fx())))
During the KI process, start with an initial a value of 0.5 and linearly decrease it to zero.
B NOTATIONS AND BACKGROUNDS
In this section, we introduce basic notations in section B.1 the definition of some normalization lay- ers in section B.2 lossless expansion in vector space in sectiioon B.3 lossless expansion for operators (layers) section B.4 and the rule of consecutive application of lossless expansion methods for consecutive layers in |scction B.4.3
B. NOTATIONS
All vectors are assumed to be column vectors. We define O, to be zero vector of dimension d. We use bold-faces fad-ffced letters for vectors, matrices, and tensors. For vector v, let v[i] be its i-th entry and v[: be its first i entries. For matrix M, ler Mii, j). M[i, and [ be its (-th) entry, i-th row, and j-th column, respectively. Morcover, let M[: i. and M|::: j] be its first rows and first j columns, respectively. We MT to denote the matrix transpose of M. We use [n] where n E 24 to denote (1, (1.....)) ,n]. We use Id to denote identity mapping. We use Concat [-] to denote horizontal concatenation.
B.2 MODEL LAYERS
In this section, give the formal definition of LayerNorm LN(-) and RMS Norm RMS(-). Definition (LayerNorm). LayerNorm LN(:; B. e) of dimension D is defined as:
LN(x; B,e) = x E[x]
C p
Var[x +e
where 3 D... Definition 2 (RMSNorm), RMS Norm RMS(-; μu, E) of dimension D is defined as:
x
RMS(X; p2
where X. D.D.
Remark. In neural networks, inputs of normalization layers are usually high dimension tensors. In this case, LayerNorm and RMSNorm normally apply the last dimension separately.
16

---

Published as a conference paper at ICLR 2024
B.3 LOSSLESS EXPANSION VECTOR SPACE
In this section, we first give the general definition of lossless expansion in vector space.
Definition 3 (Lossless expansion in vector space). Given S and T are two vector spaces where the dimensions satisfy dim(T) 2 dim(S), , a vector space expansion V S - is said to be lossless if it is invertible.
Remark. Note that the identity function Id is lossless with its inverse being itseif.
Then we give few examples of lossless vector space expansions. These examples will also be used LEMON.
Example B.3.1 (Vector average expansion Vavg). Let E IR" be a vector of dimension Ds and its average Avg(x) [[[x] DE D: x(i]. Xavg is called the average expanded of dimension Dr with Dr Ds if
Xay Vayg(x) Concat
x xT Avg(x) Avg(x)
E ER
[Dr/Dal
Dr mod Ds
Example B.3.2 (Vector zero expansion Visso). Let Rs be vector of dimension Ds. xmm called the zero expanded x of dimension Dr with Dr 2 Ds if
X
Vzero(X) = Concat
x xT.,..,,'
xx 0,
E ER. Pr.
[Dr/Dal
Dr mod Ds
Example B.3.3 (Vector circular expansion Veie). Let RD be vector of dimension Ds. Xare called the circular expanded x of dimension Dr with Dr 2 Ds if
eiv
Veire (x) Concat
x xT| [: Dy mod Ds] DDr/Ds]
RR.
Example B.3.4 (Vector random expansion Vrana). Let E RR be a vector of dimension Ds. X called the random expanded of dimension Dr with DT Ds if
X Lond Vrand (x: Concat
X ,T,,',, XT
ER.r
Dr/Ds]
where Ç E IR Dr mod De is an arbitrary vector.
Remark. (1) All vector expansion examples above follow the same pattern. Specifically, when ex- panding from dimension Dg to Dr, all vector expansion methods pad first Dr/ Ds] Dg entries by repeating x Dr/ Ds} number of times. Each method deals with the remaining Dr mod Ds entries differently. (2) The random vector Ç in vector random expansion is arbitrary, so Vavgs Veewo Veir C expansion methods for vectors. In practice, neural networks like Transformers are dealing high dimensional tensors. These tensors can essentially be thought of collections of vectors. In such scenarios, we apply the expansion methods separately the last dimension of these tensors.
Vrand- (3) Here all three examples
In the following claim, show that vectors expanded by these operators are lossless.
Claim 1. Vector average expansion Varg. vector zero expansion Vzew- vector circular expansion Van and vector random expansion Vians are all lossless expansion for vectors.
Proof. The inverse function 1 Dr + Da these vector expansion methods is
v-1 x|: Ds).
Remark. In practice, we want inverse mapping of expansion methods to be easily computed just like the example above.
17

---

Published as a conference paper at ICLR 2024
B.4 LOSSLESS EXPANSION FOR OPERATORS
We then give the definition of lossless expansion for operators. These operators apply on tensors, hence our definition of lossless operator expansion based on lossless expansion in vector space. These operators can be different layers used in Transformer architectures, including LayerNorm, convolutional layers, and fully-connected layers, etc.
Definition (Lossless expansion for operators). Consider vector spaces Sth S T and T such that dim ( S. "" and dim (S I dim(TT"k). Moreover, suppose the operator is denoted with g()) S Soui We say the operator expansion
(Vn: -los-less for g()) if there exist lossless imput vector space expansion Vin S + T and lossless output vector space expansion Vout Sour + Tou such that (xx)) = (x)))) Vx S
Remark. (1) Intuitively, lossless operator expansion can be understood as follows: when using Van losslessly expanded input, the output of the & expanded operator is also Vour losslessly expanded version of the original output. (2) For conciseness, we use 'Eig) is (Vi:: V -losseess and 'E is
(Vin, -losslsss for g(-)² interchangeably (3) We only require the vector expansions Vn and V
to be invertible, we do not have restrictions on the operator expansion E.
B.4.1 LOSSLESS EXPANSION FOR MATRIX MULTIPLICATION
Then give few examples of lossless expansion for operators. We give examples for matrix multiplication since fully-connected layers are building blocks for Transformers. We first start by introducing the following three lossless operator expansion methods for matrix multiplication suming that the input dimension is unchanged so Vi Id.
Example B.4.1 (Matrix row-average expansion Erowavg): ) Let M D matrix of dimension Ds x P and its average 1.2 Σ. Mīi, Monmay is called the row-average expanded M of dimension Dr x P with DT Ds if
M Misaay Emn,avy (M) Concat
M
MT
R DrxP
[Dr/Dal
Dr mad Da
Moreover, Eromang is (Id, -lossless for M.
Example B.4.2 (Matrix row-zero expansion Enowsem). Let M E DsxP be matrix of dimension Ds x P. M pow.erro is called the row-zero expanded M dimension Dr x P with Dr Dg if
M M Pomizero = Eyowizero (M)- = Concat
MT. ....... M 0
Op Dy mod Ds
E R DTXP
Drr/De]
Moreover, Enouzoro is (Id, -lossless for M.
Example B.4.3 (Matrix row-circular expansion ). Let M E DaxP be matrix of dimension Ds x M Moociic is called row-circular expanded M of dimension Dr x P with Dr 2 Ds if
Moo.,iirc Erowccire (M) Concat
M
M (M]: Dr mod Ds.:))
PTXP
Drr/Ds]
Moreover, Eronorg is (Id, -lssless M.
Remark. Similar vector expansion examples, these matrix row-expansion methods follow the same pattern. Specifically, when expanding the number of rows from dimension Ds to DT, these expansion methods pad first [Dr DDr/ss]] } Ds rows by repeating M DDr/Ds] number of times. Each method deals with the remaining Dr mod Ds rows differently.
The following two lossless operator expansion methods assume that the output dimension is un- changed V.n Td.
18

---

Published as a conference paper at ICLR 2024
Example B.4.4 (Matrix column-random expansion Ecotsama). Let MerMe RPDDS be matrix of dimen- sion Px Ds and Ç & (Px mod Ds) is arbitrary matrix. Col is called the column-random expanded M of dimension x DT with Dr 2 Ds if
M raldrand Ecol.rund [M; ()
Concat
M MLDr/Dsl C E RxDr
DD//Dal
where
Da/
Σ
M M.
Moreover, Ecol.rand is (Vzems Ilossssss for M.
Example B.4.5 (Matrix column-circular expansion E.ol.cir). Let M E RxDs be a matrix of di- mension P x Ds and Ms = M[::: Dr mod Ds] E IR PDDr mod Ds), M is called the column-circular expanded M of dimension P Dr with Dr 2 Ds if
M colctiee Ecoletie (M) Concat
M .... MLDr/Del M ERPx
[Dr/Ds]
where
Drr/Dsl
M +
M Dr mod Ds] M(: Dr mod Ds},
and
Dr/DDsl
Σ
M Dr mod Ds :] M[: Dr mod Ds :].
Moreover, Ecol.rnd (Veire Ic)-lossless for M.
Note that lossless matrix expansion and lossless matrix column expansion can be used together with the following claim.
Claim 2. Consider matrix column expansion Endl is (Veat; Id)-lossless for M, and matrix pansion Erow is (Id, Vloss-lssssess for M. Egot o Even and Enow Evol are both (Vents Viow -lssseess for M.
The claim is easy to prove since rows and columns are expanded independently.
B.4.2 LOSSLESS EXPANSION FOR BIAS
Note that the fully-connected layer consists of a matrix multiplication followed by a bias operator. We now give examples for the bias operator B(x; b) x
Example B.4.6 (Bias average expansion Evias,avg). Consider the bias operator B(x; b) = x + where b E Ds. Binas.ng (:5 bisan,ang) = Esias,auz [B(;; is called the average expanded B of di- mension DT with Dr 2 Ds if binas,ove = Vang (b). Moreover, Epiaravg is (Vaug, Voogs-lsssess for B.
Remark. Note that can easily extend Exnas.ong Ebias.cire and Ebias. by expanding b to Veiro (b) and V (emo(b) (b), respectively. Moreover, Eiar and Epiar,eroo are (Veires Vein -lossless and (Vzero Vzem) ) lossless for B, respectively.
B.4.3 CONSECUTIVE APPLICATION OF LOSSLESS EXPANSION FOR APORATORS
In previous sections give examples of lossless expansion methods for single operators. Now, to ensure lossless when applying expansion methods to consecutive layers/operators, we introduce the following claim:
19

---

Published as a conference paper at ICLR 2024
[own') e geon [rwni) geonn (DW)) Reton
Embeddings
L(mo Exx
LayerNorm
Multi-head Attention
(a) Small source model.
Bts e Us
Tewty geto (w49) canm,
Embeddings
of geinn mzi < gomm,
a 2n Wir 2/t Nyyly 40)
(165 wath
((wn) gonn
LayerNorm
Multi-head Attention
(b) Large target model.
Illustration of LayerNorm expansion Ln and MHA expansion EMHA- We assume d dx = dy We transpose weight matrices so that they can be considered left multiplied with vectors. The vectors in black font color indicate the intermediate values of inputs while the matrices in white color indicate parameters of the module. Biases are ignored for better illustration.
Figure 9:
Claim (Lossless of consecutive application). 1f is (Va, Vo)-lossless for and & is (Vo, Ve)- lossless for g2. Then Eg[g2] E1[ (Va, -eo-ssssl for g2 o g1.
Proof. This is easily obtained if input x is V, losslessly expanded, then the output of & [g1))()); Xmid = & (X)))()(xx), V. lossless by definition. Using the fact that (-) is (V,, Voo-lsss and the input Xmid is V. losslessly expanded, conclude the proof.
Remark. By leveraging Claim 3 we can separately apply lossless expansion methods to various layers/operators larger network. The only requirement that the output vector space expansion of one expansion method matches the input vector space expansion of the subsequent expansion method.
C DETAILS OF LEMON FOR PRE-LN TRANSFORMERS
In this section, we provide detailed explanation of applying LEMON on Pre-LN Transformer archi- tecture. By Claim 3, we can deal with different modules separately. In the following sections, we delve into the details of applying expansion methods to these modules.
C.1 WIDTH EXPANSION FOR PRE-LN TRANSFORMER BLOCKS
We first recap the Pre-LN Transformer architecture. It usually consists of (1) the embedding layer, (2) several Pre-LN Transformer blocks, (3) the final LayerNorm layer, and (4) a decoder layer.
Suppose that the hidden dimension D of the transformer is increased from D3 Dr. The head dimension d unchanged during expansion. Hence, the number of heads is increased from Ds / to Dr/d. We use W.K W WY to denote the key, query, and value weight matrix for i-th head Head, in the MHA module. We use Wo to denote the projection matrix.
We use Ebeck to denote the width expansion of Pre-LN Transformer blocks. Entock can be decom- posed into (1) LayerNorm expansion εLN, (2) MHA module expansion EMIA, and (3) MLP module expansion EMPP. We introduce these expansion methods in the following paragraphs. We provide an illustration ELN and ÊMHA in Figure 9
(1) LayerNorm expansion with &LN. We define the expansion procedure for LN as follows. We use LN:: Mamna B e²) where piant = V7) E IR. Dr , = Vaero (B) E RRD Dr and e" = =7y²c with 17 =D1/ Ds] (Ds/D) Dr) to expand the original LayerNorm layer LN(:; pa, e). The expansion lossless and the proof is given in Proposition I Moreover, ε, is (Vang, Vyem) - for LN(-).
In Figure 9 we omit and ( for better illustration.
20

---

Published as a conference paper at ICLR 2024
(2) MHA expansion with EMHA. We explain how to expand MHA as follow:
wK w W in self attention. We consider the affine transformations applied to a single token x E RDs [ in a sequence in self attention the form of (x W b
(x: W (.))) = (W and vx(x; WY, [b)) (W)) where (W/)T, (W (WW) R Ps and b," b by ERAK
During expansion, we increase the dimension of R x to xd and b ",,,,, bY unchanged. Since the number of rows for (W (W?)T, (WY)T is unchanged, we only increase the number of columns by applying column-randon expansion Expl.and defined in Pxample B.4.4 to its transpose for each head, i.e., we use {Euotrand [T C/]D" Exstrand (W:)T; } and
(W.)),,
from
{ (Eadlmnd [( for the expanded weights of W w and wy, where C C., 5/ E Rddx((Dr mod De) are random matrices. Biases are unchanged.
Heads in self attention. We increase the number of heads in circular pattern. See Figure 3b for an illustration. Note that (1) When [Dr/Ds] > 1, we can set W W[Pr/Psl differently for replicated heads break weight symmetry; (2) Addi- tionally, when DJ mod Ds * O, random matrices ç.K <..C5 can be chosen differently for replicated heads to break weight symmetry. Please see see Example B.4.4 for definitions of W W. PrrrDDs) and Ç . .C..
Projection matrix in self attention. For the projection transformation the form of W + bo where W. RDxx and bo E RR2s, we use &çolcirr and &oowayg defined in Example B.4.5 and Example B.4.1 to expand the weights and biases. Specifically, we use ECool,iiee [Crwwavg Wo]lF" E RPTXDr Dr for the expanded weight of Wo. We then use Vavg (bo) E Pr for the expanded bias of bo.
Moreover, EMHA is (Vzero: Vlog--losseess for MHA(-).
(3) MLP expansion with &MP.. Consider the MLP in the form of MMP(x) = (Wiix++ bbe))+ +
2122 where a is the non-linear activation. We explain how expand MLP follow:
For the first fully-connected layer, we increase the columns by random expansion and in- crease the rows by circular expansion. Specifically, use Evolrand [Erowcirc (Wret)) and for the expanded weight and bias.
For the second fully-connected layer, increase the columns by circular expansion and increase the rows by average expansion. Specifically, we use Ecolusie [Erowavg (Wra))] and Varg (bbz2) for the expanded weight and bias.
Moreover, EMLP is (Veeoo Varg)-lossles for MLP(-).
C.2 WIDTH EXPANSION OF OTHER LAYERS
In this section, explain how to expand the rest of the layers, i.c., embedding layers and decoder layers.
Embeddings expansion with Vang- We first average expand the embedding for each token x by adding its average, i.e., with Vavg- For Vision Transformers, we do so by adding averaged channels for patch embeddings.
Decoder layer expansion with Edec. For Vision Transformers, the decoder layer is a fully-connected layer with the form Dec(x) Wa + b. We increase the rows of the matrix by applying column- random expansion to its transpose, i.e., we use Ecotraad W dec) for the expanded weights. The bias is unchanged.
In the formulation of MHA in section 3 W w w. are right matrix multiplied with the input sequence matrix X R D3 Here the form of W, x for better illustration.
21

---

Published as a conference paper at ICLR 2024
For language models, the decoder layer is shared with the embedding layer. So we have to instead scale the weight and bias of the LayerNorm before the decoder layer by 1/ Dsl. Moreover, Edec is (Vzem: Ioss-losseess for Dec.
C.3 DEPTH EXPANSION
Depth expansion is explained in the section 4
C.4 PARAMETER CHOICES
We consider the Dr S 2Ds for better illustration There are mainly the following parameters to choose for LEMON. For the non non-divisible case, we set the random parameter Ç in the LayerNorm such that Unif( 1, 1). When using matrix column-randow expansion &c, rand for the indivisible case, Çi.j iid (0, 0.023).
Vision transformers. For the width expansion parameters of the Vision Transformers, we set Wies for indivisible case and W for divisible case W + & where D
iid randomly initialized and Di.j ~ (.000.022)
For the depth expansion parameters, set the free parameters that used to cancel out replicated neurons following the distribution (0, 0.02²).
ResNets. For the width expansion parameters of the ResNet, set Wies for indivisible case and W for divisible case to be 2 W. 4, where & E (Dsx(D)-Ds) is randomly initialized and + follow the distribution used by the default implementation.
Language models. For the width expansion parameters of BERT, we set W for indivisible and 2 for divisible case to 4 where 4 E (DD is randomly initialized and d.. ~ (0,, 0.002?).
iid
For the depth expansion parameters, set the projection matrix of the MHA block and the sec- ond colne-connected layer of the MLP block to be zero matrices. Moreover, inspired by advanced knowledge initialization (AKI) (Chen al. 2021a), we append heads/neurons from the adja-
cent l
D LEMON FOR OTHER ARCHITECTURES
Though we haven't ' included experiments for Res-Post-Norm and Post-LN blocks in main exper- iments, we show that LEMON is able to perform lossless model expansion for these scenarios. We then briefly discuss how to handle RMS norm (Zhang & Sennrich 2019), which is used in LLaMa TTuvron et al. 2023) We also discuss how to apply LEMON on convolutional neural networks.
D.1 CEsSPPST-NoREE TRANSFORMERS
We consider the Transformer with the following architecture: (1) an embedding layer, (2) several Res-Post-Nomn blocks, and (3) the final decoder layer
D.1.1 WIDTH EXPANSION
The only difference between the expansion methods of Res-Post-Norm Transformers and Pre-LN Transformers is that we zero expand embedding vector for each token with Vzero-
For the MHA and MLP modules, we use the exact same expansion introduced seccietio C. where it is (Varo. Vavg) - lossless for MHA and MLP. Consequently, our expansion is (V Varoo Vzeno lossless
'In fact only need to deal with such cases in our experiments.
This still lossless since the last layer a left-multipleed with zero matrix followed by adding a zero vector.
10 We there is final LayerNorm before the final decoder layer.
22

---

Published as a conference paper at ICLR 2024
for Res- Pess-ost-NNomm Transformer blocks. Since the last decoder expansion is (Veero Id)s-lsssless for Dec, our expansion method is strict lossless.
D.1.2 DEPTH EXPANSION
For increasing depth, we only need to set the weights and bias of the LayerNorm for each added layer to be all zeros.
D.2 POST-LN TRANSFORMERS
For Post-LN Transformers, we can only deal with divisible cases, i.e., Dy mod Ds = 0. Sup- pose Dr/ Ds n, in this case, all the embedding and outputs of modules (MIP and MHA) duplicated times and hence lossless. The only difficulty is to deal with depth expansion.
Depth expansion. Suppose we given a pre-trained Post-LN Transformer block gi (x)
LN1 (Modulei(x) + x) μi @ Norm(Modulee(xx) + x) + bj- First we expand Modulei to Module) 0,* so that it outputs zeros. Then we can create two expanded layers sf where of (x*)
1@Norm(Module) 0,* ++x*)++++++++++++ Norm(x") and 05(x)) pi ONorm(Modulee(x"" +x*)
bj. It is easy to show that o is lossless where we use the fact that Norm(Norm(x)) = Norm(x).
D.3 TRANSFORMERS WITH RMS NORM
RMS Norm Zhang Sennrich 2019 is used by foundation models like LLaMa (Touvron et al. 2023) and Baichuan Yangg et al. 2023 See Definition 2 for the definition of RMS Norm. Suppose
we want to expand the RMS Norm from dimension Dg to Dr, we use the following expansion.
RMS Norm expansion with ERMS- We use RMS( Mand: e*) where U = rm) E RR Pr and nfe with 1) = [Dr/Ds (Ds/D) Dr) to expand the original RMS Norm layer LN(-; μ4, B, c).
The expansion is (Vzero
for RMS(.). The proof is provided in PPropssitin 4
D.4 CONVOLUTIONAL NEURAL NETWORKS: RESNET
We use Conv(k (k x ki Cin, Cout) to denote convolutional layer with Ci in-channels, Com out-channels, and kernel size io x k We assume the kernel weight W RCaxC..kxt and bias b CR..
We use BN and ReLU to denote BatchNorm and ReLU, respectively. ResNet and WideResNet with more than 50 layers consist of multiple Bottleneck blocks, where there are sub-blocks: (1) Conv(D. D. Ds. x 1)-BN-ReLU, (2) Conv(Ds. Ds, 3 x 3)-BN-ReLU, and (3) Conv(Ds, Dx11)) in the residual branch.
We consider expanding ResNet to WideResNet.
Width expansion. To apply width expansion, we do the following:
(1) For the first sub-block, increase the number of out-channels of the first convolutional layer from Ds to Dr. Specifically, the expanded weight satisfies :. Wi mod Ds, :..,, :]. Vi [Dr], and b² = b/: mod Ds). Vi [Dr]. The output of the convolutional layer will be also in circular pattern in the channel dimension. This also holds true after the application of BatchNorm and ReL.U since the statistics of BatchNorm are computed within channels.
(2) For the second sub-block, increase the number of out-channels and in-channels of the first con- volutional layer from Ds to Dr. We apply the same operation to the out-channed dimension similar to (I). For the in in-channel dimension, we need to make sure that the weights of replicated channels sum up to the original weight. Specifically, suppose that the replicated channels indices are denoted C, = fili mod Ds = 2}. Then we need to set Eiec, W"i,, = Wik, :.;: for lossless ex- pansion. Moreover, we need to make sure W a, b. * W* b. Vi.j C., e [Cin] k [k] = [Caa] for symmetry breaking.
(3) For the last sub-block, increase the number of in-channels of the first convolutional layer from Ds to Dy similar to (2).
Depth expansion. For depth expansion, we simply set the weight and bias of the last BatchNorm layers in the increased layers to be zeros.
23

---

Published as a conference paper at ICLR 2024
E PROOFS
E. PROOFS FOR TRANSFORMERS WITH LAYERNORM
In this section, we first show that three main components εLN, EMIIA, and ÊMPP are lossless. Then, we prove that LEMON defined in Appendix Clis lossless.
We first start by showing that our LayerNorm expansion EIN defined in section C.I is lossless.
Proposition 1 (Lossless expansion for LayerNorm εLN). Consider LN pu. G.c) of dimension Dg where f 3 E IR Ds. Define average expanded of x E RP of dimension Dr to be Xayg = Vavg(x) E Rrr where Dr Dg- If yand rannand(1 (p) E Dr, Pixo (3) RDr,, and c* 11², where 1) = v [DT Ds] (Ds/ Dr), then
LN(Xaugi Prands Peono ()) Vzeno (LN(X; (1,
Proof. Since [[x Dr x Dr (rr/ Σ. Da x(i] (Dr mod D3 D5))
[|r] and Var [ Dr (LDT) Ds Ds Var [r]x]] (Dr mod Ds) 0) Var(x).
For 1 sis Dr Ds Ds:
X [i]
E avgl
LN LN(Xivgi Mrand Biero E [i]
M [+] + B
Var (xiig] + c
x(i mod Ds} E[x]
c malii mod Dsl + Bli mod Ds}
/ +
Veeno (IN((: M.BBCc))
For DD/DssDS sis DT:
X [i] E xxg
LN(Xivgi Prand Br c* [i]
F [i] + B3oro
/ Var[xing] +e
E[x] [[
OnC mod Ds]
"/ Var x +e
0
u4 (e.))
Hence, LN(x@ Prans Bieno E*) V zeroo (LN(Xx μ. 3,)).
Remark. When Dr is divisible by Ds, then y) = 1. Hence, explains why simply circularly expanding LayerNorm is lossless in such scenario.
Proposition | naturally leads to the following corollary.
Corollary 1. EI introduced in Desinition is (Vagz Voso)-lsssless for LN(-).
we are ready to prove that EnHaa and ÊMII are lossless. We first show that EMHA lossless in Proposition 2.
Using Claim 3)
Proposition 2 (Lossless of EMHA). EMPIA defined in section C. is (Voe,, Vavg) -lsssess for MHA.
Proof. Consider sequence input X Ex is expanded losslessly by Vieno to X E ED. We expand the source small MHA such that the target large model is MHA* EMIIA (MHA).
We first check the key, query, and value of each head Head, such that i < for the large model MHA". We denote them as K Q V RExxx, Note that biases bK b by Rdk not expanded. Hence, these outputs are identical the output of
Dz/d
24

---

Published as a conference paper at ICLR 2024
the small source model K, Q V, E Rxxx since (W.) J. (W?) (W expanded by Ec. rand. which (Vuro. -lssssssssss Consequently, Head, = Attention Q K V;) = softmax (Q; (K:)T//dK V is identical to the output of i-th head of the MHA the source small model, which is Head,,
Since heads are circularly expanded, the output MHA" is also Veire lossless.
Finally, since W is expanded by Ecol,ire and Erowarg. which is (Veine: Vang) llag)slossless. With the fact that bias bo is not expanded (unchanged), we obtain the result that EMHA is (Vzero: Vavg)-lossless for MHA.
We then show that ÉMLP is lossless in Proposition
Proposition 3 (Lossless of EMPP). This is easily obtained since the first fully-connected layer is (Vzem )eire) -lossles. Hence, the output is Veire losslessly expanded. After applying element-wise nonlinear activation, the output is still Veire losslessly expanded. Since the second fully-conneche layer is (Vam, V -losslsss,,,,, we conclude the proof that EMP. is (Vzms, Varz) -lossles for MLP.
Hence, using Proposition 2 and Proposition 3 aloog with Claim we obtain the following Corol- lary 2 and Corollary 3
Corollary 2.
The expanded Pre-LN MHA module EMHI((HH) ELN(LN) is (Varz, Varg) -lossless for MHA LN.
Proof. Since Ln is (Vavg: Vicro) -lossless for LN, and EMHA (Vzero Vavg) llossless for MHA. The result is obtained by Claim 3
Corollary 3. The expanded Pre-LN MLP module EMIP(MPP) o EIN(I.N) is (Varz, (aw)) -lossless for MLP LN.
By incorporating the residual connection, obtain the following corollary.
Corollary 4. The expanded Pre-LN modules (Pre-LN MHA/MLP) with residual connections are (Vang, Vovz)-lossle llsssesss for the original Pre-LN modules with residual connections.
Once again using Claim 3.
naturally obtain the following corollary.
Corollary 5. The width- exppnded Pre-IN Transformer layer Estoct is (Vang, -losssloss for g.
Finally, by considering the embedding layers and encoder layers, we show that LEMON is lossless.
Corollary 6. LEMON introduced in section C.I is (Id, Id)--lssless for Pre-LN Transformers, i.e., strict lossless or identical.
Proof. Since embeddings are average expanded, the output of Pre-LN Transformer blocks are av- erage expanded. Hence, outputs of the final LN before the encoder is zero expanded. Since the decoder layer expansion is (Vzem, -lossossless for Dec(.), we obtain the result that LEMON is (Id, Iosslesss
E.2 PROOFS FOR TRANSFORMERS WITH RMS NORM
In this section, we show that ÊRMS defined section D.3 is lossless.
Proposition (Lossless expansion for RMS Norm ERMS), Consider ( Ja, e) of dimension Ds where μ E Ds. Define zero expanded of Ds of dimension DT be x Veen xx) E I Dr, where DT Ds. If W'annd (1Vran(((4) Dr, and e" 1728, where
7 / Dr Ds * (Ds/Dr). then
RMS 'rro F Vieno (RMS(X XC: (),
25

---

Published as a conference paper at ICLR 2024
Proof.
For sis Dr/Dss Ds:
RMS (x Miandi E" [1]]
(x )2 + M [i]
Dr
V Dr
x/i mod Ds}
3 muli mod Ds]
Dr/Dg
Dg >
iisl (x(1])² + 17²e
Dg
Dr
x/i mod Dsl
@ mulli mod Dss
"V D5 ΣBs D5 (xli)) +
Veeoo [RM5(x [i]
For /rr/Ds]D Ds sis Dr:
RMS j Frand: c²)[1] =
zero
Prand i
Di=1 (Xtema) )2 c
Dr
DT
4 Prand
Dr E.r (xForo) + c
Vzero (RMS(x: ,e, c)][s]
Hence, RMS(X Frand c*) = Vzero (RMS(x μl c)).
Proposition 4 naturally leads to the following corollary.
Corollary 7. ERMS introduced in vection D).3 is (Vzero: Viero) -losstess for RMS(-).
ADDITIONAL EXPERIMENTS
F.1
COMPARISON WITH LIGO
Validation Accuracy (%)
80
75
65-
60-
Scratch
Ours
55
0
50
100 150
200
250
300
Training Epochs
Figure 10: We expand ViT 384) to VIT(12, 768). Our expanded model recovers the performance of the target model with 85 epochs (28.3% compared to training from scratch).
LiGO (Wang et al. 2023a) is unavailable for direct comparison due to the absence of open-source code. Hence, we compare them with their reported values. Note that our method is lossless only for Pre-LN Transformer architecture while LiGO reports their results for language models mainly on Post-LN BERT and RoBerTa. As a consequence, we compare our results with LiGO on VIT(12, 384) (ViT-Small) VT((1,, 768) (ViT-Base [I The result is shown in Figure 10
Note that DeiT without distillation exactly VIT.
26

---

Published as a conference paper at ICLR 2024
Our method is able recover the performance of the target model with 85 epochs, leading 71.67% computational saving. is higher than the reported value of 55.40% for LiGO
F.2 CONVOLUTIONNLL NEURAL NETWORKS
60.0
Validation Accuracy (%)
77.5
75.0
72.5
70.0
65.0
bert2UERT-AN Epochó0
EpochGO
62.5
scratch Epoch-00
Ours Epoch60
0
20
40
60
80
Training Epochs
We expand ResNet-50 to WideResNet- 10. Our expanded model (Ours Epoch60; red) recovers the performance of the target model within 60 epochs (33.3% compared to training from scratch). ben22BRT-AKKI (bert2BERT-AKI Epoch60; green) is unable accelerate the training compared training from scratch (scratch Epoch60; pink). Note that LEMON lossless. How- ever, the accuracy of the model expanded by LEMON decreases after epoch since there is no learning rate warm-up phase.
Figure 11:
We expand ResNet-50 WideResNet-110 assess the versatility and efficiency of LEMON in comparison to the bert2BERT-AKE method, known for its performance in the main manuscript. We utilized optimized learning scheduler with maximum rate of 0.1 (default), decaying the 20th and 40th epochs.
Results. We show the result in Figure 11 LEMON is able to recover the performance of the large network in 60 epochs, achieving 33% computational savings. Note that bertBBERT-AKI shows inferior performance compared training from scratch. We hypothesize that this might be due to lack of compatibility of bert2BERT-A -AKI with the ResNet architecture
Note that DeiT-Base (ViT-Base) has final validation accuracy of 81.00% for LiGO, which is lower than the 81.70% reported value of the official DeiT and implementation.
27

---

Published as a conference paper at ICLR 2024
80.0
Expanded, 0.1 2.1*ad defoult
5
Expanded, default
Training from scalth
Validation Acc (%)
77.5
75.0
Train loss
72.5
3
67.5
65.0
Expanded, default
62.5
Expanded. default Ir
Training from scatch
60.0
0
20 40
60
80
20
40
60
80
Training epochs
Training epochs
(a) Training loss
(b) Test accuracy
Figure 12: Training loss (a) and test accuracy (b) comparison of training from scratch with the max- imum learning rate 0.1 (Training from scratch; blue), model expanded by LEMON trained with the maximum learning rate 0.1 (Expanded, default lr; purple), and model expanded by LEMON trained with the maximum learning rate 0.01 (Expanded, 0.1 default Ir; red). Using smaller learning rate leads to smaller training loss but worse generalization performance.
Effeets of maximum learning rate. To understand how different maximum learning rates impact the performance model expansion, we conducted similar experiments. Specifically, we compared the following setups: (1) Training model from scratch with maximum learning rate of 0.1, referred to as 'Training from scratch'; (2) A model expanded using LEMON and trained with the default maximum learning rate of 0.1, denoted 'Expanded, default Ir:: and (3) A model expanded using LEMON but trained with reduced maximum learning rate of 0.01, termed 'Expanded, 0.1 default Ir'.
In line with the observations in Transformer architectures, we noticed that a smaller learning rate tends result lower training loss but potentially affects generalization performance adversely.
F.3 POST-LN BERT
2.4
Training
2.3
bert2BERT-AKI
berL2BERT-FPI
Log MLM M loss
2.2
LEMON (Ours)
2.1
2.0
1.9
1.7
Training iterations (x103)
Figure 13: We expand Post-LN (ERT 384) to BERT(12, 768). Our expanded model achieves a log validation loss of 1.67 within 137k steps (63.43% compared to 216k steps for training from scratch).
28

---

Published as a conference paper at ICLR 2024
In this section, present our experiments conducted on Post-Layer Normalization (Post-LN) BERT models to further validate the effectiveness of LEMON. Specifically, we focused on expanding BERT(6, 384) to (1, 768). We set target log validation MLM loss of 1.67 for this ex- periment.
We trained the expanded model using LEMON for 143k steps. The results, detailed in Figure 13 demonstrate that LEMON was able to achieve the targeted log validation MLM loss 1.67 within just 137k steps. This result translates a computational cost saving of 36.57%, compared training B1R 768) from scratch.
G MORE RELATED WORKS
Efficiency in deep learning can be achieved in multiple ways. In this section we provide brief overview of efficient deep learning regarding model training and inference, distinguishing it from methods addressing data efficiency (Gong et al. 2021 Wu et al. b).
Efficient deep learning. In the realm of deep learning, the drive for efficiency has led researchers to develop multitude of methods aimed at optimizing model efficiency. Techniques such as neural architecture search (NAS) (Zoph & Le 2016) [iu et al., (0188) have been employed to automate the discovery of optimal network architecture. Quantization (Rastegari et al.) 2016: (Hubara et al.) 2017) refines the numeric precision of model parameters to boost computational speed. Knowledge distillation (Hinton et al. 2015) and knowledge inheritance (Qin et al. 2021 allow target models to inherit the knowledge of their source counterparts. Neural network pruning @LeCun et al. 1989) involves removing unnecessary connections to accelerate model training inference. Finally, model growth methods (Chen et al. 2013) directly use the weights of source models to initialize the large target models.
Neural architecture search (NAS) has emerged as a promising solution for automating the pro- cess of neural architecture design, eliminating the need for labor-intensive manual designs across various deep learning tasks. Initial methodologies leveraged reinforcement learning (Zoph & Le 2016 Baker et al. 2016) and evolutionary algorithms (Real al. 2019) identify high-performing architectures. Despite their success, significant draw back was their computational demands. Ad- dressing this, DARTS (Liu et al.) (018)) introduced continuous relaxation of architectural represen- tation, allowing for search via gradient descent. However, DARTS can be challenging to optimize, and its weight-sharing approach has been criticized for potential performance degradation (Yu al., 2019 Wang et al.) (2020b). Seeking further efficiency, Mellor et al. (Mellor et al.) (02T) introduced a training-free NAS, which evaluates randomly initialized architectures, thus fully eliminating neu- ral network training during the search phase. Subsequent Waining-frree methods explored searches using Neural Tangent Kernel (NTK) (Xu et al. [021]] Chen et al. 2021b Wang et al. 2022a), linear regions (Chen al. 2021b), and criteria related to pruning (Abdelfattah et al. 2021).
When considered alongside model expansion, NAS holds potential for determining the optimal num- ber of layers and hidden dimension of the large target model.
Neural network pruning. Pruning techniques can be broadly classified based on their timing into three categories: post-hoc pruning, pruning-at-in iatiing-----niiialzzttioo methods, and moaning-during- training methods. (1) Post-hoc pruning method removes certain weights of a fully-trained neural network.
Post-hoc pruning was initially proposed accelerate model inference (LeCun al. 1989 Hassibi
et al. 1993 Han et al. 2013), while lottery ticket works (Frankle & Carbin 2018 Renda et al.- 2020 shifted towards uncovering trainable sub-networks. (2) SNIP (Lee et al. 2018) is one of the pioncering works of pruming-ar cemmingg-t-initiinn at-initialization methods that aim to find trainable sub-networks with- out any training. Subsequent research (Wang et al. 2020aa Tanaka et al. 2020 de Jorge al. 2020) Lee et al., 2019 Wang et al. (022b) introduced varying metrics for pruning at the network initializa- tion stage. (3) Finally,
methods prune or adjust DNNs throughout training. Early works incorporate explicit to (Louizos et al., (017) or é (Wen al.) 2016) regularization terms to encourage sparsity, hence mitigating performance degradation commonly associated with post-hoc pruning. More recent techniques like DST methods (Bellec et al., 2017| Mocanu et al.) 2018 Evei et al., 2020) Liu et al.) [021a] Waan et al. (2023b) allow for adaptive mask modifications during training while adhering to specified parameter constraints.
29

---

Published as a conference paper at ICLR 2024
Neural network pruning has potential synergies with model expansion, akin to the dynamics of DST. A combined approach could involve iterative increases and decreases in hidden dimensions or layers during training, potentially accelerating training speed.
30


Reviews:
[{"review_id": "siGgRsxNY0", "reviewer": "Reviewer_u5r2", "summary": "The authors proposed a lossless model expansion method which initialize scaled models using the weights of smaller pre-trained model. Specifically, the proposed method break the symmetry of replicated neurons by setting their fan-out weights to be unequal, and introduce average expansion to deal with LayerNorm for indivisible width increment. Besides, the authors explored the training recipes for the expanded models and proposed an optimized learning rate scheduler that decays more rapidly than training from scratch. Experimental results show that the proposed method can effectively expand both Vision Transformer and BERT, while significantly reducing the training overhead.", "questions": "Please refer to the weakness part.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The motivation is clear. The author focus on scaling deep neural networks in effective way by leveraging the knowledge acquired by their smaller counterparts.\n2. The paper is well organized in terms of written description. The authors provided easy-to-understand diagrams.\n3. The idea is technically feasible and the authors provide detailed proofs in appendix.", "weaknesses": "1. The challenge arising with the ‘symmetry breaking’ is described in the third paragraph of section Introduction: “the expanded model will never gain more capacity than the source model.”This statement raises confusion as training a model with smaller capacity but larger size appears to be of limited value, which incurs greater overhead but achieves limited performance. \n2. Have the considered baselines for expansion in Section 6 been confirmed to be lossless? If not, it is necessary to present the gap with the original model. If they are indeed lossless, an analysis should be provided to explain why the proposed method achieves higher validation accuracy compared to AKI, which also breaks symmetry, as shown in Figure 8.\n3. The results in Table 2 are a bit confusing. It is unclear whether model expansion or longer training duration indeed contributes to the improved performance.\n4. The novelty of this paper seems quite limited. The key idea of model expansion seems a simple extension of net2net. Are there any essential technical differences? In my opinion, a simple extension of an existing approach is insufficient for a top-tier conference.\n\n\nAfter rebuttal:\n\nThanks for the detailed response from the authors. Although the key idea is similar to net2net, there are still some new contributions in expanding layernorm and attention. Thus, I decided to raise my score.", "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": "VJEu8zHrL3", "reviewer": "Reviewer_JSXE", "summary": "**Idea**: \n* This paper introduces a method for initializing scaled models using the weights of their smaller pre-trained counterparts. The method allows for expanding neural network models in a lossless manner, increasing depth and width without sacrificing performance.\n* The paper introduces lossless layer expansion techniques, including row-average expansion, row-zero expansion, column-random expansion, and column-circular expansion.\n* The expansion procedure for LayerNorm and Multi-head Attention (MHA) modules in Pre-LN Transformer blocks is explained, showing that the expansion is lossless and preserves the properties of the original layers.\n\n**Experiments and Analysis**: \n* The method is versatile and compatible with various network structures, although the _experiments are only shown on Vision Transformers and BERT_. LEMON outperforms baselines on these architectures in terms of performance and computational cost.\n* Detailed explanations and insights into various techniques and approaches for training deep neural networks are provided, with a focus on language models.\n* The authors investigate the effects of maximum learning rate and learning rate scheduler when training expanded models.\n* LEMON is compared to a similar method called LiGO and shows better results in terms of computational saving.", "questions": "Can \"incremental\" model expansion help achieve even better performance? For example, instead of expanding from \"Model (small) pretrained --> Model (huge)\", would it be better to expand in steps as \"Model (small) pretrained --> Model (middle) --> Model (big) --> Model (huge)\"", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "### S1 - Interesting technical contributions\n* The authors provide a comprehensive and detailed exploration of lossless model expansion techniques (e.g. row-average expansion, row-zero expansion, column-random expansion, and column-circular expansion), including addressing the challenges of symmetry breaking and indivisible width increments.\n* Provide valuable insights into training recipes for expanded models, including an optimized learning rate scheduler that can further enhance performance.\n\n### S2 - Good results and experimental analysis\n* Extensive experiments with ViT and BERT are shown with a thorough investigation of the effects of maximum learning rate and learning rate scheduler when training expanded models.\n* The proposed method achieves similar performance to the original models with fewer training epochs, highlighting its efficiency and effectiveness.\n* LiGO is a similar very recent method, and LEMON shows better results in terms of computational saving.", "weaknesses": "### W1 - Experiments limited only to ViT and BERT\n* The paper could benefit from experiments on the generalizability of LEMON to other architectures beyond Vision Transformers and BERT models. For example, CNN models are completely unexplored in terms of experiments.\n* I suggest adding model expansion experiments for ResNet18 --> ResNet50 and EfficientNetB0 --> EfficientNetB4 (or other variants).\n\n### W2 - Lacks theoretical analysis/explanation of \"effect of learning rate and schedule\"\n* Sections 5.1 and 5.2 experimentally study the effect of learning rates and schedules. However, the paper lacks a theoretical analysis of why this happens. For example, why does a small learning rate lead to lower final performance? I think only experimental verification is not enough and this requires some theoretical analysis.", "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": "Lt7zRLHfZz", "reviewer": "Reviewer_pAyq", "summary": "This paper proposed a model expansion algorithm that uses the pre-trained parameters of a smaller model to initialize a larger model. The proposed algorithm allows expanding model’s width and depth to arbitrary width and depth for most Transformer variants. The algorithm ensures that the larger model has the same output as the smaller model (thus does not require calibration dataset) to preserve the small model performance while having symmetry breaking for continuous training to further optimize the larger model. The expansion algorithm alone does not ensure that the larger model can achieve the same performance as the same sized model trained from scratch using smaller training cost. The authors found right training configuration is critical for obtaining this goal.", "questions": "The suggestions are listed in weakness section.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method is simple, yet it allows expanding smaller model to arbitrary width and depth (not necessarily indivisible by the width and depth of the original models) while ensuring the output of the expanded model stays the same as smaller model and parameter symmetry breaking. \n\n2. The expansion can be performed on individual modules of a Transformer. This localized expansion ensure compatibility with different Transformer variants. \n\n3. The expansion algorithm alone does not give desired performance. The authors explored different training configuration, including learning rate and scheduler, for more desired performance and found the training configurations greatly affect the results. \n\n4. The authors performed ablation study to isolate the effect of optimized training configuration from expansion algorithms to make sure the proposed expansion algorithm indeed performs better compared to baselines.", "weaknesses": "1. In Figure 7 (c) and (d), the loss curves for BERT language modeling are still decreasing when training is stopped. It might be better to train the model till convergence to evaluate whether or not the proposed method can have the same performance as the model trained from scratch.\n\n2. Vision Transformer is a pre-norm Transformer, and in BERT language modeling, the authors also used the pre-norm variant. Since the authors claimed compatibility of the algorithm with different variants, it would be better to see the experiments on different variants (at least a post-norm variant) to verify the claim. While in Appendix, the authors show lossless expansion for other variants, it is also important to evaluate the performance metrics. \n\n3. Since this work also studies the initialization of model parameters, it might be interesting to compare the proposed idea with other initialization approach, such as Mimetic initialization (https://arxiv.org/abs/2305.09828, also mentioned in the related work)", "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": "J6oAJHIwM8", "reviewer": "Reviewer_XnGd", "summary": "This paper proposes a model expansion method that utilizes knowledge from existing smaller models. The authors analyze and design for different model structures, breaking the symmetry of repeating units by setting unequal output weights. This leads to a lossless model expansion approach. The training process is thoroughly analyzed, and the method achieves promising experimental results for both Vision and NLP Transformer models.", "questions": "- Regarding the significant drop in accuracy during the early stages of training, it may be attributed to the transfer of pre-trained parameters from one local optimum to another. Considering the requirements of certain real-time systems, such performance drop in model accuracy is unacceptable.  I think maybe we can smooth out this process by, for example, setting a mask to control the number of trainable parameters in each epoch, gradually transitioning them to avoid a drastic drop in accuracy. Overall, solving this problem could be of significant importance for many real-world applications.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper provides a comprehensive review of related work, presenting the current research status of expanding small models into larger ones.\n- The method is highly versatile, as it is designed and analyzed for different structures within the Transformer, making it applicable to commonly used Transformer architectures. Moreover, they provide practical tuning suggestions for training.\n- The analysis and observations made in the experiments are interesting.\n- The authors demonstrate significant acceleration during the training process.", "weaknesses": "- After reading this paper, I would like to know more about the practical application of the model expansion method in real-world scenarios. Please provide more examples to illustrate the effectiveness and applicability of the proposed approach.\n- How does the performance change in Figure 6b when using a learning rate larger than the default value?", "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)