Harnessing Diversity for Important Data Selection in Pretraining Large Language Models

$**Overall response:**$

Thanks for your valuable comments! Based on your comments, we have added a 7B model, 2 reference datasets, 7 downstream tasks and in total 12 additional experiments to fully utilize the 10 pages and put some experiments in the Appendix, so as to demonstrate the generalizability and robustness of $`Quad`$. We also clarify your concerns about the sensitivity of the sampling threshold of influence in our paper. We promise to include the above discussions into the final version if the paper is accepted.

$**Weakness 1:**$ The experiments are not thorough. The experiment and analysis only take up 2 pages, which fails to thoroughly discuss many important components and designs of the proposed method. For example, whether the CS function could be designed in another way? Also, while the evaluation is conducted in various benchmarks, the authors only adopt experiments in a single LLM. This also makes people doubt the generalization of their method. I suggest the author to fully utilize the 10-page space and conduct further experiments.

$**Response for Weakness 1:**$ Thanks for the valuable suggestion to improve our paper! As the reviewer suggested, to demonstrate the robustness and effectiveness of $`Quad`$, we have run new experiments to evaluate different model sizes, different downstream tasks, different candidate datasets and run various ablation studies (including the clustering algorithms, clustering numbers, sampling influence threshold, etc.). We have included the above discussion into Appendix G and H of the revised paper.

Specifically, we train a 7B model from scratch and evaluate all methods on the SlimPajama dataset, using LAMBADA as the reference dataset. In this experiment, 100B tokens out of 627B are selected for training using $`Quad`$. As shown in the table below, $`Quad`$ still demonstrates superior model performance on 7B model compared to all baselines, such as the DSIR, Semdedup, PPL, Qurating and MATES. Also, $`Quad`$ achieves similar performance with the random selection which trains the model with 400B tokens, consuming 307.6 FLOPs instead of 1063.7 FLOPs (1e19) for 400B tokens. However, $`Quad`$ only uses 100B tokens, thus greatly improving the training efficiency.

Note that our CS function is to balance the exploration (i.e., diversity) and exploitation (i.e., quality) inspired by the upper confidence bound score[1], which is a typical solution of MAB.
However, certainly our approach can use other CS functions. Here we discuss two intuitive alternatives, i.e., the exploitation-only and exploration-only methods to demonstrate the superiority of us. The former one samples some data instances from each cluster, computes their influence scores and selects instances from clusters with high influence scores. The latter one just randomly samples from different clusters without considering the data quality. As shown in the table below, we can observe that for both 1B and 7B models, our method outperforms the two baselines because we consider both the data quality and diversity to select data instances.

We have included the above discussion into Appendix H of the revised paper.

[1] Auer, P. "Finite-time Analysis of the Multiarmed Bandit Problem.", 2002.

$**Limitation 2:**$ The second limitation of this study is that the authors do not fully motivate the problem setting. While cleaning noisy data and removing near-duplicate examples are useful for LLM pretraining, I am not fully convinced why you have to select a subset of ``quality data'' from the full dataset. For example, if we do a coarse-grained filtering of the original 627B to leave around 400-500B data and use all of them to do LLM pretraining, will the model performance better?

$**Question 3:**$ Why only select ``30B tokens out of the 627B'' (line 360), what if we just do a coarse-grained filtering of the original 627B to leave around 400-500B data (e.g., via near-duplicate removal) and use all of them to do LLM pretraining, will the model performance better?}

$**Response for the second Limitation and Question 3:**$ In this response, we first discuss the advantage of using $`Quad`$ to select data instances. We then run experiments to validate the argument.

Overall, we agree with the reviewer that data cleaning is likely to enhance the model performance. However, the obtained data after cleaning is still large-scale, which is prohibitively expensive to train, but not always necessary. This is because a subset of judiciously selected data can achieve competitive or even superior model performance compared with training with all clean data. The reasons are two-fold. First, $`Quad`$ is a model-aware data selection method, which uses the data influence to measure the impact of each data instance on current pretraining progress, but coarse-grained data cleaning methods fail to capture the model's current state. Second, $`Quad`$ also considers the diversity for data selection, and thus the selected small part of data instances can well cover different characteristics of downstream tasks.

