Research map / Multi-agent systems and games / LLM agents
AI Research Agents for Machine Learning: Search, Exploration, and Generalization in MLE-bench
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Formalizes AI research agents as graph search over artifacts, introduces AIRA-dojo and an improved operator set (AIRA operators) with greedy, MCTS and evolutionary search, reaching 47.7% medal rate on MLE-bench lite and documenting validation-to-test overfitting.
Problem
Existing AI research agents on MLE-bench entangle search policy, operators, and implementation, making it hard to know what drives performance; a generalization gap between validation and test scores also misleads search.
Approach
Cast research agents as graph-based search (fitness, selection, operators, operator policy, termination), show AIDE's operators are the bottleneck rather than search policy, design improved operators (complexity cue, scoped memory, think tokens, crossover) and pair them with greedy, MCTS and evolutionary search in a unified sandboxed dojo.
Key results
- AIRA-MCTS with R1 reaches 47% average medal rate on MLE-bench lite, up from 39.6% (AIDE baseline) and 47.7% best reported.
- Selecting final node by test score instead of validation would raise medal rate by 9 to 13 absolute points, indicating systematic overfitting.
- AIRA Greedy with o3 reaches 31.6% medal rate on full MLE-bench (vs 16.9% for AIDE o1-preview).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace AIDE operators with AIRA operators while keeping greedy search | AIDE Greedy · MLE-bench Lite | any medal rate: 39.8~\% vs. 45.5~\% | Helped |
| Run AIDE Greedy in AIRA-dojo instead of the reported setup | AIDE Greedy · MLE-bench Lite | medal rate: 35.2~\% to 45.9~\% | Helped |
| Enable the Memory operator (ablation) | AIDE · MLE-bench Lite | mean medal rate | No clear effect |
| Vary MCTS UCT exploration with AIDE operators (ablation) | AIDE MCTS · MLE-bench Lite | medal rate | No clear effect |
| Replace AIRA operators with AIDE operators in MCTS (ablation) | AIRA MCTS · MLE-bench Lite | medal rate | Hurt |
| Select the final solution using test score rather than validation score | AIRA MCTS · MLE-bench Lite | medal rate: 9.4~\% | Helped |
| Select the final solution using test score rather than validation score | AIRA Evolutionary · MLE-bench Lite | medal rate: 12.4~\% | Helped |
| Select the final solution using test score rather than validation score | AIRA Greedy · MLE-bench Lite | medal rate: 15~\% | Helped |
| Select the final solution using test score rather than validation score | AIDE Greedy · MLE-bench Lite | medal rate: 16.6~\% | Helped |
| Use test score for final-node selection while retaining validation-based search | AIRA agents · MLE-bench Lite | medal rate: 9 to 13 absolute points | Helped |
| Submit the top 3 validation-ranked nodes instead of one | AIRA agents · MLE-bench Lite | medal rate: additional 10~\% | Helped |
| Extend search beyond the 24-hour budget | AIRA Greedy · MLE-bench Lite | peak medal rate: absolute improvement of 6\% | Helped |
Surprising findings
- Search policy changes yield little gain when using AIDE's operators; operator set is the bottleneck.
- AIDE's test medal rate plateaus or slightly declines over 24 hours while validation keeps improving.
Limitations
- Results depend on LLM operators; no finetuning of operator models.
- Limited to 24-hour / 1 GPU budget; scaling behavior beyond this only partially explored.
- Possible data contamination of LLMs with Kaggle tasks.
Experimental setup
Datasets: MLE-bench lite MLE-bench
Benchmarks: MLE-bench lite MLE-bench (full)
Metrics: medal success rate gold medal rate silver-or-above rate validation-test gap
Baselines: AIDE greedy (o1-preview) AIDE greedy (DeepSeek R1) AIDE greedy (o3) R&D-Agent OpenAI o3 agent GPT-5-thinking agent
Builds on: AIDE MLE-bench Monte Carlo Tree Search Inspect
Closest papers by meaning
- MLE-bench: Evaluating Machine Learning Agents on Machine Learning Engineering (ICLR 2025 · similarity 0.93)
- FML-bench: A Controlled Study of AI Research Agent Strategies from the Perspective of Search Dynamics (NeurIPS 2026 · similarity 0.91)
- MLR-Bench: Evaluating AI Agents on Open-Ended Machine Learning Research (NeurIPS 2025 · similarity 0.91)
- MARS: Modular Agent with Reflective Search for Automated AI Research (ICML 2026 · similarity 0.90)
- MLRC-Bench: Can Language Agents Solve Machine Learning Research Challenges? (NeurIPS 2025 · similarity 0.89)
- CoMind: Towards Community-Driven Agents for Machine Learning Engineering (ICLR 2026 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.