GREATS: Online Selection of High-Quality Data for LLM Training in Every Iteration

We thank the reviewer for the positive feedback!

Q. [Additional Baseline Comparison.] "Why are uniform selection and selection-via-Proxy not included as baselines?"

A. We thank the reviewer for bringing up potentially additional baselines. If we understand correctly, the “uniform selection” baseline corresponds to the “Regular” curve in our current experiments, where the training batches are sampled uniformly at random. For “Selection-via-proxy”, our understanding is that it is a general paradigm for accelerating offline data selection. We'd like to clarify that offline data selection techniques are not competitors to online batch selection methods like GREATS. Instead, these approaches are complementary and can be used in conjunction in the training pipeline. Offline methods can provide an initial high-quality dataset, while online methods can further optimize the training process by selecting the most informative batches during runtime. This potential synergy between offline and online data selection is an interesting direction for future research.

As suggested by the reviewers, we have expanded our comparisons to include reference model-based online batch selection techniques [1] and embedding similarity-based batch selection. Our results show that GREATS consistently outperforms these baselines across various hyperparameter settings, further validating its effectiveness. See Q1 in global response for a detailed discussion.

[1] Mindermann, Sören, et al. "Prioritized training on points that are learnable, worth learning, and not yet learnt." ICML 2022.

Q. [Error bar for the experiment results?]

A. We sincerely appreciate the reviewer's insightful suggestion. We have updated the experiment results for key benchmarks with error bars across 3 independent runs.

These updated results with error bars further solidify our findings. GREATS consistently outperforms all baselines across these key benchmarks, even when accounting for run-to-run variability. We've also added error bars to the loss curves for these experiments (see Q1 in global response) which show similar results. Due to the computational intensity nature of some experiments, we are in the process of extending this multiple-run analysis to the Alpaca-Samsum and OpenWebText (pretraining) settings. We're committed to providing a complete picture of GREATS' performance across all benchmarks in the revision.

Q. [Accuracy results for other datasets?]

A. We thank the reviewer for the valuable comment!

OpenWebText Pretraining: For our GPT2 pretraining on OpenWebText, we primarily focused on test loss as the most appropriate metric. This choice was made due to the nature of pretraining (where the model is not instruction-tuned) and our computational constraints. Test loss provides a direct measure of the model's ability to predict the next token, which is the core objective of language model pretraining.

Alpaca-Samsum: Since Samsum is a summarization task, we evaluated the model's performance using the ROUGE-1 score, which is a standard metric for summarization quality. Here are the additional experiment results on Samsum's test set:

As we can see, GREATS outperforms all baselines on this metric, further demonstrating its effectiveness across diverse task types.

We thank the reviewer for the positive comments!

Q. [Additional Baseline Comparison.]

A. We sincerely appreciate the reviewer's feedback and suggestions for additional baselines. We have taken these recommendations seriously and have expanded our comparisons accordingly. (1) Reference model-based online batch selection (RHO Loss) [1]: We have implemented and evaluated the RHO Loss method as suggested. (2) Similarity to validation data: Regarding the suggestion of "a classifier to detect similarity to validation data," we note that our validation sets are very small, which may not be sufficient to train a robust classifier. Instead, we implemented a method that selects the training batch based on embedding similarity to the validation data, using Sentence-BERT embeddings. This approach serves as a proxy for the suggested classifier-based method while being applicable to our small validation set scenario.

Our expanded experiments show that GREATS consistently outperforms these additional baselines across various hyperparameter settings, further validating its effectiveness. We have included a detailed discussion of these results in Q1 in global response.

[1] Mindermann, Sören, et al. "Prioritized training on points that are learnable, worth learning, and not yet learnt." ICML 2022.

Q. [Relationship to the influence function.]

A. We appreciate the reviewer's insightful comment regarding the relationship between our proposed method and influence functions. Influence functions typically assume strong convexity of the loss landscape and convergence to a local optimum. In contrast, our method makes no such assumptions, making it more suitable for the highly non-convex landscapes encountered in large language model training. While influence functions focus on the impact of training examples on the final model performance, our method explicitly considers the immediate impact on the next training step. This allows us to capture the dynamic nature of the training process and adapt to the model's current state. We will incorporate this discussion into the related work section.

Q. [Runtime complexity of GREATS?]