Overall, model-aware data selection is of great significance because (1) it is customized to the current pretraining progress, which can greatly improve the model performance with a small number of data instances; and (2) the model can be trained over these selected small amounts of data, thus greatly improving the training efficiency.

Next, we run two new experiments to validate this.

First, we would like to clarify that the SlimPajama dataset used in our experiment has already undergone multiple rounds of coarse-grained data cleaning processes [1], including short documents filtering, deduplication, minhash generation, etc. Therefore, to demonstrate the significance of the data selection method, we compare the performance of $`Quad`$ and directly training with clean data with different data sizes with 1.3B and 7B models, evaluating on the long-text generation tasks mentioned in the last response. As shown in the table below, we can observe that Quad-1.3B-100B outperforms Clean-1.3B-100B because Quad judiciously selects data that benefits the current pretraining progress while the latter one just selects 100B data randomly from the clean SlimPajama dataset without considering the current pretraining progress. Because of the similar reason, Quad-1.3B-100B achieves comparable model performance with Clean-1.3B-400B, which saves 75% of computation costs. We can also observe that training on 7B model has a similar result.

Moreover, we add another dataset Openwebmath [2] as our new reference dataset. It is based on the Slimpajama dataset, which aims at selecting data to improve the mathematical skill of the model, i.e., a specific domain. We use GSM8K, MATH, OCW, SAT, MMLU STEM as the downstream tasks and the final accuracy is computed the average score of these downstream tasks. As shown in the table below, we can observe that random selection of 400B data instances does not perform better than selecting 100B by $`Quad`$ on downstream mathematical evaluation datasets. This is because the 400B data instances contain a significant amount of information irrelevant to mathematics, which degrades the model performance, while $`Quad`$ is a model-aware method that accurately identifies the data that benefits the current pretraining progress.

We have included the above discussion into Appendix G of the revised paper.

[1] https://cerebras.ai/blog/slimpajama-a-627b-token-cleaned-and-deduplicated-version-of-redpajama

[2] Paster, Keiran, et al. "Openwebmath: An open dataset of high-quality mathematical web text." arXiv preprint arXiv:2310.06786 (2023).

$**Overall Response:**$

Thanks for your insightful comments! Based on your comments, we have added a 7B model, a reference dataset, 7 downstream tasks and 5 additional experiments to demonstrate the scalability and robustness of $`Quad`$. We additional clarify your concerns about the motivation of our problem setting and some presentation issues in our paper. We promise to include the above discussions into the final version if the paper is accepted.

$**Limitation 1:**$ The main limitation of this study is its experiment scale is too small (in terms of claiming the method for LLM pretraining). Only 30B out of 627B data are selected for training a small 1.3B decoder-only LLM. In practice, most LLMs will start with at least 7B so how the technique can generalize to reasonably sized LLMs is questionable. Also related to this limitation is the diversity of testing datasets, little (long-text) generation tasks are evaluated.

$**Response for the first limitation:**$ We thank your comments for improving our paper. Based on your valuable suggestions, we have added experiments by enlarging the model size and the diversity of downstream tasks, so as to further validate the effectiveness of $`Quad`$.

To be specific, for the model size, we enlarge the model size from 1.3B to 7B and select 100B tokens from Slimpajama dataset. As shown in the table above, we can observe that $`Quad`$ still performs better than other baselines on accuracy because we select high quality data considering the diversity. What's more, $`Quad`$ also has a good scalability compared with state-of-the-art baselines (Qurating and MATES) because we use the K-FAC technique to accelerate the influence computation and MAB framework to quickly identify the data instances that are beneficial to the downstream tasks.

For testing datasets, we add two downstream tasks (i.e., WikiText [1] and HelloBench [2]) about long-text generation tasks. As shown in the table below, on the candidate dataset Slimpajama, for the 7B model, we can observe that $`Quad`$ still outperforms other baselines because we simultaneously consider the quality and diversity in data selection. We also add five downstream tasks above the math in the next response to further demonstrate that we can well generalize to various tasks.

Moreover, we also add the Openwebmath dataset as another new reference dataset to train a 7B model on the Slimpajama dataset, which aims at selecting data to improve the mathematical skill of the model. Correspondingly, we also include GSM8K[3], MATH[4], OCW[5], SAT[6], and MMLU STEM[7] as new downstream tasks to evaluate the model's mathematical capabilities, with the final accuracy calculated as the average score across these tasks. As shown in the table below, $`Quad`$ still demonstrates superior model performance compared to all baselines.

