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

Why does RAP propose multiple plans for a task? Especially, this is not inference time compute (thoughts instead of execution) but indeed executes all the plans and compares the results in the last, which results in a significant amount of costs.

R3. As clarified in Section 3.5 of the paper, the multiple-plan approach in RAP is designed not to impose significant computational overhead but to balance exploration with efficiency. We achieve this balance through the integration of Pseudo Data Analysis and Training-Free Model Search and Hyperparameter Optimization, both of which rely solely on LLM inference. This eliminates the need for executing all candidate plans fully during the exploration phase.

To further minimize runtime costs, only plans that pass our rigorous multi-stage verification process, as described in Section 3.6, are selected for actual runtime execution. This ensures that the computational resources are spent only on the most promising solutions. As shown in Table R3 (also Section 4.3), the total search time to derive a code of a deployable model is approximately 10 minutes, which constitutes a negligible fraction of the overall model training and optimization time. Thus, the proposed framework ensures a computationally efficient implementation while preserving the robustness and adaptability of the AutoML process.

Table R3. Time costs averaged across different tasks and datasets for a single run under the constraint-free and constraint-aware settings.

We hope this clarifies our approach and addresses the reviewer’s concerns regarding computational costs.

How do judge the difference between the LLM-agent framework for Auto-ML and the general LLM-agent? It seems to me that the general LLM-agent framework can cater to AutoML tasks since they are just a small subset of real complex tasks.

R4. We appreciate the reviewer’s observation regarding the applicability of general-purpose LLM-agent frameworks for handling AutoML tasks. While it is true that generic LLM-agents are versatile and can theoretically execute certain AutoML tasks, they lack critical domain-specific capabilities essential for achieving robust and reliable performance across the structured and interdependent processes of a full AutoML pipeline. Generic LLM frameworks are typically designed to support broad problem domains and rely heavily on user-defined instructions or augmentations, often leading to inefficiencies, suboptimal performance, and increased failure rates in complex tasks like end-to-end AutoML. In contrast, AutoML-Agent is purpose-built with explicit mechanisms, such as retrieval-augmented planning, specialized agents for sub-tasks, and multi-stage verification, all tailored to address the unique challenges of AutoML.

Furthermore, as highlighted in Table R4, even when a general-purpose LLM leverages the same backbone model, they consistently underperform compared to specialized methods like DS-Agent or AutoML-Agent in both success rates and downstream task metrics. AutoML-Agent not only outperforms these frameworks but also uniquely integrates the entire pipeline into a seamless, modular workflow, significantly reducing errors and enhancing usability for diverse downstream tasks. This distinction underscores the importance of designing specialized frameworks like AutoML-Agent, which prioritize task-specific optimizations over generality, ensuring both efficiency and reliability in automating full AutoML pipelines.

Table R4. Average performance scores (from Tables 5,6, and 7) between different LLM frameworks.

Additionally, I question the motivation and necessity behind devising this AutoML-agent framework. Do non-expert users really exist and need this level of assistance? Although the limitation section acknowledges the issues of hallucination and unreliability within LLMs, I want to emphasize that systems or frameworks with bugs can be prohibitive for non-expert users, making it extremely difficult for them to debug and resolve problems. This is not necessarily a flaw of this paper, but it seems premature to propose frameworks that may not be practical at this time.

R2. We appreciate the reviewer’s thoughtful feedback and the opportunity to clarify the motivation and necessity of the AutoML-Agent framework, as well as its practicality for non-expert users.

• Motivation and Necessity: In the era of generative AI, many companies are adopting AI technologies. However, these companies often lack AI experts, leading to software engineers—who are non-experts in AI—attempting to implement such solutions and facing significant challenges. Researchers from various domains outside of AI (e.g., economics, chemistry, healthcare) are increasingly seeking to apply AI models in their work but struggle due to their lack of specialized AI expertise. Besides, the existence of non-expert users in the ML landscape is well-documented in relevanted studies [a--d]. For example, an early AutoML survey paper [a] highlights the growing demand for ML tools among stakeholders across various domains. It emphasizes that AutoML tools aim to make ML accessible to non-experts, improve efficiency, and accelerate research, addressing the pressing need for user-friendly AI solutions. Similarly, as thoroughly illustrated in Section 3 of a recent study [c], the use of LLMs as interfaces for AutoML and as components of AutoML systems offers significant opportunities for non-ML experts seeking to apply off-the-shelf data-driven solutions to their problems.
• Example Scenarios: To further illustrate the necessity and practicality of AutoML-Agent, consider an academic researcher aiming to evaluate several ML models for a novel dataset within a constrained timeline. Traditionally, this involves significant manual effort in setting up pipelines, selecting models, and debugging code—steps that can be both error-prone and time-consuming. Similarly, in industry settings, ML engineers often need to rapidly prototype models for specific business requirements, such as creating a lightweight and efficient spam detection system for mobile applications. AutoML-Agent enables such users to focus on high-level problem formulation by generating deployment-ready models that adhere to specified constraints, like latency or accuracy, directly from natural language task descriptions. This capability reduces setup time and errors, enabling users to focus on innovation rather than implementation logistics.
• Practicality and Error Mitigation: While we acknowledge the inherent limitations of LLMs, such as potential hallucination, these challenges are mitigated in our framework through:
  • Retrieval-Augmented Planning: Ensures robustness by grounding plans in external knowledge.
  • Multi-Stage Verification: Detects and corrects errors at each stage, ensuring reliable outputs.
  • Structured Prompt Parsing: Aligns the framework's capabilities with precise user requirements, reducing ambiguities.

