SPENCER: Self-Adaptive Model Distillation for Efficient Code Retrieval

Code retrieval aims to provide users with desired code snippets based on users’ natural language queries. With the development of deep learning technologies , adopting pre-trained models for this task has become mainstream. Considering the retrieval efficiency, most of the previous approaches adopt a dual-encoder for this task, which encodes the description and code snippet into representation vectors, respectively. However, the model structure of the dual-encoder tends to limit the model’s performance, since it lacks the interaction between the code snippet and description at the bottom layer of the model during training. To improve the model’s effectiveness while preserving its efficiency, we propose a framework, which adopts S elf-Ada P tive Model Distillation for E fficient C od E R etrieval, named SPENCER. SPENCER first adopts the dual-encoder to narrow the search space and then adopts the cross-encoder to improve accuracy. To improve the efficiency of SPENCER, we propose a novel model distillation technique, which can greatly reduce the inference time of the dual-encoder while maintaining the overall performance. We also propose a teaching assistant selection strategy for our model distillation, which can adaptively select the suitable teaching assistant models for different pre-trained models during the model distillation to ensure the model performance. Extensive experiments demonstrate that the combination of dual-encoder and cross-encoder improves overall performance compared to solely dual-encoder-based models for code retrieval. Besides, our model distillation technique retains over 98% of the overall performance while reducing the inference time of the dual-encoder by 70%.