PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Edan Toledo, Karen Hambardzumyan, Martin Josifoski, Rishi Hazra, Nicolas Baldwin, Alexis Audran-Reiss, Michael Kuchnik, Despoina Magka, Minqi Jiang, Alisia Maria Lupidi, Andrei Lupu, Roberta Raileanu, Kelvin Niu, Tatiana Shavrina, Jean-Christophe Gagnon-Audet, Michael Shvartsman, Shagun Sodhani, Alexander H. Miller, Abhishek Charnalia, Derek Dunfield, Carole-Jean Wu, Pontus Stenetorp, Nicola Cancedda, Jakob Nicolaus Foerster, Yoram Bachrach

arXiv:2507.02554v2cs.AIcs.LG

TL;DR

AI research agents must be designed so that performance contributions from search policies, operators, and evaluation can be distinguished on challenging machine-learning tasks. The paper formalizes agents as search algorithms, introduces AIRA-dojo and OAIRA, and evaluates Greedy, MCTS, and Evolutionary strategies. The best agent raises MLE-bench lite Kaggle medal success from 39.6% to 47.7%, while the study also identifies evaluation and compute-scaling limitations.

  • Problem

    Existing AI research agents entangle algorithm design, implementation, and compute, while AIDE provides limited insight into which components drive performance.

  • Method

    The paper formalizes agents as search algorithms, develops AIRA-dojo and OAIRA, and systematically pairs operator sets with Greedy, MCTS, and Evolutionary search policies.

  • Results

    47.7% Kaggle medal success rate is achieved on MLE-bench lite, up from 39.6%, and selecting final solutions by test rather than validation score would increase medal rate by 9 to 13% absolute.

  • Takeaways & Limitations

    The interplay between operator design and search policy is central to performance, while robust final-node selection is a promising avenue for improvement.

  • Takeaways & Limitations

    Evaluation is restricted to MLE-bench’s 24-hour and one-GPU constraints, which provides limited insight into performance scaling with greater resources.

Abstract

from arXiv · show

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.

1 Introduction

The paper frames AI research agents as search systems whose operators and search policies must be disentangled to identify performance bottlenecks. Using AIRA-dojo and improved operators, the authors achieve stronger MLE-Bench lite results while highlighting evaluation and generalization concerns.

  • AI research agents combine algorithm design, implementation, and compute usage, making controlled attribution of performance improvements difficult.
  • AIRA formalizes agents as search algorithms that navigate candidate solutions with search policies and iteratively modify them using operators.
  • AIDE’s operators, rather than its search algorithm, are identified as a bottleneck, motivating experiments with Evolutionary and Monte Carlo Tree Search agents.
  • 47.7% medal success rate is achieved by the best-performing agent, up from 39.6% on MLE-bench lite.
  • Selecting final solutions by test rather than validation score would increase medal rate by 9 to 13% absolute, indicating systematic overfitting during search.
  • AIRA-dojo provides a scalable, customizable environment for comparing operators, search policies, evaluation methods, and tasks under a comparable setup.
  • 55% performance is reached on MLE-Bench Lite when the best-performing agents are rerun with the latest AIRA-dojo version.

2 Research Agents as Search Algorithms

The paper models research agents as graph-based search algorithms that select existing artifacts, apply operators, and evaluate resulting candidates. This formulation separates search policies, operators, fitness functions, and termination rules for systematic comparison.

  • Research agents iteratively refine executable artifacts using external tools, execution feedback, and prior experimental insights.
  • A search agent operates on a directed graph whose nodes are artifacts and whose edges represent transformations between artifacts.
  • The framework specifies a fitness function, selection policy, operator set, operator policy, and termination rule.
  • Operators transform one or more selected artifacts into new artifacts and may be rule-based, LLM-based, tool-based, or composite.
  • AIDE uses greedy selection with Draft, Improve, Debug, and Memory operators, while alternative search policies retain the AIDE operator set.

3 Experiment Design