A. We appreciate the valuable suggestion! Here’s the complexity analysis for GREATS: Assuming that the batch size is $B$ and the number of model parameters is $p$, the complexity of backpropagation is $O(Bp)$. The additional cost of the "ghost" dot-product involves computing the pairwise dot-product between each data point’s activation (or output derivative), which are of dimension $O(\sqrt{p})$, and then taking the scalar product between the inner product results on activations and output derivatives. The step of computing the pairwise activation (or output derivative) dot-products has a runtime complexity of $O(B^2 \sqrt{p})$, resulting in a $B \times B$ matrix. Taking the element-wise product between the two $B \times B$ matrices has a runtime complexity of $O(B^2)$. Hence, the overall runtime complexity is $O(Bp) + O(B^2 \sqrt{p})$. In practice, since the batch size $B$ is significantly smaller than the number of parameters $p$, the additional runtime $O(B^2 \sqrt{p})$ is negligible compared to the backpropagation complexity $O(Bp)$. This analysis explains why our empirical runtime measurements (in seconds) show that GREATS performs comparably to regular training. We will incorporate this detailed complexity analysis into our paper to provide a more comprehensive understanding of the algorithm's efficiency. Thank you for highlighting this important aspect!

Q. “A minor point is that the argument in lines 28-29 is tenuous at best ("Moreover, online batch selection operates on smaller batches of data, reducing the need for cumbersome data preprocessing and enabling more efficient use of computational resources compared to static data selection methods that process the entire dataset upfront".) It is much more convenient to pre-cache the data when training these models. At training time, one really just wants to use accelerators as efficiently as possible.”

A. We appreciate the reviewer’s valuable comments. We agree with the reviewer and acknowledge that offline data selection provides the convenience for not modifying training pipeline. We will include this point and modify the Introduction accordingly. On the other hand, online batch selection offers unique advantages in terms of adaptability to the model's learning progress. This can be particularly valuable in scenarios where the relevance of data may change as training progresses. Furthermore, offline and online selection methods are not mutually exclusive and can be complementary. A hybrid approach, leveraging both pre-cached, pre-selected data and dynamic online selection, could potentially offer the best of both worlds.

Q. [Hessian approximation used in the experiments?]

A. In all of our experiments, we use identity matrix as the approximation to the Hessian matrix, and we found the performance is already good. Future work can further explore the use of more advanced techniques for Hessian approximation, such as K-FAC [1].

[1] Martens, James, and Roger Grosse. "Optimizing neural networks with kronecker-factored approximate curvature." ICML 2015.

Q. [How to choose the validation data?]

A. In the standard benchmarks such as MMLU, the validation data is being provided for each subject of size 5. For the benchmarks where the validation set size is large, we sample the validation set used in the experiments uniformly at random. We will incorporate this additional detail into our experiment setting section.

Q. [Notation]

A. We thank the reviewer for the suggestions on the notations, which we will incorporate into our revision. We note that the utility function $U^{(t)}: R^d \rightarrow R$ maps to the model performance change in $t$th step instead of the optimal batch.

Q. [Is the proposed utility function close to submodular?]

A. We sincerely appreciate the reviewer's insightful question. We conducted additional experiments and in Global response’s Figure 5 (c), we plot how the loss change in a single gradient update step, i.e., our utility function $U^{(t)}(S)$, against varying input batch sizes $|S|$. We observed that as batch sizes increase, the average loss change (computed across randomly sampled subsets of fixed sizes) indeed exhibits a "diminishing returns" trend. This behavior is consistent with the general properties of submodular functions. However, we also noted significant variance in the utilities, especially for smaller batch sizes. This high variance suggests that strict submodularity may not hold. The observed "diminishing returns" trend aligns with the intuition behind submodular functions and provides some justification for the effectiveness of our greedy selection approach. While our utility function may not be strictly submodular, we conjecture that it could potentially fall into the category of weakly submodular functions.

We will incorporate this analysis into our paper, which might open up interesting theoretical directions for future work.

Q. [Discussion on submodular subset selection techniques]

A. We sincerely appreciate the reviewer for bringing our attention to the relevant literature on submodular subset selection techniques. To the best of our knowledge, gradient coreset-based online data selection algorithms, including those utilizing submodular optimization (like [1]), generally face scalability challenges when applied to the setting of foundation models. The primary limitation stems from their requirement to compute (or even store) per-sample gradient vectors, which becomes prohibitively expensive in terms of both computation (or memory) for large-scale models with millions or billions of parameters. For instance, if we understand correctly, to optimize the resulting error in [1], the computation of $Err$ requires computing per-sample gradient vectors. An interesting future avenue is to explore whether our "ghost inner product" technique can be further extended to address the computational problems in this line of work. We will incorporate these references and discussion into our related work section in the revision.

Additional baselines: As suggested by other reviewers, we have expanded our comparisons to include reference model-based online batch selection techniques [2] and embedding similarity-based batch selection. Our results show that GREATS consistently outperforms these baselines across various hyperparameter settings, further validating its effectiveness. See Q1 in global response for a detailed discussion.

