Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Preference-Guided Refactored Tuning for Retrieval Augmented Code Generation

Xinyu Gao, Yun Xiong, Deze Wang, Zhenhan Guan, Zejian Shi, Haofen Wang, Shanshan Li

TL;DR

The paper addresses redundancy and a mismatch between retrievers and generators in retrieval-augmented code generation. It introduces RRG, adding a code refactorer that compresses retrieved content into concise, model-friendly snippets, and employs a two-stage training regime to fuse parametric and non-parametric knowledge with generator preferences. Through extensive experiments on ConCode and CodeSearchNet, RRG yields substantial gains in EM, BLEU, and CodeBLEU across diverse retrievers and generator models while reducing inference costs. The findings demonstrate strong generalization and practical potential for plug-and-play integration, with future directions toward interpretability and adaptive token-length strategies.

Abstract

Retrieval-augmented code generation utilizes Large Language Models as the generator and significantly expands their code generation capabilities by providing relevant code, documentation, and more via the retriever. The current approach suffers from two primary limitations: 1) information redundancy. The indiscriminate inclusion of redundant information can result in resource wastage and may misguide generators, affecting their effectiveness and efficiency. 2) preference gap. Due to different optimization objectives, the retriever strives to procure code with higher ground truth similarity, yet this effort does not substantially benefit the generator. The retriever and the generator may prefer different golden code, and this gap in preference results in a suboptimal design. Additionally, differences in parameterization knowledge acquired during pre-training result in varying preferences among different generators. To address these limitations, in this paper, we propose RRG (Retrieve, Refactor, Generate), a novel framework for effective and efficient code generation. This framework introduces a code refactorer module between the retriever and the generator to bridge them. The refactoring process transforms the raw retrieved code into a more concise, efficient, and model-friendly version. It eliminates redundant information and noise, reducing the input length. Consequently, the generator receives higher-quality context, enabling it to produce more accurate results with lower inference costs. We conducted comprehensive experiments on multiple datasets. In the experiments, we confirmed the existence of a preference gap between the retriever and the generator, and RRG effectively bridges this gap. Specifically, RRG achieved significant performance improvements, with increases of up to 28% on EM, 13% on BLEU, and 6.8% on CodeBLEU.

Preference-Guided Refactored Tuning for Retrieval Augmented Code Generation

TL;DR

The paper addresses redundancy and a mismatch between retrievers and generators in retrieval-augmented code generation. It introduces RRG, adding a code refactorer that compresses retrieved content into concise, model-friendly snippets, and employs a two-stage training regime to fuse parametric and non-parametric knowledge with generator preferences. Through extensive experiments on ConCode and CodeSearchNet, RRG yields substantial gains in EM, BLEU, and CodeBLEU across diverse retrievers and generator models while reducing inference costs. The findings demonstrate strong generalization and practical potential for plug-and-play integration, with future directions toward interpretability and adaptive token-length strategies.

Abstract

Retrieval-augmented code generation utilizes Large Language Models as the generator and significantly expands their code generation capabilities by providing relevant code, documentation, and more via the retriever. The current approach suffers from two primary limitations: 1) information redundancy. The indiscriminate inclusion of redundant information can result in resource wastage and may misguide generators, affecting their effectiveness and efficiency. 2) preference gap. Due to different optimization objectives, the retriever strives to procure code with higher ground truth similarity, yet this effort does not substantially benefit the generator. The retriever and the generator may prefer different golden code, and this gap in preference results in a suboptimal design. Additionally, differences in parameterization knowledge acquired during pre-training result in varying preferences among different generators. To address these limitations, in this paper, we propose RRG (Retrieve, Refactor, Generate), a novel framework for effective and efficient code generation. This framework introduces a code refactorer module between the retriever and the generator to bridge them. The refactoring process transforms the raw retrieved code into a more concise, efficient, and model-friendly version. It eliminates redundant information and noise, reducing the input length. Consequently, the generator receives higher-quality context, enabling it to produce more accurate results with lower inference costs. We conducted comprehensive experiments on multiple datasets. In the experiments, we confirmed the existence of a preference gap between the retriever and the generator, and RRG effectively bridges this gap. Specifically, RRG achieved significant performance improvements, with increases of up to 28% on EM, 13% on BLEU, and 6.8% on CodeBLEU.
Paper Structure (27 sections, 5 equations, 6 figures, 5 tables)

This paper contains 27 sections, 5 equations, 6 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Pretrained Models for Code Generation
  4. Retrieval-Augmented Generation
  5. Retrieval-Augmented Code Generation
  6. Approach
  7. Overview
  8. Stage 1: Refactored Fine-Tuning
  9. Training to refactor code
  10. Training to utilize information
  11. Stage 2: Preference-aware Tuning
  12. Experimental Setup
  13. Datasets
  14. Base Models
  15. Metrics
  16. ...and 12 more sections

Figures (6)

  • Figure 1: Comparison of two RACG frameworks. The upper shows the conventional approach, where the retrieved code is directly fed to the LLMs directly, but this holds redundant information and is not model-friendly. The lower shows the RRG framework, where the retrieved code is refactored into concise, model-friendly code, which improves the effectiveness and efficiency of the generator.
  • Figure 2: The overall architecture of our proposed RRG with Preference-Guided Refactored Tuning.
  • Figure 3: Shows the improvements brought by the RRG framework when configured with different retrievers. The generator based on CodeGPT-java.
  • Figure 4: Shows the performance of RL and NO-RL in different code tokens settings
  • Figure 5: Case comparison. Redundant info and preference gaps led the generator to produce wrong answers from raw retrieved code, whereas it generated correctly from the refactored code. Highlighted tokens significantly impact code generation, whereas misleading tokens in the raw do not affect the generator in the refactored version.
  • ...and 1 more figures