We hope this clarification adequately addresses your concerns and demonstrates the value and timeliness of our work.

References

• [a] Karmaker, Shubhra Kanti, et al. "Automl to date and beyond: Challenges and opportunities." ACM Computing Surveys (CSUR) 54.8 (2021): 1-36.
• [b] Sun, Yuan, et al. "AutoML in The Wild: Obstacles, Workarounds, and Expectations." Proceedings of the 2023 CHI Conference on Human Factors in Computing Systems. 2023.
• [c] Tornede, Alexander, et al. "AutoML in the Age of Large Language Models: Current Challenges, Future Opportunities and Risks." TMLR. 2024.
• [d] Davenport, T., and R. Bean. "Five key trends in AI and data science for 2024." MIT Sloan Management Review (2024).

The core contribution of this paper is the multi-agent framework with RAP for processing AutoML tasks in an end-to-end manner. While the ambition of AutoML is significant and the framework is claimed to be powerful, the benchmarks and tasks presented are very simple. I believe these tasks would not require such a complex framework and should be manageable for most ML practitioners.

R1. We sincerely appreciate the reviewer’s acknowledgment of the ambition and core contributions of our work. We respectfully address the concerns regarding the simplicity of the benchmarks and tasks presented.

• Benchmark Selection: The benchmarks were chosen to enable fair comparisons with existing approaches, as they have been adopted in recent studies [a, b]. While these tasks may appear straightforward to skilled ML practitioners, they pose significant challenges for automated frameworks, especially in achieving seamless integration of data handling, model selection, hyperparameter tuning, and deployment. Our approach aims to automate the entire AutoML pipeline end-to-end, significantly reducing the reliance on manual effort and extensive domain expertise. Notably, we demonstrated that a current state-of-the-art framework, DS-Agent [a], still struggles with even these benchmarks, underscoring the value of our contribution and inherent difficulties of full-pipeline code generation. That is, the simplicity of the benchmarks does not negate the complexity of automating the entire process; instead, these tasks serve to validate the framework's capacity for generalization across diverse data modalities and application domains.
• Comparison with ML Practitioners: We would like to emphasize that AutoML-Agent significantly enhances usability for non-expert users, addressing a key gap in current AutoML tools. For users with limited expertise, the framework lowers the barrier to deploying robust models, which is crucial in democratizing access to AI technologies. For ML-expert users, our framework provides value by offering improvements in efficiency and performance, even exceeding state-of-the-art models built by the human experts, such as TabPFN [c] and iTransformer [d].

References

• [a] Guo, Siyuan, et al. "DS-Agent: Automated Data Science by Empowering Large Language Models with Case-Based Reasoning." ICML 2024.
• [b] Liu, Siyi, Chen Gao, and Yong Li. "Large Language Model Agent for Hyper-Parameter Optimization." arXiv preprint arXiv:2402.01881 (2024).
• [c] Hollmann, Noah, et al. "TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second." ICLR 2023.
• [d] Liu, Yong, et al. "iTransformer: Inverted Transformers Are Effective for Time Series Forecasting." ICLR 2024.

The AutoML-Agent framework proposed in this paper relies on the collaborative work of multiple agents, which may lead to significant computational resource consumption. In environments with limited computing power, implementing this framework using open-source large language models (LLMs) might face challenges.

R1. We thank the reviewer for highlighting potential challenges in resource-constrained environments. Below, we address these concerns.

