Balanced Data Sampling for Language Model Training with Clustering (2024)

Abstract: Data plays a fundamental role in the training of Large Language Models (LLMs). While attention has been paid to the collection and composition of datasets, determining the data sampling strategy in training remains an open question. Most LLMs are trained with a simple strategy, random sampling. However, this sampling strategy ignores the unbalanced nature of training data distribution, which can be sub-optimal. In this paper, we propose ClusterClip Sampling to balance the text distribution of training data for better model training. Specifically, ClusterClip Sampling utilizes data clustering to reflect the data distribution of the training set and balances the common samples and rare samples during training based on the cluster results. A repetition clip operation is introduced to mitigate the overfitting issue led by samples from certain clusters. Extensive experiments validate the effectiveness of ClusterClip Sampling, which outperforms random sampling and other cluster-based sampling variants under various training datasets and large language models.

## introduction large language models (llms) have opened new frontiers in understanding and generating human languages @xcite @xcite . a critical aspect of training these models lies in the acquisition and organization of training data @xcite . some works focus on selecting high-quality data. these works usually perform data filtering or data cleaning from a large corpus, based on either rulebased @xcite or modelbased algorithms @xcite @xcite . on the other hand, several approaches concentrate on optimizing the composition weights of the collected data. these methods typically focus on optimizing the domain weights either using heuristics @xcite or model statistics @xcite . while much attention has been given to the collection and composition of diverse datasets for training llms @xcite @xcite @xcite , it is still unclear how the training data sampling affects the optimization of language models. the sampling methods of existing works are coarse-grained, which determines the sampling weights or sampling order of each domain. the domain is usually defined based on the data source or other metadata during the dataset collection, which is coarse-grained and inaccurate. for instance, the llama models @xcite assign domain mixture weights heuristically, including 67% commoncrawl web data, 4.5% github code, 4.5% wikipedia documents, etc. and @xcite ; @xcite improve certain abilities of llms by tuning the model on specific domains, like code or mathematical texts. however, the texts are sampled randomly in each domain, which ignores the unbalanced distribution of the expressed meanings and topics. due to the nature of the data corpus, texts with similar meanings have a long tail distribution in the training set @xcite @xcite . when using random sampling, llms can underfit rare documents and overfit common samples. it is straightforward to utilize a uniform sampling strategy that samples texts with different meanings evenly, which up-samples rare documents and down-samples common texts. however, uniform sampling will up-weight rare documents and repeat them so many times in the training, resulting in severe overfitting of trained llms on these training samples. to address the above challenges, we propose clusterclip, a cluster-based sampling strategy with clip operation to mitigate overfitting. this sampling strategy has two steps. firstly, we leverage data clustering to reflect the data distribution. based on off-the-shelf nlp tools, the semanticrelated texts could be grouped into the same cluster. by calculating the size of each semantic cluster, we can evaluate the data rarity. secondly, we perform the data sampling using the cluster information. the documents from different clusters are sampled evenly at the beginning of the training, thus encouraging the model to learn from rare documents instead of wasting computation on common texts. as the training progresses, a clip operation is applied. if certain documents are sampled too many times, these documents are clipped and no longer sampled from the dataset. thus, clusterclip rebalances the data distribution, which facilitates the model to learn rare documents but avoids severe overfitting of repeated texts by clipping. our approach is distinguished by its ability to improve model learning efficiency and generalization without relying on dataset-specific metadata or complicated optimizations. extensive experiments demonstrate the versatility and effectiveness of our proposed clusterclip sampling. we show that it consistently enhances the performance of representative llms, llama2-7b @xcite and mistral-7b @xcite , in both pre-training and supervised fine-tuning scenarios, indicating its broad applicability. moreover, we further evaluate several variants of cluster-based sampling methods and demonstrate that: (1) all the cluster-based sampling method variants outperform random sampling, indicating the effectiveness of sampling based on semantic distribution; (2) the clip operation effectively mitigates the overfitting of repeated documents, leading to large improvements in diverse tasks. (3) specializing data training order (e.g. general-to-specific or specific-to-general) can affect the performance of llms on various downstream tasks, showing the promise of progressive learning on llms; to sum up, our contributions are as follows: ## methodology we propose clusterclip sampling, a cluster-based sampling strategy to rebalance the data distribution of the training corpus to facilitate model learning and mitigate overfitting. in this section, we delve into the detailed process of data clustering and the sampling strategy using cluster information to balance the sampling probabilities of common and rare documents. ## experimental setups to validate the proposed clusterclip sampling, we conduct extensive experiments on multiple datasets and llms. in this section, we introduce the experimental setups of our experiments, including the baselines, training datasets, training hyperparameters, and evaluation setups. ## experimental results in this section, we present the experimental results, including the comparison of the proposed sampling method with baselines (sec 5.1), a detailed analysis of multiple variants of cluster-based sampling methods (sec 5.2), and the ablation study of different parts of clusterclip (sec 5.3). ## conclusion in this work, we propose clusterclip sampling based on data clustering to balance the long-tail distribution of the training set of large language models. we compare the proposed method with random sampling and other cluster-based sampling method variants in both supervised fine-tuning and pre-training. extensive experimental results across 7 datasets across diverse tasks and domains demonstrate the effectiveness of clusterclip sampling, which outperforms baselines under different training sets and models. we hope this work can instigate more research on data sampling approaches for improving language model training.