Vector Policy Optimization: Training LLMs for Diversity to Unlock Test-Time Search

Paper: 2605.22817 Authors: Ryan Bahlous-Boldi, Isha Puri, Idan Shenfeld, Akarsh Kumar, Mehul Damani, Sebastian Risi, Omar Khattab, Zhang-Wei Hong, Pulkit Agrawal Categories: cs.LG, cs.AI, cs.CL, cs.NE

The Gap

Current LLM post-training uses scalar reward optimization (PPO, DPO, GRPO), which collapses the model’s output distribution into a narrow mode. This works fine if you only need one answer. But modern inference pipelines—pass@k sampling, best-of-n selection, evolutionary search like AlphaEvolve—need the model to generate diverse candidates that a downstream reward function can filter. The mismatch is stark: we train for convergence, then ask for exploration at test time.

The problem compounds when deployment environments differ from training. A code model might face new test suites, new coding styles, or new performance constraints. A chatbot might encounter user personas not seen during RLHF. Scalar RL bakes in one reward function’s preferences, leaving the model brittle when the reward changes.

Prior work on diversity (like DPP sampling or temperature tuning) treats it as a post-hoc inference trick. This paper asks: what if we train for diversity from the start?

Problem: Scalar RL → Low-entropy policy → Poor test-time search
            |
            v
Observation: Rewards are often vector-valued in practice
            |
            v
Method: VPO optimizes for diverse solutions across reward dimensions
            |
            v
Evidence: Beats GRPO on pass@k, best@k, evolutionary search
            |
            v
Conclusion: Diversity should be the default post-training objective

The Increment

One sentence: Before VPO, we trained LLMs to maximize a single scalar reward; after VPO, we train them to produce a portfolio of solutions that cover different trade-offs in a vector reward space.

Core Mechanism

VPO starts with the observation that most “scalar” rewards are actually aggregations of vector components. Code correctness is per-test-case pass/fail. Chatbot quality might be helpfulness + safety + conciseness. Instead of collapsing these into a weighted sum during training, VPO keeps them separate.

The algorithm modifies GRPO’s advantage estimation. Standard GRPO computes advantages using a scalar reward baseline. VPO instead projects each solution onto a Pareto front in the vector reward space, then assigns advantages based on how well each solution covers an underexplored region of that front. Solutions that specialize to different reward trade-offs get positive advantages; redundant solutions near the mode get penalized.

Concretely: sample N solutions per prompt, compute their vector rewards, identify the Pareto front, cluster solutions by which reward dimension they excel at, then weight the policy gradient to encourage one solution per cluster. The model learns to hedge—output one solution optimized for reward dimension 1, another for dimension 2, and so on.

Input prompt
     |
     v
Sample N solutions from policy
     |
     v
Compute vector reward for each: [r1, r2, ..., rk]
     |
     v
Identify Pareto front in reward space
     |
     v
Cluster solutions by reward specialization
     |
     v
Assign advantages: +1 for diverse specialists, -1 for redundant modes
     |
     v
Policy gradient step pushes toward diversity

Think of VPO like training a restaurant kitchen. Scalar RL is like telling every chef to make the dish the head judge likes best—you end up with five identical plates. VPO is like having five judges with different tastes (spicy, sweet, umami, texture, presentation) and training each chef to specialize in winning over one judge. At test time, when a new judge shows up, you’ve got a better chance one of your chefs’ dishes will appeal to them. The kitchen (model) learns to maintain a repertoire, not converge to a single recipe.

Key Concepts

• Vector Reward: Instead of collapsing multiple evaluation criteria into a single number (e.g., 0.5 ** correctness + 0.3 * efficiency + 0.2 * readability), keep them as a vector [correctness, efficiency, readability]. This preserves information about trade-offs. A solution might score [0.9, 0.4, 0.7]—high correctness, low efficiency. Another might be [0.7, 0.9, 0.5]. Scalar aggregation would rank them, but vector representation lets you see they’re solving different problems. VPO exploits this: train the model to produce both types of solutions, so at test time you can pick based on which dimension matters most.

• Pareto Front: Imagine plotting all solutions in reward space. The Pareto front is the set of solutions where you can’t improve one dimension without sacrificing another. If solution A scores [0.8, 0.6] and solution B scores [0.7, 0.7], B is on the Pareto front (balanced) but A might not be if there’s a solution C at [0.8, 0.7] that dominates it. VPO uses the Pareto front to identify which solutions are genuinely offering different trade-offs versus which are just worse versions of the same strategy.

