Title: MATES: Model-Aware Data Selection for Efficient Pretraining with Data Influence Models
Venue: NeurIPS 2024
Accepted: yes
Keywords: Large Language Models

Pretraining

Data Curation
Abstract: 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.

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

Abstract

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

The power of large language models (LLMs) rises with scaling up [7; 23; 58]: pretraining models with more parameters on more data using more compute resources [23; 27]. Among these three aspects of scaling, compute is often the most restrictive factor, as current large-scale pretraining frequently demands millions of GPU hours [2; 8; 58], while the model parameters and the pretraining data amounts are determined based on the pre-allocated compute budget [23; 27].

This provides a unique opportunity to elevate the scaling law of pretraining through data selection since the available data sources, such as the web [42; 55], are orders of magnitude bigger than available compute resources and contain data of varying quality. Recent research has shown that effective data selection can improve the generalization ability of pretrained models [15; 62], enhance scaling efficiency [5], and introduce specialized capabilities [34]. These explorations mainly focus on heuristic-based methods, such as rule-based filtering [49; 50; 55], deduplication [1; 45; 57], proximity to high-quality corpora [17; 33; 61; 64], and prompting LLMs [52; 62]. Despite their success on certain datasets and models, these techniques rely heavily on the static heuristics and often overlook the evolving nature of pretraining [39]. As a result, their performance tends to be limited when applied to the real-world pretraining scenarios. [3; 33].

The data preferences of pretraining models dynamically shift as they progress through different stages of pretraining [4; 11; 39; 40]. For instance, Figure 1a illustrates how the data influence measured by the pretraining model evolves at different pretraining steps. As a result, the data quality measurement should also keep pace with the model's evolving data preferences during the pretraining. This leads us to our core research question: How can we precisely track the data influence with the pretraining model and efficiently select pretraining data based on the acquired influence?

In this paper, we introduce M odelA ware data selection with da T a influenc E model S (MATES), a new pretraining paradigm where pretraining data is selected on-the-fly by a data influence model capturing the ever-changing data preferences of the pretraining model. To track the model's data preferences, we locally probe the oracle data influence by evaluating the pretraining model's perfor-

10k

20k

30k

40k

10k

20k

30k

40k

1.000

0.413

0.505

0.317

0.413

1.000

0.522

0.389

0.505

0.522

1.000

0.623

0.317

0.389

0.623

1.000

0.4

0.6

0.8

1.0

Data Influence Correlation

(a) Preference correlation.

4

8

12

16

Training FLOPs (1e19)

44

45

46

47

48

Avg. Downstream Acc

2.3x Faster

MATES

Random

(b) Accuracy w.r.t. FLOPs.

mance on a reference task after training on individual data points. Then, we train a small data influence model using the locally probed oracle data influence, which then predicts data influence over the whole pretraining corpus and selects the most influential data for the next pretraining stage. This side-by-side learning framework ensures that the data influence model continuously adapts to the evolving data preferences of the pretraining model, providing the most valuable data accordingly.

Our pretraining experiments with 410M and 1B models on the C4 dataset [50] demonstrate that MATES can significantly outperform random selection by an average zero-shot accuracy of 1.3% (410M) and 1.1% (1B) across various downstream tasks, ranging from reading comprehension, commonsense reasoning, and question answering. MATES doubles the gains obtained by state-of-theart data selection approaches that rely on signals from reference models larger than the pretraining model. Furthermore, our model-aware data selection significantly elevates the scaling curve of pretraining models, as shown in Figure 1b, reducing the total FLOPs required to achieve certain downstream performances by more than half. Further analyses confirm the advantages of our locally probed oracle data influence and the effective approximation of this oracle with data influence models. Ablation studies demonstrate the robustness of MATES across various hyperparameter settings and different design choices of the data influence models.

We summarize our main contributions as follows:

We propose a model-aware data selection framework, MATES, where a small data influence model continuously adapts to the constantly changing data preferences of the pretraining model and selects the training data to optimize the efficacy of the pretraining process.

We effectively collect oracle data influence through local probing with the pretraining model and use a small BERT-base model to approximate it accurately.

Weempirically verify the superiority of MATES over rule-based, influence-function-based, and LLM-based selection methods and the effectiveness of the probed oracle and its approximation.

2 Related work

Early approaches on data selection relied heavily on manual intuitions. For example, T5 [50] first proposed the C4 pipeline, followed by Gopher rules [49], which utilized criteria like document length, mean word length, and the presence of harmful or stop words to curate data. Recent FineWeb dataset [46] further applied quality and repetition filters on top of these basic rules. These rule-based data selection methods have been shown effective as an initial data curation step [46], though manual intuitions may not capture the nuances of models' data preferences [36].

Data Influence

Model

Predict data

influence

Collect data

influence

Training w/

selected data

Training

Initialization

Data Influence

Model

Collect data

influence

Training w/

selected data

Training

Initialization

…

…

Training w/

random data

Objective

Maximize the target

performance of the

language model within

the given FLOPs

Predict data

influence

Deduplication is another standard approach in pretraining data selection. Specifically, Lee et al. [32] and Penedo et al. [45] explored exact string match and fuzzy MinHash to filter out duplicate sequences. SemDeDup [1] further leveraged pretrained embeddings to identify semantic duplicates, and D4 [57] introduced diversification factors into the deduplication process. These methods effectively narrow down the number of similar documents in a corpus and are often used together with quality-oriented selection techniques to improve performance.

Selecting data that is proximate to high-quality corpora can enhance the data quality as well [17; 61]. Techniques leveraging n-gram similarity [17; 64; 33] and language modeling perplexity [8; 14; 16; 61] have been adopted to evaluate how closely sequences in large datasets approximate high-quality data. As the size of pretraining data grows, the effectiveness of proximity-based methods becomes unclear, as they potentially reduce the diversity of the pretraining data and, consequently, the general capability of the pretrained models [36].

Recent advancements explore the use of LLMs to improve the pretraining data quality. For instance, QuRating [62] and Ask-LLM [52] employed LLMs like GPT-3.5 to annotate high-quality documents. Maini et al. [38] rephrased web corpora by providing LLMs with detailed prompts to balance quality and diversity. These methods leverage the capabilities of strong reference LLMs, which are often several orders of magnitude larger, to guide the pretraining of smaller models.

Influence functions [29; 59] provide a theoretical tool to assess the impact of individual data points on a model's performance. However, they face scalability challenges in the context of LLMs due to the expensive gradient calculations [21; 54]. To efficiently and robustly approximate influence functions, TRAK [44] performed Taylor approximation and gradient dimension reduction, making influence computation feasible for pretraining experiments [15]. Nevertheless, the computational cost remains prohibitive for model-aware data selection, which requires tracking the evolving data preferences of pretraining models on the fly.

On the other hand, many researchers have proposed curriculum learning strategies that dynamically adjust the data distribution in the pretraining [11; 24; 47; 51; 56; 62]. ELECTRA-style models [11] incorporated curriculum learning in their pretraining process by synchronously training the model with an auxiliary generator, which provided more and more difficult training signals for the discriminator. This implicit curriculum significantly improved the pretraining efficiency on denoising language models [4; 11; 19; 39; 65]. Other methods have explicitly designed the curriculum for pretraining data selection, such as decreasing gradient norms [56], least certainty [24; 51], and increasing expertise [62], demonstrating the benefits of adapting the pretraining data signal according to the model's ever-changing preferences.

3 Methods

This section first introduces the model-aware data selection framework MATES (§ 3.1) and then proposes a local probing technique to collect oracle data influence during pretraining (§ 3.2).

Algorithm 1 Model-Aware Data Selection

Require: Training data D t , hold-out data D h , reference data D r , pretraining model M , optimizer A , total training step T , selected size k , data influence model Θ , sampling temperature τ , update step U

3.1 Model-aware data selection framework

MATES selects the most effective data for pretraining a language model, aiming to maximize its final target performance, as illustrated in Figure 2. The target performance here can be evaluated using any downstream task or their combinations. Specifically, we leverage non-evaluation data as a reference for the model's target performance and select the pretraining data according to the reference loss.

