SPENCER: Self-Adaptive Model Distillation for Efficient Code Retrieval

TL;DR

SPENCER introduces a self-adaptive model distillation framework that combines dual-encoder and cross-encoder architectures to enhance code retrieval accuracy while significantly reducing inference time.

Contribution

The paper presents a novel self-adaptive model distillation method with a teaching assistant selection strategy to improve efficiency and performance in code retrieval models.

Findings

Achieves over 98% of the original model performance.

Reduces inference time of the dual-encoder by 70%.

Improves overall accuracy compared to solely dual-encoder models.

Abstract

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 Self-AdaPtive Model Distillation for Efficient CodE Retrieval, 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%.