STARE: Surprisal-Guided Token-Level Advantage Reweighting for Policy Entropy Stability

2026-06-17

AI Summary

This work addresses the instability and degraded exploration in large language models trained with reinforcement learning from reward validation, which often stems from policy entropy collapse. The study reveals, for the first time, a token-level credit assignment mismatch in the GRPO algorithm, identifying a four-quadrant structure and near-critical behavior between advantage estimates and surprise. Building on these insights, the authors propose STARE: a method that identifies critical tokens via intra-batch surprise quantiles, reweights their effective advantages, and incorporates a target-entropy closed-loop gating mechanism to stabilize policy entropy. Evaluated across model scales from 1.5B to 32B and three task categories, STARE enables stable training over thousands of steps, outperforming baselines such as DAPO by 4%–8% in average accuracy on AIME24/25, while jointly increasing both reflection depth and response length—demonstrating superior exploration-exploitation balance.

Abstract

Reinforcement Learning with Verifiable Rewards algorithms like GRPO have emerged as the dominant post-training paradigm for complex reasoning in LLMs, yet commonly suffer from policy entropy collapse during training. We conduct a first-order gradient analysis of token-level entropy dynamics under GRPO and identify a token-level credit assignment mismatch: the per-token entropy variation decomposes into the product of the trajectory-level advantage and an entropy sensitivity function over the next-token distribution, yielding an advantage-surprisal four-quadrant structure and a near-criticality property. Motivated by it, we propose STARE (Surprisal-guided Token-level Advantage Reweighting for policy Entropy stability), which identifies entropy-critical token subsets via batch-internal surprisal quantiles, selectively reweights their effective advantages, and incorporates a target-entropy closed-loop gate for stable entropy regulation. Across model scales from 1.5B to 32B and three task families (Short CoT, Long CoT, and Multi-Turn Tool Use), STARE sustains stable RL training over thousands of steps while maintaining policy entropy within the target band. On AIME24 and AIME25, STARE outperforms DAPO and other competitive baselines by 4%-8% in average accuracy, with reflection tokens and response length growing in tandem, indicating sustained exploration-exploitation balance that further unlocks RL training potential.Code is available at https://github.com/hp-luo/STARE.

Problem

• policy entropy collapse
• reinforcement learning
• credit assignment
• token-level entropy
• training stability

Innovation

• Surprisal-guided Reweighting
• Token-level Advantage
• Policy Entropy Stability
• GRPO
• Reinforcement Learning for LLMs