Formally, given a sizen pretraining dataset D t and the current model state M , in each iteration, the objective of data selection is to find an optimal batch B ∗ from D t to minimize the loss L over the reference data D r after training M on B ∗ :

where A ( M , B ) denotes the optimization of model M on a batch B , e.g., one-step training with Adam [28] and ℓ denotes the function to compute the model loss on an input-output pair ( x, y ) .

There are two challenges to implement this framework. First, enumerating all possible batches will exponentially increase computational complexity. Second, obtaining oracle data influence for all pretraining data points is challenging. To address these issues, we introduce two techniques: pointwise data influence and data influence parameterization.

Pointwise Data Influence. To avoid the computationally intensive task of enumerating all possible batches, a more practical workaround is to decompose the group influence into the pointwise influence [15; 44]. Following previous research [44], we aggregate all the data influences by the summation, assuming that each data point x i has an independent influence irrespective of the others:

where I M is the oracle pointwise data influence function based on model state M . I M continuously changes along with the model pretraining.

Data Influence Parameterization. Estimating oracle pointwise data influence normally involves gradient-based calculation [21; 29; 44], which is impractical to perform over millions of pretraining examples for every pretraining model state. To make the data influence collection feasible, we propose to collect the oracle data influence on a small hold-out dataset D h (sampled from the same distribution as D t ) and fine-tune a small data influence model Θ on { ( x i , I M ( x i ; D r )) | x i ∈ D h } to approximate the oracle. This data influence parameterization process transfers the costly influence computation to the small data influence model's inference.

Then, we obtain the influence prediction Θ( x i ) with the fine-tuned data influence model over all the training examples x i ∈ D t . For better efficiency, we only asynchronously update the data influence model every U steps with the new oracle I M . Therefore, one data influence model checkpoint can select the entire data subset S ∗ k for the next U steps of pretraining. The selection process uses the Gumbel-Topk algorithm [30; 62] to sample data from D t , with influence scores as weights:

where τ is the sampling temperature. The update step U is chosen to balance the data selection efficiency with the evolving data influence. Empirically, we warm up the model with a randomly sampled sizek subset from D t in the initial U steps. The overall pretraining and data selection pipeline of MATES is illustrated in Algorithm 1.

3.2 Locally probed oracle data influence

The rest of this section presents our method to estimate the oracle data influence I M . We start the derivation from the standard influence functions [29; 59] that quantify the reference loss change if one data point x i in the training data is upweighted by a small ϵ . We denote the optimal model state after the upweighting as M ∗ ϵ,x i = arg min M 1 n ∑ n j =1 L ( x j |M ) + ϵ L ( x i |M ) and simplify the optimal model under ϵ = 0 case (i.e., no upweighting) as M ∗ . Then, the oracle data influence of upweighting x i is given by:

where H M ∗ = 1 n ∑ n j =1 ∇ 2 M L ( x j | M ∗ ) is the Hessian and is positive definite. The derivation from Eq. 7 to Eq. 8 is given by building a quadratic approximation to the empirical risk around M ∗ and taking a Newton step [29]. Now consider the case that we incorporate x i into the training data, which means ϵ = 1 n , then the parameter change due to the inclusion of x i is M ∗ 1 n ,x i -M ∗ ≈ -1 H -1 ∗ ∇ M L ( x i |M ∗ ) and the influence in Eq. 8 can be further represented as:

n M

In practice, we manage to obtain the data influence based on the current model state M , while the above influence calculation still remains meaningful in the non-converged state by adding a damping term λ that ensures H M + λI is positive definite [29]. Under this assumption, the first term in Eq. 11 can be regarded as a fixed value whatever x i is, since x i is sampled from the hold-out data D h . The second term in Eq. 11 can be locally probed with the one-step training of the current model with the new x i , i.e., A ( M , x i )) . This one-step training incorporates x i into the optimization of the current model. Finally, the influence of x i on the reference loss is:

This formula means, for each x i in the hold-out data D h , we run one-step training with the current model M and evaluate the difference in reference loss before and after one-step training. To ensure that a positive influence score reflects a beneficial impact on model performance, we empirically define the negative influence, L ( D r |M ) -L ( D r | A ( M , x i )) , as our locally probed oracle data influence of x i . A full derivation of locally probed oracle data influence can be found in Appendix A.

Methods

(#FLOPs

∗

1e19)

SciQ

ARC-E

ARC-C

LogiQA

OBQA

BoolQ

PIQA

WinoGrande

Average

410M Setting:

410M model, 25B tokens

Random (6.35)

64.1

(

1.5

)

40.2

(

1.0

)

25.6

(

1.3

)

24.7

(

1.7

)

29.4

(

2.0

)

58.9

(

0.9

)

39.7

(

0.5

)

67.1

(

1.1

)

50.6

(

1.4

)

44.5

(

1.3

)

DSIR (6.35)

63.1

(

1.5

)

39.9

(

1.0

)

23.8

(

1.2

)

27.0

(

1.7

)

28.4

(

2.0

)

58.3

(

0.9

)

39.6

(

0.5

)

66.8

(

1.1

)

51.5

(

1.4

)

44.3

(

1.3

)

LESS (246.35)

64.6

(

1.5

)

42.3

(

1.0

)

23.1

(

1.2

)

25.2

(

1.7

)

30.4

(

2.1

)

55.6

(

0.9

)

41.9

(

0.5

)

67.2

(

1.1

)

51.0

(

1.4

)

44.6

(

1.4

)

SemDeDup (7.81)

63.5

(

1.5

)

42.4

(

1.0

)

24.4

(

1.3

)

27.6

(

1.8

)

30.0

(

2.1

)

58.2

(

0.9

)

40.8

(

0.5

)

67.8

(

1.1

)

52.3

(

1.4

)

45.2

(

1.4

)

DsDm(10.72)

65.4

(

1.5

)

41.7

(

1.0

)

24.7

(

1.3

)

27.5

(

1.8

)

29.0

(

2.0

)

57.5

(

0.9

)

40.3

(

0.5

)

68.1

(

1.1

)

50.1

(

1.4

)

44.9

(

1.4

)

QuRating

∗

(26.35)

64.8

(

1.5

)

42.0

(

1.0

)

25.4

(

1.3

)

25.3

(

1.7

)

30.2

(

2.1

)

58.9

(

0.9

)

40.7

(

0.5

)

67.5

(

1.1

)

52.1

(

1.4

)

45.2

(

1.4

)

MATES (Ours) (8.11)

66.0

(

1.5

)

41.8

(

1.0

)

25.0

(

1.3

)

25.7

(

1.7

)

30.8

(

2.1

)

60.6

(

0.9

)

41.0

(

0.5

)

68.7

(

1.1

)

52.7

(

1.4

)

45.8

(

1.4

)

1B Setting:

1B model, 25B tokens

Random (17.67)

65.8

(

1.5

)

43.7

(

1.0

)

25.6

(

1.3

)

27.5

(

1.8

)

31.8

(

2.1

)

60.2

(

0.9

)

43.8

(

0.5

)

68.9

(

1.1

)

50.7

(

1.4

)

46.4

(

1.4

)

DSIR (17.67)

65.8

(

1.5

)

42.6

(

1.0

)

24.7

(

1.3

)

28.7

(

1.8

)

29.2

(

2.0

)

59.7

(

0.9

)

44.2

(

0.5

)

68.3

(

1.1

)

53.2

(

1.4

)

46.3

(

1.4

)

SemDeDup (19.13)

66.8

(

1.5

)

45.5

(

1.0

)

25.3

(

1.3

)

27.6

(

1.8

)

30.6

(

2.1

)

60.2

(

0.9

)

45.3

(

0.5

)

69.7

(

1.1

)

52.5

(

1.4

)

47.1

(

1.4

)

DsDm(22.04)

68.2

(

1.5

)

45.0

