Efficient Code LLM Training via Distribution-Consistent and Diversity-Aware Data Selection

The emergence of large language models (LLMs) has significantly advanced the intelligence-driven transformation of software engineering, particularly demonstrating breakthrough capabilities in code generation and program comprehension tasks. The open-source community has developed multiple high-performance code LLMs, with notable representatives including CodeLlama Roziere et al. (2023), DeepSeek-Coder Guo et al. (2024), CodeGemma Team et al. (2024), and Qwen2.5-Coder Hui et al. (2024). CodeLlama, built upon the Llama 2 Touvron et al. (2023) architecture through continued pre-training on 500B code tokens, achieves performance comparable to commercial models in standardized benchmarks. The DeepSeek-Coder series, trained from scratch on 2T high-quality multilingual code corpora, demonstrates significant superiority over the same-scale counterparts in benchmark evaluations. The Qwen2.5-Coder series are built upon the Qwen2.5 Yang et al. (2024) architecture and continue training on a vast corpus of over 5.5 trillion tokens, surpassing closed-source models like GPT-4 Achiam et al. (2023) across multiple benchmarks and establishing new state-of-the-art records. These open-source models, through architectural innovations and training strategy optimizations, not only deliver highly customizable code intelligence solutions but also construct efficient technical infrastructure for both academic research and industrial applications.

Instruction fine-tuning has been demonstrated as a crucial approach to enhance model performance and align models with human preferences. The study Chung et al. (2024) indicates that scaling both the number of tasks and the model size through instruction fine-tuning yields performance improvements across various model classes. For instance, WizardCoder Luo et al. (2023) leverages the Evol-Instruct Xu et al. (2024) method to iteratively evolve the complexity of the CodeAlpaca Chaudhary (2023) dataset. This approach results in a fine-tuning dataset consisting of approximately 78K highly complex programming instructions, thereby improving the performance of CodeLLama-Python-34B to 73.2% on the HumanEval benchmark. Similarly, Magicoder Wei et al. (2024) proposes the OSS-Instruct approach, which generates highly diverse instruction data by using open-source code snippets. This approach successfully generated a dataset containing approximately 75K entries. Additionally, WaveCoder Yu et al. (2024) makes full use of open-source code data through a carefully designed generator-discriminator data synthesis framework to generate high quality and diverse instruction data in multi-task scenarios.

Instruction fine-tuning typically requires large amounts of data, but research such as LIMA Zhou et al. (2023a) has pointed out that data quality is more crucial than quantity. Therefore, selecting the most valuable data to improve training efficiency has become a focal point of research. DEITA Liu et al. (2024b) focuses on the complexity, quality and diversity of the data, designing a multifaceted approach to select instruction data. Based on the Evol-Instruct technique, ChatGPT is used to augment the instructions, and the instructions are then evaluated for complexity and quality by specially trained scorers. Quantifiable metrics like PPL Ankner et al. (2024), IFD Li et al. (2024b), and Superfiltering Li et al. (2024a) focus on identifying hard samples that are difficult to learn. DQ Zhou et al. (2023b) integrates data distillation and coreset selection Iyer et al. (2021) techniques, emphasizing the selection of diverse data. Some methods Chen et al. (2024); Lu et al. (2024); Xu et al. (2023) that rely on external oracles use powerful language models, such as ChatGPT, for data selection. However, due to cost constraints, utilizing external oracles is not always feasible.

Overall, many data selection algorithms are primarily designed for general tasks, and the integration of parametric models into code data selection remains underexplored.

Given a large instruction tuning dataset $D=\{x_{1},x_{2},\ldots,x_{n}\}$, where each $x_{i}=(I_{i},C_{i})$ represents an individual instruction-code pair, our goal is to select a subset $S_{\pi}^{m}\subset D$ of size $m$ using a selection strategy $\pi$. The performance of the model after fine-tuning on $S_{\pi}^{m}$ is denoted by $P(S_{\pi}^{m})$ and is used to evaluate the effectiveness of the selected subset. The optimal strategy $\pi^{*}$ under a fixed budget $m$ is defined as:

During data selection, we consider two key factors: On one hand, the selected subset $S$ should have the distribution as close as possible to that of the original dataset $D$. On the other hand, the subset $S$ should maintain the diversity. By balancing these two aspects, we aim to select a representative subset that covers edge cases in the original dataset, thereby enhancing model performance.

Compared to selecting in the discrete data space, it is more efficient and feasible to conduct data selection in the feature space. We map each data sample $x_{i}$ to the high-dimensional feature space using a feature encoder $E$, resulting in $f_{i}=E(x_{i})$ which is then normalized. After normalization, we obtain the feature set of the dataset $D$ as $F_{D}=\{f_{i}\}_{i\in[n]}$, whose distribution is denoted by $p_{F_{D}}$.