• Resource Consumption: To reduce computation costs in multi-agent systems, it is possible to sequentially process tasks by loading only a single agent at a time or to reduce the number of multi-agents, thereby adjusting parallelism. The AutoML-Agent framework is designed with modularity and flexibility at its core, enabling adaptation to diverse resource settings. Specifically, while the framework supports parallel multi-agent execution to optimize task completion time, it also allows for sequential task execution by aligning the agent roles with the system's state. This configuration significantly reduces memory and processing demands, making the framework compatible with resource-constrained environments.
• Open LLMs: The backbone model in our framework does not necessarily have to be an open-source LLM, which typically incurs significant infrastructure costs. Instead, commercial LLMs offering more cost-effective APIs can be used. As presented below, due to the relatively low cost of less than $1 (please see Section 4.3 or Table R1 below) for generating code of a deployable model using GPT-4o, we believe this low operational cost obviates the immediate need for open-source LLMs. Moreover, open-source LLMs typically demand fine-tuning or additional engineering to match GPT-4o's performance, further complicating their deployment in constrained environments. Thus, we believe that our choice balances efficiency, cost-effectiveness, and practicality for a wide range of users.

Table R1. Time and monetary costs averaged across different tasks and datasets for a single run under the constraint-free and constraint-aware settings.

The method proposed in this paper has only been experimentally validated on current SOTA models, Mixtral-8x7B and GPT-4o. Given that the performance improvements observed in downstream tasks may primarily be attributed to the inherent advantages of the large language models (LLMs) themselves rather than the proposed method, I am uncertain about the generalizability of the experimental conclusions to smaller-scale or less powerful models. To further demonstrate the robustness of the proposed method, it is suggested that the authors supplement the paper with experimental results on such models.

R2. We thank the reviewer for their thoughtful comments on the generalizability of our proposed method. Below, we provide clarifications to address these concerns.

• Attribution of Performance Gains: The observed improvements are primarily due to the methodological contributions of AutoML-Agent, as evidenced by comprehensive ablation studies (Section 4.4 and Table R2.1 below). Key components, such as plan decomposition and multi-stage verification, consistently enhance downstream task performance. These results emphasize that the gains are not solely due to the capabilities of LLMs.

Table R2.1. Downstream task performance (i.e., normalized performance score) from ablation study on different variations.

• Challenges with Smaller Models: While we acknowledge the importance of validating our framework on smaller-scale models, these models exhibit systemic limitations that hinder their utility for complex tasks requiring code execution and extensive planning, which we briefly discussed in Section 5 of the revised paper. During the early stage of development, we found that small language models, including LLaMA-2-7b, Qwen1.5-14B, gemma-7b-it, WizardLM-70B [a], and even code-specialized like CodeQwen1.5-7B failed to generate executable code for tasks requiring extensive planning or interdependent processes (e.g., full-pipeline skelenton script presented in Appendix A). Commonly, these smaller models exhibit similar issues, such as cutting the given codes, changing comments without code completion, partial completion, and returning code as it is given. As shown in Table R2.2 below, even vanilla GPT-3.5 and GPT-4o struggle with generating executable code for complex tasks, involving extensive planning or interdependent processes. Similar findings have been corroborated by prior studies [b, c], suggesting that these challenges are systemic to smaller language models and not unique to our framework.

Table R2.2. Average performance scores (from Tables 5,6, and 7) between different LLM frameworks.

We hope this explanation clarifies the design, evaluation rationale, and constraints in testing with smaller models. We are committed to extending this work to address these challenges in future iterations. Thank you again for the opportunity to clarify these points.

References

• [a] Xu, Can, et al. "WizardLM: Empowering large pre-trained language models to follow complex instructions." ICLR. 2024.
• [b] Guo, Siyuan, et al. "DS-Agent: Automated Data Science by Empowering Large Language Models with Case-Based Reasoning." ICML 2024.
• [c] Hong, Sirui, et al. "Data interpreter: An llm agent for data science." arXiv preprint arXiv:2402.18679 (2024).

Dear Reviewer XBqa,

Thank you for taking the time to thoroughly review our responses and for acknowledging that we have addressed all of your concerns satisfactorily. We greatly appreciate your recognition of our efforts and your thoughtful review. If there are any additional points or areas that require further clarification or discussion to strengthen our paper, we would be happy to address them.

Sincerely,

Submission 7839 Authors

Using LLM agent to solve autoML is not a new research problem. Specifically, what is the difference between the paper: https://arxiv.org/abs/2410.17238 Both papers use reflection, while the paper above uses a more sophisticated mechanism, e.g., MCTS, which yields a strong performance. Why don't you consider using MCTS instead? Does this work compatible with MCTS?

R4. Thank you for referencing the MCTS-based work. We acknowledge the reviewer’s concerns and appreciate the opportunity to clarify our framework’s novelty. We also agree with the reviewer that using LLM agents to solve AutoML tasks is not a novel research problem, as we discussed in the Introduction and Related Work sections. However, existing methods focus primarily on either data or modeling aspects, whereas our framework is a full-pipeline AutoML platform, which aims to output deployable model APIs given task descriptions.

Regarding the SELA paper, we would like to make the following clarifications.

