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

Authors: zichun  yu
spandan  das
chenyan  xiong
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 19:26:55

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.

Main paper: 
 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 Model-Aware data selection with daTa influencE modelS (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 performance 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. Then, a data influence model is trained to approximate data influences on the target performance of the pretraining model and select the most effective data for the next pretraining stage.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.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. 
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. 
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). (x i ; D r ) | x i ∈ D h } Fine-tune data influence model Θ on {(x i , I M (x i ; D r )) | x i ∈ D h } Select data S * k ← {Gumbel-Top-k( Θ(xi) τ ) | x i ∈ D t } end if Sample a batch of data B * from S * k M ← A(M, B * ) end for 3.1 Model-aware data selection frameworkMATES 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 size-n 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 * :EQUATIONEQUATIONwhere 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:EQUATIONwhere 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-Top-k algorithm [30; 62] to sample data from D t , with influence scores as weights:EQUATION)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 size-k subset from D t in the initial U steps. The overall pretraining and data selection pipeline of MATES is illustrated in Algorithm 1.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 asM * ϵ,xi = 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:EQUATIONwhereH 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 ofx i is M * 1 n ,xi -M * ≈ -1 n H -1 M * ∇ M L(x i | M *) and the influence in Eq. 8 can be further represented as:EQUATIONEQUATIONEQUATIONIn 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:EQUATIONThis 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. 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] , ARC-C [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. 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. 
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. 
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.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] , reducing 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. 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.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. 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. 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. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
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] . 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.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. 
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.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. 
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). 
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? 
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). 
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. 
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. 
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. 
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. 
 
 Appendix: 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 xi by ϵ as:EQUATIONand 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 xi in the training data is upweighted by a small ϵ, i.e.:EQUATION)To computedM * ϵ,x i dϵ ϵ=0, we consider the optimality condition at M * ϵ,x i :EQUATIONDifferentiating both sides with respect to ϵ:EQUATIONEQUATIONwhere HM * = 1 With the first-order Taylor expansion, we can approximate the change in model parameters:EQUATIONFor ϵ = 1 n , i.e., incorporating xi into the training data, this becomes:EQUATIONSubstitutingEQUATIONEQUATIONIn 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 HM + λI is positive definite [29] . Finally, the locally probed oracle data influence of xi is:EQUATIONTo ensure that a positive influence score reflects a beneficial impact on model performance, we empirically define the negative influence, L(Dr | M) -L(Dr | A(M, xi)), as our locally probed oracle data influence of xi. We provide all experimental configurations in Table 5 .WSD learning 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:lr(t) =      t W • η, if t < W η, if W ≤ t < S 0.5 4•(t-S)/D • η, if S ≤ t < S + D (26)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.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 wo ∈ R H to get the model prediction wo • h. The training objective is the mean squared error between the model prediction and the normalized ground truth IM.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.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). 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.Table 7 : Ablation study of pretraining 410M models with different update steps U and sampling temperatures τ between 40k to 50k steps. We show zero-shot/two-shot results. This group of experiments conducts ablation studies on the key hyperparameters. UpdateStep 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 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. 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 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 : [v1, v2, ..., vm] vector, where m means the size of probing data pool (80k in our experiment) and vi = 1/vi = 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 ARC-E/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• 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 , 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.• 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.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).Due to the high computational cost of LESS to compute the gradients for all the pretraining examples, we only report its results in the 410M setting. 


Decision: Reviews are somewhat borderline, though certainly leaning on the side of accepting since all lean positive. Reviewers praise the originality of the work, and the findings themselves, as well as aspects of the data/experiments, etc. Arguably somewhat "generic" praise though still uniformly positive. Some issues raised about the writing, evaluation, and various other small things, though these seemed fairly minor.

Reviews: review: paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: - 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.
- 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.
- The choice of evaluation benchmarks doesn’t seem obvious and feels quite debatable to me. See the questions section for more on this.
questions: questions: While 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.


