[Paper Review] Code Less, Align More: Efficient LLM Fine-tuning for Code Generation with Data Pruning

Yun-Da Tsai, Mingjie Liu+1 more|arXiv (Cornell University)|Jul 6, 2024

Mathematics, Computing, and Information Processing4 citations

TL;DR

This paper proposes a data pruning framework for efficient fine-tuning of code-specific LLMs by clustering synthetic code instruction-code pairs and selecting representative samples using multiple metrics. Experiments show that training on just 10% of the original data preserves 96.1% of HumanEval performance and improves MBPP accuracy by up to 3.5%, demonstrating significant efficiency gains without performance loss or degradation.

ABSTRACT

Recent work targeting large language models (LLMs) for code generation demonstrated that increasing the amount of training data through synthetic code generation often leads to exceptional performance. In this paper we explore data pruning methods aimed at enhancing the efficiency of model training specifically for code LLMs. We present techniques that integrate various clustering and pruning metrics to selectively reduce training data without compromising the accuracy and functionality of the generated code. We observe significant redundancies in synthetic training data generation, where our experiments demonstrate that benchmark performance can be largely preserved by training on only 10% of the data. Moreover, we observe consistent improvements in benchmark results through moderate pruning of the training data. Our experiments show that these pruning strategies not only reduce the computational resources needed but also enhance the overall quality code generation.

Motivation & Objective

• To address the high computational cost and redundancy in large-scale synthetic code datasets used for fine-tuning code LLMs.
• To develop a scalable, efficient data pruning strategy tailored specifically for code generation tasks.
• To investigate whether reducing training data through clustering and pruning can maintain or even improve model performance.
• To evaluate the impact of different clustering algorithms and pruning metrics on downstream code generation benchmarks.
• To provide a practical, low-resource alternative to full-data fine-tuning for code LLMs.

Proposed method

• Clustering data samples based on both instruction and code solution embeddings using KMeans and HDBSCAN.
• Applying dimensionality reduction via PCA to reduce embedding dimension from 1536 to 10, improving clustering speed by ~12x.
• Selecting representative samples from each cluster using multiple pruning metrics, including diversity and uncertainty scores.
• Combining instruction and code embeddings as input for clustering to improve representativeness and selection quality.
• Using a hybrid approach: clustering first, then pruning based on metric scores to ensure diversity and informativeness.
• Evaluating performance on standard benchmarks (HumanEval and MBPP) after fine-tuning on pruned datasets.

Experimental results

Research questions

• RQ1 Can significant data reduction via clustering and pruning preserve or improve code generation performance in fine-tuned LLMs?
• RQ2 How does the choice of clustering algorithm (e.g., KMeans vs. HDBSCAN) affect the effectiveness of data pruning?
• RQ3 What is the optimal balance between data size and model performance when pruning synthetic code datasets?
• RQ4 Does combining instruction and code embeddings yield better clustering and pruning outcomes than using either alone?
• RQ5 To what extent does dimensionality reduction (e.g., PCA) impact the efficiency and performance of the pruning pipeline?

Key findings

• Training on only 10% of the original synthetic dataset preserves 96.1% of the HumanEval pass@1 score and 98.5% of the MBPP pass@1 score compared to full-data training.
• Moderate pruning leads to consistent performance improvements: up to 2.7% absolute gain on HumanEval and 3.5% on MBPP compared to full-data fine-tuning.
• The clustering algorithm is more critical than the pruning metric, as KMeans with PCA achieved the best trade-off between speed and performance.
• Using both instruction and code embeddings for clustering yields slightly better results than using either alone, indicating the importance of multimodal representation.
• Applying PCA reduced embedding dimension from 1536 to 10, cutting clustering runtime from 1307 to 183 seconds, with only minor performance degradation (≤0.6% on MBPP, ≤4.3% on HumanEval).
• HDBSCAN clustering was too slow to complete within 4 hours and was therefore excluded from final evaluation, highlighting the importance of computational efficiency in the pipeline.