PaperlayerView exact revision on arXiv

Source-linked AI summary

MATES: Model-Aware Data Selection for Efficient Pretraining with Data Influence Models

Zichun Yu, Spandan Das, Chenyan Xiong

arXiv:2406.06046v2cs.CLcs.LG

TL;DR

Static data-selection methods may not track how pretraining models’ data preferences change, despite the potential of selection to improve efficiency. MATES locally probes data influence and trains an influence model to guide stage-wise selection; experiments report better downstream accuracy than random selection and reduced compute to reach target performance.

  • Problem

    Existing selection methods rely on static heuristics or larger reference models and may not capture pretraining models’ evolving data preferences.

    Source: §1 · p. 1

  • Method

    MATES locally probes oracle data influence with the pretraining model, trains a small influence model to approximate it, and selects data for the next pretraining stage.

    Source: §1 · p. 2

  • Results

    1.3% (410M) and 1.1% (1B) average zero-shot accuracy gains over random selection were reported across downstream tasks, with MATES doubling gains from larger-reference-model approaches and reducing FLOPs to target performance by more than half.

    Source: §1 · p. 2

  • Takeaways & Limitations

    MATES outperforms random, rule-based, influence-function-based, and LLM-based data selection methods in the reported pretraining experiments.

    Source: §1 · p. 2, §7 · p. 10

    Experiments use 410M- and 1B-parameter models, so how well the methods scale to production models with billions of parameters and trillions of tokens remains unclear.

    Source: §6 · p. 10

Abstract

from arXiv · show

Pretraining data selection has the potential to improve language model pretraining efficiency by utilizing higher-quality data from massive web data corpora. Current data selection methods, which rely on either hand-crafted rules or larger reference models, are conducted statically and do not capture the evolving data preferences during pretraining. In this paper, we introduce model-aware data selection with data influence models (MATES), where a data influence model continuously adapts to the evolving data preferences of the pretraining model and then selects the data most effective for the current pretraining progress. Specifically, we collect oracle data influence by locally probing the pretraining model and fine-tune a small data influence model to approximate it accurately. The data influence model then predicts data influence over the whole pretraining corpus and selects the most influential data for the next pretraining stage. Experiments of pretraining 410M and 1B models on the C4 dataset demonstrate that MATES significantly outperforms random data selection on extensive downstream tasks. It doubles the gains achieved by the state-of-the-art data selection approach that leverages larger reference models and reduces the total FLOPs required to reach certain performances by half. Further analyses validate the effectiveness of the locally probed oracle data influence and the approximation with data influence models. Our code is open-sourced at https://github.com/cxcscmu/MATES.

1 Introduction

Pretraining data selection can improve efficiency, but static selection methods may miss the changing preferences of a model as it trains. MATES adapts data selection to those preferences, and experiments report gains over random selection and larger-reference-model approaches.

1 Introduction

  • Compute is often the most restrictive scaling resource, while web data sources are much larger than available compute and vary in quality.The introduction frames data selection as an opportunity to improve pretraining scaling efficiency.Source: §1 · p. 1
  • Static heuristics and larger reference models are common selection signals, but existing methods often overlook evolving preferences during pretraining.The paper notes that these methods can be limited in real-world pretraining scenarios.Source: §1 · p. 1
  • 1.3% (410M) and 1.1% (1B) average zero-shot accuracy gains over random selection were reported across downstream tasks.MATES also doubles the gains of selection approaches relying on larger reference models.Source: §1 · p. 2
  • MATES uses a small data influence model to track the pretraining model’s changing preferences and select training data for the current stage.The framework locally probes oracle data influence and trains the influence model to approximate it.Source: §1 · p. 2

2 Related work

Prior pretraining data selection methods use manual rules, deduplication, proximity to high-quality corpora, or larger language models. Other approaches estimate influence or dynamically adjust data signals, but influence calculations can be costly to track as models change.

