Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Shubham Gandhi, Manasi Patwardhan, Lovekesh Vig, Gautam Shroff

TL;DR

BudgetMLAgent tackles the high cost of solving complex ML tasks with LLMs by introducing a multi-agent framework that leverages no-cost LLMs as the base, combined with profiling, cascades, memory recall, and occasional planning-expert lifelines. The approach demonstrates dramatic cost reductions (up to $0.054 per run for no-cost configurations) while achieving superior average success rates on MLAgentBench tasks relative to a GPT-4 single-agent baseline. Key contributions include a structured LLM profiling paradigm, a cost-aware cascading mechanism, retrieval-enabled long-term memory, and an occasional GPT-4 planning expert to recover from stagnation. The work suggests that cost-effective, multi-agent collaboration can rival expensive single-agent systems in real-world ML engineering tasks, with meaningful implications for scalable AI-assisted ML workflows.

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.

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

TL;DR

BudgetMLAgent tackles the high cost of solving complex ML tasks with LLMs by introducing a multi-agent framework that leverages no-cost LLMs as the base, combined with profiling, cascades, memory recall, and occasional planning-expert lifelines. The approach demonstrates dramatic cost reductions (up to $0.054 per run for no-cost configurations) while achieving superior average success rates on MLAgentBench tasks relative to a GPT-4 single-agent baseline. Key contributions include a structured LLM profiling paradigm, a cost-aware cascading mechanism, retrieval-enabled long-term memory, and an occasional GPT-4 planning expert to recover from stagnation. The work suggests that cost-effective, multi-agent collaboration can rival expensive single-agent systems in real-world ML engineering tasks, with meaningful implications for scalable AI-assisted ML workflows.

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

Paper Structure

This paper contains 21 sections, 1 figure, 2 tables.

Table of Contents

  1. Introduction
  2. MLAgentBench Dataset
  3. BudgetMLAgent
  4. Multi-Agent LLM Profiling
  5. LLM Cascade
  6. Ask-the-Expert Action
  7. Experimentation and Results
  8. Metrics
  9. Success Rate
  10. Cost
  11. Models and hyperparameters
  12. Baselines
  13. Results and Discussion
  14. RQ1 - Do no-cost and low-cost LLM single-agents sacrifice performance for cost savings?
  15. RQ2 - How do profiling and cascades affect performance and cost?
  16. ...and 6 more sections

Figures (1)

  • Figure 1: BudgetMLAgent defined in Section \ref{['sec:budgetmlagent']} with profiling, cascades, ask-the-expert lifelines andretrieval from logs. Note that GPT4 cascade and planning expert calls both count towards themax GPT4 calls limit. Ge: Gemini Pro (no-cost LLM Agent).