(

1.0

)

26.5

(

1.3

)

26.6

(

1.7

)

29.4

(

2.0

)

59.0

(

0.9

)

44.8

(

0.5

)

68.9

(

1.1

)

51.9

(

1.4

)

46.7

(

1.3

)

QuRating

∗

(37.67)

67.1

(

1.5

)

45.5

(

1.0

)

25.6

(

1.3

)

26.9

(

1.7

)

29.8

(

2.0

)

60.3

(

0.9

)

45.2

(

0.5

)

70.2

(

1.1

)

51.6

(

1.4

)

46.9

(

1.3

)

MATES (Ours) (19.97)

67.3

(

1.5

)

44.9

(

1.0

)

25.9

(

1.3

)

28.7

(

1.8

)

32.2

(

2.1

)

60.9

(

0.9

)

45.3

(

0.5

)

69.5

(

1.1

)

52.4

(

1.4

)

47.5

(

1.4

)

4 Experimental methodologies

Implementation Details. We pretrain 410M/1B models with Pythia [5] architecture from scratch on the C4 dataset [50] as our pretraining model M , and continuously fine-tune BERT-base [13] as our data influence model Θ . The data influence model is smaller than the pretraining model, ensuring the efficiency of data selection. More details of data influence models can be found in Appendix B.3. For model pretraining, we utilize Warmup-Stable-Decay (WSD) scheduler proposed in MiniCPM [25], which offers flexibility in the varying training length [22]. More details of WSD scheduler can be found in Appendix B.2.

For MATES selection, we sample 20% data with their influence scores as weights at each pretraining stage (10k steps). The sampling temperature τ is set to 1.0 to balance the data quality and the diversity. The update step U is also set to 10k so that the selected data is trained by one epoch. Following DsDm [15], we leverage LAMBADA [43] as our reference data D r . LAMBADA is a widely-used language modeling task and often serves as a validation task for language model pretraining [7; 8; 23]. A summary of these details can be found in Appendix B.1.

Evaluation Methods. We use lm-evaluation-harness [18] codebase to perform a holistic evaluation of the pretraining models across 9 downstream tasks, including SciQ [60], ARC-E [12], ARCC [12], LogiQA [35], OBQA [41], BoolQ [10], HellaSwag [66], PIQA [6], and WinoGrande [53]. These tasks cover the core abilities of the pretrained language model, ranging from reading comprehension, commonsense reasoning, and question answering. We report zero-shot accuracy for BoolQ and WinoGrande; otherwise, normalized zero-shot accuracy. We also report the total GPU FLOPs, including model pretraining and data selection, as both are parts of the scaling formulae [20].

The learning outcome of our data influence model is evaluated by the validation Spearman correlation between its predictions and the oracle data influence on a 10% hold-out validation set. This set is sampled from the collected data-oracle mapping { ( x i , I M ( x i ; D r )) | x i ∈ D h } .

Baselines. We compare MATES with random selection as well as the state-of-the-art pretraining data selection baselines, which include (1) DSIR [64]: proximity with Wikipedia by n-gram features. (2) SemDeDup [1]: removal of semantically duplicate sentences. (3) LESS [63]: computing cosine similarity between training and validation (LAMBADA) gradients as influential scores. 1 (4) DsDm [15]: static approximation of influence scores on LAMBADA by a well-trained proxy model [48]. (5) QuRating [62]: ranking with Llama-learned [58] quality scores identified by GPT-3.5 in terms of educational values. These baselines cover all the mainstream data selection schemes, ranging from simple heuristics, static influence functions, and LLM rating. Some recent methods, such as Ask-LLM [52], are not open-sourced yet, prohibiting direct comparisons.

2

4

6

8

Training FLOPs (1e19)

43

44

45

46

Avg. Downstream Acc

MATES

Random

(a) 410M w/ FLOPs.

10k 20k 30k 40k 50k

Training Steps

43

44

45

46

Avg. Downstream Acc

MATES

Random

(b) 410M w/ steps.

5

10

15

20

Training FLOPs (1e19)

45

46

47

48

Avg. Downstream Acc

MATES

Random

(c) 1B w/ FLOPs.

10k 20k 30k 40k 50k

Training Steps

45

46

47

48

Avg. Downstream Acc

MATES

Random

(d) 1B w/ steps.

5 Evaluation results

This section evaluates the effectiveness of MATES (§ 5.1), locally probed oracle data influence (§ 5.2), and data influence model (§ 5.3). It further presents a case study to illustrate the model's changing data preferences (§ 5.4). More ablation studies can be found in Appendix C.

5.1 Overall performance

MATES outperforms the state-of-the-art data selection approach. Table 1 presents the zeroshot evaluation of pretraining 410M/1B models with different data selection methods. MATES significantly outperforms random selection by an average downstream accuracy gain of 1.3% and 1.1% in the 410M and 1B settings, respectively. These gains are nearly double those achieved by the state-of-the-art data selection method, QuRating, which depends on larger reference models. Notably, we observe an absolute improvement of nearly 2.0% across most tasks in the 410M setting, except for ARC-C and LogiQA, where the predictions of 410M models are close to random guessing. In the 1B setting, MATES outperforms random selection across all 9 tasks, demonstrating the strong scalability and generalization capabilities of our method.

MATES selects the data with low costs. We also show a detailed breakdown of the pretraining cost of MATES in Table 2. The relative selection cost of larger models is generally smaller since their pretraining cost dominates the total FLOPs while the training and inference cost of our data influence model remains stable. Even with the 1B pretraining model, the wall clock time to collect one oracle data influence is only around 2.5 seconds on one GPU, which means we can get all the required 160k oracle scores during 50k-step pretraining on one node (8 GPUs) around 14 hours, which is significantly lower than the actual pretraining time (4 days). The inference speed of our data influence model can also be improved with data parallelism or a fast-inference framework like vLLM [31], re-

ducing the selection cost further. This breakdown analysis underscores the low expense of MATES in achieving model-aware data selection.

MATES significantly elevates the scaling curves. Figure 3 plots the performance of the pretraining models w.r.t. different FLOPs and steps. Measuring by FLOPs counts both the model pretraining and data selection costs, demonstrating the total compute expense during pretraining. Measuring by steps reflects the compute cost of the model pretraining alone, as the data selection can be trivially parallelized when more computational resources are available. At both 410M and 1B scales, MATES

-0.6

-0.3

0.0

0.3

Oracle Data Influence

0

1

2

3

4

Probability Density

Standard Deviation=0.114

Pos (46%), Neg (54%)

Val. Spearman=0.50

(a) LAMBADA.

-0.6

-0.3

0.0

0.3

Oracle Data Influence

0

1

2

3

4

Probability Density

Standard Deviation=0.068

Pos (15%), Neg (85%)

Val. Spearman=0.35

(b) ARC-E (MC).

-0.6

-0.3

0.0

0.3

Oracle Data Influence

0

1

2

3

4

Probability Density

Standard Deviation=0.085

Pos (28%), Neg (72%)

Val. Spearman=0.39

(c) ARC-E (LM).

-0.6

-0.3

0.0

0.3

Oracle Data Influence

0

1

2

3

4

Probability Density

Standard Deviation=0.110

Pos (18%), Neg (82%)

Val. Spearman=0.42

(d) FLAN.

significantly elevates the scaling curves compared to random selection, reducing the FLOPs and pretraining steps required to reach a certain downstream performance by more than half. Scaling efficiency is more evident at the 1B scale, where model pretraining dominates 88.5% of FLOPs versus 11.5% for data selection. As a result, MATES reaches the same performance as random selection with only 43.3% of the total FLOPs. A similar observation can be found in Figure 9, where the performance of MATES is comparable to or higher than the full pretraining using 3x data. These results reveal the promising potential of MATES in elevating the scaling law of foundation models.

5.2 Effectiveness of locally probed oracle data influence