Similarly, for the selected subset $S$, we define its feature set as $F_{S}=\{f_{j}\}_{j\in[m]}$, with its corresponding distribution $p_{F_{S}}$. Our objective is to find the optimal selection strategy $\pi^{*}$ as follows:

where $M(\cdot,\cdot)$ measures the distance between two feature distributions, $R(\cdot)$ evaluates the diversity of the selected subset, and $\lambda$ is a scaling factor that balances two terms. The first term in Equation 2 aims to align the distributions, while the second term ensures the diversity within the subset.

It’s challenging to directly optimize the selection strategy $\pi$ in the discrete space, so we adopt a parametric model $p_{\theta_{S}}$ to approximate $p_{F_{S}}$, where $\theta_{S}=\{\theta_{S}^{j}\}_{j\in[m]}$ are the continuous parameters. Each optimized parameter $\theta_{S}^{j}$ corresponds to the feature of a selected sample $f_{j}$, and we select the feature $f_{j}$ that is closest to $\theta_{S}^{j}$. Thus, the optimization objective is written as follows:

The key distinction between the features $F_{S}=\{f_{j}\}$ and the parameters $\theta_{S}=\{\theta_{S}^{j}\}$ is that $f_{j}$ is a discrete feature corresponding to a data sample, whereas $\theta_{S}^{j}$ is continuous in the feature space. By optimizing $\theta_{S}$, we search for ideal samples in the feature space that cover the original data distribution while maintaining dispersion.

The first term represents the distribution matching term $M(p_{F_{D}},p_{\theta_{S}})$, which measures the similarity between the feature distribution of the dataset and the parametric model. It is defined as:

where $sim(\cdot,\cdot)$ is a similarity function, such as cosine similarity, and $\tau$ is a temperature parameter that controls the sensitivity of similarity measure. This term encourages the parameters to be well-spread across the feature space, effectively covering the distribution of the original features.

The second term is the diversity regularization $R(\theta_{S})$, which promotes diversity within the selected subset by minimizing the similarity between the selected samples. It is defined as:

We jointly optimize the following loss function to achieve the goal in Equation 3, where $\lambda$ is set to 1 by default.

We optimize the loss function in Equation 7 using gradient descent. Upon completion of the optimization, we find the features $\{f_{j}\}_{j\in[m]}$ that exhibit the highest similarity to $\theta_{S}^{j}$.

Finally, we collect the data samples corresponding to these selected features to form the instruction subset $S_{\pi}^{m}$, which will be used for fine-tuning.

Algorithm 1 illustrates the implementation details of the data selection process. We use the pre-trained model all-mpnet-base-v2¹¹1https://huggingface.co/sentence-transformers/all-mpnet-base-v2 from the sentence-transformers Reimers and Gurevych (2019) library as the feature encoder. This model has been trained on over 1 billion text pairs, effectively capturing the semantic information in instruction texts. For each data sample $x_{i}=(I_{i},C_{i})$, we encode the instruction text $I_{i}$ to obtain the feature vector $f_{i}=E(I_{i})\in\mathbb{R}^{768}$, as the instruction text fully defines the task semantics. All features are then processed using L2 normalization, forming the normalized feature set $F_{D}=\{f_{i}\}_{i\in[n]}$.

Before optimizing the parametric model, the parameters $\theta_{S}$ are initialized by randomly sampling from the feature set $F_{D}$. During each iteration, to avoid memory overflow, we calculate the similarity between sample features and parameters in a batch-wise manner. Subsequently, we update $c_{i}$ according to Equation 4. Finally, we compute the loss function in Equation 7 and update the parameters $\theta_{S}$ using gradient descent.

When the optimization process is finished, we find the sample feature $f_{j}$ that exhibit the highest similarity to $\theta_{S}^{j}$ according to Equation 8. The corresponding original sample $x_{j}$ is selected for subsequent fine-tuning. This process ensures that the selected subset is highly representative and covers the edge cases in the original dataset, ultimately providing high-quality data support for the fine-tuning process.

