Learngene Tells You How to Customize: Task-Aware Parameter Initialization at Flexible Scales
Authors: Jiaze Xu
Shiyu Xia
Xu Yang
Jiaqi Lv
Xin Geng
Venue: ICML 2025 poster (2025)
Keywords: model initialization
Learngene
hypernetwork
Primary area: deep_learning->algorithms
TL;DR: Task-Aware Parameter Initialization at Flexible Scales

Abstract: Appropriate parameter initialization strategies are essential for reducing the high computational costs of training large pretrained models in various task scenarios. Graph HyperNetwork (GHN), a parameter initialization method, has recently demonstrated strong performance in initializing models.
However, GHN still faces several challenges, including limited effectiveness in initializing larger models, poor performance on smaller datasets, and the requirement of task-specific GHN training, where each new task necessitates retraining the GHN model, leading to increased computational and storage overhead.
To overcome these challenges, motivated by the recently proposed Learngene framework, we propose a novel method called **T**ask-**A**ware **L**earngene (**TAL**). Briefly, our approach pretrains a TAL model under the guidance of a well-trained model and then performs multi-task tuning to obtain a shared TAL model that enables parameter prediction based on both model architectures and task-specific characteristics.
Extensive experiments show the superiority of TAL.
Models initialized with TAL outperform those initialized using GHN method by an average of 24.39\% in terms of accuracy across Decathlon datasets.

Learngene Tells You How to Customize: Task-Aware Parameter Initialization at Flexible Scales
Jiaze Xu 12 Shiyu Xia 12 Xu Yang 12 Jiaqi Ly Xin Geng 12
Abstract
2023), and other fields. However, pretraining such models requires computational resources and training costs (Radford et al., 2021; Touvron et al., 2023; Yang et al., 2024; Peng et al., 2025), thus making it challenging and costly to obtain well-trained models in many task scenarios with resource con- straints (Mchta Rastegari, 2021; Fu et al., 2025). So that an appropriate model initialization strategy becomes erucial, as effective parameter initialization not only accelerates model convergence but also enhances model performance (He et al., 2015; Zhang et al., 2018; Wang et al., 2022; 2024; Yao et al., 2025), thereby significantly reducing overall training costs.
Appropriate parameter initialization strategies are essential for reducing the high computational costs of training large pretrained models in various task scenarios. Graph HyperNetwork (GHN), a parameter initialization method, has recently demonstrated strong performance in initializing models. However, GHN still faces several challenges, including limited effectiveness in initializing larger models, poor performance on smaller datasets, and the requirement of task-specific GHN training, where each new task necessitates retraining the GHN model, leading to increased computational and storage overhead. To overcome these challenges, motivated by the recently proposed Learngene framework, we propose a novel method called Task-Aware Learngene (TAL). Briefly, our approach pretrains a model under the guidance of a well-trained model and then performs multi-task tuning to obtain a shared TAL model that enables parameter prediction based on both model architectures and task-specific characteristics. Extensive experi- ments show the superiority TAL. Models initial- ized with TAL outperform those initialized using GHN method by an average of 24.39% in terms of accuracy across Decathlon datasets. We provide the code at https://gtthub.commmattiiuu///// Aware-Learngene.
Recently, Graph HyperNetwork (GHN) (Zhang et al., 2018; Knyazev et al., 2021; 1; 2023) has been proposed as an approach to make model pretraining more accessible by reducing com- putational costs and enabling parameter initialization for models of multiple scales. Formally, given set of neural network architectures f as training data, where each archi- tecture is represented as computational graph fo (Knyazev et al., 2021), a GHN (denoted as Hp and parameterized by 0) learns to predict the parameters of these neural net- works. The prediction process can be formulated as Wpedd = Hf O). where W represents the predicted network parameters. During training, the GHN is optimized to min- imize the loss function on target dataset 'D like ImageNet- IK(ussaakvvkky et al., 2015). The predicted parameters W pred demonstrate superior initialization performance com- pared to conventional random initialization methods, leading to reduced training time and computational costs.
Despite the advantages demonstrated by GHN, as illustrated in Fig. several significant limitations persist in their practical applications. First, although the latest GHN method LoGAH (Zhou et al., 2024) has demonstrated significant progress in handling deep neural architectures, its effective- ness diminishes when initializing larger models like ViT- base (Dosovitskiy, 2020). And most GHN methods still have considerable room for improvement in initialization accuracy across various datasets, such as ImageNet- K K(Russakovsky et al., 2015) and CIFAR-100 100 (Krizhevsky et al., 2009). Second, GHN requires independent training for each specific task, meaning that when facing different downstream applica- tions, a new GHN model has to be retrained accordingly. This requirement not only substantially increases computational overhead but also imposes additional storage burdens.
1. Introduction
Pretrained models have achieved remarkable success in com- puter vision (Dosovitskiy, 2020; Radford et al., 2021), natural language processing (Radford et al., 2019; Touvron et al.,
'School of Computer Science and Engineering, Southeast Uni- versity, Nanjing 210096, China "Key Laboratory of New Gener- ation Artificial Intelligence Technology and Its Interdisciplinary Applications (Southeast University), Ministry Education, China. Correspondence to: Jiaqi Lv <is.jiaqi. Iv eeeg@gmall.oo comi..com>.
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada, PMLR 267, 2025. Copyright 2025 by the author(s).

---

Learngene Tells You How to Customize: Task-Aware Parameter Initialization at Flexible Scales
(a) GHN
(b) Learngene
Task-Aware Learngene
High storage consumption High computationnl overhead
Ancestry model
Task
guide
GHI
GHN 2
GHN
TAL. model
Task 2
learngeme
decoder
condense
expand
initialice
leamgene
nry
nryy
tinyy
Tayy
small
anail
0.0.0
smalll
Sinal
Small
Ancesiry
Task
beaz
Task
base
Task
base
model
Task
Task 2.
hoc
Task lc
Figure 1: (a) GHN exhibits limitations in initializing large-scale models and requires high storage and computational resources for each individual task. (b) The Learngene framework condenses critical components (learngene) from large ancestry model and expands them initialize models of various sizes. (c) The Task-Aware Learngene (TAL. pipeline: first pretraining the TAL model using ancestry model, then tuning it in multi-task setting to generate task-specific initialization for models of different scales.
One recently proposed framework called Learngene (Wang al., 2022), offers inspiring insights into addressing the limitations of GHN mentioned above. As shown in Fig.1 Learngene adopts unique paradigm where it first condenses a well-trained model, termed ancestry model, into a small critical part known as learngene. Subsequently, this learngene can be expanded to initialize multiple models of varying sizes for different downstream tasks (Wang et al., 2022; Xia al., 2024a; Shi al., 2024; Feng al., 2024b;a; Xie et al., 2024). While both Learngene and GHN aim to provide effective model initialization, Learngene distinguishes itself through its ability to inherit and utilize knowledge from the ancestry model and consider the commonalities across downstream tasks, enabling better adaptation to different application scenarios.
of accuracy across Decathlon datasets (Rebuffi et al., 2017).
The main contributions of this work are as follows:
(1) We design Task-Aware Learngene (TAL), an end-to-end mechanism that effectively represents and transfers shareable knowledge across tasks in parameter prediction. TAL not only provides well-initialized parameters for larger models but also enables parameter prediction based on both model architectures and lask-specific characteristies.
(2) Extensive experiments demonstrate the superiority of TAL across various scenarios. Compared to training from scratch and previous GHN methods, models initialized with TAL achieve superior performance while substantially reducing both computational costs and storage requirements.
In this paper, connect these two lines of work, proposing novel approach called Task-Aware Learngene (TAL) that encodes shareable information to predict initial parameters for models across flexible scales. As shown in Fig. (c)), our approach involves two stages. First, we train the TAL model on a large dataset under the guidance an ancestry model to transfer knowledge. Then, we tune the trained TAL model on multi-task datasets to effectively filter and convey the task-specific knowledge from the ancestry model. Finally, the trained TAL model can predicting model parameters for various tasks, even unseen ones and supports model cus- tomization at different scales. We systematically investigate the effectiveness of TAL. Extensive experiments show the superiority of TAL. For example, ViT-small model initialized by TAL achieving 22.44% higher accuracy on ImageNet- IK without any training compared to LoGAH (Zhou et al., 2024). Moreover models initialized with TAL also outperform those initialized using LoGAH by an average of 24.39% in terms
2. Related Work
2.1. Learngene
The Learngene method (Wang et al., 2022), inspired by biological gene inheritance, focuses on extracting compact components, known as learngene, from well pretrained mod- cls (ancestry models) to initialize models. Existing methods like Vanilla-LG (Wang et al., 2022), TLEG (Xia et al., 2024a), Learngene Pool (Shi et al., 2024) and SWS (Xia al., 2024b) employ different strategies to select and expand learngene. Vanilla-LG extracts key layers as learngene and splices them with randomly initialized layers. TLEG uses linear expansion of learngene layers, while Learngene Pool refines large mod- els into multiple small models, using their layers as learngene instances to construct new models. In Task-Aware Learngene (TAL), learngene becomes the encoder part of the TAL model rather than a sub block of the ancestry model. Model generation becomes an incoding-decoding process, with
2

---

Learngene Tells You How to Customize: Task-Aware Parameter Initialization at Flexible Scales
ancestry model and multi-task knowledge injected through learngene. can process both model and task information, initializing flexible-scale models for different tasks.
h 1) = Em ()), where h,") is one-hot vector denoting an operation.
4. Task-Aware Learngene
2.2. Graph HyperNetwork
Fig.2 (a-c) illustrates the overall pipeline of Task-Aware Learngene (TAL). First, we train the TAL model on a large dataset under the guidance of an ancestry model to transfer knowledge. Then, we tune the trained TAL model on multi- task datasets effectively filter and convey the task-specific knowledge from the ancestry model to downstream models cross different tasks. Finally, the trained TAL model can predicting model parameters for various tasks, even unseen ones and supports model customization at different scales.
Graph HyperNetwork (GHN) (Zhang et al., 2018; Knyazev et al., 2021) employs hypernetwork for direct param- eter prediction. This approach has attracted significant research interest due to its superior performance and remarkable adaptability. GHN-2 (Knyazev et al., 2021) and GHN-3 (Knyazev et al., 2023) further improved the parameter prediction capabiliitees of GHN by improving the learning process process of the model computation graph. The latest LoGAH method (Zhou al., 2024) introduces low-rank approximation (LoRA) technology, allowing GHN to predict the parameters of larger models using smaller hypernetworks. This progress has greatly improved the efficiency and ability of GHN in handling large-scale model parameter prediction tasks. Our task-aware learngene (TAL) incorporates modules to process task information, enabling single TAL model to customize models of varying scales for different tasks.
TAL model structure and components. In the TAL, we adopt encoder decoder structure for model parameters predic- tion (Knyazev al., 2023; Zhou et al., 2024). We refer to the encoder part of the TAL model as learngene because it first in- herits knowledge from the ancestry model and then transfers task-specific knowledge based on different tasks. Specifi- cally, learngene receives both model configuration through model computational graph and task information from ances- try model. Based on task information, learngene can filter out task-specific knowledge which previously inherited from the ancestry model and inject it into model computational graphs, thereby producing task-specific computational graphs.
3. Background
A Graph HyperNetwork (GHN) is a neural network Hp parameterized by / and trained on a dataset DD The input of GHN Hp(() is computational graph fo of a neural network f and the output of GHN is the parameters of the model weed = (
The architecture of learngene is shown in Fig.2(d). Inspired by (Perez et al., 2018; Oreshkin et al., 2018), we introduce a task hypernet that processes task information to dynamically generate parameters for the task-specific layer, which is implemented layer acts on the model computational graph, transferring task information to it.
In (Knyazev et al., 2021), GHN Hp is trained by SGD over M training architectures ( ) and N training data samples {I,, } on the following optimization problem:
simple MLP. Then task-specific
1 NN M
In this process, task information is passed in the form of a task embedding {I, JT= for each task, which is generated by the ancestry model through the average feature extraction of the task images (Vu et al., 2020).
argmin NM Cl HDl.)) la-=
(1)
when training GHN (f), a meta-batch of m training architectures is sampled input for GHN. Meanwhile, mini-batch of training datas x is sampled and fed into the snrameter-predicter m architectures to get X predictions g. The cross-entropy loss L is computed between á and ground truth labels y of x for classification tasks. Afterward, the loss is back- pack-propagated to update the parameters of HD by gradient descent.
The task hypernet h generates task bias parameters "r, and B, of the task-specific layer.
(2)
where W lxtt and W ERA
The task-specific layers apply these bias parameters to the model's 'ss computation graph using the following formula:
The input computational graph f (V,E)) is a directed acyelic graph (DAG), where the nodes V correspond op- erations (convolution, pooling, self-attention, etc.) (Knyazev et al., 2021), while the edges E correspond to the forward pass flow of inputs through the network f. GHN takes d-dimensional node features (()) E RlVixd as input obtained using an embedding layer for each i-th node:
-
The task-specific model computational graph generated by learngene is then passed to the decoder (Zhou et al., 2024), which decodes the graph to generate the model parameters.
3

