Slicing Vision Transformer for Flexible Inference

We appreciate the Reviewer's approval and valuable comments. We respond to the Reviewer's concerns as below.

Weakness 1: Fixed scaling ratio.

Thank you for the comment. The hidden dimension of ViT has to be an integer multiple of the number of heads (e.g., 6/12) so ViT cannot support arbitrary width ratios like CNN. For instance, the theoretical maximum number of networks that ViT-B could represent is 25 if $s=0.25$. However, we do acknowledge that Scala still cannot generalize to unseen width ratios as this issue is essentially related to the architecture design of transformer and we constrain ourselves from modifying the conventional structure as it has been well-optimized in various scenarios. On the other hand, it further emphasizes that the problem we target to solve is very challenging as the intermediate subnets ($X=13$) are only trained for around 18 epochs and they are expected to match the performance of networks that are trained for 100 epochs.

Weakness 2: Results of 300 epochs.

Thanks for the advice and we will emphasize the 300-epoch results of Sec. 5.3 and 5.4 in the final version.

Question 1: Root cause.

Thanks for the great question. It is the smallest subnet that causes the issue no matter what the width is and we validate it by setting $s=0.5$ to be the smallest scaling ratio and removing Isolated Activation. The table shows that constantly activating the smallest subnet in the naive way still results in worse performance even if the width ratio is larger than 0.25.

Question 2: Noise Calibration.

Compared to standard knowledge distillation, the teacher of those subnets in Scala is dynamically changing (different widths) and will be optimized during training like the students as their weights are shared. On the one hand, the intermediate teachers serve as teacher assistants which fill the gap between $F^{l}\left (\cdot \right )$ and $F^{s}\left (\cdot \right )$ and provide simplified optimization objective for small subnets, in contrast to US-Net which utilizes $F^{l}\left (\cdot \right )$ as the teacher for all subnets and results in inferior performance. On the other hand, it also indicates that those teachers cannot provide faithful predictions at the early stages and it is essential to add Cross-Entropy loss to those subnets during training which is overlooked in previous baseline methods.

Question 3: Impact of only CE at early epochs.

Thanks for the valuable suggestion. We further conduct experiments by removing the distillation loss of subnets in the first 10 epochs (denoted by *) and there is only minor performance change in the subnets which suggests that Noise Calibration is indeed helpful at early stages.

We hope the explanations could address your concerns and we would appreciate it a lot if you could recognize the contributions of our work.

We appreciate the Reviewer's feedback. We provide further explanations to clarify the Reviewer's concerns based on several key points as below.

Weakness 1: Discussion with dynamic networks.

We thank the Reviewer for the valuable suggestion and we will add discussion with those dynamic networks in the final version, including MSDNet[1], RANets[2], GFNet[3], DVT[4], CF-ViT[5], Dyn-Perceiver[6], etc.

However, the differences between our work and this line of research are: (1) Motivation: dynamic networks are designed to reduce inference costs with dynamic computational graphs; Scala is proposed to make ViT become slimmable without touching the specific architecture design. (2) Method: dynamic networks tailor their computational graphs to accommodate varying inputs, thereby optimizing resource allocation on a per-sample basis. In contrast, Scala treats every sample equally and modulates the computational costs by adjusting the width ratios which aligns with the inherit design of ViT and its inference costs can be explicitly controlled by a single hyperparameter $r$.

Essentially, dynamic networks are agnostic with Scala and we will consider the integration of those two lines of research as our future plan.

Weakness 2: Lack of novelty.

Thank you for the comment. While we do agree that 'smaller ViTs are intrinsically the sub-networks of a larger ViT with different widths' is not a surprising observation, making ViTs become slimmable is indeed non-trivial.

According to our analysis in Sec. 2, ViTs display minimal interpolation ability which emphasizes that the problem we target to solve is very challenging as the intermediate subnets ($X=13$) are only trained for around 18 epochs and they are expected to match the performance of networks that are trained for 100 epochs.

