We thank you for your reviews and address your concerns as follows.

Q1

The author should elaborate on the learngene theory and the computational graph in more detail.

A1

We 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. Regarding 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.

node\_feat:
tensor([
    [9],   # 'input'
    [13],  # 'pos\_enc'
    [5],   # 'msa'
    [12],  # 'ln'
    [3],   # 'linear'
    [7],   # 'sum'
    [3],   # 'linear'
])
node\_info:
[
    [[0, 'input', 'input', None, False, False]],
    [[1, 'pos\_enc', 'pos\_enc', (1, 768, 14, 14), False, False]],
    [[2, 'msa', 'msa', (1, 768, 14, 14), False, False]],
    [[3, 'ln', 'ln', (1, 768), False, False]],
    [[4, 'linear', 'linear', (768, 3072), False, False]],
    [[5, 'sum', 'sum', None, False, False]],
    [[6, 'linear', 'linear', (3072, 768), False, False]],
]
edges embedding:
tensor([
    [0, 1, 1],  # input -> pos\_enc
    [1, 2, 1],  # pos\_enc -> msa
    [2, 3, 1],  # msa -> ln
    [3, 4, 1],  # ln -> linear
    [4, 5, 1],  # linear -> sum
    [5, 6, 1],  # sum -> linear
])

Hopefully, this simple example helps clarify the concept.

Q2

Why does TAL+ perform worse than TAL in some cases.

A2

The 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.

In 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.

Q3

In Formula 6, Laux needs to be further explained.

A3

The 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.

Q4

What is the difference between the Decoder and Encoder in Figure 2?

A4

We 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.

The 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.

Q5

How much additional overhead does the two-stage training incur compared to other methods?

A5

Table 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.

Q6

How is an entire ViT initialized?

A6

Our 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.

We thank you for your reviews and address your concerns as follows.

Q1

An important Task-Aware Parameter initialization baseline is not compared: Learning...(TIP 2024). The models used in this paper is too small and shallow.

A1

Thank 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.

Q2

There is limited discussion on the scalability of TAL.

A2

Our 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.

Q3

The theoretical underpinning of the method is not fully explored.

A3

We discuss a simplified case of our TAL method and provide a theoretical derivation. We reach the following conclusions:

• Convergence: Gradient-based optimization of hypernetworks converges to stationary points under standard smoothness assumptions.

• Approximation Capability: Sufficiently expressive hypernetworks can approximate optimal model parameters to arbitrary precision.

1. Problem Definition and Optimization Objective We define the following setup:

• 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)$.
• Model $M$ Also an MLP, using parameters $p$ to perform a binary (0,1) classification task and compute the loss $\mathcal{L}(p)$. The optimization objective is to train the hypernetwork $H$ to minimize the cross-entropy loss: $\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]$

2. Convergence Analysis

Theorem 1 (Convergence to Stationary Point): Assume the following conditions hold:

• The loss function $\mathcal{L}(p)$ is $\beta$-smooth
• Hypernetwork $H(\theta)$ is $L_H$-Lipschitz continuous
• The composed function $\mathcal{L}(H(\theta))$ has bounded gradients Then, using gradient descent with learning rate $\eta < \frac{2}{L_H\beta}$, after $T$ iterations: $\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}$

Proof:
By $\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}$: $\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$ Summing over $t=0,1,...,T-1$ and rearranging: $\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)}$ $\leq \frac{\mathcal{L}(H(\theta_0)) - \mathcal{L}(H(\theta^*))}{\eta\left(1 - \frac{L_H\beta\eta}{2}\right)}$ Since $\eta < \frac{2}{L_H\beta}$ implies $1 - \frac{L_H\beta\eta}{2} > 0$, and using the minimum gradient norm: $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$ $\leq \frac{\mathcal{L}(H(\theta_0)) - \mathcal{L}(H(\theta^*))}{\eta\left(1 - \frac{L_H\beta\eta}{2}\right)}$ With proper learning rate, $1 - \frac{L_H\beta\eta}{2} \geq \frac{1}{2}$, resulting in: $\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}$ This shows that as $T \to \infty$, the gradient norm approaches zero, indicating convergence to a stationary point. $\square$

3. Optimality Analysis

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: $\||H(\theta) - p^*\|| < \delta$

Proof:
According 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$

We thank you for your reviews and address your concerns as follows.

Q1

Were 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.

A1

TAL and GHN-3, LoGAH do use the same model training dataset, ViTs-1K, which ensures a fair comparison between them. The 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. Beyond 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.

Q2

Lack of Supporting Evidence for Convergence Claims.

A2

We 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.

Table: Training Loss Comparison of ViT-small with Different Initialization Methods on UCF Dataset Across Epochs

Q3

Limited Validation of Untrained Initialization Effectiveness.

A3

We 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.

Q4

Incomplete Comparison with Task Information.

A4

We 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.

Table: Performance of untrained models initialized with GHN-3 and GHN-3 with Task Initialization (GHN-3 w/ T.I.)

Q5

Insufficient Motivation for the Ancestry Model.

A5

The 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.

Q6

Ambiguity in "Learngene" Concept.

A6

The 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.

We thank you for your reviews and address your concerns as follows.

Q1

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 "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.

A1

Our 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.

Q2

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.

A2

Our 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.