AutoML-Agent: A Multi-Agent LLM Framework for Full-Pipeline AutoML

2024-10-03

arXiv.org

AI Summary

Existing AutoML systems rely heavily on expert configuration, resulting in low usability; although LLM-assisted approaches have emerged, they typically target isolated pipeline stages and fail to harness LLMs’ end-to-end reasoning capabilities. This paper introduces the first multi-agent large language model framework for full-stack AutoML—spanning data acquisition, preprocessing, model search, hyperparameter tuning, and deployment. Our approach innovatively integrates retrieval-augmented multi-stage planning, task-parallel decomposition, multi-stage program verification, and domain-adaptive prompt engineering to enable natural-language-driven fully automated machine learning. Evaluated across 14 diverse datasets and 7 downstream task categories, our framework achieves significantly higher end-to-end automation success rates. Generated models maintain high cross-domain performance, while human intervention is reduced by over 70%.

Technology Category

• Machine Learning: Auto ML and Hyperparameter Tuning
• Multiagent Systems: Multiagent Learning
• Natural Language Processing: (Large) Language Models

Application Category

• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

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.

Problem

• Automates full-pipeline AutoML for non-experts
• Enhances exploration with retrieval-augmented planning
• Improves efficiency via multi-agent parallel task execution

Innovation

• Multi-agent LLM framework for full-pipeline AutoML
• Retrieval-augmented planning strategy for optimal plans
• Multi-stage verification for reliable code generation