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

TL;DR

MLE-STAR is a novel machine learning engineering agent that combines search-based retrieval and targeted refinement to improve model selection and component exploration, significantly enhancing performance on Kaggle competitions.

Contribution

It introduces a search-guided, iterative refinement approach for LLM-based MLE agents, enabling effective task-specific model selection and deep component exploration.

Findings

Achieves medals in 64% of Kaggle competitions on MLE-bench Lite.

Outperforms the best alternative methods significantly.

Introduces a novel ensembling strategy guided by MLE-STAR.

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 64% of the Kaggle competitions on the MLE-bench Lite, significantly outperforming the best alternative.