---

Learngene Tells You How to Customize: Task-Aware Parameter Initialization at Flexible Scales
Train on large dataset
(c) Customize your model
(d) TAL's encoder (learngene)
Large Datasset
Sectario
Sccuuioo
Senazio
TAL model
learngene
decoder
Seaineeity
Tur
Zore
Foedforward
Dansse
Stft label
Ancestry model
TAL model
Layer Norm
Ancestry model
Multi-task tuning
learggene
Deseder
ombeddiing
Multi-Head
Yark entroddig
Attention
TAL. model
Layer Norn
leamgene
decodes
MLP
Task HyperNet
Tisk
Docouet
Compontionnl groph
Figure 2: (a) Training TAL model on a large dataset under the guidance a large-scale foundation model (ancestry model). (b) Tuning TAL model to multiple tasks. (c) Customizing task-specific models with flexible scale based on new task scenarios. (d) The learngene based on a transformer architecture and consists of a stack of transformer blocks.
Train TAL model large dataset. In order to inherit the knowledge of the ancestry model, we first train the TAL model on a large dataset under the guidance of the ancestry model. We adopt the method from (Shu et al., 2021), converting features extracted by the ancestry model from the images into a probability distribution map. For the training models of the TAL, also apply the same method obtain their feature probability distributions and compute the KL divergence between those of the ancestry model. This function is denoted as:
ancestry model, enabling models initialized by TAL model to handle complex tasks.
Tuning TAL model under multi-task setting, We then tune the TAL model on multiple tasks, leveraging task information to enable learngene to generate task-specific computation graphs, thereby decoding the model parameters tai loilored each task. We formulate the loss function for this part of the TAL model's training. Given the data from set of tasks (D,)T=,, here T is the total number of tasks and D, = [ . N- shows the training data for r-th task with N, samples.
aax (FnottaaaaxxFnnnnnn M) ||lofmaaaaa
(4)
Assuming there a TAL model (0) parameterized by 0 that computes the output the parameters of the model Wpred (D( for input computational graph fo of neural network. In multi-task setting, TAL model is trained by SGD over M training models { M and T training tasks (Dr JT=1 on the following optimization problem:
where Euain and Eane refer to the encoders' output of the training model and ancestry model respectively. The matrix M E gdxd'' which transforms the output dimension d of Emain to match the output dimension d' of Enne. like the param- eters of other parts of the training model, the transformation matrix's ' parameters are predicted directly by TAL. model.
Considering loss between training model's predicted label and ground truth label:
argmin TM
1
ΣΣ Σ
W. (2)0)))))))) 0))
r= 1a (E).
Lef= -C (ra(eeerrnnnn
(5)
(7)
where L is typically the cross-entropy loss and wy shows the sampling weight for 7-th task.
where (uain()) represents the training model's 's predicted label of input image data x: and ye denotes the ground truth label belonging to category Then, while one model is used as training data for TAL model, the total training loss is computed as follows:
Multi-task training allows the models predicted by TAL to inherit task-specific knowledge filtered by learngene from the ancestry model, well as the shared knowledge across tasks.
Customize models across different tasks. After training on a large dataset and tuning on multiple tasks, the TAL model can provide mask-specific, variable-sized models for both seen and unseen tasks. By simply passing the required model configuration and task information to the TAL model,
L L= (1-a:) Lets:
Through this process, the TAL model can inherit and utilize the vast amount of domain knowledge already learned in the
4

---

Learngene Tells You How to Customize: Task-Aware Parameter Initialization at Flexible Scales
Table I: Performance of models on ImageNet-IK initialized with GHN-3, LoGAH, TAL and TAL", after 75 epochs of training for all initialization methods.
datasets: the ViTs-1K dataset (Zhou et al., 2024) for TAL, GHN-3, and LoGAH methods, and the ViTs*-1K dataset for our enhanced TAL* method. Both datasets contain 1000 different ViT-style computational graphs as the architecture source. The vision tasks are evaluated on the Visual Domain Decathlon Challenge (Rebuffi et al., 2017), which comprises 10 diverse datasets: (1) ImageNet- 1K (IN-1K) (Russakosskyy et al., 2009),
MODEL
TRAINING STATE
GHN-3
LOGAH
TAL
TAL*
12-TINY
UNTRAINED
34,45
22.79
26.20
31.31
TRAINED
48.93
62.53
63.03
60.04
12-SMALL
UNTRAINED
31.03
16.78
23.28
39.22
et al., 2015). (2) CIFAR- (C1000
TRAINED
53.47
65.41
66.61
65.63
(3) Aircraft (Airc.) (Maji et., 2013), (4) Daimler pedestrian classification (DPed) & Gavrila, 2006), (5) De- scribable textures (DTD) et al., 2014), (6) German traffic signs (GSTR)(Stllkanm) et al., 20 (7) Omniglot (OGlt)(Lke et al., 2015), (8) SVHN(Netzer et al., 201 1) (9) UCF101 Dynamic Images (UCF)(Soomro et al., 2012), and (10) Flowers102 (Fiwr)(Nilsback & Zisserman, 2008). detailed dataset descriptions, please refer to Appendix .I.
12 BASE
UNTRAINED
0.10
0.10
0.10
38.72
one can instantly obtain well-initialized model parameters tailored to the task at hand.
Model architecture datasets. In GHN-1/22 (Knyazev et al., 2021) and GHN-3 (Knyazev et al., 2023), training architectures are sampled from DeepNets-IM, a dataset of I million architectures (Knyazev et al., 2021). In LoGAH works, they constructed task-specific architecture datasets: ViTs-IK for vision tasks and GPTs-1K for language tasks (Zhou et al., 2024), each containing 1K different computational graphs. For our TAL method, we adopt these datasets from LoGAH, and additionally design an enhanced vision model library denoted as ViT dataset for vision tasks. Unlike the original ViTss 1K which has 10M parameter constraint, our ViT incorporates wider and deeper model architectures. We also increase the proportion of larger models in the dataset. This design strategy leads to significant improvements in TAL's performance. The model trained with this enhanced library is denoted as TAL.*.
For language experiments, all methods utilize GPTS- 1K (Zhou et al., 2024) as the model architecture dataset. The language tasks are evaluated on four widely-used NLP benchmarks: Microsoft Research Paraphrase Corpus (MRPC) (Dolan & Brockett, 2005), Corpus of Linguistic Acceptability (CoLA) (Warstadt, 2019), Recognizing Textual Entailment (RTE) (Wang. 2018), and Internet Movie Database reviews (LMDB) (Maas et al., 2011). 1).
Baselines, We compare TAL with GHN-3 (Knyazev et al., 2023) and the latest LoGAH method (Zhou et al., 2024), which improves the design of the decoder and significantly enhances the initialized models' performance. To ensure fair comparison, we reproduce all experiments using the official source code of these methods under identical experimental settings and environment.
5. Experiments
In this section, we evaluate our proposed TAL method by predicting parameters for models of various scales across different tasks, comparing it with previous methods includ- ing GHN-3, LoGAH, and random initialization. First, we assess TAL's capability in predicting ViT model parameters on both seen and unseen visual tasks during training. Then, we examine TAL's effectiveness in predicting parameters for GPT-2 models on natural language tasks. Furthermore, we conduct comprehensive ablation studies to investigate the impact of various factors on TAL's performance. Finally, we visualize the intermediate learngene outputs to demonstrate TAL's effectiveness in handling multiple tasks.
Sampling tasks. During multi-tas task training, we sample tasks using conventional temperature-based sampling (Raffel et al., 2020) with a temperature T 2 for all methods. Tasks are sampled proportionally to p. 1/T where p.- Σ N; and N- represents the number of training samples for the 7-th task. Note that this sampling probability p.!" directly corresponds to the sampling weight wy, introduced in Formula 7.
N=
Training Details. For both TAL and TAL*, we first pretrain the hypernets on ImageNet- 1K for 75 epochs, followed by 100 epochs of multi-task training on the Decathlon Challenge datasets. For TAL*, we leverage the logits from the ancestry model as soft labels to guide the training process. All mod- els are trained using automatic mixed precision in Py PyTorch, with a cosine annealing learning rate schedule starting at Ir Ir= 3e 34-4, weight decay x= 1e-2 and predicted parameter regularization 7 =3e-5 (Knyazev et al., 2023). We use a pre- trained ViT-Base (Dosovitskiy, 2020) as the ancestry model.
5.1. Experimental Setup
Datasets. Our experiments comprehensively evaluate our proposed methods (TAL and TAL*) against existing ap- proaches (GHN-3 and LoGAH) on both vision and language tasks. Each experiment requires two types of datasets: model architecture datasets for parameter prediction and task-specific datasets for downstream evaluation.
For vision experiments, we employ two model architecture
5