This set of experiments analyzes the effectiveness of oracle data influence with different reference tasks. Besides LAMBADA used in our main experiment, we also consider taking the training sets of ARC-E [12] and FLAN [9] as the reference tasks to show the generalization ability of our method. ARC-E represents one of the knowledge-based question-answering tasks, while FLAN represents a large set of varied instruction-formatted data. For ARC-E, we construct each example either as a multiple-choice selection (i.e., outputting the correct option [A-D]) or as a language modeling task (i.e., outputting the verbalized answer) to investigate whether the task format will affect the data influence collection and parameterization.

In Figure 4, we demonstrate the oracle data influence distribution with different reference tasks, the standard deviation of the distribution, and the proportions of the data with positive/negative oracle data influence. The oracle distribution remains spread-out across all reference tasks, indicating that our oracle effectively differentiates data influences. Notably, the positive influence proportion for LAMBADA in Figure 4a is higher than others by more than 20%, suggesting that more data is deemed beneficial when LAMBADA serves as the reference. We hypothesize that this is because LAMBADA is essentially a word prediction task, which aligns more closely with the pretraining objective compared to knowledge-based or instruction-following tasks. This hypothesis is further supported by the results for ARC-E (Figure 4b/4c), where the language modeling format identifies more positively influential data points compared to the multiple-choice format.

We further validate the effectiveness of the oracle across different reference tasks by sampling 20% data with oracle scores as weights. Due to the infeasibility of obtaining oracle scores for all the pretraining data, we only run one short decay stage with the data selected by different oracle scores.

As shown in Table 3, taking ARC-E as the reference task can benefit the model's in-domain accuracy, but its generalization performance is worse than using LAMBADA. In contrast, FLAN benefits a wider range of downstream tasks due to its diverse instructions. However, there remains a trade-off between the performance of different tasks, so the average accuracy of choosing FLAN is similar to that of LAMBADA. This experiment highlights the robustness of our oracle, as it is not limited to specific reference data but generalizes effectively across multiple reference tasks.

5.3 Effectiveness of data influence model

This experiment studies the effectiveness of data influence model. As shown in Figure 4, our data influence model effectively approximates the oracle data influence across various reference tasks. Nevertheless, a higher standard deviation in the oracle distribution tends to enhance the validation Spearman correlation. This suggests that greater variability in influence scores can provide more diverse signals for the data influence model to capture, making it better approximate the oracle distribution.

Figure 5 compares MATES with static data influence models trained on influence from a 10k or a 50k random-pretrained model checkpoint. In Figure 5a, we measure the validation Spearman correlation between the predictions of data influence models and the oracle data influence probed with each pretraining checkpoint. The correlations of static data influence models are always below 0.5, while data influence models in MATES, with dynamic updates, can capture the ever-changing data preferences more and more precisely along with the model pretraining (e.g., the correlation is around 0.7 in 40k steps). The effects of model-aware data selection are directly reflected in downstream accuracy. In Figure 5b, the data selected by static data influence models will cause a notable performance drop, especially at the early pretraining stage. These observations confirm the ever-changing nature of data preferences in pretraining and the advantages of model-aware data selection to elevate the scaling curves of pretraining.

5.4 Case study

This case study demonstrates representative examples to illustrate the evolving preferences of the pretraining model in detail. As shown in Table 4, the model at the early pretraining stage (the 10k checkpoint) tends to favor learning natural narrative examples without delving too deeply into specific knowledge. At the 20k checkpoint, it appears to shift focus toward factual knowledge (e.g., pages from Wikipedia) while gradually reducing reliance on natural narratives. At the 30k checkpoint, the model shows a preference for more detailed academic text, such as official teaching slides. By the 40k checkpoint, the model may begin learning more long-tail knowledge, like Telescopic Forklift Training. These examples provide insights into the evolving nature of the model's data preferences throughout pretraining. Although they may not cover every aspect of the selected data, our observation highlights the necessity to adapt data selection strategies to different model learning stages.

6 Discussion and limitations

Combinational Measurement of Data Influence. One primary assumption in our work is that each data point contributes independently to the pretraining outcome, without accounting for its interactions with other data points. Despite a common hypothesis [15; 44], the pretraining essentially applies the long-term combinational effect of batched data on language models, and the learning of many advanced capabilities is accumulative. Efficiently and effectively measuring and learning the combinational and accumulative nature of the pretraining process can make us better understand and leverage the value of data [37].

10k

20k

30k

40k

Training Steps

0.3

0.4

0.5

0.6

0.7

Val. Spearman

MATES

10k Training

50k Training

20k

30k

40k

50k

Training Steps

43

44

45

46

47

Avg. Downstream Acc

MATES

10k Selection

50k Selection

Exploratory Scale. As an exploratory research work, our experiments are conducted at a moderate scale, with a pretraining model of 410M or 1B parameters. Although the trend from 410M to 1B indicates the robustness of our observations, it remains unclear how well our methods scale up to production-level models with billions of parameters and trillions of pretraining tokens. On the one hand, moving to that scale provides more headroom for data selection with more urgent needs for efficiency, more leniency on the relatively small compute spent on data selection, and a larger pool of candidate data. On the other hand, large-scale pretraining may yield various stability issues that require dedicated work to introduce new techniques [58; 67]. We leave the exploration of larger models to future work.

7 Conclusion

In this paper, we introduce MATES, a novel framework to enhance the efficiency and effectiveness of language model pretraining through model-aware data selection. MATES leverages a data influence model to continuously capture the evolving data preferences of the pretraining model throughout the pretraining process, thereby selecting the training data most effective for the current pretraining stage. To achieve that, we locally probe the oracle data influence on a reference task using the pretraining model and fit the data influence model on the probed oracle. Our empirical results demonstrate that MATES surpasses random, rule-based, influence-function-based, and LLM-based data selection methods on pretraining, significantly elevating the scaling curves of pretraining LLMs. Further analyses confirm the effectiveness of our locally probed oracle data influence and the accurate approximation of this oracle with data influence models. Our work successfully demonstrates the potential of model-aware data curation in pretraining, and we hope it will motivate further explorations on improving the scaling law of foundation models through better data curation techniques.

References

Amro Kamal Mohamed Abbas, Kushal Tirumala, Daniel Simig, Surya Ganguli, and Ari S. Morcos. SemDeDup: Data-efficient learning at web-scale through semantic deduplication. In Workshop on Multimodal Representation Learning: Perks and Pitfalls , 2023.

Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. GPT-4 technical report. ArXiv preprint , 2023.

Alon Albalak, Yanai Elazar, Sang Michael Xie, Shayne Longpre, Nathan Lambert, Xinyi Wang, Niklas Muennighoff, Bairu Hou, Liangming Pan, Haewon Jeong, et al. A survey on data selection for language models. ArXiv preprint , 2024.

Payal Bajaj, Chenyan Xiong, Guolin Ke, Xiaodong Liu, Di He, Saurabh Tiwary, Tie-Yan Liu, Paul Bennett, Xia Song, and Jianfeng Gao. METRO: Efficient denoising pretraining of large scale autoencoding language models with model generated signals. ArXiv preprint , 2022.

Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar van der Wal. Pythia: A suite for analyzing large language models across training and scaling. In Proc. of ICML , 2023.

Yonatan Bisk, Rowan Zellers, Ronan LeBras, Jianfeng Gao, and Yejin Choi. PIQA: Reasoning about physical commonsense in natural language. In Proc. of AAAI , 2020.

Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Proc. of NeurIPS , 2020.

Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, and et al. PaLM: Scaling language modeling with pathways. ArXiv preprint , 2022.

Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Eric Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, Albert Webson, Shixiang Shane Gu, Zhuyun Dai, Mirac Suzgun, Xinyun Chen, Aakanksha Chowdhery, Sharan Narang, Gaurav Mishra, Adams Yu, Vincent Zhao, Yanping Huang, Andrew Dai, Hongkun Yu, Slav Petrov, Ed H. Chi, Jeff Dean, Jacob Devlin, Adam Roberts, Denny Zhou, Quoc V. Le, and Jason Wei. Scaling instruction-finetuned language models. JMLR , 2024.

Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. BoolQ: Exploring the surprising difficulty of natural yes/no questions. In Proc. of NAACL-HLT , 2019.

