SelfRACG: Retrieval-Augmented Code Generation
- SelfRACG is a retrieval-augmented code generation method that uses LLM hidden states to produce internal query embeddings, aligning retrieval with the intended next code fragment.
- It incorporates an Information Need Expression module via Layer-wise LoRA and a two-stage Information Need-Guided training strategy to bridge the content gap between consecutive code fragments.
- Empirical evaluations on RepoEval and CrossCodeEval benchmarks demonstrate that SelfRACG significantly outperforms conventional retrieval methods, enhancing both generation quality and retrieval accuracy.
Searching arXiv for the specified paper and closely related RACG literature. SelfRACG is a retrieval-augmented code generation framework in which a LLM produces a dense embedding of its own information need and uses that embedding as the retrieval query, rather than relying on an external retriever driven by lexical or embedding similarity over the current code context. It is introduced to address the “content gap” in retrieval-augmented code generation: consecutive code fragments may be logically related yet diverge in tokens, APIs, and variable names, so retrieval based only on the visible current fragment often returns code similar to past content rather than code aligned with the next-step requirement. SelfRACG addresses this by combining an Information Need Expression module with a two-stage Information Need-Guided training strategy, enabling retrieval to be conditioned on the LLM’s own hidden-state representation of what comes next (Dong et al., 25 Jul 2025).
1. Problem setting and conceptual motivation
Retrieval-augmented code generation typically supplies an LLM with an incomplete code context and additional retrieved code snippets, after which the model generates the next code fragment. In the vanilla RACG setting, retrieval is driven by an external module such as BM25 or an embedding model, and the query is derived from the current context. This makes retrieval dependent on surface similarity or generic semantic similarity rather than on the generator’s anticipated continuation (Dong et al., 25 Jul 2025).
The central problem identified by SelfRACG is the “content gap” between consecutive fragments. In realistic software artifacts, one fragment may perform data preprocessing while the next performs model training; the two are part of the same workflow but can differ substantially in vocabulary and APIs. Under this condition, an external retriever that encodes only the current fragment tends to return more preprocessing-related material, even when the model is about to generate a training function. SelfRACG is formulated around the claim that the LLM’s hidden state for the next token already encodes what it “wants to do next,” and that a retrieval query extracted from this internal state should be better aligned with subsequent generation than a query based on the current fragment alone (Dong et al., 25 Jul 2025).
This framing places SelfRACG within repository-level code completion rather than generic text retrieval. A plausible implication is that the method is especially pertinent when code generation depends on transitions across stages of a pipeline, cross-file dependencies, or function/API abstractions that are not lexically recoverable from the immediately preceding context.
2. Core architecture
At a high level, SelfRACG turns one code LLM into both a generator and a retrieval embedding model specialized to its own information needs. Its architecture comprises three components: an Information Need Expression (INE) module, a retrieval module, and a generation module (Dong et al., 25 Jul 2025).
The INE module is implemented by Layer-wise LoRA (L-LoRA), a parameter-efficient adapter that operates in parallel with each transformer layer’s normal self-attention. From the last token’s hidden state, it computes a retrieval representation intended to encode the information need for the next code fragment. The retrieval module uses as the query embedding and performs Maximum Inner Product Search over a vector index of candidate code fragments. The generation module is the same LLM backbone, whose core parameters are not changed; it receives the original incomplete code together with the retrieved code fragments and generates the next fragment (Dong et al., 25 Jul 2025).
For a transformer with layers and hidden size , hidden states at layer are denoted
The standard attention projections for are
and standard scaled dot-product attention is
SelfRACG introduces a separate retrieval LoRA path:
A parallel retrieval-aware attention is then computed as
0
where 1 is the index of the last token, and the final information need representation is obtained by mean pooling across layers:
2
Retrieval scoring uses inner product similarity,
3
with MIPS over stored candidate embeddings 4 (Dong et al., 25 Jul 2025).
A key architectural point is that generation continues to use the original self-attention path, while the L-LoRA path is used only for retrieval. The paper explicitly notes that this preserves the base LLM’s generation ability (Dong et al., 25 Jul 2025).
3. Information Need Expression
In SelfRACG, “information need expression” does not mean natural-language query generation. The expression is an implicit vector embedding 5 computed from hidden states via the retrieval-aware attention path. It is not emitted as tokens; rather, it is an internal representation extracted from the last-token hidden state across layers (Dong et al., 25 Jul 2025).
The computation proceeds layer by layer. First, the model forms the ordinary 6, 7, and 8 projections from the contextual hidden states. Second, the retrieval projections 9, 0, and 1 are computed using L-LoRA. Third, the last token’s retrieval query attends over the whole sequence to produce the layer-wise representation 2. Finally, these layer-wise representations are averaged to form the query embedding 3 (Dong et al., 25 Jul 2025).
This design is significant because the hidden states are conditioned on the entire current context, including incomplete code and any additional instructions. The last token’s hidden state therefore reflects the model’s continuation dynamics rather than mere similarity to previous text. The paper explicitly argues that 4 encodes the information need for the next fragment, not just similarity to the past context (Dong et al., 25 Jul 2025).
SelfRACG also distinguishes query and candidate embeddings parametrically rather than through natural-language prompting. Query representations are computed directly from incomplete code, whereas candidate representations are computed after prepending a special prefix token indicating that the input is a retrieval candidate. This separation is described as helping the model differentiate between context embeddings and candidate embeddings (Dong et al., 25 Jul 2025).
4. Information Need-Guided training
The INE module is trained through a two-stage Information Need-Guided strategy, both stages using the same contrastive objective:
5
For a query 6, 7 is the positive candidate and 8 is the negative set (Dong et al., 25 Jul 2025).
Stage 1 is retrieval learning on GitHub. The data source is the GitHub corpus from CodeRAG-Bench, reported as approximately 9M code files and filtered to approximately 0M files after removing benchmark repositories and files shorter than 20 lines. Each file is split into fragments using a fixed 20-line interval. For each fragment 1, the positive sample 2 is the next fragment in the same file, while negatives are taken from in-batch samples. This stage is intended to teach the INE module next-step continuity in code, such as “define variables 3 use them” or “preprocess 4 train,” but remains generic with respect to any particular generator’s stylistic preferences (Dong et al., 25 Jul 2025).
Stage 2 aligns retrieval with the LLM’s own generation preferences. For each fragment 5, the base LLM generates several candidate next fragments 6 using vLLM. The immediate next-step generation 7 is treated as the positive sample, while the negative pool is formed from 8 together with in-batch negatives. The paper reports about 500k synthetic samples for this stage, generated using OpenCoder/Qwen-based LLMs. Only the INE module’s L-LoRA parameters are trained; the original weights of the base LLM are not changed (Dong et al., 25 Jul 2025).
The stated purpose of Stage 2 is to resolve permutation ambiguity and align retrieval with “this LLM’s own generation preferences.” In the authors’ formulation, among several logically valid next fragments, the model learns to retrieve the one it is most likely to generate itself (Dong et al., 25 Jul 2025). This suggests a shift from generic repository continuity modeling to model-specific retrieval alignment.
5. Retrieval-generation workflow and handling of the content gap
At inference time, SelfRACG first encodes the incomplete code fragment into a query embedding 9 using the LLM and the INE module. Candidate code fragments are pre-indexed by computing their embeddings 0 with the same INE module, but with the candidate prefix token. Retrieval then uses MIPS to return the top-1 fragments maximizing 2 (Dong et al., 25 Jul 2025).
After retrieval, the top-3 fragments are concatenated with the original context according to a prompt template similar to those used in RepoCoder and CodeRAG. The combined prompt is then passed to the same LLM backbone to generate the next code fragment. The framework is iterative, but its distinctive feature lies in query computation rather than in the outer retrieval-generation loop itself (Dong et al., 25 Jul 2025).
The contrast with vanilla RACG is direct. In vanilla RACG, the retrieval query is an embedding or sparse representation of the current code context, and the retriever is external—examples listed in the paper include OpenAI embeddings, NV-Embed, and GritLM. In SelfRACG, the retrieval query is the information need embedding extracted from the model’s hidden states for the next token, and retrieval is handled by a small adapter on top of the same LLM (Dong et al., 25 Jul 2025).
The paper further compares SelfRACG with three retrieval strategies: VanillaRACG using GritLM, a “w/ next fragment” strategy that uses the fragment after the retrieved code as extra knowledge, and a “w/ next query” strategy that has the LLM generate an explicit textual query which is then embedded with GritLM. SelfRACG is reported to outperform all three. The authors attribute this to the implicit embedding 4 being more expressive than either content-similarity embeddings of the current code or textual queries that must then be re-embedded, possibly incurring information loss (Dong et al., 25 Jul 2025).
6. Empirical evaluation
SelfRACG is evaluated on RepoEval and CrossCodeEval, two repository-level code completion benchmarks. RepoEval comprises line-level completion, API-level completion, and function-level completion. CrossCodeEval targets cross-file contextual understanding and requires modeling dependencies across multiple files (Dong et al., 25 Jul 2025).
The method is applied to OpenCoder-1.5B-Base, OpenCoder-8B-Base, and Qwen2.5-Coder-Instruct models at 3B and 7B scales. The code completion metrics are Exact Match, Edit Similarity, and Pass@1 for RepoEval-Function. Retrieval metrics are Recall@5 for 6 and MRR@10. Baselines include BM25, OpenAI text-embedding-3-small/large, GritLM-7B, and NV-Embed-v2 (Dong et al., 25 Jul 2025).
The following results are explicitly reported for OpenCoder-8B-Base:
| Setting | Mean ES | RepoEval-Function Pass@1 |
|---|---|---|
| No retrieval | 0.518 | 24.8% |
| VanillaRACG (BM25) | 0.557 | — |
| VanillaRACG (OpenAI-large) | 0.567 | — |
| VanillaRACG (NV-Embed-v2) | 0.572 | — |
| VanillaRACG (GritLM-7B) | 0.576 | 30.7% |
| SelfRACG | 0.593 | 33.8% |
Task-specific Edit Similarity results for OpenCoder-8B are also given: RepoEval-API increases from 0.650 with GritLM-7B to 0.680 with SelfRACG; RepoEval-Line from 0.689 to 0.705; and RepoEval-Function from 0.463 to 0.481. For OpenCoder-1.5B, mean ES is reported to improve from 0.545 with GritLM to 0.552 with SelfRACG. For Qwen2.5-Coder-3B, mean ES improves from 0.555 under best VanillaRACG to 0.569 with SelfRACG, and for Qwen2.5-Coder-7B from 0.556 to 0.562 (Dong et al., 25 Jul 2025).
Retrieval quality is likewise improved. On line-level completion, the paper reports Recall@1 and MRR@10 of 0.174 and 0.356 for GritLM-7B, 0.178 and 0.357 for NV-Embed-v2, and 0.237 and 0.413 for SelfRACG with OpenCoder-8B. SelfRACG with Qwen2.5-Coder-7B reaches Recall@1 = 0.244 and MRR@10 = 0.413, while Qwen2.5-Coder-3B reaches Recall@1 = 0.239 and MRR@10 = 0.411 (Dong et al., 25 Jul 2025).
The paper states that the gains are stronger for API- and function-level tasks, where the content gap is larger and next-step information needs are more abstract, and slightly smaller for line-level tasks, where next lines are more deterministic and easier to predict (Dong et al., 25 Jul 2025). This interpretation is consistent with the method’s design objective.
7. Ablations, limitations, and broader significance
Ablation studies on OpenCoder-8B line-level completion compare full SelfRACG with variants without Stage 2 and without L-LoRA. Full SelfRACG yields MRR@10 = 0.413 and ES = 0.7048. Removing Stage 2 reduces performance to MRR@10 = 0.364 and ES = 0.6912. Removing L-LoRA reduces it further to MRR@10 = 0.338 and ES = 0.6847 (Dong et al., 25 Jul 2025).
The paper interprets these results as evidence that Stage 2 preference alignment is crucial for both retrieval and generation, and that L-LoRA is more effective than full-parameter tuning under limited GPU budget because it enables larger batches and hence more negatives in the contrastive loss. It further notes that without Stage 2, retrieval can become misaligned with what the LLM actually generates, producing “knowledge conflict” in which retrieved code suggests a generic continuation rather than the one preferred by the generator (Dong et al., 25 Jul 2025).
SelfRACG is also positioned as computationally cheaper than GritLM: the paper reports 384 GPU hours versus 3072 GPU hours. At the same time, it explicitly acknowledges several limitations. Experiments are limited to models up to 8B parameters. Evaluation is restricted to retrieval-augmented code generation rather than broader RAG tasks. Performance still depends on the quality and coverage of the retrieval index and on the fragment splitting strategy, with 20-line chunks potentially misaligned with semantic units. Training still requires large GitHub corpora and synthetic Stage 2 data. The method has not been stress-tested for hallucination or security risks (Dong et al., 25 Jul 2025).
Conceptually, SelfRACG can be framed as replacing the standard RACG query function with one derived from the LLM’s internal hidden-state trajectory:
- standard RACG: 7, where 8 is the current context and 9 is an external encoder;
- SelfRACG: 0, where 1 are the LLM hidden states and 2 is the L-LoRA-based retrieval adapter (Dong et al., 25 Jul 2025).
The theoretical motivation given in the paper is that the next-token distribution 3 is computed from the last-token hidden state, so those hidden states contain the semantic information needed for continuation. A learned mapping from hidden states to retrieval embeddings should therefore align retrieval with actual future needs. This suggests a more general principle for retrieval-augmented generation: retrieval may be improved when query formation is coupled to the generator’s own continuation dynamics rather than delegated to an external similarity model alone (Dong et al., 25 Jul 2025).