- 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?
- 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?
- If the runs are multi-epoch, what are the results in the full-data (no sampling) case?
- 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?
- 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?
- 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.
limitations: limitations: The authors address the major technical limitations of their work in Section 6. No obvious negative societal impact.
paper_summary: paper_summary: It 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.
strength_weakness: strength_weakness: 1. 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.
2. 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.
3. 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.
questions: questions: 1. 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?
2. This is not critical: but looking forward to seeing this approach being employed in other data curation benchmarks, such as DataComp.
limitations: limitations: I didn't see any potential negative societal impact of their work.
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: 1. 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. 
2. 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.

[1] Raju, Ravi S., Kyle Daruwalla, and Mikko Lipasti. "Accelerating deep learning with dynamic data pruning." arXiv preprint arXiv:2111.12621 (2021).
[2] Qin, Ziheng, et al. "InfoBatch: Lossless Training Speed Up by Unbiased Dynamic Data Pruning." The Twelfth International Conference on Learning Representations.
questions: questions: 1. 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? 

2. How does the proposed method MATES perform compared to using the full dataset for training?
limitations: limitations: The authors adequately discuss the limitations of the work.
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: - **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.
  - **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.
  - **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.

- **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.

- **Missing Implementation Details**: Some critical details about the implementation are not sufficiently detailed, making the presentation somehow confusing. See the Question section.

[1] LESS: Selecting Influential Data for Targeted Instruction Tuning, ICML 2024.
questions: questions: - 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? 
- 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?
- For Figures 4 and 6, what is implied by "Val. Spearman"? Was the influence computed for the entire dataset?
limitations: limitations: Besides 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.

Comments: title: Looking Forward to Your Reply
Thank you for your response
Thank you for your rebuttal
TL;DR of the setup clarification
Code to compute our reported standard error
comment: I appreciated the authors' response. I will increase the score.
Dear Reviewer 2Zbt,

We have carefully addressed your feedback in our rebuttals and provided detailed responses to each of your comments. We believe these clarifications will help in assessing our work more comprehensively.

We would greatly appreciate it if you could review our rebuttals and provide any further feedback, given that the author-reviewer discussion will be closed on Aug. 13 at 11:59 p.m. AoE in no more than one day. We are willing to answer any further questions.

Thank you for your time and consideration. We look forward to your reply.

Best,

The Authors
We appreciate your response and recognition of our work! For Figure 1 in our supplementary pdf, we will change the label name to be “Full”.

For the computational overhead incurred by MATES, we reported the total pretraining FLOPs, including all the cost of maintaining the data influence model, in Table 1 in our original paper. Figure 3 also shows the total FLOPs-Acc curves, and MATES greatly elevates the scaling curves (even with additional data selection cost) compared to random selection. However, we agree that it is a good idea to have the detailed breakdown of the computational overhead incurred by MATES, so we provide it in the table below.

| Pythia-410M                      | #FLOPs (x 1e19) | Ratio of the total #FLOPs |
| -------------------------------- | --------------- | ------------------------- |
| Main model pretraining           | 6.35            | 78.3%                     |
| Oracle data influence collection | 0.30            | 3.7%                      |
| Data influence model training    | 9e-3            | 0.1%                      |
| Data influence model inference   | 1.46            | 18.0%                     |
| Total                            | 8.11            | 100%                      |


| Pythia-1B                        | #FLOPs (x 1e19) | Ratio of the total #FLOPs |
| -------------------------------- | --------------- | ------------------------- |
| Main model pretraining           | 17.67           | 88.5%                     |
| Oracle data influence collection | 0.84            | 4.2%                      |
| Data influence model training    | 9e-3            | 0.05%                     |
| Data influence model inference   | 1.46            | 7.3%                      |
| Total                            | 19.97           | 100%                      |

