[论文解读] Cascaded Fast and Slow Models for Efficient Semantic Code Search

Akhilesh Deepak Gotmare, Junnan Li+2 more|arXiv (Cornell University)|Oct 15, 2021

Software Engineering Research参考文献 24被引用 7

一句话总结

本文提出 CasCode，一种基于级联快速与慢速 Transformer 的框架，用于高效且准确的语义代码搜索。首先利用快速编码器通过密集向量相似性检索 top-K 代码候选，随后通过一个慢速联合分类模型对候选进行重排序以提升准确性，最终在六种编程语言的 CodeSearchNet 上实现了新的 SOTA MRR 得分 0.7795。

ABSTRACT

The goal of natural language semantic code search is to retrieve a semantically relevant code snippet from a fixed set of candidates using a natural language query. Existing approaches are neither effective nor efficient enough towards a practical semantic code search system. In this paper, we propose an efficient and accurate semantic 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 performance 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. The proposed cascaded approach 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) as opposed to the previous state-of-the-art result of 0.713 MRR on the CodeSearchNet benchmark.

研究动机与目标

• 解决现有语义代码搜索系统在实际部署场景中效率低下且性能不佳的问题。
• 克服联合分类模型在大规模代码仓库中应用时计算成本过高的问题。
• 通过结合快速检索与高精度重排序，实现在语义代码搜索中速度与准确性的平衡。
• 通过共享参数联合训练快速与慢速组件，降低生产环境中的内存开销。
• 在 CodeSearchNet 基准上超越先前最先进方法，提升检索性能。

提出的方法

• 使用快速 Transformer 编码器为自然语言查询和代码片段生成密集语义嵌入，通过预计算的向量索引实现基于相似度的高效检索。
• 应用一个独立的、较慢的二分类模型，联合编码查询-代码对以预测匹配概率，相比简单相似度度量方法提升准确性。
• 实施级联推理流程，将快速检索阶段得到的 top-K 候选由慢速分类器进行重排序。
• 提出一种共享参数变体，通过对比学习（infoNCE）和二元交叉熵损失的多任务训练，使单个 Transformer 编码器同时服务于快速与慢速阶段。
• 离线构建预计算向量索引（PL 索引）以加速在线推理，该成本不计入运行时测量。
• 在第二阶段以批处理方式处理快速阶段输出的 top-K 候选，以优化推理效率。

实验结果

研究问题

• RQ1 结合快速检索与慢速重排序的级联架构，能否同时提升语义代码搜索的效率与准确性？
• RQ2 在双阶段代码搜索系统中，共享参数训练在不损失性能的前提下，能在多大程度上降低内存开销？
• RQ3 与基于嵌入的简单检索模型相比，联合分类模型在 MRR 和推理速度方面的表现如何？
• RQ4 在重排序阶段处理的候选数 K 变化时，推理速度与检索准确率之间的权衡关系如何？
• RQ5 为何共享参数模型的多任务训练性能劣于分别微调快速与慢速模型？

主要发现

• CasCode 在六种编程语言的 CodeSearchNet 基准上实现了新的 SOTA 平均倒数排名（MRR）得分 0.7795，显著优于此前 SOTA 的 0.713。
• CasCode 的独立模型变体实现了 0.7825 的 MRR，单次查询推理时长为 0.2883 秒，在单张 A100 GPU 上可处理 3.46 个查询/秒。
• 共享参数变体将模型大小减半，同时保持强劲性能，MRR 达到 0.7686，单次查询推理时长为 0.2956 秒。
• 将重排序阶段处理的候选数从 K=100 降低至 K=10，可使推理速度提升近 3 倍（从 3.46 提升至 9.78 个查询/秒），同时 MRR 仅轻微下降（从 0.7825 降至 0.7724）。
• 尽管参数量仅为一半，共享参数模型的推理成本与独立模型几乎相同，表明其内存利用效率高。
• 共享模型的多任务训练性能略逊于分别微调的快速与慢速模型，表明在对比学习与分类目标的联合优化之间存在平衡挑战。