Diversify and Conquer: Diversity-Centric Data Selection with Iterative Refinement

## Abstract

Abstract Finetuning large language models on instruction data is an important step in enriching the knowledge learned during pre-training and improving instruction-following capabilities.
As the number of instruction datasets continues to grow, selecting the right data to achieve optimal results becomes increasingly important.
In this work, we ask a prominent question: How can we determine the optimal subset of data for effective training? While much of the existing research primarily emphasizes local criteria, such as instance quality, for subset selection, we argue that a global approach focused on data diversity is more critical.
Our approach utilizes k 𝑘 k italic_k -means clustering to ensure that the selected subset effectively represents the full dataset.
We propose an iterative refinement method inspired by active learning techniques to resample instances from clusters, with the importance and sampling weight of each cluster being reassessed in every training iteration.
This method allows us to reduce the effect of outliers and automatically filter out clusters containing low-quality data.
Through extensive evaluation across natural language reasoning, general world knowledge, code and math reasoning tasks, and by fine-tuning models from various families, we observe consistent improvements, achieving a 7% increase over the random selection and a 3.8% improvement over state-of-the-art sampling methods.
Our work highlights the significance of diversity-first sampling when finetuning LLMs to enhance performance across a broad array of evaluation tasks. Our code is available at https://github.com/for-ai/iterative-data-selection .

## 1 Introduction

Large language models are trained on vast amounts of data scraped from the internet, containing a wide range of content qualities. (Penedo et al., ; Chen et al., ; Laurenccon et al., ; Marion et al., ).
Models develop a broad understanding of language and acquire general knowledge from the unstructured data in this pretraining phase (Da et al., ; Chang et al., ) and align with user intent in the finetuned stage using instruction datasets which consists of a more structured format of question and response pairs (Chung et al., ; Taori et al., ; Li et al., ; Üstün et al., ).
Recent years have seen substantial efforts to create datasets using various manual (Conover et al., ; Köpf et al., ; Singh et al., ) and synthetic (Taori et al., ; Wang et al., ; Shimabucoro et al., ) methods, making it increasingly challenging to determine which dataset is best suited for downstream tasks.
A crucial question regarding the scalability of finetuning LLMs is: “what is the optimum subset of data that allows for efficient training and captures aspects of the data relevant to downstream tasks?”

Instances in a dataset contribute to a model’s learning process with varying degrees of impact, affecting the model’s performance and generalization (Sorscher et al., ; Chen et al., ).
While recent research has predominantly emphasized local features, such as the quality of individual instances for subset selection, we argue that prioritizing a global feature —diversity—yields greater benefits.
When selecting a subset of instances, we manage computational complexity while balancing the trade-off between diversity and representativeness (Zhou et al., ), ensuring that the subset captures the underlying data distribution (Ivison et al., ; Wang et al., ).
Preserving a high level of sample diversity during finetuning is crucial for improving generalization capabilities (Zhang et al., ; Yue et al., ).
Wang et al. () revealed that using a range of instruction datasets can boost downstream tasks.
Wang et al. () provided a theoretical analysis using determinantal point processes to underscore the significance of diversity in the selection of subsets.
However, ensuring diversity during sampling is difficult, and current methodologies fall short of fully addressing this challenge.
Most scoring-based subset selection methods prioritize sample quality and characteristics and subsequently apply a diversity filter (Liu et al., ; Xia et al., ).
Still, since diversity is inherently a global property, addressing it only in the second step limits its effectiveness because it lacks a comprehensive view of the entire collection.
This limitation often arises because assessing the data collection globally is computationally expensive (Bukharin & Zhao, ).

In this work, we propose a scalable iterative sampling and refinement method to efficiently select a subset of instruction data and maximize the diversity of samples.
We iteratively refine the sample selection using early training signals from the fine-tuning model and proceed with continued fine-tuning.
With the same training budget, we achieve substantial improvements over fixed sampling approaches and previous state-of-the-art data selection methods.
We evaluate the finetuned models on a wide range of tasks, including question answering, math, reasoning, and code, and show consistent improvements over baselines.
Overall, our experiments and analyses demonstrate that by sampling a small subset of data, we achieve performance improvements of up to 7% over random selection and 3.8% over the previous sampling methods on a wide variety of tasks.
In summary, our contributions are as follows:

- •
We systematically analyze various clustering and sampling methods and demonstrate that $k$-means clustering is particularly effective for selecting an optimal, diverse subset of instruction data, especially when paired with a quality sampling step.
- •
Our simplest variant, which involves efficiently clustering data points and randomly sampling from each cluster, already achieves performance on par with advanced state-of-the-art sampling techniques, without the need for costly LLM scoring. This supports our hypothesis on the importance of diversity and the representativeness of the sampling process.
- •
We further propose an iterative clustering algorithm that simultaneously combines the learning feedback from the training model and optimizes for diversity based on data distribution for effective instruction tuning. This method outperforms previous approaches on all downstream tasks.

We release the code and the data artifacts used in our experiments to facilitate reproducibility and future research.

## 2 Methodology

### 2.1 Static Data Selection

Given a large and diverse set of instruct data $\mathcal{D}=\{x_{1},x_{2},\dots,x_{n}\}$, we select a subset $\mathcal{D}^{\prime}$ with budget $b\in\mathbb{N}^{+}$, where $b=|\mathcal{D}^{\prime}|\ll|\mathcal{D}|$ and finetune a language model and evaluate a selection of downstream tasks.
This subset should be a representative sample of the training data, maintaining high quality and offering a diverse range of examples.
We propose to define the problem of sample selection for training data of a language model as a clustering problem with clustering objectives where we want to group similar samples together and separate dissimilar samples into different clusters.
We explore various sampling methods to ensure the inclusion of optimal samples from different clusters.