[1] Killamsetty, Krishnateja, et al. "Grad-match: Gradient matching based data subset selection for efficient deep model training." ICML 2021

[2] Mindermann, Sören, et al. "Prioritized training on points that are learnable, worth learning, and not yet learnt." ICML 2022.

We thank the reviewer for the positive comments!

Q. [Relationship between online batch selection and offline data selection / data mixture optimization techniques [1, 2]? & Additional baseline comparison]

A. We appreciate the reviewer highlighting the relevant literature on offline data selection [1] and data mixture optimization [2]. We'd like to clarify that offline data selection techniques are not direct competitors to online batch selection methods like GREATS. Instead, these approaches are complementary and can be used in conjunction in the training pipeline. Offline methods can provide an initial high-quality dataset, while online methods can further optimize the training process by selecting the most informative batches during runtime. This potential synergy between offline and online data selection is an interesting direction for future research.

Additional baselines: As suggested by other reviewers, we have expanded our comparisons to include reference model-based online batch selection techniques [3] and embedding similarity-based batch selection. Our results show that GREATS consistently outperforms these baselines across various hyperparameter settings, further validating its effectiveness. See Q1 in global response for a detailed discussion.

[1] Xie, Sang Michael, et al. "Data selection for language models via importance resampling." NeurIPS 2023.

[2] Xie, Sang Michael, et al. "Doremi: Optimizing data mixtures speeds up language model pretraining." NeurIPS 2024.

[3] Mindermann, Sören, et al. "Prioritized training on points that are learnable, worth learning, and not yet learnt." ICML 2022.

Q. [Experiment with different hyperparameters]

A. As suggested, we have performed experiments with different choices of learning rates and batch sizes. Our results show that GREATS consistently performs well across all the hyperparameter configurations we considered. This robustness suggests that GREATS can maintain its effectiveness without requiring extensive hyperparameter tuning. An ablation study on the number of validation points in Figure 3 (a)-(b) of the paper. This study demonstrates that GREATS can perform effectively even with a small number of validation points (2 to 5), which is a crucial aspect of its practical applicability.

Q. [Potential distribution shift between validation and test distribution?]

A. We appreciate the reviewer's insightful question. Subsampling validation data is one solution here. This subsampling approach effectively reduces overfitting to the validation set and introduces some variability in the selection process, which can help mitigate the impact of potential distribution mismatches.

We acknowledge that the availability of a high-quality validation set is a limitation of GREATS. However, most state-of-the-art data selection (e.g., [1, 2]) and online batch selection techniques (e.g., [3]) rely on a high-quality validation set. This is a common approach in the field to guide the selection process toward data that aligns with the target domain or task. In finetuning tasks, practitioners typically have a small set of high-quality, task-specific data that can serve as a validation set. Our experiments demonstrate that GREATS can perform well even with a very small validation set (as few as 2-5 samples), which makes the requirement more practical to fulfill in many scenarios. We stress that while GREATS does require a validation set, it avoids other limitations of existing methods, such as the need for reference models, making it more practical for large-scale training scenarios.

[1] Xia, Mengzhou, et al. "LESS: Selecting Influential Data for Targeted Instruction Tuning." ICML 2024.

[2] Xie, Sang Michael, et al. "Data selection for language models via importance resampling." NeurIPS 2023.

[3] Mindermann, Sören, et al. "Prioritized training on points that are learnable, worth learning, and not yet learnt." ICML 2022.

Q [Do the experiments use Adam optimizer?]

A. In all of our experiments, we use AdamW as the optimizer for model training. This aligns with common practices in training large language models. We acknowledge that our theory for GREATS is derived in terms of SGD. While the theory can be extended to Adam, the efficient "ghost" inner-product technique we developed is specifically tailored for SGD and is challenging to directly adapt to Adam due to its more complex update rule and maintained moments. Hence, we adopt a hybrid approach where we select the valuable batch based on the theory derived from SGD, but perform the actual training updates using AdamW. This approach allows us to benefit from the computational efficiency of our SGD-based selection method while still leveraging the optimization benefits of AdamW for training. Despite this theoretical mismatch, our experiments demonstrate that this hybrid approach performs well in practice. The consistent improvements over baselines across various models and tasks suggest that the SGD-based utility function provides a good proxy for data importance, even when Adam is used for training. We recognize the potential for further improvement by aligning the batch selection more closely with Adam-style updates.

We thank the reviewer for the nice words!

Q [Requirement of validation set]

