What Does It Take to Be a Good AI Research Agent? Studying the Role of Ideation Diversity — part 1/3

Cleaned-text code points [0, 15806). Bibliography omitted here; source.md includes it.

# What Does It Take to Be a Good AI Research Agent? Studying the Role of Ideation Diversity

###### Abstract

AI research agents offer the promise to accelerate scientific progress by automating the design, implementation, and training of machine learning models. However, the field is still in its infancy, and the key factors driving the success or failure of agent trajectories are not fully understood. We examine the role that ideation diversity plays in agent performance. First, we analyse agent trajectories on MLE-bench, a well-known benchmark to evaluate AI research agents, across different models and agent scaffolds. Our analysis reveals that different models and agent scaffolds yield varying degrees of ideation diversity, and that higher-performing agents tend to have increased ideation diversity. Further, we run a controlled experiment where we modify the degree of ideation diversity, demonstrating that higher ideation diversity results in stronger performance. Finally, we strengthen our results by examining additional evaluation metrics beyond the standard medal-based scoring of MLE-bench, showing that our findings still hold across other agent performance metrics.

## 1 Introduction

The rapid advancement of Large Language Model-based ( Brown et al., 2020 ) agents equipped with tools ( Schick et al., 2023 ) has sparked interest in the quest to develop research agents, in areas as challenging as chemistry ( Boiko et al., 2023 ) or biology ( Swanson et al., 2025 ) .

These research agents constitute an emerging paradigm in computational scientific discovery characterized by end-to-end autonomous systems capable of conducting independent research.

In particular, recent work on autonomous AI research agents ( Shen et al., 2023 ; Huang et al., 2024 ; Toledo et al., 2025 ; Zhao et al., 2025 ) improves upon automated machine learning engineering tools ( Feurer et al., 2022 ) by mirroring the cognitive processes of human researchers through a structured research pipeline: idea generation and hypothesis setup, experimental design and implementation, empirical validation, and iterative refinement. Recent advances have achieved notable milestones such as creation of the first fully autonomous AI-generated research paper accepted through peer review ( Yamada et al., 2025 ) .

Despite the potential of these recent breakthroughs in automating AI science, the field is still in its infancy and little is understood about the factors driving their successes and failures. Error analysis is substantially more complicated than in classic machine learning setups, due to the presence of long multi-step trajectories often guided by heuristic-based search algorithms ( Toledo et al., 2025 ) and leveraging tool use, which requires complex evaluation frameworks. Moreover, obtaining large-enough samples to perform meaningful analysis and ablate design choices can be computationally prohibitive.

This paper starts from the postulate that ideation diversity is a key bottleneck in AI research agents’ performance. To study this hypothesis, we face two key challenges: analyzing complex agentic trajectories at scale, and measuring and controlling ideation diversity.

We perform a first-of-its-kind, large-scale study of AI research agents’ trajectories in MLE-bench ( Chan et al., 2025 ) , a well-known benchmark of Kaggle machine learning tasks. We study 6 different LLM backbones equipped with 2 different agentic frameworks (or scaffolds ) on the 75 machine learning tasks available in MLE-bench across 10 to 20 random seeds, yielding a total of 11,000 trajectories. This corresponds to roughly 1,200,000 individual nodes in the agent scaffold search, for a total of 264,000 GPU hours.

To measure ideation diversity, we propose calculating Shannon entropy ( Shannon, 1948 ) on the distribution of model architectures that the agent plans to implement in the ideation phase. Figure 1 shows the correlation between ideation diversity and the performance in MLE-bench using our generated trajectory bank, hinting at a relation between the two.

To confirm the diversity hypothesis, we perform a controlled experiment where we remove mechanisms yielding highly diverse solutions by amending the prompt. 1 1 1 Our main results in Section 4.2 are based on decreasing ideation diversity through the prompt shown to the agent. In the Appendix A , we provide additional results where we control diversity via the sampling temperature parameter. We run the controlled experiment on a subset of MLE-bench, studying two diversity setups (control and ablated) using two agentic frameworks on 22 machine learning tasks across 10 seeds. We show that when ablating ideation diversity, the agents’ performance decreases. Finally, we strengthen this finding by examining additional evaluation metrics on MLE-bench aside from the standard score based on the Kaggle medal system.

### 1.1 Contributions

In summary, the contributions of this paper are as follows: • We propose methods to quantify and control the agent’s ideation diversity .

• We perform a first-of-its-kind, large-scale analysis on agentic trajectories . We study a total of 11,000 AI research agents’ trajectories across multiple agentic frameworks, LLM backbones, and machine learning tasks.

