MLE-World: Predicting Machine Learning Experiment Outcomes with a Learned World Model
Organizations: University of Washington · University of Illinois at Urbana-Champaign · University of Oxford · Columbia University ·
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/.
Figures & tables
Appendix figures & tables24 assets
Supplementary material from the paper’s appendix.
Appendix
| Programs | % | |
|---|---|---|
| All training programs | 497,056 | 100.0 |
| Generator model | ||
| Qwen3-Coder-Next | 209,573 | 42.2 |
| Qwen3.5-4B | 129,907 | 26.1 |
| Qwen3.5-122B-A10B | 77,239 | 15.5 |
| Qwen3.5-9B | 54,815 | 11.0 |
| Base model | Qwen3.5-4B |
|---|---|
| Method | SFT of all parameters (LlamaFactory), vision tower frozen |
| Training data | 497,056 programs (Table 5 ) |
| Chat template | qwen3_5 , thinking disabled |
| Added tokens | 82 exception class tokens ( EXC_None , EXC_ValueError , …; Appendix C ) |
| Sequence length | 12,288 tokens, packed without attention across examples |
| Epochs / steps | 5 / 6,685 |
| None ( null ) | MissingCorpusError |
|---|---|
| AcceleratorError | ModuleNotFoundError |
| ArrowInvalid | NameError |
| ArrowTypeError | NotADirectoryError |
| AssertionError | NotFittedError |
| AttributeError | NotImplementedError |
| AxisError | OSError |
| target | Terminal output | Exception class | Crash class | Validation metric (precision / recall) |
|---|---|---|---|---|
| terminal output (untrimmed) + submission flag | 59.5 | 71.4 | 64.4 | 87.3 / 83.5 |
| terminal output (trimmed) + submission flag | 59.4 | 71.7 | 64.8 | 88.7 / 84.3 |
| + exception class | 60.7 | 73.0 | 66.7 | 87.3 / 87.6 |
| + validation metric | 62.8 | 75.3 | 69.1 | 89.0 / 89.5 |
| validation metric + submission flag | – | – | – | 90.0 / 77.0 |
| MLE-World | Coder | DeepSeek | Judge (Coder) | |
| Reasoning | off | off | on (max effort) | off |
| Temperature | 0 (greedy) | 0 (greedy) | 1.0 | 1.0 |
| Top- | 1.0 | 1.0 | 0.95 | 0.95 |
| Top- | – | – | – | 40 |
| Samples per program | 1 | 1 | 1 | 3 |
| Max output tokens | 8,192 | 8,192 | 32,768 | 3,072 |
| Method | Terminal output | Exception class | Crash class | Validation metric |
|---|---|---|---|---|
| Task majority | – | 34.2 1.2 | – | 55.2 2.3 |
| Nearest program | 28.0 1.1 | 41.1 1.2 | 14.0 1.0 | 55.6 2.2 |
| Parent record | 22.3 1.0 | 34.2 1.2 | 6.7 0.7 | 48.8 2.3 |
| Random record | 11.2 0.8 | 22.2 1.0 | 3.5 0.6 | 20.2 1.8 |
| MLE-World | 77.9 1.0 | 86.4 0.8 | 76.4 1.3 | 78.1 1.9 |
| DeepSeek | Coder | MLE-World (4B + SFT) | |
| Finished before token cap | 94.9 | 91.8 | 99.7 |
| Exception class | 93.6 | 100.0 | 100.0 |
| Terminal output | 91.8 | 91.8 | 99.7 |
| Validation metric | 91.8 | 91.8 | 99.7 |
| All three fields | 91.4 | 91.8 | 99.7 |
| Median output tokens | 12,042 | 144 | 160 |
| Terminal output | Exception class | ||||||
| MLE | DS | Coder | MLE | DS | Coder | ||
| Stage | |||||||
| Draft | 561 | 78.3 | 40.9 | 16.5 | 86.5 | 55.4 | 26.6 |
| Improve | 561 | 83.2 | 38.9 | 21.4 | 89.8 | 57.2 | 29.8 |
| Debug | 578 | 72.3 | 43.5 | 15.5 | 82.9 | 59.0 | 26.3 |
| Outcome | |||||||
| Cases | Agreement | |
|---|---|---|
| All cases | 120 | 96.7 |
| Actual failures | 82 | 95.1 |
| Crashes of the same class | 40 | 93.3 |
| Clean runs | 20 | 100.0 |
| Selection rule | Mean percentile | Valid submissions |
|---|---|---|
| Improve stage, eight hours; real check after each selection | ||
| Metric greedy, | 38.58 3.52 | 38/51 |
| Metric UCB, | 45.38 2.36 | 44/51 |
| Debug stage, eight hours; real check after each selection | ||
| Greedy, | 41.88 2.10 | 48/51 |
| UCB, , | 48.37 1.82 | 48/51 |
| Setting | Mean percentile | Valid submissions | |
|---|---|---|---|
| (a) Ungated comparison at three hours | |||
| Metric only BoN reference † | 0 | 36.47 0.10 | 38/51 |
| UCB | 0.25 | 37.28 4.87 | 38/51 |
| UCB | 0.5 | 41.71 2.89 | 38/51 |
| UCB † | 1 | 38.07 1.93 | 41/51 |
| UCB | 2 | 34.82 0.92 | 35/51 |
| Mean percentile (mean iterations) | ||||
|---|---|---|---|---|
| Gated | ||||
| Policy | Ungated | |||
| Sequential | 32.31 (22.0) | 37.39 (27.0) | 20.86 (30.5) | 47.20 (37.6) |
| BoN | 36.47 (29.6) | 44.90 (49.1) | 48.30 (69.4) | 52.25 (58.0) |
| Pairwise | 40.44 (28.5) | 40.19 (44.1) | 38.60 (44.7) | 45.26 (71.6) |
| Fusion | 34.01 (24.9) | 38.90 (51.9) | 44.59 (44.6) | 39.48 (45.4) |
| Random | Confidence gate | percentile | ||||
|---|---|---|---|---|---|---|
| Policy | Valid | Pct. | Hours | Pct. | Hours | SE |
| BoN | 46/51 | 47.37 | 5.54 | 47.21 | 4.49 | |
| Evaluation point | Mean percentile | vs. Base | Valid (%) |
|---|---|---|---|
| 8 h cutoff | 41.9 3.5 | +5.1 3.4 | 88.2 3.4 |
| Final (16 h cap) | 45.0 2.3 | +8.1 6.2 | 94.1 3.4 |
| Final submission | Verified new bests | |||||
|---|---|---|---|---|---|---|
| Policy | Valid (%) | Mean pct. | Pct. if valid | First (h) | Trials with one | Per trial |
| Base AIDE | 78.4 | 36.9 | 47.0 | 0.43 | 40 | 3.4 |
| MLE-World guided | ||||||
| Sequential | 90.2 | 41.1 | 45.5 | 0.62 | 46 | 4.5 |
| + confidence gate | 96.1 | 47.7 | 49.7 | 0.68 | 49 | 4.4 |
| Pairwise | 92.2 | 42.8 | 46.4 | 1.06 | 47 | 5.1 |
| Rule | Mean percentile | mean pct. (points) | Valid (%) | valid (points) |
|---|---|---|---|---|
| Selection: pick one of the same four candidates | ||||
| Uniform over the pool | 10.9 | – | 26.1 | – |
| Screen by predicted outcome | 16.4 | +5.51 1.18 | 41.9 | +15.8 3.1 |
| + ranking by predicted metric | 17.1 | +0.68 0.79 | 42.5 | +0.5 1.3 |
| Real validation of all four † | 19.8 | +2.65 2.19 | 47.1 | +4.6 2.6 |
| Pairwise: judge the same four candidates | ||||
| Task | Teams | Gold rank | Gold pct. | Golds | Any medal | Policy pct. (range) |
| Tasks with a gold | ||||||
| detecting-insults-in-social-commentary | 50 | 5 | 90.0 | 17 | 17 | 81.3–100.0 |
| dogs-vs-cats-redux-kernels-edition | 1,315 | 12 | 99.1 | 4 | 9 | 33.3–95.5 |
| plant-pathology-2020-fgvc7 | 1,318 | 12 | 99.1 | 4 | 5 | 10.1–92.9 |
| tabular-playground-series-dec-2021 | 1,189 | 12 | 98.7 | 4 | 4 | 46.3–72.7 |
| aerial-cactus-identification | 1,221 | 12 | 87.7 † | 2 | 2 | 50.4–69.8 |
| Golds | Mean percentile | ||||||
|---|---|---|---|---|---|---|---|
| Policy | detecting- insults | Other | Any medal | Invalid | Gold trials | Other valid trials | Total |
| Base AIDE | 3 | 3 | 11 | 11 | 11.8 | 25.1 | 36.9 |
| Sequential | 3 | 3 | 11 | 5 | 11.5 | 29.5 | 41.1 |
| Pairwise | 3 | 5 | 14 | 4 | 15.4 | 27.4 | 42.8 |
| Fusion | 2 | 1 | 11 | 2 | 5.9 | 41.3 | 47.2 |
| BoN | 3 | 0 | 9 | 6 | 5.9 | 33.6 | 39.5 |
| Method | Terminal output | Exception class | Crash class | Validation metric |
|---|---|---|---|---|
| DeepSeek | 42.1 2.8 | 63.7 2.8 | 45.4 3.2 | 21.3 5.8 |
| Coder | 21.1 2.2 | 30.3 2.6 | 18.9 2.4 | 31.9 6.8 |
| MLE-World (4B, no CoT) | 53.3 2.8 | 60.3 2.8 | 56.3 3.1 | 44.7 7.3 |
| Held out task | Related training tasks | Shared structure and remaining shift |
|---|---|---|
| APTOS Blindness Detection | Plant Pathology 2020; Dog Breed | Image classification with pretrained visual encoders; retinal severity introduces ordinal labels and a new metric. |
| 3D Object Detection | Dog Breed; Aerial Cactus | Shared image, array, and PyTorch components can expose familiar import and file errors; 3D geometry and detection remain new task specific computations. |
| RANZCR Catheter/Line | Plant Pathology 2020 | Image preprocessing and visual prediction heads are shared; chest radiographs and multiple catheter/line labels change the target. |
| Method | Terminal output | Exception class | Crash class | Validation metric |
|---|---|---|---|---|
| APTOS Blindness Detection | ||||
| DeepSeek | 49.3 4.9 | 68.0 4.5 | 59.0 5.6 | 16.7 8.6 |
| Coder | 18.3 3.8 | 32.0 4.6 | 15.3 4.2 | 16.7 9.1 |
| MLE-World (4B, no CoT) | 53.0 4.9 | 60.0 4.8 | 53.6 5.6 | 50.0 12.1 |
| 3D Object Detection | ||||
| DeepSeek | 35.7 4.7 | 52.0 4.9 | 26.5 5.0 | 31.8 10.1 |
Test-Time Scaling
Test-time scaling studies how extra inference-time computation can improve model answers, decisions, and generated outputs.
See all 223 papers