Beyond Scale: The Diversity Coefficient as a Data Quality Metric for Variability in Natural Language Data

"Importantly, the paper does not adequately account for potential confounding factors. When merging two datasets to create a more diverse one, the authors do not consider that other characteristics of the combined dataset may influence performance, aside from diversity."

This is a valid concern, and we appreciate the opportunity to clarify our approach. To minimize confounding factors, we ensured that merged datasets maintained the same token count and similar distributional properties. Furthermore, we used out-of-distribution (OOD) evaluation datasets (e.g., OpenWebText2 and C4) to mitigate potential biases introduced by training data similarity.

While no experimental setup can completely eliminate confounding factors, we believe our approach provides a strong foundation for isolating the effects of diversity. We will explicitly address these methodological considerations in the revised manuscript. To address this, we designed our experiments to minimize potential confounds when merging datasets. Specifically:

Controlled Token Count: We ensured that merged datasets maintained a consistent token count with their constituent datasets. This controlled for size as a variable influencing performance.

Evaluation on OOD Data: Our evaluation was conducted on out-of-distribution (OOD) datasets, such as OpenWebText2 and C4, which are distinct from the training datasets. This approach reduces the impact of spurious similarities between training and evaluation datasets, ensuring that the observed performance improvements are likely due to increased diversity rather than content overlap.

Cross-Diversity Analysis: To further mitigate biases, we analyzed cross-diversity coefficients, which demonstrate that the individual datasets (PubMed, USPTO) and their combination have similar levels of dissimilarity with OOD datasets. This reduces the likelihood that performance differences are due to specific content alignments rather than diversity.

While no experimental design can entirely eliminate confounding factors, we believe our methodology provides strong evidence for the role of diversity in driving performance improvements. We will clarify these considerations in the revised manuscript to better reflect our efforts to control for confoun

"The evaluation is based solely on cross-entropy scores from some evaluation data, which is not a strong or convincing measure. Even if all results showed a consistent trend (which they do not), the claim that increased diversity leads to better downstream performance is weakly supported."

Cross-entropy is widely used in language modeling research as a reliable proxy for generalization performance. By evaluating on OOD datasets (OpenWebText2 and C4), we demonstrate that higher diversity improves generalization to unseen data. Given our compute limitations, we weren't able to train for more tokens to get accuracy changes on "emergent" metrics like accuracy. In addition, when finely controlling the merging of data sets to create more diverse data sets, we had to trade off in favour of scientific interpretability perfect accuracy. No scientific experiment is perfect but we believe we controlled for the most important factors with challenging OOD data sets with metrics that could measure in statistically significant ways with relatively high R^2 values what we wanted to study.

"The paper’s structure and writing quality leave much room for improvement. Greater clarity in both the organization of ideas and the precision of language would significantly enhance the paper's readability and overall impact."

Thank you for this constructive suggestion. We will revise the manuscript to improve the organization of ideas and refine the language for clarity and precision. Additionally, we will enlarge the figures and ensure all visual elements are easily interpretable.

We thank the reviewer for their thoughtful feedback and for recognizing the strengths of our work.

For calibration purposes, we would like to clarify the ICLR 2025 review rubric compared to other conferences like NeurIPS 2024. Notably:

A score of 8 on the ICLR 2025 rubric indicates an "Accept," whereas NeurIPS 2024 uses 7 for this designation. Similarly, a "Strong Accept" is represented by a score of 10 at ICLR 2025 compared to 9 at NeurIPS 2024. We address the specific points raised:

"The authors assert that their work represents a paradigm shift in data-centric machine learning through the introduction of data diversity (L. 59). This claim is factually incorrect, as many previous studies have explored data diversity. Several relevant papers, which I list at the end of this review, were overlooked."

We appreciate the reviewer highlighting the broader context of related work. While we acknowledge that data diversity has been explored in prior studies, our contribution lies in introducing a formal and quantifiable measure of diversity—the diversity coefficient—which is specifically tailored to pre-training datasets for large language models. This metric provides a systematic way to compare datasets, which was not previously available in the literature.

