MLE-World, AI 연구 에이전트가 ML 실험을 실행 전에 예측하도록 돕다

들어가며

AI 연구 에이전트는 머신러닝 프로그램을 작성하고 수정하면서 실험 결과를 바탕으로 성능을 개선할 수 있다. 하지만 후보 프로그램 하나를 평가하려면 대개 실제 모델 학습과 검증이 필요하다. 후보가 많아질수록 실행 비용도 커지기 때문에, 에이전트가 각 단계에서 탐색할 수 있는 선택지는 제한된다. MLE-World는 이 문제를 해결하기 위해 실제 실행 전에 결과를 예측하는 세계 모델을 제안한다.

핵심 내용

• 실험의 종료 출력을 예측한다. MLE-World는 Qwen3.5-4B를 미세 조정한 모델로, 프로그램 코드와 데이터 설명을 입력받아 오류 추적, 학습 로그, 검증 점수를 예상한다.
• 예측 신뢰도도 함께 제공한다. 모델은 결과뿐 아니라 예측이 맞을 가능성을 나타내는 신뢰도 점수를 생성한다. 신뢰도가 낮은 후보는 실제로 실행해 확인하고, 비교적 확신이 높은 후보는 예측만으로 우선 평가할 수 있다.
• 다섯 가지 테스트 시점 확장 정책을 사용한다. 순차적 수정, Best-of-N, 쌍별 판단, 융합, 그리고 신뢰도에서 탐색 보너스를 얻는 UCB 방식이 포함된다.
• 기존 프롬프트 방식보다 높은 일치율을 보였다. 17개 MLE-bench Lite 과제의 보류 프로그램에서 실제 종료 출력과의 일치율은 77.9%였다. 비교 대상인 더 큰 프롬프트 기반 모델은 41.1%를 기록했다.
• 실제 실행을 줄이면서 성능을 높였다. MLE-World를 활용한 테스트 시점 확장은 세계 모델이 없는 에이전트보다 관측 평균 성능을 최대 22.1% 높였고, 실제 실행 횟수도 줄였다.

의미와 영향

MLE-World의 목표는 학습과 평가를 없애는 것이 아니다. 먼저 저렴한 예측으로 후보를 비교하고, 실제 계산은 유망하거나 불확실한 후보에 집중하는 것이다. 이 구조를 사용하면 에이전트는 한 단계에서 더 많은 아이디어를 검토하면서도 모든 후보를 직접 실행해야 하는 부담을 피할 수 있다.

신뢰도 점수가 중요한 이유도 여기에 있다. 예측된 로그나 점수는 실제 실험의 증거가 아니므로 항상 오류 가능성이 있다. MLE-World는 이 불확실성을 의사결정에 반영한다. 신뢰도가 높은 예측은 탐색 범위를 넓히는 데 활용하고, 신뢰도가 낮은 예측은 실제 실행을 통해 검증한다.

이 연구는 AI 연구 에이전트용 세계 모델이 단순히 결과를 맞히는 도구를 넘어, 제한된 실행 예산을 배분하는 장치가 될 수 있음을 보여준다. 향후 실용성을 높이려면 예측 정확도뿐 아니라 모델이 자신의 불확실성을 얼마나 정확히 추정하는지도 중요해질 것이다.

출처: Hugging Face Daily Papers