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

MLE-World. Predicting Machine Learning Experiment Outcomes with a Learned World Model. By Brian Ko and four others, from the University of Washington, the University of Illinois at Urbana-Champaign, and three other institutions. Published on the twenty-seventh of October.

AI research agents can propose programs that solve machine learning tasks, execute them, and improve them. The bottleneck is the execution. Running each candidate program means actually training and evaluating ML models. Every candidate costs a training run, and that limits how many candidates an agent can explore at each step. An agent that can only afford to try a few ideas will only ever compare a few.

MLE-World is a world model, finetuned from Qwen3.5-4B, that predicts what a program would produce without executing it. Given the program's code and a description of its data, it predicts the terminal output: error traces, training logs, and validation scores. Alongside each prediction it gives a confidence score, estimating whether that prediction is correct. A prediction costs one inference call rather than a training run, and that difference is what makes test time scaling, or TTS, affordable for these agents.

On top of the world model sit five TTS policies. Each one compares several candidates per step by their predicted outputs. They are Sequential revision, Best-of-N, Pairwise judging, Fusion, and a UCB style rule that also adds an exploration bonus drawn from the confidence score. Since predictions can be wrong, there is a safeguard. A candidate is run for real only when its confidence score is low, so the expensive path is kept for the cases where the prediction is least trusted.

The evaluation covers seventeen MLE-bench Lite tasks. On held out programs, MLE-World matches the real terminal output seventy-seven point nine percent of the time. Much larger prompted models manage forty-one point one percent.

Then the agent itself. Guiding test time scaling with MLE-World gives up to twenty-two point one percent higher observed mean performance than the agent without a world model. It gets there while using less real execution, not more.

The trade on offer is simple to state. Swap a training run for a single inference call, keep real execution for the cases where the model is unsure, and an agent can afford to weigh more candidates per step. A small model that reads code and guesses its output becomes the thing that lets the search widen.