Source-linked AI summary
Retrieval Augmented Code Generation and Summarization
Md Rizwan Parvez, Wasi Uddin Ahmad, Saikat Chakraborty, Baishakhi Ray, Kai-Wei Chang
TL;DR
Code generation and summarization remain challenging because generated code can be low quality and existing models do not directly leverage available high-quality code and descriptions. REDCODER retrieves relevant code or summaries, augments the input, and supports unimodal or bimodal retrieval databases through dense retrieval and a modular generator. Across Java and Python code generation and summarization benchmarks, the framework improves reported Exact Match from 18.6 to 23.4 and BLEU-4 from 18.45 to 22.95, including when target candidates are removed from retrieval.
Problem
Existing code generation and summarization approaches often produce poor-quality code and do not directly leverage available high-quality code and descriptions during generation.
Method
REDCODER retrieves relevant code or summaries with a dense retriever, augments the original input, and generates outputs using databases containing unimodal or bimodal instances.
Results
REDCODER improves Exact Match from 18.6 to 23.4 and BLEU-4 from 18.45 to 22.95 across code generation and summarization evaluations, even after target candidates are removed.
Takeaways & Limitations
Retrieval-augmented generation can reuse relevant prior code or summaries to improve code and documentation generation across the evaluated benchmark tasks.
Takeaways & Limitations
The framework considers retrieval databases containing singleton code or summaries and parallel code-description pairs, requiring separate generator modeling considerations for these settings.
Abstract
from arXiv · showhide
Software developers write a lot of source code and documentation during software development. Intrinsically, developers often recall parts of source code or code summaries that they had written in the past while implementing software or documenting them. To mimic developers' code or summary generation behavior, we propose a retrieval augmented framework, REDCODER, that retrieves relevant code or summaries from a retrieval database and provides them as a supplement to code generation or summarization models. REDCODER has a couple of uniqueness. First, it extends the state-of-the-art dense retrieval technique to search for relevant code or summaries. Second, it can work with retrieval databases that include unimodal (only code or natural language description) or bimodal instances (code-description pairs). We conduct experiments and extensive analysis on two benchmark datasets of code generation and summarization in Java and Python, and the promising results endorse the effectiveness of our proposed retrieval augmented framework.
1 Introduction
REDCODER addresses poor code-generation and summarization quality by retrieving relevant code or summaries and augmenting the generator input. It combines dense retrieval with a modular generator framework that supports different retrieval-database formats.
- Motivation: Automated code generation and summarization remain difficult, with generated code often suffering from poor quality.The tasks require handling diverse programming-language tokens and understanding lexical, syntactic, and semantic structure.
- Approach: REDCODER retrieves existing high-quality source code or descriptions and includes them directly in the generation process.The framework is motivated by developers adapting relevant code and consulting documentation or forums when implementing software.
- Approach: The framework uses a two-step process: retrieve relevant code or summaries, then generate the target output from them alongside the original input.Figure 1 illustrates retrieval, aggregation, and generation for code generation.
- Retriever: REDCODER extends dense retrieval with separate encoders designed for programming-language and natural-language understanding.This design targets semantic retrieval beyond sparse token-level matching and avoids the large-database cost associated with cross-encoders.
- Generator: The generator supports retrieval databases containing either unimodal code or descriptions, or bimodal code-description pairs.The framework preserves model-agnostic characteristics by augmenting inputs without modifying the generator architecture.
- Results: 18.6 to 23.4 Exact Match and 18.45 to 22.95 BLEU-4 were the reported improvements for code generation and summary generation, respectively.These gains remained when the target candidate was forcefully removed from retrieved code or summaries, across Java and Python evaluations.
2 Background
The paper formulates code generation and summarization as sequence-to-sequence tasks supported by retrieval databases. It reviews sparse and dense retrieval, DPR-style training, and PLBART as the generator used in REDCODER.
- Problem Formulation: Code generation maps natural-language descriptions to source code, while code summarization maps source-code snippets to natural-language summaries.The formulation treats the two tasks as opposite input-output directions.
- Problem Formulation: For an input x, the retriever selects top-k output sequences from a database and concatenates them with x before generation.The augmented input is x′ = x ⊕ Y1 ⊕ Y2 ... ⊕ Yk.
- Retrieval Background: Sparse retrievers such as TF-IDF and BM25 rank documents through lexical matching using sparse vector representations.They compute relevance scores from query-document term overlap.
- Retrieval Background: Dense retrieval encodes queries and documents into fixed-size representations and retrieves documents using maximum inner product search.DPR uses separate query and passage encoders and trains with relevant and irrelevant passages.
- Retrieval Background: DPR training can use in-batch negatives, including BM25-selected hard negatives that lack a target-matching substring.This provides contrastive alternatives for learning query-passage relevance.
- Generator Background: PLBART is a sequence-to-sequence Transformer pretrained on source code and natural-language descriptions through denoising autoencoding.REDCODER adopts PLBART as its generator module.
3 Proposed Framework: REDCODER
REDCODER augments code generation and summarization with retrieved code or summaries through a modular retriever-generator framework. Its SCODE-R retriever uses dense encoders and in-batch negatives, while SCODE-G supports singleton or paired retrieval candidates.
- REDCODER retrieves relevant code or summaries and augments the original input before generation.The framework separates retrieval from generation, allowing different retriever and generator designs.
- 3.1 Retriever: SCODE-R: SCODE-R uses separate bidirectional Transformer encoders for source code and natural-language summaries.The encoders are built on DPR and can use CodeBERT or GraphCodeBERT representations.
- 3.1 Retriever: SCODE-R: SCODE-R fine-tunes parallel code-summary examples with the matching output as positive and other in-batch outputs as negatives.Unlike open-domain QA training, it omits BM25-derived hard negatives because non-target retrieved code or summaries may still help generation.
- 3.2 Generator: SCODE-G: SCODE-G concatenates the input with top-k retrieved sequences and feeds the augmented sequence to PLBART.The default singleton setting separates candidates with [csep] tokens; the augmented input is truncated to PLBART's maximum length of 512.
- 3.2 Generator: SCODE-G: REDCODER-EXT incorporates both members of retrieved code-summary pairs, replacing an unavailable paired sequence with an empty string.This setting supports databases containing mixtures of singleton and paired candidates.
4 Experiment Setup
The experiments evaluate REDCODER for code generation and summarization in Java and Python using CodeXGLUE and Concode data. Evaluation includes retrieval and generative baselines, large retrieval databases, and standard generation and code-correctness metrics.
- Datasets: The study evaluates code generation and summarization in Java and Python using CodeXGLUE, with Java code generation also tested on Concode.CodeXGLUE summarization data is curated from CodeSearchNet by filtering noisy examples.
- Retrieval Databases: The code retrieval database contains an average of 1.4M deduplicated Java and Python functions with paired and monolingual code.A separate summary database retains 1.1M unique summaries, and the default setup excludes target code or summaries.
- Evaluation Metrics: BLEU-4 evaluates generation and summarization quality, while CodeBLEU measures lexical, abstract-syntax-tree, and data-flow matching.Exact Match reports the percentage of outputs exactly matching their references.
- Baseline Methods: The comparison includes dense DPR and sparse BM25 retrievers, retrieval-based models, and generative models trained from scratch or pretrained.Generative baselines include RoBERTa, CodeBERT, GraphCodeBERT, GPT-style models, LSTM Seq2Seq, and Transformer systems.
- Experimental Analysis: The experiments also report retrieval performance with Mean Reciprocal Rank and examine settings where the target remains in the retrieval database.Table 6 compares SCODE-R with alternative retrievers, while Table 5 evaluates target-inclusive retrieval.
5 Results
REDCODER improves code generation and summarization by augmenting generators with retrieved code or summaries, outperforming strong baselines across evaluated tasks and languages.
- Code Generation: PLBART achieves an EM score of 18 on Concode but rarely matches target code on CodeXGLUE, indicating limited code quality without proper input supervision.BLEU and CodeBLEU scores are also low in these settings.
- Code Generation: SCODE-R significantly outperforms BM25 among retrieval-only models, while retrieved code remains relevant even when it does not exactly match the target.The targets are filtered from retrieval, so Exact Match is zero for retrieval-only models.
- Code Generation: REDCODER outperforms PLBART in code generation, with BLEU gains of 18.88, 19.54, and 5.8 on CodeXGLUE Java, CodeXGLUE Python, and Concode.REDCODER-EXT improves BLEU on Concode and all CodeXGLUE metrics by approximately 1% over REDCODER.
- Code Summarization: Retrieval-augmented generative methods outperform retrieval-only and other generative methods for code summarization.REDCODER and REDCODER-EXT surpass BM25 + PLBART by approximately 1.5 points for Python and 3.2 points for Java.
6 Analysis
Analysis examines retrieval quality, robustness to database composition and target length, retrieval-count effects, and generated-code quality through human evaluation.
- Target Retrieval: When the target remains in the retrieval database, SCODE-R performs much better and REDCODER is further enhanced across cases.The analysis also plots Recall@K for Java and Python code generation on CodeXGLUE.
- Retriever Analysis: SCODE-R performs comparably to GraphCodeBERT while retrieving significantly faster and scaling better through preindexing.Separate query and document encoders enable faster retrieval in practice.
- Target Length: SCODE-R remains consistently effective as Python target code length increases, whereas PLBART performance consistently decreases.This consistency boosts REDCODER and REDCODER-EXT relative to the generative baseline; Java shows similar results.
- Number of Retrievals: Performance typically improves with more retrievals on both tasks, with roughly 5 code candidates and 30 summaries working sufficiently well.Figure 8 presents performance as the number of retrieved candidates changes.
7 Related Works
Related work spans learning-based code generation and summarization, retrieval-based software engineering applications, and developer use of searchable code resources.
- Code Summarization: Code summarization research includes token-sequence models, tree-based models that use code structure, and retrieval-based methods.The cited literature treats source code as sequences or leverages structural properties such as syntax trees.
- Code Generation: Code generation research includes abstract syntax tree generation for syntactic correctness and pre-trained models that learn code syntax and semantics from unlabeled source code.These approaches represent two major directions described in the related work.
- Code Retrieval and Others: Software engineering applications use information retrieval for related code, API examples, libraries, and usable source code.REDCODER’s design is inspired by developers’ behavior when searching for and adapting existing resources.
8 Conclusion
The paper proposes REDCODER to automate code and documentation writing by reusing developers’ previous work, and reports performance gains from retrieved information.
- Conclusion: REDCODER reuses previously written code and documentation to automate developers’ code and documentation writing.The framework is evaluated on two benchmark datasets, with results showing a significant performance boost from retrieved information.
A Qualitative Example
The qualitative example compares PLBART with REDCODER variants on Java code generation for retrieving a MuxerStream at a specified position. Retrieval helps incorporate relevant class and argument information, while REDCODER-EXT also uses retrieved summaries to select and adapt code.
- PLBART generates basic but relevant code for the MuxerStream retrieval task.
- The retrieved rank-1 and rank-3 code contains related cPtr statements from MuxerStream and DeMuxerStream classes.
- REDCODER fuses retrieved code to produce somewhat correct MuxerStream code that uses the position argument.
- REDCODER-EXT matches the input summary with a retrieved summary, producing a MuxerStream object while borrowing throw exceptions from rank-3 code.
B Performance Difference of PLBART on CodeXGLUE and Concode
The analysis attributes performance differences between PLBART on Concode and CodeXGLUE to dataset characteristics. Concode’s preprocessing and lexical overlap make exact matching and BM25-assisted retrieval more favorable than on CodeXGLUE.
- Concode includes environment variables and methods in input summaries and replaces specific variable names with generic tokens.
- PLBART achieves good EM on Concode but not for real-code generation in CodeXGLUE, according to the authors’ suspicion.
- BM25-retrieved code has greater word overlap with Concode targets than with CodeXGLUE targets.
- BM25 plus PLBART boosts performance more on Concode than on CodeXGLUE.
- The retrieval statistics distinguish parallel and nonparallel code or summary databases, with Concode retrieved only from itself.