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

AI summary

Overview

• Research area: Machine learning engineering (MLE) agents, LLM-based "world models," and test-time scaling/search.
• Technical level: Advanced (assumes familiarity with LLM finetuning, agent scaffolds, and search policies), though the core idea is explained accessibly.
• Scope: The paper introduces MLE-World, a 4B-parameter model finetuned to predict a candidate ML program's terminal output without running it, and builds five search policies plus a confidence-based execution gate on top of those predictions, evaluated on 17 MLE-bench Lite tasks.

What This Paper Is About

AI research agents that solve machine learning tasks work by writing programs, training and evaluating them, and revising based on the output. That execution step is slow and expensive, which caps how many candidate programs an agent can try within a fixed time budget. The authors train a small language model, MLE-World, to predict what a program's terminal output would be (error trace, training logs, and validation score) from the code and data description alone, then use those cheap predictions to compare many candidates per step and only run the real thing when the model is unsure.

Key Contributions

1. An execution world model for MLE. MLE-World is a 4B model, finetuned from Qwen3.5-4B on logged code/output pairs, that reads a program, its installed packages and versions, a data overview, and the time limit, and predicts a four-field record: exception class, terminal output, validation metric, and whether a submission was written. It outperforms much larger prompted models.
2. Five world-model-guided test-time scaling (TTS) policies. Sequential revision, Best-of-N (BoN), Pairwise judging, Fusion, and a UCB-style rule that adds an exploration bonus from the model's own confidence. Each draws several candidates per step, predicts their outputs, and executes only the selected one.
3. Selective execution via a confidence score. The first decoded token is the predicted exception class; its renormalized token probabilities give p_max (confidence in the most likely outcome), u(x) = 1 - p_max (uncertainty), and p_clean (probability of a clean run). A gate triggers a real run when confidence falls below a threshold, using separate thresholds for predicted successes and failures.
4. Empirical validation on MLE-bench Lite. Accuracy measurements, cost measurements, ablations isolating the value of predicted terminal output and the exploration bonus, and gated-versus-ungated comparisons across 17 tasks with three seeds each.

Main Findings

• Prediction accuracy well above prompted baselines. On 1,700 held-out programs, MLE-World's predicted terminal output matches the real output on 77.9% of programs, versus 41.1% for DeepSeek-V4-Flash-0731 and 17.8% for Qwen3-Coder-Next. It also leads on exception class (86.4% exact match), crash class (76.4%), and validation metric within 5% relative error (78.1%).
• Much cheaper inference. Predicted outputs cost 13–1,061× less GPU inference time than the prompted models, making it practical to screen multiple candidates per step.
• Model beats corpus-copying baselines. Nearest-program retrieval reaches 28.0% terminal-output match and parent-record copying 22.3%, and task-majority reaches 34.2% on exception class — all below MLE-World, indicating it learns program-to-outcome relationships rather than memorizing task priors.
• All five policies improve observed performance. Without gating, UCB reaches 52.1 mean percentile versus 45.4 for Base AIDE; Sequential 51.3, BoN 50.2, Fusion 47.5, Pairwise 47.1.
• Gating cuts execution time while broadly holding or improving quality. Gating saves 30–40% of real execution time (e.g., Fusion 7.93 h → 4.77 h per trial; Sequential 8.29 h → 5.68 h). Gated Sequential reaches 55.4 mean percentile, the highest observed mean in the paper — the abstract reports up to 22.1% higher observed mean performance than the agent without a world model, with less real execution. UCB is the one policy that dips slightly under gating (52.1 → 51.1).
• Predicted terminal output, not just predicted scores, helps. On 150 shared states, Fusion rises from 14.0 to 20.6 mean percentile with predicted terminal output; Sequential 10.4 → 15.3; Pairwise 13.0 → 15.4.
• The exploration bonus matters. Varying only the Improve-stage β (0 vs 1) gives 45.4 mean percentile with the bonus versus 38.6 with greedy selection.
• Confidence tracks accuracy. Across quintiles of p_max on the 1,700 held-out programs, match rates for metrics, exception classes, and terminal outputs all rise with confidence — justifying spending real runs on low-confidence predictions.
• Gating is not just random skipping. In a separate 8 h control, BoN with confidence gating uses 19% less real execution than a random gate at a similar final percentile.
• Structured target fields help. Adding exception class and validation metric to the target raises terminal-output agreement from 59.4% to 62.8% and metric recall from 84.3% to 89.5% in the ablation.
• Known weak spot: timeouts. Timeouts are the least accurately predicted outcome because runtime must be estimated from model size, epochs, and data size; a missed timeout still costs a full run.
• Partial out-of-distribution signal. On three tasks excluded from finetuning, the 4B MLE-World still reaches 53.3% terminal-output accuracy versus 42.1% for DeepSeek with reasoning and 21.1% for Coder.

