BudgetMLAgent: A Cost-Effective LLM Multi-Agent system for Automating Machine Learning Tasks

TL;DR

BudgetMLAgent introduces a cost-effective multi-agent system using low-cost LLMs and expert collaboration to improve success rates in ML tasks while significantly reducing costs.

Contribution

The paper presents a novel multi-agent LLM system that combines low-cost models with expert calls, achieving better performance at a fraction of the cost compared to single GPT-4 systems.

Findings

94.2% cost reduction over GPT-4 single-agent system

Improved success rate of 32.95% versus 22.72% for GPT-4

Effective use of low-cost models with expert collaboration

Abstract

Large Language Models (LLMs) excel in diverse applications including generation of code snippets, but often struggle with generating code for complex Machine Learning (ML) tasks. Although existing LLM single-agent based systems give varying performance depending on the task complexity, they purely rely on larger and expensive models such as GPT-4. Our investigation reveals that no-cost and low-cost models such as Gemini-Pro, Mixtral and CodeLlama perform far worse than GPT-4 in a single-agent setting. With the motivation of developing a cost-efficient LLM based solution for solving ML tasks, we propose an LLM Multi-Agent based system which leverages combination of experts using profiling, efficient retrieval of past observations, LLM cascades, and ask-the-expert calls. Through empirical analysis on ML engineering tasks in the MLAgentBench benchmark, we demonstrate the effectiveness of our system, using no-cost models, namely Gemini as the base LLM, paired with GPT-4 in cascade and expert to serve occasional ask-the-expert calls for planning. With 94.2\% reduction in the cost (from \$0.931 per run cost averaged over all tasks for GPT-4 single agent system to \$0.054), our system is able to yield better average success rate of 32.95\% as compared to GPT-4 single-agent system yielding 22.72\% success rate averaged over all the tasks of MLAgentBench.