Class token knowledge distillation for efficient vision transformer

Abstract

Vision Transformer (ViT) achieves higher performance compared to Convolutional Neural Networks(CNNs) but requires more computational cost. Knowledge Distillation (KD) has demonstrated potential in compressing complex networks by transferring knowledge from a large pre-trained model to a smaller one. However, existing KD methods for ViT either employ CNNs as teachers or overlook the importance of class token ([CLS]) information. It leads to failing to effectively distill ViT’s distinct knowledge. In this paper, we propose 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 adaptively extract patch embedding information into the CLS under teacher guidance. Through these two strategies, [CLS]-KD consistently outperforms existing state-of-the-art methods on the ImageNet-1K dataset across various teacher-student settings. Moreover, the proposed method shows its generalization ability through transfer learning experiments on the CIFAR-10 and CIFAR-100 datasets.