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

TL;DR

AutoML-Agent tackles the challenge of end-to-end AutoML by introducing a multi-agent LLM framework that coordinates data retrieval, preprocessing, model search, hyperparameter optimization, evaluation, and deployment. It leverages retrieval-augmented planning to generate diverse, up-to-date plans, decomposes plans into task-specific sub-tasks for parallel execution, and employs multi-stage verification to reduce code-generation errors and ensure deployment readiness. Across seven downstream tasks and fourteen datasets, AutoML-Agent achieves higher success rates and competitive downstream performance compared with baselines, while delivering notable efficiency gains over training-based search methods. The framework demonstrates robust performance under constraint-aware and constraint-free settings and highlights practical considerations such as cost, latency, and reliability, offering a scalable path toward democratizing AutoML for non-experts.

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.

Figures (7)

• Figure 1: AutoML-Agent receives user's instructions and delivers optimized deployable models.
• Figure 2: Overview of our AutoML-Agent framework. (1) Initialization stage aims to receive a valid user instruction using request verification. (2) Planning stage focuses on extracting ML related information by parsing the user instruction into a standardized form, and uses it to devise plans accordingly. (3) Execution stage executes each action given by the devised plans. Finally, based on the best execution results, AutoML-Agent outputs codes containing deployable model to the user.
• Figure 3: An example of prompt parsing process of an instruction-response pair $\{(I_i, R_i)\}$.
• Figure 4: Performance comparison across all datasets using the SR, NPS, and CS metrics under (a) constraint-free and (b) constraint-aware settings. Higher scores indicate better results.
• Figure 5: Results of (a) ablation study and (b) hyperparameter study in the CS metric.