For clustering purposes, we consider two main clustering objectives: $k$-center and $k$-means. Both of these two objectives are metric clustering where we are given a set of points $\mathcal{D}$ with distance metric $d:\mathcal{D}\times\mathcal{D}\rightarrow\mathbb{R}_{\geq 0}$ and the goal is to pick a set of centers $\mathcal{C}=\{c_{1},\dots,c_{k}\}\subseteq\mathcal{D}$ of size at most $k$. For $k$-center, we want to pick $\mathcal{C}$ such that the maximum distance of data points to centers is minimized. More precisely, in $k$-center, we want to minimize

Figure: Figure 1: Our proposed clustering ($k$MQ) and two sampling methods: We visualize our static data selection with $k$MQ, as proposed Section 2.1 and the iterative data selection pipeline where we refine the selection criteria and resample new instances in each iteration, as proposed in Section 2.2.
Refer to caption: x1.png

$$ $\max_{x_{i}\in\mathcal{D}}d(x_{i},\mathcal{C})$ (1) $$

where $d(x_{i},\mathcal{C})=min_{c_{j}\in\mathcal{C}}d(x_{i},c_{j})$ is the distance of point $i$ to the closest center in $\mathcal{C}$. The $k$-means objective is similar to $k$-center objective but instead of looking at the $l_{\infty}$ norm of the vector that defines the distance of points to $\mathcal{C}$, we look at the $l_{2}$ norm of this vector. More precisely, in $k$-means, we want to minimize

$$ $\sum_{x_{i}\in\mathcal{D}}d^{2}(x_{i},\mathcal{C})$ $$

Based on this objective and given the set of centers $\mathcal{C}={c_{1},\ldots,c_{k}}$, we define $\mathcal{D}_{j}$ as the subset of data points in $\mathcal{D}$ that are closest to center $c_{j}$ and belong to the $j^{\text{th}}$ cluster:

$$ $\mathcal{D}_{j}=\{x_{i}\in\mathcal{D}\mid d(x_{i},c_{j})\leq d(x_{i},c_{l}) \text{ for all }l\neq j,l=1,\ldots,k\}$ (2) $$

where $d(x_{i},c_{j})$ is the distance between data point $x_{i}$ and center $c_{j}$.

Beyond the clustering, the next step concerns how to sample data from the clusters with a fixed budget of $m$. We investigate both random sampling and a more informed, quality-based sampling approach.
For the quality-based sampling, inspired by the previous approaches (Liu et al., ; Bukharin & Zhao, ), we propose $k$-means-quality ($k$MQ), where we first perform the traditional $k$-means by clustering the instruction data into $k$ centroids, in which $k\ll b$, and sample data from each cluster to form $\mathcal{D}^{\prime}$.
Note that we assign each cluster a budget proportional to its size ($b_{j}=\frac{|X_{j}|}{|X|}\cdot b$) and draw samples within each cluster based on the probability weighted by the quality score.
We use the same scoring method introduced by Liu et al. () to obtain quality scores, enabling a fair comparison of the hypotheses regarding the importance of diversity-first versus quality-first sampling.
More concretely, we sample:

$$ $\{x_{1},x_{2},\ldots,x_{b_{j}}\}\sim\text{Multinomial}(\mathcal{D}_{j},\{p(x \mid q)\}_{x\in\mathcal{D}_{j}})$ (3) $$

where $\{x_{1},x_{2},\ldots,x_{b_{j}}\}$ is the data sampled from cluster $\mathcal{D}_{j}$ with replacement, $b_{j}$ is the budget assigning to the $j^{\text{th}}$ cluster and $p(x\mid q)$ is the probability of picking $x$, weighted by its quality $q$.

Additionally, we take a systematic approach to studying the role of diversity and show the importance of the choice of $k$ in affecting downstream performance, which has been overlooked in previous works (see analysis in Section 4.3).

### 2.2 Iterative Data Selection

Figure: Algorithm 1 Iterative Data Selection Pipeline

In the previous section, we introduced a two-step approach: sampling a fixed subset of data first and finetuning a model on it.
The sampling and finetuning steps are performed independently without any information exchange between the two steps.
However, the initial stages of finetuning can offer insights into how individual data points influence the learning process (Anthony et al., ; Muldrew et al., ).
Here, we investigate whether we can improve our sampling method by incorporating early training feedback into the selection process.
We accomplish this by periodically increasing the weight of clusters from which the model learns well while decreasing the weight of clusters that the model finds difficult to generalize.

The motivation is twofold: (1) Not all data clusters possess the same level of quality and impact. We further analyze the distribution and quality scores across clusters, revealing significant disparities (see analysis in §4.4).
This indicates that some clusters are notably better quality, while others predominantly consist of low-quality data.
(2) From a curriculum learning perspective, models can develop different skills and knowledge at varying rates (Xu et al., ; Xu & Tewari, ; Feng et al., ). Increasing the selection from challenging clusters for models to learn can enhance their generalization capability.

Our iterative approach is:

1. Initialization
Given a fixed training budget of $b$, we use $k$MQ as described in the previous section to cluster and sample an initial set of instances of the size $\frac{b}{\mathcal{N}}$.
Next, we finetune the base model for one epoch by going over the sampled data once, using this checkpoint to guide the iterative selection.

2. Estimation of Sample Difficulty
Using the latest checkpoint, we perform one round of inference on the prompts on which the model is trained.
Specifically, given the prompt $x_{i}$ from the initial sampled set, we generate a new completion $y_{i}$ from the original seed data, forming the tuple $(x_{i},y_{\text{gen}},y_{\text{gold}})\in J^{i}$.
We then evaluate the quality difference between $y_{\text{gen}}$ and $y_{\text{gold}}$ using a scorer $\mathcal{S}$. We compute the score for each instance by the following:

$$ $\mathcal{S}(x_{i},y_{\text{gen}},y_{\text{gold}})=\text{score}(x_{i}\oplus y_{ \text{gold}})-\text{score}(x_{i}\oplus y_{\text{gen}})$ (4) $$

where $\oplus$ is the concatenation operator.
We explore the effectiveness of different scoring methods in section 4.2.