2 Related work

  • Rule-based curation uses hand-designed criteria, while deduplication filters exact, fuzzy, or semantic duplicates and is often combined with quality-oriented selection.Manual rules can be effective as an initial curation step but may miss nuances of model preferences.Source: §2 · pp. 2–3
  • Proximity methods select data using n-gram similarity or language-model perplexity, but may reduce diversity and general capability as corpora grow.The effectiveness of proximity-based methods at larger data scale is described as unclear.Source: §2 · p. 3
  • LLM-based methods use stronger reference models to annotate or rephrase web data, including to balance quality and diversity.The cited methods rely on reference LLMs often several orders of magnitude larger than the models they guide.Source: §2 · p. 3
  • Influence functions estimate individual data points’ impact on performance, but costly gradient calculations make them difficult to scale for tracking changing preferences.TRAK makes influence computation feasible for pretraining experiments, yet its cost remains prohibitive for on-the-fly model-aware selection.Source: §2 · p. 3
  • Curriculum methods dynamically change pretraining signals, using implicit or explicit strategies such as gradient norms, certainty, and expertise.These approaches adapt data signals to models’ changing preferences.Source: §2 · p. 3

3 Methods

MATES selects training data to reduce loss on reference data after the model is updated, while approximating pointwise influence to avoid exhaustive batch search. It locally probes oracle influence and uses an influence model to support selection across the corpus as the model changes.

3 Methods

  • MATES locally probes oracle data influence, then trains a small influence model to approximate it and select data for the next pretraining stage.The pipeline begins with a randomly sampled subset for warm-up before iterating selection and training.Source: §3 · pp. 3–4

3.1 Model-aware data selection framework

  • The framework selects a batch from the pretraining data to minimize reference-data loss after optimizing the current model on that batch.The reference data are non-evaluation data and may represent one or more downstream tasks.Source: §3.1 · p. 4
  • Pointwise data influence approximates batch influence by summing individual examples’ influences, assuming each example contributes independently.The oracle pointwise influence depends on the model state and changes during pretraining.Source: §3.1 · p. 4
  • Data influence parameterization avoids computing oracle influence over millions of examples by probing a small hold-out set and fine-tuning an influence model.The model predicts influence over the training set, with asynchronous updates every U steps and Gumbel-Top-k sampling weighted by influence scores.Source: §3.1 · pp. 4–5

3.2 Locally probed oracle data influence

  • The locally probed oracle estimates an example’s effect by one-step training from the current model and comparing reference loss before and after the update.The negative of the loss difference is defined as influence so that a positive score reflects a beneficial impact on performance.Source: §3.2 · p. 5

4 Experimental methodologies

Experiments use Pythia models pretrained on C4, with a continuously fine-tuned BERT-base data influence model guiding staged selection. They evaluate downstream accuracy, influence-model prediction quality, and total compute including selection.

4 Experimental methodologies

  • MATES pretrains 410M/1B Pythia models from scratch on C4 and continuously fine-tunes a smaller BERT-base data influence model.The smaller influence model is intended to keep data selection efficient.Source: §4 · p. 6
  • At each 10k-step stage, MATES samples 20% of the data using influence scores as weights, with temperature τ=1.0 and one-epoch updates.LAMBADA is the reference data in the main experiment.Source: §4 · p. 6
  • Evaluation covers nine downstream tasks and reports zero-shot accuracy, normalized zero-shot accuracy where applicable, and total GPU FLOPs including data selection.BoolQ and WinoGrande use zero-shot accuracy; the other tasks use normalized zero-shot accuracy.Source: §4 · p. 6
  • Influence-model quality is measured by validation Spearman correlation against oracle influence on a 10% hold-out set from the collected data-oracle mapping.Source: §4 · p. 6
  • Baselines span random selection, heuristic filtering, deduplication, gradient-based influence, proxy-model influence, and LLM-rated quality; LESS is reported only for 410M due to its gradient cost.Source: §4 · p. 6
  • MATES data selection accounts for 21.7% of total FLOPs for 410M models and 11.5% for 1B models.Source: Figure 3 · p. 7

5 Evaluation results

MATES improves downstream performance over random and reference-model selection while keeping selection costs comparatively low. The analyses examine reference-task choice, influence-model adaptation, and changing data preferences across training stages.

5.1 Overall performance

  • MATES gains average downstream accuracy of 1.3% for 410M and 1.1% for 1B over random selection, nearly doubling QuRating’s gains; 1B improves on random across all nine tasks.In 410M, gains are nearly 2.0% on most tasks, except ARC-C and LogiQA.Source: §5.1 · p. 7
  • MATES selection costs are relatively smaller for larger models; collecting one oracle score takes about 2.5 seconds on one GPU, and 160k scores take about 14 hours on eight GPUs.The paper contrasts the 14-hour collection estimate with four days of pretraining.Source: §5.1 · p. 7, Table 2 · p. 7

