Does your data spark joy? Performance gains from domain upsampling at the end of training

Introduction

Pretraining datasets for large language models (LLMs), such as Dolma (Soldaini et al., 2023), have grown to trillions of tokens. To accommodate such large scales, they are typically composed of two types of data sources. First, they contain large amounts of web scraped data processed from CommonCrawl (CC) dumps. These are typically hundreds of billions to trillions of tokens in size and contain a diverse distribution of information. However, because of their size, they are necessarily less information dense and are not as filtered. Second, LLM pretraining mixes contain datasets that either target certain domains or come from single high quality sources. These are much smaller (often less than a hundred billion tokens tokens). They are also more carefully processed and are dense with information from domains we want LLMs to be good at; however, since their sources are limited, they are often less diverse (Computer, 2023).

One of the biggest challenges to pretraining LLMs is determining the optimal strategy for mixing datasets that come from CC and smaller domain specific sources. Some previous works have opted to pretrain entirely on heavily processed CC data (Penedo et al., 2023). Others have used different heuristics to balance between CC and more domain specific datasets (Computer, 2023). However, most recent language models trained at scale disclose limited information on the contents of their pretraining data (Touvron et al., 2023; Jiang et al., 2023; 2024; Team et al., 2024). At smaller scales, there have been attempts to algorithmically optimize the data mix proportions, but these methods have not been openly validated at the scale most modern language models are trained (Xie et al., 2024). Given the sheer cost of validating data mixing strategies at this scale, there is a paucity of open research on pretraining data for LLMs.

Ideally one would conduct data mix experiments at smaller scales to identify what is a good data mix. However, this is often ineffective because large FLOP scales are required to reveal significant changes in difficult and emergent benchmarks. In fact, most LLMs trained at smaller scales register random accuracy on many important benchmarks such as MMLU (Wei et al., 2022). As a result, experiments at smaller scales can often be misleading; the variation between different data mixes on important benchmarks is often due to noise rather then dataset quality at this scale. On the other hand, it is prohibitively expensive and impractical to exhaustively characterize datasets by doing multiple training runs at the scale needed to measure above random performance on these metrics.

In this work, our goal is to characterize the utility of an alternative approach to conduct pretraining data experiments at a reasonable scale. Our strategy is to modify the data mixture at the end of training after we have already trained for enough FLOPs to measure meaningful signal on difficult benchmarks. We show that this is an effective strategy for improving LLM pretraining data mixes with experiments that are an order of magnitude cheaper than full training runs.

Contributions:

We begin with a baseline mix of publicly available datasets that achieves the same scaling of performance with FLOPs as the Llama-2 model family for a 7B model trained for 1 trillion tokens.

We introduce domain upsampling—a data intervention which upsamples domain specific datasets relative to Common Crawl at the end of training—and demonstrate that it can boost challenging metrics. In particular, we observe improvements of up to 6.90 pp on MMLU, 8.26 pp on GSM8K, and 6.17 pp on HumanEval relative to the base data mix in our training setup. This makes our performance comparable to Llama-2 (7B) but at approximately half the training FLOPs.

We ablate the percentage of training that utilizes domain upsampling and show 10%-20% is optimal for navigating the tradeoff between general language modeling capabilities and targeted benchmarks.

We show how domain upsampling can be used as a FLOP-efficient tool to characterize how individual datasets impact model capabilities. By removing a subset of math-heavy pretraining data from the datasets we upsampled at the end of training, we quantified the impact these datasets have on specific benchmarks.

Training Details

We studied domain upsampling on 7 billion parameter models trained for 1 trillion tokens. This FLOP scale was chosen so that the model performed above the noise floor on key metrics like MMLU enabling us to see the effects of data interventions on the model.

The 7B models trained for this work are decoder-only transformers using the MPT architecture in LLM Foundry (MosaicML et al., 2023). To evaluate our models we use the latest version of the Eval Gauntlet v0.3.(MosaicML et al., 2023), an evaluation framework consisting of 35 popular in context learning evaluation tasks used to evaluate LLM base models. The Gauntlet v0.3 aggregates scores on benchmarks across 6 categories. It is described in Appendix A. We use an inverse square root learning schedule similar to (Zhai et al., 2022).

Results

Here we present the experiments demonstrating the performance boost achieved by domain upsampling as well as its utility in characterizing how datasets affect challenging, emergent metrics.

To construct a baseline data mix, we grouped a set of publicly-available datasets into 4 broad categories:

Large-Scale Common Crawl: Datasets derived from Common Crawl that emphasize scale. These datasets trade off thorough quality filtering in favor of curating a large and diverse set of tokens.

Small-Scale Common Crawl: Datasets derived from Common Crawl with more extensive filtering but are smaller than large-scale Common Crawl.

Domain Specific data: Small datasets that target certain domains or are from individual sources and are of high quality (e.g. Wikipedia).

Code: Code data across a variety of programming languages.

