一句话:提出级联快慢模型并共享单编码器参数,实现高效且准确的文本到代码检索,在 CodeSearchNet 上取得 SOTA。
语义代码检索(文本到代码检索)旨在用自然语言查询从代码库中召回语义相关的代码片段。构建实用检索系统时,现有方法难以在检索速度与结果相关性之间取得最优平衡。本文提出一种级联快慢模型的文本到代码检索框架:先学习一个快速 transformer 编码器以构建可扩展索引完成快速召回,再训练一个基于分类的慢速重排序模型,对快速检索返回的 top-K 结果进行精排以提升准确性。为降低实际部署两个独立模型带来的高内存开销,作者进一步提出基于共享参数的单 transformer 编码器联合训练快慢模型。在 CodeSearchNet 基准的 6 种编程语言上,该级联方法平均 MRR 达 0.7795,优于此前最优的 0.744,且高效可扩展。
The goal of semantic code search or text-to-code search is to retrieve a semantically relevant code snippet from an existing code database using a natural language query. When constructing a practical semantic code search system, existing approaches fail to provide an optimal balance between retrieval speed and the relevance of the retrieved results. We propose an efficient and effective text-to-code search framework with cascaded fast and slow models, in which a fast transformer encoder model is learned to optimize a scalable index for fast retrieval followed by learning a slow classification-based re-ranking model to improve the accuracy of the top K results from the fast retrieval. To further reduce the high memory cost of deploying two separate models in practice, we propose to jointly train the fast and slow model based on a single transformer encoder with shared parameters. Empirically our cascaded method is not only efficient and scalable, but also achieves state-of-the-art results with an average mean reciprocal ranking (MRR) score of 0.7795 (across 6 programming languages) on the CodeSearchNet benchmark as opposed to the prior state-of-the-art result of 0.744 MRR. Our codebase can be found at this link.
@inproceedings{Gotmare0JH23,
title = {Efficient Text-to-Code Retrieval with Cascaded Fast and Slow Transformer Models},
author = {Akhilesh Deepak Gotmare and Junnan Li and Shafiq Joty and Steven C.H. Hoi},
booktitle = {FSE},
year = {2023},
doi = {10.1145/3611643.3616369}
}