We will revise the phrasing in the manuscript to ensure our claims better reflect the novelty of our contribution without overstating its scope. Additionally, we will incorporate the suggested references to situate our work more comprehensively within the existing body of research.

"The range of values produced by the metric is not intuitive, with a lower bound around 0.05 and an upper bound of 0.4, which feels inelegant."

This can be easily fixed. We can do a subtraction of the minimum and division of the maximum. We considered this but we don't think it's a substantial contribution except aesthetics rather than deep scientific understanding. We are happy to do so in the camera ready upon acceptance.

"The authors claim that higher diversity leads to better performance (Figure 2), but this relationship does not hold consistently. In some cases, performance declines with increased diversity, which the authors fail to address."

We recognize the importance of addressing exceptions to our general trend. While the majority of our experimental results demonstrate that higher diversity correlates with better generalization (as shown in Figure 2), we acknowledge that there are outlier cases where this relationship does not hold. These exceptions are noted, and we believe they are attributable to the influence of other factors (e.g., dataset content or domain-specific biases). We will explicitly discuss these outliers in the revised manuscript and provide additional context to clarify our findings. To increase robustness of our results, we included 26 regression lines and R^2 to back up our experimental findings.

Though it is true that in some cases, performance declines with increased diversity, 7/15 runs (a slight minority) have USPTO+PubMed-trained models having lower loss than only PubMed-trained models (for both OpenWebText2 and C4 evaluations), including nearly all of the runs with large datasets i.e. training for over a billion tokens (e.g. GPT2-810M-2B, GPT2-51M-1.31B, GPT2-345M-2.2B, GPT2-117M-2.2B).

In all cases, PubMed+USPTO outperforms USPTO alone (a much less diverse dataset).

For ref: In particular, the highest diversity model has the lowest loss for GPT2-810M-2B, GPT2-51M-1.31B, GPT2-345M-2.2B, GPT2-51M-550M, GPT2-117M-2.2B, LLAMA2-7B-Run4-6.36M, LLAMA2-7B-Run1-6.36M

Claim: “PRE-TRAINING IN HIGHER DIVERSITY LEADS TO BETTER EVALUATION PERFORMANCE” (Section 3.1)

Figure 2 provides robust evidence for the relationship between dataset diversity and improved evaluation performance. It includes 26 linear regression lines across various model sizes, training runs, and parameter counts, demonstrating consistent trends supported by R2 values. The evaluation was conducted on out-of-distribution (OOD) datasets, OpenWebText2 and C4, rather than validation splits from training data. This approach ensures the results reflect generalization to OOD tasks and strengthens the evidence for the causal role of diversity.

While diversity alone may not be the sole determinant of performance, our controlled experiments consistently show that higher diversity improves generalization.

Although linking diversity to performance is indeed a critical aspect of our paper, an equally significant aspect of our paper is the introduction of the diversity coefficient, which provides us a reliable measure of pre-training dataset diversity (aside from diversity’s impact on performance). Hence, we believe the merits of this latter argument are important to consider as part of the main thrust of paper, in addition to relating diversity to performance.

Q: PubMed+USPTO often has higher loss than PubMed.

In 7/15 runs (a slight minority) have USPTO+PubMed-trained models having lower loss than only PubMed-trained models (for both OpenWebText2 and C4 evaluations), including nearly all of the runs with large datasets i.e. training for over a billion tokens (e.g. GPT2-810M-2B, GPT2-51M-1.31B, GPT2-345M-2.2B, GPT2-117M-2.2B). In all cases, PubMed+USPTO outperforms USPTO alone (a much less diverse dataset). For ref: In particular, the highest diversity model has the lowest loss for GPT2-810M-2B, GPT2-51M-1.31B, GPT2-345M-2.2B, GPT2-51M-550M, GPT2-117M-2.2B, LLAMA2-7B-Run4-6.36M, LLAMA2-7B-Run1-6.36M.

Q: Highly unrepresentative datasets used for pre-training.