---

Learngenc Tells You How to Customize: Task-Aware Parameter Initialization at Flexible Scales
Table 2: Performance of untrained models on Decathlon tasks initialized with GHN-3, LoGAH, TAL and TAL*. Note that for Vase only TAL* initialization results are shown, as it the only method capable predicting parameters for base-scale models.
MODEL
METHOD
AIRC.
C100
DPED
DTD
GSTR
OGLE
SVHN
UCF
FLWR
GHN-3
3.12
34.06
85.41
6.38
87.77
0.06
10.00
2.25
7.06
LOGAH
2.58
29.16
78.83
8.30
92.12
0.26
17.32
3.94
6,96
3-TINY
TAL.
6.69
39.53
79.83
22.55
94.86
28.37
83.31
28.33
33.82
TAL.
2.04
50.80
88.88
29.26
98.87
0.11
87.37
35.76
50.00
GHN-3
3.24
35.19
87.72
6.86
89.12
0.06
10.00
3.43
10.88
LoGAH
3.21
46.33
81.19
9.20
96.82
0.31
20.82
5.02
8.33
6-TINY
TAL.
17.43
48.95
85.87
28.30
99.25
48.98
89.85
38.63
45.69
TAL*
1.98
51.01
88.06
28.62
98.89
0.18
87.27
35.71
49.71
GHN-3
3.15
31.12
85.63
6.86
85.02
0.06
10.00
3.07
9,90
LOGAH
3.15
33.80
79.86
9.41
96.52
0.18
20.56
5.53
9.12
12-TINY
TAL
16.71
44.72
84.76
28.03
99.15
28.11
88.87
39.81
45,00
TAL*
2.07
50.99
88.21
29.41
98.94
0.14
87.30
35.81
49.80
GHN-3
2.91
35.75
86.77
6.70
87.68
0.06
10.00
2.61
10.39
LoGAH
3.15
44.95
80.66
9.36
96.35
0.28
19.91
5.33
8.73
3-SMALL
TAL.
17.16
47.25
83.78
27.23
98.93
47.75
87.83
38.32
44.31
TAL.
1.95
56.05
94.20
33.14
99.57
0.03
91.24
47.23
55.78
GHN-3
3.12
35.30
86.80
7.34
90.05
0.06
10.00
2.36
12.75
LoGAH
3.18
45.93
80.60
9.95
96.97
0.26
21.04
5.38
8.63
6-SMALL
TAL.
17.85
49.88
83.91
28.67
99.30
50.29
89.82
40.83
46.47
TAL*
1.98
56.08
94.42
32.98
99.62
0.09
91.31
48.16
56.57
GHN-3
2.64
5.55
84.30
7.23
84.39
0.06
10.00
1.74
9.51
LoGAH
2.64
35.77
80.39
8.24
96.81
0.18
20.66
4.30
6.57
12-SMALL
TAL
17.16
45.63
82.16
27.87
99.18
39.66
88.55
40.98
45.10
TAL*
1.86
56.04
94.66
33.14
99.58
0.14
91.30
48.21
56.57
3-BASE
TAL*
2.22
55.99
94.10
32.93
99.41
0.09
90.99
4.39
56.08
6-BASE
TAL*
1.80
56.02
94.00
33.03
99.41
0.11
91.05
47.18
56.08
12-BASE
TAL'
2.31
56.06
94.39
32.93
99.41
0.11
90.97
47.59
55.69
5.2. Main results
architectures, only TAL* results presented since it is the only method capable of parameter prediction for base-scale models. As shown in Tab.2, untrained models initialized with TAL/ outperform the GHN-3 and LoGAH across all Decathlon tasks.
TAL/ TAL* achieves better performance ImageNet- 1K. We evaluate the performance of the TAL/TAL" the ImageNet- 1K. As shown in 1ab. the untrained model, struc- tured as ViT-Small and initialized using TAL* outperforms initialized with LoGAH by 22.44 % on ImageNet- I.. Fur- thermore, after 75 epochs training, the model initialized with PALaevieves 1.20% higher accuracy compared LoGAH initialization. Notably, among all initialization methods, only TAL* is capable of predicting parameters for ViT-Base scale models, achieving an initialization accuracy of 38.72% % These results show that TAL can effectively inherit and utilize the knowledge already learned in the ancestry model.
Models initialized with TAL/TAL* outperform those initialized by the LoGAH method during the training process. We select 12-layer ViT-tiny (12-Tiny) and 12-layer ViT-small (12-Small) as test models for further evaluation. Our TAL/TAL* initialization consistently outper- forms other initialization methods across all tasks. Notably, on DTD, UCF, and Flower datasets, models initialized with our method achieve approximately 15-25 higher accuracy after training compared other initialization approaches.
Models initialized with TAL/TAL* demonstrate strong performance without any training on Decathlon tasks. We compare TAL/TAL* with GHN-3 and LoGAH methods on Decathlon tasks using ViT models of varying archi- tectures and depths. The experiments are conducted with ViT-Tiny and ViT-Small at three different depths: 3, 6, and 12 layers, as well as ViT-Base models. Notably, for ViT-Base
The TAL method significantly reduces computational costs. We calculate the total training time for each method across the previous ten experimental datasets. As shown in Tab.4, under identical experimental conditions, the TAL method demonstrates notable efficiency advantages, reducing training time by 22% 22% compared the LoGAH
6

---

the classification of skin lesions. We select 3-layer ViT-tiny (3-Tiny) and a 6-layer ViT-small (6-Small) as test models for further evaluation. Tab.5 shows that models initialized Learngenc Tells You How to Customize: Task-Aware Parameter Initialization at Flexible Scales
Table 3: Performance of trained models on Decathlon tasks initialized with RandInit, GHN-3, LoGAH, TAL and TAL". : For models initialized with RandInit, accuracy is reported after 200 epochs of training for each task, while for models initialized with other methods, trained for 100 epochs.
MODEL
METHOD
AIRC.
C100
DPED
DTD
GSTR
OGLE
SVHN
UCF
FLWR
RANDINIT
7.80
58.64
98.74
23.09
99.58
24.12
86.38
24.03
35.10
GHN-3
4.80
55.88
98.20
20.11
97.53
9.77
83.08
24.80
32.45
12-T
LoGAH
5.61
58.93
97.47
20.32
99.23
29.22
86.13
35.81
35.78
TAL.
19.17
59.80
98.88
29.73
99.60
57.50
90.31
46.47
48.33
TAL.
7.23
58.90
99.46
33.09
99.78
24.01
90.94
49.28
58.04
RANDINIT
8.01
60.17
98.44
25.05
98.89
15.70
85.57
23.31
35.00
GHN-3
5.22
57.35
98.32
15,21
97.36
20.98
80.21
22.34
34.51
12-SMALL
LOGAH
7.20
59.98
97.45
20.21
98.95
24.55
85.10
33.86
34.51
TAL
18.69
60.49
98.89
31.01
99.77
57.01
92.02
47.75
49.71
TAL*
8.61
60.77
99.51
34.36
99.79
23.04
91.72
52.77
59.02
Table 4: Training time comparison of different methods, all experiments run on an NVIDIA RTX 4090 with time measured in hours (h).
Table 6: Performance of untrained GPT2 models on NLP tasks initialized with LoGAH and TAL.
MODEL
METHOD
MRPC
COLA
RTE
IMDB
METHOD
GHN-3
LOGAH
TAL
TAL*
3-GPT2
LoGAH
$5.88/ 75
0.70
47.65
$2.06
TAL
62.99/ 76.70
2.89
46.21
63.21
TIME(HOURS)
58.81
46.55
36.19
70.53
6-GPT2
LoGAH
59.07/72.12
1.23
46.21
53.58
61.52/73.70
47.29
62.76
Table Performance of models on unseen tasks initialized with RandInit, LoGAH and TAL, after 5 and 100 epochs of training for each task.
LoGAH
60.78/ 74.19
1.23
48.01
57.76
9-GPT2
TAL
$8.00 68.97
1.75
48.78
62.09
LOGAH
56.13/ 68.98
0.10
48.01
58,07
12-GPT2
TAL.
$2.70/ 60.53
3.05
49.82
61.49
LOGAH
57. 17.96/ 71.01
0.82
47.47
55.37
DATASET
MODEL
EPOCHS
RANDINIT
LoGAH
TAI.
AVGLACC
TAL
$8.83/ 69.98
2.75
48.02
62.39
5
82.71
87.56
88.86
3-TINY
100
89.41
91.08
90.99
E-NNIST
5
83.66
87.78
89.78
6-SMALL
100
88.98
90.68
91.56
5
29.20
42.57
47.06
T-NINY
100
59.91
60.60
61.41
FER2013
5
30.00
32.57
49.99
with TAL converge faster and achieve higher test accuracy on unseen downstream tasks.
6 M-LL
100
62.55
61.94
65.09
5
83.09
87.56
86.59
3-TINY
100
97.46
97.71
97.71
TAL shows promising results NLP tasks. We further evaluate TAL's effectiveness on NLP tasks by conducting experiments MRPC, COLA, RTE and IMDB datasets using GPT2-small models with 3, 6, 9, and 12 layers. For training setup, we individually train a separate LoGAH model for each task with 250 epochs, while in TAL method, we leverage pretrained GPT-2 model (Radford al., 2019) to provide task information and train single shared TAL model across all four tasks with only 100 epochs. Compared to LoGAH, TAL not only reduces the training time by more than 50% but also demonstrates competitive or superior performance across test tasks. Tab.6 shows TAL's notable improvement on COLA and IMDB tasks, with accuracy gains of 1.93% and 7.02 % respectively on average.
HAM10000
82.13
89.61
91.18
6-SMALL
100
97.70
97.83
97.95
method. Although TAL. requires longer training time (70.53 hours), this increased computational cost is well justified by its substantially expanded capability. TAL can predict parameters for models more than 4 times larger than previous methods, while also achieving significantly better initialization performance across most tasks.
TAL presents superior parameter prediction ability across unseen tasks. We evaluate TAL against LoGAH trained ImageNet- 1K and random initialization (RandInit) on a broader of unseen datasets. Specifically, we use three datasets from distinct fields: Fashion MNIST (Xiao et al., 2017), a dataset of fashion item images; FER2013 (Goodfel- low et al., 2013), a facial expression recognition dataset; and HAM10000m (Tschandl et al., 2018), medical dataset for
5.3. A Analysis and A Ablation
In the main experiments, high-quality model initialization is shown to significantly accelerate convergence and improve

