MLE-STAR: Machine Learning Engineering Agent via Search and Targeted Refinement
This addresses the challenge of automating ML model implementation for practitioners, offering a novel agent approach with strong performance gains.
The paper tackles the problem of building machine learning engineering agents that can automatically implement ML models via code generation, proposing MLE-STAR which uses search and targeted refinement to achieve medals in 64% of Kaggle competitions on MLE-bench Lite, significantly outperforming alternatives.
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 64% of the Kaggle competitions on the MLE-bench Lite, significantly outperforming the best alternative.