MLE-WORLD: PREDICTING MACHINE LEARNING EXPERIMENT OUTCOMES WITH A LEARNED WORLD MODEL

AI research agents can propose, execute, and improve solutions to machine learning tasks, but they face a critical bottleneck: execution cost. An agent must wait for each candidate's execution to finish before it receives feedback, which makes it hard to scale at test time by evaluating more candidates. We introduce MLE-World, a learned world model for machine learning engineering that predicts a candidate solution's terminal output (including error traces, training logs, and validation scores) from its code and a description of the task's data, without executing the program. We train MLE-World by supervised fine-tuning Qwen3.5-4B on agent trajectories that pair solution code with observed execution outputs. Evaluating a candidate then requires a single inference call to a 4B model, so an agent can screen several candidates per step and execute only one. Because these predictions are not always accurate, MLE-World also reports a confidence score for each predicted outcome, which we utilize in two ways. First, we propose a confidence gate to only execute solutions with low confidence while relying on the world model to simulate executions with high confidence. Second, because greedily selecting the candidate with the best validation score does not encourage exploration, we propose a selection rule inspired by upper confidence bounds (UCB) to add an exploration bonus based on confidence to predicted validation scores. On 17 known MLE-bench Lite tasks, all five gated policies raise mean leaderboard percentile from 36.9 to 46.449.2, and the gate cuts real execution by 2446%.