MLE-World：用学习到的世界模型预测机器学习实验结果 MLE-World: Predicting Machine Learning Experiment Outcomes with a Learned World Model
微调4B模型预测代码执行的完整终端输出与分数，让MLE智能体的测试时扩展便宜可行
前置知识
世界模型（World Model）
模型驱动 RL 与智能体研究中的世界模型指学习函数 $f_\phi$ 预测「执行某动作后环境变成什么样」，无需真正交互即可模拟、在想象中规划，通用形式 $\hat{o}_{t+1}=f_\phi(\kappa,o_t,a_t)$，经典工作如 World Models、Dreamer、MBPO。
MLE-World 整个方法就是训练这样一个模型：把「执行候选程序」当作动作，把终端输出当作下一观测，用一次 4B 前向推理替代一次真实的模型训练与评估。
测试时扩展（Test-Time Scaling, TTS）
指在推理阶段投入更多计算以换取更好结果的范式，有两条路径：一是采样更多完整轨迹（如 Best-of-N、重复采样，Large Language Monkeys），二是在单条轨迹内增加交互步数（自我修订、树搜索、更长交互）。两条路径都需要可靠的反馈信号来引导搜索方向。
在 MLE 场景中，无论哪种 TTS 都意味着更多次真实训练运行，代价极高。本文的核心主张是：世界模型提供的廉价预测反馈让 TTS 在 MLE 中第一次变得可负担。
AIDE 搜索脚手架
MLE-bench 的参考智能体框架：在以程序为节点、以三种算子为边的树上搜索——Draft 写全新程序、Debug 重写有 bug 的父节点、Improve 对有效父节点做一处原子修改。每个节点消耗一次真实执行，执行器返回训练日志/错误堆栈/验证分数，最后以验证分数最高的节点作为提交。
本文所有实验都构建在 AIDE 之上（尽管方法兼容其他脚手架），理解 Draft/Debug/Improve 三种算子才能理解五种 TTS 策略如何嵌入搜索过程。
UCB 与探索-利用权衡
来自多臂老虎机与贝叶斯优化的选择规则：不单纯选估计价值最高的动作，而是选「估计价值 + 不确定性加成」最高者，如 GP-UCB 使用 $v(s)+\beta\cdot$后验标准差。其目的是避免因估计误差而永久错过被低估的好动作。
论文把这一思想低成本化：用世界模型第一个输出 token 的类分布给出不确定性 $u(x)=1-p_{\max}(x)$ 作为探索加成，实验证明比纯贪心选择高约 6.8 个百分位点。
监督微调（SFT）与负对数似然
在配好的（输入, 目标）数据对上微调语言模型，最小化目标 token 的负对数似然 $\mathcal{L}(\phi)=-\mathbb{E}_{(x,y)\sim\mathcal{D}}\log p_\phi(y\mid x)$。训练后模型学会在给定输入的条件下生成目标格式与内容。
MLE-World 就是对 Qwen3.5-4B 做 SFT 得到的：49.7 万条真实（程序, 执行输出）对，加上 82 个异常类新增 token 的设计都依托标准 SFT 流程。
置信度校准
指模型输出的概率与实际正确率之间的一致程度：若模型声称 90% 置信的预测实际正确率接近 90%，则称校准良好。校准是「何时信任模型预测」的基础，常按置信度分位数（quintile）检验匹配率是否单调上升。
本文的执行门控（$p_{\max}<\gamma$ 才真实运行）和 UCB 探索加成都依赖 $p_{\max}$ 具备校准性；Figure 7 验证了各字段匹配率随 $p_{\max}$ 五分位单调上升，这是整套方法成立的实证前提。
研究动机
AI 研究智能体（如 AIDE）通过「写程序 → 真实执行 → 读终端输出 → 修改」的循环解决 Kaggle 式机器学习工程任务，但这个反馈回路极其昂贵：每验证一个候选程序都要完整训练并评估一次模型。论文给出的数据非常直观——在 17 个 MLE-bench Lite 任务、每 trial 12 小时预算下，Base AIDE 平均花 9.94 小时在真实执行上，其中 83.8% 的运行以崩溃告终、2.8% 超时，大部分算力浪费在注定失败的程序上；每 trial 平均提交 144.9 个程序、只有 13.4% 干净跑完。已有缓解手段各有短板：剪枝和执行前排序（如 ForeAgent 用 LLM 偏好两两比较候选）能降低总执行次数，但单次执行成本不变，而且它们只输出分数或偏好，不提供错误堆栈和训练日志——智能体拿不到可用于调试的具体反馈，想知道程序为什么崩仍得等一次真实运行。执行瓶颈因此直接限制智能体每一步能测试的想法数量：无论是采样更多候选还是做更多修订步，测试时扩展（TTS）都因为每一步背后是一次真实训练而变得难以负担。
本文的目标是本文的目标是让「想象」替代部分「实验」：训练一个执行世界模型 MLE-World，输入候选程序的代码、可导入包及版本、任务数据概览（文件树、CSV 维度、列摘要与示例值）和执行时限，不运行程序就预测其完整终端输出（错误堆栈、训练日志）、异常类型（82 类之一）、验证指标和是否写出提交文件，并附带置信度分数。在此之上实现三件事：第一，用预测输出驱动五种测试时扩展策略（Sequential 修订、Best-of-N、Pairwise 比较、Fusion 合成、UCB 探索），每步筛选 $N=4$ 个候选、只真实执行选中的那一个，使「比较多个候选」的成本从 $N$ 次训练降为 $N$ 次推理加 1 次训练；第二，用置信度门控决定何时必须真实运行——预测置信度 $p_{\max}<\gamma$ 才执行，把真实算力集中花在模型最没把握的地方；第三，实证检验这条路线能否在固定时间预算内提高提交质量，论文报告最高带来相对 22.1% 的观察均值提升，同时节省 30–40% 真实执行时间。
与已有工作不同的是，独特切入角度有三点。其一，把 LLM 世界模型从 web、计算机操作、软件工程智能体迁移到 MLE 场景：最接近的 DSWorld 模拟数据科学工作流的单个操作步骤、按操作成本决定是否模拟，而 MLE-World 预测完整程序的终端输出（含错误堆栈与训练日志），并以预测可靠性而非操作成本决定何时真实执行，直指 model-based RL 的经典难题——错误预测会误导搜索（Janner et al. 2019）。其二，与只做分数预测或偏好排序的前作（ForeAgent、preference models）不同，结构化四字段输出 $\langle$异常类, 终端输出, 指标, 提交标志$\rangle$ 同时支撑两类用途：分数用于候选筛选，错误堆栈用于 Debug 修订；消融还显示加上结构化字段后终端输出一致率从 59.4% 升至 62.8%，多任务目标互相增益。其三，把 GP-UCB 思想低成本化：异常类是 82 个新增 token 中的首个解码 token，其归一化分布免费给出不确定性 $u(x)=1-p_{\max}(x)$，同一信号既作探索加成又作执行门控，无需额外部件。
核心方法
直觉：ML 程序的执行结果计算昂贵但高度可学——代码一行之差就改变错误类型，数据维度与库版本决定崩溃方式，日志与分数由代码和数据共同决定，可从海量执行日志中蒸馏。路线四层。(1) 数据：5 个智能体模型（4B–122B）在 17 个 MLE-bench Lite 任务上跑 AIDE，保留失败轨迹，去重过滤得 497,056 条（程序, 输出）对。(2) 模型：对 Qwen3.5-4B 全参数 SFT，最小化 $\mathcal{L}(\phi)=-\mathbb{E}\log p_\phi(y\mid x)$，目标 $y=\langle c,o,m,b\rangle$（异常类、终端输出、指标、提交标志）。(3) 搜索：嵌入 AIDE，每步抽 $N=4$ 个候选各做一次 4B 前向得 $\hat{y}_i$，五种 TTS 策略之一选出唯一真实执行对象。(4) 门控：$p_{\max}(x)<\gamma$ 才真实执行，预测超时永不执行、预测新最优指标者强制执行，最终以真实验证最优者提交。整体把比较 $N$ 个想法的成本从 $N$ 次训练压到 $N$ 次推理加至多 1 次训练。
核心创新三点。第一，预测终端输出而非分数或偏好：错误堆栈与日志让世界模型同时服务两种用途——按预测指标筛选候选、按预测错误修改程序；这是与 ForeAgent 等偏好排序方法的本质区别，后者无法告诉智能体程序为何崩溃。第二，置信度免费来自输出结构：82 个异常类各绑定一个新增 token，一次贪心解码即得 82 类结果分布 $p_\phi(c\mid x)$，从中读出 $p_{\max}(x)$、$u(x)=1-p_{\max}(x)$ 与干净概率 $p_{\mathrm{clean}}(x)$；Figure 7 证明 $p_{\max}$ 校准良好（各字段匹配率随五分位单调上升），同一分布既驱动 UCB 探索加成又驱动门控 $p_{\max}<\gamma$，无需额外校准模型。第三，成本感知的差异化门控：预测失败者几秒内崩溃（中位 3 秒），真实运行便宜且换回精确错误堆栈，用高阈值 $\gamma_f$ 多放行；预测成功者是长任务（中位 115 秒），才是执行时间大头，用低阈值 $\gamma_s=0.50$ 尽量跳过；预测超时者永不执行，因其运行占满时限只为确认已知失败。
方法步骤详情
流程六步。(1) 数据：757,186 次执行按 token 流哈希去重、过滤得 498,756；排除近重复后每任务抽 100 条构成 1,700 条 held-out，其余 497,056 条训练。(2) 训练：LlamaFactory 全参数 SFT，8×H200 训 42.7 小时，5 epochs，序列 12,288；最终 loss 0.079/0.081。(3) 推理：贪心解码，按序输出异常类 token、终端输出（首尾各 2,500 字符窗口）、验证指标、提交标志。(4) 搜索步：抽 4 候选预测 $\hat{y}_i$：BoN 取预测指标最大者；Sequential 链式修订；Pairwise 由 LLM judge 两两淘汰；Fusion 合成新程序；UCB 按 $s^\star_t=\arg\max[v(s)+\beta u(x_s)]$ 选择（$\beta=0.5$）。(5) 门控：$p_{\max}<\gamma$（$\gamma_s{=}0.50$、$\gamma_f{=}0.56/0.83$）才真实运行，否则预测指导下一步。(6) 以真实验证最优者提交评分。
技术新颖性
新颖性分四层。模型层：首个面向程序级执行结果的 MLE 世界模型，预测完整终端输出而非标量分数，4B 微调模型大幅超越大得多的零样本模型（终端输出 77.9% 对 41.1%），推理成本 2.7 GPU-min、低 13–1,061 倍——执行预测是 SFT 可习得的能力，prompting 覆盖不了。信号层：把置信度压到单个输出 token 的分布上，与事后校准（温度缩放、深集成、多次采样方差）相比零额外成本且校准良好。系统层：成功/失败/超时三档差异化门控把「何时信任模型」形式化为可调 operating point——$\gamma_s\in\{0.90,0.50\}$ 对应 15.5%/17.8% 错误率与 67.7%/99.6% 运行时节省的显式权衡，由 3 小时 sweep 选定，联合优化预测误差与执行成本的设计少见。数据层：混合 5 个智能体策略的日志训练使模拟器对后续使用它的智能体更鲁棒，并证明保留失败轨迹至关重要（只学干净运行的模型在 Debug 候选上最不准）。相对 DSWorld、SWE-World，「完整程序+终端输出+可靠性门控」的组合是本文独有。
实验结果
Q1（1,700 条 held-out）：终端输出一致率 77.9%（DeepSeek 41.1%、Coder 17.8%），异常类 86.4%（57.2%），崩溃类 76.4%（48.2%），指标 5% 误差内 78.1%（18.6%）；语料基线仅 28.0–41.1%，排除背诵。成本：一次预测仅 2.7 GPU-min（基线 36.5–2,889，低 13–1,061×）；超时最弱（46.1%）。Q2（12h、17 任务×3 seeds）：五策略均值均高于 Base AIDE 45.4——UCB 52.1（+6.7±2.5），其余 47.1–51.3；153 态诊断：Fusion 14.0→20.6、提交率 +12.5 点，换真实输出无增益（20.4）；β=1 对贪心 +6.8。Q3（门控）：五策略省 30–40% 真实执行（Fusion 省 39.8%），gated Sequential 55.4 最高即相对 22.1% 提升；随机门控对照证明省在「跳对运行」（4.49h 对 5.54h）。泛化：3 个未训练任务 53.3%，高于 DeepSeek 42.1%。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 终端输出预测（held-out 程序） | 语义一致率（Coder judge，与人工标签一致 96.7%） | MLE-World 77.9%±1.0 | DeepSeek-V4-Flash（最强 prompted 基线）41.1%±1.1；语料基线 nearest program 28.0% | +36.8 个百分点 |
| 异常类型预测 | 精确类匹配（82 类） | MLE-World 86.4%±0.8 | DeepSeek 57.2%±1.2 | +29.2 个百分点 |
| 崩溃类型预测（仅崩溃样本） | 精确类匹配 | MLE-World 76.4%±1.3 | DeepSeek 48.2%±1.5 | +28.2 个百分点 |
| 验证指标预测 | 5% 相对误差内命中率 | MLE-World 78.1%±1.9 | DeepSeek 18.6%±1.7 | +59.5 个百分点 |
| 推理成本 | GPU 时间（1,700 程序） | MLE-World 2.7 GPU-min | DeepSeek 2,889 GPU-min；Coder 36.5 GPU-min | 低 1,061× / 13× |
| MLE-bench Lite 搜索（12h，无门控） | 平均排行榜百分位 | UCB 52.1±0.8 | Base AIDE 45.4±4.8 | +6.7±2.5（五个策略全部为正） |
| MLE-bench Lite 搜索（12h，加门控） | 平均排行榜百分位 / 真实执行时长 | Gated Sequential 55.4（最高）；Gated Fusion 52.1 | Base AIDE 45.4（9.94 h/trial） | 相对提升最高 22.1%；执行时间节省 30.0–39.8% |
| 未微调任务泛化（3 个新任务，300 程序） | 终端输出一致率 | MLE-World（4B，无推理）53.3%±2.8 | DeepSeek（带推理）42.1%；Coder 21.1% | +11.2 个百分点（对最强基线） |
局限与改进
作者承认四条局限。其一，主评估在分布内：held-out 程序虽经严格去重，但 17 个任务全部在微调集中出现，模型学的是「新程序在旧任务上的执行动力学」；3 个未训练任务上明显下滑（53.3% 对分布内 77.9%）。其二，训练依赖约 50 万条历史执行日志，偏好有积累的团队，零样本基线无此任务知识，对比不完全公平；带语料库访问的 Base AIDE 未测试。其三，最终分数由真实执行决定，预测错误（held-out 上 22% 终端输出、14% 异常类不一致）仍可能误导搜索。其四，超时最不准（终端输出 46.1%），漏检一次浪费整个时限。我的观察：UCB 是唯一加门控后均值下降的策略（52.1→51.1，−0.9±2.4），且 51 个 trial 中 6 个无有效提交——95.9% 的 978 个搜索步退回 $p_{\mathrm{clean}}$ 回退路径，回退设计在困难任务上脆弱；门控收益跨策略差异大（Pairwise +1.2±4.1）；结论建立在 17 个相对简单的 Kaggle 任务上（12 个从未产出奖牌），外推到研究级任务无证据。
独立分析的弱点
第一，任务级泛化不足：世界模型学的是「任务×库版本」上的执行动力学，held-out 只是新程序而非新环境，未训练任务终端输出仅 53.3%。改进：扩大任务家族多样性、检索相似任务记录、家族元学习。第二，置信度信号过窄：$p_{\max}$ 只刻画 82 类异常分布，它高不代表日志细节正确（Figure 7 最高五分位匹配率仍低于 100%），而门控恰在预测成功时最依赖日志可信度。改进：多次采样自一致性投票、token 级熵聚合。第三，超时预测需从代码推断运行时长（46.1%），漏检一次浪费整个时限并阻塞执行器。改进：运行时长回归头、相似程序历史检索、增量式世界模型（先预测前几秒日志再决定）。第四，门控阈值是静态超参：$\gamma_s/\gamma_f$ 来自一次 sweep 且跨策略不同（0.56 对 0.83）。改进：conformal prediction 校准或在线 bandit 自适应阈值。第五，UCB 覆盖缺陷：6/51 无提交，$p_{\mathrm{clean}}$ 回退与指标评审相互作用脆弱。改进：每步至少真实验证一个候选、用可用性与指标的联合判据替代概率回退。
未来方向
作者提出两个方向：可泛化的世界模型（引入推理或扩大规模——4B 无推理模型已在未训练任务上超过带推理的 DeepSeek，暗示上限更高）和改进运行时预测。可延伸方向：其一，把 MLE-World 用作 model-based RL 的模拟器直接优化智能体策略（衔接 Yang et al. 2026a），而非仅用于筛选；其二，多步 rollout 规划——在世界模型上做树搜索（MCTS/LATS 式），让智能体在想象中评估多步计划后再落地，本文的逐步单层筛选只是最保守用法；其三，为混合预测/执行搜索建立理论：$u(x)=1-p_{\max}(x)$ 目前只是经验探索代理，仿照 GP-UCB 显式建模预测误差方差、给出带 regret 保证的采集函数很自然；其四，世界模型与智能体共同演化——新轨迹回流训练、模型随搜索进度更新（类似 WebEvolver），并处理分布漂移；其五，扩展到执行有副作用的场景（数据库写入、外部 API），门控从省钱问题升级为安全问题；其六，系统研究节省的 30–40% 执行时间如何最优再投资——附录 O 显示可买更多搜索步，执行预算分配是开放问题。
复现评估
开源前景好：作者承诺发布 checkpoint、含 held-out 划分的执行语料、所有 prompt 及智能体与评估代码。训练配置完全披露：LlamaFactory 全参数 SFT，8×H200 训 42.7 小时，5 epochs，序列 12,288，lr $1\times10^{-5}$，bf16 + FlashAttention-2；最终 loss 0.079/0.081；附录 A–G 覆盖数据统计、去重规则、holdout 构造与输出解析。评估协议详尽：judge 与人工标签一致率 96.7%、门控阈值 sweep、随机门控对照、UCB 消融、按任务/阶段分解全部列出。复现门槛分层：仅评估（下载 checkpoint 跑预测）单卡可行；复现 agent 实验（51 trials × 8–12h × 多设置）需要可观 GPU 集群；最大障碍是训练数据——50 万次真实执行的收集耗时未报告，按每 trial 约 10 小时估算需数千次完整运行。总体：权重公开后属中等难度复现；资源有限团队可评估与微调自己的小模型，从零收集等量语料不现实。
论文图表
