MLE-World, 코드 실행 없이 ML 실험 결과 예측해 AI 연구 에이전트 탐색 비용 줄인다

머신러닝 과제를 푸는 AI 연구 에이전트가 후보 프로그램을 실제로 학습시키지 않고도 결과를 미리 가늠하게 하는 방법이 논문으로 공개됐다. 논문이 제안한 MLE-World는 Qwen3.5-4B를 미세조정한 월드 모델로, 프로그램의 코드와 데이터 설명만 받아 오류 기록, 학습 로그, 검증 점수 같은 최종 출력을 예측한다. 예측이 맞을 확률을 나타내는 신뢰도 점수도 함께 내놓는다.

연구팀은 MLE-bench Lite 과제 17개에서 MLE-World가 따로 떼어 둔(held out) 프로그램의 실제 최종 출력과 77.9% 일치했다고 밝혔다. 훨씬 큰 프롬프트 기반 모델들의 일치율은 41.1%였다. MLE-World로 이끈 테스트 시점 확장(TTS)은 월드 모델이 없는 에이전트보다 관측된 평균 성능이 최대 22.1% 높았고, 실제 실행은 더 적었다고 한다.

작동 방식

AI 연구 에이전트는 프로그램을 제안하고 실행해 개선하지만, 후보마다 모델을 학습하고 평가해야 해서 한 단계에서 살펴볼 수 있는 후보 수가 제한된다. MLE-World는 예측 한 번이 학습 한 번이 아니라 추론 호출 한 번의 비용으로 끝나므로, 여러 후보를 비교하는 테스트 시점 확장을 감당할 수 있게 한다는 것이 연구팀의 설명이다.

연구팀은 MLE-World의 예측으로 후보를 비교하는 TTS 정책 다섯 가지를 만들었다. 순차 수정(Sequential revision), Best-of-N, 쌍별 판정(Pairwise judging), 융합(Fusion), 그리고 신뢰도 점수로 탐색 보너스를 더하는 UCB 방식의 규칙이다. 예측이 틀릴 수 있어서 신뢰도 점수가 낮을 때는 해당 후보를 실제로 실행한다.