MLE-bench: Evaluating Machine Learning Agents on Machine Learning Engineering — part 1/3

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

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

###### Abstract

We introduce MLE-bench, a benchmark for measuring how well AI agents perform at machine learning engineering. To this end, we curate 75 ML engineering-related competitions from Kaggle, creating a diverse set of challenging tasks that test real-world ML engineering skills such as training models, preparing datasets, and running experiments. We establish human baselines for each competition using Kaggle’s publicly available leaderboards. We use open-source agent scaffolds to evaluate several frontier language models on our benchmark, finding that the best-performing setup — OpenAI’s o1-preview with AIDE scaffolding — achieves at least the level of a Kaggle bronze medal in 16.9% of competitions. In addition to our main results, we investigate various forms of resource-scaling for AI agents and the impact of contamination from pre-training. We open-source our benchmark code ( github.com/openai/mle-bench/ ) to facilitate future research in understanding the ML engineering capabilities of AI agents.

## 1 Introduction

Language models (LMs) have achieved impressive performance on many coding benchmarks ( Chen et al., 2021 ; Hendrycks et al., 2021 ; Austin et al., 2021 ; Li et al., 2022 ) and are making progress on a variety of machine learning tasks, such as architecture design and model training ( Zheng et al., 2023 ; Huang et al., 2024b ) . LMs have also been adopted into programming tools ( Kalliamvakou, 2022 ) , and progress in agent scaffolding has increasingly automated developer workflows ( cognition.ai, 2024 ; Dohmke, 2024 ) . However, while there has been a surge in development on model and agent capabilities, few benchmarks holistically measure autonomous end-to-end ML engineering.

We introduce MLE-bench, an offline Kaggle competition environment for assessing how well AI agents can perform difficult machine learning engineering (MLE) tasks. We built MLE-bench to be a robust measure of real-world progress in autonomous ML engineering agents, focusing on two main design choices: (i) selecting tasks that are challenging and representative of contemporary ML engineering work, and (ii) being able to compare evaluation results to human-level performance.

The resulting benchmark consists of 75 diverse Kaggle competitions across a variety of domains, including natural language processing, computer vision, and signal processing. Many of the competitions are contemporary challenges with real-world value, such as OpenVaccine: COVID-19 mRNA Vaccine Degradation Prediction ( Das et al., 2020 ) and the Vesuvius Challenge for deciphering ancient scrolls ( Lourenco et al., 2023 ) . The total value of prizes awarded across the 75 competitions is $1,948,016 ($25,974 per competition on average).

AI agents that autonomously solve the types of challenges in our benchmark could unlock a great acceleration in scientific progress, a prospect that is exciting but also warrants careful understanding of model progress in order to deploy advancements in a safe and controlled manner. For example, MLE-bench can be used as a measure for model autonomy in OpenAI’s Preparedness Framework ( OpenAI, 2023 ) , autonomous capabilities in Anthropic’s Responsible Scaling Policy ( Anthropic, 2023 ) , and ML R&D in Google DeepMind’s Frontier Safety Framework ( Google DeepMind, 2024 ) .

We find that, when combined with open-source scaffolds, leading LMs achieve meaningful scores on our benchmark. The best-performing agent we evaluated, o1-preview with AIDE, uses scaffolding purpose-built for Kaggle competitions and achieves a medal in 16.9% of competitions on average. We find that performance significantly improves when agents are given multiple attempts per competition; for example, o1-preview’s score doubles from 16.9% using pass@1 to 34.1% using pass@8. Similarly, GPT-4o scores 8.7% given 24 hours to attempt each competition, but 11.8% when given 100 hours. In general, we found that agents can score well on competitions that can be solved with well-known approaches but struggle to debug issues and recover from missteps.

Our contributions include: 1. MLE-bench – a benchmark of 75 offline Kaggle competitions for evaluating ML engineering capabilities of AI agents, carefully handcrafted by a team of ML engineers.

2. Large-scale evaluations of state-of-the-art models and agent frameworks, revealing new information about the prospects and limits of autonomous ML engineering agents.

3. Experiments on scaling resources for agents, including scaling agent runtime, hardware resources, and pass@k attempts, exploring performance ceilings for present-day agents.

4. Experiments investigating the relationship between dataset contamination and agent performance, as well as agent-monitoring tools to detect plagiarism and cheating.

## 2 MLE-bench

MLE-bench consists of 75 machine learning engineering tasks manually sourced from Kaggle to reflect a core set of day-to-day skills that ML engineers use in frontier labs.

