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

TL;DR

Data Agent introduces an end-to-end dynamic data selection framework that learns a sample-wise policy to accelerate training and improve performance across diverse datasets and models.

Contribution

It formulates data selection as a training-aware sequential decision process, enabling adaptive, task-agnostic, and scalable sample prioritization.

Findings

Reduces training costs by over 50% on ImageNet-1k and MMLU.

Consistently accelerates training while maintaining or improving accuracy.

Demonstrates robustness to noisy datasets and broad applicability.

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.