Inference-Time Masking for Retrieval-Augmented Code Generation
Abstract: Retrieval-Augmented Code Generation utilizes relevant code examples as auxiliary context to improve model performance in code generation tasks. However, the effectiveness of this strategy diminishes after the second iteration, as the retrieved code examples remain the same and the generated code becomes similar. To address this issue, we propose an **I**nference-time **M**asking strategy for **R**etrieval-**A**ugmented **C**ode **G**eneration (IM-RACG), where the retrieved code examples are masked before being used as auxiliary context. By masking parts of the examples, the diversity of the auxiliary context is increased and the context length is reduced effectively. Given the low information density of code, the remaining context still contains valuable information. As a result, this strategy encourages the model to generate more diverse code, leading performance to scale with the number of iterations. Experimental results on MBPP and HumanEval datasets demonstrate that IM-RACG significantly enhances all tested model's performance across, with an average improvement of approximately **4.5%** in pass rate compared to the original iterative RACG. Additionally, IM-RACG shows the greatest enhancement on MBPP using Llama-8b, with an increasement of the pass rate from 76.2% to **83.4%**.
Paper Type: Long
Research Area: NLP Applications
Research Area Keywords: code generation and understanding
Contribution Types: NLP engineering experiment
Languages Studied: English
Submission Number: 248
Loading