Initializing Variable-sized Vision Transformers from Learngene with Learnable Transformation

[Presentation]

1) The description of transformations and notations.

Thank you for pointing this out!

In the revision, we will simplify the notations of Section 3 and highlight the core idea behind the two types of transformations.

2) How to choose the start and step size in "step-wise selection".

In practice, we usually choose the starting row as 1 and calculate the step sizes as $\lfloor\frac{D_{in}^{(aux)}}{D_{in}^{(des)}}\rfloor$ and $\lfloor\frac{D_{out}^{(aux)}}{D_{out}^{(des)}}\rfloor$.

3) Difference between "random selection" and "random selection (wo)".

For random selection, we first randomly generate a set of indices, based on which we select corresponding rows/columns. During the selection process, we ensure that selecting the same rows/columns for all transformation matrices.

In the case of random selection (wo), as opposed to random selection, we opt to select distinct rows/columns for each transformation matrice (L238-239).

4) Caption of Figure 3.

In the revision, we will add more descriptions in the caption of Figure 3 to avoid any misinterpretation.

[Models larger than Aux-Net]

Please see the relevant discussions in G1 of General response.

[Comparisons between LeTra and baselines in training costs]

Thank you for pointing this out!

Compared to Scratch-1, LeTra reduces around 3.6× total training costs (12$\times$100 epochs versus 320+12$\times$1 epochs), where "320" means the training epochs of stage 1.

Additionally, in comparison to Scratch-1, we calculate GPU hours for training 12 Des-Nets and find that LeTra significantly reduces total training GPU hours by approximately 7× (1322 GPU hours for Scratch-1 versus 190 GPU hours for LeTra).

[More choices of transformation selection]

Thank you for your insightful question!

We rank the weights by magnitude and select the largest rows/columns as the target transformation matrices, a strategy we refer to as "rank". According to the table below, we observe that "rank" achieves performance comparable to LeTra (continuous selection), thereby validating the robustness of our trained transformation matrices.

Dear Reviewer p972,

Thank you very much for your constructive comments on our work.

We have made every effort to address the concerns raised. If any aspect of our response is unclear or requires further clarification, please let us know. If everything is clear, we kindly ask if you could consider improving the score.

[Presentation]

1) Simplification of notations in Sec 3.1.

Thank you for pointing this out! In the revision, we will simplify the notations of Sec 3.1 to describe the method more clearly.

2) Initialization process in Figure 3.

Thank you for raising this confusion!

During the first training stage of LeTra, we construct the Aux-Net according to Eq.(1).

For the second initialization stage of LeTra, we first select certain rows/columns from well-trained $F_{l,in}$ to form target matrices $F_{l,in}^{des}$. Subsequently, we use $F_{l,in}^{des}$ to perform matrix multiplication with learngene matrices. The multiplication results are then inserted into original learngene matrices to derive $W_{l}^{'}$.

Similarly, we select certain rows/columns from well-trained $F_{l,out}$ to construct target matrices $F_{l,out}^{des}$. Then we use $F_{l,out}^{des}$ to multiply $W_{l}^{'}$ and insert the multiplication results into $W_{l}^{'}$ to obtain $W_{l}^{des}$.

Thus, the first training stage and the second initialization stage of LeTra are not inherently interconnected. In the revision, we will add these descriptions in the caption of Figure 3 to avoid any misinterpretation.

[Experimental results]

1) Difference between Scratch-1 and Scratch-2.

Scratch-1 involves training models from scratch with sizes identical to those used in LeTra (L544-545). Scratch-2 entails training models from scratch with sizes similar to those in Scratch-1 because these models only vary in depth (L546-547). Additional model details can be found in Table 5.

2) Teacher network.

We choose LeViT-384 [48] as the teacher/ancestry model (L536-537).

3) Comparison with IMwLM [36] and Matformer [i].

We employ DeiT-base distilled as the larger pretrained model for IMwLM [36]. For Matformer [i], we refer to the results presented in Figure 4(a) of their original paper. From Figure I of our uploaded PDF, compared to IMwLM [36] which initializes smaller models from larger pretrained ones and Matformer [i], we observe that LeTra achieves superior performance. Notably, LeTra demonstrates the capability to initialize models whose sizes are independent of larger pretrained models.

[i] Kudugunta, Sneha, et al. "Matformer: Nested transformer for elastic inference.", NeurlPS 2023.

4) Comparison with scratch training in training costs.

