World reward models choose better programs without predicting execution
Predict the Reward, Not the World: World Reward Models for ML Agents
Artificial IntelligenceComputation and Language
Summary
Machine learning agents often try many program options by running them, but this takes a lot of time since most runs are discarded. The authors found that instead of predicting what will happen when a program runs, it’s better to predict which program will do better directly. They created a model called ForeCritic that learns from past results which program is likely superior. This model picks better programs much more often and speeds up the process without losing quality.
What this means in practice
- •For machine learning engineers: Speed up program search by using reward models that pick better candidates directly instead of simulating full execution.
- •For software optimization teams: Reduce the time to find effective program improvements by relying on direct reward prediction rather than costly full program runs.
Authors
Yuxiao Lee, Zherui Yang, Fan Liu, Hao Liu
Abstract
Search-based machine-learning (ML) agents improve a solution by writing candidate programs, executing them, and keeping the best. Execution gives the search reliable feedback, but a single run can take hours, and most runs end in programs the agent discards. Recent agents try to save these runs by letting a large language model act as a world model that predicts how each program would perform. However, a wrong prediction feeds the search false information. We argue that an agent does not need to predict the outcome of execution, only the decision that execution would support. We propose World Reward Models (WRMs), which learn from real execution results to choose directly which candidate will perform better, without predicting any intermediate outcome. We prove that a selector which predicts execution from the same inputs cannot outperform the best direct selector, that predicted outcomes add an error which can make a likelihood-trained world model prefer the worse program on every instance, and that cross-entropy on executed winners drives a WRM toward the optimum. Our WRM, ForeCritic, adapts a 9B typed decision model to 5k executed program pairs. On 12,438 held-out pairs from 26 ML tasks, it picks the better program 79.4% of the time, while general-purpose LLM judges reach 53--62%, and its probabilities are calibrated. Inside AIDE, ML-Master, and ForeAgent, it cuts mean task time by 4.4--7.5 times while keeping their mean scores within 0.01 or raising them.