MLE-STAR: A state-of-the-art machine learning engineering agent

The rise of machine learning has created complex engineering challenges, prompting research into using large language models (LLMs) as agents to automate these tasks. These LLM agents conceptualize ML problems as code optimization challenges, generating executable code. However, current agents often rely on familiar methods and struggle with deep exploration of specific code components. To address this, MLE-STAR was developed, a novel ML engineering agent that integrates web search and targeted code block refinement. MLE-STAR begins by searching the web for relevant models and then iteratively refines specific code blocks, identified through ablation studies, based on LLM-generated plans. The agent also employs a new strategy for ensembling multiple candidate solutions. Furthermore, MLE-STAR includes modules for debugging, checking for data leakage, and ensuring proper data usage. Evaluations on MLE-Bench-Lite demonstrated MLE-STAR's significant outperformance compared to existing alternatives, winning medals in 63% of Kaggle competitions. This success is attributed to its use of more recent models, focused refinement, and robust checking mechanisms. MLE-STAR's automated approach aims to lower the barrier to ML adoption and adapt to advancements in the field. An open-source codebase for MLE-STAR is now available.