SPS: Steering Probability Squeezing for Better Exploration in Reinforcement Learning for Large Language Models

1. Problem Statement

Reinforcement Learning (RL), particularly with verifiable rewards (RLVR), has become a dominant paradigm for training Large Language Models (LLMs) on reasoning tasks. However, current RL methods face a critical limitation: insufficient exploration.

• The Phenomenon: While RL effectively improves single-sample success rates (Pass@1), it often fails to significantly improve multi-sample performance (Pass@k), which relies on the model's ability to generate diverse, correct reasoning trajectories.
• The Root Cause (Squeezing Effect): The authors identify a fundamental "probability squeezing effect." In standard RL (e.g., GRPO), negative gradient updates applied to low-probability (incorrect) tokens do not redistribute probability mass to diverse alternatives. Instead, due to the normalization properties of the softmax function, the removed probability mass is disproportionately absorbed by the already dominant (greedy) tokens.
• Consequence: This leads to a systematic sharpening of the output distribution, causing "entropy collapse." The model converges prematurely to a narrow set of high-reward trajectories, restricting its ability to discover novel or alternative correct reasoning paths.

2. Methodology: Steering Probability Squeezing (SPS)

To address the squeezing effect, the authors propose SPS, a training paradigm that interleaves conventional RL with Inverse Reinforcement Learning (IRL).

Core Mechanism

Instead of allowing probability mass to collapse onto dominant trajectories, SPS actively "steers" this redistribution toward under-explored but potentially valuable regions of the trajectory space.

1. Interleaved Training Loop:

   • Phase 1 (RL): The model undergoes standard on-policy RL (e.g., GRPO) to generate rollouts and identify high-reward trajectories.
   • Phase 2 (IRL): The model uses the generated rollouts as "demonstrations" to perform Inverse Reinforcement Learning.
     • Objective: The IRL stage employs a forward-KL divergence loss to reshape the policy distribution. It encourages the policy to align with the empirical distribution of the rollouts ($\pi_{rollout}$), effectively redistributing probability mass away from overly dominant modes toward the diverse set of sampled trajectories.
     • Self-Supervision: Crucially, this process uses only on-policy rollouts as demonstrations. No external expert data, annotations, or prior knowledge is introduced.

2. Low-Likelihood Trajectory Emphasis (L2TE):

   • To further counteract squeezing, the IRL phase specifically prioritizes sampling rollouts from trajectories with relatively low model likelihood.
   • By amplifying the learning signal from these under-explored solutions, L2TE prevents the policy from collapsing into a single greedy mode.

3. Iterative Strategy:

   • The RL and IRL phases are repeated in a loop. The updated policy from the IRL phase is fed back into the next RL phase, enabling progressive redistribution of probability mass and sustained exploration throughout training.

3. Key Contributions

1. Theoretical Analysis of Squeezing: The authors provide a preliminary analysis and theoretical proof (via softmax normalization properties) demonstrating that standard RL objectives inherently induce a "squeezing effect," leading to systematic distributional sharpening and limiting exploration. They also identify an empirical upper bound on Pass@k performance under standard RL.
2. SPS Framework: They propose the SPS algorithm, which uniquely integrates IRL into the RL training loop to explicitly reshape trajectory distributions without external supervision.
3. Iterative Training Strategy: They introduce a mechanism to alternate between RL and IRL, ensuring that the model continues to explore alternative reasoning paths rather than converging prematurely.
4. Comprehensive Evaluation: Extensive experiments on Olympiad-level mathematical benchmarks demonstrate that SPS significantly outperforms state-of-the-art baselines in multi-sample exploration metrics.

4. Experimental Results

The method was evaluated on five Olympiad-level mathematical benchmarks (AIME-24/25, BRUMO, HMMT-FEB/NOV) using Qwen2.5-Math models (1.5B and 7B parameters) trained on subsets of the OpenR1-Math dataset.

• Pass@k Improvement: SPS consistently achieves superior Pass@128 scores compared to vanilla GRPO, DAPO, and GSPO.
  • Notable Result: On the Qwen2.5-Math-1.5B model, SPS achieved a Pass@128 score of 63.33 on the BRUMO benchmark, a +10.00 point improvement over vanilla GRPO.
• Model Scale Sensitivity: The authors observed that standard GRPO sometimes reduces Pass@128 on smaller models (1.5B) due to overfitting and squeezing, whereas SPS successfully mitigates this, improving performance across both 1.5B and 7B models.
• Diversity Metrics: Analysis of trajectory similarity shows that SPS generates more diverse reasoning paths (lower similarity scores) compared to GRPO, confirming that the performance gains stem from genuine exploration rather than just sampling efficiency.
• Efficiency: The IRL stage adds negligible computational overhead (approx. 3% of total training time per iteration) while providing significant performance gains.

5. Significance and Implications

• Redefining RL Exploration: The paper challenges the assumption that RL naturally promotes exploration. It highlights that without intervention, RL tends to collapse distributions, and proposes IRL as a principled mechanism to counteract this.
• Self-Supervised Distribution Reshaping: SPS demonstrates that models can effectively "steer" their own probability distributions using their own generated data, eliminating the need for expensive external reward models or expert demonstrations to maintain diversity.
• Scalability for Reasoning: By breaking the "squeezing" bottleneck, SPS offers a pathway to unlock the full potential of large-scale sampling (Pass@k) in reasoning tasks, which is crucial for complex problem-solving where multiple attempts are required to find a solution.
• Theoretical Insight: The work provides a deeper understanding of the learning dynamics in RLVR, offering a theoretical upper bound on exploration and explaining why entropy collapse occurs in practice.

In conclusion, SPS represents a significant step forward in aligning RL training dynamics with the goal of diverse exploration, enabling LLMs to discover a broader range of correct reasoning trajectories in complex domains like mathematics.