Kevin Clark, Minh-Thang Luong, Quoc V. Le, and Christopher D. Manning. ELECTRA: Pre-training text encoders as discriminators rather than generators. In Proc. of ICLR , 2020.

Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? Try ARC, the ai2 reasoning challenge. ArXiv preprint , 2018.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proc. of NAACL-HLT , 2019.

Nan Du, Yanping Huang, Andrew M. Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, Barret Zoph, Liam Fedus, Maarten P. Bosma, Zongwei Zhou, Tao Wang, Yu Emma Wang, Kellie Webster, Marie Pellat, Kevin Robinson, Kathleen S. Meier-Hellstern, Toju Duke, Lucas Dixon, Kun Zhang, Quoc V. Le, Yonghui Wu, Zhifeng Chen, and Claire Cui. GLaM: Efficient scaling of language models with mixture-of-experts. In Proc. of ICML , 2022.

Logan Engstrom, Axel Feldmann, and Aleksander Madry. DsDm: Model-aware dataset selection with datamodels. In Proc. of ICML , 2024.

Ruyi Gan, Ziwei Wu, Renliang Sun, Junyu Lu, Xiaojun Wu, Dixiang Zhang, Kunhao Pan, Ping Yang, Qi Yang, Jiaxing Zhang, et al. Ziya2: Data-centric learning is all LLMs need. ArXiv preprint , 2023.

Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The Pile: An 800gb dataset of diverse text for language modeling. ArXiv preprint , 2021.

Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 2023.

Linyuan Gong, Chenyan Xiong, Xiaodong Liu, Payal Bajaj, Yiqing Xie, Alvin Cheung, Jianfeng Gao, and Xia Song. Model-generated pretraining signals improves zero-shot generalization of text-to-text transformers. In Proc. of ACL , 2023.

Sachin Goyal, Pratyush Maini, Zachary C Lipton, Aditi Raghunathan, and J Zico Kolter. Scaling laws for data filtering-data curation cannot be compute agnostic. In Proc. of CVPR , 2024.

Roger Grosse, Juhan Bae, Cem Anil, Nelson Elhage, Alex Tamkin, Amirhossein Tajdini, Benoit Steiner, Dustin Li, Esin Durmus, Ethan Perez, et al. Studying large language model generalization with influence functions. ArXiv preprint , 2023.

Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben allal, Leandro Von Werra, and Martin Jaggi. Scaling laws and compute-optimal training beyond fixed training durations. In Workshop on Efficient Systems for Foundation Models II , 2024.

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katherine Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Oriol Vinyals, Jack W. Rae, and Laurent Sifre. An empirical analysis of computeoptimal large language model training. In Proc. of NeurIPS , 2022.

Xudong Hong, Sharid Loáiciga, and Asad Sayeed. A surprisal oracle for active curriculum language modeling. In Proc. of CoNLL , 2023.

Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, et al. MiniCPM: Unveiling the potential of small language models with scalable training strategies. ArXiv preprint , 2024.

Andrew Ilyas, Sung Min Park, Logan Engstrom, Guillaume Leclerc, and Aleksander Madry. Datamodels: Predicting predictions from training data. In Proc. of ICML , 2022.

Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. ArXiv preprint , 2020.

Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Proc. of ICLR , 2015.

Pang Wei Koh and Percy Liang. Understanding black-box predictions via influence functions. In Proc. of ICML , 2017.

Wouter Kool, Herke van Hoof, and Max Welling. Stochastic beams and where to find them: The gumbeltop-k trick for sampling sequences without replacement. In Proc. of ICML , 2019.

Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proc. of SOSP , 2023.

Katherine Lee, Daphne Ippolito, Andrew Nystrom, Chiyuan Zhang, Douglas Eck, Chris Callison-Burch, and Nicholas Carlini. Deduplicating training data makes language models better. In Proc. of ACL , 2022.

Jeffrey Li, Alex Fang, Georgios Smyrnis, Maor Ivgi, Matt Jordan, Samir Gadre, Hritik Bansal, Etash Guha, Sedrick Keh, Kushal Arora, et al. DataComp-LM: In search of the next generation of training sets for language models. In Proc. of NeurIPS , 2024.

Zhenghao Lin, Zhibin Gou, Yeyun Gong, Xiao Liu, Yelong Shen, Ruochen Xu, Chen Lin, Yujiu Yang, Jian Jiao, Nan Duan, et al. Rho-1: Not all tokens are what you need. ArXiv preprint , 2024.

Jian Liu, Leyang Cui, Hanmeng Liu, Dandan Huang, Yile Wang, and Yue Zhang. LogiQA: A challenge dataset for machine reading comprehension with logical reasoning. In Proc. of IJCAI , 2020.

Shayne Longpre, Gregory Yauney, Emily Reif, Katherine Lee, Adam Roberts, Barret Zoph, Denny Zhou, Jason Wei, Kevin Robinson, David Mimno, and Daphne Ippolito. A pretrainer's guide to training data: Measuring the effects of data age, domain coverage, quality, & toxicity. In Proc. of NAACL-HLT , 2024.

Scott M. Lundberg and Su-In Lee. A unified approach to interpreting model predictions. In Proc. of NeurIPS , 2017.

Pratyush Maini, Skyler Seto, Richard Bai, David Grangier, Yizhe Zhang, and Navdeep Jaitly. Rephrasing the web: A recipe for compute and data-efficient language modeling. In Proc. of ACL , 2024.

Yu Meng, Chenyan Xiong, Payal Bajaj, Saurabh Tiwary, Paul Bennett, Jiawei Han, and Xia Song. COCOLM: correcting and contrasting text sequences for language model pretraining. In Proc. of NeurIPS , 2021.

Yu Meng, Chenyan Xiong, Payal Bajaj, Saurabh Tiwary, Paul N. Bennett, Jiawei Han, and Xia Song. Pretraining text encoders with adversarial mixture of training signal generators. In Proc. of ICLR , 2022.

Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? A new dataset for open book question answering. In Proc. of EMNLP , 2018.

Arnold Overwijk, Chenyan Xiong, and Jamie Callan. ClueWeb22: 10 billion web documents with rich information. In Proc. of SIGIR , 2022.

Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Proc. of ACL , 2016.

Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guillaume Leclerc, and Aleksander Madry. TRAK: Attributing model behavior at scale. In Proc. of ICML , 2023.

Guilherme Penedo, Quentin Malartic, Daniel Hesslow, Ruxandra Cojocaru, Hamza Alobeidli, Alessandro Cappelli, Baptiste Pannier, Ebtesam Almazrouei, and Julien Launay. The RefinedWeb dataset for Falcon LLM: Outperforming curated corpora with web data only. In Proc. of NeurIPS , 2023.

Guilherme Penedo, Hynek Kydlíˇ cek, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. The FineWeb datasets: Decanting the web for the finest text data at scale. ArXiv preprint , 2024.

Ziheng Qin, Kai Wang, Zangwei Zheng, Jianyang Gu, Xiangyu Peng, Daquan Zhou, Lei Shang, Baigui Sun, Xuansong Xie, Yang You, et al. InfoBatch: Lossless training speed up by unbiased dynamic data pruning. In Proc. of ICLR , 2024.

Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, Jeffrey Dean, and Sanjay Ghemawat. Language models are unsupervised multitask learners. In Proc. of OSDI , 2019.

Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training Gopher. ArXiv preprint , 2021.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. JMLR , 2020.

Ravi S Raju, Kyle Daruwalla, and Mikko Lipasti. Accelerating deep learning with dynamic data pruning. ArXiv preprint , 2021.

Noveen Sachdeva, Benjamin Coleman, Wang-Cheng Kang, Jianmo Ni, Lichan Hong, Ed H Chi, James Caverlee, Julian McAuley, and Derek Zhiyuan Cheng. How to train data-efficient LLMs. ArXiv preprint , 2024.

Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. In Proc. of AAAI , 2020.