We use three open-source instruction datasets, including Evol-Instruct-Python-26K²²2https://huggingface.co/datasets/mlabonne/Evol-Instruct-Python-26k, CodeExercise-Python-27K³³3https://huggingface.co/datasets/codefuse-ai/CodeExercise-Python-27k, and OSS-Instruct-75K⁴⁴4https://huggingface.co/datasets/ise-uiuc/Magicoder-OSS-Instruct-75K. The Evol-Instruct-Python-26K dataset is the Python subset of the Evol-Instruct-80K⁵⁵5https://huggingface.co/datasets/nickrosh/Evol-Instruct-Code-80k-v1 dataset, and its construction follows an iterative evolution strategy. Based on the CodeAlpaca Chaudhary (2023) instruction dataset, multi-round complexity enhancement operations are applied to the original problems using ChatGPT. The CodeExercise-Python-27K dataset is generated using the Camel Li et al. (2023a) framework, covering hundreds of Python-related topics including basic syntax and data structures, algorithm applications, database queries, machine learning, and more. The OSS-Instruct-75K dataset utilizes ChatGPT to generate programming problems and their corresponding solutions. Its uniqueness lies in the use of open-source code snippets as guidance for generation. To maintain consistency in the programming language, we filter this dataset and only retain Python-related entries. Finally, we merge these three datasets to obtain a mixed dataset, Mix-Python-92K. We use this mixed dataset for training.

HumanEval/HumanEval+. HumanEval Chen et al. (2021) comprises 164 Python programming problems designed to assess the ability of code generation models. Each problem is accompanied by approximately 9.6 test cases to check whether the generated code works as expected. HumanEval has become one of the most widely used benchmarks for evaluating the performance of code LLMs, making it a key tool in the field of artificial intelligence for coding. To enhance the rigor of the evaluation, HumanEval+ Liu et al. (2023) builds upon the original dataset by significantly increasing the number of test cases through the use of LLMs and mutation strategies, resulting in a more comprehensive evaluation benchmark.

MBPP/MBPP+. MBPP Austin et al. (2021) consists of approximately 1,000 Python programming challenges sourced from a crowd of contributors, targeting beginners in programming and focusing on core principles and the usage of the standard library. Each challenge includes a description, a solution and three tests to verify the accuracy. To improve the reliability of the benchmark, MBPP+ Liu et al. (2023) extends the original dataset by incorporating a subset of hand-verified problems from the MBPP-sanitized dataset, ensuring that the tasks are well-defined and unambiguous. This enhances the benchmark’s reliability and suitability for more rigorous evaluations.

Pass@k. We use the Pass@k metric Chen et al. (2021) to enhance the reliability of our evaluation process. We count the total number of generated samples that successfully passing all test cases, denoted as $k$, to compute the Pass@k.

where $n$ is the total number of generated samples for each problem and $c$ is the number of correct generated code samples passing all test cases($n>k\geq c$). In subsequent experiments, we compute the Pass@1 (%) metric using greedy decoding.

In our experiments, we use DeepSeekCoder-Base-6.7B as the base model and conduct training on eigth NVIDIA A800-80GB GPUs using PyTorch’s Fully Sharded Data Parallel (FSDP) module for 3 epochs. Specifically, we employ the AdamW optimizer with a learning rate of 5e-5, a cosine learning rate scheduler, and 100 warmup steps. The maximum sequence length per batch is set to 4096 tokens with a global batch size of 512. To improve training efficiency, we utilize the Dynamic Pack Lv et al. (2025). Model evaluation is conducted using the EvalPlus Liu et al. (2023) library.

For optimizing the parametric model, we adopt the same hyperparameter settings as in ActiveFT Xie et al. (2023). Specifically, we use the Adam optimizer with 300 optimization iterations, a learning rate of 0.001, and set the temperature parameter $\tau$ in Equation 7 to 0.07.

We compared various sampling methods, with the experimental results presented in Table 1. when sampling 10K samples, our method outperforms all baseline approaches across all benchmarks. Specifically, our method achieves 69.5% on HumanEval, 63.4% on HumanEval+, 77.2% on MBPP, and 63.2% on MBPP+. These results demonstrate the importance of considering both distribution consistency and diversity in subset. Our approach effectively balances these two factors, enabling the identification of high-quality samples that contribute to enhancing model performance.

K-Center Sener and Savarese (2018) and DQ Zhou et al. (2023b) focus on maximizing diversity within the subset. While these approaches increase inter-sample diversity, they do not guarantee alignment with the original dataset’s distribution. Both DQ and K-Center achieve 64.6% on HumanEval, which is lower than our method. This suggests that while diversity is important, it must be balanced with distribution consistency to fully leverage the dataset’s structure. PPL and IFD prioritize the selection of complex samples, aiming to enhance model performance by focusing on challenging instances. However, these methods perform poorly compared to random sampling on several benchmarks. For instance, IFD shows a substantial performance drop on MBPP and MBPP+, indicating that complexity-focused sampling may deviate from the original distribution. DEITA combines complexity and quality assessment, yielding competitive results. However, it requires the training of two additional scoring models, resulting in higher computational costs and longer training times. In contrast, our method achieves superior performance with significantly lower computational overhead, making it more efficient and scalable.

