Uniform-Correct Policy Optimization: Breaking RLVR's Indifference to Diversity

1. Problem Statement

Reinforcement Learning with Verifiable Rewards (RLVR) has significantly improved the one-shot accuracy (Pass@1) of Large Language Models (LLMs) on logical reasoning tasks such as mathematics. However, a critical side effect has been identified: Diversity Collapse.

• The Phenomenon: While RLVR increases the probability of generating one correct answer, it often drastically reduces Pass@K (the probability that at least one of $K$ samples is correct). This indicates that the model's solution space is contracting, concentrating probability mass on a narrow subset of correct solutions while suppressing other valid reasoning paths.
• The Cause: The authors identify a structural limitation in standard RLVR objective functions (such as GRPO - Group Relative Policy Optimization). These objectives are indifferent to the distribution of probability mass within the set of correct solutions.
  • Mathematically, any policy that places 100% of its mass on the set of correct outputs $Y^+$ yields the same maximum reward value, regardless of whether this mass is concentrated on a single solution (collapsed) or evenly distributed across all valid solutions (diverse).
  • Combined with on-policy sampling dynamics, this indifference creates a self-reinforcing feedback loop: solutions sampled slightly more frequently receive larger gradient updates, further increasing their probability and sampling frequency, leading to irreversible mode collapse.

2. Methodology

A. Theoretical Analysis

The authors formalize the collapse mechanism and define the optimal policy structure:

1. Collapse Mechanism: They prove that under GRPO, sampling biases lead to unequal sample counts. Gradient updates scale with these counts, reinforcing the probability of dominant solutions and suppressing minority solutions. Once a solution falls below a sampling threshold ($\pi(y) \ll 1/K$), it receives no gradient signal and cannot recover.
2. Definition of the Optimal Policy: The authors propose that the optimal policy should be the Uniform-Correct Policy (UCP). This policy satisfies two criteria:
   • Robustness: It maximizes the remaining probability mass in the worst case when a subset of correct solutions is disabled (e.g., due to distribution shifts).
   • Entropy-Regularized Optimality: It uniquely maximizes the objective of aligning correctness with diversity (entropy) within the correct set.
   • Definition: The UCP assigns zero probability to incorrect outputs and distributes probability uniformly across all correct solutions ($\pi(y|x) = 1/|Y^+|$ for $y \in Y^+$).

B. Proposed Solution: Uniform-Correct Policy Optimization (UCPO)

To achieve the Uniform-Correct Policy, the authors propose UCPO, a modification of the GRPO objective function.

• Objective Function: UCPO extends the standard maximization of correctness by adding a conditional uniformity penalty.
  $$\mathcal{L}_{UCPO} = \underbrace{\log Z_\theta(x)}_{\text{Correctness Mass}} - \tau \underbrace{KL(u(\cdot|x) \parallel q_\theta(\cdot|x))}_{\text{Uniformity Penalty}}$$

  • $Z_\theta(x)$: Total probability mass on correct outputs.
  • $q_\theta(\cdot|x)$: The policy's conditional distribution over correct outputs.
  • $u(\cdot|x)$: The uniform distribution over the set of correct outputs.
  • $\tau$: A hyperparameter controlling the strength of the uniformity penalty.

• Gradient Mechanism:

  • Standard GRPO weights gradients based on current sampling frequency (favoring dominant modes).
  • UCPO reweights the gradient signal to compensate for this bias. It interpolates between the current conditional distribution and the uniform distribution target.
  • Inverse Weighting: The gradient estimator assigns higher weights to underrepresented correct solutions (those with low current probability) and lower weights to dominant solutions. This effectively reallocates the "advantage budget" to reinforce minority-correct answers without sacrificing total correctness mass.

3. Main Contributions

1. Identification of Indifference: The paper formally proves that standard RLVR objectives are indifferent to diversity within the correct set, allowing stochastic training dynamics to drive irreversible mode collapse.
2. Theoretical Characterization: It establishes the Uniform-Correct Policy as the unique optimal solution under both robustness and entropy-regularized criteria.
3. Algorithm Design (UCPO): It derives UCPO, a simple yet effective modification of GRPO that enforces the Uniform-Correct Policy as its unique optimum.
4. Gradient Reallocation: It introduces a mechanism to redirect gradient signals to underrepresented correct answers, breaking the self-reinforcing collapse cycle.

4. Experimental Results

The authors evaluated UCPO across three models (1.5B, 7B parameters) and five mathematical reasoning benchmarks (AIME24/25, AMC, MATH, OlympiadBench).

• Pass@K Improvement: UCPO consistently improves Pass@K (especially Pass@64) while maintaining competitive Pass@1.
  • AIME24 (7B Model): Achieved an absolute improvement of +10% in Pass@64 compared to standard GRPO.
  • Average Improvement: Across all models, UCPO improved average Pass@64 by +3.5 to +4.8 points over GRPO.
• Diversity Metrics:
  • UCPO achieved up to 45% higher equation-level diversity within the correct set compared to GRPO.
  • Unlike global entropy regularization (which often spreads mass to incorrect answers), UCPO maintains high correctness mass ($Z_\theta \approx 1$) while diversifying the correct set.
• Comparison with Baselines: UCPO outperformed other diversity-preserving methods such as Entropy Regularization (Ent-Reg), KL Covariance Control, and Focal GRPO. Global entropy regularization proved fragile, often degrading Pass@1 or failing to improve diversity within the correct set.
• Efficiency: UCPO introduces negligible computational overhead (only scalar reweighting of advantages) and requires no additional rollouts or reference models.

5. Significance

This work addresses a fundamental flaw in current RLVR training paradigms for reasoning models. By shifting the focus from merely "finding one correct answer" to "maintaining a diverse distribution of all correct answers," UCPO enables:

• Prevention of Solution Space Contraction: It ensures models do not "forget" valid reasoning paths during fine-tuning.
• Increased Robustness: A policy covering multiple correct solutions is more robust to distribution shifts or noise in the verifier.
• Provision of a Principled Objective: It moves beyond heuristic entropy bonuses to a theoretically grounded objective that explicitly addresses the structure of the optimal policy.

In summary, UCPO offers a scalable, theoretically grounded solution to the problem of diversity collapse in RLVR, enabling LLMs to achieve higher multi-sample coverage without sacrificing one-shot accuracy.