GIO: Gradient Information Optimization for Training Dataset Selection

We thank reviewer NSKq for their insightful feedback. We are pleased that they find the theoretical concept underpinning the method interesting and appreciate the algorithm’s competitive experimental results. We outline our response to your questions below:

W1/Q1 You are right that the greedy hill-climb method does not correspond exactly to the overall optimal solution. In the method section under Limitations, we note that the space is non-convex and methods like hill-climbing (gradient descent) will likely find local but not globally optimal minima. We propose alternative methods more suited to non-convex spaces like particle swarm optimization and look forward to exploring those techniques and seeing the impact on experimental results in future work. However, we acknowledge that the Limitations paragraph is phrased more as a limitation on gradient descent itself, and not a limitation between equations (1) and the hill-climbing concept of equation (2). Making that gap the focus of the Limitations paragraph rather than gradient descent specifically would make the insight and dicussion more general, and we will update this shortly in our revision. Along similar lines with your first question, submodularity unfortunately does not hold for KL divergence so we can’t bound the approximation rate directly. The closest related submodular function and technique we could discover is submodular optimization using submodular mutual information, which we evaluate (the Submodular Optimization baseline in our experiments). We find our method outperforms that technique with the submodular function, which implies our approximation rate must be less than that of the submodular function (submodular mutual information). We will add these details on this in the revision shortly. This is an extremely interesting insight and we can’t thank you enough for pointing this out and helping us improve our presentation and insight!

W2 We are unsure about your point; the concept of selecting a subset that matches a target set is commonly adopted as a data selection problem, see Yao et al (BM25 baseline), Kaushal et al., 2022 (Submodular optimization), Xie et al, Brown et al (GPT large language models) etc. Our setup using G (general data pool), X (a desired data distribution to follow) and D (data we must include) allows to flexibly represent many data selection problems. Our method is then based on incorporating all three variables. The data selection problem as you mention is to find the best subset that represents the full set, which as you correctly point can be initialized in our setup as $D=\emptyset$, $X=G$ and the algorithm can simply be applied accordingly. We investigate the performance of our model under this setup in the FashionMNIST experiment with good results, surpassing the baseline.

Q2 Absolutely, in response to your’s and another reviewers feedback we will incorporate details and guidance on how to choose an X (and D) in real world setups. Selecting X is specific to the goal of what you are data-selecting for; for example, WMT provides many dev sets of high quality which are close to the desired task, so we were confident those were a good target (which we collected from 08 to 13 to have more data). The WMT14 training dataset comes from a variety of sources (e.g. commoncrawl) of variable quality, so we felt finding the best subset that represents this full set (X=G) would not yield the most gains, when instead we have much higher quality dev sets available to be our target. In the real world, for foundation models for example which use web-scraped data as training data, choosing X to be the Wikipedia corpus (as the authors of GPT do when data-selecting) makes the most sense as this represents high quality language. Or in industry, for example, choosing real customer traffic your model may see in production makes the most sense as this will get the training set closest to the real production application of the model. We feel you have brought up an important discussion on how to practically apply this algorithm to new settings; we will provide the above examples and some further guidance on how to select X and update the revision shortly. Thank you!

Comments 1.$\Omega$ is the measure space on which the sets of data and probability densities are defined. 2. Thank you for pointing this out, we will fix it in our revision. 3. You are absolutely correct, the relaxation is $\mathbf{v} \in G$ to $\mathbf{v} \in \Omega$, $\mathbf{v_{opt}}$ is also in $\Omega$ (and then we find the closest to that in G).

We thank reviewer Km99 for their insightful review and feedback and are thrilled they find our method simple, well-motivated and theoretically grounded and demonstrates strong empirical results in multiple domains. We agree that our method is highly relevant to the current research trend in pretrained models, and look forward to our method being used to more efficiently train models on smaller data while achieving similar or superior performance.

W1 GIO is flexible enough to have many stopping criteria depending on what the goal of data selection is (just match the target with unique samples, select only up to a data size, etc) which we explore in Appendix A. For Section 4.3 Spelling Correction experiments specifically, we wanted to experiment with a different stopping criterion, and previous expertise in that particular domain implied that selecting good duplicates (i.e. implicitly weighting some data more) improves model performance. However, we acknowledge we did not discuss this decision in depth and thank the reviewer for pointing this out. We will add more background to this decision in the revision shortly.

