Technical Summary: SAGE – Shaping Anchors for Guided Exploration in RLVR of LLMs

1. Problem Statement

Recent advancements in Reinforcement Learning with Verifiable Rewards (RLVR) have demonstrated the ability to improve the pass@1 accuracy of Large Language Models (LLMs) on reasoning tasks. However, a critical limitation has emerged: while RLVR often boosts single-sample success rates, it frequently fails to yield comparable gains in pass@k (the probability of finding a correct answer among $k$ samples).

Prior analyses suggest this discrepancy arises because standard RLVR induces mode collapse (or "support shrinkage"). Instead of discovering novel reasoning strategies, the policy collapses onto a narrow subset of pre-existing, high-density reasoning modes present in the base reference model. This behavior is attributed to the structural constraints of the standard reverse-Kullback-Leibler (reverse-KL) regularization term. While reverse-KL provides necessary training stability by anchoring the policy to a reference distribution ($\pi_{ref}$), it inherently suppresses low-density trajectories, effectively preventing the model from exploring alternative, valid reasoning paths that lie outside the immediate support of the reference model.

Attempts to resolve this by removing the KL term or replacing it with forward-KL have proven suboptimal. Removing the KL term often leads to reward hacking and instability, while forward-KL, though expanding support, frequently allocates probability mass to reward-irrelevant regions, failing to improve pass@1. The central challenge is to retain the stabilizing benefits of reverse-KL while enabling controlled empirical support expansion.

2. Methodology: SAGE Framework

The authors propose SAGE (Shaping Anchors for Guided Exploration), a framework that reinterprets the reverse-KL regularization not as a static constraint to be weakened, but as a dynamic anchor whose target distribution can be reshaped to guide exploration.

Core Mechanism

Standard reverse-KL RLVR optimizes:
$$J(\theta) = \mathbb{E}_{x,y \sim \pi_\theta} [r(x, y) - \beta D_{KL}(\pi_\theta \parallel \pi_{ref})]$$
This leads to a stationary solution where the policy is an exponential tilt of the reference: $\pi_\theta \propto \pi_{ref} \exp(r/\beta)$. This multiplicative update amplifies high-density regions of $\pi_{ref}$ while suppressing low-density ones.

SAGE modifies this by introducing a guide function $q(x, y) > 0$ that reshapes the anchor distribution itself. The new objective becomes:
$$J'(\theta) = \mathbb{E}_{x,y \sim \pi_\theta} [r(x, y) - \beta \tilde{D}_{KL}(\pi_\theta \parallel q \cdot \pi_{ref})]$$
Here, $\tilde{D}_{KL}$ is a pseudo-KL divergence defined for unnormalized functions. The stationary solution becomes:
$$\pi_\theta(y | x) \propto q(x, y) \pi_{ref}(y | x) \exp\left(\frac{1}{\beta} r(x, y)\right)$$
By designing $q(x, y)$ to be large for underexplored but potentially valid reasoning modes, SAGE amplifies their probability mass without breaking the stability provided by the anchor $\pi_{ref}$.

Instantiation via Intrinsic Signals

To make $q(x, y)$ practical without requiring explicit reward information during the guide construction, the authors propose using intrinsic signals derived from the reasoning process itself:

1. Random Exploration: A baseline using stochastic perturbations.
2. Token-Level Exploration: Uses surprisal ($-\log p$) to bias toward rare lexical choices. The authors note this captures surface-level deviations but may not reflect structural reasoning differences.
3. Branch-Level Exploration (SAGE-Branch): The primary proposed variant. It uses entropy ($H_t$) as a signal for reasoning branching. The guide function is defined as:
   $$q_{Branch}(y_t | x, y_{<t}) = 1 + \gamma (H_t - \tau)_+$$
   where $H_t$ is the token-level entropy, $\tau$ is an activation threshold, and $\gamma$ controls the magnitude. This selectively amplifies exploration at high-uncertainty decision points (branching nodes) while suppressing noise in deterministic symbolic steps.

Theoretical Guarantees