Thank you for pointing this out! In comparison to Scratch-1, we calculate GPU hours for training 12 Des-Nets and find that LeTra significantly reduces total training GPU hours by approximately 7× (1322 GPU hours for Scratch-1 versus 190 GPU hours for LeTra).

5) Performance of Des-Nets without distillation in the first stage.

Firstly, it is important to note that we did not employ distillation during the fine-tuning of Des-Nets in the second stage.

Secondly, we omit the distillation process in the first stage and present the performance of Des-Nets in Figure II of our uploaded PDF.

While the distillation process in the first stage marginally enhances the final performance of Des-Nets, the most substantial performance improvement arises from our proposed initialization process using LeTra.

6) Performance of Des-Nets which are smaller than learngene.

Please see the relevant discussions in G1 of General response.

Dear Reviewer 8CMe,

We would like to express our sincere gratitude for your thoughtful and thorough comments.

We have made every effort to address the concerns raised. If any part of our response remains unclear or requires further clarification, please let us know.

1) Performance improvements with increased depth.

Thank you for pointing this out!

In Table 4, we present the results of LeTra with 2-epoch tuning (L315), which aims to demonstrate the effectiveness of our proposed depth transformation rather than performance improvements with increased depth.

Interestingly, during the early training phase (e.g., 2 epochs), model performance does not consistently improve with increased depth, whether using Scratch that trains randomly-initialized models or LeTra that trains learngene-initialized models. To validate this claim, we provide the results of Scratch with 5-epoch training and LeTra with 2-epoch tuning below.

Furthermore, extending the training of LeTra-initialized models to more epochs reveals a consistent improvement in model performance with increased depth.

2) Contribution of distillation to the final performance.

Thank you for your nice concern!

We remove the distillation process in the first stage and present the performance of LeTra in Figure II of our uploaded PDF.

Comparing "LeTra (without distillation in first stage) (5 epoch)" with "LeTra (with distillation in first stage) (5 epoch)", we observe that the distillation process in the first stage has a marginal improvement on the performance of Des-Nets.

Additionally, comparing "LeTra (without distillation in first stage) (5 epoch)" with "Scratch (5 epoch)", we find that LeTra's proposed initialization process significantly enhances Des-Nets' performance.

Consequently, we can safely conclude that the most substantial performance gains stem from our proposed initialization process using LeTra rather than the distillation process in the first stage.

Furthermore, it is important to note that we did not employ distillation during the fine-tuning of Des-Nets in the second stage.

3) Explanation of Table 1.

We acknowledge the necessity of retraining a specific task head when transferring well-trained parameters from task A to task B. However, in Table 1, both task A and task B for all baselines and LeTra involve ImageNet-1K. Therefore, we also inherit the classification head parameters from either the first stage or the pre-training stage to initialize the Des-Nets for ImageNet-1K.

4) Discussion about importance of depth and width for deep networks.

Thank you for your insightful question! While it is widely accepted that depth significantly impacts deep network design (as shown in the left table below, where model performance increases with depth), we empirically find that configuring width could also enhance model performance (as shown in the right table below). This discovery motivates us to explore transforming learngene across both depth and width dimension. Furthermore, recent studies (see [i]) emphasize the importance of simultaneously considering both width and depth dimensions in neural network design.

[i] "The shaped transformer: Attention models in the infinite depth-and-width limit.", NeurlPS 2023.

Dear Reviewer i78Y,

We greatly appreciate the concerns provided and have made every effort to address all the points raised.

Is there any unclear point in our response that needs further clarification?

Dear Reviewer i78Y,

The author's rebuttal discussion is ending soon. Can you respond to the authors' reply and see if it addresses your concerns or if you would maintain your rating? Thank you for your time and effort.

Best,

Your AC

1) Scaling Des-Net bigger than Aux-Net.

Please see the relevant discussions in G1 of General response.

2) Combination with other Learngene strategies.

Thank you for your insightful question!

We could combine LeTra with other Learngene strategies. For instance, we could replace the depth transformation strategy of LeTra with linear expansion strategy proposed by TLEG [23].

3) More complex downstream tasks.

Thank you for your nice concern!

In the revision, we are committed to expanding our experimental scope to include more diverse tasks and datasets.

4) Linear probing results.

As shown in the table below, we can observe that LeTra achieves better performance than Pre-Fin under the linear probing protocol. For example, LeTra outperforms Pre-Fin by 3.43% and 1.86% on CIFAR-100 and CIFAR-10 with Des-H12-L12.