Methodology in Plain English

The researchers start from AIDE, an existing agent scaffold that searches over a tree of programs using three moves (Draft, Debug, Improve), where every node normally costs one real training run.

They collect the logged results of roughly half a million such runs — 497,056 training programs generated by five different agent models ranging from 4B to 122B parameters, deliberately including failed and buggy runs (52.3% crash, 7.0% timeout, 40.7% clean). Each record pairs the program plus context (packages, a data overview, the time limit) with the outcome.

They then finetune all parameters of Qwen3.5-4B with standard supervised learning on those pairs for five epochs, using 8 NVIDIA H200 GPUs (42.7 hours), a 12,288-token sequence length, and a learning rate of 1e-5. Each of the 82 exception classes is added to the tokenizer as a single token, so the very first decoded token yields a full probability distribution over outcomes — and that distribution is the confidence score.

At search time, each step draws N candidates, asks MLE-World to predict each one's output, and applies one of five policies to pick or synthesize a single program. A gate then decides whether to trust the prediction or run the program for real, using a lower threshold for predicted successes (long, expensive runs) and a higher one for predicted failures (which usually crash within seconds and yield a useful error trace); predicted timeouts are never run.

Evaluation uses 17 of the 22 MLE-bench Lite tasks — the five excluded tasks are ones where the agents rarely produced clean submissions — with 100 held-out programs per task (1,700 total), a 12 h budget per trial, and three seeds. Predictions are graded against real executed outputs, with an LLM judge for terminal-output and crash-class matching (agreement with human labels: 96.7%). Search results are reported as mean leaderboard percentile rather than medals, since medal thresholds are noisy.

Why This Matters

Impact on research. The paper reframes execution cost in MLE agents as a modeling problem rather than only a scheduling or pruning problem. Prior work prunes branches, predicts performance, or ranks candidates — but those approaches leave per-execution cost unchanged and provide no error traces or training logs to debug with. MLE-World supplies full terminal output, and its confidence signal lets the agent decide when a prediction is untrustworthy — a more targeted criterion than deciding based only on how expensive an operation is.

Real-world applications:

• Automated ML engineering platforms: agents that search over model pipelines under a wall-clock budget, screening many candidate pipelines per step.
• Cheap pre-flight checks in CI for ML code: predicting whether a training script will crash, and with which exception, before allocating GPUs.
• Hyperparameter and architecture search: using predicted validation metrics and uncertainty to prioritize which configurations actually get trained.
• Research-assistant tooling for ML scientists: surfacing plausible error traces and scores for candidate experiments before committing compute.

Industry relevance. Compute is the binding constraint on agentic ML systems. A 4B model that predicts outcomes at 13–1,061× lower inference cost than large prompted models, and that saves 30–40% of real execution time through gating, directly translates into either lower spend or more search iterations per dollar. The dependence on a large corpus of logged prior executions also points to a data-moat dynamic: organizations with accumulated run logs are best positioned to train such simulators.

Future Directions

• Generalization beyond the training tasks. The main evaluation is in-distribution (new programs on tasks seen in finetuning). The authors report a promising but partial out-of-distribution result (53.3% on three excluded tasks) and name a generalizable world model as future work, suggesting added reasoning or a larger model.
• Better runtime and timeout prediction. Timeouts are the least accurate predicted outcome and a missed timeout still costs a full run; improving runtime estimation is explicitly flagged.
• Closing the loop on the corpus requirement. Training uses about 500,000 prior executions, favoring settings with accumulated logs. Whether Base AIDE with access to the corpus alone could match MLE-World is stated as untested.
• Extending confidence-gated search. The gate is policy-agnostic and the UCB bonus is described as an empirical exploration proxy rather than a formal confidence bound; both invite stronger uncertainty formulations and tuning (e.g., per-policy thresholds, which currently vary between 0.50 and 0.83).

Target Audience

Researchers and engineers working on LLM agents for code and machine learning, test-time scaling and search over programs, and learned simulators or world models for agentic settings. It is most useful to readers with some background in LLM finetuning and agent scaffolds (advanced level), but the framing — cheap predicted feedback plus confidence-gated execution — is legible to anyone building cost-constrained agent systems.

Authors’ abstract

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/.