Exploring Learngene via Stage-wise Weight Sharing for Initializing Variable-sized Models

1 Introduction

Vision Transformers (ViTs) have become increasingly popular, showcasing their remarkable performance across a wide range of vision tasks [Dosovitskiy et al., 2021; Liu et al., 2021; Wang et al., 2022b; Oquab et al., 2023; Yan et al., 2023]. In practical deployment, it is often necessary to train models of various scales to flexibly accommodate different resource constraints. These constraints may exhibit significant diversity, such as mobile devices with limited available resources and computing centers with substantial computational capabilities. Clearly, training each target model from scratch provides a straightforward solution, where weight initialization is a crucial step prior to training which aids in model convergence and affects the final quality of the trained model [Glorot and Bengio, 2010; He et al., 2015; Arpit et al., 2019; Huang et al., 2020].

∗Co-corresponding author.

Nowadays, a variety of large-scale pretrained models, developed by the research and industry community, are readily available to transfer and finetune the learned weights for diverse downstream tasks [Radford et al., 2021; He et al., 2022; Touvron et al., 2023; Oquab et al., 2023]. However, such scheme needs to reuse the original whole pretrained model parameters every time facing different downstream tasks regardless of the available resources. Unfortunately, for many pretrained model families (MAE [He et al., 2022]), even the smallest model (86MViT-Base [Dosovitskiy et al., 2021]) can be considered extremely large for some resource-constrained settings. To tackle this, developers would have to first pretrain target model to meet certain resource demand, which is time-consuming, computationally expensive and lacks the flexibility to initialize models of varying scales. Thus, how to flexibly initialize diverse models to satisfy different resource constraints arises as an important research question.

Recently, [Wang et al., 2023] proposes a novel learning paradigm known as Learngene to achieve this goal. As shown in Fig. 1(a), Learngene firstly learns one compact part termed as learngene, which contains generalizable knowledge, from a large well-trained network termed as ancestry model (Ans-Net). Subsequently, learngene is expanded to initialize variable-sized descendant models (Des-Net), after which they undergo normal fine-tuning. Based on the gradient information of Ans-Net, Vanilla-LG [Wang et al., 2022a] extracts a few high-level layers as learngene and combines them with randomly-initialized layers to build Des-Nets. TLEG [Xia et al., 2024] extracts two layers as learngene which is linearly expanded to build varying Des-Nets. LearngenePool [Shi et al., 2024] distills one large model into multiple small ones whose layers are used as learngene instances and then stitches them to build Des-Nets. However, there exist several limitations in previous works. Firstly, the extracted learngene itself only contains well-trained parameters but lacks crucial knowledge essential for the subsequent expansion process. Secondly, the learngene learning process provides insufficient guidance about how to effectively and conveniently expand the learngene, thereby constraining its potential for initialization. Thirdly, empirical performance still lags far behind the pre-training and fine-tuning approach.

To validate the importance of such knowledge and guidance, we study from an idea termed Simple-LG: we first train a vanilla 5-layer model whose layers compose learngene layers, and then expand its layer to initialize Des-Nets of required depth, as shown in Fig. 1(b). Obviously, neither the welltrained layers nor the learning process of the 5-layer model possess knowledge or guidance for the subsequent expansion process. As a result, this strategy brings severe performance degradation compared to well-trained ones. For example, it leads to a 27.1% accuracy degradation of initialized 12-layer Des-Net in ImageNet-1K [Deng et al., 2009] classification without any fine-tuning. Moreover, we display the validation loss value of these initialized models in Fig. 1(d), which delivers two important messages: 1) These initialized models are not fine-tuned, yet they attain meaningful loss value, which demonstrates the potential of expanding learngene layers to initialize models. 2) With the increasing layer number of models, i.e., the increasing number of expanding layers, the discrepancy between the validation loss of initialized models and that of well-trained ones becomes larger.

Upon closer observation, if we treat each learngene layer as an individual stage, we can recursively update each layer at its stage, which equals to sharing weights across multiple layers within each stage throughout the training, shown in left of Fig. 1(c). In this way, we seamlessly integrate stage information into each learned learngene layer, i.e., learngene layer actually contains knowledge of multiple weight-shared layers within its stage. Moreover, we could emulate the learngene expansion process via adding weight-shared layers within each stage during learngene training, thereby providing clear guidance on how to expand, namely, expanding the weightshared layers at its corresponding stage, shown in Fig. 1(c). Both stage information and expansion guidance are necessary: 1) expanding learngene layers which lacks stage information destroys the intrinsic layer connection (See Simple-LG). 2) Without expansion guidance, the position of expanded layers remains uncertain. Based on this insight, we present Stagewise Weight Sharing (SWS), a simple but effective Learngene approach for efficient model initialization. SWS divides one Transformer into multiple stages and shares the layer weights within each stage during the training process. Specifically, we design and train an auxiliary model (Aux-Net) whose layer weights are shared via SWS to obtain learngene layers. Then we can expand these layers at their corresponding stage to initialize variable-sized Des-Nets. As shown in Fig. 1(d), we observe that the validation loss of models initialized by SWS can be significantly reduced with the increasing layer numbers.

