MLE-World: Predicting Machine Learning Experiment Outcomes with a Learned World Model

AI research agents can propose, execute, and improve programs that solve machine learning tasks. However, executing each candidate program requires training and evaluating ML models, limiting how many candidates an agent can explore at each step. We introduce MLE-World, a world model finetuned from Qwen3.5-4B that predicts a program's terminal output (error traces, training logs, and validation scores) from its code and data description without executing it, together with a confidence score estimating the prediction's correctness. Because a prediction costs one inference call rather than a training run, MLE-World makes test time scaling (TTS) affordable for these agents. We build five TTS policies, guided by MLE-World, that compare several candidates per step by their predicted outputs: Sequential revision, Best-of-N, Pairwise judging, Fusion, and a UCB style rule that also adds an exploration bonus from this confidence score. Since predictions can be wrong, we run a candidate for real only when its confidence score is low. On 17 MLE-bench Lite tasks, MLE-World matches the real terminal output on 77.9% of held out programs, versus 41.1% for much larger prompted models. MLE-World guided TTS achieves up to 22.1% higher observed mean performance than the agent without a world model, with less real execution. Project webpage: https://mle-world.github.io/.

한 줄 요약

MLE-World는 Qwen3.5-4B를 기반으로 학습된 월드 모델로, ML 실험 결과를 예측해 실제 실행 없이 효율적인 검색을 가능하게 한다.

핵심 기여도

MLE-World는 4B 파라미터 모델로, 코드와 데이터 설명만으로 실제 실행 결과(에러, 로그, 검증 점수)를 77.9% 정확도로 예측한다.
5가지 TTS 정책(Sequential revision, Best-of-N, Pairwise judging, Fusion, UCB)을 제안해 예측 기반 후보 비교를 수행한다.
예측 신뢰도 점수를 사용해 실제 실행을 최소화하면서도 22.1% 더 높은 평균 성능을 달성한다.
13–1,061× 적은 GPU 추론 시간으로 기존 프롬프트 기반 모델(41.1% 정확도)을 초과한다.

핵심 아이디어

MLE-World는 기존 연구 에이전트가 ML 실험을 실행할 때 발생하는 높은 비용을 해결하기 위해, 실제 실행 없이 실험 결과를 예측하는 월드 모델을 제안한다. 이 모델은 코드와 데이터 설명만으로 터미널 출력을 예측하고, 예측의 정확도를 신뢰도 점수로 제공한다. 이는 실제 실행 대신 추론 단일 호출로 가능하며, TTS(Timescale Scaling)를 저비용으로 가능하게 한다. 핵심 아이디어는 실험 후보를 비교할 때 예측된 출력을 기반으로 선택하고, 신뢰도가 낮은 경우에만 실제 실행을 수행하는 것이다. 이는 UCB(상위 신뢰 경계) 정책과 결합되어 탐색과 활용의 균형을 유지한다.

기술적 접근법

**모델**: Qwen3.5-4B를 기반으로 학습된 MLE-World는 실패 및 성공한 실행의 코드-출력 쌍을 기반으로 학습된다.
**예측 항목**: 에러 트레이스, 훈련 로그, 검증 점수.
**TTS 정책**: 5가지 정책(Sequential revision, Best-of-N, Pairwise judging, Fusion, UCB)이 예측된 출력을 기반으로 후보를 비교하고 선택한다.
**신뢰도 기반 실행 제어**: 예측 신뢰도가 낮은 경우에만 실제 실행을 수행한다.
**실험 환경**: MLE-bench Lite 데이터셋(17개 태스크, 1,700개 테스트 프로그램)에서 평가.

주요 결과

**MLE-bench Lite**: MLE-World는 77.9%의 프로그램에서 실제 터미널 출력과 일치하는 예측을 제공한다. 이는 DeepSeek, Qwen3-Coder-Next 등 대규모 프롬프트 모델(41.1%)을 크게 앞선다.
**성능 개선**: TTS 정책을 사용한 MLE-World는 실제 실행 없이 22.1% 더 높은 평균 성능을 달성한다.
**실행 비용 절감**: 13–1,061× 적은 GPU 추론 시간으로 동일한 성능을 유지한다.

의의 및 한계

MLE-World는 ML 실험의 반복 실행 비용을 줄이고, 더 많은 후보를 탐색함으로써 에이전트의 효율성을 높인다. 특히, 실패한 실행을 예측하는 능력은 디버깅과 후속 개선에 유용하다. 그러나 실패한 실행에 대한 예측이 우수한 반면, 성공한 훈련 결과나 타임아웃 예측은 아직 명확하지 않다. 또한, 학습된 코드와 유사한 태스크에서의 전이 성능은 제한적이며, 더 넓은 평가와 학습 데이터 확장이 필요하다.

실용적 활용

MLE-World는 ML 엔지니어링, 자동화된 하이퍼파라미터 튜닝, 코드 생성 및 디버깅 도구에 적용 가능하다. 특히, 대규모 실험 반복이 필요한 산업 연구 환경에서 실행 시간과 비용을 줄이는 데 유용하다.