Andrea Schioppa, Polina Zablotskaia, David Vilar, and Artem Sokolov. Scaling up influence functions. In Proc. of AAAI , 2022.

Luca Soldaini, Rodney Kinney, Akshita Bhagia, Dustin Schwenk, David Atkinson, Russell Authur, Ben Bogin, Khyathi Chandu, Jennifer Dumas, Yanai Elazar, Valentin Hofmann, Ananya Jha, Sachin Kumar, Li Lucy, Xinxi Lyu, Nathan Lambert, Ian Magnusson, Jacob Morrison, Niklas Muennighoff, Aakanksha Naik, Crystal Nam, Matthew Peters, Abhilasha Ravichander, Kyle Richardson, Zejiang Shen, Emma Strubell, Nishant Subramani, Oyvind Tafjord, Evan Walsh, Luke Zettlemoyer, Noah Smith, Hannaneh Hajishirzi, Iz Beltagy, Dirk Groeneveld, Jesse Dodge, and Kyle Lo. Dolma: an open corpus of three trillion tokens for language model pretraining research. In Proc. of ACL , 2024.

Megh Thakkar, Tolga Bolukbasi, Sriram Ganapathy, Shikhar Vashishth, Sarath Chandar, and Partha Talukdar. Self-influence guided data reweighting for language model pre-training. In Proc. of EMNLP , 2023.

Kushal Tirumala, Daniel Simig, Armen Aghajanyan, and Ari Morcos. D4: Improving LLM pretraining via document de-duplication and diversification. In Proc. of NeurIPS , 2023.

Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. ArXiv preprint , 2023.

Sanford Weisberg and R Dennis Cook. Residuals and influence in regression. 1982.

Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsourcing multiple choice science questions. In Workshop on Noisy User-generated Text , 2017.

Guillaume Wenzek, Marie-Anne Lachaux, Alexis Conneau, Vishrav Chaudhary, Francisco Guzmán, Armand Joulin, and Edouard Grave. CCNet: Extracting high quality monolingual datasets from web crawl data. In Proc. of LREC , 2020.

Alexander Wettig, Aatmik Gupta, Saumya Malik, and Danqi Chen. QuRating: Selecting high-quality data for training language models. In Proc. of ICML , 2024.

Mengzhou Xia, Sadhika Malladi, Suchin Gururangan, Sanjeev Arora, and Danqi Chen. LESS: Selecting influential data for targeted instruction tuning. In Proc. of ICML , 2024.

Sang Michael Xie, Shibani Santurkar, Tengyu Ma, and Percy Liang. Data selection for language models via importance resampling. In Proc. of NeurIPS , 2023.

Zhenhui Xu, Linyuan Gong, Guolin Ke, Di He, Shuxin Zheng, Liwei Wang, Jiang Bian, and Tie-Yan Liu. MC-BERT: Efficient language pre-training via a meta controller. ArXiv preprint , 2020.

Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can a machine really finish your sentence? In Proc. of ACL , 2019.

Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. OPT: Open pre-trained transformer language models. ArXiv preprint , 2022.

A Full derivation of locally probed oracle data influence

This section presents a full derivation of locally probed oracle data influence. First, we define the optimal model state after upweighting one training data point x i by ϵ as:

and simplify the optimal model under ϵ = 0 case (i.e., no upweighting) as M ∗ . Following Koh and Liang [29], data influence quantifies the reference loss change before and after one data point x i in the training data is upweighted by a small ϵ , i.e.:

To compute d M ∗ ϵ,x i dϵ ∣ ∣ ∣ ∣ ϵ =0 , we consider the optimality condition at M ∗ ϵ,x i :

Differentiating both sides with respect to ϵ :

where H M ∗ = 1 n ∑ n j =1 ∇ 2 M L ( x j | M ∗ ) is the Hessian and is positive definite. Therefore,

With the first-order Taylor expansion, we can approximate the change in model parameters:

For ϵ = 1 n , i.e., incorporating x i into the training data, this becomes:

Substituting H -1 M ∗ ∇ M L ( x i | M ∗ ) in Eq. 19 with -n ( M ∗ 1 n ,x i -M ∗ ) , we get:

In practice, we manage to obtain the data influence based on the current model state M , while the above influence calculation still remains meaningful in the non-converged state by adding a damping term λ that ensures H M + λI is positive definite [29]. Finally, the locally probed oracle data influence of x i is:

To ensure that a positive influence score reflects a beneficial impact on model performance, we empirically define the negative influence, L ( D r | M ) -L ( D r | A ( M , x i )) , as our locally probed oracle data influence of x i .

B Experimental details

This section provides our experimental configurations (§ B.1), details of WSD scheduler (§ B.2), and design of data influence model (§ B.3).

B.1 Experimental configurations

We provide all experimental configurations in Table 5.

B.2 WSDscheduler

WSDlearning rate scheduler is initially proposed in MiniCPM [25] and is found to scale predictably and reliably similar to the widely-used cosine learning rate scheduler [22]. Moreover, WSD offers better flexibility than cosine for training across different lengths since its learning rate is constant during the stable stage. The learning rate of WSD scheduler is configured as follows:

where t , W , S , and D represent the number of steps now, at the end of the warmup, stable, and decay stages, respectively. η is the max learning rate. In our main experiments, we choose W = 2000, S = 50000, and D = 200. Inspired by MiniCPM [25], each checkpoint is evaluated after the short decay stage for better stability. We run all experiments on 8 A6000 GPUs, which will take 2 days for 410M models and 4 days for 1B models.

B.3 Design of data influence model

Our BERT-based data influence model averages all the hidden representations of the last model layer to obtain the sequence representation h ∈ R H , where H is the hidden size of the model. Note that BERT can only support a maximum input sequence length of 512. To deal with our pretraining sequence length of 1024, we divide one sequence into two chunks and forward them separately. Then, we average the hidden representations from both chunks to obtain the final h . This vector will be multiplied by a regression output weight w o ∈ R H to get the

model prediction w o · h . The training objective is the mean squared error between the model prediction and the normalized ground truth I M .

Note that BERT has a different tokenizer from our pretraining models, but the average sequence length of all the examples after the tokenization is almost the same. Our chunk-based design can be easily extended to longer pretraining sequences in the future.

C Additional experiments

This section presents the comparison of different data influence attribution methods (§ C.1), the ablation study on data influence models (§ C.2), and the analysis of oracle data influence (§ C.3).

C.1 Comparison of different data influence attribution methods

We compare the performance of data selection directly using our oracle with DsDm at the short decay stage in Table 6. Our locally probed oracle, utilizing one-step training, outperforms DsDm, which relies on Taylor approximation and gradient dimension reduction. The performance gap (0.8 zero-shot accuracy gain) is significant, considering the decay stage only consists of 200 steps. This improvement can be attributed to two factors: (1) we leverage the current model state to calculate the data influence rather than relying on an existing checkpoint like DsDm, and (2) we perform one-step training to obtain the oracle score, considering the training dynamics of the pretraining model and eliminating the precision loss from multiple approximation manipulations in DsDm. It is always costly to acquire the optimal data influence oracle for the entire pretraining dataset, but our local probing method with data influence parameterization beats Taylor approximation with fewer FLOPs, offering a new direction for future exploration.

C.2 Ablation study on data influence models

This group of experiments conducts ablation studies on the key hyperparameters.

2.6

2.8

3.0

3.2

Influence Score (BS = 1)

2.8

2.9

3.0

Influence Score (BS = 16)

Spearman = 0.898

(a) Batch size (BS).

0.35

0.40

0.45

LAMBADA Acc

0.43

0.44

0.45

Avg. Downstream Acc

Spearman = 0.877

(b) LAMBADA as

D

r

.

0.40

0.41

0.42

ARC-E Acc

0.43

0.44

0.45

Spearman = 0.399

(c) ARC-E as

D

r

.

0.66

0.67

0.68

PIQA Acc

0.43

0.44