The paper provides a theoretical analysis (Theorem 4.3) demonstrating that, unlike standard RLVR which is bounded by the empirical support of the reference model, SAGE can theoretically expand the empirical support. Specifically, if the guide function $q$ is sufficiently large for a reward-valid solution $y^*$ that is currently outside the empirical support of the reference policy, SAGE can assign it a non-negligible probability, effectively recovering "hidden" reasoning modes.

3. Key Contributions

1. Analysis of RLVR Limitations: The authors identify that the reverse-KL anchor, while stabilizing, is the primary structural cause of mode collapse in RLVR. They argue that simply removing or replacing the KL term is insufficient, as it disrupts the efficiency-coverage trade-off.
2. SAGE Framework: They propose a principled method to enable controllable empirical support expansion by reshaping the KL anchor distribution via a guide function $q(x, y)$, rather than removing the anchor.
3. Intrinsic Signal Integration: They demonstrate how lightweight, intrinsic signals (specifically entropy and surprisal) can be used to construct effective guide functions that require no additional forward passes or external data.
4. Empirical Validation: Extensive experiments show that SAGE variants, particularly the branch-level exploration, consistently improve both pass@1 and pass@k across diverse mathematical reasoning benchmarks.

4. Experimental Results

The authors evaluated SAGE on the AIME 2024-2025, AMC23, and MATH-500 datasets using the Qwen2.5-Math-7B-Base model and the GRPO algorithm.

• Performance Gains: The GRPO + Branch variant achieved the strongest average performance, improving pass@1 from 0.266 (vanilla GRPO) to 0.284 and pass@256 from 0.714 to 0.729. It outperformed baselines across all three datasets.
• Comparison with Alternatives:
  • GRPO w/o KL: Showed higher training rewards but significantly worse test-time performance (pass@1 and pass@256), confirming the risk of reward hacking without the KL anchor.
  • Forward-KL: Improved pass@256 slightly but failed to match the pass@1 gains of SAGE, indicating it allocates mass to off-target regions.
  • Other RL Algorithms: SAGE was successfully integrated with BNPO and DAPO, consistently resolving the trade-off where these methods typically improved pass@1 at the cost of pass@k.
• Reasoning Pattern Analysis: Qualitative analysis revealed that while standard RLVR reinforced dominant patterns like "constraint setup," SAGE significantly increased the frequency of proof by contradiction (from 0.8% in GRPO to 3.3% in GRPO+Branch), a rare but valid reasoning mode.
• Robustness: In out-of-distribution evaluations (Knights and Knaves benchmark), SAGE-trained models maintained performance under distribution shift, whereas standard GRPO models exhibited significant degradation, suggesting SAGE mitigates mode collapse.
• Generalization: The method was also effective on long-reasoning models (DeepSeek-R1-Distill-Qwen-7B) and smaller models (Qwen2.5-1.5B/3B), demonstrating robustness across different model scales and reasoning paradigms.

5. Significance and Claims

The paper claims that SAGE offers a practical and principled approach to mitigating mode sharpening in reverse-KL RLVR. By treating the KL regularization as a controllable instrument rather than a rigid constraint, SAGE enables LLMs to discover and utilize diverse reasoning trajectories that are valid but underrepresented in the base model.

The authors emphasize that SAGE is modular and orthogonal to existing techniques (such as advantage shaping or dynamic sampling), as it modifies the anchor distribution itself rather than the optimization structure. They position SAGE as a general framework for stable, KL-compatible empirical support expansion.

Limitations and Scope:
The authors modestly acknowledge that SAGE operates within the support of the fixed reference policy. It redistributes probability mass among trajectories that already have non-zero likelihood under $\pi_{ref}$; it cannot generate entirely new reasoning patterns that are impossible for the base model to produce. They note that while distillation from stronger teacher models can expand the reasoning boundary, SAGE complements this by enabling stable exploration within that boundary. The current instantiation relies on lightweight intrinsic signals, and the authors suggest that learning more expressive guide functions from stronger models is a promising direction for future work.