W2 We thank you for your suggestion on the phrasing of Section 4.3; you are absolutely correct, the experiment is selecting a set of data from a pool of mixed high and low quality data with reference to an X that is real human mistakes (very high quality), and shows that given an X of high quality data, GIO identifies and selects high quality data (73%) from the pool of mixed quality data. We acknowledge that the wording confuses readers as to which set is X and G, and will rephrase this section in our revision shortly.

Q1 The goal of the FashionMNIST experiment was primarily to show that GIO works well on other domains besides text, and also to show another use case of reducing a dataset size using GIO (X=G) under a data constraint, rather than picking from data to conform to some other target X. The experiment showed GIO was able to do so effectively; further analysis in Appendix A has some fascinating clues into why this works. Within a class most images are similar (e.g. most bags are rectangular in the “bag” class), but there are outliers and some may even look like other classes. We found that random selection selected both the “core” of the classes (e.g. rectangular bags) but also included outliers, whereas our method focused almost exclusively on choosing the “core'” of the classes first and therefore had fewer outliers. Under the 25% data constraint, it makes sense that we should focus on getting the “core” images correct and not waste data on outliers we are unlikely to get right anyway, which is why our method's data led to better performance. We showed this phenomenon visually with a tSNE and an analysis of the average distance from one class to another, and confirmed that our method selects data that is clustered closer to together (the “core” data) and more separable from other classes. We wish we could have spent further time and paper space both diving deeper into the performance of GIO on FashionMNIST and running many comparative image data selection baselines as it is an interesting domain to see GIO work visually, and look forward to expanding on this area of application in future work.

We deeply thank C7wm for their thoughtful and in-depth response. We are glad they find the approach competitive against recent comparable methods in this area and appreciate the task, domain and label agnostic nature of the algorithm. Below we outline our response to your questions:

W1 Algorithm 1 is intended as a supplement to the introduction, related work and method sections which define the variables X, D and G when we introduce the paper and problem setup (G is the general pool of data, X is the target set representing the desired target distribution and D is a set, possibly empty, of data we absolutely want to include.). v_opt is found using gradient descent until converges (line 5 of Algorithm 1 and 3.3 in the Method). However, we acknowledge that Algorithm 1 without the context of introduction and method doesn’t stand alone, as you correctly point out. We thank you for your feedback and will update Algorithm 1 to define the variables and be explicit about the gradient descent step so it is a standalone figure.

W2/Q4 We agree a study of computational complexity would be beneficial and will update the version to include them. BM25 is between O(n) and O(n2) with parallelization, Submodular optimization is O(n) and pruning is O(n log(n)) (GIO is O(S n) as noted in the paper where S is the gradient descent steps). In practice, we found submodular optimization to be fastest, followed by our GIO, pruning and finally BM25, which aligns. Training iterations are fixed between all methods. On our largest set with 35M training examples, GIO takes only about 40min end-to-end to complete (on CPUs only)! We feel the speed and simplicity of GIO makes it a great candidate for dataset filtering. We provided more details on how long it took for each experiment in Appendix C; it ranges from ~20min ~40min. We thank the reviewer for this feedback that makes our case about GIO stronger and look forward to elaborating these datapoints in our paper!

W3/Q5 We agree that FashionMNIST is a smaller dataset. The largest dataset we used (first experiment, WMT14) was 35M samples for WMT14 EN-FR which is a relatively large dataset, and GIO took ~40min to complete. The model trained on GIO selected data surpassed the performance of a model trained on all data and all competitive baselines, which we feel is a strong result that GIO performs well on large datasets. We also used 14M samples for spelling correction and 5M samples for WMT14 EN-DE. We wish we could have spent further time and paper space diving deeper into the performance of GIO on FashionMNIST, Cifar10, Imagenet etc and look forward to investigating this in future work!

W4/Q6 Selecting X is important to the success of GIO(and submodular optimization and BM25), and as we mention in the limitations section a bad choice of X could lead the model to generalize poorly. Selecting X and D is specific to the goal; for example, WMT14 provides dev sets of high quality which are close to the task, so we were confident those were a good target. For WMT14 EN-DE, given the small amount of data at $D=\emptyset$ initialization we ablated providing other Ds (D=25, D=50), and as our experiments suggest this has a positive impact (but not for EN-FR where there is plenty of data). We hope these details give you some insight into how we made these decisions. We feel you have brought up an important discussion on how to apply GIO to new settings and deeply appreciate it; we will provide the examples and further guidance on how to select D and X and update the revision shortly. Thank you!