Kaggle is a platform that hosts data science and ML competitions requiring participants to build predictive models to solve challenges, often using real-world datasets. Participants compete to achieve the best score on a metric pre-defined for each competition, and are ranked on a leaderboard against one another. Bronze, silver, and gold medals are awarded for top competition results.

### 2.1 Dataset Curation

Each sample in MLE-bench is a Kaggle competition consisting of:

• A description scraped from the “Overview" and “Data" tabs of the competition website.

• The competition dataset , in most cases using a new train-test split (more details below).

• Grading code used to evaluate submissions locally.

• A snapshot of the competition’s leaderboard used to rank submissions against humans.

To arrive at the set of competitions constituting MLE-bench, we begin with the 5673 completed Kaggle competitions listed on the Meta Kaggle dataset 1 1 1 kaggle.com/datasets/kaggle/meta-kaggle (accessed May 15th, 2024)) . We exclude Community Competitions since their quality is less rigorously vetted than other competitions. We manually screen the remaining 586 competitions for relevance to modern-day ML engineering. We exclude competitions where we cannot replicate the grading procedure or cannot recreate reasonable train-test splits. See Appendix A.1 for the full list of screening criteria.

Additionally, we manually annotate the problem type of each competition (e.g. text classification, image segmentation, etc.). We also annotate each competition with a complexity level: Low if we estimate that an experienced ML engineer can produce a sensible solution in under 2 hours excluding the time taken to train any models, Medium if it takes between 2 and 10 hours, and High if it takes more than 10 hours. See Appendix A.2 for more details.

From this process, we select 75 competitions to include in MLE-bench, comprising 22 competitions Low in complexity (30 % \% ), 38 Medium (50 % \% ), and 15 High (20 % \% ). We include an additional 7 competitions as a development split, for developing agents without over-fitting to the test set. We recommend using the Low complexity split if using all splits is too resource-intensive.

For each competition, we use the original dataset if publicly available, although Kaggle competitions often do not release the test set even after the competition ends. In such cases, we manually create new train and test splits based on the publicly available training data 2 2 2 We discuss the splits further in Appendix A.7 . . We take care to ensure that the distributions of the original and reconstructed test sets are similar by checking that the example submission scores similarly on both sets. We take the new test set to be 10% of the original train set, except for when it didn’t make sense to do so 3 3 3 e.g. doing so for the “New York City Taxi Fare Prediction” competition would result in a test set 100x larger than the original, so we opted to maintain the original train/test ratio in such cases. . Due to these measures, we expect scores on the MLE-bench competition test sets to be comparable to human scores on the competition’s leaderboard, especially on average.

Finally, we implement the grading logic for each competition based on the described evaluation metric in the competition’s description, so that submissions can be graded locally. Evaluation metrics vary by competition, from standard metrics like area under the receiver operating characteristic (AUROC) to domain-specific loss functions.

### 2.2 Metrics

##### Leaderboards

We contextualize MLE-bench performance using the leaderboards 4 4 4 Snapshots of the Private leaderboard were taken between May and August 2024. of each Kaggle competition. On Kaggle, competitions may have two leaderboards: “Public" and “Private." We found that Kaggle submissions sometimes overfit to the Public leaderboard, so we opt to use the Private leaderboard.

##### Medals

Kaggle awards bronze, silver, and gold medals to top competitors based on their performance relative to the leaderboard ( Table 1 ). Similarly, MLE-bench awards medals to agents’ submissions by comparing them against the Private leaderboard, as if the agent were participating in the competition at the time. The thresholds for bronze, silver, and gold vary depending on the number of participants in a competition such that a given medal should always reflect a similar level of achievement across different competitions. Although not all competitions on Kaggle award medals, in MLE-bench we apply the medal thresholding logic to all competitions.

##### Headline metric

To provide a singular metric for MLE-bench, we calculate the percentage of attempts that are awarded any medal (bronze and above). This is designed to be a challenging metric, with a ceiling comparable to the achievements of the very best human Kagglers after years of cumulative effort. Only nine humans have ever achieved medals on 75 different Kaggle competitions 5 5 5 According to Meta Kaggle (last accessed 23rd October 2024), Kagglers titericz , kazanova , mathurinache , lucamassaron , mikeskim , abhishek , alexxanderlarko , confirm , and coreacasa each have more than 75 unique competition medals. An impressive achievement! .

##### Raw Scores

