Source-linked AI summary

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

Cody Blakeney, Mansheej Paul, Brett W. Larsen, Sean Owen, Jonathan Frankle

arXiv:2406.03476v1cs.LGcs.CL

TL;DR

The paper addresses the costly problem of identifying useful balances between diverse CommonCrawl data and smaller, information-dense domain datasets. It evaluates end-of-training domain upsampling as a cheaper data intervention, finding strong gains on difficult benchmarks and a 10%-20% intervention duration that balances targeted and general capabilities.

  • Problem

    Large-scale training is needed to reveal meaningful effects of pretraining data on difficult benchmarks, but exhaustive data-mix experimentation is expensive and small-scale results can be misleading.

    Source: §1 · p. 2
  • Method

    The paper upsamples domain-specific datasets relative to CommonCrawl during the final training phase and varies intervention duration to study benchmark effects and tradeoffs.

    Source: §1 · p. 2, §3.2 · p. 4, §3.3 · p. 6
  • Results

    Up to 6.90 pp on MMLU, 8.26 pp on GSM8K, and 6.17 pp on HumanEval are gained over the base mix, with performance comparable to Llama-2 (7B) at approximately half the training FLOPs.

    Source: §1 · p. 2
  • Takeaways & Limitations

    End-of-training domain upsampling provides a FLOP-efficient tool for improving targeted capabilities and characterizing the utility of individual pretraining datasets at scale.

    Source: §1 · p. 2, §4 · p. 8

    The domain-upsampled mixture is not necessarily better when trained for the full duration, and an alternative full-run mixture might achieve similar results but would be expensive to find.

    Source: §3.3 · p. 7

Abstract

from arXiv · show

Pretraining datasets for large language models (LLMs) have grown to trillions of tokens composed of large amounts of CommonCrawl (CC) web scrape along with smaller, domain-specific datasets. It is expensive to understand the impact of these domain-specific datasets on model capabilities as training at large FLOP scales is required to reveal significant changes to difficult and emergent benchmarks. Given the increasing cost of experimenting with pretraining data, how does one determine the optimal balance between the diversity in general web scrapes and the information density of domain specific data? In this work, we show how to leverage the smaller domain specific datasets by upsampling them relative to CC at the end of training to drive performance improvements on difficult benchmarks. This simple technique allows us to improve up to 6.90 pp on MMLU, 8.26 pp on GSM8K, and 6.17 pp on HumanEval relative to the base data mix for a 7B model trained for 1 trillion (T) tokens, thus rivaling Llama-2 (7B)$\unicode{x2014}$a model trained for twice as long. We experiment with ablating the duration of domain upsampling from 5% to 30% of training and find that 10% to 20% percent is optimal for navigating the tradeoff between general language modeling capabilities and targeted benchmarks. We also use domain upsampling to characterize at scale the utility of individual datasets for improving various benchmarks by removing them during this final phase of training. This tool opens up the ability to experiment with the impact of different pretraining datasets at scale, but at an order of magnitude lower cost compared to full pretraining runs.

1. Introduction

LLM pretraining mixes must balance diverse but low-density CommonCrawl data against smaller, information-dense domain datasets, yet large-scale experiments are costly and small-scale benchmark signals can be misleading. The paper introduces end-of-training domain upsampling as a cheaper way to improve targeted capabilities and study dataset utility.

1. Introduction

  • CommonCrawl supplies diverse but less information-dense data, whereas smaller domain-specific datasets are more carefully processed and denser but less diverse.This contrast motivates changing the mixture near the end of training.Source: §1 · p. 1
  • Small-scale data-mix experiments can be misleading because difficult benchmarks may remain near random accuracy, while full-scale comparisons are prohibitively expensive.The paper frames this cost-signal problem as a barrier to open research on pretraining data.Source: §1 · p. 2
  • End-of-training data interventions enable pretraining-mix experiments at an order of magnitude lower cost than full training runs.The approach is intended to measure meaningful benchmark effects after sufficient training compute has already been used.Source: §1 · p. 2
  • Up to 6.90 pp on MMLU, 8.26 pp on GSM8K, and 6.17 pp on HumanEval are achieved over the base data mix, making performance comparable to Llama-2 (7B) at approximately half the training FLOPs.The comparison applies to the paper’s 7B, 1T-token training setup.Source: §1 · p. 2
  • Domain upsampling increases domain-specific data relative to CommonCrawl during the final training phase to improve challenging benchmark performance.The intervention targets the final portion of training rather than replacing the full baseline mixture.Source: §1 · p. 2
  • 10%-20% of training with domain upsampling is reported as optimal for balancing general language modeling capabilities against targeted benchmark performance.The paper varies the duration of the intervention to study this tradeoff.Source: §1 · p. 2
  • The method also characterizes individual dataset utility by removing selected math-heavy datasets during the final upsampling phase.This provides a FLOP-efficient way to connect datasets with benchmark-specific capability changes.Source: §1 · p. 2

2. Training Details

The experiments use 7B models trained for 1T tokens so key benchmarks rise above the noise floor, with decoder-only MPT models evaluated using the Eval Gauntlet. The initial data mix reaches or exceeds the Llama-2 family’s benchmark scaling relationship.