We systematically investigate the design of weight sharing and the initialization strategy. With extensive experiments, we show the superiority of SWS: 1) Compared to training from scratch, SWS achieves better performance with much less training efforts on ImageNet-1K. Take Des-B as an example, SWS performs better while reducing around 6.6× total training costs. 2) When transferring to downstream classification datasets, SWS surpasses existing Learngene methods by a large margin, e.g., +5.6% on Cars-196. 3) When directly evaluating on ImageNet-1K without any tuning after initialization, SWS outperforms existing initialization methods by a large margin, e.g., +9.4% with Des-B (86M). 4) When building variable-sized models, SWS achieves better results while reducing around 20× parameters stored to initialize and around 10× pre-training costs, in contrast to the pre-training and fine-tuning. Our main contributions are summarized as:

• We propose a simple but effective Learngene approach termed SWS for efficient model initialization, which is the first work to systematically explore the potential of weight sharing for initializing variable-sized models.
• We present the design of weight sharing and the initialization strategy, which firstly highlights the importance of stage information and expansion guidance.
• Extensive experiments demonstrate the effectiveness and efficiency of SWS, e.g., compared to training from scratch, training with compact learngenes can achieve better performance while reducing huge training costs.

2 RelatedWork

2.1 Weight Initialization

Weight initialization is a pivotal step prior to training one model and crucially affects the model performance [Glorot and Bengio, 2010; He et al., 2015; Mishkin and Matas, 2015; Arpit et al., 2019; Huang et al., 2020]. Proper initialization aids in model convergence and training efficiency [LeCun et al., 2002], while arbitrary initialization may hinder the training process [Mishkin and Matas, 2015]. Comprehensive initialization strategies have been proposed, such as default initialization from Timm library [Paszke et al., 2019], Xavier initialization [Glorot and Bengio, 2010] and Kaiming initialization [He et al., 2015]. Nowadays, a plethora of pretrained models are readily accessible, offering an excellent initialization for finetuning models across a range of downstream tasks [Bao et al., 2022; Oquab et al., 2023]. However, this approach needs to reuse the entire model for each distinct downstream task, irrespective of the available resources. Furthermore, we need to pre-train again in instances where a pretrained model of the required size is unavailable, which is extremely time-consuming and computationally expensive. Recently, [Xu et al., 2023; Samragh et al., 2023] propose to initialize small models with a larger pretrained model. By contrast, we seek to train compact learngenes once via stage-wise weight sharing and then we can initialize variable-sized models.

Phase 1: learngene learning process

Figure 2: In the first phase, we build an auxiliary model comprising multiple stages. The layer weights in each stage are shared. Note that the number of layers in each stage and the number of stages are both configurable. Then we train it via distillation. After the learngene learning process, learngene layers containing stage information and expansion guidance are adopted to initialize descendant models of variable depths in the second phase. Finally, these models are fine-tuned normally and deployed to practical scenarios with diverse resource constraints.

2.2 Learngene

Learngene proposes to firstly learn a compact part, referred to as learngene, from a large well-trained model termed as ancestry model (Ans-Net) [Wang et al., 2023; Feng et al., 2024]. Subsequently, learngene is expanded to initialize variablesized descendant models (Des-Net), after which they undergo normal fine-tuning. Vanilla-LG [Wang et al., 2022a] extracts a few high-level layers as learngene and combines them with randomly initialized layers to build Des-Nets. TLEG [Xia et al., 2024] linearly expands learngene which consists of two layers to initialize Des-Nets of varying scales. LearngenePool [Shi et al., 2024] distills one pretrained model into multiple small ones whose layers are used as learngene instances, after which they are stitched to build Des-Nets. In contrast, we firstly investigate integration of stage information into learned learngenes and explore obtaining useful guidance about how to expand learngenes from the learngene learning process, thus better initializing Des-Nets.

2.3 Weight Sharing