---

Learngene Tells You How to Customize: Task-Aware Parameter Initialization at Flexible Scales
Table 7: Performance of untrained models initialized with single-task tuning (TAL-st) and multi-task tuning (TAL) on Decathlon tasks.
MODEL
METHOD
AIRC.
C100
DPED
DTD
GSTR
OGLE
SVHN
UCF
FLWR
TA
1.95
5.99
90.00
6.65
82.08
0.03
27.22
13.78
10.20
12- 12-TINY
TAL
19.8
54.89
98.61
30.21
99.57
63.57
90.38
48.1
47.84
TAL-ST
2.79
8.66
89.56
8.67
99.29
0.09
21.84
21.84
7.84
12-SMALL
TAL
17.16
45.63
82.16
27.87
99.18
39.66
88.55
40.98
45.10
Table 8: Performance of untrained models on Decathlon datasets initialized with TAL without using task informa- tion(T.I.) learngene or ancestry model(ans-net) on De- cathlon datasets.
Datasets
Air.
C100
Died
DTD
METHOD
ANT-NET
T.I.
AvG AcC
OGie
SVHN
TAL(W/O T.I.) TAL(W/O ANS-NET) TAL
40.71
UCF
X
/
46,80
/
53.37
final test accuracy. Therefore, in our analysis, we evaluate the performance of untrained models initialized with different methods on each dataset.
-10
30
Figure 3: The computational graphs of all models generated by learngene on the Decathlon datasets. Each point repre- sents model computational graph. Different colors denote different tasks and the size of the point corresponds to the model's scale, with larger points indicating larger models.
The effect multi-task training. We investigate the impact of multi-task training by comparing TAL with TAL-st, where TAL-st sequentially fine-tunes the TAL model (pretrained on ImageNet) on each downstream task individually. We evaluate using two models: 12-layer ViT-tiny (12-Tiny) and a 12-layer ViT-small (12-Small). As shown in Tab. 7, models initialized with TAL significantly outperform TAL-st across almost all Decathlon tasks.
tasks clearly cluster together in two-dimensional space. This indicates that learngene can effectively integrate task information while incorporating task information into the model computational graph.
The effect of ancestry model and task information. As shown in Tab.8, we conduct ablation studies by removing task information (T.I.) in learngene for TAL(W/o T.I.) and removing the guidance of ancestry model for TAL(W/o ans- net). Models initialized by TAL outperform TWL(W/o T.I.) and Tnd (w/o (es-nee) by 12.66% % and 6.57 % in test accuracy across Decathlon tasks, respectively. More comprehensive evaluation results provided in Appendix .2
6. Conclusion
In this paper, we propose a novel method called Task-Aware Lear) that predicts model parameters conditioned on desired model scales and task-specific characteristics. Experimental results on various datasets demonstrated the effectiveness of TAL's ability to predict parameters. Untrained models initialized using TAL achieved significant improvements across various datasets compared to the previ- ous GHN initialization methods. Remarkably, the accuracy of these untrained models even surpassed the performance of models trained using other initialization methods.
Visualization of task-specific model computational graphs. To verify the effectiveness of learngene in dynamically encoding the model computational graph under different task conditions, we visualize the output of learngene. We use 3-12 layers of ViT-small, a total of 10 test models and apply learngene to output their task-specific computational graphs on Decathlon Challenge datasets. We use the PCA (Abdi & Williams, 2010) method to map the high-dimensional features of the model computational graph to 2D space and visualize them. The result is shown in Fig.3. As the task information changes, the model's ' learned computational graph exhibits a significant clustering effect, the learned computational graphs of the model for different
A Acknowledgements
This research was supported by the Jiangsu Science Foundation (BK20243012, BG2024036), the National Science Foundation of China (62125602, U24A20324, 92464301, 62406066), the Fundamental Research Funds
8

---

