Vector Policy Optimization (VPO): Training LLMs for Diversity to Unlock Test-Time Search
Vector Policy Optimization (VPO) is a post-training algorithm from MIT's Improbable AI Lab, MIT-IBM Computing Research Lab, and Sakana AI. It reframes RL fine-tuning around diversity rather than a single scalar optimum, so that test-time search (best-of-N, evolutionary search, MCTS) can pick the best candidate from a heterogeneous pool.
Why scalar RL fails
GRPO, DPO, and PPO compress multi-faceted rewards into one scalar, producing what is called a "single-mode" policy. When test-time search calls for variety (best@k, pass@k), the model offers near-duplicate candidates. The authors argue that, in real tasks, rewards are inherently vector-shaped: code test cases, multi-objective navigation, multi-hop QA citations, or tool-call format/parameter axes. Average optimization across these axes wastes the dimensions the problem actually has.
Core insight
Once test-time search is standard practice, RL post-training should focus on producing a diverse, capable set of solutions and defer selection to the search loop.
The VPO algorithm
- Vector reward: r ∈ ℝᵈ instead of r ∈ ℝ.
- Set-level objective: maximize R(S) = E_{w∼Dir(1)} [ max_{y∈S} wᵀ r(x,y) ] — for a randomly drawn weight vector on the simplex, at least one candidate in the set must be good.
- Multi-Answer Chains: in one autoregressive rollout, the model emits m=3 answers separated by markers; later candidates see earlier ones and can deliberately cover unexplored regions (in-context exploration).
- Drop-in GRPO replacement: each of G rollouts produces a set S⁽ᵍ⁾; K Dirichlet-sampled weights are shared across the group; the per-rollout Monte-Carlo reward is R̂(S) = (1/K) Σₖ max_{s∈S} w₍ₖ₎ᵀ r(x,s), followed by the standard GRPO z-score advantage applied uniformly to tokens.
- Compute-matched GRPO (n=24 vs n=8): Giving GRPO 3× more samples only reaches 0.763/0.765 — multi-rollout scalar search cannot substitute for diversity-oriented training.
- Goal-Conditioned GRPO on Maze: Conditioning GRPO on w ∼ Dir(1) leads to mode collapse (best@3 = best@6 = 0.205); VPO's in-context exploration is far more effective than explicit conditioning.
- VPO keeps discovering new candidates and eventually solves a subset.
- GRPO plateaus early and cannot make progress.
- Bahlous-Boldi, R., Puri, I., Shenfeld, I., et al., "Vector Policy Optimization: Training for Diversity Improves Test-Time Search," arXiv:2505.17385, 2025.
- Institutions: MIT, Improbable AI Lab, MIT-IBM Computing Research Lab, Sakana AI.
- Models: Qwen3-1.7B/4B/8B, Qwen2.5-Coder-7B.
- Domains: Maze, MuSiQue, EUREQA, ToolRL, LiveCodeBench.
- Baselines: GRPO, Multi-RLVR, Random-Weighting GRPO, Max-at-K, MaxRL, Goal-Conditioned GRPO.
Comparison table: scalar vs vector RL
| Dimension | Scalar RL (GRPO) | Vector RL (VPO) | |---|---|---| | Reward | r ∈ ℝ | r ∈ ℝᵈ | | Objective | max E[w*ᵀ r], w* fixed | max E_w[max_y wᵀ r], w ∼ Dir(1) | | Output | single answer | m=3 answers per rollout | | Effect | one optimum | Pareto-frontier coverage |
Evaluation domains
| Domain | Model | Vector reward | Scalar baseline | |---|---|---|---| | Maze | Qwen3-4B | ℝ⁴: completion + gold + diamond + lava avoidance | uniform mean | | MuSiQue | Qwen3-1.7B | ℝ⁵: 4 citation hops + answer F1 | (Σhops + 3×F1)/7 | | EUREQA | Qwen3-8B | {0,1}⁵: 5-entity exact match | uniform mean | | ToolRL | Qwen3-1.7B | ℝ⁴: format + 3 F1 axes | uniform mean |
Results: search budget amplifies VPO
MuSiQue
| Method | best@3 | best@5 | best@10 | best@30 | diversity | |---|---|---|---|---|---| | GRPO | 0.711 | 0.716 | 0.721 | 0.728 | 0.054 | | VPO | 0.742 | 0.780 | 0.809 | 0.832 | 0.587 |Diversity jumps ~10×; VPO's best@3 already exceeds GRPO's best@30.
Maze
| Method | best@3 | best@5 | best@10 | best@30 | diversity | |---|---|---|---|---|---| | GRPO | 0.432 | 0.432 | 0.432 | 0.432 | 0.003 | | VPO | 0.512 | 0.564 | 0.591 | 0.593 | 1.006 |GRPO plateaus completely; VPO nearly doubles best@30 performance.
EUREQA (causal-chain extraction)
VPO at best@30 reaches 0.279 vs GRPO 0.236, with sustained growth as k increases.ToolRL (near-ceiling task)
All methods exceed 0.9; VPO catches up by k=10 and wins on diversity (1.297 vs 0.044).Ablations
LiveCodeBench with evolutionary search
On the 32 hardest problems where both GRPO and VPO score 0 at best@30, running OpenEvolve for 200 rounds:Diversity is shown to be a prerequisite for evolutionary recombination and mutation.
When VPO does not help
When reward dimensions are near-collinear (the simplex collapses to a line). On UltraFeedback + ArmoRM-5 (5 highly correlated dimensions), VPO underperforms scalar baselines. Real multi-objective structure is required.
Takeaway
Scalar RL trains a model to answer "what is the best answer?"; vector RL trains it to answer "here is a set that covers every reasonable definition of good." As search ecosystems (AlphaEvolve, Best-of-N, MCTS) become standard, diversity at training time is no longer optional. VPO achieves this with three simple ingredients: Dirichlet sampling on the simplex, multi-answer chains, and a drop-in GRPO replacement — no architecture changes required.