用搜索+精准优化，让AI自动搞定机器学习工程任务
MLE-STAR: Machine Learning Engineering Agent via Search and Targeted Refinement
- 先搜网络找有效模型，再分组件逐步优化代码
- 在Kaggle竞赛中达成64%胜率，远超现有方法
- 适合需要高效构建机器学习系统的开发者
基于大语言模型（LLM）的机器学习工程（MLE）代理可通过代码生成自动实现机器学习模型。然而，现有方法通常过度依赖模型自身知识，采用粗粒度探索策略，一次性修改整个代码结构，限制了其在特定任务中选择有效模型以及深入探索特征工程等组件的能力。为此，我们提出MLE-STAR，一种新型的MLE代理构建方法。MLE-STAR首先利用搜索引擎从网络检索有效模型作为初始解，随后通过针对特定机器学习组件的策略迭代优化该解。这一探索过程由消融分析引导，以评估各代码块的影响。此外，我们引入一种新的集成方法，使用由MLE-STAR推荐的有效策略。实验结果表明，MLE-STAR在MLE-bench Lite的Kaggle竞赛中取得了64%的胜率，显著优于最佳对比方法。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。