Vector Policy Optimization: Training for Diversity Improves Test-Time Search
Ryan Bahlous-Boldi ⋅ Isha Puri ⋅ Idan Shenfeld ⋅ Akarsh Kumar ⋅ Mehul Damani ⋅ Sebastian Risi ⋅ Omar Khattab ⋅ Zhang-Wei Hong ⋅ Pulkit Agrawal
Abstract
Modern reasoning systems increasingly spend their compute at test time, drawing many candidate solutions and selecting among them with procedures that range from best-of-$N$ sampling to full evolutionary search. The efficiency of this search is bounded by the diversity of the candidate pool: when the proposing policy has collapsed onto a single mode, additional samples are near-duplicates and the search budget is wasted, so reaching a target quality demands ever more rollouts. The standard paradigm of LLM post-training optimizes a pre-specified scalar reward, which drives current LLMs toward these low-entropy response distributions and leaves test-time search saturating early. We propose Vector Policy Optimization (VPO), an RL algorithm that explicitly trains policies to produce diverse, competent solution sets so that each test-time sample is more informative and search reaches higher quality at a smaller budget. VPO exploits the fact that rewards are often vector-valued in practice (e.g. per-test-case correctness in code, or per-criterion ratings from multiple reward models) and is a drop-in replacement for the GRPO advantage estimator that trains the model to emit a set of solutions spanning different trade-offs in the reward space. Across four tasks, VPO reaches the quality a compute-matched GRPO attains at its full sampling budget using far fewer samples: about $10\times$ fewer on the widest-objective tasks and $3\times$ otherwise. Furthermore, VPO lets evolutionary search solve problems a scalar-trained model cannot. As test-time search becomes the dominant way reasoning models spend compute, training for diversity is a promising route to making that search efficient.
Chat is not available.
Successful Page Load