Computational Budget Should Be Considered in Data Selection

TL;DR

The paper addresses the gap that data selection methods typically ignore training compute budgets. It introduces CADS, a compute‑budget aware data selection framework formulated as a bilevel optimization where the inner loop trains under a fixed budget $C$ and the outer loop optimizes data selection; it develops a Hessian‑free, probabilistic gradient estimator and a penalty‑based relaxation that reduces inner optimization to a one‑dimensional surrogate, enabling efficient training. CADS offers two practical variants, CADS‑E and CADS‑S, for example‑level and source‑level selection, respectively, and demonstrates up to $14.42 ext{%}$ accuracy gains and $3$–$20 imes$ speedups across vision and language benchmarks. The work shows that accounting for compute budgets yields more effective data subsets, enabling robust performance under limited resources and scalable data selection for large‑scale tasks and instruction tuning.

Abstract

Data selection improves computational efficiency by choosing informative subsets of training samples. However, existing methods ignore the compute budget, treating data selection and importance evaluation independently of compute budget constraints. Yet empirical studies show no algorithm can consistently outperform others (or even random selection) across varying budgets. We therefore argue that compute budget must be integral to data-selection strategies, since different budgets impose distinct requirements on data quantity, quality, and distribution for effective training. To this end, we propose a novel Computational budget-Aware Data Selection (CADS) method and naturally formulate it into a bilevel optimization framework, where the inner loop trains the model within the constraints of the computational budget on some selected subset of training data, while the outer loop optimizes data selection based on model evaluation. Our technical contributions lie in addressing two main challenges in solving this bilevel optimization problem: the expensive Hessian matrix estimation for outer-loop gradients and the computational burden of achieving inner-loop optimality during iterations. To solve the first issue, we propose a probabilistic reparameterization strategy and compute the gradient using a Hessian-free policy gradient estimator. To address the second challenge, we transform the inner optimization problem into a penalty term in the outer objective, further discovering that we only need to estimate the minimum of a one-dimensional loss to calculate the gradient, significantly improving efficiency. Extensive experiments show that our method achieves performance gains of up to 14.42% over baselines in vision and language benchmarks.

Figures (14)

• Figure 1: Validation loss of models trained on low- and high-frequency data under varying compute budgets. Two regimes emerge: low-budget favors low-frequency data, while high-budget favors high-frequency data, showing compute’s effect on data choice.
• Figure 2: Subset size evolution during optimization with Bilevel-CADS starting from different initial subset sizes.
• Figure 2: Accuracy (%) comparison of different data selection methods with various compute budget.
• Figure 3: Evolution of the sampling ratio over training iterations under different compute budgets: 90,000, 135,000, 180,000, and 225,000 samples respectively.
• Figure 4: Sampling ratio evolution on DomainNet