We set the proportions for mixing these datasets based on a rough heuristic for the number of epochs each of these groups would be seen during the 1 trillion token training duration. Specifically, we choose 0.5 epochs for the Small-Scale Common Crawl and Domain Specific data and 1 epoch for Code. The remainder of the 1 trillion tokens are filled with Large-Scale Common Crawl. The exact proportions are in Table 2.

The rationale behind choosing these proportions is as follows: we expect the Small-Scale Common Crawl and Domain Specific data to be of high quality and we wanted them to be well represented on our 1 trillion token budget. Also, we wanted to emphasize coding ability and so we decided to sample code data at a high percentage—initial experiments indicated that a high percentage of code around 20% boosted programming and reasoning ability without negatively impacting language abilities. We then treat the Large-Scale CC as filler tokens that increase the diversity of our dataset and allow us to fill our token budget.

Importantly, since the goal of our experimental setup is to demonstrate the utility of domain upsampling at the end of training (discussed in section 3.2), we opt for choosing a reasonable heuristic for picking our initial data mix proportions without too much optimization. Table 3 and Figure 1 show the performance of this initial pretraining data mix for two 7B models trained for 0.5T and 1T tokens. This heuristic has indeed been validated by our empirical results; plotting error vs. FLOPs shows that our models lie on or below the Llama-2 scaling line on the Gauntlet v0.3 Core Average, MMLU, GSM8K, and HumanEval. Interestingly, though the overall performance scaling (as measured by Gaunlet v0.3 Core Average) is very similar, our particular data choices and mixing coefficients have led to slightly different tradeoffs. The model trained for 1T tokens outperforms the Llama-2 7B model trained for 2T tokens on GSM8K and HumanEval. This indicates that our models have better mathematical and programming ability despite being trained for half the number of tokens. We also provide a comparison to OpenLlama 7Bv2 (Geng & Liu, 2023), a 7B model that provides some open details about their data mix.

2 Domain upsampling significantly boosts performance on challenging metrics

Next, we introduce domain upsampling during the last 20% of training for our 1T token training run. For this, we start with a checkpoint at 0.8T tokens of training, change the mixing proportions of our pretraining data mix, and continue training for the remaining 0.2T tokens. The exact mixing proportions of our domain upsampled pretraining mix are in Table 4. These percentages were chosen based on the following heuristic: we hypothesize that though the Large-Scale CC adds a lot of diversity to the pretraining data mix, it is advantageous to emphasize Domain Specific data at the end of training to bias our model towards token distributions that have high information density in domains we care about. Thus, we remove Large-Scale Common Crawl from our data mix while upsampling both Domain Specific and Code subsets. We also maintain Small-Scale Common Crawl at high percentage to prevent a large distribution shift in our pretraining data.

The results of this end-of-training data intervention are shown in Table 5 and Figure 2. Domain upsampling was incredibly effective in boosting model performance relative to the initial pretraining data mix on all challenging benchmarks. Given the large amount of code and math related data in the domain upsampled data mix, it is perhaps unsurprising that this intervention led to GSM8K and HumanEval scores that are approximately 10pp higher than Llama-2 (7B) despite the model being trained for half the total number of tokens. Additionally, this did not come at a cost to general language modeling capabilities; it led to an overall model performance improvement as measured by Gauntlet v0.3 Core Average. In fact, it improved world knowledge—as measured by MMLU and the Gauntlet v0.3 subset—relative to the base data mix, bringing us closer to Llama-2 (7B) performance on these metrics. There was only a small 1pp tradeoff in the Language Understanding subset.

Overall, this across the board improvement on challenging benchmarks establishes the efficacy of domain upsampling as a pretraining data intervention for improving model performance. Importantly, even using simple heuristics for choosing the new data mix proportions has strong positive effects, leaving opportunity for further improvement with better tuned mixing proportions.

The success of domain upsampling for the last 20% of training raises the question: are the improvements from an end-of-training data intervention or are they from overall better data mix proportions? Phrased another way, are the data mix proportions in Table 4 better than our initial data mix and would training a model for 1T tokens with this data mix lead to better performance? In this section, we provide evidence that this is not the case and in fact, treating domain upsampling as an end-of-training data intervention helps us better tradeoff domain specific improvements and general language modeling capabilities.

To identify when in training this intervention should be applied, we ablate our previous experiment by performing domain upsampling for the last 5%, 10%, 20%, and 30% of training. The results of this experiment are shown in Figure 3 and Table 6. Note, while the math and programming related benchmarks, such as HumanEval, GSM8K and related Gauntlet v0.3 subscores, continue to improve as we increase the fraction of training that uses domain upsampling, other benchmarks reach optimal performance at 20% or less. For example, MMLU peaks at 10% and Gauntlet v0.3 Core Average peaks at 20%. Thus, as we increase the fraction of training with domain upsampling beyond 20%, improvements on math and coding benchmarks come at the cost of performance on general language modeling abilities.

