Source-linked AI summary
AIDE: AI-Driven Exploration in the Space of Code
Zhengyao Jiang, Dominik Schmidt, Dhruv Srikanth, Dixing Xu, Ian Kaplan, Deniss Jacenko, Yuxiang Wu
TL;DR
Machine learning engineering requires costly trial-and-error, motivating automation that can search beyond predefined configurations. AIDE treats engineering as stateless code optimization with tree search, and achieves strong performance across Kaggle, MLE-Bench, and RE-Bench evaluations. Its benchmark scope is tempered by possible test-set mismatch and training-data contamination.
Problem
Machine learning engineering is time-consuming because stochastic data and optimization processes require extensive trial-and-error.
Method
AIDE uses an LLM-powered agent that explores machine-learning code through tree search, evaluating, summarizing, and refining candidate solutions.
Results
AIDE achieves superior performance across Kaggle tasks and research-oriented benchmarks, including MLE-Bench and RE-Bench.
Takeaways & Limitations
AIDE provides a principled way to combine iterative LLM prompting with tree-based exploration of code solutions.
Takeaways & Limitations
Kaggle scores may not be directly comparable because test sets can differ, and language-model training data may contaminate evaluations.
Abstract
from arXiv · showhide
Machine learning, the foundation of modern artificial intelligence, has driven innovations that have fundamentally transformed the world. Yet, behind advancements lies a complex and often tedious process requiring labor and compute intensive iteration and experimentation. Engineers and scientists developing machine learning models spend much of their time on trial-and-error tasks instead of conceptualizing innovative solutions or research hypotheses. To address this challenge, we introduce AI-Driven Exploration (AIDE), a machine learning engineering agent powered by large language models (LLMs). AIDE frames machine learning engineering as a code optimization problem, and formulates trial-and-error as a tree search in the space of potential solutions. By strategically reusing and refining promising solutions, AIDE effectively trades computational resources for enhanced performance, achieving state-of-the-art results on multiple machine learning engineering benchmarks, including our Kaggle evaluations, OpenAI MLE-Bench and METRs RE-Bench.
1 Introduction
AIDE addresses the time-consuming, stochastic trial-and-error involved in machine learning engineering by searching directly over code with an LLM-powered agent. It organizes candidate solutions as a tree, reuses promising solutions, and has been evaluated across increasingly challenging benchmarks.
- Machine learning engineering is time-consuming because stochastic data and optimization processes require extensive trial-and-error.
- Code-space optimization searches flexible, domain-informed solutions rather than predefined configurations, potentially improving sample efficiency for compute-bound tasks.
- AIDE automates machine learning engineering with an LLM-powered agent that organizes historical solutions in a tree and incrementally refines individual nodes.A hard-coded tree-search algorithm accumulates improvements guided by automated evaluations.
- AIDE was evaluated on tabular Kaggle tasks, deep-learning MLE-Bench tasks, and time-constrained AI research tasks in RE-Bench.
- The paper first formally specifies AIDE and then analyzes empirical evaluations from both the authors and independent benchmarks.
2 Preliminaries
AIDE models LLM-driven iterative problem solving as stateless optimization over independently evaluable candidate solutions. This replaces history-as-state reasoning with direct comparison of solutions and supports tree-search methods.
- General-purpose LLM agents often model tasks as POMDPs, treating the entire interaction history as the state.Appending all historical data can eventually fill the model’s context window and limit scalability.
- AIDE defines a solution space S, a stateless objective h, and the goal of finding the solution s that maximizes h(s).
- Independent evaluation lets AIDE directly compare candidate solutions instead of unrolling a single long-horizon decision process.This perspective aligns with optimization methods such as tree search.
3 Methodology
AIDE searches machine-learning code by repeatedly generating, evaluating, recording, summarizing, and selecting solutions in a tree. Its policy separates drafting, debugging, and improvement while concise summaries preserve useful history without unbounded prompts.
- Overall design: Combining search, coding, and summarization enables systematic code-space exploration while avoiding an ever-increasing prompt history.
- High-level algorithm: AIDE’s loop proposes a solution, evaluates it, adds it and its score to the tree, selects the next base node, and returns the best solution found.
- Solution tree: Each tree node is a script, each edge records an improvement attempt, and the search policy selects which solution to improve next.The coding operator drafts, debugs, or refines scripts using summarized tree context.
- Search policy: The search policy drafts initial solutions, debugs buggy nodes within a bounded depth, and otherwise improves the best non-buggy solution.These heuristics encourage diverse initialization and limit debugging attempts.
- Coding operator: The coding operator uses specialized prompts for drafting new programs, debugging execution failures, and making atomic changes whose effects are measurable.Examples include changing optimizers or adding regularization.
- Solution tree: Figure 1 depicts scripts as tree nodes, coding-operator transitions as arrows, and branches ending in bugs or improved and optimal solutions.
- Summarization and prompts: A summarization operator extracts relevant improvement ideas, metrics, hyperparameters, and debugging hints to guide later proposals without saturating prompts.Each revision remains stateless while prior information still informs subsequent code generation.
- Summarization and prompts: Machine-learning prompts also include a lightweight data preview containing metadata such as dataset size, column names, and data splits.The preview is not a complete exploratory data analysis pipeline.
4 Evaluation
AIDE was evaluated on tabular Kaggle tasks, MLE-Bench competitions, and AI R&D environments using human-relative and agent-comparison metrics. Results were competitive across settings, while performance varied by task and evaluation limitations affect interpretation.
- 4.1 Weco Kaggle Benchmark: AIDE achieved 51.38% Exceeds % of humans and surpassed the human median in 50% of 16 Weco-Kaggle Lite tabular tasks.H2O AutoML and LangChain AutoGPT scored 35.34% and 32.34%, respectively, on Exceeds % of humans.
- 4.1 Weco Kaggle Benchmark: The Kaggle evaluation may not be directly comparable to private-test performance because its manually defined holdout test set can differ from Kaggle’s private test set, and model-training contamination is possible.These limitations can introduce variance in percentiles and affect interpretation of the benchmark results.
- 4.2 AIDE in MLE-Bench: On MLE-Bench, AIDE with o1-preview earned medals in 16.9% of competitions at pass@1, nearly four times the rate of OpenHands.The benchmark covers 75 real Kaggle competitions, and AIDE was reported as the top-performing agent framework with state-of-the-art language models.
- 4.2 AIDE in MLE-Bench: On MLE-Bench Lite, adding AIDE to o1-preview increased valid submissions from 63.6% ± 4.5% to 92.4% ± 2.6% and above-median solutions from 13.6% to 59.1% ± 4.5%.Gold-medal achievement rose from 6.1% ± 2.6% to 21.2% ± 6.9%, while overall medals rose from 7.6% ± 2.6% to 36.4% ± 7.9%; all improvements were statistically significant at p < 0.01.
- 4.3 AIDE in RE-Bench: On seven RE-Bench AI R&D tasks, AIDE outperformed humans within six hours, but humans eventually caught up and AIDE struggled with larger codebases or multi-step improvements.AIDE exceeded human performance in Optimize a Kernel, while it repeated local patches in Agent for Rust CodeContests.
5 Related Work
Prior work automates machine learning search over predefined spaces, while newer LLM agents enable more flexible code-space exploration. Related agent research spans general-purpose planning systems and specialized domain agents, including automated model and architecture search.
- LLM Agents: General-purpose agents interleave natural-language reasoning with task execution across information retrieval and multimodal processing.ReAct and HuggingGPT exemplify broad task-oriented agent frameworks.
- Automated Machine Learning: AutoML automates model selection, hyperparameter tuning, and pipeline configuration by searching predefined model spaces.Examples include Bayesian optimization, genetic programming, meta-learning, and ensemble-based systems.
- Neural Architecture Search: Neural Architecture Search automatically designs neural network topologies using reinforcement learning, evolution, or differentiable optimization.Differentiable approaches reduce computational cost by optimizing over a relaxed search space.
6 Conclusion
AIDE is presented as an LLM agent for machine learning engineering that drafts, debugs, and refines solutions through tree-based code exploration. The paper reports strong performance on Kaggle and research-oriented benchmarks, with third-party results extending to additional AI R&D tasks.
- Conclusion: AIDE systematically drafts, debugs, and refines code solutions through tree-based exploration.The approach combines iterative LLM prompting with structured exploration of code solutions.
- Conclusion: AIDE achieves superior performance on Kaggle tasks and more research-oriented benchmarks.The conclusion reports these results without specifying a single aggregate metric.
- Conclusion: Third-party experiments show AIDE generalizing to neural architecture search, Triton Kernel optimization, and other AI R&D tasks.These examples extend beyond the tabular machine learning tasks for which AIDE was developed.
A Baseline Specifications
The supplied passage identifies Table 4 as reporting baseline hyperparameters, but provides no hyperparameter values or further specification.
- A Baseline Specifications: Table 4 is labeled as reporting baseline hyperparameters.No individual hyperparameters are included in the supplied passage.
A.1 H2O AutoML Baseline
The H2O AutoML baseline searches predefined algorithms and hyperparameter combinations, then evaluates individual models and stacked ensembles. The comparison also includes an AutoGPT setup with task descriptions, file tools, Python evaluations, and iterative human interaction for ChatGPT.
- H2O AutoML: H2O AutoML searches a predefined set of algorithms, including forests, generalized linear models, gradient boosting, neural networks, and stacked ensembles.The supplied list names DRF/XRT, GLM, H2O GBMs, DeepLearning, and stacked ensembles.
- H2O AutoML: H2O AutoML performs random search over predefined hyperparameter combinations instead of exhaustive grid search.This avoids the computational expense of exhaustive enumeration.
- H2O AutoML: It evaluates individual models and stacked ensembles with cross-validation, ranking them by task-appropriate metrics such as accuracy, AUC, or RMSE.Stacked ensembles combine predictions from the best-performing models.
- AutoGPT Baseline: The AutoGPT baseline receives a competition task descriptor and tools for file operations and Python REPL evaluations.The descriptor includes dataset, file, evaluation, and submission information.
- AutoGPT Baseline: Its prescribed workflow iterates over models or feature selections, trains models, generates predictions, and saves scripts and submissions.The rules prohibit modifying the train and test files and changing column-name capitalization.
- Human ChatGPT Baseline: The human ChatGPT competitor solves competitions using overview and data-tab information, with iterative interaction addressing hallucinations and outdated packages.The comparison uses gpt-4-0125-preview for ChatGPT rather than AutoGPT.
B Analysis of AIDE
Aggregated code complexity generally increases with AIDE’s iterative steps, indicating progressively more elaborate generated solutions.
- Aggregated code complexity generally rises as AIDE takes more iterative steps.The measures combine LOC, LLOC, Volume, N1, and MI; complexity dips slightly at first before steadily increasing.
B.2 Cost Analysis
AIDE’s per-task inference costs are usually moderate, though extensive prompting makes some tasks more expensive.
- Approximately $2.50 per task is the upper observed cost, while most tasks remain below $1.50.The benchmark used GPT-4 Turbo (gpt-4-0125-preview) and early-2024 pricing data.
- The reported costs reflect token usage and minimal manual intervention.
C Weco Kaggle Benchmark
The Weco-Kaggle benchmark covers tabular competitions with varying evaluation availability, GPU requirements, and compressed data sizes.
- The benchmark lists competitions with local evaluation and Kaggle-submittable indicators.The table also records whether GPU resources are needed and gives compressed data sizes.