2. Training Details

  • 7B models are trained for 1T tokens so MMLU and other key metrics provide signal above the noise floor for evaluating data interventions.The study uses this scale to expose meaningful effects of pretraining-mixture changes.Source: §2 · p. 2
  • The models are decoder-only transformers using the MPT architecture, and evaluation uses Eval Gauntlet v0.3 across 35 in-context learning tasks.The Gauntlet aggregates benchmarks across six categories.Source: §2 · p. 2
  • The baseline 7B models trained for 0.5T and 1T tokens lie on or below the Llama-2 scaling line on key benchmarks.The scaling relationship is estimated from Llama-2 models by regressing log error against log training FLOPs.Source: Figure 1 · p. 3

3. Results

Across four experiments, the authors establish a baseline data mix, show that end-of-training domain upsampling improves difficult benchmarks, identify a 10%-20% duration trade-off, and use the intervention to assess individual datasets at lower cost.

3.1 Baseline data mix achieves Llama-2 scaling

  • The baseline mix combines large- and small-scale Common Crawl, domain-specific data, and code, using heuristic proportions for a 1T-token training run.Large-Scale Common Crawl fills the remaining budget, while code receives an elevated proportion to emphasize programming ability.Source: §3.1 · p. 3, Table 2 · p. 4
  • The 1T-token baseline 7B model outperforms Llama-2 7B trained for 2T tokens on GSM8K, HumanEval, and selected Gauntlet v0.3 subsets.The comparison indicates stronger mathematical and programming ability despite half the training tokens, while overall scaling remains similar.Source: §3.1 · pp. 3–4, Table 3 · p. 4

3.2 Domain upsampling significantly boosts performance on challenging metrics

  • 6.90 pp on MMLU, 8.26 pp on GSM8K, 6.17 pp on HumanEval, and 3.95 pp on Gauntlet v0.3 improve with domain upsampling at no additional FLOP cost.The intervention applies during the final 20% or 200B tokens and improves overall Gauntlet v0.3 performance, with only a small 1pp Language Understanding trade-off.Source: Figure 2 · p. 5, §3.2 · p. 5, Table 5 · p. 5

3.3 Changing the duration of domain upsampling enables us to navigate the trade-off between targeting specific domains and general purpose language models

  • 10%-20% final-phase domain upsampling provides the best trade-off, because GSM8K and HumanEval keep improving while MMLU peaks at 10% and Gauntlet Core Average at 20%.Beyond 20%, gains on math and coding benchmarks come at the cost of general language modeling performance, so the upsampled mix is not preferred for the full run.Source: Figure 3 · p. 6, §3.3 · pp. 6–7

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

  • Removing math-specific datasets during the final 10% causes mathematical knowledge and reasoning performance to fall below the domain-upsampled model and even below the baseline.All Gauntlet v0.3 subcategories except programming also decline, demonstrating a cheaper way to characterize dataset contributions at benchmark-relevant scale.Source: §3.4 · pp. 7–8, Table 7 · p. 7

4. Discussion

The paper introduces domain upsampling as a cheaper way to test pretraining-data interventions at scale, where small-budget experiments often fail to reveal transferable capability changes. It uses end-of-training data changes to improve targeted performance, navigate domain-versus-general tradeoffs, and assess individual datasets.

4. Discussion

  • Small-compute data-mix observations often fail to transfer because important model capabilities emerge only at larger scales.This makes conventional pretraining-data exploration costly and potentially misleading.Source: §4 · p. 8
  • Domain upsampling makes targeted changes to the data mix at the end of training and achieves Llama-2 (7B) performance with half the training budget.Source: §4 · p. 8
  • Varying the duration of domain upsampling exposes the tradeoff between targeting specific domains and preserving general-purpose language modeling.The approach is presented as a way to study this tradeoff rather than selecting a universally optimal mixture.Source: §4 · p. 8
  • End-of-training data changes can cheaply characterize the impact of individual math-focused datasets and support FLOP-efficient pretraining-data studies.The authors also describe the method as a platform for testing data interventions at scale and release models and intermediate checkpoints as research artifacts.Source: §4 · p. 8

A. Gauntlet v0.3

Gauntlet v0.3 organizes evaluation into six capability categories rather than one monolithic score. Its aggregation rescales metrics above their random baselines, enabling finer-grained comparisons of model capabilities and dataset effects.

A. Gauntlet v0.3

  • Gauntlet v0.3 groups benchmarks into six competencies: world knowledge, commonsense reasoning, language understanding, symbolic problem solving, reading comprehension, and programming.The categories span factual knowledge, commonsense and symbolic reasoning, language, passage-based answering, and code generation.Source: A. Gauntlet v0.3 · p. 12
  • The programming category quantifies the ability to generate code from docstring descriptions.This category is distinct from the knowledge, reasoning, language, and reading-comprehension groups.Source: A. Gauntlet v0.3 · p. 12
  • Scores are aggregated after subtracting each metric’s random baseline and rescaling accuracy above chance to a 0–1 range.For a four-option task scoring 30%, the normalized value is (0.3 − 0.25)/(1 − 0.25) = 0.0667.Source: A. Gauntlet v0.3 · p. 12
  • The Gauntlet v0.3 metrics are 0-shot unless otherwise denoted.Table 8 lists the benchmarks included in each category.Source: A. Gauntlet v0.3 · p. 12, Table 8 · p. 13
Loading 2406.03476v1…