3. Resampling
By aggregating and normalizing the scores of samples within each cluster, we modify the sampling weight of each cluster in the next iteration. The goal is to assign a higher weight to the clusters containing higher-quality data while reducing the number of instances selected from lower-quality clusters. We define the score and weight of the $j^{\text{th}}$ cluster as follows:

$$ $s_{j}=\frac{1}{|D_{j}|}\sum_{i=1}^{|\mathcal{D}_{j}|}\mathcal{S}(x_{i},y_{ \text{gen}},y_{\text{gold}})$ (5) $$

$$ $w^{it}_{j}=\frac{s_{j}}{\sum_{c=1}^{k}s_{c}}w^{it-1}_{j}$ (6) $$

where $s_{j}$ is the score of the $j^{\text{th}}$ cluster, $w^{it}_{j}$ is the weight of the $j^{\text{th}}$ cluster at iteration $it$. $it$ is the iteration number, where $it\in\{0,1,\ldots,\mathcal{N}\}$. $\mathcal{N}$ is the maximum number of iterations and $k$ is the total number of clusters.

We adjust the cluster weights and select $\frac{b}{\mathcal{N}}$ new candidates based on these updated weights.
We then train the model and return to step 2.
This process continues until the entire training budget is utilized.
The pseudocode summarizing our iterative data selection method is shown in Algorithm 1.

## 3 Experiments

### 3.1 Training setup

Source Datasets  We focus on two large and widely used instruction datasets that include prompts on a diverse set of topics: Alpaca (Taori et al., ) and WizardLM (Xu et al., ). The Alpaca dataset includes 52K prompts and uses the self-instruct framework to evolve seed human instruction datasets and generate a large collection. WizardLM includes 196K prompts where they used Evol-Instruct to automatically augment instruction tuning datasets (Alpaca, ShareGPT) to make their instructions more complex (in-depth evolution) and more diverse (in-breadth evolution).

**Table 1: Detailed information of our evaluation settings. For each evaluation dataset, we present the number of few-shot examples and metric adopted for evaluation.**
| Evalset | Metric | # shots |
| --- | --- | --- |
| MMLU | acc | 5 |
| GSM8k | acc | 5 |
| HellaSwag | acc-norm | 10 |
| ARC | acc-norm | 25 |
| TruthfulQA | acc | 0 |
| HumanEval | pass@10 | 0 |

Encoding data points  We use Cohere English embedding (embed-english-v3.0) to embed the instruction datasets. Note that we encode both the prompts and completions.
To study the impact of the embedding model, in Section 4.3 we experiment with other models to encode instances in our training pool, namely OpenAI embedding (text-embedding-3-large) and Llama-2-7B model (using the last hidden state of the last token).

Training Recipes  For all experiments, we finetune the llama-2-7B base model (Touvron et al., ).
We train for 3 epochs to achieve convergence and optimal instruction-following performance.
We use an AdamW optimizer (Loshchilov & Hutter, ), with a learning rate of 1e-5 and 1,000 warming-up steps. The maximum token size is 2048, and the effective batch size is 64.
Additionally, in section 4.5 we study the transferability of our findings to other base models and experiment with fine-tuning Mistral (Jiang et al., ) and Llama-3-8B (Dubey et al., ).

### 3.2 Evaluation setup

To present a comprehensive overview of the performance of our method, we conduct a comprehensive
evaluation of our approaches and the established baselines across a range of LLM benchmarks.

Natural Language Reasoning We use HellaSwag (Zellers et al., ), and TruthfulQA (Lin et al., ). HellaSwag is a test of commonsense inference. TruthfulQA measures a model’s propensity to reproduce falsehoods.

World Knowledge We evaluate on MMLU (Hendrycks et al., ) and ARC (Clark et al., ).
MMLU consists of a range of multiple-choice academic questions. ARC is a set of grade-school science questions.

Code Generation We use the extensively utilized HumanEval (Chen et al., ) benchmark consisting
of 164 coding problems to evaluate LLMs’ code-writing capabilities at the function level by reporting the pass@10 metric.

Math Reasoning We use GSM8k (Cobbe et al., ) to evaluate the mathematical abilities of models; GSM8k contains 1319 grade school math test data. We adopt 8-shot testing and report the exact matching.

### 3.3 Baselines

We implement two strong data selection methods, Deita (Liu et al., ) and QDIT (Bukharin & Zhao, ) and compare our methods against them.
Additionally, we explore other clustering and sampling methods: $k$-center clustering ($k$-Center), where $k$ equals the number of data points, $k$-means-closest ($k$M-Closest), which selects samples based on the closest distance, and $k$-means-random ($k$M-Random), which selects randomly from each cluster, both with the same budget as our proposed approach $k$MQ.
We also compare our methods to the random selection of data points.

**Table 2: Data selection performance of $k$MQ and baseline methods. All methods sample 10K (5%) from the full WizardLM (196k) dataset. kMQ-$k$ denotes $k$-means-quality with $k$ clustering centroids. For both $k$M-Closest and $k$M-Random, we show the results of the optimal $k$.**
|  | MMLU | GSM8K | HellaSwag | ARC | TruthfulQA | HumanEval | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Random | 42.4 | 13.3 | 79.9 | 53.6 | 44.8 | 28.5 | 43.8 |
| Deita (Liu et al., ) | 44.1 | 15.6 | 80.1 | 54.3 | 44.9 | 30.8 | 45.0 |
| QDIT (Bukharin & Zhao, ) | 43.3 | 14.5 | 81.1 | 54.4 | 45.2 | 32.7 | 45.2 |
| \hdashline<br>$k$-Center | 41.5 | 11.8 | 79.2 | 51.7 | 43.5 | 28.4 | 42.7 |
| $k$M-Closest | 42.1 | 14.2 | 80.4 | 54.9 | 44.6 | 31.2 | 44.6 |
| $k$M-Random | 43.2 | 15.4 | 81.0 | 55.5 | 44.8 | 31.2 | 45.2 |
| $k$MQ | 45.9 | 16.2 | 81.2 | 55.3 | 45.5 | 33.0 | 46.2 |
| Iterative $k$MQ | 46.1 | 18.4 | 80.1 | 56.0 | 46.3 | 34.3 | 46.9 |