We have included the above discussion into Appendix G and I of the revised paper.

[1] Merity, Stephen, et al. "Pointer sentinel mixture models." arXiv preprint arXiv:1609.07843 (2016).

[2] Que, Haoran, et al. "Hellobench: Evaluating long text generation capabilities of large language models." arXiv preprint arXiv:2409.16191 (2024).

[3] Cobbe, Karl, et al. "Training verifiers to solve math word problems." arXiv preprint arXiv:2110.14168 (2021).

[4] Hendrycks, Dan, et al. "Measuring mathematical problem solving with the math dataset." arXiv preprint arXiv:2103.03874 (2021).

[5] Lewkowycz, Aitor, et al. "Solving quantitative reasoning problems with language models." NeurIPS, 2022.

[6] Azerbayev, Zhangir, et al. "Llemma: An open language model for mathematics." arXiv preprint arXiv:2310.10631 (2023).

[7] Hendrycks, Dan, et al. "Measuring massive multitask language understanding." arXiv preprint arXiv:2009.03300 (2020).

Thanks for the responses. They do improve the paper and thus I increase the score by 2 for the additional experiments and discussions. However, I still think my core questions are not addressed.

For the topic of "pre-training", the experiment scale in this paper (even considering the 7B experiments in the rebuttal) is not sufficient. How well the Quad method generalize to larger model is questionable. When the model becomes large enough, do we still need to select a subset of good examples instead of focusing on getting more high quality data? I understand due to practical reasons (e.g., compute budgets), it's hard to do large scale experiments in some settings, but again, for the "pre-training" research, this is a challenge we have to face.

Second, the authors claim one advantage of Quad is "model-aware data selection" that is customized to the downstream tasks. However, I think the goal of "pre-training" is to get a good foundation (unaligned) model that is good for later-stage post-training/alignment. If we just want to get good scores for a few downstream tasks, we don't do it in the pre-training stage.

Third, related to the problem setting, in real-world, when will people do 7B-scale model pre-training from scratch? Usually people either start from a pre-trained checkpoint then do post-training/alignment (for a few specific tasks), or do continuous pre-training/tailpatch training (for knowledge enhancement). If we really need a small scale pre-train ckpt, most likely we will distill a large pre-train ckpt to a smaller model. So how useful this problem setting should be further discussed.

Finally, it would be interesting to show Quad-7B-400B and compare with Clean-7B-400B and see the performance gap.

To summarize, I appreciate authors' due diligence on improving the paper quality, but I still think there are many questions need to be answered in this work.

Dear reviewer $`R56B`$

Thank you for your thoughtful efforts in reviewing our work! As the discussion period nears its end, we kindly hope you can take a moment to review our rebuttal. Looking forward to your feedback!

Best Wishes!

ICLR Authors

Thanks for the detailed response and I appreciate authors' efforts on addressing some of my concerns. I go over more related literatures in the "pre-training data selection" field and agree that 7B scale experiments are a generally acceptable scale in academia. Therefore, I will increase the score by 1 and am OK to see this paper accepted. Hopefully it can benefit the community and inspire for more data-centric work in the field.

Below are some of my add-on thoughts and suggestions on this work.

First, I think the authors should further emphasize the continuous pre-training part and conduct more experiments. The current experiments on math benchmark is not really sufficient. It only tests the model's math capabilities on 5 datasets. In some sense, this is already very close to the post-training stage and we are aligning model for more math use cases. The ideal experiment setting is to check model performance on tasks that require new knowledge from continuous pretrain dataset, meanwhile reporting the model performance using pretrain metrics like perplexity or few-shot ICL numbers. This leads to another question, when your reference dataset is huge, will the influence score computation very small (or even possible if large dataset is coupled with large model).