We acknowledge that the datasets used (PubMed, USPTO, PubMed+USPTO) are not anymore representative of modern pre-training corpora like C4 or The Pile. However, these datasets were intentionally chosen for their distinct, non-overlapping domains, allowing us to isolate and highlight diversity contrasts effectively. By selecting single-source datasets, we minimized confounding factors and created a clearer signal on how diversity impacts performance.

While larger, more representative datasets like C4 or RefinedWeb could be explored, they are typically filtered web scrapes with significant domain overlap. Combining such datasets is unlikely to provide as much variation in diversity as distinct single-source datasets like PubMed and USPTO. This lack of sufficient diversity “stimulus” would likely obscure the effects of diversity on downstream performance.

We note that PubMed and USPTO are subsets of The Pile, a widely used pre-training corpus, making them relevant for evaluating diversity’s role. Expanding the dataset selection remains a valuable direction for future work, but it introduces challenges in identifying unbiased evaluation corpora that do not arbitrarily favor specific datasets, potentially confounding the results.

Only 3 datasets considered. This is way too small for linear regression. Should be more variation in datasets, rather than in model size and dataset size. Would be much more convincing with one model/dataset size, but many many datasets.

Already addressed.

Expanding the set of datasets has the drawback that it becomes harder to find an evaluation dataset which is not biased toward any of these datasets, i.e. which is not arbitrarily similar (or dissimilar) from one of the used training datasets, irrespective of the dataset’s diversity. This creates a confounder of performance, as less diverse datasets might perform as well as or better than more diverse datasets merely due to an arbitrary similarity in the distributions of the training and eval data, obscuring the effect of dataset diversity.

Q: Choice of evaluation metric of PPL on C4/etc. is unsubstantiated. Is its superior performance just because it’s more similar to C4/etc.?

Perplexity (PPL) is widely used to evaluate general language modeling performance and strongly correlates with downstream task success. By evaluating on OpenWebText2 and C4—distinct from the training data—we demonstrate generalization to OOD tasks. Using our cross-diversity coefficient as a proxy for dataset similarity, we find that PubMed and USPTO are similarly distant from these evaluation datasets. This supports our conclusion that performance improvements are driven by diversity rather than dataset similarity.

Additionally, perplexity on web-scraped datasets has been shown to predict downstream performance across a range of tasks, further validating its use as an evaluation metric. Future work could incorporate task-specific benchmarks for additional insights.

Quoted Reviewer Questions and Responses

Line 122, which t is used to compute the FIM matrix for sequences of a batch? Are all the steps from the sequence averaged? Or is FIM only built based on the last step?

Thank you for raising this question. The $t$ in the equation at line 122 refers to each step within a sequence. The Fisher Information Matrix (FIM) is computed by averaging over all steps within each sequence, and the resulting averages are then aggregated across the batch.

I suggest the authors rewrite some pieces of text. There are places where authors use phrases like "by them," "they...," and it is not clear what the referents are, which makes the reading and understanding of the paper more involved.

We appreciate this constructive feedback and apologize for any lack of clarity. We will rewrite these ambiguous references in the final version to ensure the text is more precise and easier to follow.

Task2Vec should be fully and better described in the main paper (including how it was used before, its intuition, etc.).

We agree that a more comprehensive explanation of Task2Vec would strengthen the paper. We will expand its description in the main body, incorporating details about its prior uses, the intuition behind its design, and its specific relevance to this work. This additional context will draw on the discussion currently provided in Appendix C.

Section I.2 should mention the data used to pre-train GPT-2.

Thank you for this suggestion. In Section I.2, we will specify that the pre-trained GPT-2 model used in our Task2Vec embedding procedure was trained on OpenAI’s WebText dataset, as described in the original GPT-2 release.

Minor Comments

Figure 1 does not seem to add much to understanding. Maybe an algorithm would be more useful?

We appreciate this perspective. While we believe the visualization in Figure 1 is valuable for readers unfamiliar with elements of the diversity coefficient computation, we agree that including an algorithmic breakdown could provide additional clarity. We will consider adding the algorithm as an appendix while retaining the figure for accessibility.

All figures are too small to read.

Thank you for this observation. We will enlarge all figures in the final version to improve readability and presentation quality.