Experiments use AIRA-dojo to provide isolated, reproducible execution and compare search policies on MLE-bench lite. The design fixes computational resources, evaluates medal success, and uses repeated seeds to address benchmark variance.

  • AIRA-dojo supplies abstractions for operators and policies, tasks with evaluation criteria, and a framework for systematically studying agentic policies.
  • Four search policies—MCTS, Evolutionary, Greedy, and the authors’ AIDE implementation—are evaluated.
  • The infrastructure’s isolation and stable execution create a reproducible testbed for benchmarking agents across parallel runs and diverse tasks.
  • MLE-bench lite contains 22 tasks selected from the full 75-task Kaggle benchmark, enabling more seeds per task because existing methods show high variance.
  • Medal Success Rate is the percentage of attempts earning bronze, silver, or gold under task-specific percentile thresholds.
  • Each evaluation runs in an isolated environment with constrained compute, including one GPU, 24 logical CPU cores, 100 GB of RAM, and 1 TB of scratch storage.
  • Agents receive a 24-hour wall-clock window, with a maximum of 4 hours per code execution.

4 AIRA

AIRA introduces the OAIRA operator set and combines it with Greedy, MCTS, and Evolutionary search policies. The design adds adaptive complexity, scoped memory, and explicit reasoning support while keeping evaluation and termination criteria consistent.

  • 4.1 Operators OAIRA: OAIRA is a new operator set based on OAIDE, designed to improve context management, structured reasoning, and strategic diversity.
  • 4.1 Operators OAIRA: Prompt-adaptive complexity uses the node’s child count to guide artifact complexity for Draft and Improve operations.
  • 4.1 Operators OAIRA: Scoped memory retrieves sibling memories for Draft and Improve to promote diversity, while Debug accesses the full ancestral debug chain.
  • 4.1 Operators OAIRA: Think Tokens explicitly encourage structured reasoning and reflection, producing an average 2× increase in completion tokens for OAIRA operators.
  • 4.2 Agents: AIRA agents pair OAIRA with distinct search policies while sharing the same 5-fold cross-validation proxy fitness and wall-clock-or-artifact termination criterion.
  • 4.2 Agents: AIRAgreedy uses greedy search with OAIRA, so its improvement over AIDEgreedy directly reflects the operator-set difference.
  • 4.2 Agents: AIRAmcts uses MCTS with OAIRA, selecting nodes by UCT, expanding leaves, debugging buggy children, and back-propagating proxy fitness.
  • 4.2 Agents: AIRAevo maintains a fixed-size population and uses parent selection, reproduction with Improve or Crossover, and replacement of least-fit individuals.

5 Experiments and Results

Experiments show that operator design and search-policy interaction strongly affect agent performance, while validation–test mismatch limits apparent gains. Improved operators, advanced search, and final-node selection each contribute to stronger MLE-bench Lite results.

  • 5.1 Analyzing the Performance of the Current SoTA: Memory ablation produced nearly identical mean medal rates, indicating memory was not a driving factor behind AIDE’s performance.The controlled comparison evaluated AIDE with and without its memory operator.
  • 5.1 Analyzing the Performance of the Current SoTA: AIDE’s search-level exploration produced only marginal medal-rate differences across MCTS exploration constants when using AIDE’s operators.The same limitation appeared when AIDE’s operators were used with evolutionary search.
  • 5.1 Analyzing the Performance of the Current SoTA: AIDE’s validation performance continued improving while held-out test performance plateaued or slightly decreased over time, indicating overfitting.The anytime profile covered a 24-hour search window.
  • 5.2 AIRA Beyond Greedy: 45.5% versus 39.8%: AIRAgreedy outperformed AIDEgreedy with the same search policy, a 14% relative improvement attributable to the operator-set comparison.Using improved operators, all tested search policies outperformed AIDEgreedy; AIRAmcts reached a 47% average medal rate with R1.
  • 5.3 The Generalization Gap: Searching with a Proxy Evaluation: 9 to 13 absolute points: selecting the best final node using test scores improved every agent over validation-only selection.For AIRAmcts and AIRAevo, oracle final-node selection eliminated the gap between Val/Val and Test/Test; for greedy agents, it closed more than 60% of the gap.
  • 5.3 The Generalization Gap: Searching with a Proxy Evaluation: With as little as 3 top-k submissions, agents achieved an additional 10% of performance, and top validation nodes were informative of the best-performing nodes.The top-k strategy selected nodes by validation score and reported the highest test score among them.

6 Related Works