To the best of our knowledge, we are the first to propose Isolated Activation, which seems simple but is very effective as we have abundant analysis (Sec. 3) to support it. Moreover, the knowledge distillation in Scala is different from the conventional setting. Instead of choosing the full model as the teacher for all subnets, we enable the intermediate networks to be the teacher as well which fills the model gap between $F^{l}\left (\cdot \right )$ and $F^{s}\left (\cdot \right )$ and provide simplified optimization objective for small subnets. This is non-trivial as we find adding a teacher network to the baseline method US-Net will result in performance decreases in the smaller subnets (please see our reply to Reviewer rLHb) which further supports our motivation to propose Progressive Knowledge Transfer. Moreover, Noise Calibration is critical in our setting as well because those intermediate teachers cannot provide faithful predictions at the early stages and it is essential to add Cross-Entropy loss to those subnets during training which is overlooked in previous baseline methods.

Based on these factors, Scala can solve this challenging problem in a simple but effective way.

We hope the explanations could address your concerns and we would appreciate it a lot if you could recognize the contributions of our work.

[1] Huang G, Chen D, Li T, et al. Multi-scale dense networks for resource efficient image classification[J]. arXiv preprint arXiv:1703.09844, 2017.
[2] Yang L, Han Y, Chen X, et al. Resolution adaptive networks for efficient inference[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2020: 2369-2378.
[3] Wang Y, Lv K, Huang R, et al. Glance and focus: a dynamic approach to reducing spatial redundancy in image classification[J]. Advances in Neural Information Processing Systems, 2020, 33: 2432-2444.
[4] Wang Y, Huang R, Song S, et al. Not all images are worth 16x16 words: Dynamic transformers for efficient image recognition[J]. Advances in neural information processing systems, 2021, 34: 11960-11973.
[5] Chen M, Lin M, Li K, et al. Cf-vit: A general coarse-to-fine method for vision transformer[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2023, 37(6): 7042-7052.
[6] Han Y, Han D, Liu Z, et al. Dynamic perceiver for efficient visual recognition[C]//Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023: 5992-6002.

Dear reviewer, please have a look at the rebuttal and indicate whether it addressed some of your concerns.

I highly appreciate the authors' response. However, I still think the lack of discussions and comparisons with dynamic networks in the original submission is an important limitation. Moreover, although the authors have provided some clarification, Isolated Activation seems like simply an engineering trick, and the proposed knowledge distillation method seems too incremental compared to existing works. These techniques may be insufficient as significant scientific contributions.

I increase my rating to 4. However, I still vote to reject this paper. The major reasons are (1) lack of discussions on important related works, and (2) weak technical novelty.

We appreciate the Reviewer's comments to point out the confusing description and we make the response as below.

Weakness 1: Phrase.

Thank you for your insightful suggestion. We will revise our presentation in the final version to prevent any potential misunderstanding. Our motivation for adopting the term "scalable representation learning" stems from the observation that different methods (e.g., DeiT, MAE, DINOv2) learn distinct representations that are not inherently slimmable (see Appendix: A.3). Scala is proposed to learn the original representations while incorporating slimmability. Therefore, we initially chose "scalable representation learning" to describe our method. However, we acknowledge that the term "scalable" may lead to confusion, and we will revise it in the final version to enhance clarity.

We hope the explanations could address your concerns and we would appreciate it a lot if you could recognize the contributions of our work.

We sincerely appreciate the Reviewer’s detailed comments and constructive suggestions for us to improve our work. We make the response as below.

Weakness 1: Phrase.

Thanks for the great suggestion and we will modify the presentation in our final version. Our motivation for adopting 'scalable representation learning' is that different methods (DeiT, MAE, DINOv2) will learn different representations, which are not naturally slimmable (see Sec A.3). Scala is proposed to learn the original representations but with the slimmable ability, so we adopt 'scalable representation learning' to describe our method. However, we do agree that the term 'scalable' may result in misunderstanding and we will modify it in the final version.

Weakness 2: Limited baselines.

Thanks for the comment.

SortedNet scales the network through multiple directions and results in irregularities in transformer architectures, compared to Scala which aligns with the inherent design of ViT to vary from widths. As their vision experiments are conducted on small-scale CIFAR datasets, we intended to reproduce their method on ImageNet but did not find released implementations.

MatFormer only scales the FFN block in transformer so it offers a minor computational adjustment space and we adapt their method (written in JAX) on DeiT-S for a fair comparison. The results are shown in Fig.1 of the PDF file (reply to all reviewers) and Scala achieves comparable performance with it (better in most cases) when $s=0.5$ with a larger adjustment scope.

Weakness 3: Scope of experiments.

We thank the Reviewer for the valuable suggestion. Indeed, extending Scala to NLP and a larger network is part of our original plan but we are struggling with computation resources as both of the experiments demand much larger training costs. We will definitely consider it as our future plan.

Weakness 4: Teacher and training costs.

As explained in Weakness 1, Scala can be regarded as a tool to make existing methods become slimmable while trying to maintain the original representations. Although we can expect to learn similar representations of methods trained with supervised learning on ImageNet-1K (DeiT), some self-supervised learning methods (DINOv2) are trained with gigantic undisclosed data which is almost impossible to reproduce their performance without utilizing the pre-trained model. So we simply use the pre-trained network (same size as the student) as the teacher to imitate the original model to inherit the original representations, instead of utilizing a larger network which usually leads to better performance.

For US-Net, we do not introduce teacher networks following their original design but we do compare with Separate Training with Distillation in Sec A.8 (Tab. 10), where Scala still outperforms the baseline at smaller widths. Here we further re-implement US-Net by adding the teacher and the results are shown in the table below. Although obtaining a better full model, US-Net exhibits worse performance at smaller width ratios because it utilizes the full network as the teacher for all subnets and this phenomenon verifies our motivation of proposing Progressive Knowledge Transfer.

Assuming to deliver 13 models in the end, we compare the training time (100 Epoch) of Scala with US-Net, Separate Training, and Separate Training with Distillation on 8 A100 GPUs and we do not include the teacher training time as Scala is aimed to scale down existing works where the same size pre-trained model can be easily downloaded. The difference between US-Net and Scala is not large as the transformer architecture has been well-optimized on GPU and we do observe a significant time gap between Scala and ST/STD as they have to train 13 models iteratively. Moreover, Scala can be configured to deliver 25 models without an increase in training time as we sample 4 networks at each iteration in all scenarios which further highlights our strengths.

Weakness 5: Interpolation.

We did not fully address the interpolation problem as this issue is essentially related to the architecture design of transformer and we constrain ourselves from modifying the conventional structure as it has been well-optimized in various scenarios. On the other hand, it further emphasizes that the problem we target to solve is very challenging as the intermediate subnets ($X=13$) are only trained for around 18 epochs and they are expected to match the performance of networks that are trained for 100 epochs.

Thanks for carefully reading our paper and we will move Fast Interpolation into the main paper.

Weakness 6: DINOv2.

We follow the previous setting except for using DINOv2-B as the teacher which means we still train the model in a supervised manner. We have tried to adopt the self-supervised learning objective from DINOv2 but found the training cost too large, and it is almost impossible to reproduce the performance following the original learning objective with ImageNet-1K since DINOv2 is trained on the private dataset LVD-142M. However, we re-conduct the experiment by removing all the CE loss to alleviate dataset bias and find this choice results in surprising generalization ability by inheriting the fruitful knowledge of DINOv2 (please see the reply to all reviewers).

Weakness 7: Ablation.

Thanks for the advice and we will add more ablation in our final version.

We hope the explanations could address your concerns and we would appreciate it a lot if you could recognize the contributions of our work.

Dear authors, thank you for your rebuttal. FYI: after the author discussion phase closed, there is still a bit over a week of discussion between ACs and Reviewers.