Line 751, "this more sophisticated aggregation method"—which aggregation method?

We appreciate this comment and recognize the lack of specificity. We will revise the text to explicitly state and clarify the method to ensure clarity for the reader.

We thank the reviewer for their thoughtful feedback and for recognizing the strengths of our work.

For calibration purposes, we would like to clarify the ICLR 2025 review rubric compared to other conferences like NeurIPS 2024. Notably:

A score of 8 on the ICLR 2025 rubric indicates an "Accept," whereas NeurIPS 2024 uses 7 for this designation. Similarly, a "Strong Accept" is represented by a score of 10 at ICLR 2025 compared to 9 at NeurIPS 2024. We address the specific points raised:

While the authors evaluate performance against pre-training data diversity (section 3.1), the metric for performance is cross-entropy loss in LM, I wonder whether specific task evaluation on NLP tasks would make sense here (many NLP benchmarks used to benchmark LLMs on tasks like question answering, GLUE, etc.).

Thank you for this thoughtful comment. We agree that evaluating specific downstream tasks can provide valuable insights in certain contexts. However, in this work, we focus on evaluating pre-trained base models that are not designed for instruction-following or fine-tuning on specific tasks. For the following reasons, we believe cross-entropy loss is the appropriate metric for our study:

Scope of the Work: This paper focuses on understanding the impact of pre-training data diversity on base model performance. Cross-entropy loss directly measures the model’s ability to predict general natural language, which aligns with our goal of providing insights into pre-training dynamics. Evaluating task-specific benchmarks, while interesting, extends beyond the scope of this work.

Applicability to Non-Instructive Models: The models studied here are pre-trained base models without instruction-tuning or task-specific fine-tuning. Benchmarks such as GLUE or question answering are designed for fine-tuned or instruction-following models. As such, they may not provide meaningful insights into the pre-training stage we are evaluating.

General Predictive Power of Cross-Entropy: Cross-entropy loss is a widely recognized and robust metric for evaluating the quality of pre-trained language models. Prior work (e.g., Chinchilla scaling laws) has shown that cross-entropy loss on a general text corpus strongly correlates with downstream task performance, making it a reliable proxy for assessing the pre-trained base model's potential.

We will ensure the paper clearly communicates the rationale for using cross-entropy loss as our primary metric and highlights the scope and focus of this study. Thank you for this valuable feedback.

The Vendi Score seems to be another approach to compute diversity, why was it not included as a baseline in the main experiments (at least in main Table 1)?

Thank you for raising this point. The Vendi Score was developed independently and concurrently with our work, as reflected in the publication timelines. We later connected with its authors and respect their contributions.

Our focus is on demonstrating through causal experiments (e.g., Figure 2) that increased pre-training data diversity improves general text performance, including on out-of-distribution tasks. This contribution goes beyond metric validation, which is the primary focus of the Vendi Score.

While we could include Vendi Score comparisons, we believe it would not meaningfully strengthen our core findings. We are open to discussing its relevance further if needed. Thank you for your constructive feedback.

In addition, the Vendi Score does not inherently define a similarity function for comparing data points, such as text, requiring users to specify this aspect. Additionally, it has not been tested on text datasets, as its experiments focus on tabular and image data. This limits its direct comparability to our diversity coefficient, which is specifically designed for and validated on natural language datasets.

We used Table 1 to showcase the diversity coefficient’s results and to validate its alignment with intuitive notions of diversity in natural language datasets. Our goal was to demonstrate diversity in PT leads to better performance:

Even if we had known about the Vendi Score during our experiments, its inclusion would not have been necessary for several reasons.

Focus on Natural Language: The Vendi Score was not tested on text data, focusing instead on tabular and image data. Our work specifically addresses natural language datasets, making the metrics not directly comparable.

Undefined Similarity Function: The Vendi Score relies on a user-defined similarity function, while our diversity coefficient incorporates a predefined approach tailored to textual data.

Core Contribution: The primary contribution of our work is the demonstration, through causal experiments (e.g., Figure 2), that increased diversity leads to improved performance on general text tasks, including out-of-distribution generalization. Benchmarking against the Vendi Score would not strengthen or clarify this key result.

