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.
Paper
Similar papers
Peer review
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.
Strengths
- 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. - 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. - Extensive experiments have been conducted to validate the method's effectiveness.
Weaknesses
- **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
- 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?
Rating
5
Confidence
4
Soundness
2
Presentation
2
Contribution
3
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.
Code to compute our reported standard error
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.
Thank you for your response
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).
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.
Strengths
1. 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. 2. The authors demonstrate the importance of adapting to fluctuating data influence in data selection methods.
Weaknesses
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
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?
Rating
5
Confidence
3
Soundness
2
Presentation
3
Contribution
3
Limitations
The authors adequately discuss the limitations of the work.
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 appreciated the authors' response. I will increase the score.
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.
Strengths
* 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. * Quality: Empirically, their results show promise to enable efficient and improved pretraining. * Clarity: Their writing and their result are clear to me. * 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.
Weaknesses
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
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.
Rating
5
Confidence
5
Soundness
3
Presentation
2
Contribution
3
Limitations
I didn't see any potential negative societal impact of their work.
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.
Strengths
- Better trained models on MATES’ curated data vs. other data curation approaches - 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. - Thorough experiments: various experiments and ablations are included, e.g., various relevant baseline comparisons, ablating the effect of various hyper-parameters in MATES, etc.
Weaknesses
- 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
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.
Rating
6
Confidence
5
Soundness
3
Presentation
2
Contribution
3
Limitations
The authors address the major technical limitations of their work in Section 6. No obvious negative societal impact.
Looking Forward to Your Reply
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
TL;DR of the setup clarification
| 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 |
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.
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.
Thank you for your rebuttal
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.
Decision
Accept (poster)