GIO: Gradient Information Optimization for Training Dataset Selection
Authors: Dante Everaert
Christopher Potts
Venue: ICLR 2024 spotlight (2024)
Keywords: data selection
data-centric AI
information theory
kl divergence
gradient
natural language processing
computer vision
Primary area: general machine learning (i.e., none of the above)
TL;DR: We introduce a domain- and task-agnostic, scalable and efficient KL divergence-based method for training dataset selection and show it can achieve outstanding results with very small train sets on a variety of text tasks and an image task.

Abstract: It is often advantageous to train models on a subset of the available train examples, because the examples are of variable quality or because one would like to train with fewer examples, without sacrificing performance. We present Gradient Information Optimization (GIO), a scalable, task-agnostic approach to this data selection problem that requires only a small set of (unlabeled) examples representing a target distribution. GIO begins from a natural, information-theoretic objective that is intractable in practice. Our contribution is in showing that it can be made highly scalable through a simple relaxation of the objective and a highly efficient implementation. In experiments with machine translation, spelling correction, and image recognition, we show that GIO delivers outstanding results with very small train sets. These findings are robust to different representation models and hyperparameters for GIO itself. GIO is task- and domain-agnostic and can be applied out-of-the-box to new datasets and domains. We open source a pip-installable implementation of the algorithm as "pip install grad-info-opt".

