Technical Summary: Back to Basics: Revisiting Exploration in Reinforcement Learning for LLM Reasoning via Generative Probabilities

1. Problem Statement

Reinforcement Learning with Verifiable Rewards (RLVR) has become a dominant paradigm for enhancing reasoning in Large Language Models (LLMs), particularly for generating long Chain-of-Thought (CoT) paths. However, standard policy optimization methods, such as Group Relative Policy Optimization (GRPO), suffer from entropy collapse and mode collapse.

The core issue stems from the reward-weighted likelihood maximization objective. During optimization, this objective disproportionately amplifies the probability mass of high-reward (correct) trajectories. Consequently, the model converges to low-entropy policies that suppress valid alternative reasoning chains, leading to a severe lack of output diversity. While existing solutions like entropy regularization or dynamic clipping attempt to mitigate this, they often remain limited to local modifications within the reward maximization framework and fail to fundamentally improve the modeling of diverse reasoning paths.

2. Methodology: ProGRPO

The authors propose ProGRPO (Probabilistic based GRPO), a novel paradigm that re-examines the construction of the Advantage function from the perspective of generative probabilities. The core innovation is an Advantage Re-weighting Mechanism (ARM) designed to equilibrate confidence levels across all correct responses.

2.1. Advantage Re-weighting Mechanism (ARM)

Instead of modifying the reward signal directly (which could distort updates if all answers in a group are correct or incorrect), ProGRPO adjusts the Advantage ($\tilde{A}_i$) by incorporating internal model signals:

• Prompt Confidence ($c_\theta(q_i)$): The model's confidence in the input prompt, calculated as the exponential of the mean log-probability of low-probability tokens in the prompt.
• Answer Confidence ($c_\theta(o_i | q_i)$): The model's confidence in generating a specific answer $o_i$ for prompt $q_i$, similarly calculated over low-probability tokens in the response.

The re-weighted advantage is defined as:
$$\tilde{A}_i = A_i + \alpha (c_\theta(q_i) - c_\theta(o_i | q_i))$$
where $A_i$ is the standard GRPO advantage and $\alpha$ is a weighting coefficient. This mechanism dynamically reshapes the reward signal: it attenuates gradient updates for over-confident reasoning paths and redistributes probability mass toward under-explored correct solutions.

2.2. Low-Probability Token Length Normalization

The authors observe that applying length normalization over an entire sequence dilutes the reward signal with "trivial" high-confidence tokens. Predictive uncertainty is typically concentrated in a small fraction of generation steps (approx. 20%).

• Strategy: The method identifies a critical subset of tokens ($T^{low}$) representing the positions with the highest predictive uncertainty.
• Implementation: Confidence scores for both prompts and answers are calculated exclusively over this informative subset, preserving meaningful variations that reflect reasoning quality rather than surface-level certainty.

2.3. Objective Function

The final objective function integrates these re-weighted advantages into the standard GRPO loss, utilizing a clipped surrogate objective to ensure training stability while encouraging diverse exploration.

3. Key Contributions

• Methodological Innovation: The introduction of ProGRPO, which extends GRPO by incorporating a confidence-aware Advantage Re-weighting Mechanism (ARM). This allows for targeted exploration without compromising training stability.
• Broad Effectiveness: Validation across diverse benchmarks (Mathematics and Code Generation) using multiple model scales (Qwen2.5-7B/32B, DeepSeek-R1-Distill-Qwen-1.5B/7B). The method consistently outperforms mainstream baselines like GRPO and FlowRL.
• OOD Robustness: Demonstration of superior Out-of-Distribution (OOD) adaptability, maintaining robust performance on unseen data distributions (e.g., GPQA, MMLU-Pro).
• Significant Performance Gains: The method achieves a superior trade-off between exploration and exploitation, significantly enhancing both accuracy and output diversity.

4. Experimental Results

Experiments were conducted on mathematical reasoning (AIME, AMC, MATH, OlympiadBench) and code generation (LiveCodeBench, CodeForces, HumanEval+) benchmarks.

• Mathematical Reasoning:
  • On Qwen2.5-7B, ProGRPO outperforms GRPO by 5.7% in Pass@1 and 13.9% in Pass@32.
  • On Qwen2.5-32B, it improves Pass@1 by 4.8% over GRPO.
  • On DeepSeek-R1-Distill-Qwen-1.5B, Pass@1 improved from 49.4% (GRPO) to 58.3% (ProGRPO).
• Code Generation:
  • On CodeForces, ProGRPO achieved a rating of 1422.49, exceeding GRPO by nearly +180 points and FlowRL by +293 points.
  • On HumanEval+, it achieved an Avg@16 score of 84.01%, the best among compared methods.
• Diversity and Entropy:
  • Training entropy analysis shows that while GRPO leads to entropy collapse, ProGRPO exhibits a stable rise in entropy, indicating sustained exploration.
  • Diversity metrics (Self-BLEU, Semantic Cosine) on AIME 2024 show significantly lower redundancy in ProGRPO outputs compared to GRPO, confirming the generation of structurally diverse reasoning paths.
• Ablation Studies:
  • The hyperparameter $\alpha$ (weighting coefficient) was found to be critical. An $\alpha$ of 0.3 provided the optimal balance between performance and stability. Values of 0 (pure GRPO) yielded weaker results, while higher values (0.7, 1.0) degraded performance by over-emphasizing confidence signals.

5. Significance and Claims

The paper claims that ProGRPO addresses the fundamental Exploration-Exploitation trade-off in LLM reasoning tasks. By shifting the focus from simple reward maximization to balancing generative probabilities and confidence, the method:

1. Mitigates Entropy Collapse: Prevents the model from converging to a single dominant reasoning path.
2. Enhances Diversity: Generates a wider variety of valid reasoning chains, evidenced by significant gains in Pass@k metrics (where $k > 1$).
3. Improves Generalization: Maintains robust performance on OOD datasets, suggesting the learned reasoning capabilities are more transferable.

The authors position this work as a "back to basics" approach, revisiting the probabilistic nature of LLM generation to solve the limitations of current RLVR paradigms, paving the way for more stable and exploratory reinforcement learning in reasoning models.