In summary, our method establishes distribution consistency as the foundation while incorporating diversity constraints, constructing a more comprehensive and efficient data selection strategy. This approach enables the identification of high-quality training samples and consistently improves the model’s overall performance.

We also compared the sampling efficiency of different methods, with the results presented in Table 1. The sampling process can be divided into two main phases: (1) the processing phase, which involves data preprocessing steps such as feature extraction and model scoring, and (2) the sampling phase, where specific strategies are applied for sample selection.

Among all compared methods, DQ exhibits the lowest sampling efficiency. Its processing phase requires nearly 20 hours to partition the dataset into non-overlapping bins, a duration that even exceeds the model training time. Although DQ achieves competitive performance, its high time cost during the sampling process represents a significant disadvantage. In contrast, the DEITA, PPL, and IFD methods show better efficiency but still require forward inference for each data sample. Specifically, DEITA relies on two 13B-parameter scoring models, while PPL and IFD depend on the model being trained for scoring. As a result, their sampling time increases linearly with the size of the dataset, resulting in poor scalability. The K-Center method reduces sampling time to 42.5 minutes by iteratively computing the similarity between candidate samples and the selected subset. However, its reliance on distance-based computations fundamentally limits further efficiency improvements. In contrast, our method achieves the highest sampling efficiency, requiring only 13.5 minutes. This advantage is attributed to the learnable parametric model for data selection, which significantly reduces computational overhead by optimizing the loss function in Equation 7 that ensures both distribution consistency with the original dataset and diversity within the selected subset.

In summary, our method maintains model performance while significantly improving sampling efficiency and scalability. Compared to other methods, our approach minimizes sampling time, enabling the efficient processing of large-scale datasets.

To investigate the impact of sampling quantity on model performance, we compared the experimental results with 5K, 10K, 15K, 20K, 25K, and the full 92K dataset, as shown in Figure 2. The experimental results reveal that the performance curves of all methods generally exhibit an initial increase followed by a decline, indicating the presence of low-quality data within the dataset. Moderate sampling (e.g., 10K data) effectively filters out low-quality samples, thus improving model performance, while excessive sampling (more than 25K data) may introduce noise, leading to performance degradation. When trained on the full dataset, the model achieved 67.1% on HumanEval and 74.9% on MBPP.

The experimental results demonstrate that our method consistently outperforms others across different sampling quantities, particularly peaking at 10K samples with 69.5% on HumanEval and 77.2% on MBPP. Compared to training on the full dataset, our method represents improvements of 2.4% and 2.3% respectively. These results validate the effectiveness of our data selection strategy, showing that introducing the parametric model into the data selection process can significantly enhance model performance.

As the sampling quantity increases from 5K to 20K, PPL and IFD show performance improvements on both HumanEval and MBPP benchmarks, indicating that “non-complex yet important” data samples also play a key role in improving performance. The improvement observed in these methods suggests that strategies focused solely on complex data sampling fail to yield the best performance, as they neglect simpler yet impactful samples. In contrast, our method considers both distribution consistency and diversity, ensuring that the sampled dataset not only includes complex data but also retains simpler yet essential samples, resulting in a more balanced performance gains.

The performance curves of DEITA, DQ and K-Center exhibit similar trends, with performance declining after exceeding 15K samples. Notably, K-Center maintains relatively good performance on the MBPP benchmark but fluctuates significantly on the HumanEval benchmark. This reflects the limitations of diversity-driven sampling strategies in programming semantic understanding. Purely diversity-based approaches neglect distribution consistency, potentially resulting in unstable performance on different benchmarks.

Overall, our method achieves stable and superior performance across different sampling quantities by balancing both distribution consistency and diversity. In particular, with a 10K sample size, the model reaches optimal performance. This demonstrates that our approach not only excels in improving performance but also effectively avoids the noise introduced by excessive sampling, thus enhancing the model’s generalization ability. Our data selection strategy enables the identification of high-quality samples, significantly boosting the model’s overall performance.

We compared the performance of different code LLMs, as shown in Table 2. WizardCoder-CL was trained on CodeLlama-Python using 78K instruction samples, while WaveCoder-DS and Magicoder-DS were both trained on DeepSeekCoder-Base with 20K and 75K samples, respectively. Despite these models utilizing much larger training datasets, our approach achieves superior performance across all benchmarks using only 10K carefully selected samples.

These results highlight that data quality is more important than quantity. By developing an effective data selection strategy, we can significantly enhance model performance. Our experiments confirm that employing the parametric model for data selection constitutes an efficient approach. It not only identifies high-quality training samples but also substantially improves training efficiency by reducing the number of unnecessary or low-quality samples. Compared to models relying on massive training data, our method delivers better performance with far fewer data, demonstrating that our data selection strategy ensures competitive performance while effectively reducing computational resource consumption.