• We show that the choice of agentic scaffold significantly influences ideation diversity . Our results further reveal a significant correlation between ideation diversity and agent performance on MLE-bench tasks.

• Through a controlled experimental design , we establish a causal relationship, showing that increasing ideation diversity leads to improved performance on MLE-bench tasks.

• We confirm that these findings are robust when evaluated with alternative performance metrics .

## 2 Research Agents and MLE-bench

##### Agents

Broadly, Wooldridge and Jennings (1995) define an agent as a computer system that is situated in some environment and that is capable of autonomous action in this environment in order to meet its design objectives. In the particular context of recent work on research agents based on generative AI and LLMs ( Schick et al., 2023 ; Boiko et al., 2023 ; Shen et al., 2023 ; Huang et al., 2024 ; Swanson et al., 2025 ; Toledo et al., 2025 ; Zhao et al., 2025 ) , we specifically refer to agent systems implemented using two main components: 1. a model backbone (typically, an LLM), which processes observations from the environment as prompts and emits text-based actions, and 2. an outer loop making use of the model backbone to interface with the environment . This outer loop orchestrating the LLM actions is usually referred to as agentic frameworks or agentic scaffolds in the literature ( Wu et al., 2022 ) . The environments for AI research agents allow the agent to create and run code among other tools ( Schick et al., 2023 ) .

##### MLE-Bench

( Chan et al., 2025 ) (Machine Learning Engineering Benchmark) is a well-known evaluation framework designed to assess autonomous AI agents’ capabilities in solving real-world machine learning problems. The benchmark is constructed from 75 tasks sourced from Kaggle competition datasets, providing a diverse collection of machine learning challenges that span computer vision, natural language processing, time series forecasting, tabular data analysis, and multimodal learning domains.

Each MLE-bench task comprises standardized components including problem documentation, training datasets, held-out test sets, sample submission formats, and automated evaluation protocols. The evaluation methodology follows Kaggle’s competitive framework, utilizing problem-specific metrics such as accuracy, F1-score, or RMSE, with performance assessed through leaderboard ranking systems.

AI research agents interact with MLE-bench through a standardized API that mirrors real-world ML development workflows. Agents must autonomously perform the complete machine learning pipeline: data exploration, feature engineering, model selection, hyperparameter optimization, and submission generation within computational constraints. The benchmark employs stratified sampling and cross-validation methodologies to ensure robust performance assessment, with final rankings determined through holdout test set evaluation to prevent overfitting to validation metrics.

Given the same dataset, there are multiple ways of defining how to interface with the benchmark, mediated by the agentic scaffold of choice. Recent work ( Toledo et al., 2025 ) proposed a tree-based agentic scaffold to tackle MLE-bench. In Section 3.1.3 , we describe the scaffolds studied in this work with more detail.

Figure 2 shows an example of a tree-search-based AI research agent attempting the Jigsaw Toxic Comment Classification Challenge MLE-bench task. The AI research agent would execute the following computational workflow: idea generation (e.g., leveraging pre-trained convolutional neural network ( LeCun and Bengio, 1998 ) features with linear classifiers), hypothesis setup (e.g., establishing baselines using CIFAR-100 ( Krizhevsky et al., 2009 ) embeddings with logistic regression), and implementation and experimentation (e.g., tuning learning rate for logistics regression). The next step would be to analyze results (e.g. look at confusion matrix) and finally make a submission to the grader. Based on the leaderboard rank and experiment analysis, the agent would then propose the next set of ideas and hypotheses (e.g. changing classifier from logistic regression to random forest or tuning a CNN end-to-end).

These idea generation, implementation, experimentation and submission steps are iterated upon to improve the agent’s leaderboard rank. This paper focuses on the idea generation step.

## 3 Methods

In this section, we briefly introduce the methodology (including data, metrics, agentic orchestrations or scaffolds, and LLM backbones) used in both our data analysis, in Section 4.1 , and the controlled experiment, in Section 4.2 .

### 3.1 General Setup

#### 3.1.1 Dataset

For the trajectory analysis, we use agent trajectories on MLE-bench. For the controlled experiment, we focus on MLE-bench lite, a curated subset of 22 tasks selected from the full benchmark.

#### 3.1.2 Metrics

In line with the benchmark guidelines, for both our data analysis and controlled experiment, we assess each agent’s performance using the Medal Success Rate (henceforth referred to as medal rate). Specifically, for each task, agents earn a bronze, silverho, or gold medal according to task-specific percentile thresholds. We report the percentage of attempts in which an agent secures a medal. Later, in Section 4.3 , we incorporate additional metrics.

