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

Patara Trirat, Wonyong Jeong, Sung Ju Hwang

TL;DR: This paper introduces AutoML-Agent, a novel multi-agent framework for building automated machine learning pipelines from natural language instructions, with informed planning and verification strategies.

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.

Lay Summary:

Building machine learning systems can be difficult, especially for people without a background in AI. Many existing tools still require coding and technical skills, making it hard for non-experts to use them. This paper tackles that problem by creating a tool called AutoML-Agent, which helps users build AI systems just by describing their task in plain language.

AutoML-Agent works like a team of smart assistants. Each assistant (or "agent") has a special role: some focus on understanding the data, others on choosing the right model, tuning it, and writing the code. The system takes in a user's request, such as "Build a model to detect spam messages," and automatically plans, builds, tests, and delivers a ready-to-use AI model.

To make this process reliable and efficient, AutoML-Agent combines smart planning strategies with checks at each step to avoid mistakes. The result is a powerful tool that builds high-performing AI systems with minimal human effort, opening the door for more people to create AI solutions in their own fields.

Link To Code: https://github.com/deepauto-ai/automl-agent

Primary Area: Applications

Keywords: Automated Machine Learning, Multi-Agent Framework, Large Language Models

Submission Number: 8647