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

📄 arXiv: 2506.15692v3 📥 PDF

作者: Jaehyun Nam, Jinsung Yoon, Jiefeng Chen, Jinwoo Shin, Sercan Ö. Arık, Tomas Pfister

分类: cs.LG

发布日期: 2025-05-27 (更新: 2025-08-28)

💡 一句话要点

提出MLE-STAR以解决机器学习工程代理的深度探索问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 机器学习工程 大型语言模型 模型优化 特征工程 自动化开发 深度探索 消融研究

📋 核心要点

  1. 现有的机器学习工程代理方法过于依赖大型语言模型的知识，且探索策略粗糙，限制了任务特定模型的选择和组件的深度探索。
  2. MLE-STAR通过搜索引擎获取外部知识，形成初步解决方案，并针对特定机器学习组件进行迭代优化，提升模型性能。
  3. 实验结果显示，MLE-STAR在64%的Kaggle竞赛中获得奖牌，显著超越了其他方法，展示了其有效性和优势。

📝 摘要（中文）

基于大型语言模型（LLMs）的机器学习工程代理（MLE）能够通过代码生成自动实现机器学习模型。然而，现有方法往往过于依赖LLM的内在知识，并采用粗略的探索策略，导致无法有效选择特定任务模型并深入探索特定组件。为此，本文提出了MLE-STAR，一种新颖的MLE代理构建方法。MLE-STAR首先利用搜索引擎从网络中检索有效模型，形成初步解决方案，然后通过针对特定机器学习组件的多种策略进行迭代优化。该探索过程受到消融研究的指导，分析各个代码块的影响。此外，我们还引入了一种新颖的集成方法。实验结果表明，MLE-STAR在MLE-bench Lite的64% Kaggle竞赛中获得奖牌，显著优于最佳替代方案。

🔬 方法详解

问题定义：本文旨在解决现有机器学习工程代理在模型选择和组件探索中的不足，尤其是对特定任务模型的选择和特征工程的深入实验。

核心思路：MLE-STAR的核心思路是结合外部知识和深度探索策略，通过搜索引擎获取有效模型，并针对特定组件进行细致的迭代优化。这样的设计能够有效克服现有方法的局限性。

技术框架：MLE-STAR的整体架构包括两个主要阶段：首先，通过搜索引擎检索有效模型形成初步解决方案；其次，利用消融研究指导的策略对模型进行迭代优化，探索不同的机器学习组件。

关键创新：MLE-STAR的关键创新在于其结合了外部知识的检索与针对特定组件的深度探索策略，这与现有方法的粗略修改策略形成了鲜明对比。

关键设计：在设计上，MLE-STAR采用了消融研究分析各个代码块的影响，并引入了一种新颖的集成方法，以提升模型的整体性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

MLE-STAR在MLE-bench Lite的64% Kaggle竞赛中获得奖牌，显示出其卓越的性能，显著优于最佳替代方案，证明了其在机器学习工程中的有效性和创新性。

🎯 应用场景

MLE-STAR的研究成果在机器学习工程领域具有广泛的应用潜力，尤其是在自动化模型开发、特征工程和模型优化等方面。其方法可以帮助数据科学家更高效地构建和优化机器学习模型，提升工作效率和模型性能，未来可能在工业界和学术界产生深远影响。

📄 摘要（原文）

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.