The data selection cost of MATES only accounts for 21.8% (Pythia-410M) and 11.6% (Pythia-1B) of the total pretraining FLOPs, compared to the state-of-the-art method, QuRating, whose selection cost accounts for 75.9% (Pythia-410M) and 53.1% (Pythia-1B). The selection cost ratio of larger models is generally smaller since their pretraining cost dominates the total FLOPs while the training and the inference cost of our data influence model remains stable. The inference speed can also be improved with a fast-inference framework like vLLM [1]. We will include this new breakdown analysis in the next version of our paper.

For the C4 dataset, we basically followed DsDm’s open-source version [2] to facilitate a fair and reproducible comparison. We also verified the effectiveness of MATES in the recent FineWeb [3] dataset (please refer to our general response 3.3), where MATES still demonstrates superior performance compared to the state-of-the-art data selection methods, such as the FineWeb-Edu classifier [3] and the fasttext-oh-eli5 classifier [4]. We will add this result in the next version.

Thank you once again for your insightful reviews. Please do not hesitate to reach out if any part of our response requires further clarification.

[1]: Kwon, Woosuk, et al. "Efficient memory management for large language model serving with pagedattention." Proceedings of the 29th Symposium on Operating Systems Principles. 2023.

[2]: Engstrom, Logan, Axel Feldmann, and Aleksander Madry. "Dsdm: Model-aware dataset selection with datamodels." arXiv preprint arXiv:2401.12926 (2024).

[3]: Penedo, Guilherme, et al. "The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale." arXiv preprint arXiv:2406.17557 (2024).

[4]: Li, Jeffrey, et al. "DataComp-LM: In search of the next generation of training sets for language models." arXiv preprint arXiv:2406.11794 (2024).
I appreciate the authors' comprehensive explanation and their efforts to address my previous concerns. The experiments conducted for full training are promising and effectively demonstrate the practical application of MATES. However, I suggest that for Figure 1 in the supplementary materials, it would be more accurate to label it as "Full" instead of "Random."

Additionally, I recommend that the authors provide a detailed analysis of the computational overhead associated with data selection, including the training of the data influence model, in comparison to the pre-training costs. This would resolve the concern about the computational costs of this manuscript.

While the approach of pretraining Pythia on C4 from scratch seems somewhat unconventional, I acknowledge that this paper opens up a promising direction for exploring how the influence of data evolves during different phases of pretraining. With the addition of the suggested clarifications, I believe this paper will make a valuable contribution to the community. Accordingly, I am inclined to raise my score to 5.
Thank you very much for your response and recognition of our work! We will clarify the experimental setup and add detailed descriptions about our method and results in the next version.
Thanks for providing additional experiments and clearing my concerns about initialization, evaluation, and computation. Results in your attachment look promising. Thanks for providing examples that evolve with model preference. They are interesting and love to see more follow-up exploration.

Overall, I think this paper is interesting and shows good empirical results, but some places lack detailed descriptions. I will increase my score to 5. Thanks for your response!
Thank you for your detailed review and follow-up questions. We believe the reduced margin on 1B is because our selection ratio and candidate data pool are not optimal for the 1B setting. We don't have enough compute resources to intensively tune these hyperparameters since our work is at an exploratory scale, as mentioned in our limitations section.

For the selection ratio, we **fix it to 20%** in our main table. However, if we vary it from **20% to 10%** for 1B at the 50k decay stage, we can observe some gains (also shown in Figure 3a in our supplementary pdf):

|      | SciQ     | ARC-E    | ARC-C    | LogiQA   | OBQA     | BoolQ    | HellaSwag | PIQA     | WinoGrande | Average  |
| :--- | :------- | :------- | :------- | :------- | :------- | :------- | :-------- | :------- | :--------- | :------- |
| 10%  | **67.8** | 44.4     | 25.5     | **28.9** | **32.6** | **60.9** | **47.4**  | **70.5** | **52.4**   | **47.8** |
| 20%  | 67.3     | **44.9** | **25.9** | 28.7     | 32.2     | **60.9** | 45.3      | 69.5     | **52.4**   | 47.5     |

