Adaptive class token knowledge distillation for efficient vision transformer

Abstract

The Vision Transformer (ViT) outperforms Convolutional Neural Networks (CNNs) but at the cost of significantly higher computational demands. Knowledge Distillation (KD) has shown promise in compressing complex networks by transferring knowledge from a large pre-trained model to a smaller one. However, current KD methods for ViT often rely on CNNs as teachers or neglect the importance of class token ([CLS]) information, resulting in ineffective distillation of ViT’s unique knowledge. In this paper, we propose Adaptive Class token Knowledge Distillation ([CLS]-KD), which fully exploits information from the class token and patches in ViT. For class embedding (CLS) distillation, the intermediate CLS of the student model is aligned with the corresponding CLS of the teacher model through a projector. Furthermore, we introduce CLS-patch attention map distillation, where an attention map between the CLS and patch embeddings is generated and matched at each layer. This empowers the student model to learn how to dynamically extract patch embedding information into the CLS under teacher guidance. Finally, we propose Adaptive Layer-wise Distillation (ALD) to mitigate the imbalance in distillation effects varying with the depth of layers. This method assigns greater weight to the losses in layers where the training discrepancies between the teacher and student models are larger during distillation. Through these strategies, [CLS]-KD consistently surpasses existing state-of-the-art methods on the ImageNet-1K dataset across various teacher–student configurations. Furthermore, the proposed method demonstrates its generalization capability through transfer learning experiments on the CIFAR-10, CIFAR-100, and CALTECH-256 datasets.