Data Agent: Learning to Select Data via End-to-End Dynamic Optimization

Suorong Yang; Fangjian Su; Hai Gan; Ziqi Ye; Jie Li; Baile Xu; Furao Shen; Soujanya Poria

cs.LGMar 8, 2026

Data Agent: Learning to Select Data via End-to-End Dynamic Optimization

Authors: Suorong Yang, Fangjian Su, Hai Gan, Ziqi Ye, Jie Li, Baile Xu, Furao Shen, Soujanya Poria

Organizations: National University of Singapore · Nanjing University · Nanyang Technological University

Abstract

Dynamic Data selection aims to accelerate training by prioritizing informative samples during online training. However, existing methods typically rely on task-specific handcrafted metrics or static/snapshot-based criteria to estimate sample importance, limiting scalability across learning paradigms and making it difficult to capture the evolving utility of data throughout training. To address this challenge, we propose Data Agent, an end-to-end dynamic data selection framework that formulates data selection as a training-aware sequential decision-making problem. The agent learns a sample-wise selection policy that co-evolves with model optimization, guided by a composite reward that integrates loss-based difficulty and confidence-based uncertainty signals. The reward signals capture complementary objectives of optimization impact and information gain, together with a tuning-free adaptive weighting mechanism that balances these signals over training. Extensive experiments across a wide range of datasets and architectures demonstrate that Data Agent consistently accelerates training while preserving or improving performance, e.g., reducing costs by over 50% on ImageNet-1k and MMLU with lossless performance. Moreover, its dataset-agnostic formulation and modular reward make it plug-and-play across tasks and scenarios, e.g., robustness to noisy datasets, highlighting its potential in real-world scenarios. Code is available at https://github.com/Jackbrocp/Data-Agent.

Figures & tables

Figure 1: (a) End-to-end dynamic data selection. Existing methods often rely on handcrafted, task-specific static heuristics to estimate sample importance, limiting the scalability across learning paradigms. In contrast, our framework formulates data selection as a learning problem and jointly optimizes it with model training in a plug-and-play manner, forming a closed-loop, training-aware selection process. (b) Illustration of data points prioritized by difficulty and uncertainty signals. The uncertainty signal concentrates on the inter-cluster boundaries and transitional regions, while the difficulty signal focuses more on the sparse cluster areas.

Figure 2: The framework of the proposed Data Agent. At each training stage, the agent observes the model state and derives reward signals from standard forward passes. These signals are combined using an adaptive weighting mechanism to guide a PPO-based actor-critic agent, which learns the selection policy. The selected data is used in subsequent training, forming a closed-loop training pipeline where data selection evolves alongside model optimization. Notably, the modular reward design enables the framework to be easily adapted to various learning paradigms.

Dataset	CIFAR-10			CIFAR-100			Tiny-ImageNet
Whole Dataset	95.6			78.2			45.0
Selection Ratio (%)	30	50	70	30	50	70	30	50	70
Static selection methods
Random	90.2	92.3	93.9	69.7	72.1	73.8	29.8	37.2	42.2
EL2N ( Paul et al., 2021 )	91.6	95.0	95.2	69.5	72.1	77.2	26.6	37.1	44.0
GraNd ( Paul et al., 2021 )	91.2	94.6	95.3	68.8	71.4	74.6	29.7	36.3	43.2

Table 1: Comparison with state-of-the-art baselines. All methods are trained using ResNet-18 on CIFAR-10/100 and ResNet-50 on Tiny-ImageNet. Random* refers to randomly selecting samples in each epoch. Some results are from ( Qin et al., 2023 ) .

Method	Herding	EL2N	GraNd	Forgetting	RL-Selector	SSP	Moderate	CLIP-Sel	UCB	Infobatch	Ours	Whole Dataset
Acc. (%)	71.1	72.3	71.0	72.5	73.4	70.0	73.1	73.2	75.8	76.5	76.8	76.4
Time (h)	10.5	10.5	10.5	10.5	10.5	10.5	10.5	10.5	10.5	10.5	10.5	17.5
Overhead (h)	$>$ 17.5	$>$ 17.5	$>$ 17.5	$>$ 17.5	$>$ 17.5	$>$ 24.0	$>$ 17.5	$>$ 1.6	0.03	0.0028	0.125	0.0
Overall (n*h)	$>$ 224.0	$>$ 224.0	$>$ 224.0	$>$ 224.0	$>$ 224.0	$>$ 276.0	$>$ 224.0	$>$ 96.8	84.0	84.0	85.0	140.0

Table 2: Results on ImageNet-1k with a 60% selection ratio using ResNet-50 on an 8-A100 server. We report wall-clock time (h) and total GPU time (GPU hours). Note that, due to high computational and memory costs ( Xia et al., 2023 ) , Glister and CG-Score are not reported. Some results are from ( Qin et al., 2023 ) .

Selection Ratio	70%	80%	90%	100%
Random*	37.5	37.9	38.5	39.6
Ours	38.5	39.0	39.6	39.6

