Enhancing High-Quality Code Generation in Large Language Models with Comparative Prefix-Tuning

TL;DR

This work introduces comparative prefix-tuning, a lightweight method that trains a single prefix to steer large language models toward high-quality code generation without compromising functional correctness. By using paired high-quality and low-quality code samples and a sequence-level ranking loss, the approach emphasizes quality-relevant patterns while constraining changes to preserve behavior, aided by a data-construction pipeline and masked losses. Experiments on Code Llama 7B (and generalization to Phi-2 and Starcoder2) show substantial improvements in pylint-based code quality (often over 30% mean gains) with maintained or improved functional correctness on APPS and HumanEval benchmarks, plus a favorable human evaluation. The method is computationally efficient (≈0.05% trainable parameters) and generalizes across models, suggesting practical utility for producing maintainable, standards-compliant code in real-world coding tasks.

Abstract

Large Language Models (LLMs) have been widely adopted in commercial code completion engines, significantly enhancing coding efficiency and productivity. However, LLMs may generate code with quality issues that violate coding standards and best practices, such as poor code style and maintainability, even when the code is functionally correct. This necessitates additional effort from developers to improve the code, potentially negating the efficiency gains provided by LLMs. To address this problem, we propose a novel comparative prefix-tuning method for controllable high-quality code generation. Our method introduces a single, property-specific prefix that is prepended to the activations of the LLM, serving as a lightweight alternative to fine-tuning. Unlike existing methods that require training multiple prefixes, our approach trains only one prefix and leverages pairs of high-quality and low-quality code samples, introducing a sequence-level ranking loss to guide the model's training. This comparative approach enables the model to better understand the differences between high-quality and low-quality code, focusing on aspects that impact code quality. Additionally, we design a data construction pipeline to collect and annotate pairs of high-quality and low-quality code, facilitating effective training. Extensive experiments on the Code Llama 7B model demonstrate that our method improves code quality by over 100% in certain task categories, while maintaining functional correctness. We also conduct ablation studies and generalization experiments, confirming the effectiveness of our method's components and its strong generalization capability.

Figures (13)

• Figure 1: Visualization of our proposed task (③) vs. code review (①), code refinement (②), and code quality analysis (④) tasks.
• Figure 2: The overall framework of our comparative prefix-tuning method for improving code quality in LLMs. The prefix is optimized using pairs of high-quality and low-quality code samples, guiding the model to generate code that adheres to coding standards and best practices.
• Figure 3: Comparison of model architectures: our method (d) vs. three related methods (a, b and c) for high-quality code generation
• Figure 4: The self-attention computation process with prefix optimization. The prefixes modify the self-attention mechanism by incorporating attribute-specific information into the key and value vectors.
• Figure 5: A data instance consisting of an instruction $I$, a high-quality code sample $x_a$, a low-quality code sample $x_b$, and their corresponding mask vectors $m_a$, $m_b$.