Latent Dynamics · 2026-10-28
MLE-World: Predicting Machine Learning Experiment Outcomes with a Learned World Model
Brian Ko, Han Wang, Hejia Geng, Li Erran Li, Sherry Yang
AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.
TL;DR
MLE-World is a learned world model fine-tuned from Qwen3.5-4B to predict a program’s terminal output (error traces, training logs, validation score) from its code and data description without executing it, plus a confidence score. It enables test-time scaling for MLE agents via five world-model-guided policies with a confidence gate for selective real execution, improving held-out matching and observed mean performance while reducing real executions.
Why it matters
Executing candidate programs in MLE agents is expensive, limiting how many candidates can be explored per step. By predicting terminal outputs and confidence in one inference call, MLE-World lets agents compare multiple candidates cheaply and only run real executions when predictions are unreliable, enabling search with more feedback under fixed budgets.
Method
- Fine-tune Qwen3.5-4B with supervised finetuning to predict four structured fields from a program and context: exception class, terminal output, validation metric, and submission flag; add a confidence score from first-token probabilities (p_max).
- Use MLE-World predictions to guide five test-time scaling (TTS) policies for candidate selection: Sequential revision, Best-of-N, Pairwise judging, Fusion, and UCB with an exploration bonus from outcome uncertainty.
- Add a confidence gate: after selecting a candidate, run it for real only when prediction confidence is missing/low (different thresholds for predicted successes vs failures), otherwise use predicted feedback for next steps.
Limitation
The paper states that recognizing execution failures transfers more clearly than predicting successful training outcomes and timeouts, which remains less established; it also says a broader evaluation and training-family ablations are needed.
Abstract (from arXiv)
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/.
Related papers
- Predict the Reward, Not the World: World Reward Models for ML Agents
- ANTE: Anticipate, Then Execute -- Learning Internal World Models for Coding Agents from Verified Beliefs
- RoboJEPA: Scaling Robotic Latent World Models
- Tev: Scaling Long-Horizon Agentic Environments with Dense Progress Rewards
- World Models Dream of Success: Diagnosing and Repairing Failure Insensitivity in Robot World Models
Publication record
Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.
First public release: 2026-10-27 · Indexed by Pub-AI: 2026-10-28 UTC
MLE-World: Predicting Machine Learning Experiment Outcomes with a Learned World Model
Preprint posted: 2026-10-27 · Original record ↗
- arxiv
- 2610.37175
Explicit source versions (1)
- v1 · 2026-10-27
- Preprint available
Source date: 2026-10-27 · First detected: 2026-10-28 UTC · Historical baseline
Source evidence (2)
- Official evidence ↗ · RSS guid=oai:arXiv.org:2610.37175v1; arxiv:announce_type=new; item pubDate (official metadata)
- Official evidence ↗ · atom:published (official metadata)
Abstract · arxiv
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/.