• Novelty and Differences: We believe that our work complements rather than competes with MCTS-based methods. The novelty of our framework lies in RAP's integration with LLMs within an end-to-end AutoML context. RAP leverages LLMs’ inherent language understanding and API retrieval capabilities to adaptively define the search space without significant computational overhead, an area where MCTS-based approaches may face challenges. While SELA focuses on iterative pipeline optimization within a tree structure, AutoML-Agent prioritizes modularity, flexibility, and real-time retrieval capabilities for full-pipeline development. These differences highlight our framework’s contribution to advancing AutoML in distinct yet complementary directions.
• Rationale for Not Using MCTS: The decision not to adopt MCTS is rooted in our framework’s design principles. In this work, we aim for full-pipeline automation and end-to-end model deployment, requiring flexibility across a wide range of tasks. Additionally, one of our key motivations was to reduce the computational overhead associated with the search process; employing MCTS would have increased search costs. Instead, we address the problem by facilitating communication between agents (relying solely on LLM inferences) and enhancing planning quality through retrieval-augmented knowledge.
• Compatibility with MCTS: The structure of AutoML-Agent could theoretically integrate MCTS into the planning and execution stages. While theoretically feasible, integrating MCTS in AutoML-Agent would introduce computational overhead in the search process.

We would also like to point out that the paper you suggested became available approximately three weeks (on October 22) after the ICLR submission deadline (on October 2). Nevertheless, we have included a comparison with SELA in Table 1 of the revised version of the paper. Therefore, we believe that our framework's novelty and contributions remain valid as presented.

What is the reason for fine-tuning mixtral for prompt parser?

R5. The reason for fine-tuning is to ensure that the Prompt Agent generates a precise JSON format required for subsequent processes, which necessitates instruction tuning of an LLM. While the GPT-4 family is overly complex for this task, models like Llama-2, Mistral, and Mixtral perform effectively, with Mixtral yielding the best results. As discussed in Section 3.3, without instruction tuning, the LLM outputs JSON key-value pairs that, while semantically valid, are less predictable. The JSON schema details are available at /prompt_agent/WizardLAMP/template_schema.json. Fortunately, due to recent updates in OpenAI’s GPT models, we can now also use the proposed JSON schema with OpenAI APIs for more seamless integration.

We appreciate the reviewers' detailed feedback and are pleased to address your questions and comments as follows.

The evaluation is only performed on synthetic datasets; no public benchmarks are tested. Public datasets like IMDB, Cora, and Citeseer are used for autoML tasks. The synthetic dataset I mentioned is the one for testing the performance of autoML algorithms. Since in 4.1, no evidence shows public benchmarks for evaluating autoML algorithms are given. Please specify any benchmarks if you do use some.

R1. Thank you for this suggestion. However, this seems like a misunderstanding, since no synthetic datasets were used in our experiments. As presented in Table 2 and Section A.2, all datasets used for the experiments are public datasets collected from various sources. Additionally, Cora and Citeseer were also used in the experiments for node classification tasks.

I am unclear about how the manager agent verifies the plans if no plans are executed and tests with generated code. Can you explain a bit?

R2. Throughout the planning and search process, all agents interact solely through message passing to achieve their goals. The Agent Manager does not verify the plans, as it is responsible for creating them. Instead, it verifies the executed results provided by the Model Agent(s) against any constraints specified by the user. For each Model Agent to execute a given plan, it receives the plan, relevant instructions, knowledge from the Agent Manager, and anticipated data analysis results from the Data Agent. Based on this information, the Model Agent performs model search and hyperparameter optimization (HPO). The outcome is a textual summary predicting the top-k models with associated optimal hyperparameters and performance scores, which is executed only after verification.

The generated plans seem homogeneous, similar, and not diverse enough, making it easy for all to fail if one fails (Appendix C2). Therefore, I am not sure about the effectiveness of the plan selection, given that the manager selects the plan based on the "simulated" result from the data and model agent.

R3. As described above, the Agent Manager does not directly verify the plans but instead evaluates the executed results, which can vary even with similar plans depending on how the Data and Model agents perform on each plan. The diversity of the plans can be easily controlled by adjusting the model hyperparameter temperature. We have provided additional planning results in the anonymous repository at /example_plans. Furthermore, we propose a multi-stage verification process that accommodates cases where all plans fail by using recorded feedback through reflection. This component enhances the success rate of any revised plans.

Dear Reviewer fRNo,

As the discussion period for ICLR ends on November 26th (AoE), we would greatly appreciate it if you could review the rebuttal we submitted, which we hope addresses all the concerns you raised, before the discussion period closes.

Thank you again for your time and valuable insights. Please don’t hesitate to let us know if any additional clarifications are needed.

Best regards,

Submission 7839 Authors