Diversify and Conquer: Diversity-Centric Data Selection with Iterative Refinement (2409.11378v1)

Abstract: Finetuning LLMs on instruction data is crucial for enhancing pre-trained knowledge and improving instruction-following capabilities. As instruction datasets proliferate, selecting optimal data for effective training becomes increasingly important. This work addresses the question: How can we determine the optimal subset of data for effective training? While existing research often emphasizes local criteria like instance quality for subset selection, we argue that a global approach focused on data diversity is more critical. Our method employs k-means clustering to ensure the selected subset effectively represents the full dataset. We propose an iterative refinement method inspired by active learning techniques to resample instances from clusters, reassessing each cluster's importance and sampling weight in every training iteration. This approach reduces the effect of outliers and automatically filters 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 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.

Diversity-Centric Data Selection for LLMs

The paper "Diversify and Conquer: Diversity-Centric Data Selection with Iterative Refinement" presents a novel approach for selecting optimal subsets of data for fine-tuning LLMs. The authors argue that instead of focusing solely on local criteria such as instance quality, a global approach that emphasizes data diversity can yield more significant benefits. Through an innovative method combining $k$-means clustering and iterative refinement, the authors demonstrate consistent improvements across various tasks, showcasing the efficacy of their approach.

Methodology

The authors propose a twofold methodology: a static data selection approach using $k$-means clustering and an iterative refinement process inspired by active learning principles.

Static Data Selection

The first step involves selecting a subset $\mathcal{D}'$ from a large dataset $\mathcal{D}$ by employing $k$-means clustering. The goal is to ensure that the selected subset is representative of the entire dataset. The $k$-means algorithm groups similar data points into clusters, from which samples are then selected either randomly or based on quality scores. The quality scores are derived using an approach inspired by previous works, such as Deita and QDIT, which employ LLMs to assess the quality of data points. The static data selection with $k$-means-quality (kMQ) has shown to outperform previous state-of-the-art methods, highlighting the importance of diversity and representativeness in data sampling.

Iterative Refinement

Building upon the static selection, the authors introduce an iterative refinement process wherein early training signals from the fine-tuning model are used to resample instances. This iterative method allows for continuous reassessment of cluster importance and sampling weights across multiple training iterations. By leveraging model feedback, the approach automatically filters out low-quality clusters and fine-tunes the sampling process to improve overall model performance progressively. The iterative refinement significantly enhances the fine-tuning process, yielding notable improvements over both random sampling and fixed sampling methods.

Experimental Results

The authors rigorously evaluate their method across several benchmarks, including natural language reasoning (HellaSwag, TruthfulQA), world knowledge (MMLU, ARC), code generation (HumanEval), and math reasoning (GSM8K). By fine-tuning models from various families, including Llama-2-7B, Mistral-7B, and Llama-3-8B, the authors consistently observe substantial performance gains. For instance, their iterative kMQ approach achieves a 7% increase over random selection and a 3.8% improvement over state-of-the-art sampling methods.

Implications and Future Directions

The implications of this research are multifaceted. Practically, the proposed method provides a scalable and efficient solution for fine-tuning LLMs by focusing on data diversity and leveraging iterative refinement. Theoretically, the work underscores the significance of global properties, such as diversity, in data selection processes, challenging the conventional emphasis on local criteria like instance quality.

Future developments may explore extending the iterative sampling approach to other phases of model development, such as pre-training and alignment. Additionally, investigating alternative feedback mechanisms and leveraging more advanced quality scorer models can further refine the selection process. The potential integration of curriculum learning principles and the exploration of other data characteristics could also enhance the efficacy of this method.

Conclusion

The "Diversify and Conquer" approach offers a robust framework for data selection in the context of fine-tuning LLMs. By prioritizing data diversity and implementing an iterative refinement process, the authors provide a compelling alternative to traditional data selection methods. Their extensive empirical evaluations demonstrate the method's superior performance, marking a significant advancement in the field of machine learning and AI.

Overall, this paper presents a significant contribution to the ongoing research on optimizing the fine-tuning of LLMs, emphasizing the critical role of diversity in achieving high-performance outcomes. The insights and methodologies introduced here are likely to inform and inspire future research and applications in AI and machine learning.