Title: AI Research Agents for Machine Learning: Search, Exploration, and Generalization in MLE-bench
Venue: NeurIPS 2025
Accepted: yes
Keywords: LLM Agents; Code generation; AutoML; LLM-driven search; ML Engineering Agents; ML Research Agents; AI Research Agents
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a framework to construct LLM-based code agents for solving MLE-bench tasks. The primary contributions of the paper are algorithmic and empirical. Key contributions include the computational environment for agents (AIRA-Dojo), the agent design framework (AIRA) and the experimental results showing improved performance on MLE-bench (lite). AIRA builds on ideas in AIDE. The main technical ideas involve inducing a search space over programs using LLM-based operators like DRAFT, DEBUG, IMPROVE, MEMORY, etc. The sandbox infrastructure creates the environment to safely and reliably conduct the (long-running) search. Experiments demonstrate a new state of the art on the benchmark. The paper also includes a careful empirical study of the search algorithm's design choices.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n- The paper studies an important and interesting problem. Progress here would have a large impact and be of interest to the community.\n- The paper is reasonably well written. Key concepts are illustrated nicely.\n- The proposed algorithmic approach is intuitively clear and technically sound. I wasn't able to find any major issues with the technical setup. The induced search space is clearly described and the operators well described in the appendix (with the exception of CROSSOVER).\n- The experiments are well-constructed and nicely detailed. The primary result (Figure 1, Figure 3) shows that AIRA performs better than the previous SOTA (AIDE) on the MLE-bench lite benchmark. The careful investigation into the underlying components of performance in Section 5.2 and 5.3 contains valuable insights into the underlying performance factors. The appendix contains a good amount of additional detail.\n\n**Weaknesses**\n\n- The paper discusses its limitations reasonably well in Sec 7.1 I agree that leakage of the benchmark during the (unseen) base LLM training might be a concern. While AIDE is a good choice of baseline, I was unable to find any evaluations of SOTA reasoning and code LLMs (Gemini 2.5, Opus 4, etc.). (This would be zero-shot, without any search beyond the TTS used by the LLM itself during inference). While I don't expect to see outperformance, it would strengthen the paper's empirical standing and make it easier to assess its relative impact in the broader (and rapidly growing) area of using LLMs to construct search-based agents.\n\n- The algorithmic novelty is a little low. The paper mostly builds on the ideas in AIDE and uses well-known planning algorithms. However, this is a minor issue given the careful and detailed empirical investigation.\n\n- The paper does not include an error analysis of the generated programs or any code quality metrics. It would be useful to better understand the error distribution during program generation (beyond those having to do with inference infrastructure).\n\nQUESTIONS:\n1. What is the performance of SOTA reasoning and code LLMs (that only do test-time search) on this task? How does it compare to AIDE and AIRA?\n\n2. What kinds of mistakes or errors were observed in the programs generated on test datasets? What is the subjective and / or objective code quality?\n\n3. What is the implementation of CROSSOVER from Line 106? I was unable to find it in the paper or appendix.\n\n4. On Line 184, what figures are being referenced? (\"In the figures, we take ...\"). (I think it refers to the prior AIDE eval metrics reported in the graphs but it was confusing on first reading.)\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-01"}

comments: SUMMARY:
The paper presents an analysis of AI research agents that solve the MLE benchmark, a set of Kaggle competitions. The authors discuss the important details of designing such agents, dividing the process into two major parts: high-level actions that are available to agents, and a search policy. The paper also investigates how the choice of proxy metric affects the agent's performance, and whether the gap between validation and test accuracies can be closed. Additionally, the authors propose to extend the action space of the current approach, which moderately boosts metrics.

STRENGTHS AND WEAKNESSES:
**Quality**: This paper presents a comprehensive empirical analysis on AI research agents, with adequate experiment design that supports the claims. To strengthen the paper, an ablation for each additional operator might be included. The reported results seem reliable, with confidence intervals provided.

**Clarity**: The paper is easy to follow and clearly written, the difference between the approaches is highlighted, and each experiment is clearly presented.

**Significance and Originality**: The paper outlines major directions for improvement for future work, and therefore is beneficial for the scientific community and agentic research in particular.

