MLE-bench: Evaluating Machine Learning Agents on Machine Learning Engineering

Key points

• Problem: Language models perform well on many coding benchmarks and are progressing on ML tasks such as architecture design and model training, but few benchmarks holistically measure autonomous end-to-end ML engineering. The authors aimed for tasks that are challenging and representative of contemporary ML engineering work, with evaluation results that can be compared to human-level performance.
• Input and output: The agent receives a Kaggle competition's description, dataset, and submission format, trains models in an offline environment, and outputs a submission.csv file. The submission is scored by grading code that re-implements the competition metric and is converted into a medal by comparing it against the Private leaderboard as if the agent had entered the competition at the time. This is a benchmark paper for evaluating ML engineering agents.
• Method: From 5673 completed competitions in the Meta Kaggle dataset, the authors dropped Community Competitions, manually screened the remaining 586, and selected 75, labeled by the time an experienced ML engineer would need for a sensible solution (excluding model training) as Low (22, under 2 hours), Medium (38, 2 to 10 hours), or High (15, over 10 hours). Where no public test set existed, they carved a new test set out of the public training data (at a 10% ratio unless otherwise stated), and the headline metric is the share of attempts that earn any medal (bronze or above) under Kaggle's medal thresholds, which depend on the number of participating teams. Agents may not write predictions directly from their own knowledge instead of using a model, or view solutions online, a log inspection tool and the Dolos plagiarism detector check this, and attempts with over 60% code similarity to any of the top 50 Kaggle notebooks are disqualified and flagged for further review.
• Implementation: In the main experiments, each competition attempt runs for up to 24 hours in an Ubuntu 20.04 Docker container holding the dataset, a submission validation server, and Python packages, on a machine with 36 vCPUs, 440GB of RAM, and one Nvidia A10 GPU, and seeds default to 3, with 16 for o1-preview (AIDE) and 36 for GPT-4o (AIDE). The three scaffolds are AIDE, MLAB (the ResearchAgent from MLAgentBench), and OpenHands (the CodeActAgent from the OpenHands platform), and four models were compared on AIDE: o1-preview, GPT-4o, Claude 3.5 Sonnet, and Llama 3.1 405B. According to the README, the GitHub repository openai/mle-bench contains dataset preparation scripts, per-competition grading scripts, the evaluated agents, the setup for the base Docker image (mlebench-env) used by agents, and a rule violation detector and a plagiarism detector, and it recommends a Lite evaluation on only the 22 Low complexity competitions when cost is a concern.
• Results: In Table 2, among the four models on the AIDE scaffold, o1-preview earned a medal in 16.9% of competitions on average, the highest, followed by GPT-4o at 8.7%, Claude 3.5 Sonnet at 7.6%, and Llama 3.1 405B at 3.0%, while GPT-4o scored 0.8% with MLAB and 4.4% with OpenHands. In the attempts experiment, o1-preview rose from 16.9% at pass@1 to 34.1% at pass@8 (the share of competitions with a medal in at least one of 8 attempts), and in the experiment that extended the time limit per competition to 100 hours and the maximum number of AIDE nodes to 5,000, GPT-4o (AIDE) rose from 8.7% (at 24 hours) to 11.8%. In the Table 3 hardware experiment GPT-4o (AIDE) scored similarly with no GPU (9.1%), one GPU (8.7%), and two GPUs (10.2%), and in Table 4 it also scored similarly on the original descriptions (8.5%) and on obfuscated descriptions run with 10 seeds (8.4%).
• Limits: Because the data are public Kaggle competitions, models may have trained on competition details, solutions, and even the datasets including the test sets, and the authors state that although their experiments found no systematic contamination effect for GPT-4o, reuse of high-level strategies is hard to detect and there are no guarantees about future models. They note that MLE-bench selects competitions with clear problem statements, clean datasets, and clear metrics, so it covers part of the capabilities needed for real-world AI research and development, where the problem statement often may not even be clear and figuring out the dataset and metrics is part of the problem. One run of the main setup takes 24 hours times 75 competitions, or 1800 GPU hours, and o1-preview (AIDE) used on average 127.5M input tokens and 15.0M output tokens for one seed of 75 competitions.

Paper links

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