The paper situates its approach among LLM-based search, automated machine learning, and research-agent frameworks. Unlike predefined-space AutoML and NAS methods, it targets open-ended research tasks and addresses HPC scalability through Apptainer.

  • Scaling Search with LLMs: LLM-based generators are increasingly combined with best-of-N, beam search, MCTS, and evolutionary search to increase test-time compute across coding, mathematics, and planning.
  • Automating ML Engineering and Scientific Discovery: Traditional AutoML and NAS search predefined, expert-designed configuration spaces, whereas recent LLM advances enable more open-ended machine-learning engineering.
  • AI Research Frameworks: Existing machine-learning and research-engineering benchmarks generally provide their own frameworks, while this approach most closely resembles Inspect's minimally assumptive abstraction.
  • AI Research Frameworks: AIRA-dojo focuses on long-running research tasks and uses Apptainer because Docker's root-privilege requirements and limited HPC integration hinder scalability.

7 Conclusion

The paper frames AI research agents through the interaction of search policies and operators, finding that operator capability can bottleneck performance improvements. It reports a new MLE-bench state of the art while identifying scaling, operator design, fine-tuning, and data contamination as directions for future work.

  • AI research agents are analyzed along two axes: search policy and operators, enabling systematic study of their interplay.The framework separates how candidate solutions are selected from how existing solutions are modified.
  • The enhanced operator set combined with Greedy, MCTS, and Evolutionary Search produced a new state of the art on MLE-bench.
  • 39.6% to 47.7%: the success rate of winning a Kaggle medal increased with the best-performing agent.
  • Limitations and Future Work: The study leaves several performance dimensions for future work, including agentic operators and LLM fine-tuning.Possible extensions include full-fledged agents as operators and supervised fine-tuning or reinforcement learning.
  • Limitations and Future Work: Under MLE-bench’s 24-hour and 1-GPU constraints, evaluating how search policies and operators scale to greater resources remains future work.The paper notes that these restrictions provide limited insight into scaling behavior.
  • Limitations and Future Work: Possible contamination from information about evaluated or similar Kaggle tasks in LLM training data remains an unresolved concern.The paper identifies continuously generating fresh and novel tasks as a research challenge.

C The Effect of Compute: Searching Beyond 24h

Extending search beyond 24 hours reveals evolving agent rankings, continued gains for AIRA agents, and eventual overfitting that limits long-horizon performance. The results support jointly considering compute availability, search strategy, operator design, and the validation–test gap.

  • C The Effect of Compute: Searching Beyond 24h: After 15 hours, agent rankings diverge substantially, so the 24-hour evaluation window provides limited insight into longer-term capabilities.The study aims to develop agents that can use computational resources beyond the benchmark’s 24-hour limit.
  • C The Effect of Compute: Searching Beyond 24h: AIRA agents continue improving beyond 24 hours, whereas AIDEgreedy plateaus during the extended 90-hour search.Notable behavioral differences emerge after roughly 15 hours.
  • C The Effect of Compute: Searching Beyond 24h: Around 50 hours, AIRAmcts begins to overfit and its performance deteriorates, while AIRAgreedy continues improving before later declining from overfitting.AIRAgreedy reaches approximately 53% medal rate, a 6-point absolute and 12% relative improvement over its best 24-hour performance.
  • C The Effect of Compute: Searching Beyond 24h: The validation–test gap grows over longer searches, and AIRAgreedy and AIRAmcts eventually reach a gap similar to AIDEgreedy.Validation performance continues climbing while overfitting becomes more severe.
  • C The Effect of Compute: Searching Beyond 24h: Long-horizon evaluation highlights the need to jointly consider search strategy, operator design, computational resources, and the generalization gap.The authors identify the generalization gap as having a fundamental impact on agent performance.

D Evaluation on the Full MLE-bench Set

The full-benchmark evaluation estimates performance across the complete MLE-bench task set and complexity-based subsets using 20 random seeds. AIRAgreedy achieves a 31.6% overall medal rate and outperforms the baseline across all reported splits.

  • D Evaluation on the Full MLE-bench Set: 31.6% overall medal rate is achieved by AIRAgreedy on the full MLE-bench evaluation.Average performance and confidence intervals are computed over 20 random seeds.
  • D Evaluation on the Full MLE-bench Set: AIRAgreedy outperforms the baseline on every reported complexity-based and MLE-bench-30 split.The reported splits include the overall benchmark, complexity-based subsets, and the MLE-bench-30 subset.
  • D Evaluation on the Full MLE-bench Set: Table 2 reports mean ± standard error medal rates across agents, including reproduced prior-work results and system-card values for OpenAI agents.AIDE values come from the MLE-Bench repository, with a separate procedure for MLE-Bench-30.
  • D Evaluation on the Full MLE-bench Set: Table 3 provides per-task medal rates grouped by task complexity.This complements the aggregate split-level results in Table 2.

