GIO: Gradient Information Optimization for Training Dataset Selection

TL;DR

GIO tackles the data-selection problem by seeking a subset V of a large pool G that preserves information about a target distribution X, formalized as minimizing the KL divergence between p_X and p_{D∪V}. It introduces a scalable gradient-based relaxation that optimizes over continuous representations and then maps the optimum back to a discrete subset from G, enabling out-of-the-box applicability across text, vision, and beyond. Across machine translation, spelling correction, and image recognition, Gio achieves performance comparable to or exceeding models trained on full data while using substantially smaller subsets, and shows robustness to embedding models and quantization parameters. The authors provide theoretical checks (self-consistency, negative-consistency, quantization-consistency) and an open-source pipeline, underscoring Gio’s practical impact for efficient, high-quality data selection in resource-constrained settings.

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

Figures (11)

• Figure 1: Visualization of the Quantization-Explosion Process. From left to 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 2: The leftmost graph shows $X$ and $G$, which come from the same distribution. The second graph shows that Gio recovers nearly all of $G$ (self consistency). The right two graphs compare Gio 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 a range of points outside the circle, which is suboptimal (fourth graph).
• 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 a uniform start greatly improves the stability of the algorithm and convergence.
• Figure 4: Sample plots for 2D example (Self-Consistency). Left: KL divergence over iterations, Right: Chosen data and target data
• 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.