[논문 리뷰] AutoML-Agent: A Multi-Agent LLM Framework for Full-Pipeline AutoML

Patara Trirat, Wonyong Jeong외 1명|arXiv (Cornell University)|2024. 10. 03.

Machine Learning and Data Classification인용 수 4

한 줄 요약

이 논문은 검색 증강 계획과 다단계 검증을 통해 조율되는 전문화된 LLM 에이전트를 활용하여 데이터 검색 및 전처리부터 모델 배포에 이르기까지 전체 기계학습 파이프라인을 자동화하는 다에이전트 LLM 프레임워크인 AutoML-Agent를 제안한다. 이 프레임워크는 14개 데이터셋과 7개 작업에서 평균 96.2%의 성공률을 기록하며, 기존 기준 대비 성공률과 효율성 면에서 뛰어나며, MCTS 기반 방법에 비해 검색 시간을 최대 90% 감소시켰다.

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 large language models (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.

연구 동기 및 목표

• 기존의 LLM 기반 AutoML 시스템이 하이퍼파라미터 튜닝이나 모델 선택과 같은 고립된 파이프라인 단계에만 집중하는 데에 한계가 있음을 해결하기 위해 종단 간 자동화를 달성하고자 한다.
• 다양한 지식 기반 계획 탐색을 가능하게 하기 위해 검색 증강 전략을 도입함으로써 전체 파이프라인 AutoML의 계획 복잡도를 줄이고자 한다.
• 역할 기반 에이전트 분해와 다단계 검증을 통해 코드 생성 정확도를 향상시키고 자율적 ML 파이프라인 구축 시 환상(hallucination)을 줄이고자 한다.
• 데이터, 모델, 평가 단계를 하나의 프레임워크에 통합함으로써 비전문가 사용자를 대상으로 효율적이고 확장 가능하며 배포 가능한 AutoML를 실현하고자 한다.
• 다에이전트 협업이 검색 증강 계획과 검증을 통해 다양한 기계학습 작업 전반에서 성공률과 계산 효율성을 크게 향상시킨다는 것을 입증하고자 한다.

제안 방법

• AutoML-Agent는 전문화된 LLM 에이전트(예: 데이터, 모델, 검증 에이전트 등)가 병렬로 작동하는 다에이전트 아키텍처를 활용하여 데이터 전처리, 신경망 아키텍처 탐색, 하이퍼파라미터 튜닝과 같은 하위 작업을 분해하고 해결한다.
• 관련 지식(예: 데이터셋, 이전 파이프라인 등)을 검색하여 다수의 후보 계획을 생성하는 검색 증강 계획 전략을 도입함으로써 최적의 해법 탐색 범위를 넓힌다.
• 각 계획은 원자적 하위 작업으로 분해되며, 역할 기반 프롬프팅을 통해 전문화된 에이전트에 할당되어 병렬 실행이 가능해져 전체 파이프라인 실행 시간을 단축시킨다.
• 다단계 검증 프로세스는 중간 결과를 검증하고 코드 생성 에이전트를 안내함으로써 반복 피드백을 통해 환상과 오류를 최소화하고 정확성을 확보한다.
• 종단 간 파이프라인 품질 평가를 위해 성공률, 정규화된 성능, 종합 점수를 조합한 종합 점수 체계를 사용한다.
• 유효하지 않은 계획을 조기에 필터링하고 비용 및 시간 최적화를 통해 제약 조건을 고려한 실행을 지원함으로써 제약 조건이 없는 설정 대비 총 실행 시간을 4.8% 감소시키고 비용을 18.5% 절감한다.

실험 결과

연구 질문

• RQ1 다에이전트 LLM 프레임워크는 전문가 수준의 설정 없이도 데이터 검색부터 배포까지 전체 AutoML 파이프라인을 효과적으로 자동화할 수 있는가?
• RQ2 검색 증강 계획은 단일 계획 또는 히وري스틱 기반 접근 방식에 비해 성공률와 탐색된 파이프라인의 다양성을 향상시키는가?
• RQ3 역할 기반 에이전트 분해와 병렬 실행은 전체 파이프라인 AutoML의 시간과 계산 비용을 크게 줄일 수 있는가?
• RQ4 다단계 검증 전략은 코드 환상과 생성된 ML 파이프라인의 신뢰성을 얼마나 줄이고 향상시키는가?
• RQ5 다양한 실제 데이터셋과 작업에서 AutoML-Agent는 최신 기준 대비 성공률, 성능, 효율성 면에서 어떻게 비교되는가?

주요 결과

• AutoML-Agent는 14개 데이터셋과 7개의 후행 작업에서 평균 96.2%의 성공률를 기록하였으며, DS-Agent(66.1%)와 제약 조건 고려 GPT-3.5(7.7%)와 같은 기준 대비 뚜렷한 성능 우월성을 보였다.
• MCTS 기반 SELA 대비 평균 검색 시간을 90% 감소시켜 2037.18초에서 249.91초로 단축하면서도 정규화된 성능 점수는 유지 또는 향상시켰다.
• 각 작업당 5개의 계획을 생성한 결과, 이미지 및 텍스트 분류 작업에서는 100%의 성공률를 기록하였고, 노드 분류 작업에서는 90.6%의 종합 점수를 확보하여 다양한 작업에 걸쳐 뛰어난 강건성을 입증하였다.
• 다단계 검증 전략은 코드 신뢰성을 향상시켜 환상과 잘못된 종속성 처리를 줄였으며, 코드 생성 및 배포 성공률 96.9%로 이를 입증하였다.
• 리더보드 비교에서 AutoML-Agent는 8개 대회 중 7개에서 Agent K를 앞서며 중앙값 리더보드 순위 88.5를 기록했고, Agent K의 86.5보다 높았으며, 5개의 과제에서 성능이 낮은 편이었지만 전체적으로 유의미한 우월성을 보였다.
• 제약 조건 고려 설정에서는 총 비용을 18.5% 절감($0.32 대비 $0.27)하고 시간을 4.8% 감소시켜(512.35초 대비 538.25초) 성능을 희생시키지 않은 채 효율성 향상을 입증하였다.