Research map / Multi-agent systems and games / LLM agents
MLE-STAR: Machine Learning Engineering Agent via Search and Targeted Refinement
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes MLE-STAR, an LLM-based ML engineering agent that retrieves state-of-the-art models via web search, refines one targeted code block at a time chosen by ablation, and ensembles candidate solutions with LLM-proposed ensemble strategies, evaluated on MLE-bench Lite.
Problem
Existing LLM-based MLE agents over-rely on the LLM's internal knowledge, selecting outdated or simple models, and explore by rewriting the entire code at once, so they pivot prematurely across pipeline components instead of deeply exploring one.
Approach
Use web search as a tool to retrieve task-specific SOTA models for the initial solution, then run nested loops: an ablation study picks the most impactful code block, an inner loop iteratively refines only that block with LLM-planned strategies using past results as feedback, and a separate ensemble-planning loop iteratively designs ways to merge multiple solutions instead of simple voting.
Key results
- On MLE-bench Lite, MLE-STAR with Gemini-2.5-Pro achieves 63.6% any-medal rate (36.4% gold), versus 25.8% for the best AIDE baseline with Gemini-2.0-Flash.
- With Gemini-2.0-Flash, MLE-STAR reaches 43.9% any-medal rate, improving AIDE's 25.8% and surpassing o1-preview AIDE's 36.6%.
- The proposed LLM-planned average/ensemble strategy raises any-medal rate from 37.9% (no ensemble) to 43.9% with Gemini-2.0-Flash.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use MLE-STAR instead of AIDE | Gemini-2.0-Flash · MLE-bench Lite | Any Medal (%): 25.8\stdv{5.4} -> 43.9\stdv{6.2} | Helped |
| use Gemini-2.5-Pro instead of Gemini-2.0-Flash | MLE-STAR · MLE-bench Lite | Any Medal (%): 43.9\stdv{6.2} -> 63.6\stdv{6.0} | Helped |
| use MLE-STAR instead of DS-Agent | wild-blueberry-yield · wild-blueberry-yield | MAE: 213 -> 166 | Helped |
| use MLE-STAR instead of DS-Agent | media-campaign-cost · media-campaign-cost | RMLSE: 0.2964 -> 0.2911 | Helped |
| use MLE-STAR instead of DS-Agent | spaceship-titanic · spaceship-titanic | Accuracy: 0.7982 -> 0.8091 | Helped |
| use MLE-STAR instead of DS-Agent | enzyme-substrate · enzyme-substrate | AUROC: 0.8727 -> 0.9101 | Helped |
| add the proposed ensemble strategy (ablation) | MLE-STAR · MLE-bench Lite | Any Medal (%): 37.9\stdv{6.0} -> 43.9\stdv{6.2} | Helped |
| use the proposed ensemble strategy instead of Best-of-N selection (ablation) | MLE-STAR · MLE-bench Lite | Any Medal (%): 42.4\stdv{6.1} -> 43.9\stdv{6.2} | Helped |
| use the proposed ensemble strategy instead of average ensemble (ablation) | MLE-STAR · MLE-bench Lite | Gold (%): 25.8\stdv{9.4} -> 30.3\stdv{5.7} | Helped |
| add the data leakage checker (ablation) | MLE-STAR · spaceship-titanic | Test accuracy: 0.8033 -> 0.7343 | Helped |
| add the data leakage checker (ablation) | MLE-STAR · spaceship-titanic | Validation accuracy: 0.8188 -> 0.8677 | Hurt |
| add the data usage checker (ablation) | MLE-STAR · nomad2018-predicting | RMSLE: 0.0591 -> 0.0559 | Helped |
Surprising findings
- LLMs often proposed logistic regression even for text classification tasks, reflecting bias toward familiar methods over up-to-date models.
- Generated solutions sometimes used test-set statistics for preprocessing or target-derived features, causing test accuracy to drop despite validation gains without the leakage checker.
Limitations
- Kaggle competitions are public, so LLMs may have seen related discussions and solutions, risking contamination.
- Evaluated mainly on 22 MLE-bench Lite competitions, with DS-Agent comparison limited to four tabular tasks.
Experimental setup
Datasets: MLE-bench Lite Kaggle competitions (tabular tasks from DS-Agent)
Benchmarks: MLE-bench Lite
Metrics: any medal rate gold medal rate valid submission rate above median rate
Baselines: AIDE MLAB OpenHands DS-Agent AutoGluon
Builds on: AIDE DS-Agent MLE-bench
Closest papers by meaning
- MLE-Dojo: Interactive Environments for Empowering LLM Agents in Machine Learning Engineering (NeurIPS 2025 · similarity 0.89)
- CoMind: Towards Community-Driven Agents for Machine Learning Engineering (ICLR 2026 · similarity 0.89)
- MLE-bench: Evaluating Machine Learning Agents on Machine Learning Engineering (ICLR 2025 · similarity 0.89)
- MLE-Smith: Scaling MLE Tasks with Automated Multi-agent Pipeline (ICLR 2026 · similarity 0.88)
- ML-Agent: Reinforcing LLM Agents for Autonomous Machine Learning Engineering (ICML 2026 · similarity 0.88)
- AI Research Agents for Machine Learning: Search, Exploration, and Generalization in MLE-bench (NeurIPS 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.