MLE-bench: Evaluating Machine Learning Agents on Machine Learning Engineering

이 논문은 AI agent의 end-to-end ML 엔지니어링 능력을 전체로 재는 benchmark가 적은 문제를, Kaggle 대회 75개를 offline 환경으로 옮기고 사람 leaderboard의 메달 기준으로 채점하는 MLE-bench(machine learning engineering benchmark)로 풀었다.

이 논문은 한마디로 Kaggle(데이터 과학 경진대회 platform) 대회 75개를 offline 환경으로 옮기고, AI agent가 model을 학습해 낸 제출물을 사람 leaderboard의 메달 threshold로 채점하는 benchmark 논문이다. 대회마다 설명, dataset(대개 새로 나눈 train-test split), grading 코드, Private leaderboard snapshot을 두고, 대표 지표는 동메달 이상을 받은 시도의 비율이다. LLM을 감싸 도구 호출과 실행 loop를 맡는 scaffold 세 개를 GPT-4o로 비교하고, 가장 나은 AIDE(Kaggle 대회용으로 해법 tree를 탐색하는 scaffold) 위에서 모델 네 개를 비교한 결과 o1-preview와 AIDE 조합이 대회의 평균 16.9%에서 메달을 받아 가장 높았다. 저자들은 contamination(공개 해법을 학습해 benchmark 점수가 부풀려지는 현상) 실험에서 성적이 memorization 때문에 체계적으로 부풀려졌다는 증거를 찾지 못했다고 보고했다.

핵심 요약

• 문제: 언어 모델은 많은 coding benchmark에서 높은 성능을 내고 architecture 설계와 model 학습 같은 ML 작업에서도 발전하고 있지만, 자율적인 end-to-end ML 엔지니어링을 전체로 재는 benchmark는 적었다. 저자들은 현재의 ML 엔지니어링 일을 대표하는 어려운 task를 고르고, 평가 결과를 사람 수준의 성과와 비교할 수 있는 측정을 목표로 삼았다.
• 입력과 출력: agent는 Kaggle 대회의 설명, dataset, 제출 형식을 받아 offline 환경에서 model을 학습하고 submission.csv 파일을 낸다. 제출물은 대회 평가 지표를 다시 구현한 grading 코드로 채점되고, agent가 그 대회에 당시 참가했다고 보고 Private leaderboard와 비교해 메달로 환산된다. 논문 종류는 ML 엔지니어링 agent를 평가하는 benchmark 논문이다.
• 방법: Meta Kaggle dataset의 완료된 대회 5673개에서 Community 대회를 빼고 남은 586개를 손으로 검토해 75개를 골랐고, 경험 있는 ML 엔지니어가 쓸 만한 해법을 내는 데 드는 시간(model 학습 시간 제외)으로 Low 22개(2시간 미만), Medium 38개(2에서 10시간), High 15개(10시간 초과)로 나눴다. 공개 test set이 없는 대회는 공개 training 데이터에서 새 test set을 떼어 냈고(별도 표시가 없으면 10% 비율), 대표 지표는 참가 팀 수에 따라 달라지는 Kaggle 메달 threshold에서 동메달 이상을 받은 시도의 비율이다. agent는 model 대신 자기 지식으로 예측을 직접 쓰거나 온라인 해법을 볼 수 없으며, log 검사 도구와 표절 검사 도구 Dolos가 이를 점검하고, Kaggle 상위 notebook 50개 가운데 하나와 코드 유사도가 60%를 넘는 시도는 실격 처리하고 추가 검토 대상으로 표시한다.
• 구현: 기본 실험에서 각 대회 시도는 dataset, 제출 검증 server, Python package가 든 Ubuntu 20.04 Docker container에서 최대 24시간 동안 실행되고, 시도마다 가상 CPU 36개, 메모리 440GB, Nvidia A10 GPU 한 장이 주어지며, seed는 기본 3개이고 o1-preview(AIDE)는 16개, GPT-4o(AIDE)는 36개였다. scaffold는 AIDE, MLAB(MLAgentBench의 ResearchAgent), OpenHands(OpenHands platform의 CodeActAgent) 세 가지를 썼고, AIDE 위에서 o1-preview, GPT-4o, Claude 3.5 Sonnet, Llama 3.1 405B 네 모델을 비교했다. README에 따르면 GitHub 저장소 openai/mle-bench에는 dataset 준비 script, 대회별 grading script, 평가한 agent, agent가 쓰는 기반 Docker image(mlebench-env) 설정, 규칙 위반 탐지 도구와 표절 탐지 도구가 있고, 비용이 부담되면 Low 난이도 22개 대회만 쓰는 Lite 평가를 권한다.
• 결과: 표 2에서 AIDE scaffold 위의 네 모델 가운데 o1-preview가 대회의 평균 16.9%에서 메달을 받아 가장 높았고 GPT-4o는 8.7%, Claude 3.5 Sonnet은 7.6%, Llama 3.1 405B는 3.0%였으며, 같은 GPT-4o를 MLAB과 OpenHands에서 실행하면 0.8%와 4.4%였다. 시도 횟수 실험에서 o1-preview는 pass@1 16.9%에서 pass@8(8번 시도 가운데 한 번이라도 메달을 받은 대회 비율) 34.1%로 올랐고, 대회당 시간 제한을 100시간으로, AIDE의 최대 node 수를 5,000으로 늘린 실험에서 GPT-4o(AIDE)는 8.7%(24시간 기준)에서 11.8%로 올랐다. 표 3의 hardware 실험에서 GPT-4o(AIDE)의 메달 비율은 GPU 없음 9.1%, GPU 한 장 8.7%, GPU 두 장 10.2%로 비슷했고, 표 4에서도 원래 설명으로 8.5%, 출처를 알아보기 어렵게 고쳐 쓴 설명(seed 10개로 실행)으로 8.4%를 받아 비슷했다.
• 한계: 데이터가 공개 Kaggle 대회라 model이 대회 설명, 해법, test set을 포함한 dataset까지 학습했을 수 있고, 저자들은 GPT-4o 실험에서 contamination의 체계적 효과를 찾지 못했지만 고수준 전략의 재사용은 잡아내기 어렵고 이후 model에 대한 보장은 없다고 밝혔다. 저자들은 MLE-bench가 문제 정의가 분명하고 dataset이 깔끔하며 지표가 분명한 Kaggle 대회를 골랐으므로, 문제 정의가 흐릿한 경우가 많고 dataset과 지표를 정하는 일까지 과제인 실제 AI 연구개발 능력의 일부를 다룬다고 밝혔다. 기본 설정 한 번에 24시간 곱하기 75개 대회로 1800 GPU 시간이 들고, o1-preview(AIDE)는 seed 하나의 75개 대회에서 평균 입력 token 127.5M개와 출력 token 15.0M개를 썼다.

논문 링크

외부 연구를 정리한 자료입니다. HDATF가 발표한 논문이나 제품 성능을 측정한 결과는 아닙니다.