Automated machine learning (AutoML) accelerates AI development by automating tasks in the development pipeline, such as optimal model search and hyperparameter tuning. Existing AutoML systems often require technical expertise to set up complex tools, which is in general time-consuming and requires a large amount of human effort. Therefore, recent works have started exploiting large language models (LLM) to lessen such burden and increase the usability of AutoML frameworks via a natural language interface, allowing non-expert users to build their data-driven solutions. These methods, however, are usually designed only for a particular process in the AI development pipeline and do not efficiently use the inherent capacity of the LLMs. This paper proposes AutoML-Agent, a novel multi-agent framework tailored for full-pipeline AutoML, i.e., from data retrieval to model deployment. AutoML-Agent takes user's task descriptions, facilitates collaboration between specialized LLM agents, and delivers deployment-ready models. Unlike existing work, instead of devising a single plan, we introduce a retrieval-augmented planning strategy to enhance exploration to search for more optimal plans. We also decompose each plan into sub-tasks (e.g., data preprocessing and neural network design) each of which is solved by a specialized agent we build via prompting executing in parallel, making the search process more efficient. Moreover, we propose a multi-stage verification to verify executed results and guide the code generation LLM in implementing successful solutions. Extensive experiments on seven downstream tasks using fourteen datasets show that AutoML-Agent achieves a higher success rate in automating the full AutoML process, yielding systems with good performance throughout the diverse domains.
翻译:自动化机器学习(AutoML)通过自动化开发流程中的任务(如最优模型搜索与超参数调优)来加速人工智能开发。现有的AutoML系统通常需要技术专业知识来配置复杂工具,这往往耗时且需大量人工投入。因此,近期研究开始利用大语言模型(LLM)来减轻此类负担,并通过自然语言界面提升AutoML框架的易用性,使非专业用户能够构建其数据驱动的解决方案。然而,这些方法通常仅针对AI开发流程中的特定环节设计,未能有效利用LLM的固有能力。本文提出AutoML-Agent,一种专为全流程AutoML(即从数据获取到模型部署)设计的新型多智能体框架。AutoML-Agent接收用户的任务描述,促进专用LLM智能体间的协作,并交付可直接部署的模型。与现有工作不同,我们未采用单一规划策略,而是引入检索增强的规划机制以增强探索性,从而搜索更优的规划方案。同时,我们将每个规划分解为子任务(如数据预处理与神经网络设计),每个子任务由我们通过提示构建的专用智能体并行执行,从而提升搜索效率。此外,我们提出多阶段验证机制,用于校验执行结果并指导代码生成LLM实现成功解决方案。在七个下游任务、十四个数据集上的大量实验表明,AutoML-Agent在全流程AutoML自动化中实现了更高的成功率,在不同领域均能生成性能良好的系统。