Learngene Tells You How to Customize: Task-Aware Parameter Initialization at Flexible Scales
for the Central Universities (2242025K30024), Jiangsu Province Science Foundation for Youths (BK20241297), Taihu Lake Innovation Fund for the School of Future Technology of Southeast University, and the Big Data Computing Center of Southeast University.
Knyazev, B., Drozdzal, M., Taylor, G. W., and Romero Sori- ano, A. Parameter prediction for unseen deep architectures. Advances in Neural Information Processing Systems, 34: 29433-29442 2021.
Knyazev, B., Hwang, D., and Lacoste-Julien, S. Can we scale transformers to predict parameters of diverse imagenet models? In International Conference on Machine Learning, pp. 17243-17259, PMLR, 2023.
Impact Statement
This paper presents work whose goal is advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Krizhevsky, A.,, Hinton, G., al. Learning multiple layers of features from tiny images. 2009.
Lake, B. M., Salakhutdinov, R., and Tenenbaum, B. Human-level concept learning through probabilistic program induction. Science, 350(6266); 1332--1333338 2015.
References
Abdi, H. and Williams, L. J. Principal component analysis. Wiley interdisciplinary reviews: computational statistics, 2(4))433-4555 2010
Maas, A., Daly, R. E.,, Pham, T., Huang. D., Ng. A. Y., and Potts, C. Learning word vectors for sentiment analysis. In Proceedings of the 49th annual meeting of the association for computational linguistics: Human language technologies, pp. 142-150, 2011.
Cimpoi, M., Maji, S., Kokkinos, I., Mohamed, S., and Vedaldi, A. Describing textures in the wild. Proceed- ings of the IEEE conference on computer vision and pattern recognition, pp. 3606-3613, 2014.
Maji, S., Rahtu, E., Kannala, J., Blaschko, M., and Vedaldi, Fine-gran grained visual classification of aircraft. arXiv preprint 130.vv3306.5555 2013.
Dolan, B. and Brockett, C. Automatically constructing corpus of sentential paraphrases. In Third international workshop on paraphrasing (IWP2005), 2005.
Mehta, S. and Rastegari, M. Mobilevit: light-weight, general-purpose, and mobile-friendly vision transformer. arXiv preprint arXiv:21 10.02178. 2021.
Dosovitskiy, A. An image is worth 16x16 words: Trans- formers for image recognition at scale. arXiv preprint arXiv:2010.1192: 2020.
Munder, S. and Gavrila, D. M. An experimental study on pedestrian classification. IEEE transactions on pattern analysis and machine intelligence, 1888333883,,,, 2006.
Feng, F., Xie, Y., Wang, J., and Geng. X. Wave: Weight template for adaptive initialization of variable-sized models. arXiv preprint arXiv:2406. /7503, 2024a.
Feng, F. Xie, Y., Yang. X., Wang, J., and Geng, X. Redefiningi creative; dictionary: Towards an enhanced semantic understanding of creative generation. arXiv preprint arXiv:2410.24100 2024b.
Netzer, Y., Wang, T., Coates, A., Bissacco, A., Wu, B., Ng, A. Y., et al. Reading digits in natural images with unsupervised feature learning. In NIPS workshop on deep learning and unsupervised feature learning, volume 2011, pp. Granada, 2011.
Fu, J. Jiang. Y.. Chen, J., Fan, J., Geng, X., and Yang, X. Speculative ensemble: Fast large language model ensem- ble via speculation. arXiv preprint arXiv:2502.01662, 2025.
Nilsback, M.-E. and Zisserman, A. Automated flower classification over a large number of classes. In 2008 Sixth Indian conference on computer vision, graphics & iimage processing, pp. 722-729. IEEE, 2008
Goodfellow, J., Erhan, D., Carrier, P. L., Courville, A., Mirza, M., Hamner, B., Cukierski, W., Tang, Y., Thaler, D., Lee, D.-H., al. Challenges representation learning: A report on three machine learning contests. In Neural information processing: 20th international conference, ICONIP 2013, daegu, korea, november 3-7, 2013. Proceedings, Part 1I1 20, pp. 117-124. Springer, 2013.
Oreshkin, B., Rodríguez López, P., and Lacoste, A. Tadam: Task dependent adaptive metric for improved few-shot learning. Advances in neural information processing systems, 3018.1...11
Peng. Y., Zhang, G., Zhang, M., You, Z., Liu, J., Zhu, Q., Yang, K., Xu, X., Geng, X., and Yang. X. Lmm-rl: Em- powering 3b Imms with strong reasoning abilities through two-stage rule- based rl. arXiv preprint arXiv:2503.075.. 2025.
He, K., Zhang, X., Ren, S., and Sun, J. Delving deep into rec- tifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pp. (026-1034, 2015.

---

Learngene Tells You How to Customize: Task-Aware Parameter Initialization at Flexible Scales
Perez, E., Strub, F., De Vries, H., Dumoulin, V., and Courville, A. Film: Visual reasoning with general conditioning layer. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
Vu, T., Wang, T., Munkhdalai, T., T,, Sordoni, A., Trischler, A., Mattarella-Micke, A., Maji, S., and lyyer, M. Exploring and predicting transferability across nlp tasks. arXiv preprint arXiv:200s.007700 2020.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 18)):9 2019.
Wang, A. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint drXiv:180..07461,2000 2018.
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al. Learning transferable visual models from natural language supervision. International conference on machine learning, pp. 8748-8763 PMLR, 2021.
Wang, Q.-F., Geng, X.,
Lin, S.-X., Xia, S.-Y., Qi, L.,, and Xu, N. Learngene: From open-world your learning task. In Proceedings the AAAI Conference on Artificial Intelligence, volume 36, pp. 8557-856$ 2022.
Wang, S., Feng, K., Li, C.,, Yuan, Y., and Wang, G. Learning generate parameters of convnets for unseen image data. IEEE Transactions on Image Processing, 2024.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21 (140):1-67, 2020.
Warstadt, A. Neural network acceptability judgments. arXiv preprint arXiv: 12471, 2019.
Rebuffi, S.-A., Bilen, H., and Vedaldi, A. Learning multiple visual domains with residual adapters. Advances in neural information processing systems, 30, 2017.
Xia, S., Zhang, M., Yang, X., Chen, R., Chen, H., and
Geng, X. Transformer as linear expansion of learngene. Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 16014-16022, 2024a.
Russakovsky, O., Deng. J., Su, H., Krause, J., Satheesh, S., Ma, S., Huang, Z., Karpathy, A., Khosla, A., Bernstein, M., et al. Imagenet large scale visual recognition challenge. International journaal of computer vision, 115: 211 211-252,2015.
Xia, Y., Zhu, W., Yang, X., and Geng. X. Exploring learn- gene via stage-wise weight sharing for initializing variable- sized models. arXiv preprint arXiv:2404. 16897, 2024b.
Xiao, H., Rasul, K., and Vollgraf, R. Fashion-mnis novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv: 1708.07777 2017.
Shi, B., Xia, S., Yang, X., Chen, H., Kou, Z., and Geng, X. Building variable-sized models via learngene pool. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 14946 4996-11454, 2024.
Xie, Y., Feng, F., Wang. J., Geng, X., and Rui, Y. Kind: Knowledge integration and diversion in diffusion models. arXiv preprint arXiv:2408.07333
Shu, C., Liu, Y., Gao, J., Yan, Z., and Shen, C. Channel- wise knowledge distillation for dense prediction. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 531153111111 13520, 2021.
Yang, X., Peng, Y., Ma, H., Xu, S., Zhang, C., Han, Y., and Zhang. H. Lever Im: configuring in-context sequence to lever large vision language models. Advances in Neural In- formation Processing Systems, 37:100341-10036 2024.
Soomro, K., Zamir, A. R., and Shah, M. A dataset of 101 human action classes from videos in the wild. Center for Research in Computer Vision, 2(11):1-7, 2012.
Yao, J., Zhang, Z., and Xu, Z.-Q. J. An analysis for reasoning bias of language models with small initialization. arXiv preprint arXiv:2502.04375, 2025.
Stallkamp, J., Schlipsing, M., Salmen, J., and Igel, C. Man vs. computer: Benchmarking machine learning algorithms for traffic sign recognition. Neural networks, 32:323-332 2012.
Zhang, C., Ren, M., and Urtasun, R. Graph hypernet- works for neural architecture search. arXiv preprint arXhv:1810.05749 2018.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A,, Lacroix, T., Rozière, B., Goyal, N., Hambro, E.,, Azhar, F., et al. Llama: Open and efficient foundation 13971, 2023.
Zhou, X., Knyazev, B., Jolicoeur- A., and Fu, J. Logah: Predicting 774 mwwliii--aaaaaaaan transformers using graph hypernetworks with 1/100 parameterss. arXiv preprint arXiv:2405. 16287, 2024.
language models. arXiv preprint
Tschandl, P., Rosendahl, C., and Kittler, H. The ham 1 10000 dataset, a large collection of multi-source dermatoscopic images of common pigmented skin lesions. scientific data. 2018: 5: 180161. Search in, 2, 2018.
10

---

Learngene Tells You How to Customize: Task-Aware Parameter Initialization at Flexible Scales
.I. Appendix 1
The Visual Domain Decathlon Challenge tests the ability of visual recognition algorithms to handle images from different visual domains. It includes 10 datasets in total:
1. ImageNet-1K (IN-1K) is the largest dataset in the Decathlon Challenge, containing 1,000 categories and 1.2 million images.
2. CIFAR-100 (C100) contains 60,000 32 x 32 color images for 100 object categories.
3. Aircraft (Airc.) contains 100 images for each of 100 different aircraft model variants, such as the Boeing 737-400 and the Airbus A310.
4. Daimler Pedestrian Classification (DPed) consists 50,000 grayscale pedestrian and non-pedestrian images, cropped and resized to 18 x 36 pixels.
5. Describable Textures (DTD) is texture database consisting of 5,640 images, organized into 47 categories such as bubbly, cracked andmarbled.
6. German Traffic Signs (GTSRB) contains cropped images for 43 common traffic sign categories in different image resolutions.
7. Omniglot (OGlt) consists of 1,623 different handwritten characters from 50 unique alphabets.
8. SVHN is real-world digit recognition dataset with around 70,000 32 32 images.
9. UCF101 Dynamic Images (UCF) is an action recognition dataset of realistic human action videos, collected from YouTube. It contains 13,320 videos across 101 action categories. In the Decathlon Challenge, the videos converted into images using Dynamic Image encoding, which summarizes each video into image based ranking principle.
10. Flowers10% (Flwr) a fine-grained classification task with 102 flower categories from the UK, each consisting of 40 to 258 images.
The detailed statistics of the datasets be found htttt:::/www.raatttt wwtppo.........////eeeeeeeeeeeee
.2. Appendix 2
Here we present the detailed results of the analysis and ablation studies for TAL (w/o I.T.) and TAL (w/o ans-ant) methods. The initialization performance of TAL (w/o I.T.) and TAL (w/o ans-ant) methods on the Decathlon dataset is shown in Tab.9.
Table Performance of untrained models on Decathlon datasets initialized with TAL without using task informationn(..)) in learngene or ancestry model(ans-net) on Decathlon datasets.
MODEL
METHOD
AIRC.
C100
DPED
DTD
GSTR
OGLE
SVHN
UCF
FLWR
AVG
TAL(W/O)
4.05
17.61
72.65
11.12
65.9
1.76
35.77
2.25
9.61
24.52
3-TINY
TAL(W/O ANS-NET)
3.39
49.92
92.74
20.80
99.39
0.18
90.2
11.37
50.69
46.52
TAL(W/O L.T.)
12.18
40.62
72.89
23.30
99.20
52.57
88.77
2.87
34.22
47.40
6-TINY
TAL(W/O ANS-NET)
3.27
49.86
92.38
20.05
99.35
0.17
90.34
12.09
51.18
46.52
TAL(W/O T.L.)
10.23
36.09
65.83
17.13
96.00
31.28
83,92
3.02
25.10
40.96
12-TINY
TAL(W/O ANS-NET)
2.88
49.45
93.01
21.22
99.31
0.15
90.43
12.19
$1.08
46.64
TAL(W/O T.L.)
12.54
40.65
80.24
20.69
98.99
49.85
85.99
4.35
36.86
47.79
3-SMALL
TAL(W/O ANS-NET)
2.88
50.76
92.76
21.97
99.35
0.09
90.51
11.53
52.35
46.91
TAL(W/O T.L.)
13.68
42.43
77.19
22.61
99.30
52.65
89,28
3.28
37.75
48.69
6-SMALL
TAL(W/O ANS-NET)
3.12
50.91
93.10
22.29
99.29
0.14
90.57
12.81
52.25
47.16
TAL(W/O T.L.)
11.73
37.67
60.17
16.97
55.30
23.35
76.45
2.97
29.22
34.87
12-SMALL
TAL(W/O ANS-NET)
2.88
50.86
93.52
21.76
99.34
0.12
90.64
12.19
$2.16
47.05

---