Published as a conference paper at ICLR 2024
GIO: GRADIENT INFORMATION OPTIMIZATION FOR TRAINING DATASET SELECTION
Dante Everaert Amazon Search Science danteev@amazon, com dante. www..eveeertrr.e @gmaaertttmmaaa com
Christopher Potts
Stanford University
egpot sraoottssettanfoddd..oo edu
ABSTRACT
It is often advantageous to train models on a subset of the available train exam- ples, because the examples of variable quality or because one would like train with fewer examples, without sacrificing performance. We present Gradient Information Optimization (GIO), scalable, task-agnostic approach to this data lection problem that requires only small set of (unlabeled) examples representing target distribution. G1O begins from natural, that is intractable in practice. Our contribution is in showing that it can be made highly scalable through a simple relaxation of the objective and highly efficient implementation. In experiments with machine translation, spelling correction, and image recognition, we show that GIO delivers outstanding results with very small train sets. These findings are robust to different representation models and hyper- parameters for G1O itself. G1O is task- and domain-agnostic and can be applied out-of-the-box new datasets and domains. We open source pip-installable implementation of the algorithm as "pip install grad-info-opt"!
objective
INTRODUCTION
In situations in which one has very large train set available, it is often advantageous to train systems on subset of the data. In the simplest case, the train set may be so large as to run up against resource constraints, and the question arises whether performance goals can be reached with less effort (e.g. Touvron et al., 2021). can also be the case that the train examples are known to be of variable quality, say, because they were harvested from diverse websites (Luccioni and Viviano, 2021), annotated by crowdworkers (Karpinska et al., 2021), or created by a synthetic data generation process (Edunov et al., 2018). In this case, the goal is to identify reliable subset of examples.
This is the data selection problem that we address in the current paper. The end goal is to select a subset of the available train examples that leads to models that perhaps even better than) those trained on all the examples. To achieve this goal, we propose Gradient Information Optimization (G10), a highly scalable, task-agnostic approach to data selection that is based in information theory. Our method assumes access to (potentially small) set of examples X that represent the desired data distribution and (presumably very large) set of potential train examples G. Our method derives set V C G that has as much information content as possible about the target distribution X. The method begins from the natural intuition that we want V to minimize the average KI. divergence from X, and the novelty of the approach lies in making this computationally tractable by relying on properties of the derivative of the KL divergence and implementing the method extremely efficiently. Crucially, our method works in any continuous representation space, is task- and domain-agnostic, and requires no labels on examples.
least as performant (and
We motivate GIO with diverse set of experiments. We first explore machine translation using the WMT14 dataset and Transformer-bart models. In this case, G is the WMT14 dataset and X is the dev set. These experiments show that, using G10, can trained on the full WMT14 corpus with only a fraction of the example in G, which represents very large efficiency gains. We then turn to spelling correction. In this case, the set G is generated by
surpass the performance of model
pip install grad-info-opt. Appendix B for details. Also available at: https: // /ig thub. com daeveraert aeeeeerrerr//rraieeneee grwdientddengggeeeggoooggg information aet.iiiitiaiiii

---

Published as a conference paper at ICLR 2024
noisy synthetic process and the target distribution X is set of actual spelling errors. Here, we using GIO to home in on realistic train examples. Our results show that we can do this extremely effectively. Finally, we apply G10 to an image recognition task (FashionMNIST) and show again that our method can reduce the size of the train sets chosen without large drops in performance, this time operating with representations of images. In this case, trust the train set G to represent the space accurately, and our goal is simply to select a useful subset of G. Thus, in this case X = G. Finally, we discuss expanding GIO, report on a wide range of robustness experiments and empirical analyses of how and why the method works in practice, and publish pip-installable package.
2 RELATED WORK
Active learning. Active learning methods (e.g. Sener and Savarese, 2018; Gal et al., 2017: Kirsch et al., 2019) can be cast as data selection methods in our sense. In active learning, one iteratively chooses new unlabeled training examples to label, with the goal of efficiently creating a powerful train set. By contrast, G1O makes no use of labels and is oriented towards the goal of identifying a subset existing cases to use for training. Additionally, active learning is most suited to classification problems, whereas GIO works with any arbitrary task.
Heuristic. GIO is closer to recent methods in which one uses a large language model to generate a large number of candidate texts and then extracts a subset of them based specific criteria. For example, Brown et al. (2020) develop heuristic method to filter CommonCrawl based on a trained classifier's probability that datapoints are high quality. Similarly, Wenzek et al. (2020) develop pipeline to clean CommonCrawl based principally on the perplexity of an LM trained on high quality text, and Xie et al. (2023) develop sampling technique based on approximate n n-gram counts.
Like GIO, these heuristic methods aim to select a subset of data that is higher quality and more relevant. However, they are either highly tailored to their particular tasks or they require very large numbers of examples (to develop classífiers or construct target prohabilities). By contrast, GIO is task- and domain-agnosee comparatively few gold examples X to serve as the target distribution.
it can be applied plug-and-play to new task and dataset, and requires
Similarity Search. Methods using vector or n-gram similarity search can also be used for data selection at scale (e.g. Johnson et al., 2017; Bernbardsson, 2017; Santhanam et al., 2022). The technique would index G and X and retrieve the top-k datapoints from G for each point in X. Like our method, similarity search works in continuous space. However, similarity search can be prone to selecting suboptimal points; we review such a case in detail in Section 3.4. Additionally, similarity search does not have a natural stopping criterion and requires data size to be chosen arbitrarily. Is 10% data enough? 20%? We don't 't know a priori. And if the data in G is far away from X, similarity search will still choose it up to the desired data size. Recently, Yao et al. (2022) use BM25 retrieval method for data selection, with strong results. However, BM25 operates on a bag-of bag-of-words model, which can make it challenging when the target set is small, and like any similarity search, requires data size to be chosen arbitrarily beforehand. Further, this method only applies to text tasks, whereas GIO applies to any task with continuous representation.
Data Pruning. Data pruning has shown promise for the data selection problem (e.g. Paul et al., 2021). Different works in data pruning use varied approaches. but generally iteratively identify and add optimal samples from training data. However, most works in data pruning (e.g. Paul et al., 2021; Yang et al., 2023; Saseendran et al., 2019) apply only to classification. However, recently Sorscher et al. (2022) develop a self-pruning method with strong results that can be applied to any task with continuous representation, like Gio, using clustering based selection. Unlike G1O and similarity search however, self-pruning does not consider any desired target distribution X.
Submodular Optimization. Submodular optimization methods have also been proposed for data selection (Kaushal et al., 2022), where an optimizer optimizes a submodular function between a general set G and a target set X. Like GIO, submodular optimization takes into account a desired target distribution X, has natural stopping criterion that does not require data size to be chosen arbitrarily. Unlike G1O however, submodular optimization is restricted only to submodular functions, which have certain assumptions that may not hold true. For example, submodular functions assume adding each extra datapoint diminishes the return of that datapoint, which is not necessarily the case when constructing dataset. For example, if we have already have selected data but have unexplored
2

---

Published as a conference paper at ICLR 2024
Algorithm I Gradient Information Optimization
Quantize X, D, G using K-means and pick the cluster centroids X Des Ge as the new points while Not Stopping Criterion do Gradient-Descend to find Vopt
V1 < previous Vepts X or random
D. We explore different techniques our experiments
Perform Vk+1 f- Vk Ove DRr(Px. PD.U(v.)) until converged to Vope
#
Update De:
Va - argininy; ||v; Vopill
The closest point G, to Vopr
De t + [vb}
Remove Vo from Ge
end while
Explode: Select points from full D and which belong the chosen centroids' (D&UV, clusters
region where X has a mode, adding data in that region does not exhibit diminishing return. Gto, on the other hand, makes no assumptions on the functions can use.
Overall, previous work in data selection is typically tailored to specific domain like NLP (e.g. Wenzek al., 2020) or image recognition (e.g- Gal et al., 2017), and makes assumptions about the data available, for example, that the target set X is large enough to construct an LM (e.g. Wenzek et al., 2020), or that it has labels (c.g. Sener and Savarese, 2018). In addition, many of these methods discrete approximations and heuristics (e.g. Xie et al., 2023; Yang et al., 2020). In this work, provide a general, pearettcally-motivvne data selection method that works with large or small X and can be applied out-of-the -f--t---oooox to any domain (image, text, etc) without needing labels.
GRADIENT INFORMATION OPTIMIZATION: METHOD
We formulate data selection as maximizing information content and outline the natural algorithm for this objective, which is infeasible. We then introduce optimizations which enable the algorithm to work at scale, and conduct tests to show the algorithm consistent and robust different scenarios.
3.1 ABSTRACT FORMULATION OF THE DATA SELECTION PROBLEM
We assume that all examples are represented in continuous space. We have set of train examples G and a target ideal state X. We allow also that there may be existing train examples D that we definitely want include in our train set, though 1) can be empty. Our goal is to identify subset V of G such that the set D UUF contains the most information about X.
In this setting, it is natural to take an cottemmation-thhoredie approach. Let px (x) be the distribution of target X, and let pDuv(x) (x) be the distribution of selected data D V. The information content of D U V about X is the negative KL divergence from p.x (x) to PDuv (x) (Kullback and Liebler, 1951). In this context, the general objective of data selection is as follows:
Choose data C such that /. px px(x) log
(x(x)
dx minimized pDuv(x)
(I)
The implication is that data selection method which gives the minimum KL divergence will also give the best performance (assuming we are correct that X represents the task to be solved).
3.2
NAIVE APPROACH
natural approach to hill-climb on the KI. divergence objective (1). Given existing data and points Vi...., Vą of G, recompute the distribution PDu(v,) 3 (x) for each va. pick the one that gives the minimum KL divergence, and add to our selected set D:
+ + argmin / px (x(xx) log (x(x) dx
v,EG J
PDu(v.
3

---

Published as a conference paper at ICLR 2024
Quantization- Quantization-Explaiinn Process
Originl
Seacted Dara
Seleted Oegine
Dinwnicce
Dimessinn
Dinwe 3
Visualization of the Quantization-Explee Process. From left right: original data (400 points), representative K-means centroids (50 points) of the original data (Quantization), selected centroids after data selection, original data represented by the selected centroids (Explosion)
Figure 1:
Unfortunately, this algorithm is intractable in practice. We need to construct a new distribution PDu(v;) (x) and compute KL divergence for every E G, at each step, Therefore, the complexity at each iteration O[G)| C). where C' is the cost of computation for the KL divergence. For a dataset of only IM and 0.1s Is per iteration, it would take 70 days complete the algorithm. The method is also prone to adding the same point multiple times.
3.3 GRADIENT INFORMATION OPTIMIZATION
GIO addresses the shortcomings of (2) with a combination of mathematical and implementational optimizations. The method is described in Algorithm 1.
First, instead of calculating divergence for each point, we use the derivative of the KL divergence to find the optimal point. We rewrite Ppu(v,) (x) = g(x, vr), a function of only and since D is not changing, and thus the optimization term in each iteration becomes:
argmin / px (x) log v.EG
Px (x(x) dx
g(x, v:)
We can relax the constraint that v, G the space of all possible v and solve this integral minimization for the optimal Vopr- Since px is unchanging and the integral implicitly removes x as variable, the integral defines a functional [g(v)] Therefore, we partially differentiate with respect to v and do gradient descent with the partials Vv F(9l to solve for Vopr- All together, this becomes:
(// p(x) log
p(x)
Vk- Vb 7 Ov
dx
g(x, vи.)
Once we have Vopr, we find the nearest v G to add to D, as the closest v, G is the solution to (3). We assume G locally dense for the extrema of the integral in (4); see Appendix A.2 for details.
The complexity at each iteration, for S gradient descent steps, is O(S C*) which does not increase with G. Therefore, when |G| S, as common in practice, the derivative trick is faster than the naive algorithm. We time both algorithms in Section 3.4 and show the derivative trick is 80% faster.
Second, even its most efficient, an algorithm that adds point-by-poin becomes intractable. There- fore, we use coantization-explosst Vassilvitskii, 2007) and pick the centroids μ:- We then perform the algorithm using the cluster centroids μe instead of the original data. Finally, we use chosen cluster centroids and explode back to the original data based on cluster membership. Figure provides an overview of the process.
process. First, we cluster the data with K-means (Arthur and
Third, to compute the KL Divergence in high-dimensional spaces, we use the k-Nearest Neighbors approximation for continuous KL divergence proposed by Wang et al. (2009), and modify to be an average across all points to bypass 0 gradient problems (details and proof of modification are in the Appendix A. Let D = m, [ = n and d be the dimensionality:
m
1 m m
I m
(Px || PD)
Σ d logux(0)
log (i)
+
log
(5)
k(n - 1)
Where PA (i) is the distance from point X; to the keth nearest point D and (r()) is the distance from point X; to the 2th nearest Xjzi- We automatic differentiation compute the derivative.
4

---

Published as a conference paper at ICLR 2024
Seli-Consistenc
GIO vs. Similarity Search
All trcaing 2era
Selecad Besk
Selcis Data
Daia
The leftmost graph shows X and G. which come from the same distribution. The second graph shows that GIO recovers nearly all G (self consistency). The right two graphs compare G1O with similarity search. Points within the circle formed by X are more ideal than points outside. By considering the distribution, GIO selects nearly all points inside before terminating (third graph). By comparison, in order to pick points within the circle, similarity search also picks range of points outside the circle, which is suboptimal (fourth graph).
Figure 2:
We can stop when the KL divergence increases, reset G and allow the algorithm to pick again, among a variety of criteria. We explore several in our experiments and list additional criteria in Appendix B.2. Unlike data selection methods that make data size a hyperparameter (e.g. Yao et al., 2022), GTO provides a natural stopping criterion (KL divergence). Finally, initializing D from a uniform start rather than empty leads to same optimal points but a smoother convergence; see Appendix A.3.
Limitations. We derived GIO from the natural information-theoreti objective, however, we can use any arbitrary statistical distance in the GIO framework. For example, in situations where G' is close to X with the exception of a large gap somewhere, the statistical distance max px (x) - pD may be better suited. We also use gradient descent to iteratively find Vopr. but we know the space is non-convex. Therefore, replacing gradient descent with a method like particle swarm optimization (Kennedy and Eberhart, 1995) may lead to hetter selected data. Finally, in practice it is important to ensure that X reasonably represents the space model might be used on. A narrow X could make a model trained on G10-selected data perform poorly when confronted with inputs that lie outside X. Methods like starting from subset of training data, which we explore, or adding uniform points to X to encourage generalization, should be explored. We leave these improvements to future work.
3.4 ANALYTIC CHECKS
G1O is self-consistent. We define self-consistency follows: if both G and X come from the same distribution, i.e., (x) = px (x), a good data selection method should choose all of G. We show G1o is self-consistente let X be 100 points from normal distribution centered at (3, 4) and let G be another 100 points from the same distribution (Figure 2, 1st graph). We run G1O on this setup; G10 selects 96% of G before termination, showing GIO is self-consistent (Figure 2, 2nd graph).
GIO is negative-consistent. We define negative consistency as follows: if G is very far from X, i.e. d(px (x). pa (x)) > O. good data selection method should not choose any of G. Most data selection methods that rely on choosing a desired data size as a stopping criteria (e.g. Yao et al., 2022; Xie al., 2023, similarity search)
not negative consistent; they will select data regardless of how close or far the data may be from X. We show GIO is negative-consisten with the following setup: let X be the same as above, but this time let G be 100 points centered far away (300, 400). We run GIO on this setup: G1O terminates without adding any points from G, showing it is negative-consistent.
Quantization in GIO is consistent with the original space. Quantizing the space with K-means should not change the distribution of data. We show this with the following setup: let X be 400 points from a normal distribution centered at (3, 4). We quantize X using K-means with K=50, and compute the KL divergence DxL || X which should be near if the distributions are close. The KL divergence is 0.44, showing the quantization in GIO is consistent with the original space.
The derivative trick is 80% faster. We benchmark the wall-clock time between the naive hill-climb method and G1o with the derivative trick, with 100 points in X and 2000 points in G spread uniformly, and run the algorithm for 100 iterations. The regular hill-climb method takes 1 1369s, whereas the derivative trick takes 257s, representing an 80% speedup.
5

---

Published as a conference paper at ICLR 2024
Table Machine translation results. Training data sizes and BLEU scores of models trained full data, G10-selected data, comparative methods and random subsets for various initialization states. Bold the best score in each initialization, italic is the best score overall. GIO outperforms a model trained with the full EN-FR with only 40% of the data and outperforms all comparative methods in 10/12 evaluations. achieves 99% of the performance EN-DE with only 60% of the data.
EN-FR
EN-DE
Init. %
System
Train Size
Dev Test
WMT
DKL
Train Size
Dev Test
WMT
DKL
14
14
Ours
34,2
41.2
156
22.1
24.3
148
BM25
33.9
41.0
172
22.6
24.9
175
0
Pruning
5.6M
33.3
40.3
152
701K
22.0
24.2
177
Submod.
33.8
40.8
170
22.0
24.5
164
Random
33.1
40.0
194
21.9
24.0
183
Ours
34.8
42.2
166
23.9
27.0
159
B.M25
34.6
42.0
179
22.9
26.3
178
25
Pruning
14M
34.1
41.4
174
1.7M
23.0
26.3
171
Submod.
34.4
41.7
181
22.7
25.5
172
Random
34.3
41.4
195
23.0
26.7
182
Ours BM25
34.7
42.3
172
24.2
27.9
164
34.3
42.1
185
23.7
27.9
178
50
Pruning
21M
34.2
41.9
183
2.5M
23.6
27.1
179
Submod.
34.4
41.7
185
21.9
24.5
176
Random
34.1
41.7
195
24.0
27.3
181
100
Full
35M
41.8
188
4M
28.2
180
GTO selects based on a distribution; similarity search does not. Let X be circle of 2D points encapsulating region of space, and let be 2000 uniformly distributed points (Figure 2, 3rd and 4th graphs). The ideal points are the points within the region encapsulated by X and should be chosen with preference over points outside. Figure 2 shows the points selected by semantic search (4th graph); in order to get the data inside the circle, it also picks the data outside the circle. Figure 2 also shows the GIO-selected points (3nd graph); by considering the distribution all of rather than the simple Euclidean location of each X, G1O selects mostly points which within the circle, desired.
EXPERIMENTS
We perform four of experiments to validate GIO. First, replicate the setup of Vaswani al. (2017) on WMT14 14 and show that using G10 can surpass the performance of the full model with fraction of the data. Next, we demonstrate G1O robust different choices of embedding models and quantization. Third, we use a spelling correction task to show that G1o selects the highest quality data from pool of -uxlityy synthetic data. Finally, we show GIO reduces the training set size of FashionMNIST image task without a big drop in performance. We show that GIO achieves the lowest KL divergence compared to alternatives, and that this correlates with model performance.
4.1 MACHINE TRANSLATION EXPERIMENTS
Our first set of experiments seeks to show that GIO can pick data from general corpus meet or exceed the performance of a model trained on the full corpus.
Data and Methods We use Transformer Big from Vaswani et (2017), trained for 300k iterations with the same hyperparameters. We use the same processed WMT14 training data. We report the BLEU score (Papineni al., 2002) on the WMT14 set.
Details of each experiment are in Appendix C From Vaswani al. (2017)

---

Published as a conference paper at ICLR 2024
Table 2: Training data sizes and BLEU scores of the base version (K=1500 and MPNet) and the variants of K and embedding model. BLEU scores of the variants vary by only 0.4% on average from the base, indicating GIO robust different quantization and embedding models
EN-FR
EN-DE
System
Train
Dev
WMT
DKL
Train Size
Dev
WMT
DKL
Size
Test
14
Test
14
Base (MPNet, K=1500)
5.6M
34.2
41.2
156
701K
22.1
24.3
148
MiniLM Variant
5.7M
34.0
41.1
737K
22.3
24.6
K=1000 Variant
5.6M
33,9
41.2
169
701K
22.1
24.3
150
K=3000 Variant
5.7M
34.2
41.3
138
718K
22.3
24.6
133
Average Variance from Base
1.9%
0.5%
0.2%
3.7%
0.6%
0.8%
We apply GIO to select subset of data from the WMT14 train set using the inputs only (as our method makes no use of labels). the training data we can select from. For target state X, we collect the dev sets for WMTOS-WMTI3. extract 3K pairs and report BLEU on this held out dev set, and use the remaining R 12K pairs as X. For initial state D, we consider starting from an empty set, 25% random subset of train data, and 50% random subset of train data, and report results for each setting. We use MPNet- Base-BBse-V V2 model (Song et al., 2020) to embed the input sentences in a continuous vector space and use K=1500 for quantization. We compare this embedding model and quantization amount to other settings in our robustness experiments (Section 4.2). As our stopping criteria, we stop when the KI. divergence increases. We also deduplicate the data pairs before training.
We compare GIO to a random subset of data of the same size. In addition, we compare against several competitive baselines with different approaches: Yao et al.'s (Yao al., 2022) recent similarity search approach of BM25 retrieval, the recent data pruning self-pruning method introduced by Sorscher al. (2022), and the submodular information optimization with mutual information as the function to optimize. To keep the setup equal, we also initialize each of the baselines from a 0%, 25% and 50% random subset and run those algorithms to have the same size as the G10-selected data. For submodular optimization and self-pruning, we use the same clustering GTO is run with.
Results We find that GIO outperforms the random baseline at every initialization, data selection method should always outperform candomly-selectedd subset of the same size. Table 1 shows the BLEU score on dev and WMT14 test sets and demonstrates G1O always outperforms random. BM25, submodular optimization and data pruning only outperform random sometimes.
G10 outperforms the EN-FR model trained on the full data using only 40% of the data. At initialization of 25% and 50%, model trained on GiO-selected data outperforms the full Vaswani et al. (2017) model trained on all data by +0.4 and +0.6 BLEU, respectively. In addition, a model trained on the Gro-selected data at 0% initialization achieves 99% of the performance of the full model with only 16% of the data. It outperforms all comparative methods at all initializations. In EN-DE, it gets 99% of the performance with 60% of the data, and 88% of the performance with only 18% of the data.
GIO outperforms all comparative methods in 10/12 of the evaluations. GIO always matches or outperforms the comparative methods with initializations of 25% 50%, by an average of +0.7 BLEU on WMT14 14 and +0.6 BLEU on Dev Test, and only falls short with 0% initialization in EN-DE.
GIO has the lowest KL divergence in 5/6 tests (Table 1), which correlates with model performance. The implication of (1) is that a dataset with lower KL divergence between train and target will perform the best. From Table the average Spearman rank correlation coefficient between KL divergence and best performance is 0.83 and the median is 1, showing a high degree of correlation between a dataset that minimizes KL divergence and model performance, and thus confirming the implication.
In summary, G1O leads to the lowest KL divergence between train and target set out of all the methods, which correlates with model performance and confirms the theory (1). Notably, a model trained with Gro-selected data outperforms a model trained on the full data in EN-FR despite using only 40%
WWe use Fairseq's (Ott al., 2019; FAIR, 2022) scripts
7

---

Published as a conference paper at ICLR 2024
Table Spelling correction results. Training data size, % high quality data and KL Divergence of the full data, G1o, BM25, submodular optimization, data pruning and random data. Bold is the best score overall. GIO selects 73% high quality data which outperforms all other methods.
System
Train Size
% High Quality
DKL
Ours
73%
224
BM25
55%
264
Pruning Submod.
3.6M
54%
241
59%
245
Random
50%
284
Full
14.7M
50%
280
of the total data and came to within 99% of the full model in EN-DE using only 60% of the data. G1O outperforms the random baseline at all initializations and outperforms all comparative methods in 10/12 evaluations. Overall, these experiments show GIO can achieve and surpass the performance of model trained on full data and comparable baselines, by explicitly optimizing for KL Divergence.
4.2 ROBUSTNESS
G1O above relies on two approximations to work: an embedding model, and K-means to quantize the space into representative samples of the full data. In this section, we show that G1O is robust to different embedding models and different values of K. The results are summarized in Table 2.
GIO works with different embedding models. G1O should be robust to different text embedding models. We change the embedding model from MPNet-Basse V2 to MiniL.M-L12-v (Wang et al., 2020), which has different architecture, training, and produces embeddings of a different size. We then rerun the 0% initialization experiments end-to-end with the new embeddings for both EN-DE and EN-FR. Table 2 shows that using MiniLM in G1O results in roughly similar selected data size (4.4% difference on average) and virtually identical performance (0.7% difference on average), demonstrating that GTO is robust to different embedding models.
G1O works with different choices of K. G10 should also be robust to varying amounts of quan- tization. We decrease the value of K from 1500 to 1000 and increase to 3000 and rerun the 0% initialization experiments end-to-end for both new values of K, EN-FR and EN-DE. For K=1000 due to the coarser grain, GIO selects more data, therefore we sample from the selected data the same amount as K=1500 in order to maintain parity. Table shows that performance is virtually identical between the different values of K (0.4% difference average), demonstrating G10 is robust to different values of K. In general, higher values K have lower KL Divergence and slightly better performance, which is expected the quantization is more fine grained.
4.3 SPELLING CORRECTION EXPERIMENTS
In this section, we set up a problem with a pool of high and low quality synthetic candidate train examples and show GIO selects mostly high quality data. In addition, we set a new state of the art on the challenging BEA4660 spelling correction benchmark of Jayanthi et al. (2020); see Appendix C.
Data and Methods We follow the setup of Jayanthi et al. (2020) and collect 15M samples from the 1 Billion Word benchmark corpus and deduplicate. To create high quality data, we use the best noising technique (prob) from Jayanthi et al. (2020) and noise half the data. For low quality data, we use the "word" method with high replacement rate (70%) and noise the other half, and mix the two sets.
We apply G1O to select subset of data from the training set. G is the training data we can select from. For target state X, Jayanthi et al. (2020) provide 40k real spelling mistakes and corrections from the BEA grammar correction corpus. For initial state D, we start from empty set. For the embedding model, we use MPNet and K=1500 for guantization. As our stopping criteria, we experiment with a new scheme: first the algorithm runs until the KL divergence increases, then reset G and allow the algorithm to pick again from the training data, until the KL divergence decreases. As before, we compare GIO to random subset of the same size and pruning, submodular and BM25 methods.
8

---

Published as a conference paper at ICLR 2024
Table 4: Image recognition results. Training data sizes and accuracy of models trained on G10- selected data and random subset. Bold is the best score between ours and random. G10 gives the best performance under the reduction training data size. Full model is provided for comparison
System
Size Train/Valid
Accuracy²
DKL
Ours 15,000/1,700
92.2%
771
Random
15,000/1,700
91.5%
740
Full
56,300/3,700
94.5%
739
Results G1O selects high quality data. A good data selection method should select mostly from the high quality data. Table 3 shows G1O selects 73% high quality data, compared to 59% for submodular optimization, 55% for BM25 and 54% for pruning. GIO's KL Divergence lower than comparative methods and random, indicating KL divergence is also an indicator of data quality in this setup.
4.4 IMAGE RECOGNITION
For our fourth set of experiments, seek to show that GIO works well in domains outside of NLP. We focus on the FashionMNIST (Xiao et al., 2017) image recognition problem and show that we can GIO to dramatically reduce train set sizes without big drops in performance.
Data and Methods The MNSSI task has 10 classes of 28x28x1 images. There are 60,000 images in the training set, and 10,000 images in the test set. Our task will be to select subset no more than 25% of the total data that best approximates the training data. We then finetune the Resnet50 model (He et al., 2015) for 5 epochs with Adam (Kingma and Ba, 2015) (LR=5e-5) with the chosen data to do FashionMNIST classification. We split the data into train and validation sets, and pick the best checkpoint by validation loss. We report the accuracy on the test set.
We apply GIO to select a subset of data from the training set. We use the training set as both G and also as our target set X. We start D from an empty set. We use the normalized and normed vector format of the images themselves in the algorithm, and use K=1000 for quantization. As our stopping criteria, we run the algorithm until we get 250 clusters (250 iterations), which is 255 of the data. We also report results on random sample of 25% of the data as comparison.
Results G1o outperforms simple random subset by +0.7%. G1O in this setup is optimized to pick the images which add the most information about the entire training set. Table 4 shows training on GIO-selected data only dropped performance by 2.3% from the full model, compared to a drop of 3.0% for random subset of the same size.
5 CONCLUSION
We presented GIO, a task- and domain-agnostic data selection method that works in any continuous space with few assumptions. GIO begins with the natural objective of minimizing KI. divergence between a target set and selected set, and uses the gradient of KL divergence and an efficient implementation to find and select the data points that optimize that objective at scale. G10 selected high quality data, selected data that outperformed models trained on full data and on recent data selection techniques, and was able to effectively reduce the training set size under a given resource constraint without big drops in performance. Current models consume large quantities of data, and we hope GIO can help improve the performance of models while using less data and fewer resources. Additionally, with large quantities of synthetic and scraped data of variable quality available, we hope G1o can help home in on high quality data, For example, G10 can be used to select high quality synthetic data output by large language models for particular task, or choose training data for LLMs that aligns with a certain intent. Improvements and changes to the statistics and optimization in GIO and applications of GIO to varied domains and tasks are promising directions for future work.
'Results averaged over runs 'See Appendix A.1 for why this not 0

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
We would like to acknowledge and thank the broader Amazon Search Spelling Correction and Autocomplete teams for their support during this work. We would like thank in particular Parivesh Priye, Jingfeng Zhang, Xuan Guo, Alex Klementiev and Rohit Patki for their help in reviewing and giving invaluable feedback and insights over the course of developing the idea. We are grateful to all those who supported us and gave us feedback throughout, and apologize in advance if we missed anyone in this section.
REPRODUCIBILITY
We include all the necessary code to reproduce both our method as well as all the baselines in the supplementary materials under the package cnaflienniiforraatiooooiiiiiii detailed instructions on the usage of all the code, including comparative method setup, as well as details on all hyperparameters used and testing setup in Appendix C.
Further, we include
REFERENCES
David Arthur and Sergei Vassilvitskii. k-means++: the advantages of careful seeding. In SODA '07: Proceedings of the eighteenth annual ACM-SIAM symposium on Discrete algorithms, pages 1027-1035. Association for Computing Machinery, January 2007. doi: (0.3555//833883/28884)) URL https: /dl...../// acm. wwt.:aamm........a.a.. btg//oi///e 5555
1283494.
Erik Bernhardsson. Annoy (approximate nearest neighbors oh yeah), 2017. URL https://
github. wwtps/oottt.o///ccooooooooooottt annoy. Apache-2.0 2.0 license.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.
Christopher Bryant, Mariano Felice, Øistein E. Andersen, and Ted Briscoe. The BEA-2019 shared task grammatical eиrо correction. In Proceedings of the Fourteenth Workshop Innovative Use of NLP for Building Educational Applications, pages 52-75, Florence, Italy, August 2019, Association for Computational Linguistics. doi: 10. 1118655/11/1194440 URL https:
aclanthology. org/W19 9-199-144 6
Sergey Edunov, Myle Ott, Michael Auli, and David Grangier. Understanding back-translation at scale. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pages 489-500, Brussels, Belgium, October-November 2018. Association for Computational Linguistics. doi: 10.18853/v/DD18 1045. URL http///ttts////caaaaaaaaao..o...
FAIR. 2022. URL htt.atppee/ggttat....//aaaaoobbp http://twww.ccccceeeceoeooooooee license.
MIT
Yarin Gal, Riashat Islam, and Zoubin Ghahramani. Deep bayesian active learning with image data. arXiv preprint
2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. arXiv preprint 151v03151.. 2015.
Amithash K J. Using resnet for fashion mnist in pytorch, 2020. URL. https:
//github. htt.tt//w..ccom/..aaaataaaa///tr Dew.eerrnneeeeeeeiiiniie Beceng____prrrree wwwriiiets/bliii s/blb//
ipynb. Public Google Colab.
Sai Muralidhar Jayanthi. Neuspell: neural spelling correction toolkit, 2021.
URL https:
htt.hhww..oo//eeeeuuee.ee 1//euspellleeeeee1eeee 1. MIT license.
10

---

Published as a conference paper at ICLR 2024
Sai Muralidhar Jayanthi, Danish Pruthi, and Graham Neubig. NeuSpell: A neural spelling
correction toolkit. Proceedings of the 2020 Conference on Empirical Methods in Nat-
ural Language Processing: System Demonstrations, pages 158-164, Online, October 2020. Association for Computational Linguistics. doi: 10. 10.18653/v 102.33/w1/220.ooo// http///000m0/eemmmsssss fitts://acclattee logy. wtw.llggoyooooooeeeeegeeeegg wwwp///........o.ooo eeppeeGemlpppdmmem 21.
URL
Edwin T. Jaynes. Prior probabilities. IEEE Transactions on Systems Science and Cybernetics, 4(3): 227-241, 1968.
Jeff Johnson, Mathhis Douze, and Herve Jegou. Billion-scale similarity search with gpus. arXiv preprint /002:/7702.88777 2017.
Marzena Karpinska, Nader Akoury, and Mohit lyyer. The perils of using Mechanical Turk to evaluate open-ended text generation. Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing. pages 1265-1285, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10. 10.186/3///202...mmmmpp 97. URL https: / htt///aaaannttoooooooooo org/ 2021. emnip-main.99
Vishal Kaushal, Ganesh Ramakrishnan, and Rishabh Iyer. Submodlib: A submodular optimization library, 2022.
J. Kennedy and R. Eberhart, Particle swarm optimization. In Proceedings of ICNN'95 IIneenational Conference on Neural Networks, volume 4, pages 1942-1948 vol.4, 1995. doi: 10.1109/ICNN. 1995.488968.
Diederik P. Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv: 1412.6441,,69800 2015.
Andreas Kirsch, Joost van Amersfoort, and Yarin Gal. Batchbald: Efficient and diverse batch acquisition for deep bayesian active learning. In H. Wallach, H. Larochelle, A. Beygelzimer. F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neu- ral Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https: httpss////pooceeengg...eeo 45c.eea2222225500222222222 www.5aaawwaace0ee0000e00 Paper. paper.pdf.
cc htt.prps..cpppppprrrra...... /tsp//////pppe///////L..
Solomon Kullback and Richard Liebler. On information and sufficiency, Annals of Mathematics, 22 (1):79-86, 1951. doi: 10. 2/11///////1177////
Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Veselin Stoyanov, and Luke Zettlemoyer. BART: Denoising sequence-to-sequenee pre-training for natural language generation, translation, and comprehension. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 7871 78810, Online, July 2020. Association for Computational Linguistics. doi: a:tp//220...1mmmmmmm acl-main 703.
URL
https: htttps:/aacaathhat htt///aaaaantntooooooooeaa
Alexandra Luccioni and Joseph Viviano. What's ss in the box? analysis of undesirable con- tent in the Common Crawl corpus. In Proceedings of the 59th Annual Meeting of the Asso- ciation for Computational Linguistics and the Ith International Joint Conference Natural Language Processing (Volume 2: Short Papers), pages 182-189, Online, August 2021. As- sociation for Computational Linguistics. doi: 12653//1122221.acl-hhorff //aclanthology org
URL https:
Myle Ott. compound_ssple bplui__leeussshhhh___oe_ee bleu.sh, 2019. URL http://ggg.....ttt...m////////// aea.ea3ceeee75223349977
MIT license.
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics (Demonstrationss) pages 48-53, Minneapolis, Minnesota. June 2019. Association for Computational Linguistics. doi: 10. 10.18653/11/11944000 URL https: / aclanthology. org/N1 4009.
11

---

Published as a conference paper at ICLR 2024
Pankaj.
Fashion mnist with pytorch (93% accuracy), 2019. URL https://wwit www.:///////.gggggg.e.. kaggle. com/ code/ mattraa///aaaiiiomissss--ii-ieie license.
accuurayy. Apache-2.0
Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: a method for automatic
evaluation of machine translation. In Proceedings of the 40th Anmual Meeting of the Association for Computational Linguistics, pages 311-318, Philadelphia, Pennsylvania, USA, July 2002. Association for Computational Linguistics. doi: 10.31 10.3115/10773333.073335 URL https:
//aclanthol log. wwwp:ogggpP21000
Mansheej Paul, Surya Ganguli, and Gintare Karolina Dziugaite. Deep learning on a data diet: Finding important examples early in training. In M. Ranzato, A. Beygelz- imer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural In- formation Processing Systems, volume 34, pages 20596-20607. Curran Associates, Inc.. 2021. URL https: http:::///ccceeeeeeeeeeeee wwwrips. wttpppperriieeeeppeeeeeeeeeeeeee files//aperr2201/ .pff.
file/ file/ http://ssseeeeeeeeeaaaaaserrrree
Keshav Santhanam, Omar Khattab, Jon Saad-Falcon, Christopher Potts, and Matei Zaharia. Col- Effective and efficient retrieval via lightweight late interaction. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages
BERTv2:
Seattle, United States, July URL 2022. Association for Computational Linguistics. doi: 118653//120222.aaac--m https: /aaaaaaaaaaatooooooooaoaa org/o0/g//00 2022. naacl-ma main. 272.
Arun Thundyill Saseendran, Lovish Setia, Viren Chhabria, Debrup Chakraborty, and Aneek Barman Roy. Impact of data pruning machine learning algorithm performance, 2019.
Ozan Sener and Silvio Savarese. Active learning for convolutional neural networks: A core-set approach. In International Conference on Learning Representations, 2018. URL https:// openreview www. forum? www.oooroooooioooooo Hidk RM.
Rico Sennrich. Subword neural machine translation, 2021. URL https::/githu..c...
besennrich/suwwore nmt. MIT license.
Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In Proceedings of the 54th Annual Meeting of the Association for Com- putational Linguistics (Volume 1: Long Papers), pages 1715-1725, Berlin, Germany, Au- gust 2016. Association for Computational Linguistics. doi: 10.18653/1/P16-1166 URL https: http:://aaaaathoooyooooooooo ort.gyy-rrr//11111.2
Kaitao Song, Xu Tan, Tao Qin, Jianfeng Lu, and Tie- Yan Liu. Mpnet: Masked and permuted pre-training for language understanding. arXiv preprint arXiv:2004.09297 2020.
Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Ari Morcos. Beyond neural scaling laws: beating power law scaling via data pruning. In S. Koyejo. S. Mo- hamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural In- formation P Processing Systems, volume 35, pages 19523-19536. Curran Associates, Inc., 2022. URL https: htttp:://prooceedingg. neurips cetpepp.cceppppeerrrrrrr..... httpappeerriileepppe/////// 1rll://00.cabbbleeeee0//33333333
Paper won--renneeeeeeeeeeeeeee Sore pdf..
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herve Jegou. Training data-efficient image transformers and distillation through attention. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 10347-10357. PMLR, 18-24 Jul 2021. URL https: / http//ppeccess...c/oo/p mlr. www.lsppress..em.ee / / httesssssvv1tttttt/eeee v119//tuvvoo221l....... hht.r. html.
Ashish Vaswani, Noam Shazeer. Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gamett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https: http:////occeeeeeeeeeee neurips. ce/ cc/paper happppaer_f__ supper_flles/_e pappr/2017/ file/ 3.le/ff5eee wfw.ee22345444 44ee24235eddeee 7dee fbd053e1 9
Paper. .pdf.
12

---

Published as a conference paper at ICLR 2024
Qing Wang. Sanjeev R. Kulkarni, and Sergio Verdu. Divergence estimation for multidimensional densities via k-nearest-neighee distances. IEEE Transactions on Information Theory, 55(5): 2392-2405, 2009, doi: 10.1 00.110//TT.220992016660
Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang, and Ming Zhou. Minilm: Deep self-attention distillation for task-agnostic compression of pre-trained transformers, 2020.
Guillaume Wenzek, Marie-Anne Lachaux, Alexis Conneau, Vishrav Chaudhary, Francisco Guzmán, Armand Joulin, and Edouard Grave. CCNet: Extracting high quality monolingual datasets from web crawl data. In Proceedings of the Twelfth Language Resources and Evaluation Conference, pages 4003-4012 Marseille, France, May 2020, European Language Resources Association. ISBN 199-10099554-34--2 URL https: / aclantholoogy org/2020 ...00...r 1rec-1- 1.44..
Han Xiao, Kashif Rasul, and Roland Vollgraf. Spelling correction foreign language. arXiv preprint arXiv:1708.077 07747, 2017.
Sang Michael Xie, Shibani Santurkar, Tengyu Ma, and Percy Liang. Data selection for language models via importance resampling. arXiv preprint arXiv:2302.03169, 2023.
Kai-Cheng Yang, Onur Varol, Pik-Mai Hui, and Filippo Menezer. Scalable and generalizable social bot detection through data selection. Proceedings of the AAAI Conference on Artificial Intelligence, 34901):1096-1100 Apr. 2020. doi: 10.16099aaii 134.:aa..i3.1............ 54600. URL https:
//ojs . aaai org/ index.ndxx. php/ AAAI/ SeA./aarticeevvew//55. 5400.
Shuo Yang, Zeke Xie, Hanyu Peng. Min Xu, Mingming Sun, and Ping Li. Dataset pruning: Reducing training data by examining generalization influence, 2023.
Xingcheng Yao and Zongmeng Zhang. NIp from scratch without large-scale pretraining, 2021. URL https:/ / / t thub. com homm/yaxxiii wwwooxinnchhng/nen TLM. MIT license.
Xingcheng Yao, Yanan Zheng, Xiaocong Yang, and Zhilin Yang. NLP from scratch without large-scale pretraining: A simple and efficient framework. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings Machine Learning Research, pages 25438-25451. PMLR, 17-23 Jul 2022. URL https: /proceed ngs. ml r. press/ v162/ya022c. html.
Yingbo Zhou, Utkarsh Porwal, and Roberto Konow. Spelling correction foreign language. arXiv preprint 1oXiv:/705.07311 2019.
13

---

Published as a conference paper at ICLR 2024
Appendix
Table of Contents
A Theoretical Work
15
A.1 0 Gradient Problem and KL Divergence Modification
15
A.1.1 KL Divergence Estimator: Recap
15
A.1.2 0 Gradient Problem
15
A.1.3
KL. Divergence Modification
16
A.2 Local Density Assumption
16
A.3
Uniform Start
B Algorithm
18
B.1 Usage
18
B.2 Stopping Criteria
20
B.3 Initial in Gradient Descent
21
Experiments
22
C.I Analytic Checks
22
C.2
WMT14
................................
22
C.2.1 Data
22
C.2.2 G1O
23
C.2.3 Baselines
24
C.2.4 Training
25
C.2.5 Testing
26
Robustness
26
C.4 Speller
27
C.4.1 Data
27
C.4.2 G1O
27
C.4.3 Baselines
29
C.4.4 Training
29
C.4.5 Testing
29
C.5 FashionMNIST
................................
30
C.5.1 Data
................................
C.5.2 G1O
................................
31
C.5.3 Training and Testing
32
D Further FaShionMNIST
33
D.1 G1o KL Divergence vs. Random KL Divergence
33
D.2 Label Distribution
34
Correlation Between Performance Divergence
................................
35
E Guidance for Choosing X, D, Algorithm Parameters
36
E.1
X
36
E.2 D.
36
E.3 Algorithm
36
14

---

Published as a conference paper at ICLR 2024
A THEORETICAL WORK
A.1 O GRADIENT PROBLEM AND KL DIVERGENCE MODIFICATION
A.1.1 KL DIVERGENCE ESTIMATOR: RECAP
Wang et al. (2009) propose KI. divergence estimator based on k-nearest neighbors (kNN) of points drawn from the probability density functions. We recap their derivation to provide necessary context for the eventual modification: Let X { be samples drawn from distribution p and = (YI...mm) be samples drawn from distribution q. Then the kNN estimate of p at X; is:
k
P& (X.) 1 (d)))((()
(6)
Where (i) is the distance from X, to the k-nearest Xjsti, and (d) is the volume of the unit ball the d-dimensional space. Likewise, the kNN estimate of q at X; is:
k
du (X,)
(1)
Where (i) is the distance from X; to the k-nearest Y, and ea (d) is the volume of the unit ball in the d dimensional space. They also propose, by the law of large numbers, that the estimate of KL divergence (p q) is:
DR 11 = 1 Σ
Pkl (X:)
log
(8)
qx (X.)
==1
Putting (6), (7) and (8) together, we get the Wang et al. (2009) kNN estimator for KI. divergence:
1 log
m (1)
DeL(e 11 9)
log
(n 1) p
(ddjug(()
(9)
[dg (1) d log (i)) log n
A.1.2 0 GRADIENT PROBLEM
However, GIO uses the gradient of the estimator to find Vopt by computing 2. DEL (px || PDu(y)). in other words, finding how the value the estimator changes with change in v. We then run into the 0 Gradient Problem with the following scenario. Suppose is far from all points in X. Then, adding v does not change the value of (+) (the distance from X; to the kth nearest point in D) as the closest points in D each point in X the same before and after adding v. Therefore, the only term that changes in the KL divergence estimator is m, constant, and when computing the derivative, the constant goes O. Altogether, this means that for points v that are far from X and their kth nearest neighbors in D, the gradient will be 0 and we cannot do gradient descent. Formally:
Suppose we have X, D and new point and are estimating Dxa (px || PDU(V)). In DRL, vi( is the distance from X; to the kth nearest neighbor in D, and (x(() is the distance from X; to the k-nearest X X Suppose that Vx E X, Vy E D, yk is the kth nearest y D for & < K E {L....). Then, the partial gradient of DRa. with respect to calculated in G1O is:
|| > x Y& || + e for small e, where
DeL. (px PDU(V)) = d Σ log (1) OV
(10)
i=1
As no other terms depend on v. Next, let f(v) = log u (i). We are then fundamentally computing
f(v f(v) However, because Vx E X, Vy D,
lim
v x-
yv ll + e, then vx
with v and (i) with + are the same, and therefore f(v + E) = f(v). But then:
ʃ(v () f(v) c
lim
lim = 0 c-0
(11)
And we have the Gradient Problem. Note: division and addition symbols above are taken element- wise in the vector space.
15

---

Published as a conference paper at ICLR 2024
A.1.3 KL DIVERGENCE MODIFICATION
We now modify the KL divergence estimator overcome the gradient problem. In order to ensure that PA (i) with v and (i) with v + are never the same, we take the average of estimators from to l m (m the size of D), i.. calculate the divergence estimator from to every point in D. In fact, Wang et al. (2009) also recommend taking an average across different values of k, as method to reduce variance. We formally derive this modification. Let vx: be the distance from X; to the kth nearest neighbor in D, and pr(i) be the distance from X; to the l-nearest X jztt- The mean of the estimator (9) across all values of ranging from m is:
m
1. m
DKL l1 g)
log vx log pi(3)] log
m
kin
(12)
1
1.m d vl dlog p (1) log k(n - 1)
m
Given this modified estimator, we show that we no longer have the 0 gradient problem. Let us have the same setup, i.e. Vx E Vy E D, X-VV
Yx 6, and further suppose this true
for all k = I...m. Thus, the derivative is:
m+1
Ov DeL (px IIDUV))))
log UA (1(1)
(13)
Ov
From before, we have shown that for all k == 1...m we arrive at a gradient. However, this expression also has the term log + (6), which mandatorily describes the distance from X, to v, this value describes the distance from X; to the furthest point in D U (v}. Therefore, we also have that
X; (v + e) * X; || and therefore vm++1 (i) with v cannot equal 11) with v + e. As
result, ʃ(v + +) + f(v) and lim ( f O. desired, and avoid the 0 gradient problem. This is the estimator that GIO uses in practice.
We note that with the modified estimator, DKL (px || px) * 0, as the first probability estimator is averaged over & but the second probability estimator In general, this is not an issue, as we are concerned with minimizing the KI. divergence, but not concerned with the magnitude of the values themselves.
A.2 LOCAL DENSITY ASSUMPTION
The derivative trick in GIO finds Vopt and then finds the closest v, G' to add to the selected data, as the closest V, E the solution to the original objective (3):
argmin / Px (x) log
px (x)
g(x, v;) dx
(14)
This holds true under the assumption of local density of the extrema of that integral, or more specifically, the minima of that integral. We outline the assumption:
Local Density Assumption: Let Vopt represent the global minimum of / px (x( log px(x) dx.
g(x, v;)
Then, for the closest v; E G to Vopt. Vb. to also be the solution to the original (14), we assume Evb E G s.t. / px(x)) log (x(x) dx < px g(x, vo)
(x(x) dx VV G.
(x) (x) log
v(x,v,)
(15)
In essence, for the closest V, G to Vopt to also be the best solution, there needs to exist a point in G close enough to the global minimum such that that point represents the greatest decrease in divergence it were added over all other points. the space of X well and has many points in it, which is true if G is large enough, then this assumption is likely satisfied. We note that additionally, we assumed V Vopt to represent the global minimum, but outlined in Section limitations, since gradient descent non-convex space this
practice, this is almost always true; if G covers
16

---

Published as a conference paper at ICLR 2024
is not guaranteed to be true. We leave improvements of finding more global Vopt future work, as described Section 3 limitations. Even with gradient descent, it may be possible to attain more global values by simultaneously gradient-descending different parts of the space and picking the final V; E G based on lowest KI. divergence.
Note that even if the local density assumption is violated, the results are not necessarily catastrophic. Even if we add a point that is the second or third best (or nth best, up to a limit), we are likely to continue attain very close KL divergence values if we followed the ideal optimization trajectory defined by (14).
A.3 UNIFORM START
In order to smooth the convergence of the algorithm when starting from an empty set D ( we propose beginning with uniform random spread of points as a good initialization basis. Intuitively, the uniform prior is the completely uninformative prior (Jaynes, 1968), and therefore does not change the values of Vopt- We prove this property formally. Let Uxx) A be uniform distribution with fixed. x-independent probability A. Then, let us initialize divergence optimization objective is mixture of the uniform prior and whichever distribution of new points the algorithm is adding: 1 + 0z g(x, v). In the regular version (4), the gradient to find Vopt is:
~ U (xx), and then the resulting KL
e1
a / px Px(x) log
Px (x)
/ px(x))
① gV v) logg(x, v)dx / px Px(xx)
dx= g(x, v)
dx
(16)
Ov
Ōv
g(x, v)
With the uniform start expressed as mixture, the gradient to find Vopt is:
① / Px(
log
(x) 1 og g( v)
dx
px(x)
① log
{ A+
1
c2 oolx dx
dv n
(17)
/ px((x)
By
9x v)
dx
B g(x v)
Where we divide by the constant og on numerator and denominator and set B = A C2. constant. The only difference between the gradient in (16) and (17) is that the denominator g(x, v) has constant added to it. But then, it is possible to create x- and v-independent bijection from (16) to (17) via f g(x, v) + gx, v) + B. Since we are optimizing over v and implicitly x defined by Ωt neither of which appear in the mapping. then the value Vopt that optimizes (16) also optimizes (17), and therefore adding a uniform start does not change the value of v Vopt. as desired. Intuitively, function does
A222 another
the uniform prior adds a constant B to g in the derivative, and adding a constant not change the location of the extrema or the shape of that function. We give an example (from self-consistency experiment):
Divergaacce vs. herations
KL Divergence vs. herations
5.94
2.4
6.8
22-
2.0
5.8
633
5.4
0.2
10
52
(a) With uniform start
(b) No uniform start
Figure 3: Left: KL divergence over iterations with uniform start, Right: KL divergence over iterations without uniform start. They follow the same trajectory (except the erratic behavior at the beginning), but adding uniform start greatly improves the stability of the algorithm and convergence.
17

---

Published as a conference paper at ICLR 2024
<L Divergance terations
Target Data and Selecaed Data
seovted bars
Target Data
20
2.0
Figure 4: Sample plots for 2D example (Self-Consistencyy) Left: KI. divergence over iterations, Right: Chosen data and target data
B ALGORITHM
B.1 USAGE
The code package for G1O is available on github gradient-in hiteet--inooommtioooottititmte zation. An example use in 2D (the self-consistency test, in this example):
From GIO import GIOKL
import numpy as np import numpy as jnp
import matpletlib. pyplot plt
Create some data def g)
mean = [3,4]
cov 110.5, 01, 10,0.511 np.randem. see
x, = np. random. mult multivariiat__nrra.. (mean, cov, 100) .T
return jnp. ([x for i in range (len (x)
def getXTest ():
mean [3,4]
cov = [[0.5,0], [0,0,]]
np. randen.. (Sod
x, random. www.mmiaaatat_eneeeeeeeaeaeeee neemall (mean, cov, 100) .I
return anp.array (I[x[1), for range (len (x) ]
getX()
test getXTest
Initialize class
gio kl GIOKL. GIOKL (uniform Low=0, uniform high=8, \
uniform_start_sizenn dim-2)
Ferform Algorithm w, kivsi gio_kl. fit (X_test, x, normalize-Palse) W[100:] # Remove the uniform start
Plot results
plt.plot (kl_divs)
pit.title Divergence vs. Iterations")
pit.xlabel ("Iterations") plt..yabel ("KL. Divergence") plt.shox()
18

---

Published as a conference paper at ICLR 2024
plt.clf()
plt..catter [[each[0] for each Wl, leach for each in w), \
label="Selected Data") www.cattter([cach[00 for each in X], [each [1] for each in x],
labre="Target Data")
plt.title Data and Selected Data")
plt.xlabel ("Dimensso 1*)
2*)
(et.de pit.show
The above code will print the KL divergence at each iteration and produce the plots in Figure 4.
A complex example using the quamtization-explosionn technique big data with Spark:
from GIO import GIOKL
import jax.numpy as inp
import matplotlib, pyplot ppplot import pyspark. sql. functions
pit
Initialize class gio_kl GIOKL.GIOKI =20, dim=768)
uniform_high-i, uniform_start_sie
Read cata
train_df, target_df = gio_kl. read data_frrmm____ee___a_ad crom_ SEATHTTRAIN PATE_ TARGET)
Quantize data
modee_oooaiin model _x, transform_train, transformed_X
gia_k1 qua target, k-1500)
X = jnp. (model X. c.custerrenterr())
train jap. array (model train. clsttetcl
data [íi, each. ()) for i, each in
enumcrate wtw.:lltttaai...stttooooooo spark. www.aakkccccattttaaaaaaaaaaa (data=data,) \
())
roide_d = gio.kl..park.
schema=("id", "centroid"!)
Perform the Algorithm
W, kl kivdivs, [lo_kl.fe fit (train, x, max_iter=300, k-5,
arop__rrttrioo=ssseeeeeete increase j_mn",="jump", 1r=0.01)
Explode back to original data and write resulting data sel__selecttons_eee gio_kl.explode (W, centroids_df, transform_train) full sewections df
write www.eeeeeeen....seeeeee...eeeeee (o ("_c0",, F_col("_el") .wwitt
oppiion ("delimiter", "(t") cssv (OUTPUT_PATH)
# Plot results
plt.plot (kl_divs)
plt.title ("KL Divergence vs. Itcrations")
plt.xlabell ("Iterations")
plt..yabel ("KL. Divergence") sho.
The main function, GIOKL. fit, takes the following arguments:
t train: training data as jop array (jпp is almost identical to numpy) [M, D] shape
X: target data as a jnp array [N, D) shape
D: initial data as a jnp array, default None. Use None to initialize from 0 (uniform) or a subset of training data
k: kth nearest neighbor use in the KI. divergence estimation, default 5
19

---

stop_criterion string for the stopping criterion, one of the following: 'increase', 'min_kl' diza_sizee.. Default is Published as a conference paper at ICLR 2024
max iter: maximum iterations for the algorithm. One iteration adds one point (cluster)
'max_resets', mni__iffference' requential_increaee_ooo 'increase'
min fff fference: the minimum difference between prior and current KL. divergence for mmin difference'stonn stop criterion only. Default is 0
resets_allowe whether if KL divergence increases, resetting G to the full train is allowed (allows the algorithm to pick duplicates). Must be set true if the stop criterion is max_resees! Default is False
max max_resets: the number of resets allowed for the 'max_resets' stop criterion only (a reset resets G to the full train set and allows the algorithm to pick duplicates). Default 2
max data size: the maximum size of data to be selected for the 'data_size stop criterion only, as percentage (of total data) between and Default is min_kl: the minimum kl divergence for the 'min ktos criterion only. Default is
max sts_seuueniaal
increases for the aal_inceesssss the maximum number of sequential KI, divergence stop criterion only. Default is
randon_init__en the percent of training data initialize the algorithm from. Default is
probability any given iteration extend the gradient descent iterations by 3x, to find potentially better extrema. Higher values come at the cost of efficiency. Default is 0
random wttprrrrttttt__teeeeeeeee
scale_factooe factor to scale the gradient by, or ''utt', Default 'auto', which is recommended
v v_init: how initialize gradients descent, one of the following: 'mean', 'prev_opt" 'jump'. Default is 'mean'
grad_desc_iterr the number of iterations to use in gradient descent. Default is 50
Seccaard_eareste For_xy: discard nearest the xy calculation of KL. divergence, for use when X and the train set the same, comes the cost of efficiency. Default False
normali le: Whether to normalize the uniform start values. Use when the values of X and Train are normalized, as when using embeddings generated by MPNet or MiniLM, for example. Default is True
1r: Learning rate for gradient descent. Default is 0.01
B.2 STOPPING CRITERIA
We discuss several possible stopping criteria implemented in the code package and outline the pros and cons of each.
The strictest stopping criterion estopping_crittrriorss criterion==increasss immediately when the KL divergence increases from the previous value. This stopping criterion makes the most theoretical sense, as a KL divergence increase indicates the point being added does not add any information about the target X. This is the criterion we use for our text experiments. In practice however, may wise to allow some tolerance of KL divergence increase, which we discuss under the Sequential Increase Tolerance item.
Strict:
to stop
Convergence Tolerance: Closely related to the strictest stopping criterion is to specify tolerance of the difference between the prior and current KL divergence and terminate when the difference falls below this tolerance (st (stopping__) weriinngnrrierriiiooooo "min di fffeeence "), The tolerance is similar to tolerance arguments for gradient descent algorithms and is designed to heuristically identify when the algorithm has converged and no longer decreasing by more than the specified amount.
The algorithm can terminate when reaches pre-determined KL divergence (stoppin wnw.criterrrr iettreonniiiiee "min_k1") In practice, however, may be difficult to know beforehand what a good value of KI. divergence may be, particularly high dimensions.
Minimum KL Divergence:
Maximum Data Size: The algorithm can terminate when reaches a certain data size. This is particularly useful stopping criterion when data size/resource constraints the primary
20

---

areas of Published as a conference paper at ICLR 2024
reason for data selection (stopping_crrle criterion="max data_size" does not use intrinsic properties of the data that was selected (e.g. KI divergence). it is the least theoretically- eeoreticallymmotivateddd
However, it
motivated stopping criterion. We use this stopping criterion for the FashionMNIST
task in Section 4.4 (see Appendix C.5 for more details)
Sequential Increase Tolerance: Instead of stopping when the KL divergence increases as in the strict criterion, we can allow the algorithm a certain amount of sequential increases in KL divergence before terminated. For example, value of 3 would tell the algorithm to terminate after 3 consecutive points added each increased the KL divergence. If the KL divergence increases, then decreases again. the algorithm can continue. This criterion can allow the algorithm to attain better minima. Primarily, the space is non-convex, and allowing temporary increases in KI. divergence can enable the algorithm to get over certain "humps" and descend into more ideal space. However, it could potentially lead to divergence in KL divergence, and also add many suboptimal points if KL divergence increases happen frequently or the tolerance is set too high.
GIO restricts points chosen to only be unique points. However, in some cases, allowing the algorithm to choose duplicate points can be beneficial. Intuitively, this would allow the algorithm to weight points in certain areas of the space it feels are most important. In practice, we can let the algorithm run until it reaches the chosen stopping criterion, then reset G and allow the algorithm to pick again. We use this concept our spelling correction experiments. In the code package, these options are given by the arguments resets_a sss_llllleeeee (True or False) and res_tss which specifies how many resets are allowed until full termination.
B.3 INITIAL GRADIENT DESCENT
We discuss the initial value of v in the gradient descent implemented in the code package and outline the pros and cons of each.
Mean: The simplest initialization for vo is to set it equal to the mean of the target X at each iteration of G1O (v_ (w_init= "mmem="))) This is motivated by the fact that the optimal Vopt lies somewhere in the space of X, and taking the mean will make Vo close to X and is a good place to start searching. In unimodal symmetric distributions, the mean will essentially be the optimal point to add. In multimodal distributions, the mean provides neutral starting point from gradient descent by which it can choose which mode to strike out for, based on the gradient. However, particularly in high dimensions, the mean may be far from the optimal point. Additionally, in scenarios where the modes are far apart and the distribution (and therefore the mean) is skewed one section, starting from the skewed mean can miss the other sections of the space.
Previous Optimal Point: We can also start Vo Vept. the previous optimal point found by gradient descent
"). This is motivated by the fact that adding new point is unlikely to alter the distribution much, and therefore the next optimal point is likely close to the previous optimal point. This is the setting use in text tasks. However in multimodal distributions with modes that are far apart, this may settle for only one or a few modes and not are biasing the gradient descent to
explore the rest of the space, the space. At the first iteration we can start from the mean.
Instead of a deterministic starting point, we can set Vo to random value from the target X at each point in the iteration. This setting will explore various parts of the space, and algorithms better suited non-convex spaces (e.g. particle swarm optimization (Kennedy and Eberhart, 1995)) also utilize stochastic methods as an instrument to achieve better convergence. Therefore, this setting can yield more diverse values and potentially better optima, with the risk of not being able to fully explore a particular area of the space. We use this technique in the FashionMNIST task, as in that task we know that the 10 classes of images will be multimodally distributed and the modes relatively delineated from each other, a setup which could cause issues with the previous two settings.
Random Jump:
combination of the methods, and additional methods, could perform even better and we leave this ex exporaaion to future work.
21

---

Published as a conference paper at ICLR 2024
Original Data PDF, Quantized Data and PDF
Drigin
Goxtecd
Demamaiinn
Dineniin1
Figure 5: Visualization of quantization consistency. Left to right: Original data and kernel density- estimated PDF, Quantized data and kernel density-estimated PDF. The distributions are very close between the quantized space and original space.
C EXPERIMENTS
provide details of our experiments. All code to replicate experiments is included the github repository http://shhhnooaaarrrrrrdrrrrriit details on usage.
see Appendix B for
C.I ANALYTIC CHECKS
We provide the sample code files to run the analytic checks in our code package at the path gradient- seamaaioroppiiizzziiooooeennnnee
Self-Consisteney: sel.rrsssssesccoosssttttr.y.ee
Quantization-Connisteen quanttiatiionnooosssttttaa
None of these code files require arguments, usage is python FILE py. Figure shows visual- ization the quantization consistency by plotting the estimated probability density functions of the original and quantized space. Figure 4 shows visualization of self consistency plotting the target and selected data.
C.2 WMT14
This experiment aims demonstrate using G1O well-known dataset with well-known setup, and shows GIO can achieve similar and even superior performance with significantly less data. See Section 4.1 for details the results of this experiment. We follow Fairseq's 's recommendations for replicating Vaswani et al. (2017) setup.
C.2.1 DATA
We load and preprocess the data using httppr//ggtht..ccm//cccoooeeeeee fairseq/blob. http:///oooo.oeeeeaaaaaaaaaeeeea EN-DE and httppp//gggttn.cccccceeeeeeeeees
sh for for EN-FR. These scripts handle
examples httpt/taaalatatttto..ooooooo http://annllaaaeoooorrrrreeeeeer downloading, preprocessing, tokenizing and preparing the train, valid and test data. For our experiments, combine the train and valid data from this to use all of ours, competitive methods and random setups, and recut new validation sets for each experiment.
For the dev sets WMT08- -MM00--WMT33-- from the provided official dev sets of ht / /www s statmt. org/ 14t cranslation-F
target X, collect the dev sets for EN-FR and EN-DE from
html and preprocess them using the same techniques and tokenizers given by the Fairseq processing scripts above. We then shuffle and cut 3k from the resulting combined dev sets be used Dev Test set, and keep the remaining 12k target X.
22

---

Published as a conference paper at ICLR 2024
C.2.2 G1O
We first generate embeddings for the train and dev sets using MPNet-Base-V2 (Song al., 2020) on the input side of the data, using the following code from the code package:
from GIO Import
gen_embeas
v2")
gen_embeds www.rggeerrrrebeeeeeeeee
(PATH INNPIT_TTUD_LLOO SAS.
OUTPUT_PATHO
We then paste the original "input It output" with the embeddings file to create file of the format "input M output \t embedding" and use this as our CSV in GIO. We used AWS p3dn.24xlarg machines to generate the embeddings, which have 8 NVIDIA Tesla V100 GPUs and as benchmark, takes roughly 4 hours to generate 15M embeddings (therefore approx. 8 hours for EN-FR and approx. hour for EN-DE). This process is highly parallelizable for speed, and additionally, more lightweight models like MiniLM (which takes roughly half the time) and even non-nee non-neural methods like Sentence2 Vec can used under speed and resource constraints. Across all initializations, we use the following parameters:
K: 1500
k in DKL: 5
Max iterations: 1000
Stopping Criterion: increase
init: prev_opt
Resets Allowed: false
Iterations in Gradient Descent: 50
Gradient Descent Learning Rate: 0.01
uniform start of 20 points, spread from -1 to 1 with 768 dimensions. For 25% and 50% initialization, we start with a random subsample of 375 and 750 clusters respectively, out of the 1500. The following is the quantization and main method signatures for 0% initialization:
For 0% initialization,
Initialize class
ki GIOKL.GICKI ()
Read data train, target gioghh.reedddddaffffooosse///vr1 PATH_TARGET)
Quantize data
model_traing model_X, X, aaansform_ttaine transformed. X
target, k=1500)
jnp. anray (model_x, (lu.lu/(eccuuttrreeterrss)))
train
(model attrrraymddtttttiicccttttttteee
())
data = [(i, each.tolist()) for i, each in
enumerate www.kpa..ccatteeaaooooooo......e wawaatt.reetttataaaaaaaaaa (data=data,)
htl.trainn.custtereeteeesss ()))
kl. renos gio_kl.
schema= "id", "centroid") )
Perform the Algorithm
W, k1 divs, fit (train, X, max_iter-1000, k-5,
stop_oriterioonn"ccceeee p_ivit="prov_pp"./
resets_ aieets_allowed-Faase grad_ desc bresc_iter500 1r=0.01)
w[20:] Remove uniform start
Explode back to original data and write resulting data sull_selections gio full_selections__c select (F.col (ooccol"""" (" ("_c0"), F.col ("_ai") wwite \ .option ("delimiten "\t") csvv (OUTPUT__PAH))
blookk.ae kll.xppldde (R, caansform_traise centroids_df)
23

---

Published as a conference paper at ICLR 2024
WMT14 EN
WMT14 ENIDE
GIO
GIO Randon Full
BM25
Random Full
RL Divergence
KL Divergance
Terations
Iterations
Figure 6:
G1O surpasses the KL divergence of the most competitive baseline, BM25, at around 80 iterations and ends significantly lower
For 0% initialization, remove the variable D (both from when it is assigned to random sample and from the main fit function call). For 50% initialization, increase 375 to 750.
We performed this algorithm in Spark environment instances. For EN-DE, the algorithm took 40min end-to-end. For EN-FR, the algorithm took 50min end-to-end. We saved the initial K means quantized data for reuse in the different initialization states and during experimentation, which saves 10min for DE and 20min for FR on subsequent runs of the algorithm. In Figure 6, we show the KL divergence graphs over the iterations for EN-FR and EN-DE at 0% initialization, and include lines indicating the divergence for BM25, Pruning, Submodular Optimization, Random and Full baselines.
cluster of 30 AWS r4.8xlarge CPU-only
C.2.3 BASELINES
For the random baselines, after rumning GIO, we take random sample the same size as that chosen by GIO. For BM25, we also select sample of the same size as that chosen by GIO. Yao et al. (2022) provide a Github implementation of their data selection technique at https: /tte:://gtttttt /tgtttb.tt/./ggg/ggg///.g..ghh.t com/ yaoxingcheng/TLLY (Yao and Zhang, 2021). We use the data_ citta_selection py, for example:
python szc/data_selection
--inoex_naee example pm25 _m25_s source --
-oorreefille SOURCE. csv \
-arget_file TARGET csv
--oupput_dir /eeample_datt
-outppt__aaee selected.csv
--top_k K
The source file is the (input only) data in this case the collected dev sets. Top K is the parameter that tells BM25 how much data to retrieve for each data point X. Since BM25 can output score of 0 and discard data, we cannot calculate collect enough data to match the desired
select from, G, and the target file is the (input only) X,
as [G] and typically, we need to overestimate size. We also need to pre and post process our data to conform to their code, which we address and provide scripts for.
Preproccessing. Their implementation requires a CSV with columns text,id, but data given with the WMT14 scripts is in plain text format with ID. We provide script imi zat i on wwwn/eepeeiiiee experiments sontts/bat /ttmmmmmaaaa_eeeee 1ece225________ seerptssmmak. csv.pt which transforms single-column file inputs into the format needed by the Yao al. (2022) code:
at
python make_csv.pp csv.py INPUT FILE TO TO_OUTPUT_FP FILE
Postprocessing. Their implementation will output a CSV with the format chosen chwsee_daotuuoeeeoe data,input ID.rank. The rank is the value of "K" for which this data retrieved. For example, if
24

---

Published as a conference paper at ICLR 2024
this was the 5th most relevant to data point, then the rank will read 5. As addressed, need to overestimate K to collect enough samples. Inevitably, we end up with too many sam- ples, and need to trim down the data by removing the least relevant retrievals. We provide a script at gradiens enadieettiit inn-innormee rmat inmaatioon-ie on-opt im zat wo//epeereeeeeee ments/ 1rr25_ssrritss get_less. . py which takes in the file output by their code, output path, and value of rank to filter; we keep the values under the specified rank value, and discard any over. This is a trial-and-error process, repeatedly trying different values of K until we get the desired data size. Usage is:
ent-info
python get_less. py INPUT TO_OUTPUT_FLL FILE SEILLTER-VALUE
Finally, we need to recover the input-output pairs, as the output file only has input-side data. We provide a script at gradient wet--informme on-opt imi zat on/ cettipee/eprrrr ienrimmentss b/mbm25 scripts/ get_pairs, py that retrieves the pairs from input side given the input data and original data of the format input \t output:
python get_pairs.pp py ORIGINAL INPUT_FIIA TO_OUTPUT_FILE
For submodular optimization, we the same cluster centers GIO, the LazyGreedy optimizer and facility location mutual information function, and run to have the same data size as the G10-selected data. We provide a script at gradient -inf-nfo cnnformatio oon-iooooon opt imi zat i wwoioonnepeeeeeeeeeeeeeeeeeeeeee inveriin iimente submod/ s ubmod. py which takes three inputs and an output argument. The first input should be parquet file with one column of array<double> for the data and the second input should be the same for the target X. The third input should be the threshold data size, determined by
Desired Size Cluster Size, and the output should also parquet file with a column array<double>.
Total Data Size
python submed.pu py DATA CLUSTERS CLUSTERS THRESHOLD ab TO_ SELOUUTPUE POT_ROLDERR
For self-pruning, we use the same cluster centers as GIO. We adapt their tech- nique in
htt//expriimeeee wwrpennntt/iiinn////eee/e/ wwwuninngpruuuuu prune. py which takes three imputs and one output argument. The first input should be a parquet file of the cluster centers and the second input should be the transformed dataframe from Spark's KMeans. The third imput should be the threshold, which is calculated by the formula Desired Data Size should be output folder.
The output
Cluster Size
python prune.pp INPUT_CLJSTER
TRANSFORMED DF THRESHOLD PO_CUTPUT_FOLDDS
C.2.4 TRAINING
We preprocess and train using Fairseq (Ott et al., 2019). We cut 3k pairs from the data for validation and keep the rest training data. We use the subword-nnt pack- age to apply BPE Sennrich et (2016); Sennrich (2021) using the codes already com- puted and provided by the Fairseq scripts download data (Appendix C.2.1). We then wwrreee-preproce..
to preprocess the data into Fairseq-readable format. We then follow the Vaswani et al. (2017) setup and train for 300k iterations, saving every 15k iterations. We use Fairseq's built frransoomerr__ vasrranaawaiii _wt_ Lemt_en___ _as____________ fr_birg architecure for EN-FR and trans former_ Former_vaswand vaswani wmt fertt_____r_i__ for EN-DE. The pre-built architectures preset es- sentially all of the model parameters described by Vaswani et al. (2017). We followed the recommen- dation in https: / //git/ //github. com/ ffacebokkeeese aaaarrc/e fai wtwairsrr/iiie rrse//isuues/////eeess /34 set the LR at 0.0007 for close replication of Vaswani et al. (2017) with the Fairseq framework. We varied the max tokens parameter ensure approximately 25k tokens per batch, as recommended in the paper (Vaswani et al., 2017).
EN-FR CLI
(which fairseq-train) DATA_ PREPROCREDDEE
arch transformer_ Saaniii__ttt_____ wwt__n__att__eeeee _n__r_big
share-all-embeddings
optimizer adam --lr 0. 0.0007 \
-dropout 0.1
-max-tokens 3512
--1 -1r-scheduler inverse_sqrt --weight --weight-decay 0.0
25

---

Published as a conference paper at ICLR 2024
--criterion label smoothed cross entropy --1abel-smnotthing 0.1
--fp16 --save-dir SAVE_DIR -aa----t--t--taaa------a------ 15000
--max-update 300000 --clip-ncrm 0. -warmug-updates 4000
-warmup-init-1r 1c-07 --nin-r 1c-09 --adam-betas 10.9,0.90)
EN-DE CLI
(which fairseq-train) DATA_PREPROCESSED
arch 2ransformer__eevaaL_mmt..ee
share-all-embeddingg
optimizer adam --lr 0.0007
dropout 0.3
max-tckens 3512
-- -Ir-scheduler inverse_sqrt --weight-decay 0.0
-criterion label_ htteee__moochec_ccsss_eteeeetee
--label-s -n--abbe--mmoohinng 0.1
--pp1 -iveee-dir SAVE_ -save--nttrrvaa--padeeee 15000
-max-update 300000 --clip-nerm
0. --wa ruup-updates 4000
Searmup-init-11 1e-07 --min-lr 1e-09 --adam-betas (0.9, 9, 0. 98)
We trained on one AWS 24xlare,,, which has NVIDIA Tesla V100 GPUs and takes about 10 hours train the full 300k iterations.
C.2.5 TESTING
We use the script provided by Fairseq moderators in http::::///sttt.i.ub...b... como:.commmnn com/myleott/ daCea 8ee Beeeeeee5e2 75826034b3 generate the predictions on the preprocessed test data from the best model by validation loss, we use checkpoint_ best.pt model produced by Fairseq and the command:
98d5c16 6 (Ott, 2019) for the evaluation process. In order to
fair Fairseq--enerate generate
TEST TEST_DATA --path MODEL -ebeam 4 --lenpen 0.6
--remcve-bpe outfile
The outfile then gets fed into the script provided by the Fairseq moderators give the BLEU score.
C.3 ROBUSTNESS
The robustness experiments follow a nearly identical setup to the above WMT14 setup. We outline the differences:
MiniLM. To use MiniLM to generate the embeddings for G1O, we replace "all Pall-mpnet: mmpett-base-ee2"" -bsse-v2" with Fall-MiniNM-112-111 class instantiation:
the GenerateFEmbeddings
From G10 import generate_texx_embeedii as gte gen eeneems gte. Generateembeddiings ("ail--tiniiMM1122-11 gen embeds. generate embeddings (PATH 2 INPUT SIDE OF DATA, CUTPUT PATH)
MiniLM takes roughly half the time than MPNet to generate embeddings. On AWS p3dn.24xlarg. machine with NVIDIA Tesla V100 GPUs, it takes roughly 30min for EN-DE and roughly hours for EN-FR. In addition, alter the instantiation of the GIOKL class specify dimensions of 384:
gio kl GIOKL. GIOKL (dim=384)
All else stays the
the WMT14 setup.
K=1000 and K=3000, For the experiments altering the value of K, we reuse all the embeddings generated with MPNet for the main experiments. The only difference is to the quantize method
signature with the difference values of K:
model train model_x laansforme_ttraie trans
weo_ki.qaantiii i (train, target k- k=1000) Or k-3000
26

---

Published as a conference paper at ICLR 2024
All else stays the same as the WMT14 setup.
When computing the variance from the base setup for data size, we omit the K=1000 data size since (as covered in Section 4.2) need to artificially subsample to keep the selected data size the same as in the base setup. This is due to the coarser grain K=1000 selecting more data. Interestingly, K=3000 and MiniLM selects roughly the same data as the base setup.
C.4 SPELLER
This experiment aims to demonstrate using GIO mix of high and low quality data, and aims to show GIO selects the high quality data. We use the spelling correction domain for this problem, as the process of noising training data to create training pairs makes it easy to create high and low quality synthetic data by varying the noise. We determine the % of high quality data selected by each method, and also train a model and show spelling performance of our method, random subset, BM25, pruning and submodular optimization and show the results of G10-selected data are better than that trained on the full mix of low and high quality data Table 6. We generally follow the recent work of Jayanthi et al. (2020) for this problem.
C.4.1 DATA
Jayanthi et al. (2020) use the Billion Word Corpus and extract roughly 1M pairs for creating synthetic data. In order to show our method at scale, we download the full Billion Word Corpus from https: /ww.. sta tmt. woooorogoooooo. Im-benchharrd rk and extract 15M data. We add the 1M pairs used by Jayanthi et al. (2020) and deduplicate on label, which results in roughly 14.7M data. When we have the labels, we need to noise them to create the inputs. Jayanthi et al. (2020) implement and use 3 different synthetic data noising methods: random, word and prob. They show "prob" is the best method, followed by "word" and then "random". We use the "prob" method on half the data to create our high quality input-output pairs, and use the "word" method on the remaining half. The "word" method takes in probability that a word will be replaced by a misspelled version of contained in a lookup table, if it exists. Jayanthi al. (2020) use probability of 20%, but make very low quality data, we use a high probability of 70%. They provide a Github implementation of their methods and data at https: //githttb com/ hewspelllll//ppp weeeennnneuueeeee /nnuspeullll (Jayanthi, 2021), which we use to noise the data accordingly. We give an example of the quality difference in Table 5.
Table 5: Right: High Quality "Prob" data. Left: Low Quality "Word" data. High quality data represents reasonable spelling mistakes, low quality data is barely intelligible. GIO selects 73% from the high quality data.
High Quality "Prob"
Low Quality "Word"
is still big poblem though
Is sttel a beig probelma thouigh It appered thiy whir serioulsly ingerd, whe BBC sayd
It apoeared they were not serousay injured, the BBC said
Were they aarried Feb
Were thwey marreid ong Feb More th uus spe tuning inf tome vthe airwaves thanx ef
More of us are tuning in to tha irwaves than ever
The high quality "prob" data is more reflective of real spelling mistakes, and in fact the "prob" method uses context matrix to make applicable misspellings in context. The "word" method at 70%, on the other hand, severely mangles the sentences to the point where we barely make out what they mean, making for very low quality data to train with. After generating half with prob and half with word, we recombine and shuffle create our training data. For our target set, Jayanthi al. (2020) provide 40k pairs of real mistakes by humans taken from the BEA grammar correction corpus (Bryant et al., 2019), which we use as our ideal target X.
C.4.2 GIO
We first generate embeddings for the train and target sets using MPNet-Batr V2 (Song et al., 2020) on the input side of the data; Appendix C.2 for an example, use the same process in
27

---

Published as a conference paper at ICLR 2024
Spelling Correction
GIO
703
BM25
Random
Full
102
Reset
KI Divergenke
093
192
300
230
200
750
Iteraions
Figure 7: Resetting enables the algorithm to make further progress reducing KL divergence.
WMT experiments. On an AWS p3dn, 24xlarge machine with 8 NVIDIA Tesla V100 GPUs, it takes about hours to generate the embeddings for the 15M data. As we mention in Section 4.3, experiment with new scheme. Instead of stopping on an increase, we allow the algorithm to reset G once when the KI. divergence increases, and pick again until the divergence increases second time. We hope this will allow the algorithm to pick not only high quality data, but emphasize (by picking duplicates) points that are also most relevant X. We use the following parameters algorithm:
K: 1500
in DKL:
Max iterations: 1000
Stopping Criterion: increase
init: prev prev_opt opt
Resets Allowed: True
Maximum Resets:
Iterations in Gradient Descent: 50
Gradient Descent Learning Rate: 0.01
The method as follows:
Initialize class
gio_k1 GIOKL.GIOKI.
Read cata
train, target gio kl, read data from_ esv (PATH TRAIN, PATH TARGET)
Quantize data
model_train model_x, transform_train transformed_x
gio_kl.quantize (train, target, k=1500)
x = jmp.array wtr.ey/mmnell..cttstteeetteeerr
train = jnp.
())
data [ii, () for i, each in \
cnumcrate wwoeeelttaa..cllsseeeeeeeerrrr ()))
roids_ = gio_kl wtw.:aaa...cceeaaaaaaatc
(data=data, \
schema= ("id", "centroid"!)
Perform the Algorithm
28

---

Published as a conference paper at ICLR 2024
M, k1 divs, gio kl. (train, x, iter=e100 k=5, stog_cciterioo"""neeeee""" resets_allowed-in max_resets=1,
y_init="prov_opt?" grad_desc_itor-500 1r=0.01)
w[20:] Remove uniform start
Exploce back to original data and write resulting data full sellliseeettonssc. gio kl. ex.explede (R, transform_traie centroids_dhh full www.eeeeeeesseeeeeeeeeeeeeeeeese of select (F col "(t") csv (OUTPUT_PATH))
F. ("___"""")))) write
.option ("delimitens
We performed this algorithm in the same Spark environment on cluster of 30 AWS r4.8xlarge CPU-only instances, which took around 1h15m complete. In Figure 7, show the KL divergence graphs over the iterations and include lines indicating the divergence for BM25, Pruning, Submodular Optimization, Random and Full baselines. We mark where the algorithm reset we can see the KL divergence further decreases.
C.4.3 BASELINES
For the random baseline, after running GIO, we take a random sample of the same size as that chosen by GIo. For BM25, submodular optimization and pruning, please Appendix C.2.3, follow exactly the same process as in our WMT experiments, using the same scripts but on different data. We also train a model on the full data for comparison.
C.4.4 TRAINING
We preprocess and train using Fairseq (Ott et al., 2019). We 3k pairs from the data for validation and keep the rest as training data. We use the subword-mn -mtt package (Sennrich, 2021) to learn and apply a 40k BPE from the combined souree and target data. We then use rseq-prepr cess to preprocess the data into Fairseq-readable format. We train with seq2seq setup proposed by Zhou et al. (2019) that learns to map input (noised) to output labels (regular, no noise). We use the BART base architecture (Lewis et al., 2020) when developing BART is naturally suited to the spelling correction task. We train for 50k iterations and pick the best model checkpoint by validation loss, saving every 5000 iterations. We use 4096 max tokens per batch, the Adam (Kingma and Ba, 2015) optimizer with LR of 0.0001, and generally keep the rest of the settings the same in the BART paper (Lewis al., 2020). For method and all baselines, we train twice and take the average.
seq2seq model, since the denoising objective they used
Speller CLI
(which fairseq-train) SATA_PREPROCESSE
-aarch bart_base
chare-all-emeeddingg
-optimieer adam --lr 0. 0.0001 -ropout
- 0.1
-max-tokens 4096 \
--1 -rrsscheduler inverse_sqrt
-fp16 --svee-dir SAVE_DIR - 5000
-max-update 50000 --clip-norm 0.0 --wa --warmup-updates 4000
-warmup- -warmpp-nit--1r init-lr 1e-07 --nin-rr 1e-09 --adam- betam 1 (0.9, 0. 98)
We trained AWS p3dn. 24xlarge which has NVIDIA Tesla V100 GPUs and takes about hours to train the full 50k iterations.
C.4.5 TESTING
As test data, Jayanthi et al. (2020) provide several sets. We cannot use the BEA 60k they provide used 40k from this target X. Instead, Jayanthi et al. (2020) provide an more challenging set 4660 intentionally ambiguous pairs curated from BEA60k (the 4660 pairs have overlap with our X), and this test set. We tokenize and preprocess this test set with
29

---

Published as a conference paper at ICLR 2024
Table 6: Spelling correction results. Training data size, accuracy/correction rate/FI scores, % of high quality data and KL. Divergence of the full data, the Gio-selected data, competitive methods and random subset. Bold is the best score overall. GIO selects 73% high quality data, outperforms all other methods and sets new state-of-the-art spelling correction.
System
Train Size
Accuracy
Correction Rate
% High Quality
DKL.
Ours
95.9
99.6
97.7
73%
224
BM25
95.6
99.3
97.4
55%
264
Pruning
3.6M
95.7
99.6
97.6
54%
241
Submod.
95.7
99.5
97.6
59%
245
Random
95.9
99,4
97.6
50%
284
Full
14.7M
95.7
99.5
97.6
50%
280
the same codes our training data. For metrics, Jayanthi et al. (2020) report word-level accuracy and word-level correction rate. Word-level accuracy is the fraction of words that the model gets correct, and word-level correction rate is the fraction of words that should be corrected that do get corrected. We also report the F1 score to arbitrarily correct every single word randomly and get 100% correction rate but 0% accuracy, or keep every word the same and get near 100% accuracy but 0% correction rate.
balance between these two metrics. theory is possible
An additional complexity arises with our particular use of seq2seq framework. The models of Jayanthi et al. (2020) use classification objective to classify each input word, which means the output and imput sequence will mandatorily have the same number of words. However, this property not necessarily guaranteed for a seq2seq setup with subword tokenization, and after inferencing models test data words do not match. For fairness, when evaluating the models we remove the few sequences that have a misalignment from the evaluation. In order to generate the predictions on the preprocessed test data from the best model by validation loss, we use the checkpoint_ best,pt model produced by Fairseq and the command:
small number (roughly 1%) where the input and output number of
fairseq
fairsee--enerate generate TEST DATA -path MODEL --beam 1
remove-bpe outfile
grop "H outfile cut -f3- > hypotheses
grep "T outfile cut -f2- > targets
grep "g outfile cut -f2- sources
We then feed the hypotheses, targets and sources into script we provide at pradient-informatie opt zaaizzttinnne wet.oo//xppemmmetee hepperrrmmntssspeleeee. ssssspele_srsriieeeee calculate_scoree py which computes the word-level accuracy and correction rate.
Results G10 outperforms random, all comparative methods, and the full model. G10 outperforms all comparative basclines accuracy, correction rate and overall FI, and matches random baseline on accuracy and outperforms on correction rate and overall F1. It outperforms the full model by +0.2pps for accuracy and +0.1pps for correction rate, despite using only 24% of the data. We set in correction rate and overall Fl score on BEA4660 over the best model reported by Jayanthi et al. (2020) (+2.3 pps Fl and +7.1 pps correction rate, Table 2 BERT model inJayanthi et al. (2020)). See Table 6 for details.
new state-of-the-
C.5 FASHIONMNIST
This experiment aims to demonstrate using G1O can effectively reduce the training data size under a resource constraint without big drops in performance. We use the FashionMNIST image recognition task (Xiao et al., 2017) to demonstrate using GIO outside language tasks. Our goal to reduce the train set size to 25% of the total size. Experiment results are available at Section 4.4.
C.5.1 DATA
The FashionMNIST task has 60k 28x28x1 images for training and 10k images for test- ing. We load the images in CSV format from https: http///w.kkkgtg/ com/datasets.
30

---

Published as a conference paper at ICLR 2024
(Pankaj. 2019). In dataloader, resize im- ages to 224x224 for Resnet50 using skimage's 's resize function, normalize each channel, and addi- tionally normalize the values to fall within 0 and 1. The dataloader code is available in the main file we publish at gradiems eeiient--- - n cenformme orma i on-opt- i imm zat i on /eepeexprriiie mentss/ experiments/ mmistt
zalando-reseale sallancd-reeearcchEEsiioooaa
fashion. py. We based the dataloader code code provided at https: /tt::://wwww kaggle. com/ code / /ppnkajj///////p/ /aanaaj/ffshiii fashiog on mnist with waath-pptorcs. 93 acccuacyy,, We use the vec- GIO. G the training set, but this case, is also the
format of the images themselves training set. Since our goal is to reduce the training set size, we want to select subset of G that contains the most information about the whole set, so the target set is the whole itself.
C.5.2 GIO
In order read the images using the code package, do simple processing to make the image CSV into a string representing an array, the script is available iondient-informe ieptatiiooonninnnnnee optimiiation//oiooiiiiiaiiiciiit wttp////xpeiiieteeeeeeee fassiiionn mnist/css to to_image.ppi py; first argument is input, second is output. We also normalize each channel and norm each image. For the algorithm itself, we stop at 25% data size, and K=1000. With this problem, we know that there are 10 classes and the images are likely to be clustered relatively distinctly, by class. Therefore, for the initial in gradient descent, we choose the random jump scheme, as we will therefore be able to cover more of the space. We use the following parameters in the algorithm:
K: 1000
k in DKL:
Max iterations: 300
Stopping Criterion: data_size
Maximum Data Size: 0.25
init: jump
Iterations in Gradient Descent: 50
Gradient Descent Learning Rate: 0.01
The method is as follows:
Initialize class
gio_k1 G1OKL.
Read data
train, target gio_ki kl, read data data_fron, from_ eev [PATH_TRAIN, PATH_ TARGET)
Quantizc data
trdel model_x, transform_train, transformed_X
quanki. (train, target, k=1000)
X jnp.array (mmdddee__..cussterceeetrr))))))
train = jnp.array. (wwod/l_taainncllllll wtw.aii.cluuserrreeeeerrr ())
data (i, toch.tolistt) for i, each in
enumerate (model train. claatenncssttreeteeers ())
reias di = gio kl. spark. www.eeeteeeeeeeaaaaeaeaeeeee
(data=data, \
schema- (idd,, "centroid"!)
Ferform the Algorithm
M, divs, gio_kl. fit (train, x, max_iter=300, k-5,
stop_critee wwtpetterriinnaattiiiiaaaa
v_init=""ump", resets_all assts_alloowdd=allse, grad grad_desc_iter-ft 1r=0.01)
w w[20:] Remove uniform start
Expl ode back to original data and write resulting data sil_seeections gio_k1.exploo (w, transform_train centroids_dd)
31

---

Published as a conference paper at ICLR 2024
FashionMNIST
GIO
Random Full
850
Diverrennce
800
780
Iterations
Figure 8: GIO approaches the KL divergence of the full data
full_select.onn_dffeelll (F .F.col ("_c0"), F.col("_c1") write optiion ("delimiter", "(t") csv (OUTPUT_
We performed this algorithm in the same Spark environment on cluster of 30 AWS r4.8xlarge CPU-only instances, which took around 20m complete. We show the KL divergence iterations. In this case, the minimum KL divergence is not 0, but 739. We explain why the KL. divergence is not 0 in Appendix A.1. The graph of KI. divergence is shown in Figure 8.
the
Baselines. For the random baseline, after running GIO, we take a random sample of the same size that chosen by Gio. We also train model full data for comparison, but the goal in this setup is not to match or exceed the full model, but minimize the drop in performance under the resource constraint.
C.5.3 TRAINING AND TESTING
We finetune pretrained Resnet50 model with the data. For GIO and random, we cut 1,700 validation and leave the remaining 15,000 (25%) as training data. For the full model, we follow this percentage and cut 3,700 as validation and keep 56,300 as the training data. In order to get the images back into the correct CSV format for the dataloader from string string-array format, we provide a script at gradi ent ent inoo-iatffrrtt.t.mmmmmmmmm on- on-opt mi zat on /eepeerrrr ments/ /t: ma ge_tooe___ooooooooooooooooo csv py; first argument is input, second is output. We edit the architecture of the Resnet50 model on the input layer to take channel instead of 3, and edit the output layer output to 10 classes. We finetune for 5 epochs with Adam and small learning rate of 0.00005 so not to overwrite the existing weights in the pretrained model. We use batch size 100 and take the best checkpoint by validation loss and report the test score on this. The code to read and train a model is available at gradi grrdient-iing eent-information iittioonoot imi zat inn/ fwweentss/aahhoon_mmir st /ettttrainnn.. i.r.iii......................... py. We based this training script on code provided Exper ments! ipynb (J, 2020). Usage:
at
python train.py TRAIN_CSV TEST_CSV VALID_CSV
For testing, the above code will output the test accuracy at each epoch. We train on one AWS p3dn.24x 24xlare machine, which has 8 NVIDIA Tesla V100 GPUs and takes around 15min train.
32

---

Published as a conference paper at ICLR 2024
Random
120
Traser
Peasce
Cost
Cops
Siat
-25
Sheaker
Adi bodt
-50
-320
-10
-35
-50
-15
25
30
75
Figure 9: Left: tSNE of Gio, Right: tSNE of Random. GIO is more separable and dense in each class, indicating G1O chooses more core data and fewer outliers
GO Bags
Random Bags
Figure 10: Left: Chosen "Bag" images of G1o, Right: Chosen "Bag" images of Random. Rectangular shaped bags form the core of the full bag class (by visual inspection). The bags chosen by G1O are consistent and similar, and are all rectangular in shape and similar sizes. Random bags, however, also include variety of more uncommon outlier shapes, including a purse and long, thin bag. Under a data constraint, it is more important get the core rectangular bag correct, as G10 chooses.
D FURTHER FASHIONMNIST
D.1 G1O KL DIVERGENCE VS. RANDOM KL. DIVERGENCE
It is interesting that despite achieving higher KL divergence than the random baseline, GIO performs better. We hypothesize that this is due to the fact that G1O is optimizing explicitly for choosing images that add the most information about the whole, i.e., discarding outlier data and focusing on "core" " data under the size constraint, which leads to higher KL divergence than random but better performance, the model will be less confused. For example, in the "bag" category, conceivably good portion of the bags look very similar, but there is also a portion that are more strange. Random would choose both the core portion of bags that looks similar but also include some of the stranger examples. GIO, on the other hand, will focus almost exclusively on choosing the core portion, as best represents, the bag category overall. In short,
that is the set that adds the most information, outliers are less important to know under size constraint, and the core is more important to get right, which GIO chooses. We validate this hypothesis three ways.
2D Embeddings. First, we perform a (SNE dimensionality reduction on the random set and G10 set in Figure 9. The diagrams show that G1O selects more core, densely clustered data and leads to better separability between classes. It does not choose much outlier data for each class that is spread around the space and could confuse the model. Random, on the other hand, is generally less separable and more sparse, and includes potentially confusing points in each class that are spread around the space and overlapping with other classes, which could confuse the model. In short, Gro-selected data is more dense and separable than random and is likely to lead to better classification performance, but will have higher KL divergence from full data than the random data because of this effect on the distribution.
33

---

Published as a conference paper at ICLR 2024
Table 7: Minimum Euclidean distance between the mean of each class and the closest other class (i.e. closest neighbor). Higher is better, as it indicates the class is further apart from others and therefore more separable. G1O higher in 8 out of the 10 classes, indicating G1o-selected data is more separable than random images.
Class
GIO
Random
922
962
1123
1044
604
600
1123
1044
604
600
1102
1005
723
664
7
1102
1005
1269
1361
1660
1570
Table 8: Distribution of labels per class. Despite only considering inputs, G10 also recovers the label distribution the target X.
Class % of Total
4
5
14.1%
10.2%
13.3%
10%
8.5%
10.9%
7.1%
11.2%
10.6%
Visualization. As an additional sanity check, we visually inspect the images chosen by G1O and random images in Figure 10. We show the bag class example. G1O selects more consistent bags that conform to each other. They are all rectangular and similar in size, and represent the core of bag images, as regular rectangular bags form samples from the full set). Random images, on the other hand, contain more strange images like triangular bags and tote bags and are less consistent. Under data constraint, GIO focuses on the core rectangular bags that best represent the class and discards outliers, which leads to better performance.
of the class (as determined by visualizing
Separability. Finally, we examine the separability between classes as an indicator whether G10 selects images that are more core and less spread/confusing. We compute the Euclidean mean for each class and compute the minimum distance between that mean to the means of all other classes, for GiO-sclected data and random images. Table 7 shows the results, and demonstrates that G1o classes are further away from each other out the 10 classes. This implies GIO selects more distinct and less confusing points, and likely explains the improved classification performance random images.
D.2 LABEL DISTRIBUTION
One of the core assumptions of the method is that we can utilize only the input side of data without consideration for labels. By reconstructing an approximate X using G with only the input side, we implicitly reconstruct the label side of X as well, and therefore there is no need to consider the labels explicitly. We verify this assumption with FashionMNIST. as has distinct, quantifiable labels unlike the text generation tasks. Table 8 shows the distribution of labels recovered by GTO, and shows that we approximately recover an even 10% for each class, which the distribution of the full X. Therefore, GIO recovers the distribution of labels in X implicitly, despite not including them in the algorithm, as desired. This is not to say labels are not important in data selection. We note that for the algorithm work well, there should exist quality relation between input and label. Specifically, we require that either input and label are good quality, or input and label are bad quality. good quality input maps wrong bad quality label, may recover bad labels.
34

---

Published as a conference paper at ICLR 2024
KL. Divergence
Clasitication Ero Fane
Class Riation Engr Rata
Best Rt
Emtr
Eme
Divergence
Figure 11: Left: KL divergence and Classification Error Rate over different data sizes. The shape of the two curves are very similar. Right: KL divergence and Classification Error Rate. High 12 = 0.98 indicates KI divergence accurately predict performance.
Table 9: Performance of the model using the different datasizes chosen as the algorithm progresses.
Iterations/Data %
Size Train/Valid
Accuracy
DKL
50/5%
3,000/1700
82.7%
793
100/10%
6,000/1700
88.1%
781
150/15%
9,000/1700
90.0%
775
200/20%
12,000/1700
90.6%
772
250/25%
15,000/1,700
92.1%
771
D.3 CORRELATION BETWEEN PERFORMANCE AND KL DIVERGENCE
We know that under the FashionMNIST problem setup for reducing training set size, a data size was chosen beforehand, which in practice would represent budget. However, beyond budget concerns, we endeavored to find out if we could use the KL divergence to predict what a good, convergent data size would be. Perhaps 25% more than enough and we could do just as well with less, or perhaps 25% is too little to get good performance, and we might want to increase our resources before experimenting.
this setup, we know the KL divergence values at each data size (each iteration), and we want to see whether the shape of the KL divergence curve could indicate the shape of the performance curve at the different data sizes. We repeat the same experiment described Section 4.4, but using 50 iterations (5% data), 100 iterations (10% data), 150 iterations (15% data) and 200 iterations (20% data). To keep the validation set at 1700, if there is not enough data left after selecting the right train size, we supplement the set with training data points from the general set. We plot the results of both KI. divergence and performance (classification error rate) in Figure 11 and describe the results in Table 9. In Figure 11, additionally plot error rate function of KL divergence and compute the coefficient of determination R² We find that the shape of classification error vs. data size and KL divergence vs. data size are very similar. Additionally, 1i2 0.98, which indicates high degree of correlation between KL divergence with data size and performance. Theoretically. this implies that we can use the KL divergence pick an ideal data size we are under budget constraints: where the KI. divergence curve flattens, we are unlikely to receive much gain in performance from increasing the data size further and can use that data size with the expectation of decent performance. We only have few datapoints to support this hypothesis, however, and we leave exploration of this property the algorithm to future work.
35

---

Published as a conference paper at ICLR 2024
E
GUIDANCE FOR CHOOSING X, D, ALGORITHM PARAMETERS
We provide a brief note on choosing X, D and various algorithm parameters.
E.1 X
The target set must be chosen carefully, as this is the distribution the algorithm will match to. If misrepresentative or biased, the chosen train set will be as well. For systems where the model may be used in production whose values are known (e.g. on search engine's 's queries), it makes sense to pick X as representative sample of that production traffic and align model training distribution to that. If the algorithm being used to select high quality, aligned, or unbiased data, it makes sense for to be a carefully curated set of
data.
is equally important to make X as broad as possible within the desired target. X misses a portion of the space (e.g. misses long search queries, for example), then the chosen data will also not capture that area of the space and the resulting model will be prone to poor generalization.
E.2 D
The initialization is versatile. It can be set that known to be high quality, like data, or it can be initialized as a random subset of G. Generally, initializing D) as a random subset of G will give better results when the data size is smaller overall (e.g. our EN-DE WMT experiment), but matters less when the data size is large (e.g our EN-FR WMT experiment).
E.3 ALGORITHM
We have covered in B many considerations for various aspects of the algorithm, in particular the stopping criterion. In general, we recommend trying different parameters and observing the intial convergence of the KL Divergence, as this is cheap to do once the K Means has taken place (we recommend saving those so doesn't 't need to be recomputed). For some tasks and data (e.g. WMT) the convergence is smooth and strict criterion is good enough. For other tasks (e.g. FashionMNIST) the convergence has some ups and downs and serves to choose another criterion like max_ Leos__squeetalliiiie innceeaseseeeess or simply terminating the algorithm when the KL divergence bottoms out. the algorithm diverges, we recommend increasing the uniform start size as well as decreasing the learning rate if it is set high.
For the i_in parameter, "mean" will reset the algorithm to the mean of X to start gradient descent every time. We recommend setting higher learning rate with this parameter, algorithm to reach different points in the space (since starts from the same point every time). This value is good for when the target distribution may be tightly clustered together (e.g. WMT). Similarly, "prev_opt" will start from the previous optimal point and is suitable for tightly clustered distributions where expect the new optimal point not to be far from the previous one. We recommend medium learning rate for this value. On the other hand, using the "jump" value is more suitable for exploring the space and adding diversity, particularly for multimodal or spread-out distributions (c.g. FashionMNIST), We recommend a medium learning rate with this value as well.
this enables the
As for resets like in the spelling correction experiments, we recommend enabling this when it is believed that having multiple copies of the same data point is useful for learning. This may be when we want to weight certain datapoints or regions of the space more (which will let the algorithm choose on it's is own how best to weight the space).
36


Reviews:
[{"review_id": "W0ItgCCO8d", "reviewer": "Reviewer_u6CK", "summary": "The paper presents a new approach for data selection called Gradient Information Optimization, based on a natural information-theoretic objective of minimizing the KL divergence between the sampled and target distribution. The paper then evaluates the approach thoroughly in a number of different settings.", "questions": "Other clarifications:\n- In 4.2, \"GIO works with different embedding models\": are the selected examples themselves similar when using different embeddings?\nMore broadly, why should GIO be invariant to this choice though? For example, it's plausible that certain embeddings focus on certain features more than others, distorting the distribution.\n- Did you try using simpler feature spaces, such as the n-grams one considered in Xie et al.?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Data selection is an increasingly important problem, and the paper presents a new sound approach\n- Overall, very polished and well-written; communicates their contributions clearly. Also, placed well within the context of related work\n- Mostly thorough evaluation of proposed approach, including self-consistency checks, ablations, and different domains (language and vision)", "weaknesses": "- It's not clear to me if the lack of hyperparameter/stopping criterion (which is argued as a strength) is real and if actually a strength. \nFor example, in the self consistency section: there is still an implicit hyperparameter for GIO  when the authors denote \"d(pX(x),pG(x))≫0\". How do you define what is too big of a distance? So it seems as if it's a bit misleading to present the lack of the dataset size parameter as a strength, if the hyperparameter just shows up in a different form. Also, intuitively, I think it feels that there should naturally be a hyperparameter, as the optimal dataset really depends on the different constraints you have (whether that be size or something else).\n\n- Lack of comparison to a similar recent approach: How does the approach here compare (both conceptually and empirically) to that of Xie et al. [1]? They also use a notion of KL distance over a feature space. Not sure if it's quite right to list their approach under one of the \"heuristic\" methos.\n\n[1] Xie, Santurkar, Ma, Liang. Data selection for Language models via Importance Sampling.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YKNcpHDzrd", "reviewer": "Reviewer_C7wm", "summary": "This paper proposes GIO, a new method for selecting subset of the training data in order to reduce data/compute cost while retaining performance similar to the original dataset. It is a task and domain agnostic scheme that requires a small set of unlabeled examples representing target distribution. It minimizes the KL divergence between the target distribution and subset in an iterative manner along with some approximations that make the procedure tractable. Empirical experiments show that this approach is competitive to the existing baselines such as submodular optimization and data pruning.", "questions": "-  In table 4, why is D_KL lower for the random method than the proposed scheme, also random D_KL is much closer to Full dataset?\t\n-  Table 4 does not include other baselines for subset selection in the classification setting?\n-  How do you define high quality in Table 3?\n-  Have you compared the training cost of the proposed method against other baselines?\n- How does the method perform when scaled to larger datasets? FashionMNIST seems too small a dataset (even random subset is not that bad in performance)?\n-  Have you done any ablations to see the impact of target state X and the initial state D? How does one go about initializing these for a new task/domain?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "- task agnostic and domain agnostic approach to dataset pruning \n- does not require labels to apply GIO", "weaknesses": "- Algorithm 1 is incomplete on its own. It does not define the inputs X, D, G. How does one find v_opt? Which equation are you referring to? Ideally, one should be able to infer the algorithm by reading this procedure.\n- No benchmarking for the compute cost of the proposed method against baselines (submodular optimization, BM25, etc.). It is hard to judge how computationally expensive this approach is compared to various baselines?\n- Its unclear if the proposed method could scale to large datasets (for instance on Image recognition task, only FashionMNIST is used)\n- It could be challenging to select target state X and initial state D", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "g4sRG1LhyF", "reviewer": "Reviewer_Km99", "summary": "This paper proposes a scheme to perform training data subset selection, connected to ideas in information theory by minimizing the (approximate KNN) KL divergence between the selected subset and the original dataset. However, this is intractable in practice, and the authors approximate this by finding the optimal data point (not necessarily contained within the training data) to add and then projecting this onto the nearest point contained within the training set. This only scales with the number of gradient steps (and not the size of the original dataset), which is much more computationally efficient. They run this gradient trick over K-means clusters rather than individual data points to make the problem more tractable.\n\nThe authors evaluate the proposed approach through multiple different qualitative and quantitative metrics: (1) self-consistency - coming from the same distribution, (2) negative-consistency - essentially ignoring outliers in the data distribution, (3) the selected subset should maintain similar values of KL divergence given a quantization of the data (i.e. K-means).", "questions": "1. Why were the other baselines not included in the FashionMNIST experiments?\n2. See the weaknesses section regarding the stopping criteria in Section 4.3.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "* The authors propose various methods/approximations to select a subset of training data with greater computational efficiency, and these approximations are tied to an underlying naive approach based on information theory. In addition, the method is simple and easy to understand.\n\n* The paper is well-motivated and solves an important issue; improving the efficiency of pretraining (through shrinking the dataset size) makes the training of large models more available to researchers at lower costs.\n\n* This approach sees good empirical results on multiple machine translation tasks and an image classification task, and the authors' intuitions are also demonstrated in synthetic experiments.\n\n* The authors also provide ablations to demonstrate the proposed method is not very sensitive to different underlying embeddings or the hyperparameters of the K-means quantization.", "weaknesses": "* I’m a bit confused as to the stopping criteria in Section 4.3. This seems to come out of nowhere and be rather ad-hoc. Why do you introduce the second stage of reselecting instances from G until the KL divergence begins to decrease? Why couldn't you simply run the first stage here?\n\n* Not a big issue, but I also think the presentation in Section 4.3 is a bit confusing. The first sentence of the section seems to claim that GIO selects a subset of high-quality data with respect to a target set that is a combination of both high and low-quality data (which is confusing as something that minimizes the KL divergence should select both). However, this isn’t the experiment being run; the target set $X$ here is only a dataset of high-quality data.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rqlUc5g9Vh", "reviewer": "Reviewer_NSKq", "summary": "The authors propose a method called gradient information optimization (GIO) to select a subset of a large training dataset that preserves model performance. GIO is formulated as a problem of minimizing the KL divergence between the subset and the target distribution. As an algorithm for this minimization problem, the authors propose an algorithm using the gradient of the KL-divergence based on the greedy method.The authors experimentally demonstrate that the proposed method outperforms existing methods on the tasks of NLP and image recognition.", "questions": "1. Could you discuss the gap between the solutions of equations (1) and (2), on the theoretical side or on the experimental side? For example, on the theoretical side, if submodularity holds for the objective function, the approximation rate is bounded. On the experimental side, could you evaluate the rate of approximation between the solutions obtained by equation (1) and equation (2), and algorithm 1?\n2. Could you discuss what real-world situations where target distributions can be obtained? Could you also discuss what kind of data should be set as target distributions? For example, in the experiment, to improve the performance of the WMT14 test set, the authors set the WMT08-13 dev set as the target instead of the WMT14 training set. As a result, they report that the performance of the WMT14 test set actually improves. Could you discuss this result, for example?\n\nComments\n1. $\\Omega$ is not defined.\n2. In Algorithm 1, $D_c \\gets D_c + \\\\{\\mathbf v_b\\\\}$ should be $D_c \\gets D_c \\cup \\\\{\\mathbf v_b\\\\}$.\n3. The definition of relaxed $\\mathbf v$ and $\\mathbf v_{opt}$ is ambiguous. $\\mathbf v \\in \\Omega$?\n\n\n## Post-rebuttal\nMy main concerns were (1) about the gap between Eq.1 and Eq.2 and (2) about the usefulness of the proposed method in real-world tasks.\n\nThere is a gap between the optimization problem (i.e., minimization of KL-divergence; \"information-theoretic approach\") and the greedy algorithm in this paper.\nThe relationship between them was not clearly discussed.\nIn their response, the authors clarified that the objective function lacks certain properties, such as submodularity. \nUnfortunately, they couldn't demonstrate the (empirical) approximation ratio achieved by the greedy algorithm. \nHowever, it is acknowledged that the greedy algorithm is one of the natural approaches to discrete optimization.\n\nIn their experiments, the proposed method performs well compared to baseline methods, including the one using submodular optimization.\nFurthermore, the authors' response highlighted that the proposed method outperformed existing methods in an experiment more closely resembling realistic scenarios.\nFurthermore, this paper proposes an effective algorithm, which has actually been applied to more than 10M data, while the minimization of KL-divergence by naive greedy algorithm is intractable.\nThe response to C7wm states that the proposed method is computable in ~40min for a dataset containing 35M data.\n\nBased on the above points, I agree that this paper holds significant value for the community, prompting me to revise my score from 5 to 8.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "It seems reasonable to select the subset so as to minimize the KL divergence.The authors propose a feasible algorithm for this minimization problem.Experimental results suggest that their proposed method works better than existing methods for machine translation and spell-correct tasks.", "weaknesses": "1. There is a gap between the two equations, i.e., the solution obtained by the greedy method in equation (2) and the optimal solution in equation (1) are different. However, the authors did not discuss the gap between the two solutions.\n2. The subset selection given the target distribution is not sufficiently motivated and GIO seems not to be suitable for the natural setting of subset slection, $D=\\emptyset$, $G=X$. For this setting, $V=X$ (selecting all data) is the obvious optimal solution to equation (1) as the authors state in Section 3.4. \n3. It is not clear whether $\\mathbf v_{opt}$ obtained by the gradient method and the optimal solution of equation (3) match. From appendix A.2., the authors assume that the two solutions match if the data density is sufficient. However, from Algorithm 1, the data are selected only from k-means centroids(G_c). Their assumption does not seem to hold, especially in high-dimensional spaces.There is not enough discussion about this.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)