Table 4: Object detection mAP (%) on MS-COCO using YOLOv8 ( Varghese and M., 2024 ) .

Selection Ratio	70%	80%	90%	100%
Random*	45.1	45.3	45.7	45.4
Ours	46.4	46.5	46.6	45.4

Table 5: Segmentation mIoU (%) on ADE20K using UperNet ( Xiao et al., 2018 ) .

Selection Ratio	60%	70%	80%	90%	100%
ImageNet-Hard	14.7	15.3	15.6	15.7	14.6
ImageNet-R	38.0	38.2	38.5	38.7	36.2
ImageNet-O	15.7	15.6	15.7	15.8	13.2

Table 6: Generalization performance of ResNet-50 trained with our method on ImageNet-Hard/R/O. We report AUPR (%) on ImageNet-O and accuracy (%) on others.

Method	Noisy Dataset		Corrupted Dataset
Method	20%	30%	20%	30%
Random	17.8	23.9	20.0	25.9
Random*	32.5	36.1	36.5	38.3
CLIP-Sel	26.1	33.1	26.1	32.1
Ours	38.1	41.9	41.9	46.7

Table 7: Robustness to noisy and corrupted data on Tiny-ImageNet using ResNet-50. The noisy ratio is 20%, and the selection ratios are 20% and 30%. Some results are from ( Yang et al., 2024 ) .

$R_{\text{diff}}$	$R_{\text{conf}}$	$r$	30%	50%	70%
			41.5	42.8	43.1
✓			42.1	45.0	48.2
	✓		42.2	45.8	48.4
✓	✓		42.5	46.1	48.6
✓	✓	✓	44.9	47.0	49.4

Table 8: Effect of the difficulty rewards

R_{\text{diff}}

, uncertainty

R_{\text{conf}}

, and adaptive reward weighting

r

on T-ImageNet using R-50.

Panel of Figure 4 — Figure 4: Effect of the RL agent on CIFAR-100 and Tiny-ImageNet under different selection ratios.

Data Selection

Data selection chooses informative, representative, or reliable examples to improve model training and evaluation while reducing data, compute, or labeling costs.

See all 80 papers

Rethinking Representativeness and Diversity in Dynamic Data Selection
Mar 5, 2026Yuzhe Zhou, Zhenglin Hua, Haiyun Guo +1Diversity-Aware Sampling Data Selection
Southeast University, Nanjing, China · Chinese Academy of Sciences, Beijing, China
Beyond What to Select: A Plug-and-play Oscillatory Data-Volume Scheduling for Efficient Model Training
May 14, 2026Suorong Yang, Hanqi Zhu, Hai Gan +4Deep Learning Optimization Data Selection
1National University of Singapore · 2Nanjing University · 3Hokkaido University +1
Adaptive Data Dropout: Towards Self-Regulated Learning in Deep Neural Networks
Apr 14, 2026Amar Gahir, Varshil Patel, Shreyank N GowdaTraining Data Selection Adaptive Sampling
School of Computer Science, University of Nottingham, Nottingham, United Kingdom

Mar 5, 2026J/K move · Enter open · S save

Rethinking Representativeness and Diversity in Dynamic Data Selection

Yuzhe Zhou, Zhenglin Hua, Haiyun Guo, Yuheng Jia

Southeast University, Nanjing, China · Chinese Academy of Sciences, Beijing, China

Diversity-Aware Sampling Data Selection Training Data Selection Data-Efficient Learning

Open paper AskPDF arXiv

Dynamic data selection accelerates training by sampling a changing subset of the dataset while preserving accuracy. We rethink two core notions underlying sample evaluation: representativeness and diversity. Instead of local geometric centrality, we define representativeness as coverage of dataset-level common or high-frequency feature factors. Instead of within-subset dispersion, we define diversity at the process level, requiring the selection trajectory to gradually include complementary rare factors over training. Based on this view, we propose a dynamic selection framework with three components. First, we score representativeness in a plug-in feature space to prioritize samples covering frequent factors. We instantiate this with a sparse autoencoder trained on the target dataset, using sparse unit activations to summarize both individual samples and dataset-wide factor statistics. Second, we realize process-level diversity by combining rare-factor sampling with a Usage-Frequency Penalty that promotes sample rotation, provably discourages monopoly, and reduces gradient bias. Third, we couple the two-dimensional scoring with a smooth scheduler that transitions selection from core-pattern consolidation to rare-factor exploration, without extra gradients, influence estimates, or second-order computations on the training model. Extensive experiments on five benchmarks across vision and text tasks demonstrate improved accuracy-efficiency trade-offs across models. Our method matches or exceeds full-data accuracy with over 2x training acceleration. Code will be released.

Abstract

Data Selection

Explore similar work

Rethinking Representativeness and Diversity in Dynamic Data Selection

Beyond What to Select: A Plug-and-play Oscillatory Data-Volume Scheduling for Efficient Model Training

Adaptive Data Dropout: Towards Self-Regulated Learning in Deep Neural Networks