ArchPilot: A Proxy-Guided Multi-Agent Approach for Machine Learning Engineering

TL;DR

ArchPilot is a multi-agent system that improves machine learning architecture search efficiency by using proxy evaluations and adaptive search, reducing the need for costly full training runs.

Contribution

It introduces a novel multi-agent framework with proxy-guided evaluation and a Monte Carlo Tree Search-based orchestration for scalable ML architecture search.

Findings

Outperforms state-of-the-art methods like AIDE and ML-Master on MLE-Bench.

Reduces reliance on full training runs, saving computational resources.

Demonstrates effective architecture search with limited budgets.

Abstract

Recent LLM-based agents have demonstrated strong capabilities in automated ML engineering. However, they heavily rely on repeated full training runs to evaluate candidate solutions, resulting in significant computational overhead, limited scalability to large search spaces, and slow iteration cycles. To address these challenges, we introduce ArchPilot, a multi-agent system that integrates architecture generation, proxy-based evaluation, and adaptive search into a unified framework. ArchPilot consists of three specialized agents: an orchestration agent that coordinates the search process using a Monte Carlo Tree Search (MCTS)-inspired novel algorithm with a restart mechanism and manages memory of previous candidates; a generation agent that iteratively generates, improves, and debugs candidate architectures; and an evaluation agent that executes proxy training runs, generates and optimizes proxy functions, and aggregates the proxy scores into a fidelity-aware performance metric. This multi-agent collaboration allows ArchPilot to prioritize high-potential candidates with minimal reliance on expensive full training runs, facilitating efficient ML engineering under limited budgets. Experiments on MLE-Bench demonstrate that ArchPilot outperforms SOTA baselines such as AIDE and ML-Master, validating the effectiveness of our multi-agent system.