[论文解读] AutoML-Agent: A Multi-Agent LLM Framework for Full-Pipeline AutoML

Patara Trirat, Wonyong Jeong+1 more|arXiv (Cornell University)|Oct 3, 2024

Machine Learning and Data Classification被引用 4

一句话总结

该论文提出 AutoML-Agent，一种多智能体 LLM 框架，通过检索增强规划与多阶段验证，实现从数据检索、预处理到模型部署的完整机器学习流水线自动化。其在 14 个数据集和 7 项任务上实现了 96.2% 的平均成功率，优于基线方法，在成功率与效率方面表现更优，且相比 MCTS 方法将搜索时间减少了高达 90%。

ABSTRACT

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.

研究动机与目标

• 解决现有基于 LLM 的 AutoML 系统仅聚焦于孤立流水线阶段（如超参数调优或模型选择）而非端到端自动化的局限性。
• 通过引入检索增强策略，降低完整流水线 AutoML 的规划复杂度，以探索多样化、知识驱动的计划。
• 通过角色特定的智能体分解与多阶段验证，提升代码生成准确性，减少自主 ML 流水线构建中的幻觉现象。
• 通过在单一框架中集成数据、模型与评估阶段，实现高效、可扩展且可部署的 AutoML，服务于非专家用户。
• 证明多智能体协作结合检索增强规划与验证，能显著提升多样化 ML 任务中的成功率与计算效率。

提出的方法

• AutoML-Agent 采用多智能体架构，其中专用 LLM 智能体（如数据、模型、验证智能体）并行运行，以分解并解决子任务，如数据预处理、神经架构搜索与超参数调优。
• 引入检索增强规划策略，从数据集或先前流水线等来源检索相关知识，生成多个候选计划，提升对最优解的探索能力。
• 每个计划被分解为原子子任务，通过角色特定提示分配给专用智能体，实现并行执行，从而缩短整体流水线时间。
• 通过多阶段验证流程对中间结果进行验证，并指导代码生成智能体，借助迭代反馈减少幻觉，确保正确性。
• 框架采用综合评分系统，结合成功率、归一化性能与综合评分，评估端到端流水线质量。
• 通过早期过滤无效计划并优化成本与时间，支持约束感知执行，相比无约束设置，总运行时间减少 4.8%，成本降低 18.5%。

实验结果

研究问题

• RQ1 多智能体 LLM 框架能否在无需专家级配置的情况下，有效自动化从数据检索到部署的完整 AutoML 流水线？
• RQ2 与单计划或启发式方法相比，检索增强规划是否能提升探索流水线的成功率与多样性？
• RQ3 角色特定的智能体分解与并行执行是否能显著降低完整流水线 AutoML 的时间和计算成本？
• RQ4 多阶段验证在多大程度上可减少代码幻觉并提升生成 ML 流水线的可靠性？
• RQ5 在多样化的真实世界数据集与任务中，AutoML-Agent 与最先进基线方法相比，在成功率、性能与效率方面表现如何？

主要发现

• AutoML-Agent 在 14 个数据集与 7 项下游任务上实现了 96.2% 的平均成功率，显著优于基线方法（如 DS-Agent 的 66.1% 与约束感知 GPT-3.5 的 7.7%）。
• 与基于 MCTS 的 SELA 相比，AutoML-Agent 将平均搜索时间减少了 90%（从 2037.18 秒降至 249.91 秒），同时保持或提升了归一化性能得分。
• 每项任务生成五个计划时，AutoML-Agent 在图像与文本分类任务中实现了 100% 的成功率，在节点分类任务中获得 90.6% 的综合得分，展现出跨任务的鲁棒性。
• 多阶段验证策略提升了代码可靠性，减少了幻觉现象并确保了依赖关系的正确处理，代码生成与部署的成功率达到 96.9%。
• 在排行榜对比中，AutoML-Agent 在 8 项竞赛中的 7 项优于 Agent K，中位数排行榜分位数为 88.5，高于 Agent K 的 86.5，且在 5/8 的任务中性能更低。
• 约束感知设置使总成本降低 18.5%（0.32 美元 vs. 0.27 美元），时间减少 4.8%（512.35 秒 vs. 538.25 秒），在不牺牲性能的前提下实现了效率提升。