This apparent trade-off indicates that the domain upsampling data mix proportions are not incontrovertibly better than the initial data mix, and training with it for the full 1T token duration would not lead to a better general purpose language model. We do not rule out that there is an alternate mix that achieves similar performance as the 20% domain upsampling experiment when trained for the full training duration. However, finding such a mix is expensive to iterate on for the full training run. Thus, the strength of domain upsampling is that it gives us a tool to navigate this tradeoff between targeted domains and general language modeling abilities with experiments that are an order of magnitude cheaper.

4 Domain upsampling is a FLOP-efficient tool to characterize how individual datasets impact model capabilities

Having observed that upsampling code and our domain specific datasets for a small percentage of training leads to significant improvements on difficult and emergent tasks, we explore the question: how does one attribute improvements to specific subsets of these data? Notably, as can be seen in Figure 3 and Table 5, GSM8K scores—a task measuring math and reasoning abilities—improves monotonically as duration of domain upsampling is increased. We hypothesize, given the quantity of math related data in our high-quality datasets that these may be responsible for some or all of this improvement. To quantify the impact of these datasets we repeat our experiment, applying domain upsampling for the last 10% of the training duration. We keep our dataset proportions identical to those in Table 4, but remove the math related subsets. We present the results in Table 7.

We observe that not only do the the mathematical knowledge and reasoning skills, as measured by MMLU (which contains STEM subsets) & GSM8k, not reach the same level of performance as the model trained using domain upsampling that included them, but in fact performance is worse then the baseline model with no domain upsampling. Moreover, every Gauntlet v0.3 subcategory score for the domain upsampling sans-math with the exception of programming is lower than the baseline model. From this we can draw the conclusion that these specific datasets are responsible for the majority of the mathematical knowledge and reasoning capabilities in both the base model and the domain upsampled variant.

With this observation we have successfully done something which generally would be considerably more expensive. That is, we have measured the impact of pretraining datasets at a scale where difficult and emergent tasks can be reliably measured, but at an order of magnitude fewer training FLOPS. We believe application of domain upsampling opens up the ability for researchers to experiment with their pretraining datasets in a tractable way as compared to full pretraining runs.

Discussion

Pretraining LLMs has become an increasingly costly and clandestine endeavor given the scale of compute required for each experiment. This problem is exacerbated by the multi-faceted decision space presented to practitioners, especially in the selection of pretraining data. Since many important model capabilities emerge with scale, trying to explore this design space at small compute budgets is often ineffective: observations made about the effects of the pretraining data mix typically do not transfer to larger models or training budgets.

In this work, we consider a baseline data mix of publicly available datasets that achieves or exceeds the scaling of the Llama-2 family of models on key benchmarks. Next, we take a crucial first step towards making experimentation with pretraining datasets cheaper. We introduce domain upsampling, a method that can strongly impact the performance of the model by making targeted changes to the data mix at the end of training. This enables us to achieve the performance of Llama-2 (7B) but with half the training budget. By varying the duration of domain upsampling, we demonstrate how to navigate the tradeoff between targeting specific domains and making general purpose language models.

Finally, we show how making changes to the data mix only during the domain upsampling period enabled us to cheaply characterize the impact of several math-focused datasets, and we see many opportunities to use this method as a general tool for studying pretraining data in a FLOP-efficient manner. It also creates a platform to test data interventions at scale: instead of testing possible dataset optimization algorithms at small scales and hoping they will generalize, we can test them at the end of training to effectively measure their impact at scale. By bringing down the cost of experimentation we have made pretraining data experiments more accessible, and we will release our models and intermediate checkpoints as research artifacts to the community as a resource to unlock further insights into LLM pretraining data.

References

Appendix A Gauntlet v0.3

The Gauntlet v0.3 is a aggregation of benchmark developed by Mosaic Research. Rather than reporting a monolithic metric in which all scores are aggregated together, the individual benchmarks were grouped into six broad competencies corresponding to different capabilities we want our LLMs to have:

World Knowledge: Measures the model’s factual knowledge across a range of subjects.

Commonsense Reasoning: Evaluates the model’s ability to do basic reasoning tasks that require commonsense knowledge of objects, their properties, and their behaviors.

Language Understanding: Assesses the model’s ability to understand structure and properties of language.

Symbolic Problem Solving: Tests the model’s ability to solve a diverse range of symbolic tasks including arithmetic, logical reasoning, algorithms, and algebra.

Reading Comprehension: Measures a model’s ability to answer questions based on information in a passage of text.

Programming: Quantifies the ability to generate code from docstring descriptions.

These divisions allow for more fine-grained comparison between models and is especially useful for understanding how datasets affect different capabilities of the model. The random baseline of each metric was subtracted out before aggregating. For example, if the metric is 4-option multiple choice questions giving a random baseline of 25% and the model achieves 30% then this would be aggregated as $(0.3-0.25)/(1-0.25)=0.0667$, essentially rescaling accuracy above change to be between 0 and 1. If the random baseline is approximately 0, then the metric is reported as is. Table 8 list the benchmarks in each category.