We thank the reviewer for their thoughtful feedback and for recognizing the strengths of our work.

For calibration purposes, we would like to clarify the ICLR 2025 review rubric compared to other conferences like NeurIPS 2024. Notably:

A score of 8 on the ICLR 2025 rubric indicates an "Accept," whereas NeurIPS 2024 uses 7 for this designation. Similarly, a "Strong Accept" is represented by a score of 10 at ICLR 2025 compared to 9 at NeurIPS 2024. We address the specific points raised:

The use of cosine distance between task2vec embeddings as a diversity measure is not well motivated. Can we consider other textual embedding methods? What are the tradeoffs?

We are grateful for this constructive comment, which allows us to clarify our work further. While some of our experiments use synthetic datasets to test specific properties of the diversity coefficient (e.g., the relationship between vocabulary size and diversity), other experiments directly evaluate real natural language datasets.

Figure 3: This demonstrates the application of the diversity coefficient to real-world datasets, such as PubMed, USPTO, and Pile-CC, highlighting how the metric aligns with intuitive understandings of variability in natural datasets. For example, it distinguishes between datasets with domain-specific focus (PubMed, USPTO) and broader coverage (Pile-CC), showcasing its relevance to real-world language.

Figure 4: Although synthetic, these experiments capture a key property of natural language: larger vocabularies correspond to increased diversity. This provides causal insight into the diversity coefficient's ability to measure meaningful aspects of language variability.

We would be happy to enhance the paper by explicitly discussing how Figures 3 and 4 collectively validate the diversity coefficient's representation of real natural language. Thank you again for this helpful feedback, which will help us improve the clarity and impact of our work.

"Use of cosine distance between Task2Vec embeddings as a diversity measure is not well motivated. What about other textual embedding methods?"

Thank you for raising this point. We chose Task2Vec embeddings for their demonstrated effectiveness in capturing semantic relationships across various domains and data types, including language and images (Achille et al., 2019). This versatility and robustness make Task2Vec particularly well-suited for our diversity coefficient, as it provides a reliable foundation for comparing the variability of datasets.

While alternative embedding methods could also be explored, we found Task2Vec to be effective and straightforward for our purposes, which is why it was selected for this study. We appreciate the suggestion to consider additional methods and are open to incorporating comparisons in future work.

The experiments on synthetic datasets are valuable, however the conclusions could be further strengthened by a study utilizing human annotation for diversity.

Thank you for this thoughtful suggestion. While we appreciate the potential value of human annotations, we believe they are not necessary in this context for several reasons:

Potential Noise in Human Annotations: Diversity is inherently subjective, and human evaluations often introduce variability and noise due to differences in individual interpretations of diversity. This could reduce the reliability and consistency of conclusions drawn from such studies.

Redundancy with Existing Experiments: Our paper already includes intuitive and systematic experiments that validate the diversity coefficient:

Figure 3 demonstrates how the diversity coefficient aligns with expected variability across real-world datasets (e.g., PubMed, USPTO, and Pile-CC). Figure 4 shows causal relationships in synthetic datasets, such as the impact of increasing vocabulary size on diversity, further supporting the validity of our metric. These results strongly suggest that the diversity coefficient captures meaningful aspects of variability, making additional human annotations redundant.

High Cost of Human Evaluations: Conducting human evaluation studies for diversity across large-scale datasets would be prohibitively expensive and time-intensive, particularly given the scale of modern natural language datasets. One of the key strengths of the diversity coefficient is its ability to provide a scalable and objective alternative to such costly manual efforts.

We appreciate your thoughtful feedback and will consider adding a discussion in the paper to clarify these points. Thank you for helping us better articulate the robustness of our existing validation and the practical advantages of our approach.

[1] Task2Vec: Task Embedding for Meta-Learning. Alessandro Achille, Michael Lam, Rahul Tewari, Avinash Ravichandran, Subhransu Maji, Charless Fowlkes, Stefano Soatto, Pietro Perona