Q1 Yes, this is interesting and we did an investigation in Appendix D. It makes sense that random has the closest KL divergence to target X, as here X=G and random is a subset of G. Here is a summary of the findings: We found that random selection selected both the “core” of the classes (e.g. rectangular bags) but also included outliers (e.g. triangular bags), whereas GIO focused almost exclusively on choosing the “core'” of the classes first and therefore had fewer outliers. We confirmed this with a tSNE and an analysis on separability of classes which showed GIO data clustered closer within classes, the “core”, and was more separable, therefore better performance. We hope you find this analysis as interesting as we did and encourage Appendix D for more information and some diagrams.

Q2 The goal of the FashionMNIST experiment was primarily to show that GIO works well on other domains besides text, and also to show another use case of reducing a dataset size using GIO (X=G) under a data constraint, rather than picking from data to conform to some other target X. The experiment showed GIO was able to do so effectively. We wish we could have spent further time and paper space both diving deeper into the performance of GIO on FashionMNIST and running many comparative image data selection baselines, and look forward to expanding on this area of application in future work!

We thank reviewer u6CK for their insightful review and feedback. We are glad to hear they have recognized GIO as a new sound approach to an important problem and appreciate their acknowledgement of the paper writing quality and thorough evaluation of the approach.

W1 The presence/lack of hyperparameters is an interesting subject. Our method can choose to stop on its own based on theoretically optimal criteria: whenever the KL divergence stops decreasing, any new datapoint selected is no longer adding any information about X and the algorithm and selected data has reached it’s optimal point. We believe this is a strength compared to many data selection methods which require a stopping criterion, most commonly desired data size, to be chosen beforehand. How can we know what is the best data size? Maybe we chose too little, or too much. By using the KL divergence itself as a stopping criteria, the algorithm can find that optimal state. That said, we acknowledge that there are problems where data size should be chosen beforehand, e.g. a certain compute budget, in which case choosing a data size beforehand like Xie et al. do is both necessary and important (and a primary motivator for their method). Our algorithm can also terminate early at a desired data size if that is a concern, or even finish before that point if the theoretical optimal state is achieved earlier. In short, we believe GIO’s ability to use both a theoretical optimal state to terminate, or a predetermined criterion, is a strength of the method. We acknowledge that this point could be expanded in the paper and presented as a strength of our method, not a weakness of other methods, and will include that shortly in our revision.

W2 We selected a diverse breadth of peer-reviewed data selection approaches to compare our method to. We selected BM25 for its information-theoretical foundation and identical problem setup (matching distributions using information theory with a desired target set X), Self-Pruning for both it’s strong empirical results in the category of data pruning and it’s ability to apply to any domain beyond text (like our method), and submodular optimization as a very similar method optimizing information on vector spaces with a desired target X, but using a different class and concept of functions. We believe these baselines are representative of the available accepted methods in various categories of data selection and have shown strong results. Xie et al’s method does use probability theory but also some heuristics fundamental to the method (vector size), which is why we initially classified the method under a heuristic method. However, considering that the concept and theory is well fleshed out and used as a background for their algorithm, we will reclassify their method under n-gram information-theoretical methods along with BM25 and update this shortly in our revision. We thank the reviewer for pointing out the merits of Xie et al’s method.

Q1 This is an excellent point. Broadly speaking, GIO only seeks to find the data in G that best represents the target X, both of which are represented by the same model in the same embedding space. As long as en embedding model is a) deterministic in the whole space and b) consistent for both X and G then the result should hold no matter how good or bad the embedding model is. If an embedding model is bad, so long as it is bad for both G and X in the same way, then the algorithm still works. Based on the reviewer's feedback, we decided to try a radical experiment: if we embedded our text with a terrible encoder, would GIO still work? We decided to repeat the EN-DE@25% translation task but this time, used a Spanish language encoder designed exclusively for Spanish to encode our English data. Here are the results of the translation:

Q2 We considered using simpler feature spaces like n-grams used by Xie et al. Xie et al’s decision to operate on n-grams has the distinct advantage that it scales more computationally efficiently to very large datasets than embeddings, which would require more storage for example. However, we designed our algorithm to be domain and task-agnostic and for this reason, as well as the advantages of semantic similarity with embeddings can capture more complex relationships, we designed our algorithm to operate on any arbitrary vector space from any domain (and demonstrated this with successful results on both text and image tasks).