E Implementation Details

The implementation uses a single-agent, resource-bounded environment with configurable search procedures and LLM-driven artifact-generation prompts. MCTS maintains visit counts and mean proxy fitness through incremental back-propagation without simulated roll-outs.

  • E Implementation Details: Experiments use one NVIDIA H200 GPU, 24 CPU cores, 100 GB RAM, and 1 TB local scratch space.Jobs have a 24-hour wall-clock cap, 4-hour execution limit, and 5-minute grace period.
  • E Implementation Details: MCTS uses five children and uct_c=0.25, while evolutionary search samples five candidates per generation.A debug cycle is capped at 10 nodes or 12 hours, whichever comes first.
  • E Implementation Details: Each MCTS node stores a visit count N(v) and empirical mean fitness Q(v), initialized from the newly evaluated leaf’s proxy fitness.The leaf value is back-propagated along the selected path.
  • E Implementation Details: Incremental updates maintain Q(u) as the mean of all propagated leaf fitnesses, with repeated visits driving convergence toward expected value.No simulated roll-outs are performed; leaf values are taken directly from F.
  • E Implementation Details: The agent prompts require one idea and one end-to-end Python implementation, including five-fold cross-validation and saved test predictions.Prompts also constrain runtime, discourage exploratory data analysis, and encourage efficient processing for large datasets.
  • E Implementation Details: Prompt variants request simple, normal, or advanced ideas with increasing complexity and compute, while keeping evaluation consistent and considering previously explored solutions.The environment supplies task, data, compute, package, and prior-solution context to the agent.

F Completion Tokens Per Method

The paper examines both operator-token usage and the reliability of agent comparisons across seeds, tasks, and validation settings. It reports longer operator reasoning chains, substantial ranking instability with few seeds, and smaller test–validation gaps for AIRA operators than for AIDEgreedy.

  • F Completion Tokens Per Method: AIRA operator changes produce substantially longer thinking chains, measured by average completion tokens per operator.Figure 10 averages token counts across independent MLE-bench lite runs.
  • F Completion Tokens Per Method: At least three seeds may be insufficient for reliable rankings; the authors recommend 10 seeds per competition and preferably 20.They also recommend stratified bootstrapping and higher-seed evaluation on representative subsets such as MLE-bench Lite.
  • F Completion Tokens Per Method: Few evaluation seeds can dramatically change agent rankings, making observed performance differences potential statistical artifacts.The analysis estimates underlying performance distributions from 20 seeds and contrasts them with rankings from smaller samples.
  • F Completion Tokens Per Method: Between-task variability accounts for a substantial share of performance variance, while some methods’ medal outcomes are similarly variable within tasks.Figure 12 also shows distributions of consistently successful, consistently failing, and inconsistent outcomes across tasks.
  • F Completion Tokens Per Method: AIRA-operator agents exhibit a smaller test–validation gap than AIDEgreedy in their performance profiles.The comparison covers AIRAmcts, AIRAevo, and AIRAgreedy against the AIDEgreedy profile.

K Per-Task Results

This section reports non-aggregated, per-task results from the MLE-bench Lite suite, summarized in Fig. 14.

  • The section presents non-aggregated results for individual tasks in the MLE-bench Lite suite.
  • The reported results are obtained from the MLE-bench Lite suite described in Section 5.2.
  • Figure 14 summarizes the per-task results.

L Sample Search Trees

The paper illustrates search-tree behavior across methods using per-task samples, score encodings, and markers for medal-winning, above-median, and buggy solutions. It also compares perceived validation performance with held-out test medal rates to expose generalization gaps.

  • Sample Search Trees: Search-tree samples come from several methods applied to the spooky-author-identification task.
  • Sample Search Trees: Node colors encode validation scores, while node labels report test scores.
  • Sample Search Trees: Medal emojis identify medal-winning nodes, ok emojis mark solutions above the human-leaderboard median, and red nodes indicate bugs.
  • Generalization: Figure 13 compares mean agent-reported validation and held-out test medal rates over 24 hours across search policies using AIRA operators.
  • Generalization: The widening band in Figure 13 represents the generalization gap between apparent validation gains and held-out test performance.
  • Per-Task Performance: Figure 14 reports per-task performance, averaging R1 methods over 20 seeds and o3 methods over 10 seeds per task.
Loading 2507.02554v2…