Weight sharing is a parameter-efficient model compression strategy [Dabre and Fujita, 2019; Lan et al., 2020; Takase and Kiyono, 2021; Zhang et al., 2022; DING et al., 2023], which effectively alleviates over-parameterization problem [Bai et al., 2019; Kovaleva et al., 2019] in large pretrained Transformers [Devlin et al., 2018]. Different from existing works, we seek to initialize variable-sized models using learngene learned via stage-wise weight sharing, which to our knowledge remains unexplored in the literature.

3 Approach

Fig. 2 depicts the overall pipeline of SWS. In phase 1, we design an auxiliary Transformer model (Aux-Net) comprising several distinct stages, where the layer weights in each stage are shared. We train the Aux-Net through distilling from the ancestry model (Ans-Net) to help learn learngene layers and note that the weights of only one layer is trained in each stage due to the sharing mechanism. In phase 2, the welltrained learngene layers containing stage information as well as expansion guidance are adopted to initialize descendant models (Des-Net) of variable depths. Finally, these Des-Nets are fine-tuned normally without the restriction of stage-wise weight sharing. Next, we firstly introduce some preliminaries.

3.1 Preliminaries

Thanks to the modular design of modern vision transformer (ViT) [Dosovitskiy et al., 2021; Touvron et al., 2021], a typical ViT of L layers equipped with parameters θ can be defined as a composition of functions: fθ = fL ◦ · · · ◦ f1, where fθ : X → Y transforms the inputs in an input space X to the output space Y, fi means the function of the i-th layer and ◦ indicates the composition. Each layer contains Multi-head Self-Attention (MSA) and Multi-Layer Perceptron (MLP) block, where Layer Normalization (LN) [Ba et al., 2016] and residual connections are used before and after each block. The basic idea of weight sharing involves sharing parameters across layers, which is a simple but effective strategy to improve parameter efficiency. For multi-stage weight sharing, weight sharing in the m-th stage can be defined as a recursive update of one shared layer:

where Hm denotes the function of m-th stage, Lm denotes the number of weight-shared layers in the m-th stage, Zi denotes the representations in the i-th layer of m-th stage and θm represents the parameters of Hm, i.e., shared weights of the Lm layers. Note that the parameters of layers in each stage are shared, i.e., the number of updated parameters in each stage during training equals to that of one layer, but the parameters between stages are not shared. Thus ViT with parameters θ can be further defined as fθ = HM ◦ · · · ◦ H1 in multi-stage weight sharing setting, where M denotes the number of stages.

3.2 Stage-wiseWeight Sharing of Learngene

As discussed before, we could seamlessly integrate stage information into each learned learngene layer via sharing weights across multiple layers within each stage throughout the training. Drawing from this insight, we propose to share the parameters of one learngene layer to form multiple weightshared layers in each stage. Thus, the number of learngene layers equals to that of stages. By configuring the number of stages M, we can obtain different number of learngenes θlg = {θ1, ..., θM}. To ensure clarity, taking the 16-layer ViT-B (114M) [Dosovitskiy et al., 2021] and M = 5 as an example, θlg comprises about 36M parameters which is approximately equivalent to the parameter numbers of five layers and only 36M parameters are updated during the learngene learning process. In the following, we detail the layer assignment strategy and learngene learning process.

Layer Assignment Strategy. Given the number of stages M, there are two options to assign the number of layers within each stage: Balanced and Unbalanced sharing mechanism. As shown in Fig. 3(a), Balanced sharing ensures uniformity in the layer numbers across most stages, whereas Unbalanced sharing sets uneven layer numbers within most stages. However, as Balanced sharing is more aligned with the existing weight sharing principle [Lan et al., 2020; Takase and Kiyono, 2021], we will show in Section 4.3 that it achieves more stable and better learngenes than Unbalanced sharing. In this case, we take Balanced sharing as the default sharing mechanism in SWS.

Learngene Learning Process. As the learngene layer is the Transformer layer, but some other components like the patch projection and task-specific head are also required to compose a complete Transformer model. Therefore, we also add them to build the Aux-Net, after which we train it through distillation. Specifically, we consider adopting predictionbased distillation [Hinton et al., 2015] to condense knowledge from the Ans-Net, which is achieved by minimizing crossentropy loss between the probability distributions over the output predictions of the Ans-Net and those of the Aux-Net. Overall, one distillation loss is defined as:

where CE(·, ·) denotes soft cross-entropy loss, pt denotes the output logits of the pretrained Ans-Net (e.g., Levit-384 [Graham et al., 2021]), ps denotes the output logits of the Aux-Net, τ denotes the temperature value of distillation and ϕ denotes the softmax function. Moreover, we can seamlessly incorporate advanced distillation techniques [Zhang et al., 2022; Ren et al., 2023; Ji et al., 2023; Li et al., 2024; Li et al., 2025] into our training process. Besides distillation, we also introduce one classification loss:

where y denotes ground-truth label. Therefore, our total training loss is defined as:

where α denotes the trade-off. Noteworthy, the weight sharing constraint always exists during training, i.e., although AuxNet contains L layers, only θlg = {θ1, ..., θM} which contains parameters of M layers are updated.

3.3 Initialization with Learngene

After obtaining learngene which is composed of well-trained {θ1, ..., θM}, we can flexibly initialize Des-Nets of varying depths Lds , fitting diverse resource constraints with much less fine-tuning efforts. In the following, we describe the initialization strategy and the initialization order.

Initialization Strategy. During the learngene learning process, sharing weights in a stage-wise fashion provides important expansion guidance for initializing Des-Nets, i.e., expanding the weight-shared layers at its corresponding stage to initialize the target Des-Net. Specifically, we present two initialization strategies: Cyclic Initialization and Randomassigned Initialization. As shown in Fig. 3(b), we illustrate the initialization process of a 5-layer Des-Net with 3 learngene layers as an example. Cyclic Initialization involves the sequential selection of learngene layers in a specific order to initialize corresponding layers of the Des-Net. In contrast, Random-assigned strategy initializes the Des-Net by randomly selecting learngene layers. However, as Cyclic Initialization is more aligned with the learngene learning process (i.e., keeping the stage order the same as the Aux-Net), we will show in Section 4.3 that it achieves more stable and better performance than Random-assigned initialization. In this case, we take Cyclic Initialization as the default initialization strategy.

Initialization Order. Besides, different initialization orders lead to differences in initialization focus. For example, in the case of "Front-Last-Mid", such initialization focuses more on the front and last part of the Des-Net, as shown in Fig. 3(c). With Cyclic Initialization, different initialization orders bring similar performance of Des-Net. Therefore, we take "FrontMid-Last" as the default initialization order in SWS.

4 Experiments

4.1 Experimental Setup

We perform experiments on ImageNet-1K [Deng et al., 2009] and several downstream datasets including CIFAR-10, CIFAR100 [Krizhevsky et al., 2009], Food-101 [Bossard et al., 2014] and Cars-196 [Krause et al., 2013]. Model performance is measured by Top-1 classification accuracy (Top-1(%)). And we report Params(M) and FLOPs(G) as the number of model parameters and indicators of theoretical complexity of model.

In the first phase, we set two variants of Aux-Net as Aux-S/B where we adopt SWS on DeiT-S/B [Touvron et al., 2021], after which we train Aux-S/B on ImageNet-1K for 150/100 epochs to obtain learngenes, respectively. Specifically, we set 5 stages for Aux-S/B, where the shared layer number in each stage is 3,3,4,3 and 3. We choose Levit-384 [Graham et al., 2021] as the ancestry model. In the second phase, we set two variants of Des-Net as Des-S/B where we change the layer numbers based on DeiT-S/B, e.g., we name the 6-layer Des-S as DesS-6. Then we initialize Des-S/B with cyclic initialization and fine-tune them for 10 epochs, except that 15 epochs for DesS-13, Des-S-14 and Des-S-15 for better performance. Source code is available at https://github.com/AlphaXia/SWS.

4.2 Main Results

SWS achieves better performance while reducing huge training efforts in contrast to from scratch training on ImageNet-1K. We report ImageNet-1K classification performance of 20 different Des-Nets in Fig. 4, where "Scratch" denotes training from scratch, "TLEG" denotes linearly expanding learngenes to initialize [Xia et al., 2024]. Compared to Scratch, SWS can achieve better performance and significantly improve training efficiency. Take 10 Des-Bs as an example, SWS performs better while reducing around 6.6× total training costs (10×100 epochs vs. 100+10×5 epochs), compared to training each Des-B from scratch for 100 epochs. For each Des-B, SWS can reduce around 20× training costs. In some cases such as from Des-B-10 to Des-B-15, SWS performs better only after 1 epoch tuning, which demonstrates the effectiveness of initialization via SWS. Compared to TLEG, SWS performs better and further enhances the efficiency. Take 10 Des-Bs as an example, SWS performs better while reducing around 2.5× total training costs (100+10×40 epochs vs. 100+10×10 epochs). In a nutshell, the efficiency of SWS becomes more obvious with the number of Des-Nets increasing as we only need to train learngenes once.