0.45

Spearman = 0.284

(d) PIQA as

D

r

.

Update Step and Sampling Temperature. Table 7 shows the impact of hyperparameters on downstream performance. We initialize the model with the 40k-step MATES checkpoint and adopt different hyperparameters in the 40k-50k training. Decreasing the update step U from 10k to 5k and 2.5k leads to little fluctuations since the model preferences may not dramatically change within 5k steps. Varying the sampling temperature τ to 2.0 and 0.0 causes decreased performances. The zero temperature extremely up-weights high-quality data, reducing data diversity, while higher temperatures like 2.0 do not sufficiently emphasize data quality. Our observations highlight the necessity of balancing quality and diversity in the data selection.

Data Selection Ratio. As shown in Figure 6, MATES shows the consistent gains compared to random selection using either low/high-selection ratio, ranging from 1/200 to 1/2. We also find the optimal sampling rate of a larger model (1B) is smaller than that of a smaller model (410M). We hypothesize that a larger model may require a more aggressive selection ratio, such as 10%, since it is more robust to the subtle distribution change of the training data. However, too low (1/200) or high selection ratio (1/2) does not perform as good, as a low ratio may harm the diversity and a high ratio does not leverage the strength of the data influence enough.

Oracle Label Amount and Data Influence Model Scale. Figure 7a demonstrates the impact of oracle label amounts by comparing the validation Spearman correlation of different data influence models at 40k steps. The data influence model is initialized from either the pretrained BERT or its

1/200

1/10

1/5

1/2

1

Selection Ratio

0.45

0.46

0.47

Avg. Downstream Acc

410M

1B

last checkpoint at 30k steps. We observe that continuous fine-tuning from 30k steps requires less than half of the oracle data compared to training from the pretrained BERT, which can significantly reduce the cost of collecting new oracle data. Furthermore, Figure 7b studies the impact of parameter counts in the data influence model. Generally, the approximation becomes more accurate as the number of parameters increases until BERT-large, which may have become saturated with the current learning algorithm using the available oracle label. We have proved the effectiveness of our data influence models through extensive experiments and thus, leave the exploration of better data influence parameterization algorithms and stronger data influence models to future work.

C.3 Analysis of oracle data influence

To further demonstrate the stability of our oracle data influence, we enlarge the one-step training batch size (BS) from 1 to 16. Following Ilyas et al. [26], we utilize LASSO regression to separate each data's influence score in the BS = 16 setup and calculate the Spearman correlation between BS = 1 and BS = 16 scores. More details of LASSO regression can be found in Appendix C.3. Figure 8a illustrates that the oracle data influence is not sensitive to the batch size as long as they correspond to the same model state. We calculate batch-level data influences, where one

1 1020

40

80

Number of Training Data (k)

0.55

0.60

0.65

0.70

Val. Spearman

Init from Pre-trained

(a) Oracle label amount.

5 10

50100 300

#Parameters (M) of BERT

0.55

0.60

0.65

0.70

Validation

44.5

45.0

45.5

46.0

Avg. Downstream Acc

Downstream

(b) Model scale.

1

3

5

7

9

11

13

15

17

19

Training FLOPs (1e19)

43.0

44.0

45.0

45.8

46.0

Avg. Downstream Acc

MATES

Full

(a) 410M w/ FLOPs.

10k

30k

50k

70k

90k

110k

130k

150k

Training Steps

43.0

44.0

45.0

45.8

46.0

Avg. Downstream Acc

MATES

Full

(b) 410M w/ steps.

3

8

13

18

23

28

33

38

43

48

53

58

Training FLOPs (1e19)

45.0

46.0

47.0

47.5

48.0

Avg. Downstream Acc

MATES

Full

(c) 1B w/ FLOPs.

10k

30k

50k

70k

90k

110k

130k

150k

Training Steps

45.0

46.0

47.0

47.5

48.0

Avg. Downstream Acc

MATES

Full

(d) 1B w/ steps.

data point can appear on multiple batches so that separating individual data influence from batch influences is a regression problem. Formally, the input of the regressor is a m -dimensional binary vector v : [ v 1 , v 2 , ..., v m ] vector, where m means the size of probing data pool (80k in our experiment) and v i = 1 / v i = 0 denotes the inclusion/exclusion of the i -th data in the current batch; the output is the oracle data influence score probed by one-batch one-step training. In our experiment, we first collect 80k input-output pairs with batch size equal to 16 to train a LASSO regressor and then acquire a m -dimensional individual data influence vector from the regressor's prediction. This vector represents the influence score shown in the y-axis in Figure 8a. The idea of adopting LASSO regression is inspired by the linear datamodels in DsDm [15] that separate individual data influence from the collection of batch data influences.

