Slicing Vision Transformer for Flexible Inference

Summary

The paper targets scaling down Vision Transformers (ViT) to fit environments with dynamically changing resource constraints. The authors propose Scala, a framework enabling a single network to represent multiple smaller ViTs with flexible inference capability by activating various subnets during training. Scala introduces Isolated Activation to disentangle the smallest sub-network and uses Scale Coordination to provide stable and accurate learning objectives. Empirical validations on different tasks show that Scala achieves scalable representation with one-shot training, matching the performance of Separate Training without modifying the original ViT structure. Scala demonstrates an average improvement of 1.6% on ImageNet-1K compared to previous methods, using fewer parameters.

Strengths

1. The problem is important in practice. 2. The experimental results seem decent.

Weaknesses

1. My major concern is that, the same aim of adapting ViTs to dynamically changing resource constraints, can also be achieved by multi-exit networks, e.g., [*1, *2, *3]. However, the paper does not discuss these highly relevant works or compare with them. Hence, I vote for rejection. 2. The method seems to lack novelty. 'smaller ViTs are intrinsically the sub-networks of a larger ViT with different widths' is not a surprising observation. The key techniques (e.g., Isolated Activation and Knowledge Distillation) are not new (naive or have been widely adopted). [*1] Huang, Gao, et al. "Multi-Scale Dense Networks for Resource Efficient Image Classification." International Conference on Learning Representations. 2018. [*2] Wang, Yulin, et al. "Not all images are worth 16x16 words: Dynamic transformers for efficient image recognition." Advances in neural information processing systems 34 (2021): 11960-11973. [*3] Han, Yizeng, et al. "Dynamic perceiver for efficient visual recognition." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

Questions

Please refer to Weaknesses.

Rating

4

Confidence

5

Soundness

2

Presentation

2

Contribution

2

Limitations

The authors have addressed the limitations and potential negative societal impacts of their work.

Summary

The paper presents Scala, a novel framework for scalable representation learning developed from US-Net. It identifies the issues of directly applying US-Net to ViTs and proposes solutions including Isolated Activation, Scale Coordination, and Stable Sampling. These innovations enable Scala to output several sub-networks in one-shot learning. Extensive experiments on various network architectures and datasets demonstrate that the sub-networks produced by Scala consistently outperform those generated by separate training, with significantly reduced training time.

Strengths

Originality: Scala addresses the limitations of US-Net and successfully applies the concept of scaling to ViT backbones. This is a significant step in the adaptation of scaling methods for more complex network architectures. Quality: The paper supports its claims with extensive experimental results, providing strong evidence for the effectiveness of Scala. Clarity: The paper is clearly written and well-organized, making it accessible and easy to follow. Significance: Scala has the potential to influence future research directions in scaling ViTs.

Weaknesses

Originality: The novelty of Scala is somewhat constrained. For instance, Noise Calibration does not show a distinct difference from standard knowledge distillation. Essentially, Scala integrates US-Net with an alternative activation for the smallest subnet and fixed scaling ratios. Quality: The authors might consider emphasizing results from a more standard 300-epoch ViT training schedule to align with common practices in the field. Clarity: No further issues. Significance: The challenge of scaling ViTs with arbitrary ratios remains unresolved.

Questions

1. Regarding the issue shown in Fig. 4, is it always the smallest subnet that causes the issue, or does it occur with subnets having a scaling ratio near 0.25? 2. Can the authors clarify any differences between Noise Calibration and standard knowledge distillation? 3. What would be the impact if the distillation part were discarded and only Cross-Entropy (CE) loss were used for the initial epochs?

Rating

7

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

The novelty and significance of Scala are somewhat limited, as discussed in the weaknesses section. However, the extensive experimental results provide a robust foundation for the claims made in the paper. Overall, the work is well-executed and makes a valuable contribution to the field, justifying a recommendation for weak acceptance.

Summary

The paper introduces Scala, a novel framework designed to effectively scale down Vision Transformers (ViTs) for use in environments with fluctuating resource constraints. The key insight is that smaller ViTs can function as sub-networks within a larger ViT, differing mainly in width. Scala enables a singular network architecture that can emulate multiple smaller ViTs, thereby offering versatile inference capabilities while maintaining the structural principles of ViTs. The framework uniquely incorporates multiple sub-networks during its training phase, utilizes Isolated Activation to differentiate the smallest sub-network, and implements Scale Coordination to streamline the learning objectives for each sub-network, aiming for simplicity, stability, and accuracy. The empirical results across various tasks confirm that Scala can learn scalable representations efficiently with a single training iteration, maintaining the integrity of the original ViT architecture and achieving performance on par with networks trained separately.

Strengths

The proposed Scala framework aims to enhance Vision Transformers (ViTs) by enabling them to learn scalable representations suitable for flexible inference. This is achieved through two key innovations: Isolated Activation, which effectively disentangles the representation of the smallest subnet to maintain clarity and specificity, and Scale Coordination, which ensures that each subnet within the larger network receives simplified, consistent, and accurate signals. These mechanisms are designed to optimize the performance and scalability of ViTs, addressing common challenges in adapting these architectures to varied and dynamic operational contexts.