QUESTIONS:
[1] In AIRAEVO explanation, it is said that there exists a Crossover operation. How is it implemented?

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 4
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper frames AI research tasks as a search problem and breaks down the design of AI research agents, such as AIDE, into two core components: search policies and operators. The authors build on this formulation to conduct a thorough ablation study, systematically evaluating the impact of different design choices for these components. The best configuration identified in the study achieves state-of-the-art results on the challenging MLE-Bench, improving the success rate for earning a Kaggle medal from 39.8% to 47% on MLE-Bench Lite. This result underscores the practical significance of their findings and highlights the value of carefully designing and evaluating agent components to advance automated machine learning research.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n\n1. The paper achieves a notable advancement by establishing a new state-of-the-art performance on MLE-Bench Lite, boosting the Kaggle medal success rate from 39.8% to 47.8%. This improvement highlights the practical impact of the proposed approach and sets a strong benchmark for future work in automated machine learning.\n\n2. The authors provide an insightful perspective by decomposing the design of AI research agents into two key components: search policies and operators. Their comprehensive ablation studies offer valuable insights into how these components contribute to agent performance, paving the way for more systematic design of AI research agents.\n\n3. The release of the codebase will be a significant strength, enabling other researchers to build upon this work and explore new designs for AI research agents. This commitment to open science enhances the paper’s impact and fosters further innovation in the field.\n\n### Weaknesses\n\n1. A primary concern is the absence of statistical significance tests to support the authors’ claims. For instance, in Line 263, the authors conclude that the search policy has minimal impact based on the small performance margin shown in Figure 4a. However, in Line 289, they argue that a slightly larger margin in Figure 3a suggests a meaningful interplay between search policy and operator design. Both conclusions are questionable, as the margins may fall within the error bars (as seen in Figure 3a) and lack statistical validation. Furthermore, the exploration of operator designs is limited, with only a brief comparison between OAI and DOJO, which is confounded by the use of different models, undermining the conclusiveness of the analysis. Including rigorous statistical tests and a broader exploration of operator designs would strengthen the paper’s findings.\n\n2. The analysis in Section 5.3 raises concerns about the interpretation of the validation-test gap. The authors suggest that the model accessing such data represents an upper bound on performance. But both the validation and test sets are equally held out from the training data, with the only difference being whether they are included during the AI research agent search. This framing misrepresents the gap’s significance, as the test set is not inherently \"ground-truth held out data\". Additionally, the method described in Line 321 relies on test set data to report the best-performing scores, which is not feasible in real-world scenarios where test data is inaccessible during model development. I suggest removing this section from the paper.\n\nQUESTIONS:\n- Figure 1a’s x-axis label “any medal” is confusing and unclear.\n- Line 93 states that “fitness function F is defined for each node by the operator that generates or modifies it,” but Line 211 claims “each agent employs the same proxy-fitness function F.” These statements appear contradictory. Which is correct?\n- Line 95 notes that “all operators are LLM-driven except for the MEMORY operator, which is defined by hand.” Does “LLM-driven” mean the operators use LLMs during execution or were generated by LLMs? Additionally, does \"the MEMORY operator is defined by hand\" imply that all other operators are LLM-generated (I guess not)?\n- Line 105’s description of the MEMORY operator is unclear. If it does not involve an LLM, how is it implemented? If it does, this contradicts Line 95’s claim.\n- Line 184 indicates that [DOJO] results use the R1 model, while [OAI] results use the O1 model. The use of different models introduces a confounder, making it difficult to draw conclusive comparisons.\n- Line 194 suggests that a larger number of children (nc) requires more advanced code complexity. Why is this the case? Additionally, how can a DRAFT node have children, given its role in initiating solutions?\n- Line 259 omits the exploration constant value c=0.5.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper provides a new perspective where the design of agents is 1) Search policy, 2) Operator that modifies the current solution. \nA formalization is provided where AIDE is re-interpreted under this new lens. \nThey also provided AI Research Agent (AIRA) Dojo, a software/framework that allows people to implement new operators and search policies.\nThe paper also found a generalization gap.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. The paper is indeed well-written and easy to follow.\n2. The re-framing is a good attempt to provide insight into this field.\n\nWeaknesses:\n\nI believe this is a borderline paper. I can either vote to accept or reject, but will not take a strong stance to vote either way. The paper is well-written, and the experiments are solid. However, the paper puts a lot of emphasis on the **reframing** of AIDE and agents. Then, whether I personally is convinced by the **reframing** is the key to whether I should vote to accept or reject this paper. \n\nI do not find this reframing convincing. My biggest confusion is -- for an agent that has access to tools (such as terminal/internet), what is the search policy and operator? The paper seems to want to merge two separate conceptual models: agent and inference-time search. I think they are different. The authors seem to think they are the same. I'm open to having a longer discussion with the authors during the rebuttal.\n\nThe paper can also use better citations, such as [1] and [2]. The overfitting phenomenon was also discovered by Anthropic last year [3], but proper citation was not provided.\n\n[1] https://deepmind.google/discover/blog/alphaevolve-a-gemini-powered-coding-agent-for-designing-advanced-algorithms/ and the paper: Novikov, A., Vũ, N., Eisenberger, M., Dupont, E., Huang, P. S., Wagner, A. Z., ... & Balog, M. (2025). AlphaEvolve: A coding agent for scientific and algorithmic discovery. arXiv preprint arXiv:2506.13131.\n\n[2] Pan, Jiayi, Xiuyu Li, Long Lian, Charlie Snell, Yifei Zhou, Adam Yala, Trevor Darrell, Kurt Keutzer, and Alane Suhr. \"Learning adaptive parallel reasoning with language models.\" arXiv preprint arXiv:2504.15466 (2025).\n\n[3] Pan, Jane, He He, Samuel R. Bowman, and Shi Feng. \"Spontaneous Reward Hacking in Iterative Self-Refinement.\" arXiv preprint arXiv:2407.04549 (2024).\n\nQUESTIONS:\nN/A\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 1, "ORIGINALITY": 2, "DATE": "2025-07-07"}