AutoML-Agent: A Multi-Agent LLM Framework for Full-Pipeline AutoML (2410.02958v1)

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.

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

The paper introduces AutoML-Agent, an advanced framework designed to automate the entire machine learning pipeline, from data retrieval to model deployment. This approach aims to democratize AI development by enabling non-experts to create data-driven solutions without requiring extensive technical knowledge. AutoML-Agent achieves this by leveraging a novel multi-agent system, where each agent specializes in distinct tasks, facilitating efficient and accurate pipeline construction through collaboration.

Framework Description

AutoML-Agent utilizes a set of LLM-based agents, each with specific roles, to automate various stages of the machine learning pipeline:

• Agent Manager: Orchestrates the overall process, including interaction with users, plan generation, task distribution, and verification of results.
• Prompt Agent: Converts user instructions into standardized JSON objects, ensuring consistency across the system.
• Data Agent: Handles data retrieval, preprocessing, and analysis, providing the necessary data insights to the Model Agent.
• Model Agent: Focuses on model selection, hyperparameter optimization, and performance profiling, generating candidate models for evaluation.
• Operation Agent: Implements the final solution by writing and deploying executable code based on selected models.

Core Methodologies

The framework introduces several innovative strategies:

• Retrieval-Augmented Planning: Generates multiple potential plans by integrating past knowledge and real-time retrieval from APIs, enhancing exploration for optimal solutions.
• Plan Decomposition: Breaks down complex plans into manageable sub-tasks, allowing agents to focus on specific responsibilities aligned with their expertise.
• Prompting-Based Execution: Utilizes LLMs' in-context learning to simulate execution without actual code runs, reducing computational overhead.
• Multi-Stage Verification: Incorporates request, execution, and implementation verifications to ensure accuracy and adherence to user requirements.

Empirical Results

The framework underwent extensive validation with experiments across multiple datasets and domains, including image, text, tabular, graph, and time series modalities. AutoML-Agent consistently outperformed existing AutoML systems and general-purpose LLMs like GPT-3.5 and GPT-4, achieving notably higher success rates and downstream performance.

• Success Rate (SR): AutoML-Agent demonstrated robust performance, particularly in the constraint-aware setting, with an average SR of 87.1%. This indicates its effectiveness in generating models that meet specified constraints.
• Normalized Performance Score (NPS): AutoML-Agent achieved superior performance on various tasks, suggesting it effectively generates high-quality models tailored to specific datasets and user requirements.
• Comprehensive Score (CS): Combining success and performance metrics, AutoML-Agent showcased outstanding overall efficiency.

Implications and Future Directions

The introduction of AutoML-Agent represents a significant advancement in making AI accessible to non-experts by automating complex machine learning tasks within a single, integrated framework. The practical applications are vast, potentially transforming industries by reducing the time and expertise required to develop sophisticated AI solutions.

Theoretically, the framework opens avenues for further exploration in retrieval-augmented techniques and multi-agent cooperation in LLMs. Future research may explore enhancements in code generation fidelity and adaptation to novel machine learning paradigms, including reinforcement learning and recommendation systems.

In conclusion, AutoML-Agent sets a new benchmark in AutoML by effectively using LLMs in a structured, multi-agent framework to deliver comprehensive, end-to-end automation processes capable of serving a diverse range of AI applications.