## 4 Results and Discussion

### 4.1 Main Findings

Table 3.3 presents the performance of the proposed methods for instruction data selection compared to several baselines across various tasks.
Our first observation is that by clustering data points using the $k$-means method and randomly sampling instances ($k$M-Random sampling) we already outperform random sampling and achieve comparable results to strong baselines: Deita and QDIT.
This is significant because this sampling method is significantly more efficient than both Deita and QDIT and does not depend on costly LLMs for scoring.
The success of this simple and efficient method highlights the impact of prioritizing diversity in sampling.

Next, we observe that by replacing the random selection step with the quality-based approach ($k$MQ) we can improve model performance on all downstream tasks.
$k$MQ outperforms strong sampling approaches, Deita (Liu et al., ) and QDIT (Bukharin & Zhao, ), on all tasks.
Next, we observe that the iterative sampling approach (Iterative $k$MQ), which leverages early training signals to refine the selected subset, outperforms all previous baselines on most tasks. This suggests that the iterative process of resampling and finetuning based on cluster performance can effectively identify and prioritize high-quality instruction data, leading to better task performance.

Overall, our findings highlight the impact of a diversity-focused sampling approach, which selects a compact yet representative subset of the data through clustering and weighted sampling from the clusters.
We find that it is also crucial to consider a feedback loop from the finetuning model and understand how it perceives and learns from the data.
By incorporating this feedback we ensure that the sampling process aligns with the model’s learning behavior for optimal results.

Figure: Figure 2: Comparison of iterative selection approach using different sample-scoring methods: perplexity, GPT-4, reward model. Note that both random and $k$MQ selection methods use 10% of data and train for three epochs. The iterative feedback runs are performed with the same budget at iteration 3, ensuring a fair comparison. Iterative sampling using a reward model achieves the best performance.
Refer to caption: x2.png

### 4.2 Comparing different scoring methods in iterative feedback

To study the impact of how we score samples during training in our iterative selection approach, we compare three methods: calculating the perplexity score of generations, using GPT-4 to obtain a quality score, and using a reward model’s(^1^11We use FsfairX-LLaMA3-RM-v0.1 (Xiong et al., ).) output.
In Figure 2 we observe that all three variants effectively improve the average performance over random selection.
It is important to note that during the first and second iterations, the iterative methods have been exposed to fewer data points compared to the random and $k$MQ baselines.
It is only at the third iteration that all methods have had the opportunity to process an equal amount of data.
While both perplexity-based and GPT-4-based scoring achieve similar performance to $k$MQ and improve over random sampling, the reward model variant largely outperforms a single-run $k$MQ.
For this experiment, we arbitrarily selected an iteration value of 3, which can be modified in future experiments.

### 4.3 Impact of Number of Clusters

Figure: Figure 3: Average performance on downstream tasks (bar plots) for different number of clusters $k$. There is a correlation between downstream performance and both Silhouette and Elbow scores. The silhouette score is an efficient and effective proxy to estimate the number of clusters eliminating the need to explore the hyperparameter space.
Refer to caption: x3.png

In $k$-means data selection, an important question is how to choose the appropriate value for the parameter $k$ (the number of clusters). Increasing the value of $k$ results in more fine-grained clusters and by ensuring that we sample from each cluster, we can increase the diversity of the selected subset.
However, overly large values of $k$ would also inevitably create outlier clusters that consist entirely of low-quality, noisy data. Since we ensure each cluster is represented in the final selection, this results in noise being included.
There is no one-size-fits-all answer, as the optimal $k$ depends on the characteristics of the pool of data.
Exploring the optimal parameter value is costly, as it must be conducted with each new dataset.
Here we use established heuristics in the clustering literature to guide this decision and study the correlation of these metrics with downstream performance of language models. Namely we investigate two methods:

Elbow method is a popular approach (Ahmed et al., ), where the objective value is plotted against different values of $k$. The goal is to identify the elbow point, where increasing $k$ yields diminishing returns in the performance.

Silhouette Score (Vardakas et al., ) provides another perspective by evaluating how well each data point fits within its assigned cluster (cohesion) compared to other clusters (separation), ranging from -1 (poor fit) to 1 (excellent fit). A high score indicates the object is similar to others in its cluster and dissimilar to those in neighboring clusters.

Although both approaches for identifying the ideal number of clusters are frequently employed, the Silhouette score is generally preferred to the Elbow method in $k$-means clustering due to its clear interpretability, robustness to noise and outliers, and suitability for high-dimensional data.
More importantly, the Elbow method is a post-hoc evaluation metric after the instruction tuning is done and is more expensive; while Silhouette score can be computed prior to any sampling and training and is very cheap.

We study how the choice of $k$ affects performance on downstream tasks and if we can identify an optimal $k$ based on the dataset’s properties.
To investigate this, we first fix our training pool (using WizardLM) and run a series of experiments with different numbers of clusters $k$.
For each value $k$, we cluster the training candidates and sample from the clusters to create subsets of instruction data. We then finetune a model on each of these subsets.
A full evaluation is conducted for every finetuned model (see detailed results in Appendix B).

Figure: Figure 4: The percentage of clusters with an aggregated quality score below the threshold of 0.3.
Refer to caption: x4.png

Figure 3 provides a summary of the results, we reported the average performance over all tasks and observe that the average performance changes dramatically when we change the number of clusters.
This is expected since we rely on the clustering step to group data points that are similar and distinct from other clusters.
Remarkably, we observe a correlation between performance on downstream tasks and the Silhouette score.
We find that the Silhouette score can be used to estimate the number of clusters required before performing the expensive pipeline of clustering, sampling, finetuning, and evaluation.
This estimation step enables us to adapt our approach efficiently to new datasets and collections, ensuring optimal performance and reducing computational costs associated with trial-and-error methods to find the best hyperparameters.