We also report the raw score achieved by models on each competition. This is useful to track competition-specific progress, though it is hard to aggregate scores across competitions since each competition uses different metrics.

### 2.3 Setup

MLE-bench is designed to be agnostic to the methods used to solve it, requiring only a CSV file to be submitted to each competition for grading. Nevertheless, we encourage developers to report certain details when evaluating their agents on this benchmark. Specifically, developers should mention the models and scaffolding used, whether the agent had internet access, available hardware, runtime, the inclusion of any partial or complete solutions to Kaggle competitions in the agent’s prompts, and any other significant deviations from our experimental setup described in Section 3 .

##### Validating submissions

Real-life Kaggle competitions often allow participants to make up to 5 submissions a day to the Public leaderboard, which checks the validity of their submission and provides a score. Analogously, we allow agents to access a local validation server to check the validity of their submission, though the tool does not provide a score (our tool uses the grader to check if a submission is valid, or provides an error message in the case of invalid submissions). There are no restrictions on how often agents can use this tool.

#### 2.3.1 Rules

Submissions must be produced by a model separate from the agent; the agent is forbidden from writing predictions directly to its submission file using its own knowledge of the world. This prevents agents from simply recalling labels from their pre-training data that it may have memorized, and ensures the agent engages in ML engineering. Agents are also forbidden from viewing solutions online, which can often be found on Kaggle or GitHub.

##### Rule-breaking detection

To ensure agents adhere to the rules, we provide a tool that inspects agent logs using GPT-4o. Specifically, the tool checks if the agent has broken the rules by manually writing the submission file without using a model, calling another external LLM API for assistance, or attempting to access unauthorized resources. Further details can be found in Appendix A.3 .

##### Plagiarism detection

To prevent plagiarism, we use the source code plagiarism detection tool Dolos ( Maertens et al., 2024 ) to compare the agent’s submitted code against the top 50 associated notebooks from the relevant Kaggle competition. These notebooks are publicly available on Kaggle and often contain successful solutions. For the purpose of our benchmark, we disqualify any attempts where the agent submits code with a high similarity score (over 60%) to any notebook and flag them for further review.

These rules are designed to prevent cheating. We further discuss the risk of score inflation via train-time contamination in Section 4 , and the limits of our mitigations in Section 6 .

## 3 Experiments and Results

In our experiments, we run agents in an Ubuntu 20.04 Docker container containing the dataset, validation server, and Python packages that might be helpful for ML engineering. Containers are executed in a secure cluster environment. For each of the 75 competitions, agents have a maximum of 24 hours to produce a submission. On each run, agents have access to a machine with 36 vCPUs, 440GB RAM, 4095 GiB SSD, and a single Nvidia A10 GPU. We repeat all experiments with 3 seeds (that is, 3 runs per competition) to compute the mean and standard error unless otherwise specified. Full details of our execution environment and scaffolds can be found in Appendices A.5 and A.6 .

### 3.1 Main experiment

##### Varying scaffolding

To determine the best-performing scaffold, we evaluate GPT-4o 6 6 6 gpt-4o-2024-08-06 using three open-source scaffolds: AIDE ( Schmidt et al., 2024 ) , ResearchAgent (referred to as “MLAB") from MLAgentBench ( Huang et al., 2024b ) , and CodeActAgent (referred to as “OpenHands") from the OpenHands platform ( Wang et al., 2024 ) . We make minor modifications to each scaffold to enhance their performance on the benchmark (details in Appendix A.6 ), and report results in Table 2 .

We find that GPT-4o (AIDE) achieves more medals on average than both MLAB and OpenHands (8.7% vs. 0.8% and 4.4% respectively), despite making a similar number of valid submissions. Notably, AIDE is purpose-built for Kaggle competitions, whereas the other scaffolds are general-purpose. See Figure 2 for a snippet of each scaffold’s trajectories.

##### Varying models

Taking the best-performing scaffold (AIDE) from the previous experiment, we experiment with changing the underlying model. We evaluate four different models 7 7 7 We attempted to evaluate Gemini-1.5-Pro (Gemini-1.5-Pro-002) but found API calls repeatedly blocked due to model outputs being flagged for recitation. with AIDE: o1-preview and GPT-4o (OpenAI), Claude 3.5 Sonnet 8 8 8 claude-3-5-sonnet-20240620 (Anthropic), and Llama 3.1 405B 9 9 9 meta-llama/llama-3.1-405b-instruct served by https://openrouter.ai/ (Meta).