DEVR: Train an Efficient Vision-RWKV Model with Improved Knowledge Distillation
摘要
Convolutional Neural Networks (CNNs) and Transformers have dominated the field of computer vision for some time, each interpreting images in their unique ways. Combining their strengths to enhance model performance and reducing the high computational costs associated with Transformers have been central challenges in recent research. To address these two issues, we proposed an efficient Vision-RWKV model with Knowledge Distillation (DEVR). Specifically, we selected RWKV for its linear complexity in place of the Transformer architecture and redesigned RWKV block to better suit vision tasks. Additionally, we proposed a distillation loss function and a training diagram that divides RWKV to 4 stages, employing extra projection layers and knowledge distillation to align RWKV’s features with CNN’s feature space. As shown in Fig. 1, our evaluations show that DEVR outperforms both vanilla VRWKV and Transformer-based model, offering faster speeds and lower computational costs. For example, DEVR-T achieve top-1 accuracy of 76.6%, which is 4.4% higher than DeiT-B and 1.5% than vanilla VRWKV. In downstream tasks such as detection and segmentation, DEVR maintains considerable performance while keeping resource consumption and speed optimal. Our code will be released on https://github.com/Maccusl/DEVR-Distillation-Efficient-Vison-RWKV.git