Learngenc Tells You How to Customize: Task-Aware Parameter Initialization at Flexible Scales
.3. Appendix 3
We discuss simplified case of our TAL method and provide theoretical derivation to complete this missing part.
1. Problem Definition and Optimization Objective
We define the following setup:
Hypernetwork 6 - RR A multilayer perceptron (MLP) that maps from parameter space θ model parameter space Rd, generating parameters y- (0).
Model M Also MLP, using parameters p perform a binary (0.1) classification task and compute the loss C(p).
The optimization objective is to train the hypernetwork H/ to minimize the cross-entropy loss:
y))og(1 fM (x:H())))
Convergence Analysis
Theorem (Convergence to Stationary Point): Assume the following conditions hold:
The loss function [(p) 3-smooth
Hyperneworr H() is Li-Lisschitz continuous
The composed function (0) has bounded gradients
Then, using gradient descent with learning rate 17 T. after T iterations:
3 (00) -(H()))
TT
Proof:
By B-smoothness of L and LH-Lipschitz continuity of H, the composite function ((H is (Lmoohsmooth. For (Ln3))smoote function, when using gradient descent with learning rate 17< LNB
LH08 2
((H()))-((((++))))
(Vo))((())))
Summing over -=0,1.... T- and rearranging:
T-1
(o)( ((H()))) (OT))
17
Lu
-
Luc
2- Since 7 Lnz implies 1 Ln >
O, and using the minimum gradient norm:
T-1 (oL(H()))) Σ
min
(oL(H(()))
((((H)
L.g00
12

---

Learngenc Tells You How to Customize: Task-Aware Parameter Initialization at Flexible Scales
With proper learning rate, Lu 2
resulting in:
[ 2(C(H( Tn
This shows that as 7 00, the gradient norm approaches zero, indicating convergence to stationary point. D
Corollary 1 (Convergence Rate): Under the conditions of Theorem 1, the gradient descent method converges to stationary point at rate of (1/vTT).
Optimality Analysis
Theorem (Universal Approximation; If the hypernetwork H a sufficiently wide and deep MLP, then for any >>> and any target parameter p" E ERA, there exists parameter 9 such that:
||"10)=pp|\<
Proof:
According to the universal approximation theorem, sufficiently wide MLP approximate any continuous function compact domain to arbitrary precision. Treating the mapping from a fixed input to the target parameter vector p' as constant
function, there exists an MLP architecture for Il and parameters 0 such that (0) || < 8 for any d > 0.
Corollary 2 (Approximation of Optimal Loss): Under the conditions of Theorem 3, for any c 0, there exists a hypernetwork H and parameters / such that:
((H(0))
Conclusion
Our analysis of hypernetwork optimization for binary classification with MLPs has established:
Gradient-bant optimization of hyperetworks converges to stationary points at rate O(I/VT) under standard smoothness assumptions.
Convergence:
Approximation Capability: Sufficiently expressive hypernetworks can approximate optimal model parameters to arbitrary precision.
13


Reviews:
[{"review_id": "U9aofeFTzo", "reviewer": "Reviewer_FFPW", "summary": "This paper presents Task-Aware Learngene (TAL), designed to initialize large models via parameter prediction. To accomplish this, the authors first employ an encoder-decoder architecture for the TAL model and train it under the supervision of an ancestry model to facilitate knowledge transfer. Subsequently, with the aim of improving the multi-task generalization ability of downstream models, the authors fine-tune the trained TAL model on a multi-task dataset, thereby acquiring multi-task knowledge from the ancestry model. Comparative evaluations against alternative methods (such as GHN-3 and LoGAH) underscore the efficacy of TAL across various scenarios.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your reviews and address your concerns as follows. \n### Q1\nWere all methods trained on the same dataset? If so, why does only TAL+ perform well? If not, the comparison is unfair, undermining the claim's validity.\n### A1\nTAL and GHN-3, LoGAH do use the same model training dataset, ViTs-1K, which ensures a fair comparison between them. \nThe results in Table 1 and Table 2 of the paper clearly demonstrate the significant advantage of TAL under this fair comparison, as TAL outperforms the comparison method GHN-3/LoGAH on almost all datasets for the task of image classification.\nBeyond this, TAL+ introduces the augmented dataset ViTs+-1K to further enhance large model initialization, making the use of an augmented dataset one of our key contributions.\n### Q2\nLack of Supporting Evidence for Convergence Claims.\n### A2\nWe use 12-layer ViT-Small initialized with five different methods and train it on Decathlon datasets. We plot the training loss versus epochs to observe convergence speed. The table below presents results for the UCF dataset, where we observe that ViT models initialized with TAL/TAL+ methods demonstrate significantly faster convergence. Similar trends are observed across other datasets. The corresponding convergence plots for all datasets will be included in the appendix, as images cannot be displayed in the response.\n|**Epoch**|**10**|**20**|**30**|**40**|**50**|**60**|**70**|**80**|**90**|**100**|\n|---|---|---|---|---|---|---|---|---|---|---|\n|RandInit|59.9801|52.7643|45.5558|33.7087|17.9540|6.8132|2.9012|1.9813|1.6555|1.0367|\n|GHN-3|123.9892|114.4703|106.5648|96.6006|87.0042|75.6672|60.6773|47.0442|33.3951|23.7440|\n|LoGAH|116.4734|109.3229|101.4643|94.3091|86.6136|76.6659|65.2919|50.5523|37.1268|23.7927|\n|TAL|8.4604|7.0408|3.1283|2.2561|1.8384|1.5562|1.0961|0.8141|0.9439|0.7467|\n|TAL+|2.0044|2.7601|0.9048|0.8046|0.7405|0.5503|0.6014|0.3757|0.4805|0.4478|  \n\n**Table:**  Training Loss Comparison of ViT-small with Different Initialization Methods on UCF Dataset Across Epochs\n### Q3\nLimited Validation of Untrained Initialization Effectiveness.\n### A3\nWe demonstrate TAL's performance on unseen tasks in Table 5 of Section 5, encompassing tasks across diverse domains including Fashion MNIST (clothing and fashion items), FER2013 (facial expression recognition), and HAM10000 (dermatological skin lesion classification). The results indicate that models initialized with TAL consistently outperform those using random initialization and the LoGAH method.\n### Q4\nIncomplete Comparison with Task Information.\n### A4\nWe augment the GHN-3 method with the same task information utilized in our TAL approach to evaluate its effectiveness. We initialize 12-layer ViT-Tiny and 12-layer ViT-Small using both standard GHN-3 and our enhanced GHN-3 with task information (GHN-3 w/ T.I.), then assessed initialization quality on the Decathlon datasets. Results demonstrate that incorporating task information also improve GHN-3's performance.\n| Model      | Method         | Airc. | C100  | DPed  | DTD   | GSTR  | OGle  | SVHN  | UCF   | Flwr  | Avg\\_Acc |\n|------------|--------------|------|------|------|------|------|------|------|------|------|--------|\n| **12-Tiny** | GHN-3        | 3.24 | 35.19 | 87.72 | 6.86  | 89.12 | 0.06  | 10.00 | 3.43  | 10.88 | 27.39  |\n|            | GHN-3 w/ T.I. | 5.31 | 22.03 | 84.66 | 11.49 | 94.01 | 0.06  | 10.00 | 1.28  | 21.27 | 27.79 **(+0.40)** |\n| **12-Small** | GHN-3        | 2.64 | 5.55  | 84.30 | 7.23  | 84.39 | 0.06  | 10.00 | 1.74  | 9.51  | 22.82  |\n|            | GHN-3 w/ T.I. | 6.60 | 21.84 | 82.70 | 14.10 | 91.52 | 0.06  | 10.00 | 1.23  | 32.25 | 28.92 **(+6.10)** |\n\n**Table:** Performance of untrained models initialized with GHN-3 and GHN-3 with Task Initialization (GHN-3 w/ T.I.)\n### Q5\nInsufficient Motivation for the Ancestry Model.\n### A5\nThe role of the ancestry model is clearly articulated in Section 4. It provides both the soft labels necessary for Stage 1 training and task-specific labeling information for multitask scenarios. The ablation experiments in Table 8 of Section 5 quantitatively demonstrate the ancestry model's significant contribution to performance, confirming its critical importance in enhancing parameter initialization quality.\n### Q6\nAmbiguity in \"Learngene\" Concept.\n### A6\nThe essence of the Learngene framework lies in condensing critical knowledge from an ancestry model to initialize downstream models. Our implementation employs an encoder-decoder architecture, departing from previous Learngene methods that relied on manual parameter extraction and heuristic designs. Despite innovations in implementation techniques, our approach maintains the basic paradigm of Learngene while providing greater flexibility and adaptability. In our experimental analysis, we conduct a detailed ablation study on the ancestry model, clearly demonstrating its significant contribution to performance improvement.", "claims_and_evidence": "In the right column of lines 36–40, the authors state that “…its effectiveness diminishes when initializing larger models like ViT-base,” a claim supported by Table 1. However, as indicated in the same table, the proposed TAL method also fails to initialize the ViT-base model on untrained configurations, achieving a very low-performance value of 0.1, identical to GHN-3 and LoGAH. While the TAL+ model outperforms other methods, its improvement appears to stem from training on an enhanced dataset (lines 236–244). My concern centers on whether all methods in Table 1 were trained on the same dataset. If so, why does only the TAL+ method perform well, while the performance of other methods is nearly random? If not, this comparison is problematic, as it evaluates methods trained on different datasets, thereby undermining the validity of the claim.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "The paper does not contain any formal proofs or theorems.", "experimental_designs_or_analyses": "I have reviewed the experimental design and analyses outlined in Section 5, and the following concerns arise: \\\n1.\tLack of Supporting Evidence for Convergence Claims: The authors assert that “Models initialized with TAL/TAL+ converge faster and outperform …” (line 327). However, this claim is not substantiated by detailed experimental results, figures, or tables. While higher accuracy values are presented, no evidence is provided to demonstrate faster convergence, which is a critical aspect of the claim. \\\n2.\tLimited Validation of Untrained Initialization Effectiveness: The effectiveness of TAL/TAL+ in initializing models without training has only been validated on the Decathlon tasks. This narrow scope raises questions about the method’s broader applicability, as it has not been tested on other NLP tasks or diverse domains. Such validation is essential to establish the generalizability of the approach. \\\n3.\tIncomplete Comparison with Task Information: Task information appears to be a feature that could also benefit other methods, such as GHN-3. However, the experiments in Table 8 focus solely on the proposed TAL method without comparing it to other methods that incorporate task information. This omission limits the ability to assess whether the observed improvements are unique to TAL or could be achieved by enhancing existing methods with similar information.", "relation_to_broader_scientific_literature": "Task-Aware Learngene (TAL) builds upon and extends existing research in model initialization, knowledge distillation, and multi-task learning. Several prior works, such as GHN-3 and LoGAH, have explored ways to generate or predict parameters for a great diversity of models. TAL differs with these earlier studies by leveraging an encoder-decoder structure and an ancestry model for knowledge transfer, as well as its explicit use of multi-task fine-tuning to enhance downstream generalization. By effectively integrating knowledge from an ancestry model and refining it through multi-task datasets, TAL demonstrates that parameter prediction for large models can achieve strong performance across diverse situations.", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "Strengths: \\\n1.\tThe proposed TAL effectively integrates multi-task knowledge from an ancestry model via an encoder-decoder structure and multi-task fine-tuning, facilitating robust parameter prediction for diverse downstream tasks. \\\n2.\tEmpirical comparisons with methods like GHN-3 and LoGAH suggest TAL’s consistent performance improvements and highlight its potential for broader applications.\n\nWeaknesses: \\\nIn addition to the experimental issues noted above, I have a few further concerns: \\\n1.\tInsufficient Motivation for the Ancestry Model: The rationale for introducing the ancestry model has not been fully articulated. A more detailed justification of its role and necessity would strengthen the paper’s foundation. \\\n2.\tAmbiguity in the Concept of “Learngene”: The term “Learngene” is somewhat confusing in this context. In prior literature, it appears to refer to “critical components from the ancestry model used to initialize models of various sizes.” However, in this paper, it seems to function more as a “critical component for a parameter prediction model, which is learned through multi-task tuning and conditioned on task embeddings.” This shift in interpretation raises questions about whether the observed performance improvements are truly attributable to the “Learngene” concept or are instead driven by the multi-task tuning process and the incorporation of task embeddings.", "other_comments_or_suggestions": "None.", "questions_for_authors": "None.", "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": "y1rdj093kY", "reviewer": "Reviewer_7VCD", "summary": "Authors propose TAL, an encoder-decoder method to generate parameters for initializing models of various sizes given a model architecture and a task embedding.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your reviews and address your concerns as follows. \n### Q1\n\nThe method relies on a single ancestry model, which might be problematic as the latest model might serve as a better ancestry model, and then requires retraining. And since the entire framework requires a certain level of \"pre-training\" before it can serve as an initialization generator, this inherently brings a conflict. It is hard to achieve \"train once and use forever\". Maybe authors can show that when serving as initialization, the choice of ancestry model isn't that important.\n\n### A1\n\nOur approach allows flexibility in choosing the ancestry model based on the TAL training datasets, including the consideration of using the most recent model. It should be clarified that our approach does not involve any retraining of the ancestry model, instead we retrain the TAL encoder-decoder framework, not the ancestry model itself. \n### Q2\n\nAuthors seem to include ViTs of various depths in their experiments as a demonstration for need of models of different sizes. What would be a potential use case for models with various depths? (i.e. wider models or deeper models) As far as I am concerned, a certain depth-to-width ratio is usually adopted and people rarely change that during application.\n\n### A2\n\nOur approach employs a structured scaling strategy rather than a fixed depth-to-width ratio. We sample layers from 3-12 and hidden dimensions with a step size of 32, ranging from 192-768 across all model depths. Attention heads are selected to ensure divisibility with hidden dimensions, maintaining architectural integrity. This parameterization offers greater freedom in depth and width combinations while still preserving architectural coherence, enabling the generation of models with diverse parameter counts tailored to various computational budgets. This flexibility addresses deployment needs across resource-constrained devices to high-performance environments.", "claims_and_evidence": "1. Yes", "methods_and_evaluation_criteria": "1. The method relies on a single ancestry model, which might be problematic as the latest model might serve as a better ancestry model, and then requires retraining. And since the entire framework requires a certain level of \"pretraining\" before it can serve as an initialization generator, this inherently brings a conflict. It is hard to achieve \"train once and use forever\". Maybe authors can show that when serving as initialization, the choice of ancestry model isn't that important.\n2. Authors seem to include ViTs of various depths in their experiments as a demonstration for need of models of different sizes. What would be a potential use case for models with various depths? (i.e. wider models or deeper models) As far as I am concerned, a certain depth-to-width ratio is usually adopted and people rarely change that during application.", "theoretical_claims": "n/a", "experimental_designs_or_analyses": "1. Experiments seem to be reasonable and demonstrating TAL's effectiveness over GHN-based method.", "relation_to_broader_scientific_literature": "n/a", "essential_references_not_discussed": "n/a", "other_strengths_and_weaknesses": "n/a", "other_comments_or_suggestions": "n/a", "questions_for_authors": "n/a", "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": "TMxe9gaNFB", "reviewer": "Reviewer_AdW5", "summary": "This paper addresses the high computational and storage overheads involved in training large pretrained models by focusing on effective parameter initialization. Building on recent advances in Graph HyperNetworks (GHN) and the Learngene framework, the authors propose a novel method called Task-Aware Learngene (TAL). TAL is designed to capture shareable, task-specific knowledge from a well-trained “ancestry” model and use it to predict initial parameters for models of varying scales.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your reviews and address your concerns as follows. \n### Q1\nAn important Task-Aware Parameter initialization baseline is not compared: Learning...(TIP 2024). The models used in this paper is too small and shallow. \n### A1\nThank you for pointing this out. We will cite it in the relevant section. As for the model size predicted by TAL+, our framework has been tested on 12-layer ViT-Base models with 86M parameters, far exceeding ResNet-34’s 22M. It also handles deeper architectures, including the 16-layer ViT-Base (115M), which achieves 31.52% accuracy on ImageNet-1K without further training. Scaling up to an 18-layer ViT-Base (129M), the model maintains a competitive 22.25% accuracy. These experiments validate TAL’s effectiveness on larger models.\n### Q2\nThere is limited discussion on the scalability of TAL.\n### A2\nOur experiments cover various model sizes, including ViT-Base (12 layers). Indeed, TAL is designed for scalability and can support larger models. We are currently testing it on large language models (e.g., LLaMA) and will present the results in future work.\n### Q3\nThe theoretical underpinning of the method is not fully explored.\n### A3\nWe discuss a simplified case of our TAL method and provide a theoretical derivation. We reach the following conclusions:\n\n- **Convergence:** Gradient-based optimization of hypernetworks converges to stationary points under standard smoothness assumptions.\n \n- **Approximation Capability:** Sufficiently expressive hypernetworks can approximate optimal model parameters to arbitrary precision.\n\n**1. Problem Definition and Optimization Objective**\nWe define the following setup:\n- **Hypernetwork** $H: \\Theta \\rightarrow \\mathbb{R}^d$  A multilayer perceptron (MLP) that maps from parameter space $\\Theta$ to model parameter space $\\mathbb{R}^d$, generating parameters $p = H(\\theta)$.\n- **Model** $M$  Also an MLP, using parameters $p$ to perform a binary (0,1) classification task and compute the loss $\\mathcal{L}(p)$.\nThe optimization objective is to train the hypernetwork $H$ to minimize the cross-entropy loss:\n$$\\min_{\\theta} \\mathcal{L}(H(\\theta)) = \\mathbb{E}_{(x,y) \\sim \\mathcal{D}} \\left[ -y\\log(f_M(x; H(\\theta))) - (1-y)\\log(1-f_M(x; H(\\theta))) \\right]$$\n\n**2. Convergence Analysis**\n\n**Theorem 1 (Convergence to Stationary Point):** Assume the following conditions hold:\n- The loss function $\\mathcal{L}(p)$ is $\\beta$-smooth\n- Hypernetwork $H(\\theta)$ is $L_H$-Lipschitz continuous\n- The composed function $\\mathcal{L}(H(\\theta))$ has bounded gradients\nThen, using gradient descent with learning rate $\\eta < \\frac{2}{L_H\\beta}$, after $T$ iterations:\n$$\\min_{t=0,1,...,T-1} \\||\\nabla_{\\theta} \\mathcal{L}(H(\\theta_t))\\||^2 \\leq \\frac{2(\\mathcal{L}(H(\\theta_0)) - \\mathcal{L}(H(\\theta^*)))}{T\\eta}$$\n\n**Proof:**  \nBy $\\beta$-smoothness of $\\mathcal{L}$ and $L_H$-Lipschitz continuity of $H$, the composite function $\\mathcal{L}(H(\\theta))$ is $(L_H\\beta)$-smooth. For a $(L_H\\beta)$-smooth function, when using gradient descent with learning rate $\\eta < \\frac{2}{L_H\\beta}$:\n$$\\mathcal{L}(H(\\theta_t)) - \\mathcal{L}(H(\\theta_{t+1})) \\geq \\eta\\left(1 - \\frac{L_H\\beta\\eta}{2}\\right)\\||\\nabla_{\\theta}\\mathcal{L}(H(\\theta_t))\\||^2$$\nSumming over $t=0,1,...,T-1$ and rearranging:\n$$\\sum_{t=0}^{T-1}\\||\\nabla_{\\theta}\\mathcal{L}(H(\\theta_t))\\||^2 \\leq \\frac{\\mathcal{L}(H(\\theta_0)) - \\mathcal{L}(H(\\theta_T))}{\\eta\\left(1 - \\frac{L_H\\beta\\eta}{2}\\right)}$$\n$$\\leq \\frac{\\mathcal{L}(H(\\theta_0)) - \\mathcal{L}(H(\\theta^*))}{\\eta\\left(1 - \\frac{L_H\\beta\\eta}{2}\\right)}$$\nSince $\\eta < \\frac{2}{L_H\\beta}$ implies $1 - \\frac{L_H\\beta\\eta}{2} > 0$, and using the minimum gradient norm:\n$$T \\cdot \\min_{t=0,1,...,T-1} \\||\\nabla_{\\theta}\\mathcal{L}(H(\\theta_t))\\||^2 \\leq \\sum_{t=0}^{T-1}\\||\\nabla_{\\theta}\\mathcal{L}(H(\\theta_t))\\||^2$$\n$$\\leq \\frac{\\mathcal{L}(H(\\theta_0)) - \\mathcal{L}(H(\\theta^*))}{\\eta\\left(1 - \\frac{L_H\\beta\\eta}{2}\\right)}$$\nWith proper learning rate, $1 - \\frac{L_H\\beta\\eta}{2} \\geq \\frac{1}{2}$, resulting in:\n$$\\min_{t=0,1,...,T-1} \\||\\nabla_{\\theta} \\mathcal{L}(H(\\theta_t))\\||^2 \\leq \\frac{2(\\mathcal{L}(H(\\theta_0)) - \\mathcal{L}(H(\\theta^*)))}{T\\eta}$$\nThis shows that as $T \\to \\infty$, the gradient norm approaches zero, indicating convergence to a stationary point. $\\square$\n\n**3. Optimality Analysis**\n\n**Theorem 2 (Universal Approximation):** If the hypernetwork $H$ is a sufficiently wide and deep MLP, then for any $\\delta > 0$ and any target parameter $p^* \\in \\mathbb{R}^d$, there exists a parameter $\\theta$ such that:\n$$\\||H(\\theta) - p^*\\|| < \\delta$$\n\n**Proof:**  \nAccording to the universal approximation theorem, a sufficiently wide MLP can approximate any continuous function on a compact domain to arbitrary precision. Treating the mapping from a fixed input to the target parameter vector $p^*$ as a constant function, there exists an MLP architecture for $H$ and parameters $\\theta$ such that $\\||H(\\theta) - p^*\\|| < \\delta$ for any $\\delta > 0$. $\\square$", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "The proposed TAL method leverages an encoder–decoder architecture where the encoder (learngene) extracts and transfers task-specific knowledge from an ancestry model, and a task hypernetwork generates task-specific bias parameters for a task-specific layer.\nThe method is evaluated on both vision and language tasks using standard model architecture datasets (e.g., ViTs-1K, ViTs+-1K, GPTS-1K) and downstream benchmarks (ImageNet-1K, Decathlon, MRPC, COLA, RTE, IMDB).", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "Experiments are conducted on multiple vision and language tasks with a variety of model scales (e.g., ViT-Tiny, ViT-Small, ViT-Base, and GPT-2 models with different depths).\nHowever, an important Task-Aware Parameter initialization baseline is not compared:  Learning to Generate Parameters of ConvNets for Unseen Image Data (TIP 2024)", "relation_to_broader_scientific_literature": "AL builds on prior work in hypernetworks and Task-Aware Parameter  initialization methods.", "essential_references_not_discussed": "An important Task-Aware Parameter  initialization works ins not discussed:  Learning to Generate Parameters of ConvNets for Unseen Image Data (TIP 2024)", "other_strengths_and_weaknesses": "The proposed method is sound and interesting.\nHowever, there are several weaknesses:\n\nAn important and recent Task-Aware Parameter initialization works [1] is not compared.\n\nThere is limited discussion on the scalability of TAL when applied to extremely large models, which could be crucial for real-world applications.\n\nThe models used in this paper is too small and shallow. Deeper network like ResNet-34 is encourgaed to be used for testing, as in previous works [1].\n\nThe theoretical underpinning of the method, including convergence properties and optimality guarantees, is not fully explored.\n\n[1]   Learning to Generate Parameters of ConvNets for Unseen Image Data (TIP 2024)", "other_comments_or_suggestions": "N.A", "questions_for_authors": "N/A", "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": "4b8Wzgo0Qo", "reviewer": "Reviewer_56qr", "summary": "This paper proposes a novel parameter initialization method called TAL, aiming to enhance the initialization effect of models for different tasks. Building upon the previous GHN and Learngene frameworks, TAL addresses their limitations. Although the GHN method is effective, it performs inadequately when dealing with large-scale models and requires retraining for each new task. In contrast, by integrating the Learngene framework with a task-aware mechanism, TAL is able to share knowledge across multiple tasks, thereby improving the accuracy of parameter initialization and eliminating the need for separate training for each task. The experimental results demonstrate that TAL outperforms methods such as GHN and LoGAH in both visual and natural language processing tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your reviews and address your concerns as follows. \n### Q1\nThe author should elaborate on the learngene theory and the computational graph in more detail.\n### A1\nWe elaborate on the concept of Learngene in the introduction and related work of the paper. The essence of the Learngene framework lies in condensing critical knowledge from an ancestry model to initialize downstream models. Our implementation uses an encoder-decoder architecture to inherit knowledge from the ancestry model and generate downstream models.\nRegarding computational graphs, we provide a detailed description in lines 155-164 of the paper and add relevant examples in the appendix to enhance the reader's understanding of the concept. The computation graphs primarily include node\\_feat, node\\_info, and edges\\_embedding. To clarify, we provide an example using the ViT model to illustrate how computation graphs are constructed.\n\n```python\nnode\\_feat:\ntensor([\n    [9],   # 'input'\n    [13],  # 'pos\\_enc'\n    [5],   # 'msa'\n    [12],  # 'ln'\n    [3],   # 'linear'\n    [7],   # 'sum'\n    [3],   # 'linear'\n])\nnode\\_info:\n[\n    [[0, 'input', 'input', None, False, False]],\n    [[1, 'pos\\_enc', 'pos\\_enc', (1, 768, 14, 14), False, False]],\n    [[2, 'msa', 'msa', (1, 768, 14, 14), False, False]],\n    [[3, 'ln', 'ln', (1, 768), False, False]],\n    [[4, 'linear', 'linear', (768, 3072), False, False]],\n    [[5, 'sum', 'sum', None, False, False]],\n    [[6, 'linear', 'linear', (3072, 768), False, False]],\n]\nedges embedding:\ntensor([\n    [0, 1, 1],  # input -> pos\\_enc\n    [1, 2, 1],  # pos\\_enc -> msa\n    [2, 3, 1],  # msa -> ln\n    [3, 4, 1],  # ln -> linear\n    [4, 5, 1],  # linear -> sum\n    [5, 6, 1],  # sum -> linear\n])\n```\nHopefully, this simple example helps clarify the concept.\n### Q2\nWhy does TAL+ perform worse than TAL in some cases.\n### A2\nThe primary distinction between TAL and TAL+ lies in the size of the model library used during training. TAL+ leverages a larger model library, which generally improves performance, as evidenced by its superior results in 7 out of 9 tasks in Decathlon datasets. However, in certain cases, such as the second row of 12-Tiny and 12-Small in Table 1, TAL+ underperforms TAL. This discrepancy arises from differences in sample distribution—since TAL+ introduces more large models into training, the proportion of Tiny-sized samples decreases significantly. This imbalance affects TAL+’s ability to adapt to smaller models.\n\nIn Table 2, TAL+ performs worse than TAL on Airc and OGle. Both datasets have high class diversity but limited samples per class, with OGle being a few-shot task comprising 1623 unique characters from 50 alphabets. The key factor here is model capacity—TAL+ models trained on a larger model library tend to have more parameters, which can be less effective in scenarios with limited training samples. The increased capacity of TAL+ models may hinder generalization on these high-diversity, low-resource classification tasks. Despite these limitations, TAL+ is capable of handling larger models and achieves superior performance on most tasks.\n### Q3\nIn Formula 6, Laux needs to be further explained.\n### A3\nThe specific expression and calculation of Laux in Formula 6 is described in detail in Formula 4. Briefly, we use the output of the ancestry model as a soft label to align the model's output initalized by TAL.\n### Q4 \nWhat is the difference between the Decoder and Encoder in Figure 2?\n### A4\nWe introduce Encoder and Decoder in section4 of the paper. The Encoder learns and characterizes the model's computational graph, capturing its structure and parameter relationships. It adopts a Transformer architecture to efficiently handle long-range dependencies in graph structures.  \n\nThe Decoder decodes the learned graph representations to generate the parameter matrices for target models. These matrices are then trimmed or copied to match the required model shape. The Decoder uses MLPs to convert graph representations into parameter values.\n### Q5\nHow much additional overhead does the two-stage training incur compared to other methods?\n### A5\nTable 4 shows the computational overhead across methods, with TAL's total time of 36.19 hours including both training stages (28.47 hours for stage 1 and 7.72 hours for stage 2). TAL requires fewer training epochs in stage 2 than comparison methods, resulting in less total training time.\n### Q6\nHow is an entire ViT initialized? \n### A6\nOur approach initializes the ViT model layer by layer. During the TAL forward pass, the encoder processes the computational graph to learn node embeddings, then the decoder generates parameters through internal loops where the decoder is called multiple times for different parameter groups of different layers (attention weights, MLP weights, etc.). Despite being sequential, this initialization process completes in under 0.1 seconds for an entire ViT model, preserving architectural information while maintaining efficiency.", "claims_and_evidence": "The argument of this paper is that the latest graph hypernetwork methods, such as LoGAH, exhibit poor initialization accuracy on IN-1K and CIFAR-100 when initializing larger models. Furthermore, the author presents the accuracy on IN-1K and other datasets in the experimental section, reasonably addressing the proposed issues.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are reasonable for the current problem.", "theoretical_claims": "The theoretical basis is the learngene. Additionally, the concept of the computational graph is also mentioned. It is necessary for the author to elaborate on these concepts in more detail.", "experimental_designs_or_analyses": "The experimental design is relatively reasonable, covering not only diverse visual tasks but also language tasks, which demonstrates the certain universality of the method proposed in this paper.", "relation_to_broader_scientific_literature": "Previous methods, such as LoGAH, show poor initialization accuracy on IN-1K and CIFAR-100 when initializing larger models. The TAL proposed by the author performs better on datasets like IN-1K.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "The experiments are sufficient, but the analysis of the experiments is rather hasty. For example, why does TAL+ perform worse than TAL in some cases (the second row of 12-Tiny in Table 1)?", "other_comments_or_suggestions": "1 In Formula 6, Laux needs to be further explained.\n\n2 It is recommended that the author optimize the logical structure of the text and provide a clearer elaboration of the two core concepts, learngene and computational graph.", "questions_for_authors": "1 What is the difference between the Decoder and Encoder in Figure 2?\n\n2 Why does TAL+ perform worse than TAL in some cases, such as the second row of 12-Tiny and the second row of 12-Small in Table 1? Additionally, why does TAL+ far outperform the other three methods in Table 1? And why does TAL+ perform much worse than TAL on AIRC. and OGLE in Table 2?\n\n3 How much additional overhead does the two-stage training (Figure 2ab) incur compared to other comparative methods?\n\n4 How is an entire ViT initialized? Is it initialized layer by layer or are the parameters of the entire network initialized directly?", "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:

Meta-review: This paper proposes an approach for initializing parameters for specific tasks by predicting parameters.  In particular, the authors use an encoder-decoder model to transfer information across tasks.  Three experts reviewed this paper with the majority arguing for acceptance, with the lone dissent both saying they are somewhat unfamiliar with the topic and not providing an update to their review or a discussion post-rebuttal.  Most comments focused on clarity questions or requests for additional comparisons, which were provided in the rebuttal with most reviewers noting that they were sufficient to address their concerns.  As such, the AC finds insufficient justification to overturn the majority recommendation to accept, and the authors are encouraged to consider the comments made by reviewers to improve their paper.
Decision: Accept (poster)