Technical Summary: Sharpness-Guided Group Relative Policy Optimization via Probability Shaping

Problem Statement

Reinforcement Learning with Verifiable Rewards (RLVR) has emerged as a primary driver for improving the reasoning capabilities of Large Language Models (LLMs) in domains such as mathematics, logic, and tool-augmented question answering. Group Relative Policy Optimization (GRPO) is the dominant optimizer in this setting due to its simplicity and lack of dependency on an explicit value function. However, current RLVR implementations often lack robust control over generalization. The authors identify that training can lead to "sharp" updates—large, unstable gradient steps—which theoretically correlates with poorer generalization performance. The core problem addressed is how to stabilize GRPO training and improve generalization without sacrificing the learning signal or introducing complex architectural changes.

Methodology: GRPO-SG

The paper proposes Sharpness-Guided GRPO (GRPO-SG), a token-weighted variant of GRPO designed to reduce the sharpness of optimization updates. The methodology is grounded in a theoretical generalization view and implemented via a confidence-based probability shaping mechanism.

Theoretical Foundation

The authors revisit GRPO through a robustness-based generalization lens. They establish that the generalization loss is upper-bounded by the sum of the empirical training loss and a sharpness surrogate measured by the gradient norm (Theorem 3.3).

• Key Insight: Smaller gradient norms correspond to less sharp updates, implying better generalization.
• Derivation: By analyzing the gradient of the empirical loss, the authors show that the gradient norm is influenced by token-level weights. Specifically, the product of a token weight ($w_{i,t}$) and the term $(1 - \pi_\theta(o_{i,t}))$ (representing uncertainty) directly impacts the magnitude of the gradient.

Algorithmic Implementation

GRPO-SG modifies the standard GRPO objective by introducing a token weight $w_{i,t}$ that acts as a confidence signal.

1. Confidence Signal: The weight is derived from the selected-token logit ($h_{i,t}$) produced by the current policy. This logit serves as a proxy for the model's confidence in the token it selected.
2. Probability Shaping: A monotone, confidence-aware mapping function $\omega(\cdot)$ is applied to the logit to generate the weight. The specific formulation (Eq. 20) is:
   $$w_{i,t} = \text{clip}\left(\alpha \cdot \left[\sigma\left(\frac{\text{sg}[h_{i,t}]}{\tau}\right) - \mu\right], L, U\right)$$
   • $\text{sg}[\cdot]$ is a stop-gradient operator to prevent the model from gaming the weights.
   • $\sigma$ is a sigmoid function, and $\tau, \mu, \alpha, L, U$ are hyperparameters controlling smoothness and bounds.
3. Mechanism:
   • High-Confidence Tokens: Tokens with high logits (high confidence) receive higher weights, preserving the learning signal on critical reasoning steps.
   • Low-Confidence Tokens: Tokens with low logits receive lower weights. Since low-confidence tokens often induce larger gradients (due to the $(1-\pi)$ term), downweighting them suppresses "sharp" updates that could destabilize training.
4. Integration: These weights are integrated directly into the GRPO surrogate objective, effectively downweighting tokens likely to cause overly large gradients while maintaining the standard PPO-style clipping and advantage normalization.

Key Contributions

1. Generalization View of GRPO: The paper provides a theoretical link between RLVR training and a robust generalization objective, characterizing generalization via a sharpness surrogate (gradient norm).
2. GRPO-SG Algorithm: It introduces a simple, effective modification to GRPO that uses token confidence (selected-token logits) to shape update magnitudes, reducing sharpness without requiring auxiliary models or complex reward shaping.
3. Empirical Validation: The method is validated across diverse RLVR benchmarks, demonstrating consistent improvements over standard GRPO and other token-level control baselines (e.g., 80/20, AR, Lopti).

Experimental Results

Experiments were conducted on three distinct RLVR settings: Mathematical Reasoning, Agentic Question Answering (tool-augmented), and Logic Puzzles (Knights & Knaves).

• Mathematical Reasoning: On benchmarks including OlympiadBench, Minerva, and MATH-500, GRPO-SG achieved average improvements of +2.7% to +6.0% over standard GRPO. For instance, on the DSR training recipe with Qwen2.5-7B, the average gain was 6.0%.
• Agentic QA: In knowledge-augmented QA tasks (using VT-Search), GRPO-SG nearly doubled the exact-match accuracy for the Qwen2.5-3B model (from 13.84 to 27.29) and provided significant gains on the Qwen3-4B model.
• Logic Puzzles: On the K&K Logic Puzzles benchmark, GRPO-SG improved the average accuracy of Qwen2.5-3B from 0.42 to 0.63 (a 50% relative gain) and Qwen2.5-7B from 0.75 to 0.91.
• Training Dynamics: Beyond accuracy, GRPO-SG exhibited smoother gradient-norm trajectories and fewer outlier spikes compared to GRPO (Figure 1), consistent with the theoretical goal of reducing sharpness. Training rewards were also higher and more stable.

Significance and Claims

The paper claims that GRPO-SG is a simple and effective generalization-oriented upgrade to GRPO for RLVR.

• Simplicity: The method requires no additional value models, human labels, or complex architectural changes; it only modifies the token weighting within the existing GRPO objective.
• Stability: By explicitly controlling update sharpness via token confidence, the method stabilizes the training process, leading to more robust reasoning capabilities.
• Generality: The improvements are consistent across different model sizes (3B to 7B+), different model families (Qwen, Mistral, LLaMA), and diverse reasoning domains (math, logic, agentic search).

The authors conclude that explicitly controlling update sharpness through confidence-based token weighting is a viable strategy to improve the stability and performance of RL-trained LLMs, offering a practical path to better generalization in verifiable reward settings.