• Advantage Shaping for Diversity: In standard RL, advantage = (this solution’s reward) - (average reward). Positive advantage → reinforce, negative → suppress. VPO modifies this: advantage = (how much this solution covers an underexplored part of the Pareto front). If the model already outputs three solutions optimized for reward dimension 1, the fourth one gets a negative advantage even if its absolute reward is high. This pushes the model toward spreading its probability mass across the reward space rather than piling it onto one mode.

Framework Shift

Before (Scalar RL):                  After (VPO):

Prompt                               Prompt
  |                                    |
  v                                    v
Policy outputs                       Policy outputs
distribution                         distribution
  |                                    |
  v                                    v
Sample N solutions                   Sample N solutions
  |                                    |
  v                                    v
Compute scalar reward                Compute vector reward
r = w1*r1 + w2*r2                    [r1, r2, r3, ...]
  |                                    |
  v                                    v
Advantage = r - baseline             Cluster by specialization
  |                                    |
  v                                    v
Gradient pushes toward               Advantage = diversity score
high-reward mode                       |
  |                                    v
  v                                  Gradient pushes toward
Low-entropy policy                   covering Pareto front
(one dominant solution)                |
                                       v
                                     High-entropy policy
                                     (portfolio of specialists)

From optimizing a single peak to cultivating a ridge—the core shift is treating diversity as the objective, not a side effect.

Expert Assessment

Problem choice: This is a real gap. The field has been scaling test-time compute (o1, AlphaCode, self-consistency decoding) while still training models with scalar RL that fights against the diversity those methods need. The timing is sharp—inference-time search is becoming standard, so this problem will only grow. Not manufactured.

Method maturity: VPO is conceptually clean but the implementation has moving parts. The Pareto front computation and clustering add overhead. The paper doesn’t deeply explore failure modes—what happens when reward dimensions are correlated, or when the vector structure is noisy? The core insight (train for diversity by preserving vector structure) is strong, but I’d want to see ablations on the clustering scheme and sensitivity to hyperparameters. It’s not brute force, but it’s also not a one-line change.

Experimental integrity: Baselines are fair—GRPO, PPO, DPO are the right comparisons. The tasks (code generation, instruction following, math, creative writing) cover different reward structures. The numbers are convincing, especially the evolutionary search results where GRPO models fail completely. One concern: the paper doesn’t report wall-clock training time or memory overhead. If VPO is 3x slower than GRPO, that’s a tax practitioners need to know about. Also, the “vector reward” framing assumes you have access to per-dimension scores, which isn’t always true (e.g., human preference data is often scalar).

Writing quality: The intro and method sections are crisp. The related work section is thin—doesn’t engage deeply with multi-objective RL or quality-diversity algorithms from evolutionary computation, which have been doing this for years. The results section front-loads the wins but buries the analysis of when VPO doesn’t help (low search budgets, highly correlated reward dimensions). Rewriting Section 5 to lead with failure modes and then show where VPO shines would make the claims more credible.

Verdict: weak accept — Addresses a real and growing problem with a principled method, strong empirical results, but needs more transparency on computational costs and failure modes.

Takeaways

For practitioners: If you’re building systems that use pass@k, best-of-n, or any form of test-time search, VPO’s core idea transfers immediately—train your model to produce diverse candidates, not just high-reward ones. Even if you can’t implement VPO exactly, you can approximate it: sample multiple solutions during training, penalize redundancy, reward coverage of different solution types.

For researchers: The vector reward framing is the real contribution. Most RL setups already have vector structure (per-example rewards, multi-task objectives, ensemble of reward models) but we collapse it into a scalar for convenience. VPO shows that preserving that structure during training pays off at test time. This generalizes beyond LLMs—any domain where deployment rewards differ from training rewards (robotics, recommendation systems, drug design) could benefit.

Specific technique: The advantage shaping scheme (positive advantage for Pareto-optimal solutions in underexplored regions, negative for redundant modes) is a drop-in replacement for standard advantage estimation. You can graft it onto PPO, DPO, or any policy gradient method. The clustering step is the key—without it, you’re just doing multi-objective RL, which can still collapse to a single compromise solution.