SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In Section 4.1 we find that by training a 6.7B parameter model on data selected by D4, we reach the final perplexity of a baseline model using  $20\%$  fewer model updates. In our particular setup, this translates to saving approximately 4300 GPU hours - we will refer to this as the naive efficiency gain as it does not account for the cost of computing the selection metric.
To demonstrate our method's practicality, we must ensure the cost of selecting data is significantly less than this. As described in Section 3.4 selecting data via D4 involves: first, embedding documents via a 125M OPT model; second, computing K-Means indices + distance to indices. The first step is completed on a single machine with 96 CPU cores in approximately one day. Given the two orders of magnitude difference between the prices of CPU and GPU cores we consider this cost negligible. For the second step, embedding 400B tokens with a 125M parameter model takes approximately 888 GPU hours, using the same A100 GPUs. Subtracting this from the naive efficiency gain of 4300 GPU hours, we arrive at an overall efficiency gain of 3412 GPU hours. This is how much compute D4 saved us in practice when training our single 6.7B parameter model. In Figure 4 we redo this calculation for different model
Figure 4: Naive and overall efficiency gain of data selection via D4 relative to the total cost of training as a function of model size on Instruct + Answers perplexity at  $R = 0.25$ .
Source: https://aws.amazon.com/ec2/pricing/on-demand/
increases with model size. Based on this, we can conservatively estimate that D4 would have overall efficiency gains of  $20\%$  for LLama-65B [50] and  $22\%$  for OPT-175B [59].
4.4 Analysis of D4
4.4.1 Why does data selection hurt performance on web snapshots?
Figure 5: Left: Train-test similarity across validation sets. X-axis denotes the name of the validation set (refer to Section 3.4 for more information about each validation set), and y-axis denotes the cosine distance to the nearest neighbor in the training set for the 1.3B OPT 40B baseline (the green triangle denotes mean, and the yellow bar denotes median). We observe that web-snapshots validation sets are closest to points in the training set and are disproportionately affected by data selection. Right: Analysis of C4 validation set. (Top): Histogram of cosine distance to nearest neighbor in train. For each bin, we show the mean original perplexity (middle) and mean difference in perplexity after data selection (bottom). "Easy" (low original ppl) points close to the training set are generally the points most affected by data selection.
While we observe consistent average perplexity improvements, Section A.3 demonstrates that this perplexity improvement varies greatly across validation sets. More importantly, data selection always impairs performance on web snapshot validation sets such as CC-dedup, CommonCrawl, and C4. To investigate why this occurs, we embed each validation set into the same embedding space as the training set and search for the nearest neighbors to validation points in the training set for our 1.3B baseline model. In the left plot of Figure 5 we show that validation sets drawn from the same distribution as web-snapshots are substantially closer to training set compared to other validation sets. The right plot of Figure 5 shows that data selection disproportionately affects web-snapshot validation sets. In the top-right plot, we see that web validation sets reside in regions of the embedding space which are sparsified as a result of data selection (e.g. regions of space close to cluster centroids in the training set), and in the bottom-right plot we see that these points are also the most affected by data selection, since their perplexity after data selection significantly increases. Moreover, the middle-right plot shows that these validation points have the lowest perplexity before pruning indicating that these points are "easy" points, perhaps due to their proximity to the training set.
Given that some of our validation sets are extremely close to the training set, we question whether they are still strong indicators of generalization. In fact, in Figure 6 we find evidence of a slight inverse relationship between perplexity on web snapshots and more robust indicators of LM ability, such as perplexity on instruction-tuned datasets and downstream accuracy. In contrast, we observe that perplexity on Instruct+Answers is positively correlated with downstream accuracy, suggesting that validation perplexity on instruction tuned data is a better measure of model quality. For this reason, we group most of our results in Section 4 into Web Snapshots and Non-web Snapshots (which consists of Web-Derived + Web-Independent from Figure 5, see Section A.1.4 for a full-list of validation set names).
Figure 6: Correlation between (left): negative Instruct+Answers perplexity and negative web snapshot perplexity, (middle): Downstream accuracy and negative web snapshot perplexity, (right): Downstream accuracy and negative Instruct+Answers perplexity. Each point is one training configuration (1.3B OPT model, 40B tokens), with the only change being the data selection method and pretraining seed. Web snapshot perplexity is slightly negatively correlated with stronger indicators of LM ability.
4.4.2 Importance of re-clustering between SemDeDup and SSL Prototypes
As mentioned in Section 3.4 we hypothesize that sparsifying dense regions of space containing excessive semantic duplicates improves the clustering quality and is, therefore, critical to the performance of D4. To isolate the effect of re-clustering on D4, we run experiments with a version of D4 where we remove the re-clustering step (e.g. we keep the original clustering). As shown in Figure 7, omitting the re-clustering step significantly worsens performance, and we observe in the rightmost plot of Figure 7 that SemDeDup indeed removes extremely dense clusters surrounding centroids (e.g. duplicate-driven clusters). We analyze this in more depth in Section A.9
Figure 7: Investigating the necessity of the re-clustering step in D4. We see that re-clustering improves perplexity across Web snapshots (left), Non-web snapshots (middle-left), and Instruct + Answers (middle-right). Right: Empirical CDF of mean distance to centroid, with and without re-clustering. Re-clustering removes duplicate driven clusters (clusters with low mean distance to centroid).
5 Summary and Limitations
We introduced D4, a method for data curation on LLMs that improves training efficiency by  $20\%$  across multiple model scales, with larger gains at increased model scale. We also demonstrated that, in contrast to common practice, repeating data via epoching can be beneficial for LLM training, but only if the data subset is intelligently selected. While we have shown encouraging efficiency gains and performance improvements via D4, our work has several limitations and many future directions.
Mixing different training distributions: While we chose one data distribution to both select data and train on, modern LLM setups usually mix different data sources. Our method is likely complimentary to such pipelines: practitioners may use D4 to diversify and de-duplicate individual data sources and then mix data sources to provide additional diversity in their training dataset. We leave exploring the efficacy of D4 on a mix of training distributions as future work, but expect that this will yield further gains by reducing redundancy across datasets as well as within datasets.
Model scale: Due to compute limitations, the largest models we evaluated were 6.7B parameters trained on 100B tokens. While, to our knowledge, this is the largest to date application of embedding based data curation approaches, further investigation at model scales exceeding 100B would be very interesting, particularly in light of our observation that the efficiency gain grows with model scale.
The authors would like to thank many people who helped bring this work to fruition: Srini Iyer, Yuchen Zhang, Todor Mihaylov, Jacob Xu Moya Chen, Mansheej Paul, Mitchell Wortsman, Amro Abbas, Aaditya Singh, Myra Cheng, and Matthew Leavitt. The authors would also like to thank Surya Ganguli, Mona Diab, and Xian Li for initial brainstorming and are grateful for help with compute infrastructure given by Henry Estela and Victoria Lin. Lastly, the authors would like to thank anonymous reviewers for improving the quality and writing of this paper.