### 4.4 Analyzing Cluster Quality

In our approaches, we rely on $k$-means clustering to ensure high diversity, but there is a risk that some clusters may consist solely of noise.
To understand how this varies with different values of $k$, we use a reward model to evaluate the quality of each cluster with a score between 0 and 1.
Figure 4 shows the number of clusters with average quality scores below a certain threshold (0.3) for different values of $k$.
We observe that by increasing the number of clusters, the percentage of the clusters dominated by low-quality data also increases.
This increases the likelihood of sampling low-quality data when attempting to ensure that every cluster is represented in the final selection.
In our iterative sampling approach, we adjust cluster weights during each training iteration and prevent noisy clusters from being over-represented in the sampled data.

**Table 3: Performance of our best sampling methods on downstream tasks for two base models: Llama-3-8B and Mistral-7B. We sample 10K (5%) from WizardLM (196k). The selection is performed with Llama-2.**
|  | MMLU | GSM8K | HellaSwag | ARC | TruthfulQA | HumanEval | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Mistral-7B |  |  |  |  |  |  |  |
| Random | 58.2 | 26.2 | 82.4 | 60.1 | 60.5 | 26.3 | 52.3 |
| \hdashline<br>$k$MQ | 59.1 | 31.0 | 83.3 | 60.2 | 64.7 | 28.4 | 54.5 |
| Iterative $k$MQ | 59.6 | 32.2 | 83.5 | 60.1 | 66.8 | 29.7 | 55.3 |
| Llama-3-8B |  |  |  |  |  |  |  |
| Random | 65.1 | 38.4 | 83.3 | 60.6 | 55.1 | 56.7 | 59.9 |
| \hdashline<br>$k$MQ | 67.2 | 40.1 | 83.5 | 61.3 | 57.3 | 57.6 | 61.2 |
| Iterative $k$MQ | 66.0 | 36.7 | 83.3 | 61.0 | 56.4 | 54.2 | 59.6 |

### 4.5 Transferability of Results

We conduct experiments with two additional base models, Mistral-7B and Llama-3 8B (Jiang et al., ; Dubey et al., ), to assess whether our findings generalize to other model families and more powerful models.
Our results in Section 4.4 demonstrate that the effectiveness of iterative refinement remains valid for the Mistral-7B model, which exhibits more robust performance.
However, the evaluation results for Llama-3 are mixed across different benchmarks.
We observe improvements on average with $k$MQ sampling and a slight decrease in performance with iterative sampling especially in reasoning tasks.
We hypothesize that Llama-2 differs from Mistral in its training data, model parameters, and training strategies.
Consequently, using Llama-2 as a scorer reveals novel data points from which Mistral can benefit. However, Llama-3, a more advanced model than its predecessors with extended training as one of the primary distinctions, uncovers fewer new, valuable data points for further learning.
This highlights that the quality scorer’s effectiveness can vary, sometimes proving more beneficial and other times less so, depending on the base model for which we are sampling.

While the iterative refinement pipeline can select a dataset restricted to certain models, we do not view this as a limitation. The primary contribution of this work is to propose a function that takes a fixed dataset and model as input and outputs the most valuable subset for learning. This approach aligns with similar works (Ilyas et al., ; Thrush et al., ). Specifically, the task is to extract a subset of data that leverages early reward signals to enhance the targeted model’s post-training performance.

## 5 Related Work

Data selection for LLMs.
Previous works on data selection can be broadly categorized into two key approaches: (1) removing undesirable examples, for instance, low-quality (Raffel et al., ; Marion et al., ), toxic (Raffel et al., ), or duplicated instances (Zhang et al., ; Abbas et al., ).
(2) identifying the most optimal subset of data.
While the definition of an optimal subset varies across different works, the shared goal is to use a small portion of the data while still maintaining strong performance.
This subset selection approach has often been done by aiming for selecting high-quality instances through a proxy: manual curation (Zhou et al., ), selecting instances from human-authored datasets (Wang et al., ), or hand-selecting datasets encouraging complexity and diversity (Ivison et al., ).
More recently, a line of work has used language models to assess the quality of each data point and select the best ones.
Xia et al. () estimates data influence and performs a low-rank gradient similarity search using a gradient datastore.
Liu et al. () scores instances using a combination of complexity and quality scores using an LLM and selects the final subset using diversity-based filtering.
While individual sample quality is a crucial factor, prioritizing this local criterion can limit the diversity of the final selection. However, diversity in instances and tasks is essential for training high-performant models (Wei et al., ; Gudibande et al., ).
Our work differs from these studies by examining what constitutes an optimal subset from a global perspective and by prioritizing representativeness.
Closest to our work, Bukharin & Zhao () emphasized quality by encoding all data points in the selection pool using an embedding model and selecting the final subset based on pairwise cosine similarity and a quality score from an LLM. In contrast, our approach offers a significantly more efficient method for subset selection, while also achieving improved performance.
Our experiment covers multiple dimensions, including various base models, different encoding and scoring methods, and extensive ablation studies with recommendations for efficient parameter selection.

Active learning and language models.
Active learning is based on the fundamental premise that “not all data is equal”. This approach aims to identify the most informative data for pretraining or adapting language models for specific tasks or capabilities, as well as pinpointing the most valuable data for learning. Margatina et al. () explored active learning for selecting in-context examples in few-shot learning, demonstrating that similar examples outperform uncertain or diverse in-context examples. Muldrew et al. () proposed active preference learning, combining iterative data acquisition with a DPO (Direct Preference Optimization) loop to reduce the frequency of querying human annotators (Oracle). Their acquisition method relies on the model’s entropy during generation. Our approach generalizes active instruction tuning (Kung et al., ) to instance-level data selection, allowing for the co-evolution of the LLMs and instruction data using an external reward signal.

## 6 Conclusion

