Inside OpenAI’s MLE-Bench: A New Benchmark for Evaluating Machine Learning Engineering Capabilities of AI Agents

Author(s): Jesus Rodriguez

Originally published on Towards AI.

I recently started an AI-focused educational newsletter, that already has over 170,000 subscribers. TheSequence is a no-BS (meaning no hype, no news, etc) ML-oriented newsletter that takes 5 minutes to read. The goal is to keep you up to date with machine learning projects, research papers, and concepts. Please give it a try by subscribing below:

Coding the engineering are one of the areas that has been at the frontiers of generative AI. One of the ultimate manifestations of this proposition is AI writing AI code. But how good is AI in traditional machine learning(ML) engineering tasks such as training or validation. This is the purpose of a new work proposed by OpenAI with MLE-Bench, a benchmark to evaluate AI agents in ML engineering tasks.

MLE-Bench is a new benchmark introduced by OpenAI to evaluate the performance of AI agents on complex machine learning engineering tasks. The benchmark is specifically designed to assess how well AI agents can perform real-world MLE work, such as training models, preparing datasets, and running experiments.

MLE-Bench leverages the challenges and structure of Kaggle competitions. Kaggle, a popular platform for data science and ML competitions, provides a rich source of diverse and challenging tasks that mirror real-world MLE problems. The benchmark comprises 75 carefully selected Kaggle competitions, covering various domains such as natural language processing, computer vision, and signal processing.

The selection of Kaggle competitions for MLE-Bench is driven by several key factors:

· Real-World Relevance: Kaggle competitions often deal with actual problems faced by ML practitioners, offering a practical gauge of an agent’s ability to handle real-world scenarios.

· Task Diversity: The 75 competitions included in MLE-Bench span 15 different problem categories, ensuring a broad assessment of ML engineering skills.

· Human Baselines: Kaggle’s public leaderboards offer a valuable benchmark for comparison. By evaluating agents against human performance, researchers can gain a clearer understanding of how agent capabilities measure up to human expertise.

Components of MLE-Bench

Each competition in MLE-Bench is structured to simulate a real Kaggle competition environment:

1. Competition Description: Each competition includes a description scraped from the original Kaggle competition page, providing the agent with an overview of the task and the dataset.

2. Dataset: MLE-Bench uses either the original dataset from the Kaggle competition or a newly created train-test split if the original test set is not publicly available. The authors took care to ensure that new splits closely mirrored the original data distribution.

3. Grading Code: Each competition has a dedicated grading code that allows for local evaluation of submissions based on the competition’s defined metric.

4. Leaderboard: A snapshot of the private leaderboard from the corresponding Kaggle competition is used to rank agent submissions against human performance.

MLE-Bench Evaluation Metrics

MLE-Bench utilizes several metrics to assess agent performance:

1. Leaderboards: Agent performance is contextualized using the private leaderboards from each Kaggle competition. This allows for a direct comparison between agents and human participants.

2. Medals: MLE-Bench adopts Kaggle’s medal system (bronze, silver, and gold), awarded based on the agent’s performance relative to the private leaderboard. The medal thresholds are adjusted based on the number of participants in each competition to ensure a consistent measure of achievement across competitions.

3. Headline Metric: To provide an overall measure of success, MLE-Bench uses the percentage of competition attempts in which an agent achieves any medal (bronze or above). This demanding metric is intended to reflect the level of expertise demonstrated by the best human Kagglers, who have earned medals in numerous competitions.

4. Raw Scores: In addition to medals, MLE-Bench reports the raw score achieved by agents on each competition. This provides a more granular understanding of agent performance on specific tasks, though it is difficult to aggregate these scores across competitions due to variations in metrics.

MLE-Bench Findings: Promising Results and Challenges

Initial experiments with MLE-Bench highlight both the potential of AI agents in MLE and the remaining challenges.

