[Paper Review] LLM-Assisted Code Cleaning For Training Accurate Code Generators

Naman Jain, Tianjun Zhang+4 more|arXiv (Cornell University)|Nov 25, 2023

Software Engineering Research4 citations

TL;DR

This paper proposes a novel LLM-assisted data-cleaning pipeline that improves code generation performance by transforming training data to be more readable and modular. By renaming variables, decomposing complex code into helper functions, and inserting natural-language plans, the method boosts CodeLLaMa-7B's accuracy by up to 30% compared to original datasets, even outperforming larger closed-source models like AlphaCode when trained on just 15% of the cleaned data.

ABSTRACT

Natural language to code generation is an important application area of LLMs and has received wide attention from the community. The majority of relevant studies have exclusively concentrated on increasing the quantity and functional correctness of training sets while disregarding other stylistic elements of programs. More recently, data quality has garnered a lot of interest and multiple works have showcased its importance for improving performance. In this work, we investigate data quality for code and find that making the code more structured and readable leads to improved code generation performance of the system. We build a novel data-cleaning pipeline that uses these principles to transform existing programs by 1.) renaming variables, 2.) modularizing and decomposing complex code into smaller helper sub-functions, and 3.) inserting natural-language based plans via LLM based transformations. We evaluate our approach on two challenging algorithmic code generation benchmarks and find that fine-tuning CodeLLaMa-7B on our transformed modularized programs improves the performance by up to 30% compared to fine-tuning on the original dataset. Additionally, we demonstrate improved performance from using a smaller amount of higher-quality data, finding that a model fine-tuned on the entire original dataset is outperformed by a model trained on 15% of our cleaned dataset. Even in comparison to closed-source models, our models outperform the much larger AlphaCoder models.

Motivation & Objective

• Improve code generation performance by enhancing data quality beyond functional correctness.
• Investigate how readability, modularity, and structural clarity in training data affect model performance.
• Develop an automated, LLM-driven pipeline to transform existing code datasets into higher-quality, more structured versions while preserving functional correctness.
• Demonstrate that high-quality data can improve data efficiency, reducing the need for large-scale training sets.
• Establish a framework for data-cleaning that identifies and applies key programming best practices as training data improvements.

Proposed method

• Use instruction-tuned LLMs to perform three iterative transformations: variable renaming for contextual relevance, code modularization into smaller helper functions, and insertion of top-down natural-language plans as comments.
• Apply each transformation using natural language prompts and input-output examples to ensure functional equivalence between original and transformed programs.
• Leverage the relative ease of editing over generating code from scratch, using LLMs to refine existing programs rather than generate from scratch.
• Construct parallel datasets by transforming programs from benchmark datasets (APPS and Code-Contests) into cleaned, modularized, and planned versions.
• Fine-tune CodeLLaMa-7B on both original and transformed datasets to evaluate performance gains.
• Use input-output examples as verification to ensure that transformed programs produce the same outputs as original programs.

Figure 1: The overview of our code cleaning approach . We apply instruction-tuned LLMs to transform existing datasets by providing natural language prompts and use input-output examples to maintain function equivalence between original and transformed programs. Our cleaning approach works in three s

Experimental results

Research questions

• RQ1 How does improving code readability and modularity through data cleaning affect the performance of code generation models?
• RQ2 Can a smaller, higher-quality dataset outperform a larger, lower-quality dataset in fine-tuning code generation models?
• RQ3 To what extent do LLM-based transformations such as variable renaming, modularization, and planning annotations improve model generalization and functional correctness?
• RQ4 Does the inclusion of natural-language plans in code improve model performance compared to raw code alone?
• RQ5 Can data-cleaning techniques that emphasize programming best practices lead to better data efficiency in code generation tasks?

Key findings

• Fine-tuning CodeLLaMa-7B on the cleaned, modularized dataset improved performance by up to 30% compared to fine-tuning on the original dataset.
• A model fine-tuned on just 15% of the cleaned dataset outperformed a model fine-tuned on the entire original dataset, demonstrating significant data efficiency gains.
• The cleaned dataset, which includes better-named variables, modular functions, and top-down planning comments, led to more structured and readable code in the model's outputs.
• Even when compared to larger closed-source models like AlphaCode, the fine-tuned CodeLLaMa-7B model trained on the cleaned data achieved superior performance on algorithmic code generation benchmarks.
• The LLM-based cleaning pipeline successfully transformed existing programs into higher-quality versions while preserving functional correctness, as verified via input-output testing.
• The method demonstrates that data quality—particularly in structure and readability—has a measurable and substantial impact on code generation performance, independent of dataset size.

Figure 3: Effect of quality on data-efficiency of the model . Finetuning on 15% of clean $\mathcal{D}_{modular}$ dataset results in similar performance as finetuning on the entire $\mathcal{D}_{original}$ dataset.