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

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.

• The paper presents a formalization of AI research agents as search algorithms that iteratively refine ML pipelines using diverse operators.
• It introduces AIRA-dojo, a containerized, reproducible platform that enables scalable and parallel ML experimentation on Kaggle tasks.
• Empirical results demonstrate that enhanced operator design, rather than search policy alone, significantly boosts medal rates and mitigates overfitting.

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

This paper presents a systematic paper of AI research agents designed to automate the end-to-end process of ML model development, with a focus on the MLE-bench benchmark—a suite of real-world Kaggle competitions. The authors formalize research agents as search algorithms operating over a space of candidate solutions, where each solution is iteratively refined using a set of operators. The work introduces AIRA-dojo, a scalable and customizable environment for evaluating such agents, and provides a detailed empirical analysis of the interplay between search strategies, operator design, and evaluation methodology.

Formalization of Research Agents as Search Algorithms

The core contribution is the formalization of AI research agents as graph-based search algorithms. Each agent maintains a search graph where nodes represent candidate artifacts (e.g., codebases or ML pipelines), and edges correspond to transformations applied via operators. The search process is governed by:

• Fitness Function: Proxy metric (e.g., cross-validation score) used to evaluate candidate solutions.
• Selection Policy: Heuristic for choosing which nodes to expand (e.g., greedy, MCTS, evolutionary).
• Operator Set: LLM-driven or hand-crafted functions for generating, improving, debugging, or recombining solutions.
• Operator Policy: Rules for selecting which operator to apply.
• Termination Rule: Stopping criteria based on time or resource constraints.

This abstraction enables controlled experimentation with different agent designs, isolating the effects of search policy and operator set.

AIRA-dojo: Infrastructure for Agentic ML Research

AIRA-dojo is introduced as a robust, containerized environment supporting reproducible and scalable agent evaluation. Key features include:

• Jupyter-based execution: Agents interact with the environment via Jupyter notebooks, enabling arbitrary code execution and shell access.
• Resource isolation: Apptainer containers enforce strict limits on compute, memory, and storage, ensuring reproducibility and preventing cross-agent interference.
• Superimage: A standardized container image with pre-installed ML libraries and CUDA support, facilitating rapid agent deployment.

This infrastructure supports high-throughput, parallel experimentation and is designed to be compatible with HPC clusters, addressing limitations of Docker-based solutions.

Operator Design and Search Policy Interplay

A central empirical finding is that the operator set, rather than the search policy, is often the primary bottleneck in agent performance. The authors compare several search strategies—greedy, Monte Carlo Tree Search (MCTS), and evolutionary algorithms—using both the baseline AIDE operator set and a newly proposed AIRA operator set. The AIRA operators introduce:

• Prompt-adaptive complexity: Dynamic adjustment of solution complexity based on search graph context.
• Scoped memory: Contextual retrieval of relevant prior solutions to promote diversity and avoid mode collapse.
• Think tokens: Explicit encouragement of structured reasoning and reflection in LLM completions.

When advanced search policies are paired with the baseline AIDE operators, no significant performance gains are observed. However, the introduction of the AIRA operator set yields substantial improvements, especially when combined with non-greedy search strategies.

Empirical Results on MLE-bench Lite

Experiments are conducted on MLE-bench Lite, a curated subset of 22 Kaggle tasks. Agents are evaluated by their success rate in achieving Kaggle medals (bronze, silver, gold) within a 24-hour compute window (1 H200 GPU, 24 CPUs, 100GB RAM). Key results include:

• State-of-the-art performance: The best agent (AIRA operators + MCTS) achieves a 47.7% medal rate, up from the previous 39.6% SOTA.
• Operator impact: Switching from AIDE to AIRA operators with a greedy policy increases the medal rate from 39.8% to 45.5%.
• Search policy impact: With improved operators, advanced search policies (MCTS, evolutionary) further boost performance, particularly in gold medal attainment.
• Generalization gap: Systematic overfitting is observed; selecting final solutions by test score (oracle) rather than validation score would increase medal rates by 9–13 percentage points, depending on the search strategy.

The analysis also demonstrates that agent rankings can change over longer compute horizons, and that robust final-node selection strategies (e.g., top-k selection) can partially mitigate the generalization gap.

Implementation Considerations

• LLM Integration: All operators are LLM-driven (DeepSeek R1, OpenAI o3, GPT-4o), with self-hosted inference for throughput and rate-limit avoidance.
• Resource Management: Strict isolation and reproducibility are enforced via Apptainer containers and standardized compute quotas.
• Evaluation Protocols: Multiple seeds per task are used to ensure statistical reliability, and stratified bootstrapping is recommended for confidence intervals.
• Scalability: The infrastructure supports long-running experiments (up to 5 days), with checkpointing to handle hardware failures.

Implications and Future Directions

The findings underscore the necessity of jointly optimizing operator design, search policy, and evaluation methodology for effective agentic ML automation. The operator set is a critical lever for performance, and advanced search strategies only yield benefits when paired with sufficiently expressive and diverse operators. The persistent generalization gap highlights the need for improved validation protocols and regularization techniques.

Future research directions include:

• Agentic operators: Incorporating more sophisticated agents (e.g., ideation agents, SWE-Agents) as operators.
• LLM finetuning: Exploring supervised or reinforcement learning to enhance operator effectiveness.
• Scaling studies: Investigating agent performance under relaxed compute constraints and over longer time horizons.
• Benchmark development: Addressing data contamination by curating novel, unseen tasks for evaluation.

Conclusion

This work provides a rigorous framework and empirical foundation for the development and evaluation of AI research agents in ML engineering. By disentangling the contributions of search policy and operator design, and by providing robust infrastructure for experimentation, the paper advances the state of automated ML research and sets a clear agenda for future work in agentic scientific discovery.