[论文解读] Generation-Augmented Query Expansion For Code Retrieval

Dong Li, Yelong Shen+4 more|arXiv (Cornell University)|Dec 20, 2022

Topic Modeling被引用 4

一句话总结

本文提出了一种新型两阶段框架——生成增强代码检索（GACR），通过从自然语言查询生成代码片段并利用其扩展查询，从而增强代码检索性能。通过融合原始文档与语义丰富的生成代码，GACR在CodeSearchNet基准测试中对六种编程语言均实现了最先进性能，显著优于现有检索模型。

ABSTRACT

Pre-trained language models have achieved promising success in code retrieval tasks, where a natural language documentation query is given to find the most relevant existing code snippet. However, existing models focus only on optimizing the documentation code pairs by embedding them into latent space, without the association of external knowledge. In this paper, we propose a generation-augmented query expansion framework. Inspired by the human retrieval process - sketching an answer before searching, in this work, we utilize the powerful code generation model to benefit the code retrieval task. Specifically, we demonstrate that rather than merely retrieving the target code snippet according to the documentation query, it would be helpful to augment the documentation query with its generation counterpart - generated code snippets from the code generation model. To the best of our knowledge, this is the first attempt that leverages the code generation model to enhance the code retrieval task. We achieve new state-of-the-art results on the CodeSearchNet benchmark and surpass the baselines significantly.

研究动机与目标

• 为解决现有代码检索模型仅依赖嵌入相似性而未利用外部知识的局限性。
• 探究代码生成模型是否可通过提供语义更丰富的查询扩展来提升检索性能。
• 设计一种融合机制，有效结合自然语言文档与生成代码，以提升表征学习效果。
• 证明生成增强查询扩展在弥合自然语言与代码之间语义鸿沟方面的有效性。

提出的方法

• 利用预训练的代码生成模型（例如，Codex）从给定的自然语言文档查询中生成代码片段。
• 通过将原始自然语言描述与生成的代码片段相结合，构建扩展后的查询。
• 采用双表征注意力机制，学习并融合来自自然语言和生成代码表征的特征。
• 在增强后的查询上训练检索模型（例如，GraphCodeBERT），以提升与相关代码片段的语义匹配能力。
• 探索多种变体，包括单个和多个生成的代码片段，以评估不同配置下的鲁棒性与性能。
• 使用对比学习与微调优化增强查询-代码对上的检索模型。

实验结果

研究问题

• RQ1 代码生成模型能否提升代码检索任务中自然语言查询的语义表达能力？
• RQ2 与仅使用原始查询相比，将生成代码与文档查询融合后对检索性能有何影响？
• RQ3 在代码检索的查询扩展中，自然语言与生成代码表征的最佳结合方式是什么？
• RQ4 生成代码的质量是否显著影响检索准确率？与使用生成器的中间隐藏表征相比，表现如何？
• RQ5 所提出方法在不同编程语言和查询长度下是否具有良好的泛化能力？

主要发现

• GACR在CodeSearchNet基准测试中实现了最先进性能，在所有六种评估的编程语言中均超越了GraphCodeBERT和CodeRetriever等先前模型。
• 在32个词元的查询长度下，GACR在CodeSearchNet基准测试中的整体检索准确率达到0.796，优于基线模型。
• 使用生成代码进行查询扩展相比仅使用原始文档查询，检索性能相对提升了10.5%。
• 将生成代码作为增强源的效果优于使用生成器的中间隐藏表征，相对提升达14.7%，超过基于嵌入的增强方法。
• 性能提升在所有编程语言中均保持一致，其中Python和Go的提升最为显著。
• 消融实验证实，双表征注意力机制能有效捕捉跨模态语义，且在不同查询长度下模型保持鲁棒性。