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

Key points

• Problem: AI research agents automate the design, implementation, and training of ML models, but the field is still in its infancy and the key factors behind the success or failure of agent trajectories are not fully understood. Long multi-step trajectories that are often guided by heuristic search algorithms, together with tool use, make error analysis far more complicated than in classic ML setups, and collecting samples large enough for meaningful analysis and design ablations can be computationally prohibitive. The authors start from the hypothesis that ideation diversity is a key bottleneck in AI research agent performance.
• Input and output: This is an empirical study that analyzes existing agent trajectories at scale and runs a controlled experiment. The input is agent trajectories on MLE-bench tasks from several LLM backbones and scaffolds, and the outputs are a diversity score, the Shannon entropy (base 2) of the model architectures extracted from the initial ideas (at most five) produced by the Draft operator (which generates the initial solutions), its correlation with medal rate and other performance metrics, and the performance change when diversity is altered through the system prompt.
• Method: The analyzed scaffolds are AIDE (AI-Driven Exploration, a greedy tree search over Python solutions), AIRA Greedy, and AIRA MCTS, all of which grow a tree of Python solution nodes with Draft, Debug, and Improve operators. In the controlled experiment, the baseline agents (AIRA Greedy and AIRA MCTS) have three diversity mechanisms by default: sibling memory that puts descriptions of sibling nodes' solutions into the context, prompt-adaptive complexity that asks for minimal complexity for the first idea, moderate for the next two, and advanced for the last two, and a request for diversity in the system prompt. The low-diversity agents remove prompt-adaptive complexity and the diversity request in the system prompt and reuse sibling memory to ask, through the system prompt, for similar ideas, and by changing only the diversity-related parts of the prompt the authors intend to leave other aspects such as implementation quality as they are.
• Implementation: The analysis backbones are o3, gpt-oss 20B and 120B, Llama Maverick, Devstral, and CWM (an open-weights LLM from Meta for research on code generation with world models), and 6 backbones with 2 scaffolds were run on the 75 MLE-bench tasks with 10 to 20 seeds each, giving 11,000 trajectories, roughly 1,200,000 search nodes, and 264,000 GPU hours. The controlled experiment ran AIRA Greedy and AIRA MCTS with the full-sized DeepSeek R1 on MLE-bench lite (22 tasks) with 10 seeds each, and all LLMs used a 128K-token context. Error bars in the controlled experiment are 95% confidence intervals from stratified bootstrapping with the rliable library.
• Results: Across the 75 tasks, architecture entropy correlated with medal rate at Pearson r 0.57, with average normalized score at 0.72, and with percentile at 0.66, and agents on o3 and gpt-oss 20B and 120B used 3.5 distinct architectures on average in their 5 initial ideas versus 2.8 on average for Llama Maverick, Devstral, and CWM. In the DeepSeek R1 controlled experiment, reducing diversity lowered the MLE-bench lite medal rate from 45.5 to 38.6 for AIRA Greedy (6.9 points) and from 47.0 to 38.6 for AIRA MCTS (8.4 points), and valid submission rate fell from 98% to 92% and 90%, respectively. The valid-submission drop came mainly from two text normalization competitions (English and Russian), where low-diversity agents repeatedly tried to implement the same T5 model (a text-to-text transformer) and consistently failed with timeouts, and the authors estimate that these two competitions account for around half of the medal-rate decline.
• Limits: The findings come from MLE-bench only, and the authors hypothesize that they are likely to generalize to other ML tasks, leaving other benchmarks to future work. Medal criteria change with the number of teams in a competition, in about 30% of MLE-bench competitions the relative gap between the best score and the bronze threshold is under 3%, agents are evaluated on custom test sets that differ from Kaggle's private test sets, and agent performance drops sharply on competitions after 2022. The authors state that second-order effects of modifying the system prompt are hard to track, and in an appendix experiment that varied temperature for AIRA Greedy with all diversity mechanisms removed, metrics other than Elo, including medal rate, showed no effect.

Paper links

External research summaries. These are not HDATF publications or measured product results.