When transferring to downstream classification datasets, SWS presents competitive results. We compare SWS against pre-training and fine-tuning (PF), Scratch, Vanilla- LG [Wang et al., 2022a] and TLEG [Xia et al., 2024] on 4 classification datasets. As shown in Fig. 5, we observe that SWS consistently outperforms several baselines, which verifies the effectiveness of initializing with learngenes trained via SWS. Take Des-B-12 as an example, SWS consistently outperforms PF by 1.12%, 0.35%, 1.65% and 2.28% respectively on CIFAR-100, CIFAR-10, Food-101 and Cars-196.

SWS significantly outperforms initialization baselines when directly evaluating on ImageNet-1K without any tuning after initialization. To validate the initialization quality of SWS, we compare SWS against Trained, Vanilla-LG [Wang et al., 2022a], TLEG [Xia et al., 2024] and IMwLM [Xu et al., 2023] on ImageNet-1K, where Trained means models trained from scratch for 100 epochs and IMwLM means initializing small models from a larger model. From Table 1, we observe that SWS significantly outperforms all baselines by a large margin. For example, SWS outperforms IMwLM by 52.3%, 35.8%, 19.7% and 9.4% respectively on Des-B-9, Des-B10, Des-B-11 and Des-B-12. Notably, we also find that directly initializing via SWS without any tuning can achieve comparable performance with well-trained models. For example, SWS outperforms Trained by 1.0%, 1.0% and 1.1% respectively on Des-B-13, Des-B-14 and Des-B-15, which verifies that the stage-wise information has been preserved to learngenes.

When initializing variable-sized models, SWS notably reduces the parameters stored to initialize compared with pre-training and fine-tuning (PF). We report the results of Des-Bs initialized from learngenes and those initialized from pretrained parameters whose number equals to that of target model. From Table 2, SWS achieves better performance and significantly reduces 20× (759.3M vs. 37.0M) parameters stored to initialize, in constrast to PF. Furthermore, SWS only needs to train learngene once while PF requires pretraining each Des-B individually, thereby significantly reducing the pre-training costs. Specifically, SWS reduces pre-training costs by 10× (10×100 epochs vs. 1×100 epochs) compared to PF. It is noteworthy that the efficiency of SWS becomes more obvious with the increasing number of Des-Nets.

4.3 Ablation and Analysis

The effect of stage-wise weight sharing. As shown in Table 3, we observe that SWS significantly outperforms SimpleLG when directly evaluating on ImageNet-1K without any tuning after initialization, which demonstrates the importance of stage information and expansion guidance for initializing Des-Nets. For example, SWS outperforms Simple-LG by 18.1% and 21.7% respectively on Des-B-10 and Des-B-11.

The effect of different stage-wise weight sharing strategies. We present the performance of Aux-B to show the quality of trained learngenes. From Table 5, we observe that MS-WS6 outperforms MS-WS4 and MS-WS5 fails to converge, which reflects balanced sharing is more stable and better than unbalanced one. Moreover, we find that MS-WS4 outperforms MS-WS3, which demonstrates starting from a weight-sharing stage is better. Also, MS-WS5 achieves comparable performance in contrast to Mini-DeiT [Zhang et al., 2022].

The effect of different initialization strategies. From Table 4, we find that performance of Des-Nets with Cyclic Initialization (CI) outperforms that with Random-assigned Initialization (RI) when directly evaluating on ImageNet-1K without any tuning after initialization. For example, CI-order1 outperforms RI by 5.3% and 7.3% respectively on Des-B-9 and Des-B-10.

The effect of learngenes on Des-Nets. We train Aux-S for more epochs. As shown in Table 6, we find that better performance of Des-S-12 can be consistently achieved with better learngenes. For example, performance on CIFAR-100 can be improved from 87.40% to 88.54% with better learngenes.

5 Conclusion

In this paper, we proposed a well-motivated and highly effective Learngene approach termed SWS to initialize variablesized Transformers, enabling adaptation to diverse resource constraints. Experimental results under various initialization settings demonstrated the effectiveness and efficiency of SWS.

Acknowledgements

This research is supported by the National Science Foundation of China (62125602, 62076063, 62206048), Natural Science Foundation of Jiangsu Province (BK20220819), Young Elite Scientists Sponsorship Program of Jiangsu Association for Science and Technology Tj-2022-027 and the Big Data Computing Center of Southeast University.