Learngene Tells You How to Customize: Task-Aware Parameter Prediction at Flexible Scales

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

Q 1:

The related works need improvement. Clearly stating the differences and contributions make the paper review evaluation better.

A 1:

We supplement this in the revised version on page 9 and 10. The added content as follows:

Parameter Prediction: Our task-aware learngene (TAL) incorporates modules to process task information, enabling a single TAL model to customize models of varying scales for different tasks.

Learngene: In task-aware learngene (TAL), learngene is no longer a sub block of the ancestry model, the encoder part of TAL model is regarded as learngene. Model generation becomes a process of encoding and decoding, with the ancestry model and multi-task common knowledge injected through the learngene. TAL model can handle both model and task information, initializing models of flexible scales for different tasks.

Q 2:

The baselines are limited. Comparing the latest one is good but diverse baselines from different aspects make the results stronger.

A 2:

We have added relevant experiments for the new baseline GHN-3 in the revised version on page 6.

Q 3:

The datasets can be somewhat out-of-dated. Considering using challenging datasets would make the results and conclusions stronger, especially in an era of foundation models.

A 3:

We train TAL-GPT2 on NLP tasks using the same approach employed for vision tasks, enabling parameter prediction for GPT-2 across various scales tailored to different tasks. Currently, the TAL-GPT2 model is still under training.

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

Q 1:

Limited contribution: It does look like a minor incremental work over existing hypernetwork methods GHN-2, LoGAH where in the main different is adding task-specific layers/information.

A 1:

Our TAL method is based on the Learngene framework and attempts to leverage large-scale pretrained model's knowledge to create descendant models tailored for downstream tasks, which differs from graph hypernetworks (GHN). In terms of method, we incorporate modules to process task information, enabling a single TAL model to customize models of varying scales for different tasks.

Q 2:

It is not very clear from the paper what is the specific "task-specific" information that is added to the learngene that helps improve the downstream model performance.

A 2:

We explained task information in Section 3.2, which is generated by the ancestry model through the average feature extraction of the task images.

Q 3:

Explanation of results/experiments: Table 2 and Table 3 has lot of experiments comparing TAL with LoGAH and showing several cases of improvements. The improvements range from slight to large improvements and in some cases negative improvements. Authors can dig deeper into what's the underlying reasoning behind this improvement/regressions.

A 3:

In our experiments (see Table 2 and Table 3), the TAL method achieves significantly higher average accuracy across multiple tasks compared to the baseline methods. Specifically, with initialized parameters, TAL outperforms the comparison methods on 7 out of 9 tasks and across all tested model sizes. A common challenge in multi-task training is that some tasks struggle to surpass the performance of single-task training[1-3]. Here are some intuitive analyses to this situation with TAL:

a. Reasons for Performance Regressions: TAL employs a multi-task training approach that, while enabling knowledge sharing, can sometimes lead to negative transfer between tasks. Specifically, for tasks highly similar to ImageNet in terms of data distribution and features, LoGAH v2 optimizes model parameters more effectively through single-task distillation. This targeted optimization allows LoGAH v2 to perform better on these similar tasks. In contrast, TAL’s multi-task setup does not sufficiently optimize these similar tasks, resulting in inferior performance compared to LoGAH v2 in certain cases.

b. Reasons for Performance Improvement: For tasks that show significant performance gains, these tasks differ substantially from ImageNet. Consequently, LoGAH v2’s single-task tuning approach is less effective. TAL leverages multi-task sharing to effectively utilize shared knowledge across tasks and employs task embedding mechanisms to filter and inject task-specific knowledge.

Q 4:

Shallow Analysis of Catastrophic Forgetting: The paper mentions TAL addresses catastrophic forgetting, but it lacks specific experiments or metrics to evaluate this claim. There is little evidence showing that TAL explicitly mitigates forgetting across tasks in sequential or continual learning contexts.

A 4:

Our work does not aim to address catastrophic forgetting. TAL, is designed to customize models of varying sizes for different tasks. Specifically, the descendant models predicted by TAL are for single tasks. In the TAL model, we added specific modules processes task information, which successfully distinguishes different tasks effectively when handling multiple tasks, and there is no catastrophic forgetting problem.

Q 5:

More details around the TSL and Task HyperNet layers. What's the specific input to this? How does this help achieve better representation of the computation graph, etc.

A 5:

We introduces TSL is a simple MLP, and Task HyperNet layers are specifically introduced in Formula 2 of Section 3.2 of the original paper. To make it easier to understand their functions, in the revised version we added a new formula in Section 3.2 on page 4. The added content as follows:

The task-specific layers apply these bias parameters to the model's computation graph using the following formula:

$

f^G_{\tau} = \gamma_{\tau} \times f^G + \beta_{\tau}

$

References

[1] Naveen Arivazhagan, Ankur Bapna, Orhan Firat, Dmitry Lepikhin, Melvin Johnson, Maxim Krikun, Mia Xu Chen, Yuan Cao, George Foster, Colin Cherry, et al. Massively multilingual neural machine translation in the wild: Findings and challenges. arXiv preprint arXiv: 1907.05019, 2019.

[2] Rabeeh Karimi Mahabadi, Sebastian Ruder, Mostafa Dehghani, and James Henderson. Parameter-efficient multi-task fine-tuning for transformers via shared hypernetworks. arXiv preprint arXiv:2106.04489, 2021.

[3]Vandenhende, Simon, et al. "Multi-task learning for dense prediction tasks: A survey." IEEE transactions on pattern analysis and machine intelligence 44.7 (2021): 3614-3633.

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

Q 1:

The abstract mentions reducing serving costs and latency, but the experiments lack clarity on how the TAL framework actually predicts a smaller parameter set to create a more compact model.

A 1:

We provide a detailed description of how TAL predicts model parameters of various scales for different tasks in Section 3 (Methods), so we do not repeat this in Section 4 (Experiments).

Q 2:

The concept of "learngene" remains unclear and appears insufficiently differentiated from the standard pretraining-then-finetuning approach.

A 2:

We introduce the Learngene concept comprehensively in Section 1 (Introduction). Instead of fine-tuning the entire pretrained model, Learngene inherits a compact component, known as learngene, from the pretrained model to create descendant models of various sizes tailored for different tasks. In Section 3 (Method), we explain that in TAL, learngene is the encoder of the TAL model and how TAL predicts model parameters of various scales for different tasks. Additionally, in Section 5 (Related Work), we detail the existing works of Learngene.

Q 3:

The background on GHN is unclear. Lines 128–132 seem just a general training process. Besides, GHN is trained on M different architectures and N data samples, but the nature of these training data samples is not explained.

A 3:

We provide a introduction to Graph HyperNetworks (GHN) in Section 1 (Introduction). In section 2 (Preliminaries), we formulate the GHN training process in details. In the datasets subsection of Section 4.1 (Experimental Setup), we describe the model datasets (referring to M different architectures) and image datasets (referring to N data samples) used to train the TAL model in our experiments. Additionally, in Section 5 (Related Work), we detail the existing works of GHN.

Q 4:

The paper makes a significant claim about handling different model scales in the introduction, but there is little detail or experimental evidence on this. It remains unclear how TAL predicts or adjusts parameters for models of varying scales.

A 4:

In all our experiments, we initialized models of various sizes using TAL and provided detailed results demonstrating TAL's ability to predict parameters for different scales.

Q 5:

The experimental results need a clearer, more structured presentation. For instance, starting with an overview of the experiments would provide context, followed by details on the setup, datasets, and baselines. Most importantly, each experiment should include an explanation of its objectives and expected outcomes, helping to guide readers through the data. This approach would make it easier to interpret the results and understand how the numbers and tables support the authors' claims. Besides, the abbreviation should be used properly.

A 5:

We added an overview of the experiments in the revised version on page 5 and 6. As for the details on the setup, datasets, baselines, objectives and expected outcomes for each experiment, and the meanings of various abbreviations, we clearly provide them in the section 4 (Experiments) of the original paper.

Q 6:

The term "Computation Graphs" is introduced but confusing me. Can the authors provide examples?

A 6:

Our paper utilizes computational graphs based on those used in GHN-3 and LoGAH methods. The computation graphs primarily include node_feat, node_info, and edges_embedding. To clarify, we provide an example using the Vision Transformer (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.

Q 7:

Training TAL on a Large Dataset: It is unclear what a "large dataset" here means. Does this refer to a pretraining dataset containing diverse domains and concepts, or something else?

A 7:

The "large dataset" mentioned in Section 3 refers to a big dataset in a certain field. The TAL model is trained on this dataset to leverage substantial knowledge from the pretrained model. Specifically, in our experiments, we used ImageNet-1K as the large dataset.

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

Q 1:

While the TAL mechanism improves model flexibility and performance, it also significantly increases the computational and resource costs in multi-task scenarios.

A 1:

All experiments are run on an NVIDIA RTX 4090. The table below shows the training time for each method in Section 4.2 of the original paper using a single RTX 4090. The modest increase in TAL's training time is acceptable.

Q 2:

The comparative experiments in the paper mainly focus on vision tasks, lacking experiments in other domains. Demonstrating similar performance improvements in other fields would enhance the generalizability and persuasiveness of the method.

A 2:

We train TAL-GPT2 on NLP tasks using the same approach employed for vision tasks, enabling parameter prediction for GPT-2 across various scales tailored to different tasks. Currently, the TAL-GPT2 model is still under training.

Q 3:

The paper lacks an in-depth discussion on the setting of sampling weights in multi-task training.

A 3:

The sampling weight w_{\tau}\ in section 3.2 of the original paper is determined by the sampling probability p_{\tau}\ described in Section 4.1. Specifically, the sampling weight w_{\tau} \ is directly proportional to the task sampling probability p_{\tau}\ , which indicates each task's likelihood of being selected during training[1]. This ensures that tasks with higher p_{\tau}\ have a greater impact on model updates, allowing the TAL model to effectively learn each task's characteristics.

Q 4:

The method and experiments are inconsitent in motivation. Specifically, the method aims to enhance the expression of gene data, which are more likely to be sequence and graph structure, while the evaluation focus on images that are not corresponds to structure modeling.

A 4:

Our TAL method is based on the Learngene framework which inspired by biological gene inheritance, focuses on extracting compact components, known as learngene, from a large-scale pretrained model (referred to as ancestry model) and using them to initialize descendant models. In TAL, learngene is the encoder of the TAL model. Our goal is to leverage the knowledge of ancestral models to create descendant models tailored for downstream tasks. As shown in Figure 2 of the original paper, our approach does not involve "gene data". Instead, we use model computation graphs as input to the TAL model and train it on various image datasets. This enables the TAL model to customize models of flexible scales for different tasks.

Q 5:

The core technique is derived from graph hypernetworks directly. However, it lacks of the novel contribution for customize graph hypernetwork for the main task in this work.

A 5:

Our TAL method is based on the Learngene framework and attempts to leverage large-scale pretrained model's knowledge to create descendant models tailored for downstream tasks, which differs from graph hypernetworks (GHN). In terms of method, we incorporate modules to process task information, enabling a single TAL model to customize models of varying scales for different tasks.

Q 6:

Can TAL mitigate the inherent conflicts among different tasks?

A 6:

Yes, TAL effectively mitigates inherent conflicts among tasks. As detailed in Section 3, we designed task-specific layers to handle each task uniquely. Table 2 and table 3 demonstrate that TAL significantly improves multi-task performance compared to the previous single-task LoGAH training. Additionally, the visualization in Section 4.3 shows that the learngene outputs of TAL models at various scales cluster by task, indicating that TAL successfully reduces conflicts between tasks.

References

[1] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1–67, 2020.