Slicing Vision Transformer for Flexible Inference

TL;DR

Scala enables a single Vision Transformer to represent multiple smaller width variants via width-based slicing, addressing dynamic resource constraints with flexible inference. It introduces Isolated Activation to separate the smallest subnet’s representation and Scale Coordination to provide each subnet with stable, simplified learning signals, including Progressive Knowledge Transfer, Stable Sampling, and Noisy Calibration. With one-shot training, Scala matches Separate Training performance while using far fewer parameters and delivers meaningful gains on ImageNet-1K compared with prior work. The approach extends to transfer learning and dense prediction, demonstrating robustness across tasks. Overall, Scala presents a practical paradigm shift for deploying slimmable ViTs on edge devices, albeit with higher training time than standard single-model training.

Abstract

Vision Transformers (ViT) is known for its scalability. In this work, we target to scale down a ViT to fit in an environment with dynamic-changing resource constraints. We observe that smaller ViTs are intrinsically the sub-networks of a larger ViT with different widths. Thus, we propose a general framework, named Scala, to enable a single network to represent multiple smaller ViTs with flexible inference capability, which aligns with the inherent design of ViT to vary from widths. Concretely, Scala activates several subnets during training, introduces Isolated Activation to disentangle the smallest sub-network from other subnets, and leverages Scale Coordination to ensure each sub-network receives simplified, steady, and accurate learning objectives. Comprehensive empirical validations on different tasks demonstrate that with only one-shot training, Scala learns slimmable representation without modifying the original ViT structure and matches the performance of Separate Training. Compared with the prior art, Scala achieves an average improvement of 1.6% on ImageNet-1K with fewer parameters.

Figures (19)

• Figure 1: Illustration of different means to slice the ViT architecture. Irregular slicing cai2019onceyu2020bignaschen2021autoformer results in unconventional structures while uniform slicing yu2019universally aligns with the inherent design of ViT to vary from widths.
• Figure 2: The available uniform slicing method US-Net yu2019universally lags behind Separate Training (ST) remarkably on ViTs. Performance gaps with ST are shown.
• Figure 3: Evaluating US-Net over CNN and ViT at unseen width ratios to examine the interpolation (denoted as $\star$) and extrapolation (denoted as $+$) abilities on ImageNet-1K.
• Figure 4: We train US-Net on ViT-S without constantly activating the smallest subnet (denoted as *) and observe an average performance gain of 1.8$\%$ at other width ratios on ImageNet-1K.
• Figure 5: Comparisons of Scala with different slicing granularity and Separate Training (ST). Improvements over ST are shown.