Second, the practical meaning of problem setting -- "when will people do 7B-scale model pre-training from scratch?" is still not really answered. For big tech companies like Google, Meta, Nvidia, they will not really train such a small LLM from scratch. Instead, they will have computes to train an O(100B) large model and distill it into smaller models. For example, Gemma-2 2B and 9B models are distilled from 27B model (https://arxiv.org/html/2408.00118v1), llama 3.2 1B model is distilled from llama 3.1 8B model, and Nvidia distills llama-3.1 8B to llama-3.1-Minitron 4B (https://developer.nvidia.com/blog/how-to-prune-and-distill-llama-3-1-8b-to-an-nvidia-llama-3-1-minitron-4b-model). I personally think more discussions on continuous pre-training is more meaningful and valuable.

Finally, I would suggest authors expanding its K-FAC discussion with more discussions in the appendix, maybe adding some pseudo code.

To summarize, even this work has some limitations and weaknesses, I agree it follows the subfield norms and standards and thus can be accepted.

$**Overall response.**$

Thanks for your insightful feedback! Based on your comments, we have added a 7B model and 6 additional experiments to demonstrate the robustness and effectiveness of $`Quad`$. We also clarify your concerns about hyperparameter tuning, the DSIR baseline, significance test, etc. We promise to include the above discussions into the final version if the paper is accepted.

$**Question 1:**$ Could the authors report the results of statistical significance tests?

$**Response:**$

Thanks for your suggestions to improve our paper! Based on the reviewer's advice, we run the experiment six times and report the statistics significance tests. The results are as follows:

Based on the above results, we perform a significance test utilizing t-tests. We make the following two hypothesis. (1) Null hypothesis ($H_0$): The average accuracy of $`Quad`$ does not exceed the average of baselines; and (2) Alternative hypothesis ($H_1$): The average performance of $`Quad`$ exceeds that of the baselines. We set the confidence level as 99%, with a significance level of $\alpha=0.01$. The p-value represents the probability of observing the test statistic under the assumption that $H_0$ is true. A small p-value indicates a lower likelihood of observing the current result if $H_0$ holds.

For each baseline, we can observe $p<\alpha$, and thus we reject $H_0$ and accept $H_1$, indicating that the average performance of $`Quad`$ is greater than that of baselines. Since the confidence level of t-test surpasses 99%, we can conclude that the superiority of $`Quad`$ compared to other baselines is statistically significant. We have included the above discussion in Appendix C of the revised paper.

$**Question 2:**$ How does the selection of clustering algorithm and number of cluster affect the performance of the proposed method?

$**Response:**$

Thanks for your suggestion. In response to your question, we add an ablation study regarding the clustering algorithms and number of clusters.

Firstly, we add an ablation study by varying the number of clusters, as shown in the table below. Using the Slimpajama dataset, the Elbow [1] method selected $k=10,000$ as the optimal number of clusters in our original paper. From the following table we can observe that the model consistently performs well when the cluster number is close to 10,000. Conversely, a very small number of clusters (i.e., $k=1,000$) leads to poor performance due to the high variance of influence scores in each cluster. Thus, the sampled instances cannot well represent the cluster. Similarly, when the cluster number is too high(i.e., $k=500,000$), there will be many clusters that are in fact contain similar data instances, and thus it is relatively hard to explore diverse clusters, thereby leading to the performance degradation.

Moreover, we evaluate the performance of several typical clustering methods including GMM[2], DBSCAN[3] and MeanShift[4]. Considering that the clustering results are affected by the parameters of clustering algorithms, we use different methods to select proper parameters. For GMM, we can use the AIC score [5] to determine the appropriate number of components. For DBSCAN, there are 2 key parameters: (1) $eps$(the radius of a neighborhood w.r.t. some data points) and (2) $minPts$ (a data point is considered as a core point if at least $minPts$ data points are within $eps$ of it). They can be set using the method in [6]. Mean-Shift is a centroid-based method that updates the centroids to be the mean of the points within a given region. The size of the region is controlled by $bandwidth$, which can be set by the estimation of the bandwidth [7]. In this set of experiments, the experimental settings of pretraining 1.3B model are consistent with those reported in our paper. The above result shows that $`Quad`$ is not sensitive to the clustering algorithms.

We have included the above discussion into Appendix D of the revised paper.

$**Overall response:**$

Thank you for the valuable comments to improve our paper! Based on your comments, we have added a candidate dataset, 2 reference datasets, 7 downstream tasks and 5 additional experiments to demonstrate the generalizability and robustness of $`Quad`$. We also clarify your concerns in the weak points. We promise to include the above discussions into the final version if the paper is accepted.

$**Question 1:**$ Clustering Quality and Sensitivity

$**Response:**$

We appreciate your comments to improve our paper. In line with the reviewer's suggestion, we have incorporated an ablation study regarding the number of clusters and provided a discussion on how the results are affected by varying the number of clusters.

Firstly, in our paper, we used the well-known Elbow [1] algorithm to select the optimal number of clusters for k-means algorithm used in $`Quad`$. The basic idea is to select the best $k$ based on the metric Within-Cluster Sum of Squares (WCSS), which is the sum of squared distances between each data instance and its cluster center, i.e., WCSS=$\sum_{i=1}^k\sum_{x \in C_i}\|x-\mu_i \|$. Elbow begins with a small $k$, and with $k$ increasing, WCSS first decreases rapidly and then slows down. Then, we identify the ``elbow point'' where the decreasing rate becomes slow as the best $k$.

Secondly, we add an ablation study by varying the number of clusters, as shown in the below table. Using the Slimpajama dataset, the Elbow method selected $k=10000$ in our original paper. From the below table, we can observe that the model consistently performs well when the cluster number is close to 10000. However, a very small number of clusters (i.e., $k=1000$) leads to poor performance due to the high variance of influence scores in each cluster. Thus, the sampled instances cannot well represent the cluster. Similarly, when the cluster number is too high (i.e., $k=500000$), there will be many clusters that are in fact contain similar data instances, and thus it is relatively hard to explore diverse clusters, thereby leading to the performance degradation. We have included the above discussion into Appendix D of the revised paper.

$**Question 2:**$ Generalizability and Robustness

$**Response:**$

Thanks for your feedback to improve our paper. To show the generalizability and robustness of our method, we add FLAN [2] (the mixture of multiple NLP tasks) and Openwebmath [3] (a math-related dataset) as reference datasets, associated with 7 new downstream tasks. We also add one candidate dataset and a 7B model.

Specifically, we first use FLAN as a new reference dataset to train a 1.3B model on a new candidate dataset (i.e., FineWeb [4]). In this experiment, 100B tokens are selected for training using $`Quad`$ and other baselines. As shown in the table below, $`Quad`$ still demonstrates superior accuracy compared to all baselines, such as the DSIR, Semdedup, PPL, Qurating and MATES. For example, we can observe that $`Quad`$ has an improvement of 0.97% on model accuracy compared with the Qurating, which is a state-of-the-art baseline. In addition, we add two downstream tasks (i.e., WikiText [5], HelloBench [6], and for WikiText, the lower the score, the better the model performance) about long-text generation tasks. As shown in the table below, on the candidate dataset Slimpajama, for the 7B model, we can observe that $`Quad`$ still outperforms other baselines because we simultaneously consider the quality and diversity in data selection.

Moreover, we also add the Openwebmath dataset as another new reference dataset to train a 7B model on the Slimpajama dataset, which aims at selecting data to improve the mathematical skill of the model. Correspondingly, we also include GSM8K[7], MATH[8], OCW[9], SAT[10], and MMLU STEM[11] as new downstream tasks to evaluate the model's mathematical capabilities, with the final accuracy calculated as the average score across these tasks. As shown in the table below, $`Quad`$ still demonstrates superior model performance compared to all baselines.

This experiment shows that our method can generalize well to various types of reference datasets and downstream tasks.

We have included the above discussion in Appendix G and I of the revised paper.

Dear reviewer 1sUi,

Greetings from the authors!

We appreciate your insightful feedback. Regarding the clustering, additional tasks and datasets as well as other concerns, we have conducted multiple experiments and included these experiments and analysis in the revised paper. We will include them in the final version and provide the link to our GitHub repository upon acceptance.

If you have any question for our paper, please feel free to point out and we will try to address it as soon as possible. We would like to express our sincere gratitude for your valuable comments again. Thanks and looking forward to your reply!

Dear reviewer $`1sUi`$,

Wishing you a happy and blessed Thanksgiving!

We would like to express our gratitude to your valuable feedback. We have carefully considered all suggestions and updated our submission accordingly.

If you have any question for our paper, please feel free to point out and we will try to address it as soon as possible. We would like to express our sincere gratitude for your valuable comments again. Thanks and looking forward to your reply!

Best wishes!

Authors

Dear reviewer $`1sUi`$

Thank you for your thoughtful efforts in reviewing our work! As the discussion period nears its end, we kindly hope you can take a moment to review our rebuttal. Looking forward to your feedback!

Best Wishes!

ICLR Authors