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

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 LLMs (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.

• The paper introduces a multi-agent framework that automates each stage of the ML pipeline, enhancing efficiency and accessibility.
• It leverages a retrieval-augmented planning strategy to decompose tasks into specialized sub-tasks, ensuring precise and efficient execution.
• Experimental validation demonstrates a 100% success rate in deployable model generation under constraint-free scenarios and superior performance in constraint-aware settings.

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

The paper "AutoML-Agent: A Multi-Agent LLM Framework for Full-Pipeline AutoML" (2410.02958) introduces a novel multi-agent framework designed to automate the entire process of machine learning model development, from data retrieval to model deployment, using LLMs. This framework, termed AutoML-Agent, aims to enhance the accessibility and efficiency of automated machine learning (AutoML) by allowing users to interact with a natural language interface.

Framework Overview

AutoML-Agent is structured around several specialized agents that collaborate to fulfill user instructions. The framework consists of the following key components:

• Agent Manager ($\mathcal{A}_{mgr}$): Acts as the core interface with users, devises plans, assigns tasks to agents, and verifies execution results.
• Prompt Agent ($\mathcal{A}_{p}$): Parses user instructions into a standardized format to streamline communication between agents.
• Data Agent ($\mathcal{A}_{d}$): Handles data-related tasks, including retrieval, preprocessing, and analysis.
• Model Agent ($\mathcal{A}_{m}$): Focuses on model search, hyperparameter optimization, and candidate ranking.
• Operation Agent ($\mathcal{A}_{o}$): Implements the final solution by producing deployment-ready code.

  Figure 1: 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 models to the user.

Retrieval-Augmented Planning and Execution

To tackle the complexities of full-pipeline AutoML, AutoML-Agent introduces a retrieval-augmented planning (RAP) strategy. This strategy generates multiple plans based on past knowledge and current research, enabling parallel exploration of different solutions. Each plan is decomposed into sub-tasks tailored to the roles of the agents, allowing for efficient and focused execution without the need for further model training.

The multi-stage verification system ensures the reliability of the model by checking the validity of user requests, execution outcomes, and implementation results. This iterative feedback mechanism helps refine the generated solutions and adapt them to the user's requirements.

Figure 2: 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.

Experimental Validation

The AutoML-Agent framework was evaluated across a variety of datasets and machine learning tasks, demonstrating superior performance compared to existing methods such as GPT-3.5, GPT-4, and DS-Agent. In particular, AutoML-Agent achieved a 100% success rate in generating deployable models under constraint-free scenarios and consistently outperformed baselines in constraint-aware settings.

Figure 3: Results of (a) ablation and (b) hyperparameter studies in the CS metric.

Implications and Future Work

AutoML-Agent represents a significant advancement in making the power of AutoML accessible beyond expert practitioners, enabling efficient development of AI models even for users with limited technical knowledge. Its framework demonstrates the potential to further democratize AI development by simplifying complex processes.

Future research may focus on extending the framework's capabilities to support more diverse types of machine learning tasks, such as reinforcement learning and recommendation systems, as well as improving its robustness to operate independently of specific LLM backbones. Addressing these areas will expand the applicability and effectiveness of AutoML-Agent in complex real-world environments.

Conclusion

AutoML-Agent offers a comprehensive and efficient solution for automating the full pipeline of machine learning development, leveraging the capabilities of LLMs through a multi-agent framework. By addressing planning complexity and implementation accuracy, it sets a new standard for accessible AI-driven solutions. As the framework evolves, it is poised to further simplify the creation of sophisticated machine learning models, enhancing the potential for widespread AI innovation.