Weaknesses

1. Recent papers[1,2,3] with "Scalable" usually scale ViT to billion size with large scale datasets like DFN, JFT, and Datacomp. Therefore, I suggest authors should reconsider if the experiments can support "Scalable". [1] Zhai, X., Kolesnikov, A., Houlsby, N., & Beyer, L. (2022). Scaling vision transformers. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 12104-12113). [2] El-Nouby, A., Klein, M., Zhai, S., Bautista, M. A., Toshev, A., Shankar, V., ... & Joulin, A. (2024). Scalable pre-training of large autoregressive image models. arXiv preprint arXiv:2401.08541. [3] Dehghani, M., Djolonga, J., Mustafa, B., Padlewski, P., Heek, J., Gilmer, J., ... & Houlsby, N. (2023, July). Scaling vision transformers to 22 billion parameters. In International Conference on Machine Learning (pp. 7480-7512). PMLR.

Questions

please refer the weakness.

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

This paper addresses the limitations in conclusion

Summary

This paper advances an approach for training Vision Transfomers (ViTs) such that at inference time they can be dynamically adjusted to fit different budget constraints with reduced drops of performance. To this end, the authors introduce Scala, a framework that allows a single network to encapsulate and train simultaneously multiple sub-networks of different capacities and widths. The methodological backbone of this work are the Universally slimmable networks (US-Net) [37], originally devised for CNNs. The authors identify and analyze a few flaws of US-Nets: difficulty to generalize to ViTs, small interpolation and extrapolation ability to sub-network size unseen during training, impact of sustained activation of the smallest sub-network that coupled with the sandwich rule for selecting sub-networks during training leads to an over-emphasis on it at the expense of the other sub-networks. The authors propose two simple strategies towards such a method for ViTs: (i) Isolated activation that separates the smallest sub-network from the other sub-networks; (ii) scale coordination consisting of a set of heuristics to ensure that each sub-network gets simple, accurate and stable learning objectives: (a) progressive knowledge transfer from larger networks to smaller ones in gradual decrease of capacity, (b) stable sampling of intermediate width ratios to avoid large variations in capacities in the sandwich, (c) noise calibration, essentially a composite loss of supervised cross-entropy and distillation from the bigger sub-network. Scala is evaluated on several settings on the ImageNet-1k dataset with ViT-Ti/S/B, hybrid CNN-ViT architectures, lightweight networks, but also for dense prediction on semantic segmentation and self-supervised pre-training with interesting results. The baselines used here were Separate Training, Autoformer and US-Net.

Strengths

### Significance - the paper deals with a challenging and useful task for deploying ViT models into different operational settings with different computational constraints without retraining or distilling specific architectures each time - although a computational overhead is expected for such methods, the main components of Scala are relatively simple and make sense - Scala achieves good performance with a higher boost in the low parameter regime ### Originality - the proposed contributions are somehow incremental as they are improving the US-Net prior work, but do have some novelty and they are simple. ### Clarity - in general this work is well argued and easy to follow. The authors construct well the arguments regarding the challenges when going from CNNs to ViT with US-Net and how to construct their Scala approach. ### Quality - the paper offers several experiments and studies in the main paper and in the appendix (longer training, fast interpolation, ablation of components) that are well thought and improve the understanding of the method. - I appreciate the experiments beyond image classification, on semantic segmentation, as well as the self-supervised pretraining and subsequent linear probing on a downstream task.

Weaknesses

### "Scalable" naming - I think that the framing of the method as _"scalable representation learning"_ is quite confusing as it's not representative for this task, it's not a name used by other related works. Importantly, it can be easily mistaken with most works that use "scalable" for depicting the ability/property of a system (method, architecture) to handle a growing amount of data, parameters, and the potential to accommodate this growth. In other words "scalable" is rather used for depicting scaling up, whereas this work depicts the property of the proposed approach to accommodate sub-networks of different lower sizes/scales from the original. - maybe other names us in related works would be more appropriate here: slimmable, elastic, modular, flexibile inference, etc. ### Limited baselines and related work - some relevant related works dealing with tranformer networks are either just briefly mentioned, e.g., Matformer [18], or not mentioned at all, e.g., SortedNet [a], Early exit [b] - One of the main baselines, US-Net is originally designed for CNNs and, as the authors mentioned, moving to ViTs is not straightforward. Matformer is criticized for the limited number of models produced, but can be considered in the several experiments with X=4 sub-networks. Matformer and SortedNet could be included in the experimental evaluation ### Scope of experiments - While the authors considered several settings for computer vision tasks (image classification, segmentation, light architectures), transformer architectures are also encountered in NLP (as mentioned by the authors in L56). In such cases the original models can have much more parameters and elastic inference for lower computational budgets would be of high interest. - It would be useful to include an experiment from NLP in the style of those from Matformer or SortedNet. - The biggest architectures used here is a ViT-B (~86M params). Extending experiments to larger modern architectures would be definitely useful and interesting. ### Clarity - it's not always clear in the text and cost estimations that Scala needs a pre-trained full network as teacher for the distillation. This add some cost in compute and time in the end. Besides it's not clear whether US-Net also needs and uses a pre-trained teacher in the reported results. - in the intro, the authors mention that they address the issue of minimal interpolation ability of ViTs. Results from Table 2 show that the interpolation abilities of ViTs with Scala are still very low. However the fast interpolation strategy from $\S$A.2 is actually interesting for practical settings even though not fully solving this issue. It might be worth moving up in the main paper. - the idea of the transferability experiment ($\S$5.4) with DINOv2 is nice. From the description it is not clear whether DINOv2 was used as teacher for the distillation or also as supervised pre-training on ImageNet-1k? Or the pre-training on ImageNet-1K was done in a supervised manner as in previous experiments? - the ablation experiment from Table 6 is nice. However the presentation with removing one component at once offers only a partial understanding of the contributions of each module. Different configurations with different modules in on/off mode should give a better global understanding. **References:** [a] Valipour et al., SortedNet: A Scalable and Generalized Framework for Training Modular Deep Neural Networks, arXiv 2023 [b] Xin et al., Deebert: Dynamic early exiting for accelerating bert inference, ACL 2020