In this paper, we present a novel approach to selecting a subset of data and optimizing the fine-tuning of language models.
Our method involved a scalable sampling technique that maximizes diversity and efficiency in subset selection.
Through our proposed $k$-means-quality ($k$MQ) algorithm and iterative selection process, we demonstrated significant performance improvements over strong baselines while maintaining a limited training budget.
Our contributions include an efficient instruction selection algorithm, the release of our encoded instruction dataset, and a systematic analysis of our method’s effectiveness across a range of tasks. Our method outperforms existing baselines, achieving up to 7% improvement in a wide range of evaluation tasks.

By addressing the challenge of optimal instruct data selection, our work paves the way for more efficient and effective finetuning of language models, making them more accessible and affordable for deployment, especially in resource-constrained settings.
We believe that our findings will contribute significantly to the ongoing research in language model optimization and their real-world applications.

## 7 Limitations and Future Work

While our proposed method has shown promising results, there are a few limitations to consider. Our evaluation focused on a specific set of tasks, and future work can aim to validate our method’s effectiveness across a broader range of language models and tasks, including data selection in the pre-training stage and alignment (Yu et al., ; Muldrew et al., ). Furthermore, our iterative selection process relies on early training signals, and we only presented this as a pilot study to encourage further research. Future work could explore alternative model feedback mechanisms to refine the selected instruction data subsets, especially in mitigating the potential for reward hacking in the iterative refinement process (Pan et al., ).

Finally, while we considered diversity and difficulty crucial factors, other characteristics of instruction data could be explored to enhance the finetuning process further.
Addressing these limitations and extending this research will contribute to more robust and adaptable language models, capable of excelling in a wide range of real-world applications.

Broader Impact  If the data selection process fails to capture important aspects of the full dataset, it could lead to biased or inconsistent outputs from the finetuned models. There are also broader societal risks around the misuse of large language models for generating misinformation, perpetuating biases, or enabling privacy violations that could be exacerbated by making these models more accessible through efficient finetuning techniques.

## Acknowledgements

We would like to thank the Cohere For AI team for their valuable feedback and for providing generous computing resources for conducting and analyzing our experiments.
In our figures, we use “Robot” icon by Andi Nur Abdillah, BEDESCHI LEONARDO, “iterative process” by cARTo, “growth” and “decrease” by archer7 from thenounproject.com CC BY 3.0.

## Appendix

## Appendix A Training Details

### A.1 Hyperparameters

For supervised fine-tuning, our training hyperparameters are presented in table 4.

**Table 4: Our training hyperparameters.**
| Parameter | Value |
| --- | --- |
| Precision | BFloat16 |
| Epochs | 3 |
| Selected Portion | 10% |
| Gradient Accumulation Step | 8 |
| Batch Size | 64 |
| Max Seq. Length | 4096 |
| K-means Random Seed | 42 |

### A.2 Computational Cost