5.2 Effectiveness of locally probed oracle data influence

  • Oracle influence distributions remain spread across reference tasks; LAMBADA has over 20% more positively influential data, while ARC-E language modeling yields more positive examples than its multiple-choice format.The authors hypothesize that LAMBADA aligns more closely with the pretraining objective because it is a word-prediction task.Source: Figure 4 · p. 8, §5.2 · p. 8

5.1 Overall performance

  • At 1B, MATES reaches random-selection performance with 43.3% of total FLOPs, reducing the FLOPs and steps needed for a target performance by more than half.At this scale, pretraining accounts for 88.5% of FLOPs and data selection for 11.5%.Source: §5.1 · p. 8

5.2 Effectiveness of locally probed oracle data influence

  • ARC-E reference data benefits in-domain accuracy but generalizes worse than LAMBADA, while FLAN benefits a wider range of tasks and has similar average accuracy to LAMBADA.The oracle is evaluated by sampling data weighted by oracle scores in one short decay stage, since scoring the full corpus is infeasible.Source: §5.2 · pp. 8–9, Table 3 · p. 8

5.3 Effectiveness of data influence model

  • The influence model approximates oracle influence across reference tasks, and higher oracle-distribution standard deviation tends to accompany higher validation Spearman correlation.The authors suggest greater influence-score variability supplies more diverse signals for approximation.Source: §5.3 · p. 9
  • Dynamic influence models capture changing oracle preferences more precisely, reaching around 0.7 validation Spearman correlation at 40k steps while static models remain below 0.5.Static-model-selected data causes a notable downstream performance drop, especially early in pretraining.Source: Figure 5 · p. 9, §5.3 · p. 9

5.4 Case study

  • Examples shift from natural narratives at 10k steps to factual knowledge at 20k, academic text at 30k, and potentially long-tail knowledge at 40k.The authors note that these examples may not cover every aspect of selected data.Source: §5.4 · p. 9

6 Discussion and limitations

The discussion identifies two boundaries: influence is measured under an independence assumption that omits interactions among data points, and experiments remain at moderate model scale.

Combinational Measurement of Data Influence

  • The influence measurement assumes each data point contributes independently, although pretraining combines batched data effects and advanced capabilities can accumulate.The authors identify measuring these combinational and accumulative effects as a direction for better understanding and leveraging data value.Source: §6 · p. 9
  • The 1B case study ranks excerpts from 80k randomly sampled data points across checkpoints, with lower influence rank indicating higher influence.The caption describes the table's scope but does not state a comparative outcome.Source: Table 4 · p. 10

Exploratory Scale

  • Experiments cover 410M- and 1B-parameter models, leaving performance at production scale with billions of parameters and trillions of tokens unclear.The authors note that larger-scale pretraining may also bring stability issues requiring dedicated work.Source: §6 · p. 10

7 Conclusion

MATES adapts data selection to changing pretraining preferences using locally probed oracle influence and a fitted influence model. The paper reports gains over several selection baselines and concludes that model-aware curation has potential to improve pretraining scaling.

7 Conclusion

  • MATES tracks evolving pretraining preferences with a data influence model and selects data expected to be effective for the current stage.It locally probes oracle influence on a reference task and fits the influence model to those probes.Source: §7 · p. 10
  • MATES surpasses random, rule-based, influence-function-based, and LLM-based selection methods in pretraining experiments.Further analyses support the effectiveness of locally probed oracle influence and its approximation by data influence models.Source: §7 · p. 10
  • The paper presents model-aware data curation as a potential route to improving foundation-model scaling through better data curation.The authors frame this as a motivation for further exploration.Source: §7 · p. 10

A Full derivation of locally probed oracle data influence

The derivation defines oracle influence by comparing reference loss before and after upweighting one training example. It uses the resulting optimal model state to form an influence estimate that remains usable away from convergence with damping.

A Full derivation of locally probed oracle data influence

  • The derivation defines the optimal model state after one training example is upweighted by ε, relative to the unweighted optimum M*.The upweighted state minimizes the training objective with an added ε-weighted loss for x_i.Source: A Full derivation of locally probed… · p. 15
  • For a non-converged model state, damping λ makes H_M + λI positive definite so the locally probed influence calculation remains meaningful.The derivation also uses a first-order Taylor expansion to approximate parameter change.Source: A Full derivation of locally probed… · p. 15
  • Oracle influence is the negative change in reference loss between the current model and the model after upweighting x_i, so positive scores indicate beneficial impact.The influence quantity follows the reference-loss-change definition described alongside the upweighting setup.Source: A Full derivation of locally probed… · p. 15

