MLE-STAR: Machine Learning Engineering Agent via Search and Targeted Refinement

Key points

• Problem: Existing machine learning engineering (MLE) agents lean on what the LLM already knows, so they drift toward familiar methods such as scikit-learn and, for example, propose logistic regression even for a text classification competition. They also rewrite the whole code structure in each iteration, so they move on to other steps before exploring one component such as feature engineering in depth.
• Input and output: This is an agent system for Kaggle-style ML tasks. The input is a task description (task type, data modality, score function) plus datasets that may span several files, and the output is a single Python script that trains a model and prints its hold-out validation score. At the end a separate agent adds test-data loading and submission file creation to that script.
• Method: A retriever agent first uses Google Search to find task-specific models with example code, writes a simple script for each, and merges them in order of validation score into an initial solution. In each outer loop an ablation-study script picks the code block with the largest effect on performance, and in the inner loop a planner proposes several block-only refinements using earlier scores and keeps the best one. Solutions built in parallel are merged with ensemble strategies that the agent proposes and refines itself, while a debugging agent, a data leakage checker and a data usage checker fix runtime errors, test-information leaks and unused data.
• Implementation: The default LLM is Gemini-2.0-Flash, with extra runs on Gemini-2.5-Pro and Claude-Sonnet-4; the setup retrieves 4 models, runs 4 inner loops and 4 outer loops, builds 2 parallel solutions, explores ensembles for 5 rounds, and caps each task at 24 hours. On 96 vCPUs, 360 GB of memory and 8 V100 GPUs with 16 GB each, one solution took 14.1 hours on average, against 15.4 hours for AIDE. The appendix states that example outputs such as the retriever output, the ablation code and the final solution code are placed in the authors' GitHub repository.
• Results: On the 22 MLE-bench Lite competitions with 3 seeds, MLE-STAR with Gemini-2.5-Pro won a medal in 63.6% of competitions, with gold in 36.4%. With the same Gemini-2.0-Flash, the medal rate rose from 25.8% for AIDE to 43.9% and the above-median rate from 39.4% to 63.6%. Under the same setting the medal rate was 37.9% without ensembling and 43.9% with the proposed ensemble, whose gold rate of 30.3% beat the 25.8% of a simple average ensemble.
• Limits: Kaggle competitions are public, so the LLM may have been trained on related discussions; as evidence against this the authors report that a Gemini-2.5-Pro judge rated all 75 pairs formed from 25 top discussions in 7 competitions and the final solutions as novel. Evaluation uses the 22-competition Lite split, the low-complexity part of the 75-competition MLE-bench, and the Claude-Sonnet-4 runs and the DS-Agent comparison each cover only 4 tasks. The AIDE results with o1-preview and GPT-4o are taken from the MLE-bench repository and use 16 and 36 seeds, against 3 for MLE-STAR.

Paper links

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