AI Research Agents for Machine Learning: Search, Exploration, and Generalization in MLE-bench

2025-07-03

AI Summary

This study addresses the limited capability of AI research agents in solving realistic machine learning problems on the MLE-bench benchmark—particularly MLE-bench Lite—due to bottlenecks in model search, exploration, and generalization. Method: We propose a co-optimization framework integrating search strategies with a learnable operator set, modeling the agent as a programmable, differentiable search process. The atomic operator set comprehensively covers data preprocessing, architecture selection, and hyperparameter optimization. We unify greedy search, Monte Carlo Tree Search (MCTS), and evolutionary algorithms to enable end-to-end automated model design, training, and evaluation on Kaggle competition tasks. Contribution/Results: Our framework significantly improves the success rate of earning Kaggle medals from 39.6% to 47.7%, establishing a new state-of-the-art. It is the first work to empirically demonstrate that tight coupling between search strategy and the operator space is critical for AutoML performance.

Technology Category

• Search and Optimization: Learning to Search
• Machine Learning: Auto ML and Hyperparameter Tuning
• Multiagent Systems: Multiagent Learning

Application Category

• Search and Retrieval-Augmented AI: Agentic search
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• User Modeling, Personalization and Recommendation: ML for personalized search and recommendations

Abstract

AI research agents are demonstrating great potential to accelerate scientific progress by automating the design, implementation, and training of machine learning models. We focus on methods for improving agents' performance on MLE-bench, a challenging benchmark where agents compete in Kaggle competitions to solve real-world machine learning problems. We formalize AI research agents as search policies that navigate a space of candidate solutions, iteratively modifying them using operators. By designing and systematically varying different operator sets and search policies (Greedy, MCTS, Evolutionary), we show that their interplay is critical for achieving high performance. Our best pairing of search strategy and operator set achieves a state-of-the-art result on MLE-bench lite, increasing the success rate of achieving a Kaggle medal from 39.6% to 47.7%. Our investigation underscores the importance of jointly considering the search strategy, operator design, and evaluation methodology in advancing automated machine learning.

Problem

• Improving AI research agents' performance on MLE-bench
• Exploring search policies and operator sets for ML models
• Achieving state-of-the-art results in automated machine learning

Innovation

• AI agents automate ML model design and training
• Search policies optimize candidate solution navigation
• Operator-search interplay boosts benchmark performance