# Generation-Augmented Query Expansion For Code Retrieval

arXiv: https://arxiv.org/abs/2212.10692 · HTML: https://mlsanity.com/paper/2212.10692/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 26 references that are in the map.
- First posted on arXiv 20 Dec 2022.

## Builds on (most cited first, 18 of 18)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [Competition-Level Code Generation with AlphaCode](https://mlsanity.com/paper/2203.07814/) (2022)
- [CodeBERT: A Pre-Trained Model for Programming and Natural Languages](https://mlsanity.com/paper/2002.08155/) (2020)
- [CodeGen: An Open Large Language Model for Code with Multi-Turn Program Synthesis](https://mlsanity.com/paper/2203.13474/) (2022)
- [CodeSearchNet Challenge: Evaluating the State of Semantic Code Search](https://mlsanity.com/paper/1909.09436/) (2019)
- [GraphCodeBERT: Pre-training Code Representations with Data Flow](https://mlsanity.com/paper/2009.08366/) (2020)
- [UniXcoder: Unified Cross-Modal Pre-training for Code Representation](https://mlsanity.com/paper/2203.03850/) (2022)
- [Generation-Augmented Retrieval for Open-domain Question Answering](https://mlsanity.com/paper/2009.08553/) (2020)
- [Retrieval Augmented Code Generation and Summarization](https://mlsanity.com/paper/2108.11601/) (2021)
- [Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup](https://mlsanity.com/paper/2101.06983/) (2021)
- [ReACC: A Retrieval-Augmented Code Completion Framework](https://mlsanity.com/paper/2203.07722/) (2022)
- [A Retrieve-and-Edit Framework for Predicting Structured Outputs](https://mlsanity.com/paper/1812.01194/) (2018)
- [Retrieval-Based Neural Code Generation](https://mlsanity.com/paper/1808.10025/) (2018)

## Built on by (most cited first, 6 of 6)

- [Large Language Models for Software Engineering: A Systematic Literature Review](https://mlsanity.com/paper/2308.10620/) (2023)
- [RepoCoder: Repository-Level Code Completion Through Iterative Retrieval and Generation](https://mlsanity.com/paper/2303.12570/) (2023)
- [Agents in Software Engineering: Survey, Landscape, and Vision](https://mlsanity.com/paper/2409.09030/) (2024)
- [Revisiting Recommendation Loss Functions through Contrastive Learning (Technical Report)](https://mlsanity.com/paper/2312.08520/) (2023)
- [(Debiased) Contrastive Learning Loss for Recommendation (Technical Report)](https://mlsanity.com/paper/2312.08517/) (2023)
- [Accurate, Efficient, and Explainable Deep Learning Approaches for Environmental Science Problems](https://mlsanity.com/paper/2605.19366/) (2026)

## Abstract (arXiv)

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.