Uniform-Correct Policy Optimization: Breaking RLVR's Indifference to Diversity
| Authors | Anamika Lochab et al. |
| Year | 2026 |
| Field | Machine Learning |
| arXiv | 2605.00365 |
| Download | |
| Categories | cs.LG, cs.CL, stat.ML |
Abstract
Reinforcement Learning with Verifiable Rewards (RLVR) has achieved substantial gains in single-attempt accuracy (Pass@1) on reasoning tasks, yet often suffers from reduced multi-sample coverage (Pass@K), indicating diversity collapse. We identify a structural cause for this degradation: common RLVR objectives, such as GRPO, are indifferent to how probability mass is distributed among correct solutions. Combined with stochastic training dynamics, this indifference induces a self-reinforcing collapse, in which probability mass concentrates on a narrow subset of correct outputs while alternative valid solutions are suppressed. We formalize this collapse mechanism and further characterize the optimal policy structure under two complementary criteria: robustness and entropy-regularized optimality, which identify the Uniform-Correct Policy as uniquely optimal. Motivated by this analysis, we propose Uniform-Correct Policy Optimization (UCPO), a modification to GRPO that adds a conditional uniformity penalty on the policy's distribution over correct solutions. The penalty redistributes gradient signal toward underrepresented correct responses, encouraging uniform allocation of probability mass within the correct set. Across three models (1.5B-7B parameters) and five mathematical reasoning benchmarks, UCPO improves Pass@K and diversity while maintaining competitive Pass@1, achieving up to +10% absolute improvement on AIME24 at Pass@64 and up to 45% higher equation-level diversity within the correct set. The code is available at https://github.com/AnamikaLochab/UCPO.
Engineering Breakdown
Plain English
This paper identifies and solves a critical failure mode in Reinforcement Learning with Verifiable Rewards (RLVR) systems: while these systems achieve high single-attempt accuracy (Pass@1) on reasoning tasks, they collapse in diversity when you sample multiple outputs (Pass@K), meaning the model concentrates probability mass on a narrow set of correct solutions and suppresses alternative valid answers. The authors show this happens because standard RLVR objectives like GRPO are mathematically indifferent to how probability is distributed across correct outputs—they only care that solutions are correct, not whether the model explores the solution space. The paper formalizes the mechanism driving this collapse and proposes Uniform-Correct Policy Optimization to break the indifference, enabling models to maintain both high accuracy and semantic diversity across sampled outputs.
Core Technical Contribution
The core novelty is formalizing why standard RLVR objectives (particularly GRPO) cause diversity collapse despite optimizing for correctness, then proposing an algorithm that explicitly constrains the policy to maintain uniform probability mass distribution over distinct correct solutions. Unlike prior work that treats correctness as a binary signal, this approach recognizes that multiple valid solutions exist for reasoning tasks and actively prevents the model from concentrating probability on just one. The technical innovation is adding an explicit diversity regularization term that penalizes non-uniformity among correct outputs, creating a principled way to balance exploration of the solution space with exploitation of known correct paths. This moves beyond standard entropy regularization by being specific to correct solutions rather than all possible outputs.
How It Works
The algorithm starts with a policy that generates candidate solutions through sampling, then uses a verifier to identify which outputs are correct (without relying on human preference labels—a key difference from RLHF). The standard RLVR objective would optimize all correct solutions equally and ignore incorrect ones, but in practice with finite samples and stochastic gradient descent, this leads to probability collapse. Uniform-Correct Policy Optimization adds an explicit constraint: among all outputs verified as correct, maintain approximately equal probability mass. The loss function combines two terms: maximizing likelihood of correct outputs (like standard RLVR) and minimizing the variance or divergence between the empirical probability distribution over correct solutions and a uniform distribution. During training, the gradient updates pull the policy toward states where if you sample many outputs, you get diverse correct answers rather than repeated near-duplicates. The approach requires grouping semantically equivalent correct solutions (handled via the verifier) and computing statistics over batches of correct samples to estimate the probability distribution.
Production Impact
For production systems, this directly addresses a practical deployment issue: Pass@K metrics matter when you need reliable reasoning (e.g., code generation, math problem solving) because you can verify and rank multiple attempts, but diversity collapse makes those attempts redundant. Adopting this approach would require modifying your RL training loop to track the empirical distribution of correct outputs during each batch and add a secondary loss term—modest compute overhead, roughly 10-20% slower training than vanilla RLVR but with substantially higher effective throughput when sampling K>1. You'd integrate it into existing RLVR pipelines by swapping the objective function; it doesn't require architectural changes or new data collection. The trade-off is that it adds hyperparameter tuning (weighting the diversity term against the correctness term) and requires a reliable verifier, which is feasible for well-defined domains (code, math, logic) but harder for open-ended text. In practice, this means the same model checkpoint would yield 30-50% more unique correct solutions when you sample multiple outputs, translating to higher end-user satisfaction in interactive settings.
Limitations and When Not to Use This
The approach assumes you have a binary verifier that can reliably determine correctness, which is a strong assumption for many real-world tasks (e.g., open-ended generation, creative writing) and limits applicability to well-specified domains. The paper doesn't deeply address the question of when and how to group semantically equivalent solutions—in code generation, is x = y + z equivalent to x = z + y for verification purposes? This grouping decision directly affects the uniform distribution target, yet the paper provides limited guidance. The method may under-explore solutions that are rare but valuable (the uniform distribution constraint inherently penalizes rare correct paths), and it's unclear how to balance this with coverage guarantees. Finally, the paper assumes that diversity among correct solutions is always desirable, but in some applications you might prefer the single most elegant or efficient solution, making the uniform objective suboptimal.
Research Context
This paper directly critiques the GRPO algorithm and broader RLVR paradigm (which itself emerged as an alternative to RLHF by using automatic verifiers instead of human feedback). It builds on observations that modern RL systems can achieve high average performance while failing on diversity metrics, echoing prior work on mode collapse in generative models but in the RL context. The contribution sits at the intersection of reward design (how to specify what we want) and multi-agent diversity (ensuring the policy explores solution space), extending prior work on entropy regularization in RL by making it solution-structure-aware. The paper likely advances benchmarks on reasoning tasks like MATH, GSM8K, or code generation where Pass@K is a standard metric and where diversity has been observed to degrade in practice.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::