A state-of-the-art machine learning engineering agent

Despite their promising initial strides, current MLE agents face several limitations that curtail their efficacy. First, their heavy reliance on pre-existing LLM knowledge often leads to a bias towards familiar and frequently used methods (e.g., the scikit-learn library for tabular data), overlooking potentially superior task-specific approaches. Furthermore, these agents typically employ an exploration strategy that modifies the entire code structure simultaneously in each iteration. This frequently causes agents to prematurely shift focus to other stages (e.g., model selection or hyperparameter tuning) because they lack the capacity for deep, iterative exploration within specific pipeline components, such as exhaustively experimenting with different feature engineering options.

In our recent paper, we introduce MLE-STAR, a novel ML engineering agent that integrates web search and targeted code block refinement. Unlike alternatives, MLE-STAR tackles ML challenges by first searching the web for proper models to get a solid foundation. It then carefully improves this foundation by testing which parts of the code are most important. MLE-STAR also utilizes a new method to blend several models together for even better results. This approach is very successful — it won medals in 63% of the Kaggle competitions in MLE-Bench-Lite, significantly outperforming the alternatives.