A. We acknowledge that the availability of a high-quality validation set is a limitation of GREATS. However, most state-of-the-art data selection (e.g., [1, 2]) and online batch selection techniques (e.g., [3]) rely on a high-quality validation set. This is a common approach in the field to guide the selection process toward data that aligns with the target domain or task. In finetuning tasks, practitioners typically have a small set of high-quality, task-specific data that can serve as a validation set. Validation data for pretraining: We agree that for pretraining scenarios, obtaining a representative validation set can be more challenging. However, even in these cases, it's often possible to curate a small set of diverse, high-quality samples that help for data selection [4, 5, 6]. Our experiments demonstrate that GREATS can perform well even with a very small validation set (as few as 2-5 samples), which makes the requirement more practical to fulfill in many scenarios. Furthermore, GREATS selects points based on their potential to improve performance on the validation set instead of simply filtering data points that are similar to the validation data. It's worth noting that while GREATS does require a validation set, it avoids other limitations of existing methods, such as the need for reference models, making it more practical for large-scale training scenarios.

[1] Xia, Mengzhou, et al. "LESS: Selecting Influential Data for Targeted Instruction Tuning." ICML 2024

[2] Xie, Sang Michael, et al. "Data selection for language models via importance resampling." NeurIPS 2023

[3] Mindermann, Sören, et al. "Prioritized training on points that are learnable, worth learning, and not yet learnt." ICML 2022

[4] Tirumala, Kushal, et al. "D4: Improving llm pretraining via document de-duplication and diversification." NeurIPS 2024

[5] Maini, Pratyush, et al. "Rephrasing the Web: A Recipe for Compute and Data-Efficient Language Modeling."

[6] FineWeb: decanting the web for the finest text data at scale.

Q [Approximation accuracy of Taylor expansion] “While the use of Taylor expansions for approximation is efficient, the accuracy of these approximations can be affected by factors like learning rate and non-linearity of the model. Further analysis and discussion on the potential limitations of these approximations would strengthen the paper.”

A. We thank the reviewer for the valuable suggestions. We conducted additional experiments and examined two key correlations in Figure 5 in Global response: (a) The correlation between the actual validation loss change in a single gradient update step and the loss change predicted by the sum of training gradients taking dot product with the validation point (i.e., the first term in Equation (4)). (b) The correlation between the actual validation loss change in a single gradient update step and the loss change predicted by both the gradient dot product sum and the Hessian interaction term, as shown in Equation (4). As we can see, if we do not incorporate the interaction term, the Pearson correlation coefficient is approximately 0.75. Figure (b) shows that the correlation coefficient improved by including the Hessian interaction term to approximately 0.84. These results demonstrate that our approximations capture a significant portion of the actual loss change, with the inclusion of the Hessian interaction term providing a notable improvement in accuracy.

Q [Comparison with reference-model-based methods]

A. We have conducted additional experiments using the RHO-Loss method from [1] as a representative baseline for reference-model-based approaches. Due to time and computational constraints, we used Llama-3.1-8B-Instruct as the reference model (with Llama2-7B as the target model), prioritizing training on data points with high RHO loss ($L(x; f_{tgt}) - L(x; f_{ref})$ where $f_{tgt}$ is the target model and $f_{ref}$ is the reference model). Surprisingly, the RHO-Loss method performed poorly, showing the worst performance across almost all settings. We hypothesize two main reasons here: (1) Noise and scale discrepancy: The loss value from the current target model $L(x; f_{tgt})$ is highly noisy and may be of a different scale compared to the reference model. This discrepancy could result in $L(x; f_{ref})$ having minimal impact on the batch selection results. (2) Reference model limitations: Due to computational and time constraints, we used a relatively small pretrained model as the reference. This may not adequately reflect the "learnability" or "difficulty" of the training examples, which is crucial for effective data selection. These findings highlight a key advantage of GREATS: it offers a more efficient and flexible alternative that can perform well without relying on reference models, which can be computationally expensive or may not be available in practice.

[1] Mindermann, Sören, et al. "Prioritized training on points that are learnable, worth learning, and not yet learnt." ICML 2022.

Q [Experiment with additional baselines & different hyperparameters]

A. As suggested, we have performed experiments with additional baselines and different choices of learning rates / batch sizes. The results are in Global response Q1. The results show that GREATS consistently performs well across all the hyperparameter configurations we considered. This robustness suggests that GREATS can maintain its effectiveness without requiring extensive hyperparameter tuning. An ablation study on the number of validation points in Figure 3 (a)-(b) of the paper. This study demonstrates that GREATS can perform effectively even with a small number of validation points (2 to 5), which is a crucial aspect of its practical applicability.