B Experimental details

The appendix outlines experimental configurations, the WSD learning-rate schedule, and the BERT-based influence model design, followed by a section of additional experiments.

B Experimental details

  • The experimental-details appendix is organized around configurations, the WSD scheduler, and the data influence model design.It also introduces additional experiments covering attribution comparisons, influence-model ablations, and oracle-influence analysis.Source: B Experimental details · p. 16, C Additional experiments · p. 17

B.1 Experimental configurations

  • Experimental configurations are tabulated in Table 5, whose caption identifies it as a configuration table.The supplied caption does not state configuration values.Source: §B.1 · p. 16

B.2 WSD scheduler

  • WSD uses warmup, stable, and decay stages; its learning rate stays constant during the stable stage, offering flexibility across training lengths.The main experiments use W=2000, S=50000, and D=200, and evaluate checkpoints after a short decay stage.Source: §B.2 · p. 16

B.3 Design of data influence model

  • The BERT-based influence model averages final-layer hidden representations into a sequence vector and trains its prediction with mean squared error against normalized oracle influence.For 1024-token pretraining sequences, it processes two chunks separately and averages their representations; the authors say this design can extend to longer sequences.Source: §B.3 · pp. 16–17

C Additional experiments

  • The additional-experiments section covers comparisons of attribution methods, influence-model ablations, and analysis of oracle influence.Source: C Additional experiments · p. 17

C.1 Comparison of different data influence attribution methods

The locally probed oracle outperforms DsDm in the short decay-stage comparison. The authors attribute the advantage to using the current model state and one-step training rather than an existing checkpoint and multiple approximations.

C.1 Comparison of different data influence attribution methods

  • 0.8 zero-shot accuracy gain separates the locally probed oracle from DsDm in the 410M short-decay comparison.The decay stage consists of 200 steps.Source: Table 6 · p. 17, §C.1 · p. 17
  • The oracle uses the current model state, whereas DsDm relies on an existing checkpoint and Taylor approximation with gradient dimension reduction.The authors connect the gap to tracking the current model and avoiding precision loss from multiple approximation manipulations.Source: §C.1 · p. 17
  • Collecting optimal influence scores across the full dataset remains costly, but local probing beats Taylor approximation with fewer FLOPs.Source: §C.1 · p. 17

C.2 Ablation study on data influence models

The ablations examine selection hyperparameters and the data influence model, then analyze oracle stability and reference-task alignment. Results highlight quality–diversity trade-offs and variation in data requirements and task correlations.

C.3 Analysis of oracle data influence

  • Oracle influence is not sensitive to batch size when scores correspond to the same model state.For batch size 16, the analysis uses LASSO regression to separate individual-data influence from batch-level influence.Source: Figure 8 · p. 18, §C.3 · pp. 18–19
  • LAMBADA accuracy correlates with average downstream accuracy at 0.877, whereas ARC-E and PIQA correlations are not high.The authors suggest ARC-E and PIQA may be less accurate reference tasks for target performance.Source: Figure 8 · p. 18, §C.3 · p. 19

C.2 Ablation study on data influence models

  • Changing the update step from 10k to 5k or 2.5k causes little fluctuation, but temperatures of 2.0 and 0.0 reduce performance.The authors attribute the temperature results to insufficient quality emphasis at 2.0 and reduced diversity at 0.0.Source: Table 7 · p. 17, §C.2 · p. 18
  • MATES consistently outperforms random selection across ratios from 1/200 to 1/2, while ratios that are too low or high underperform.The authors associate low ratios with reduced diversity and high ratios with weaker use of influence scores.Source: §C.2 · p. 18, Figure 6 · p. 18
  • Continuous fine-tuning from the 30k-step checkpoint needs less than half the oracle data required when training from pretrained BERT.Approximation accuracy generally improves with model size up to BERT-large, where it may saturate under the current algorithm and available labels.Source: Figure 7 · p. 18, §C.2 · p. 18

Related papers

Loading 2406.06046v2…