Questions

This paper takes an interesting direction of study: how to train ViTs such that a fine-grained elasticity in terms of sub-network sizes is possible at runtime. The proposed Scala approach is well described, makes sense and achieves good results in several computer vision settings. I do have a few concerns related to the phrasing of this type of works ("scalable representation learning") which can be confusing, the absence of larger architectures and of recent relevant baselines. My current rating is mildly positive (rather on the fence though) and I'm looking forward for the rebuttal. Here are a few questions and suggestions that could be potentially addressed in the rebuttal or in future versions of this work (please note that suggested experiments are not necessarily expected to be conducted for the short rebuttal period): 1. Please clarify the points raised in the clarity section: use of teacher model for Scala and US-Net, implementation of transferability experiment. 2. Comparison of training cost between Scala (including teacher training), US-Net and Separate Training baselines. 3. Add a discussion of differences and when possible experimental comparison with Matformer and SortedNet baselines on image classification or semantic segmentation. 4. Extension of experiments to NLP architectures and tasks in the style of SortedNed, Matformer

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

The authors addressed some of the limitations in the conclusion section.

Deadline Approaching

Dear Reviewers: Thank you so much for delivering your valuable comments to us. We have tried our best to provide a detailed response based on your concerns. We would like to kindly remind you that Aug 13th is the deadline for the discussion phase and your opinions regarding our response matter a lot to us. Thank you so much for helping us improve our paper. Sincerely, Authors

Post-rebuttal comments

I would like to thank the authors for the detailed and informative rebuttal. I imagine they invested significant effort into clarifying the concerns from the 4 reviewers and I'm confident they will improve this work. I've read the other reviews and the responses of the authors to them and skimmed through the paper again. The strong points I see in this rebuttal are: - several clarifications around the implementation, in particular on the use of pre-trained teacher across settings: both on the regular setting but also on the transferability task where DINOv2 was used as teacher and not as a self-supervised objective - the addition of a comparison with a more recent and related work, Matformer. Pity that the code of SortedNet is not available. - more detailed information about training costs w.r.t. to the main baseline US-Net, that originally does not benefit from a pretrained teacher. The authors report results for US-Net + KD in the rebuttal and it still under-performs compared to Scala. - new results using DINOv2 as teacher without cross-entropy loss showing the effectiveness of the option of using foundation models for this setting. On the downside the evaluation is still limited to mid- to small-sized architectures (the biggest architecture used here is a ViT-B with 86M parameters) and to computer vision settings. I can relate to the argument of the authors regarding computational cost. However given their results with DINOv2, one may devise a similar strategy for distilling a NLP foundation model to a different task than next-token prediction, e.g., text-classification. Hopefully the authors will find a solution in this direction for the paper update. **Wrap-up**. I think the submission has improved over the rebuttal with the new information and results and I encourage the authors to include the new findings in the main paper, and to release implementation code (as they stated in the NeurIPS paper checklist, to help other works to build upon and compare against Scala). At this point I'm rather leaning for a positive recommendation for this work. I do not have other questions.

Kind Reminder of Deadline

Dear Reviewer, Thanks for your efforts so far in our paper review. Currently, most of the concerns from other Reviewers have been well-addressed and we are eager to know whether our response has resolved your concerns. Concretely, we have added the discussion between Scala and dynamic networks which shows that our work is agnostic with this line of research. Moreover, we have clarified the contribution of our work and further expanded the application scenarios of Scala to foundation models like DINOv2. Due to the coming discussion deadline, we would like to kindly remind the Reviewer if our response through the whole rebuttal has addressed any of your concerns and helped the Reviewer to reevaluate this work. We really appreciate it if you could give us any feedback and your opinions are rather important to us. Thank you very much for your time! Sincerely, Authors