We also investigate the traits of our reference task, LAMBADA, when collecting oracle data influences. Specifically, we collect the task accuracy at each model checkpoint and measure the Spearman correlation between a single task and average downstream accuracy. As shown in Figure 8b, LAMBADA has a significantly positive Spearman correlation (0.877) with average downstream accuracy. In contrast, the correlation between ARCE/PIQA and average downstream accuracy is not high in Figure 8c/8d, although they are parts of the evaluation tasks. These results imply that ARC-E/PIQA may not reflect the target performance as accurately as LAMBADA when acting as the reference.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The abstract and introduction clearly state the contributions, such as the introduction of a model-aware data selection framework, the use of a small data influence model, and the empirical validation of the proposed method. These claims are consistent with the experimental results and discussions provided in the paper.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The paper discusses several limitations, such as the computational constraints that limit the model size to 1B parameters and the need for more advanced data sampling strategies.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The paper provides a detailed and complete theoretical proof of oracle data influence proposed in the method.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The paper includes detailed descriptions of the experimental setup, including the datasets used, the models employed, and the training procedures. It also releases code, data, and model checkpoints to facilitate reproducibility.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The paper releases code, data, and model checkpoints, providing sufficient instructions to reproduce the main experimental results.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/public/ guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https://nips.cc/public/ guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The paper specifies details such as the datasets, model architectures, hyperparameters, and the training procedures used in the experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: The paper includes standard errors in the experimental results, ensuring that the reported performance gains are robust and reliable.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: The paper provides details on the type of compute resources used, including the specific hardware, memory, and the approximate time required for execution of the experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research conducted in the paper conforms, in every respect, with the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: The paper discusses both positive and negative societal impacts, including the importance of high-quality data in preventing bias and ensuring the responsible use of language models.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: The paper credits the original creators of datasets and models used in the experiments and provides references to the licenses and terms of use.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: The paper introduces new assets, such as the data influence model, and provides documentation alongside them.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing or research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing or research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces MATES, a method termed \"model-aware data selection with data influence models\". MATES is designed to optimize data selection for large language model (LLM) pre-training efficiently. The method dynamically considers various influence models during different stages of pre-training, tailored to the current state of the training model. Experiments validating this approach have been conducted using the C4 dataset with Pythia models.\n\nSTRENGTHS:\n- This paper reveals an important observation that the influence of data changes throughout the training process. It underscores the necessity for model-aware data selection methods to adapt accordingly.\n- The utilization of the Gumbel-Top-k algorithm is particularly interesting and well-suited for balancing data quality and diversity. This diversity is crucial as it helps mitigate the independent influence assumption of the training data.\n- Extensive experiments have been conducted to validate the method's effectiveness.\n\nWEAKNESSES:\n- **Gap between Motivation and Experiment Design**: While the authors have compared MATES with multiple baselines concerning data selection, there remains a gap regarding the necessity of data selection for pretraining as opposed to full training or no continuous pretraining.\n  - **Full Training**: The results and computational costs of pretraining the entire C4 dataset have not been demonstrated. If the performance significantly drops or the computational cost reduction is minimal, then data selection may not be necessary initially.\n  - **No Continuous Pretraining**: While MATES focuses on pretraining, the chosen model, Pythia, is already pretrained on the Pile dataset (i.e., pretraining a pretrained model). If Pythia already performs adequately on those tasks, then data selection might not be necessary initially.\n\n- **Significance of the Results**: The results presented in Table 1 are not significantly robust. For instance, for Pythia-1B, only 5 out of 18 tasks show improvements greater than one **standard error**, suggesting that the p-values are likely higher than 0.05. In some recent similar works [1], the improvement over random selection can be up to seven times the **standard deviation**. Given that MATES also adds computational costs for influence computation, the significance of these results is arguable.\n\n- **Missing Implementation Details**: Some critical details about the implementation are not sufficiently detailed, making the presentation somehow confusing. See the Question section.\n\n[1] LESS: Selecting Influential Data for Targeted Instruction Tuning, ICML 2024.\n\nQUESTIONS:\n- For results presented in Table 1 and Figure 3, what is the amount of sampled data used for computing $I_m$? How to make the computation of $I_m$ efficient, considering that for each sampled data, one pass of the Adam and recomputation of the reference task loss are required? \n- In relation to Figure 5(a), how is LASSO regression employed to \"separate\" influence scores? Does this imply that data points within the same batch can have varying influences?\n- For Figures 4 and 6, what is implied by \"Val. Spearman\"? Was the influence computed for the entire dataset?\n\nLIMITATIONS:\nBesides the limitations discussed by the authors themselves, an additional limitation of MATES is that it uses a single pass of the Adam to approximate the change in loss. This approach assumes that each data point will be trained exactly once. While this assumption might be valid for certain pretraining scenarios, it does not hold for fine-tuning, where all data are used multiple times.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces model-aware data selection with data influence models (MATES) that selects high-quality data for pre-training large language models (LLMs). MATES addresses the issue that existing data selection methods don’t adapt to evolving data preferences during pre-training. MATES continuously adapts to the main model’s evolving data preferences by uses a small data influence model to select the most effective pre-training data for the next stage. The data influence model also reduces the cost of computing the influence scores on-the-fly. Experiments on Pythia and the C4 dataset show that MATES outperforms random data selection and existing static data selection approaches. Further analysis validates the ever-changing data preferences of the pre-training models and the effectiveness of our data influence models to capture it.\n\nSTRENGTHS:\n1. The proposed method outperforms existing static data selection baselines on the Pythia-410 model. Compared to random selection, it can achieve a certain accuracy with less GPU FLOPS.\n2. The authors demonstrate the importance of adapting to fluctuating data influence in data selection methods.\n\nWEAKNESSES:\n1. The authors only compare MATES to other static baseline methods on the Pythia-410M model. For most of the rest experiments, the authors only compare to random selection. \n2. Missing references. There are existing dynamic data pruning methods [1][2] that the authors may want to include in the discussion of related works. These works also noticed the influence of data points changes as the training proceeds.\n\n[1] Raju, Ravi S., Kyle Daruwalla, and Mikko Lipasti. \"Accelerating deep learning with dynamic data pruning.\" arXiv preprint arXiv:2111.12621 (2021).\n[2] Qin, Ziheng, et al. \"InfoBatch: Lossless Training Speed Up by Unbiased Dynamic Data Pruning.\" The Twelfth International Conference on Learning Representations.\n\nQUESTIONS:\n1. The authors only compare against a full set of baseline methods in the Pythia-410M experiments. In the Pythia-1B experiments, the authors only compare against the random selection baseline. Why are the results of other baselines omitted in the Pythia-1B experiment? How do other baseline methods such as SemDedup perform in the Pythia 1B experiments? \n\n2. How does the proposed method MATES perform compared to using the full dataset for training?\n\nLIMITATIONS:\nThe authors adequately discuss the limitations of the work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIt is important to carefully select data during the pretraining stage, as the pretraining data are often obtained from web crawling and can be extensive and noisy. Existing methods for data selection include heuristic-based approaches, clustering-based techniques, and the use of influence models. However, these methods result in static selected datasets and do not adapt to the evolving training process. This paper introduces a novel approach by developing an influence model that co-evolves with the main pretraining model. The selected data is no longer static; now, it is dynamic and sensitive to the model. The study demonstrates that continuously adapting data preferences enables efficient training, and the influence model is both small and accurate. Empirical evidence presented in the framework, MATES, proves its better performance compared to other methods, showcasing advantages in both performance and efficiency.\n\nSTRENGTHS:\n* Originality: This paper introduces a new data selection framework in the pretraining stage by learning a smaller data influence model and co-evolving with the main model's data preference.\n* Quality: Empirically, their results show promise to enable efficient and improved pretraining.\n* Clarity: Their writing and their result are clear to me.\n* Significance: Data curation is crucial to denoise massive pretraining datasets. This paper presents a new approach to selecting data points by learning model preferences. Their empirical results support their claim, and I believe their notion is significant to the field.\n\nWEAKNESSES:\n1. In Figure 2, it appears that the data influence model needs to be reinitialized at each iteration. It's not clear why we need to reset it instead of continuing to learn with the main model.\n2. It's unclear how the influence model includes and excludes data points and converts them into labels for learning the influence score. Could the authors clarify this? Additionally, what is the performance of the influence model across training epochs and how do we evaluate it? I believe there is an imbalance in the labels.\n3. Although the learning curve indicates efficient and improved performance, it's uncertain how much computational burden is involved in maintaining the trained influence model. It would be helpful to include these details.\n\nQUESTIONS:\n1. It's an interesting paper, and I am curious about the changes in data preference. Can the authors provide some examples of what the model likes to learn at the beginning but discards later, or vice versa?\n2. This is not critical: but looking forward to seeing this approach being employed in other data curation benchmarks, such as DataComp.\n\nLIMITATIONS:\nI didn't see any potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a new method “MATES,” which aims to select pretraining data using a reference high-quality data source. MATES uses an estimated influence function to iteratively select the most influential datapoints. Experiments on the Pythia model + dataset show good promise over existing data curation methodologies.\n\nSTRENGTHS:\n- Better trained models on MATES’ curated data vs. other data curation approaches\n- The distillation of influence scores into a smaller proxy model (BERT-base in the paper) allows MATES to perform iterative, model-state dependent data selection vs. other baselines which perform one-step, model-state-independent data selection.\n- Thorough experiments: various experiments and ablations are included, e.g., various relevant baseline comparisons, ablating the effect of various hyper-parameters in MATES, etc.\n\nWEAKNESSES:\n- I found the paper writing to be quite loose and missing context in various places. I had to read a lot of the text multiple times as well as filling in the missing context myself.\n- While I appreciate the extensive experiments covered in the paper, I feel the paper misses a few important results and/or ablations. See the questions section for more on this.\n- The choice of evaluation benchmarks doesn’t seem obvious and feels quite debatable to me. See the questions section for more on this.\n\nQUESTIONS:\nWhile I completely understand the cost of experiments included in the paper, I believe a few critical questions/extra experiments might greatly improve the paper. Please feel free to ignore any of the following questions in the interest of compute.\n\n\n- I couldn’t find details on which dataset was used. While the abstract mentions C4, the main text doesn’t specify this, making me question if the Pythia dataset was used instead?\n- The above question also makes way for a deeper question: since the selection ratio is fixed to 20% (of the entire dataset, I assume), this translates to how many *tokens sampled*? Are the training runs multi-epoch or is the sampled data greater than the training budget?\n- If the runs are multi-epoch, what are the results in the full-data (no sampling) case?\n- How does MATES compare to other approaches when the sampling rate is not equal to 20%? Is MATES also useful in the super low/high-sampling rate regime?\n- What if a different target dataset is used vs. only using lambada? Can we use a mixture of target datasets? How does the design of the target dataset/mixture affect different evals?\n- Finally, the evaluation benchmarks used in the paper seem to be cherry-picked based on the “high-level” proximity to the kind of data present in lambada. The paper would greatly benefit from using more diverse evaluation scenarios like coding, math, knowledge, etc.\n\nLIMITATIONS:\nThe authors address the major technical limitations of their work in Section 6. No obvious negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-16"}