MLE-STAR: Machine Learning Engineering Agent via Search and Targeted Refinement

이 논문은 ML 엔지니어링 agent가 LLM이 이미 아는 익숙한 모델에 기대고 코드 전체를 한 번에 고치던 문제를, 웹 검색으로 과제에 맞는 모델을 찾고 ablation으로 고른 code block 하나만 반복해 다듬는 방식으로 풀었다.

이 논문은 한마디로 웹 검색과 ablation 기반 code block 수정으로 Kaggle식 ML 과제의 해법 코드를 만드는 MLE agent MLE-STAR다. Google Search로 찾은 모델들로 초기 해법을 만들고, ablation study로 성능 영향이 가장 큰 code block을 골라 그 부분만 여러 계획으로 고친 뒤, 병렬로 만든 해법을 agent가 제안한 ensemble 전략으로 합친다. MLE-bench Lite 22개 대회에서 Gemini-2.5-Pro로 63.6%의 대회에서 메달을 땄고, Gemini-2.0-Flash 기준으로 AIDE의 메달 비율 25.8%를 43.9%로 올렸다.

핵심 요약

• 문제: 기존 MLE(머신러닝 엔지니어링) agent는 LLM이 이미 아는 지식에 기대어 scikit-learn 같은 익숙한 방법으로 쏠리고, 텍스트 분류 대회에도 logistic regression을 제안하는 식으로 과제에 맞는 최신 모델을 놓쳤다. 또 매 반복마다 코드 전체를 한 번에 고쳐서, feature engineering 같은 한 component를 깊게 탐색하기 전에 다른 단계로 넘어갔다.
• 입력과 출력: 이 논문은 Kaggle식 ML 과제를 푸는 agent 시스템이다. 입력은 과제 유형, 데이터 형태, 평가 지표를 담은 과제 설명과 여러 파일로 된 데이터셋이고, 출력은 모델을 학습하고 hold-out 검증 점수를 출력하는 단일 Python 스크립트다. 마지막에 별도 agent가 이 스크립트에 test 데이터 로딩과 submission 파일 생성을 붙인다.
• 방법: retriever agent가 먼저 Google Search로 과제에 맞는 모델과 예제 코드를 찾고, 모델마다 간단한 스크립트를 만든 뒤 검증 점수가 좋은 순서로 합쳐 초기 해법을 만든다. outer loop마다 ablation study 스크립트로 성능에 가장 크게 영향을 준 code block을 고르고, inner loop에서 planner가 이전 시도 점수를 보며 그 block만 고치는 계획을 여러 번 내 가장 좋은 것을 남긴다. 병렬로 만든 해법들은 agent가 스스로 제안하고 다듬는 ensemble 전략으로 합치고, debugging agent, data leakage checker, data usage checker가 실행 오류, test 정보 누출, 데이터 누락을 고친다.
• 구현: 기본 LLM은 Gemini-2.0-Flash이고 Gemini-2.5-Pro와 Claude-Sonnet-4로도 실험했으며, 모델 검색 4개, inner loop 4번, outer loop 4번, 병렬 해법 2개, ensemble 탐색 5라운드에 과제당 24시간 제한을 뒀다. 96 vCPU, 360 GB 메모리, 16 GB V100 GPU 8장에서 해법 하나에 평균 14.1시간이 걸렸고 AIDE는 15.4시간이었다. 부록은 retriever 출력, ablation 코드, 최종 해법 코드 같은 예시 산출물을 저자 GitHub 저장소에 두었다고 적었다.
• 결과: MLE-bench Lite 22개 대회(seed 3개)에서 Gemini-2.5-Pro를 쓴 MLE-STAR는 63.6%의 대회에서 메달을 땄고 금메달 비율은 36.4%였다. 같은 Gemini-2.0-Flash에서 메달 비율은 AIDE의 25.8%에서 43.9%로, 중앙값 이상 비율은 39.4%에서 63.6%로 올랐다. 같은 조건에서 ensemble 없이 37.9%이던 메달 비율은 제안한 ensemble로 43.9%가 됐고, 금메달 비율은 단순 평균 ensemble의 25.8%보다 높은 30.3%였다.
• 한계: Kaggle 대회가 공개돼 있어 LLM이 관련 토론 글을 학습했을 위험이 있고, 저자들은 7개 대회의 상위 토론 글 25개와 최종 해법을 짝지은 75쌍을 Gemini-2.5-Pro judge가 모두 새롭다고 판정한 결과를 근거로 들었다. 평가는 MLE-bench 75개 대회 가운데 난이도가 낮은 Lite 22개에서 했고, Claude-Sonnet-4 실험과 DS-Agent 비교는 각각 과제 4개로 범위가 좁다. AIDE의 o1-preview와 GPT-4o 결과는 MLE-bench 저장소에서 가져온 값이며 seed 수가 16개와 36개로 MLE-STAR의 3개와 다르다.

논문 링크

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