Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SelfRACG: Enabling LLMs to Self-Express and Retrieve for Code Generation

Qian Dong, Jia Chen, Qingyao Ai, Hongning Wang, Haitao Li, Yi Wu, Yao Hu, Yiqun Liu, Shaoping Ma

TL;DR

SelfRACG addresses the misalignment between externally retrieved code fragments and an LLM's next-step information needs in RACG. It introduces an Information Need Expression (INE) module and an Information Need-Guided (ING) two-stage training strategy, enabling the LLM to self-express retrieval requirements and align retrieval with generation while preserving generation quality. The approach uses Layer-wise Low-Rank Adaptation (L-LoRA) to create retrieval-specific attention in parallel with standard self-attention, and a contrastive ING objective over GitHub data and LLM-synthesized fragments. Empirical results on RepoEval and CrossCodeEval show improved code completion and retrieval performance with far lower training costs (roughly 1/8 of strong baselines) and lightweight deployment, indicating practical advantages for retrieval-augmented code generation. Overall, SelfRACG delivers a scalable, self-guided RACG framework with strong empirical gains and efficiency gains, paving the way for broader application of self-expressive retrieval in code and beyond.

Abstract

Existing retrieval-augmented code generation (RACG) methods typically use an external retrieval module to fetch semantically similar code snippets used for generating subsequent fragments. However, even for consecutive code fragments, the content often diverges due to logical progression, resulting in a content gap. This gap undermines the performance of current RACG methods, as \textit{external} retrieval modules based on content matching fail to infer the specific information need of LLMs to generate the next code fragment. Therefore, we propose \textbf{SelfRACG}, a novel paradigm that enables large language models (LLMs) to \textbf{Self}-express their information needs to enhance \textbf{RACG}. Specifically, SelfRACG includes an information need expression module and a two-stage information need-guided training strategy, which encourages LLMs to express their information need. Extensive experiments demonstrate that SelfRACG can retrieve external knowledge that better aligns with the LLM's own information needs, resulting in superior generation performance compared to vanilla RACG.

SelfRACG: Enabling LLMs to Self-Express and Retrieve for Code Generation

TL;DR

SelfRACG addresses the misalignment between externally retrieved code fragments and an LLM's next-step information needs in RACG. It introduces an Information Need Expression (INE) module and an Information Need-Guided (ING) two-stage training strategy, enabling the LLM to self-express retrieval requirements and align retrieval with generation while preserving generation quality. The approach uses Layer-wise Low-Rank Adaptation (L-LoRA) to create retrieval-specific attention in parallel with standard self-attention, and a contrastive ING objective over GitHub data and LLM-synthesized fragments. Empirical results on RepoEval and CrossCodeEval show improved code completion and retrieval performance with far lower training costs (roughly 1/8 of strong baselines) and lightweight deployment, indicating practical advantages for retrieval-augmented code generation. Overall, SelfRACG delivers a scalable, self-guided RACG framework with strong empirical gains and efficiency gains, paving the way for broader application of self-expressive retrieval in code and beyond.

Abstract

Existing retrieval-augmented code generation (RACG) methods typically use an external retrieval module to fetch semantically similar code snippets used for generating subsequent fragments. However, even for consecutive code fragments, the content often diverges due to logical progression, resulting in a content gap. This gap undermines the performance of current RACG methods, as \textit{external} retrieval modules based on content matching fail to infer the specific information need of LLMs to generate the next code fragment. Therefore, we propose \textbf{SelfRACG}, a novel paradigm that enables large language models (LLMs) to \textbf{Self}-express their information needs to enhance \textbf{RACG}. Specifically, SelfRACG includes an information need expression module and a two-stage information need-guided training strategy, which encourages LLMs to express their information need. Extensive experiments demonstrate that SelfRACG can retrieve external knowledge that better aligns with the LLM's own information needs, resulting in superior generation performance compared to vanilla RACG.

Paper Structure

This paper contains 28 sections, 7 equations, 5 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Code Large Language Models
  4. Retrieval-Augmented Code Generation
  5. LLM-based Embedding Models
  6. Method
  7. Information Need Expression Module
  8. Information Need-Guided Training
  9. Inference
  10. Experimental Setup
  11. Baselines
  12. Benchmarks
  13. Evaluation Metrics
  14. Code Completion Metrics
  15. Code Retrieval Metrics
  16. ...and 13 more sections

Figures (5)

  • Figure 1: Comparison of vanilla RACG vs. SelfRACG.
  • Figure 2: The workflow of SelfRACG. We use the $l$-th layer as an example and provide a detailed explanation of the relationship between retrieval-aware attention and self-attention within the gray box at the left of figure.
  • Figure 3: The illustration of two stage in ING training.
  • Figure 4: The generation performance with different number of retrieved code fragments.
  • Figure 5: Prompt template used in SelfRACG.