Ds-agent: Automated data science by empowering large language models with case-based reasoning

Motivation

Known: LLM 기반 에이전트는 다양한 작업을 자동화할 수 있지만, 데이터 사이언스 태스크에서 합리적인 실험 계획 생성과 hallucination 문제로 인해 높은 완료율을 달성하지 못한다. Kaggle은 데이터 사이언트들의 풍부한 지식과 코드를 보유한 대규모 플랫폼이다.
Gap: 기존 LLM 에이전트들은 데이터 사이언스 작업에서 정상적인 실험 계획을 생성하지 못하며, 이를 해결하기 위한 finetuning은 막대한 계산 자원과 레이블 수집 시간을 요구한다. 따라서 Kaggle의 전문가 지식을 효율적으로 활용하면서도 저자원 환경에 적응할 수 있는 방법이 필요하다.
Why: 자동화된 데이터 사이언스는 전문 지식 없이도 데이터 인사이트에 접근할 수 있게 하여 민주화를 실현한다. 효율적이고 비용 효과적인 솔루션은 실제 배포와 광범위한 사용을 가능하게 한다.
Approach: DS-Agent는 개발 단계에서 CBR 프레임워크로 Kaggle 사례를 검색·재사용·피드백에 따라 수정하며, 배포 단계에서는 단순화된 CBR로 과거 성공 사례를 직접 코드로 생성한다.

Figure 4. Success rate of four different agents in the development

개발 단계 성능: GPT-4를 사용한 DS-Agent는 12개 태스크에서 100% 성공률 달성
배포 단계 성능: GPT-3.5와 GPT-4에서 각각 85%, 99% one pass rate로 최고 성능 달성 (기준 방법: 56%, 60%)
오픈소스 LLM 성능: Mixtral-8x7b-Instruct의 one pass rate를 6%에서 31%로 5배 이상 향상
비용 효율성: GPT-4 표준 시나리오 $1.60/실행에서 저자원 환경 $0.135/실행으로 감소, GPT-3.5는 $0.06에서 $0.0045로 감소
평균 one pass rate 개선: 대체 LLM들에 대해 36% 평균 개선 달성

Figure 2. Comparison between (a) RAG based LLMs and (b) CBR

Retriever 설계: 태스크 τ와 이전 피드백 l_{t-1}을 기반으로 case bank에서 유사한 과거 사례를 검색하는 확률 분포 p_R 구현
LLM 기반 생성: 검색된 사례 c, 태스크 τ, 피드백 l_{t-1}을 입력으로 하여 솔루션 y를 생성하는 조건부 분포 p_LLM 활용
Evaluator 구현: 생성된 솔루션의 실행 결과를 평가하여 피드백 l을 생성하는 p_E 구성
반복 루프: 식 (1)의 CBR 공식에 따라 피드백과 검색된 사례를 주변화하여 일관된 성능 개선 달성
Case bank 관리: 최적 평가 성능을 보인 솔루션을 case bank에 저장하여 향후 재사용 가능하게 함
배포 단계 적응: 개발 단계에서 수집한 성공 사례를 배포 단계에서 검색하고 최소한의 수정으로 재사용
RAG 대비 개선: RAG는 단일 변수 c만 주변화하지만, CBR은 피드백 l_{t-1}도 함께 고려하여 반복적 개선 가능

평가 제한: 12개 개발 태스크와 18개 배포 태스크로 제한적이며, 더 다양한 데이터 사이언스 시나리오에 대한 평가 필요
배포 단계 가정: 배포 태스크가 개발 단계와 유사한 분포를 가정하므로, 크게 다른 분포의 태스크에 대한 성능 미지수
Case bank 의존성: Kaggle에서 수집한 사례 품질과 양에 크게 의존하므로, 사례 부족 시 성능 저하 가능성
LLM 비용 문제: GPT-4 기반 개발 단계에는 여전히 비교적 높은 비용 ($1.60/실행) 소요
후속 연구: (1) 더 효과적인 case 검색 및 랭킹 메커니즘 개발, (2) 도메인 외(out-of-distribution) 태스크에 대한 적응 능력 향상, (3) 더 큰 scale의 case bank 구축 및 관리 방안 연구