With the 10% ratio, MATES (Pythia-1B) can have a 6/0/4 #Win/#Tie/#Loss over the best baselines. This indicates that we still have headroom to improve MATES with a more suitable selection ratio in the full training run (not only at the decay stage due to the rebuttal time limit).

For the selection pool, we **fix it to 125B tokens for all models**, but we believe a larger pool will generally benefit 1B models. Recent DataComp-LM [1] also shows that the improvement trends observed on a smaller scale (410M) align with the larger scale (1B and 7B) **when the larger models have a larger selection pool (e.g., 15T tokens at most)**. We plan to verify the effectiveness of MATES on this standard DataComp-LM benchmark and will update the results with a more suitable ratio and pool in the next version.

Additionally, for your reference, we report the #Win/#Tie/#Loss of all our baseline methods in their original papers:

| Method   | Baselines                          | #Win/#Tie/#Loss over the best baseline as reported in their papers |
| -------- | ---------------------------------- | ------------------------------------------------------------ |
| DSIR     | Random, Manual, Heuristics         | 4/0/5                                                        |
| SemDeDup | Random, NearDup                    | 2/0/1                                                        |
| DsDm     | Random, Classifier, DSIR, SemDeDup | 5/1/9                                                        |
| QuRating | Random, DSIR                       | 5/0/6; If we only compare MATES with Random and DSIR, the #Win/#Tie/#Loss is 8/1/1 |

Note that all these baseline papers have already been accepted at the top-tier conferences. This underscores the point that it is reasonable for an effective method to not outperform the best baseline in every task but rather to achieve a solid average performance gain across tasks. We recognize that each method has its unique strengths. For instance, QuRating is specifically optimized for educational data and thus is likely to significantly elevate the model performance on the knowledge QA tasks. On the other hand, MATES stands out due to achieving the highest overall performance while utilizing few additional FLOPs. 

Thank you once again for raising this question. Please do not hesitate to reach out if any part of our response requires further clarification.

[1]: Li, Jeffrey, et al. "DataComp-LM: In search of the next generation of training sets for language models." arXiv preprint arXiv:2406.11794 (2024).
I appreciate the authors for conducting additional experiments and providing the response. The response answers my question about missing references and full dataset training. However, according to the additional 1B zero-shot results provided in Table 1 of the newly uploaded pdf, the proposed MATES wins the best baselines on 3 out of the 10 datasets, draws with the best baselines on 2 out of the 10 datasets, and loses to the best baselines on 5 out of the 10 datasets. This seems to be worse compared to the 410M zero-shot results (wins on 7 / 10, loses on 3 / 10). Does this suggest the advantage of MATES degrades as the model size increases? I will adjust my score if this question is clear.
| Name                                   | Value                                                        |
| -------------------------------------- | ------------------------------------------------------------ |
| Pretraining dataset                    | C4, trained from scratch                                     |
| Amount of collected oracle data        | 80k at 10k steps and 20k at 20k, 30k, 40k, 50k steps         |
| Initialization of data influence model | Initialized with pretrained BERT at 10k steps and continuously fine-tuned at 20k, 30k, 40k, 50k steps |
| Validation set of data influence model | 10% sampled from the collected oracle data                   |
For your reference (Weakness 2), we provide the relevant code block to compute `stderr` from `lm-evaluation-harness`. 

```python
def sample_stddev(arr):

  mu = mean(arr)

  return math.sqrt(sum([(x - mu) ** 2 for x in arr]) / (len(arr) - 1))

def mean_stderr(arr):

  return sample_stddev(arr) / math.sqrt(len(arr))
```

The `arr` in the code is passed with the binary accuracy list across all the examples, where 1 denotes model answers correctly and 0 denotes model answers wrong. The return value of `mean_stderr` only shows the variability of the model’s accuracy across all the examples, not across different random seeds as LESS did.