1. Agent Performance: The best-performing agent, o1-preview combined with the AIDE scaffolding, achieved a medal in 16.9% of the competitions. This result is noteworthy, considering that achieving the “Kaggle Grandmaster” status requires 5 gold medals.

2. Impact of Attempts: Allowing agents multiple attempts per competition led to significant performance improvement. Both o1-preview and GPT-4o showed a consistent increase in the percentage of medals earned as the number of attempts increased.

3. Compute Scaling: Interestingly, varying the amount of computational resources provided to GPT-4o did not result in significant performance differences. This suggests that current agents may not effectively adapt their strategies based on available hardware.

4. Time Scaling: Providing GPT-4o with a longer time limit (100 hours) per competition allowed the agent to gradually improve its solutions and achieve more medals over time.

The experiments also revealed areas where current agents struggle:

Valid Submission Creation: Despite access to a validation server that checks submission validity, many agents failed to consistently produce valid submissions. This indicates a need for improvement in how agents utilize feedback and validation tools.

Resource Management: Agents often exhibited poor resource management, executing commands that overloaded system resources. This points to a lack of awareness regarding time and computational constraints, hindering their ability to efficiently solve tasks.

The use of public Kaggle competitions in MLE-Bench raises concerns about potential contamination. Contamination refers to artificially inflated benchmark performance due to models memorizing or leveraging information from publicly available solutions during training.

To address these concerns, MLE-Bench incorporates several mitigation strategies:

1. Rule-Breaking Detection: A tool utilizing GPT-4o is employed to analyze agent logs, identifying potential rule violations such as manual submission creation or attempts to access unauthorized resources.

2. Plagiarism Detection: The Dolos code plagiarism detection tool is used to compare agent-generated code against the top-performing notebooks from the corresponding Kaggle competition. Submissions with a high similarity score are flagged for review, ensuring that agents do not simply copy existing solutions.

3. Familiarity Analysis: Researchers investigated the potential impact of model familiarity with competition discussions by measuring the correlation between GPT-4o’s familiarity with winning solutions and its performance. The analysis found no correlation, suggesting that GPT-4o’s performance is not systematically inflated due to prior exposure to competition information.

4. Obfuscated Competition Descriptions: To further examine the influence of contamination, an experiment was conducted using rewritten competition descriptions designed to obscure their origin. GPT-4o’s performance remained consistent between the original and obfuscated descriptions, indicating that the model does not rely heavily on recognizing specific competitions or their associated solutions.

What Makes MLE-Bench Special?

MLE-Bench represents a significant step towards evaluating and understanding the capabilities of AI agents for ML engineering tasks. The benchmark has several important implications:

Measuring Agent Progress: MLE-Bench provides a robust framework for tracking the progress of AI agents in autonomously performing complex MLE tasks.

Understanding Acceleration Risks: The benchmark helps assess the potential of AI agents to accelerate scientific discovery. While this acceleration holds great promise, it also raises concerns about potential misuse or unintended consequences if agents are not properly aligned with human values and safety guidelines.

Guiding Responsible Development: By highlighting the capabilities and limitations of current agents, MLE-Bench informs the responsible development of AI systems.

OpenAI acknowledges several limitations in MLE-Bench, including the potential for subtle forms of contamination that are difficult to detect and the limited coverage of the full spectrum of AI research and development activities. They encourage the development of additional evaluations that address these limitations and further our understanding of agent capabilities in ML engineering.

MLE-Bench is open-sourced, encouraging transparency and collaboration in AI research. As AI agents continue to evolve, benchmarks like MLE-Bench will play a critical role in ensuring that the development of advanced AI technologies remains aligned with human goals and values.

Join thousands of data leaders on the AI newsletter. Join over 80,000 subscribers and keep up to date with the latest developments in AI. From research to projects and ideas. If you are building an AI startup, an AI-related product, or a service, we invite you to consider becoming a sponsor.

Published via Towards AI