[Paper Review] Generation-Augmented Query Expansion For Code Retrieval

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

Topic Modeling4 citations

TL;DR

This paper proposes Generation-Augmented Code Retrieval (GACR), a novel two-stage framework that enhances code retrieval by generating code snippets from natural language queries and using them to expand the query. By fusing the original documentation with semantically rich generated code, GACR achieves state-of-the-art performance on the CodeSearchNet benchmark across six programming languages, significantly outperforming existing retrieval models.

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.

Motivation & Objective

• To address the limitation of existing code retrieval models that rely solely on embedding similarity without leveraging external knowledge.
• To investigate whether code generation models can improve retrieval performance by providing semantically richer query expansions.
• To design a fusion mechanism that effectively combines natural language documentation and generated code for improved representation learning.
• To demonstrate the effectiveness of generation-augmented query expansion in bridging the semantic gap between natural language and code.

Proposed method

• Utilize a pre-trained code generation model (e.g., Codex) to generate code snippets from a given natural language documentation query.
• Construct an expanded query by combining the original natural language description with the generated code snippet.
• Employ a dual representation attention mechanism to learn and fuse features from both the natural language and the generated code representations.
• Train a retrieval model (e.g., GraphCodeBERT) on the augmented queries to improve semantic matching with relevant code snippets.
• Explore multiple variants, including single and multiple generated code snippets, to assess robustness and performance under different configurations.
• Use contrastive learning and fine-tuning to optimize the retrieval model on the augmented query–code pairs.

Experimental results

Research questions

• RQ1 Can code generation models improve the semantic expressiveness of natural language queries in code retrieval tasks?
• RQ2 How does the fusion of generated code with documentation queries affect retrieval performance compared to using only the original query?
• RQ3 What is the optimal way to combine natural language and generated code representations for query expansion in code retrieval?
• RQ4 Does the quality of generated code significantly impact retrieval accuracy, and how does it compare to using intermediate hidden representations from the generator?
• RQ5 How does the proposed method generalize across different programming languages and query lengths?

Key findings

• GACR achieves state-of-the-art performance on the CodeSearchNet benchmark, surpassing previous models like GraphCodeBERT and CodeRetriever across all six evaluated programming languages.
• The model achieves an overall retrieval accuracy of 0.796 on the CodeSearchNet benchmark with a 32-token query length, improving over the baseline.
• Using generated code for query expansion leads to a 10.5% relative improvement in retrieval performance compared to using only the original documentation query.
• The use of generated code as an augmentation source outperforms using intermediate hidden representations from the generator, with a 14.7% relative improvement over embedding-based augmentation.
• The performance gain is consistent across all programming languages, with the highest improvement observed in Python and Go.
• Ablation studies confirm that the dual representation attention mechanism effectively captures cross-modal semantics, and the model remains robust under varying query lengths.