#### 3.1.3 Agentic Scaffolds

Following recent work ( Toledo et al., 2025 ) , we formalize AI research agents as search algorithms composed of a search policy, used to navigate the space of candidate solutions to a task, and a set of operators, which modify existing solutions to generate new candidate solutions.

Enabled by this formalization, we study a range of agentic structures, specifically focusing on: (1) AIDE ( Jiang et al., 2025a ) , an LLM-driven agent that approaches problem-solving as a tree-search over the domain of Python solutions, utilizing a Greedy policy. (2) AIRA GREEDY \text{AIRA}_{\text{GREEDY}} ( Toledo et al., 2025 ) , another greedy tree-based search policy, with a different design for operators, memory scope, and prompts, and (3) AIRA MCTS \text{AIRA}_{\text{MCTS}} ( Toledo et al., 2025 ) , utilizing Monte Carlo Tree Search (MCTS ( Coulom, 2006 ; Kocsis and Szepesvári, 2006 ; Browne et al., 2012 ) ) for its search policy, in contrast to its greedy counterparts.

In all three agentic scaffolds, the process results in trees where each node represents a Python code solution, created by one of the following operators: 1. Draft , which generates the initial population of solutions; 2. Debug , which identifies and corrects errors within a given node; and 3. Improve , which enhances the solution of a given node to increase its performance according to evaluation criteria.

Additionally, the memory configuration dictates how each operator is selectively provided with previously produced artifacts, with well-scoped memory preventing issues such as context overload, mode collapse, and debug loops.

#### 3.1.4 LLM Backbones

For the data analysis, we use the following LLM backbones for the agents studied: o3 ( Jaech et al., 2024 ) , gpt-oss ( OpenAI, 2025 ) (20B and 120B), Llama Maverick ( Team, 2025b ) , Devstral ( Team, 2025a ) and CWM ( FAIR CodeGen Team, 2025 ) . Those represent different model sizes and architectures.

We conduct the controlled experiment discussed in Section 3.3 with the full-sized DeepSeek R1 model ( DeepSeek-AI et al., 2025 ) . All of the LLMs above use a 128K-token context window to ensure input coverage without truncation.

### 3.2 Measuring Ideation Diversity

Diversity can manifest in many aspects of machine learning engineering, such as data preprocessing, feature engineering, model development, and validation. In this analysis, our focus is limited to examining the diversity of machine learning models trained by agents.

All AI research agents in our study begin their exploration by generating at maximum five initial ideas to solve the task at hand (exactly five for greedy searches, and up to five for MCTS), using the Draft operator. To measure ideation diversity, we compare agents by extracting two pieces of information from these five initial ideas. First, we extract the high-level ML approach or architecture used by our agent (for example CNN ( LeCun and Bengio, 1998 ) , Transformer ( Vaswani et al., 2017 ) , Decision Trees); and second, we also extract the specific model employed by the agent, with variants grouped together (e.g., EfficientNet-B4 is grouped as EfficientNet ( Tan and Le, 2019 ) ).

We study whether the design of the agent has a significant impact on the diversity of ideas generated by comparing the distribution of models used by AI research agents.

To quantify diversity, we leverage the model architectures that the agent intends to train. From the distribution of model architectures, we compute the Shannon entropy (in base 2), quantifying the average uncertainty (and therefore diversity) of the model architecture used by the AI research agent.

### 3.3 Controlling Ideation Diversity

As part of our experimentation, we control the level of diversity using the system prompt in the LLM behind the agent. We compare two levels of diversity, the baseline agents and the agents with ablated diversity.

#### 3.3.1 Baseline agents

We run baseline (or control) agents (with 2 scaffolds, AIRA Greedy \text{AIRA}_{\text{Greedy}} , and AIRA MCTS \text{AIRA}_{\text{MCTS}} ), which, by default, are equipped with three mechanisms to enhance diversity. 1. Sibling memory , which provides to a new draft node the memory of its siblings, by including in the context descriptions of the solutions devised by the sibling nodes. 2. Prompt-adaptive complexity , which is a dynamic complexity cue within the system prompt aiming to guide the complexity of artifacts generated by the agents. For the first initial idea, we ask the agent to come up with an idea of minimal complexity. For the next two initial ideas, the system prompt asks for moderate complexity, and advanced complexity for the last two initial ideas. 3. Mention of diversity in the system prompt , asking the base model to come up with different aspects of the solution every time.

#### 3.3.2 Agents with ablated diversity

We remove prompt-adaptive complexity and the mention of diversity in the system prompt and we reuse sibling memory to request from the agent in the system prompt to come up with similar ideas.