We also utilised Deepspeed-Zero3 (Rasley et al., ) for better efficiency training. Models are finetuned with combination of TPU and GPU. For TPU, we used the code provided by young-geng/EasyLM(^2^22[young-geng/EasyLM](https://github.com/young-geng/EasyLM/tree/main)) and done with TPUv3-32 nodes. For GPU, 2 A100-80GB are used across the fine-tuning.

## Appendix B Impact of Number of Clusters

**Table 5: Performance of models trained on different number of data clusters $k$. We sample 10K (5%) for each experiment. Silhouette score correlates with downstream tasks and is an efficient proxy for estimating the number of clusters before sampling.**
| Method | MMLU | GSM8K | HellaSwag | ARC | TruthfulQA | HumanEval | Avg. | Silhouette Score |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| $k$MQ-64 | 43.1 | 13.9 | 80.2 | 54.3 | 44.8 | 29.5 | 44.3 | 17.4 |
| $k$MQ-128 | 43.4 | 12.8 | 79.9 | 54.1 | 45.0 | 28.4 | 43.9 | 15.6 |
| $k$MQ-256 | 42.3 | 13.1 | 80.0 | 53.2 | 44.3 | 28.1 | 43.5 | 14.1 |
| $k$MQ-512 | 46.4 | 17.0 | 81.2 | 55.3 | 45.5 | 33.0 | 46.4 | 16.8 |
| $k$MQ-1024 | 45.6 | 17.8 | 81.6 | 55.5 | 44.9 | 34.1 | 46.6 | 18.2 |
| $k$MQ-2048 | 46.0 | 20.2 | 82.1 | 55.5 | 45.0 | 37.2 | 47.7 | 20.3 |
| $k$MQ-4096 | 44.2 | 15.2 | 79.1 | 54.3 | 42.0 | 33.2 | 44.7 | 20.0 |
| $k$MQ-8192 | 44.1 | 13.6 | 78.9 | 54.2 | 41.6 | 31.8 | 43.0 | 18.7 |

**Table 6: Additional experiments on Alpaca dataset (52k). We sample 5K (10%) for each experiment. kMQ-$k$ denotes $k$-means-quality with $k$ clustering centroids. For both $k$M-Closest and $k$M-Random, we show the results of the optimal $k$ among all choices of $k$.**
|  | MMLU | GSM8K | HellaSwag | ARC | TruthfulQA | HumanEval | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Random | 38.2 | 9.1 | 79.1 | 51.3 | 41.1 | 20.5 | 39.9 |
| Deita | 39.4 | 10.7 | 79.4 | 51.2 | 41.7 | 22.9 | 40.9 |
| QDIT | 38.7 | 11.3 | 79.8 | 51.6 | 42.6 | 25.6 | 41.6 |
| \hdashline<br>$k$-Center | 37.3 | 8.1 | 79.0 | 50.7 | 41.0 | 12.8 | 38.2 |
| $k$M-Closest | 40.1 | 10.3 | 79.3 | 51.2 | 42.5 | 24.3 | 41.3 |
| $k$M-Random | 39.6 | 11.4 | 79.1 | 51.2 | 42.8 | 25.1 | 41.5 |
| $k$MQ-64 | 41.3 | 12.6 | 79.7 | 51.1 | 43.4 | 25.3 | 42.2 |
| $k$MQ-256 | 39.5 | 12.3 | 79.1 | 51.0 | 42.7 | 26.0 | 41.8 |
| $k$MQ-1024 | 37.3 | 11.2 | 78.6 | 51.2 | 41.5 | 22.3 | 40.4 |

## Appendix C Scorer Details

For perplexity, we pass the $x_{i}\oplus y_{gen}$ and $x_{i}\oplus y_{gold}$ to the model to compute the perplexity scores. The scorer with Perplexity is as follows:

$$ $\mathcal{S}(x_{i},y_{\text{gen}},y_{\text{gold}})=-\log(\frac{PPL(x_{i}\oplus y _{\text{gen}})}{PPL(x_{i}\oplus y_{\text{gold}})})$ (7) $$

For GPT-4 direct scoring, we give the two completions to GPT-4 and ask it to give a rating between 1 and 5.
We use the template as shown in Figure 6 to prompt GPT-4 for being the LLM-as-a-judge and by replacing the reward scoring ($R$) by the GPT score in Equation 4. The template is inspired by Zheng et al. (). For the reward model, we use an off-the-shelf model based on Llama-3(^3^33[FsfairX-LLaMA3-RM-v0.1](https://huggingface.co/sfairXC/FsfairX-LLaMA3-RM-v0.1)).

Figure: Figure 6: Prompt template for requesting a response evaluation from GPT-4-turbo, where variables ${instruction} and ${response} are replaced with examples in our dataset.

## Peer Reviews

confidence: 4
contribution: 1
initial_rating: 1
presentation: 2
questions: See weaknesses
soundness: 1
strengths: The paper is easy to understand and well-written, and the problem of selecting diverse and high quality samples for instruction tuning is important. However, I have several concerns regarding the design choices of the methodology and the experiments.
summary: The authors propose a data subset selection strategy that considers both diversity and quality to select samples for finetuning LLM's. Specifically, the strategy clusters the points and selects samples with high quality/hardness scores from each cluster.
weaknesses: ## Weaknesses
- **[Major]** The key sampling strategy used in this work is a reincarnation of strategies that have been used for a while in data pruning and active learning. For example, [1] (work not cited in this paper) first clusters the unlabeled set and chooses the samples with the highest uncertainty from each cluster in a round-robin fashion though it is very straightforward to substitute uncertainty with quality. 
- **[Major]** There are many missing baselines in the experiments. Many active learning and subset selection strategies can be used for this setting from the core ML community [2,3,4] many of which have been applied to the problem of supervised finetuning [5]. There are also subset selection works specifically for supervised finetuning that are missing, for example [6]. At least a few of these should be included as baselines in this work, and discussion of all of these works should be included in the related works section. 
- **[Major]** Additional ablations are necessary. What happens if we only use quality without any diversity prefiltering? 
- **[Major]** In most of the experiments, evaluation is only done on a single subset size. Evaluation on multiple subset sizes is fairly standard in data selection literature even at large scales [5,6], and should be included here as well.
- **[Major]** The strategy first clusters the points and then samples based on quality. However, this does not guarantee diversity in the final subset. Using quality based sampling in the second stage can still induce redundancies **within a cluster**. There is nothing to prevent two identical high quality samples that belong to the same cluster from being selected. Filtering based on some score before applying diversity sampling prevents this (this has been done in [7]). 
- **[Minor]** Why is the sampling strategy randomized? From my understanding, sampling is performed only once for each iteration and scores must be computed for each sample to compute its probability so there doesn't seem to be any benefit for this in terms of computational cost. 
- **[Minor]** In Algorithm 1, why is $\mathcal{D}'$ grown incrementally? In other words, a new set of points is sampled from $\mathcal{D}/\mathcal{D'}$ and added to $\mathcal{D}'$ at each round, but shouldn't all of $\mathcal{D'}$ be resampled? It seems as though points that the LLM has already been trained on will be easy and assigned a low score based on Eq. 4.  
- **[Minor]** This claim "The success of this simple and efficient method highlights the impact of prioritizing diversity in sampling." does not make sense given that the baseline method QDIT also is a diversity-based selection approach. 

I know I have listed several points and I understand that it may not be possible to address all of them in the short rebuttal period, so I would encourage the authors to prioritize addressing the **[Major]** points.

## References
- [1] Batch Active Learning at Scale (https://arxiv.org/pdf/2107.14263)
- [2] Deep Batch Active Learning by Diverse, Uncertain Gradient Lower Bounds (https://arxiv.org/abs/1906.03671)
- [3] D2 Pruning: Message Passing for Balancing Diversity and Difficulty in Data Pruning (https://arxiv.org/abs/2310.07931)
- [4] Active Learning for Convolutional Neural Networks: A Core-Set Approach (https://arxiv.org/abs/1708.00489)
- [5] An Experimental Design Framework for Label-Efficient Supervised Finetuning of Large Language Models (https://arxiv.org/abs/2401.06692)
- [6] Diversity Measurement and Subset Selection for Instruction Tuning Datasets (https://arxiv.org/pdf/2402.02318)
- [7] Submodularity in Data Subset Selection and Active Learning (https://proceedings.mlr.press/v37/wei15.html)

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. Please refer to Weaknesses.
soundness: 3
strengths: 1. The proposed data selection method achives a better downstream performance than evaluated baselines, which showcases the importance of diversity in data selection.

2. The downstream analysis datasets are extensive.
summary: This study addresses the challenge of selecting optimal data subsets for fine-tuning LLMs on instruction datasets by focusing on data diversity rather than solely on local quality criteria. Utilizing a k-means clustering method, the approach aims to represent the entire dataset effectively, with an iterative refinement inspired by active learning to continuously adjust sampling weights and filter out low-quality data clusters. The experiemnts illustrate that this diversity-first approach leads to an improvement over random sampling showcasing the value of diverse data selection for enhancing LLM performance across various tasks.
weaknesses: 1. I think the baseline comparison is not complete. For example, there is a paper "One-Shot Learning as Instruction Data Prospector for Large Language Models" published on ACL 2024 which is higly related to this topic but not been included as baseline in this paper.

2. The experiment still cannot persuade me that the diversity is the most imporant factor in selecting instruciton data. In Table 3, the proposed sampling method seems marginally better than the Random selection. In some datasets, it is even worse than Random selection, such as HumanEval on Llama-3.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: - For k-means, why is k-means++ initialization not used, instead of using 42 as random seed?
- I am interested in seeing some run-time analysis, as usually the size of the instruction tuning dataset is large. 
- Can the authors try to run experiments using facility location alone, and then using that instead of clustering? It may be important to note that different choice of kernel in facility location can result in different performance (Cosine v/s RBF as done in [3]). Note that one has to do tuning, as they've also suggested a way to tune the resulting submodular function. 
- Can authors also try to compare against DPP as done in [4] ?
soundness: 3
strengths: Paper proposes on how to do unsupervised tuning of the cluster center hyperparameter. I also appreciate the use of reward model as a score function, which goes beyond the classic loss based score such as perplexity. Experiments and ablations are conducted over several datasets. As far as this study is concerned for clustering based solutions, I like this work.
summary: This paper introduces use of clustering for data subset selection using clustering followed by quality score sampling for efficient instruction tuning. Paper proposes multiple variants of the clustering based method, and how to do unsupervised tuning of the cluster center hyperparameter. I also appreciate the use of reward model as a score function, which goes beyond the classic loss based score such as perplexity. Experiments and ablations are conducted over several datasets, including transferability, showing robustness of the proposed approach.
weaknesses: A major weakness of this work is lack of literature review. For example, use case of diversity and some sort of quality score has been used in Machine Learning for a long time, from active learning [1], to curriculum learning [2, 6] (also see common followups to [6] which are on the similar lines, such as [7]), to more recent LLM instruction tuning based works [3, 4, 5]. However, paper fails to acknowledge these works except [3]. Submodular functions have been used for the purposes of inducing diversity for a very long time, which also includes determinantal point process. However, no discussion on submodularity has been provided in the related works. 

I believe [5] has done somewhat detailed study on submodular functions, its usecase with quality scores (uncertainty, confidence) and I think the perplexity based score function may have some correlation between confidence based score function.

 To summarize, I think this is a good paper when it comes to the clustering based methods, but it needs to benchmark against the submodular methods as well.  

References - 

[1] Submodularity in Data Subset Selection and Active Learning
[2] Curriculum Learning by Dynamic Instance Hardness 
[3] An Experimental Design Framework for Label-Efficient Supervised Finetuning of Large Language Models
[4] Diversity Measurement and Subset Selection for Instruction Tuning Datasets
[5] SmallToLarge (S2L): Scalable Data Selection for Fine-tuning Large Language Models by Summarizing Training Trajectories of Small Models
[6] Coresets for Data-efficient Training of Machine Learning Models
[7] Towards Sustainable Learning: Coresets for Data-efficient Deep Learning

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: None.
soundness: 3
strengths: 1. Even though the claim that data diversity is critical is well-known [1,2], the motivation of this paper is clear.
2. The proposed method is sound. Using active learning is a straightforward idea to select the optimum subset of data.
3. This paper is well-written and easy to follow. The released code and the data may facilitate reproducibility and future research.

Overall, this paper attempts to address an important question and proposes an effective method that achieves better performance than the compared methods.

[1] LIMA: Less Is More for Alignment

[2] Data Diversity Matters for Robust Instruction Tuning
summary: Finetuning LLMs on instruction data is an important step to improve instruction following capabilities. This paper attempts to address how to determine the optimal subset of data for effective training. Unlike previous works that focused on local criteria (e.g., data quality), this paper claims that data diversity is more critical for effective training. This paper utilizes k-means clustering and proposes an iterative refinement method inspired by active learning techniques to resample instances from clusters. Experiments show that the proposed method achieves a 7% increase over the random selection and a 3.8% improvement over state-of-the-art sampling methods.
weaknesses: 1. Baselines are weak: Even though the author attempts to demonstrate that data diversity is more important than data quality, the quality-based data selection method chosen by the authors is insufficient to support their claims, as they have selected only one method based on quality. (i.e., Deita,). There are still lots of published papers focusing on data quality [1,2,3]. Meanwhile, the improved performance compared with Random Selection is also not convincing.
2. Evaluation is weak: Even though the paper claims that finetuning LLMs on instruction data is an important step to improve instruction following capabilities, there is no evaluation of the quality of the generated response (e.g., MT-Bench [4]). The evaluation in the paper mainly focuses on understanding tasks instead of instruction-following tasks.

[1] From Quantity to Quality: Boosting {LLM} Performance with Self-Guided Data Selection for Instruction Tuning

[2] Clustering and Ranking: Diversity-preserved Instruction Selection through Expert-aligned Quality Estimation

[3] LESS: Selecting Influential Data for Targeted Instruction Tuning

[4] Judging LLM-as-a-judge with MT-bench and chatbot arena

Meta-review: This paper proposes considering both quality and diversity in the data selection process, to select instructional data better for follow-up tasks. The motivation is overall clear. Besides, the technical details of the proposed method are easy to follow. 

This paper has some flaws that cannot be addressed by slight modifications. First, it lacks sufficient comparison with previous work in terms of content and ideas, which makes the novelty and its technical contribution unclear. Second, the comparison method is not comprehensive enough. The experimental results are not convincing enough. Third, although the technical details of the paper are easy to follow, it does not provide an intuitive and convincing explanation of why the proposed method actually works.

AC carefully reviewed the paper, read the reviewers’ comments, the authors’ responses, and organized the discussion. The current form needs further polishes to reach acceptance standards. Therefore, the final recommendation is "reject".

Decision: Reject