Agentic system for better dataset search improves machine learning training
From Intent to Improvement: Agentic Dataset Search for Better Model Training
Artificial IntelligenceInformation Retrieval
Summary
Finding the right data to train machine learning models is hard because many data collections online lack clear descriptions or important details. The authors designed AgenticDS, a new method that carefully checks the meaning and requirements of your search before picking datasets. It uses clear rules and verifying steps to make sure the chosen data really matches your needs and can improve model training. Tests show AgenticDS finds better datasets than other methods and helps models perform better in math, programming, and medical tasks.
What this means in practice
- •For machine learning engineers: Find higher-quality training datasets that meet specific constraints to improve model accuracy in diverse domains like mathematics and medicine.
- •For data platform teams: Enhance dataset search tools with a verifiable, agentic approach that enforces query constraints and provides explanations for dataset relevance.
Authors
Qiaosheng Chen, Kaijia Huang, Fanqi Zhou, Gong Cheng
Abstract
Finding training data in million-scale machine learning repositories is a Web search problem where text relevance is not enough. A query couples a semantic training goal with constraints such as task, license, and training stage, and the metadata that would verify them are often missing. Text retrieval ranks dataset cards by similarity but cannot enforce hard constraints or follow the relations that reach a sparsely described dataset. Free-form tool-calling agents are more flexible, but the same model that picks a dataset also declares it a match, and nothing checks that the query's constraints are actually met. Neither approach establishes that the data it returns will improve model training. To address it, we propose AgenticDS, a contract-governed agentic framework in which every language model decision instantiates a typed object before deterministic code consumes it. The framework compiles a query into a typed search plan, executes it through admissible deterministic operators over hybrid retrieval and bounded graph expansion, and verifies the provenance-grounded evidence before accepting a ranking or emitting a typed diagnosis that permits only local repairs. On HuDS-Bench, a new benchmark of 200 queries with 24,908 graded judgments pooled over 15 systems, AgenticDS improves NDCG@10 by 25.7\% over the strongest dense--sparse hybrid and by 18.9\% over a commercial LLM agent. A controlled search-to-train study further shows that the datasets it retrieves give the best downstream accuracy in mathematics, code, and medicine. Code and data are publicly available at https://github.com/nju-websoft/AgenticDS.