16
0

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

Main:11 Pages
24 Figures
Bibliography:3 Pages
9 Tables
Appendix:27 Pages
Abstract

Agents based on large language models (LLMs) for machine learning engineering (MLE) can automatically implement ML models via code generation. However, existing approaches to build such agents often rely heavily on inherent LLM knowledge and employ coarse exploration strategies that modify the entire code structure at once. This limits their ability to select effective task-specific models and perform deep exploration within specific components, such as experimenting extensively with feature engineering options. To overcome these, we propose MLE-STAR, a novel approach to build MLE agents. MLE-STAR first leverages external knowledge by using a search engine to retrieve effective models from the web, forming an initial solution, then iteratively refines it by exploring various strategies targeting specific ML components. This exploration is guided by ablation studies analyzing the impact of individual code blocks. Furthermore, we introduce a novel ensembling method using an effective strategy suggested by MLE-STAR. Our experimental results show that MLE-STAR achieves medals in 44% of the Kaggle competitions on the MLE-bench, significantly outperforming the best alternative.

View on arXiv
@article{nam2025_2506.15692,
  title={ MLE-STAR: Machine Learning Engineering Agent via Search and Targeted Refinement },
  author={ Jaehyun Nam and Jinsung Yoon and Jiefeng Chen and Jinwoo Shin and Sercan Ö. Arık and Tomas Pfister },
  journal={arXiv preprint arXiv:2506.15692},
  year={ 2025 }
}
Comments on this paper