MLE-World: Predicting Machine Learning Experiment Outcomes with a Learned World Model
作者: Brian Ko, Han Wang, Hejia Geng, Li Erran Li, Sherry Yang
分类: cs.LG, cs.AI
发布日期: 2026-10-27
🔗 代码/项目: PROJECT_PAGE
💡 一句话要点
提出MLE-World以预测机器学习实验结果,提升测试时间效率
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)
关键词: 机器学习 世界模型 测试时间缩放 程序预测 自动化优化
📋 核心要点
- 现有方法在执行候选程序时需要大量的训练和评估,限制了探索的候选数量。
- MLE-World通过预测程序输出而无需执行,显著提高了测试时间效率,支持多种TTS策略。
- 在17个MLE-bench Lite任务中,MLE-World的预测准确率达到77.9%,显著优于其他大型模型。
📝 摘要(中文)
AI研究代理可以提出、执行和改进解决机器学习任务的程序。然而,每个候选程序的执行需要训练和评估ML模型,这限制了代理在每一步可以探索的候选数量。我们提出MLE-World,这是一个基于Qwen3.5-4B微调的世界模型,能够在不执行程序的情况下,根据代码和数据描述预测程序的终端输出(错误跟踪、训练日志和验证分数),并提供一个置信度评分来估计预测的正确性。由于预测只需一次推理调用而非训练运行,MLE-World使得测试时间缩放(TTS)对这些代理变得可负担。我们构建了五种TTS策略,利用MLE-World指导,通过预测输出比较多个候选程序。MLE-World在17个MLE-bench Lite任务中,匹配真实终端输出的准确率为77.9%,而更大模型的准确率仅为41.1%。
🔬 方法详解
问题定义:本论文旨在解决AI研究代理在执行候选程序时面临的高成本问题,现有方法需要大量训练和评估,限制了候选程序的探索数量。
核心思路:MLE-World通过微调的世界模型预测程序的终端输出,避免了实际执行的开销,从而提高了测试时间效率。
技术框架:MLE-World的整体架构包括输入代码和数据描述,生成预测输出和置信度评分。基于这些预测,构建了五种TTS策略以比较候选程序。
关键创新:MLE-World的主要创新在于其能够在不执行程序的情况下,准确预测程序的输出,并提供置信度评分,显著提升了测试时间缩放的可行性。
关键设计:在模型设计中,MLE-World使用了微调的Qwen3.5-4B模型,采用了特定的损失函数来优化预测准确性,并通过置信度评分来决定是否执行候选程序。
📊 实验亮点
MLE-World在17个MLE-bench Lite任务中,预测准确率达到77.9%,显著高于41.1%的大型模型。同时,MLE-World指导的TTS策略实现了高达22.1%的性能提升,减少了实际执行的需求。
🎯 应用场景
该研究的潜在应用领域包括自动化机器学习、程序生成和优化等。MLE-World能够帮助AI代理更高效地探索和评估候选程序,从而加速机器学习模型的开发和优化,具有重要的实际价值和未来影响。
📄 摘要(原文)
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/.