AdaViT: Task Adaptive ViT Compression Based on Module Replacement

Large pre-trained visual models, such as Vision Transformer (ViT), have shown excellent results in a variety of visual processing tasks. However, the large parameter sizes of these models make them difficult to deploy in applications that require limited resources for fast real-time inference. Existing model compression methods compress pre-trained ViTs into models of the same structure, but the compressed models fail to maximise performance on different tasks and consume significant computational and storage resources to adapt to downstream tasks. Pre-trained ViT models learn rich knowledge from large datasets and require different knowledge for different tasks, so removing task-specific redundant knowledge from pre-trained models is the key to achieve model compression while maximising model performance on different tasks. Based on this idea, we propose a novel ViT compression method, AdaViT. Firstly, a lightweight module is introduced, which adds less than 2% of extra parameters to realise the adaptive tuning of the model to different tasks to show the best prediction, and its effectiveness is proved by experiments. Second, inspired by the idea of knowledge distillation, we propose a new module replacement compression method, which effectively compresses the ViT by gradually replacing the Transformer module in the original ViT, and achieves task-oriented adaptive ViT compression through the combination of adaptive modules and model replacement methods. We evaluate AdaViT on several visual classification tasks and compare it with other ViT compression methods, demonstrating the effectiveness of task-adaptive ViT compression.