STARE: Reponderação de Vantagem em Nível de Token Guiada por Surprisal para Estabilidade da Entropia da Política

Resumo

Algoritmos de Aprendizagem por Reforço com Recompensas Verificáveis, como o GRPO, emergiram como o paradigma dominante de pós-treinamento para raciocínio complexo em LLMs, mas comumente sofrem de colapso de entropia da política durante o treinamento. Realizamos uma análise de gradiente de primeira ordem da dinâmica de entropia em nível de tokens sob o GRPO e identificamos uma incompatibilidade de atribuição de crédito em nível de token: a variação de entropia por token se decompõe no produto da vantagem em nível de trajetória e uma função de sensibilidade à entropia sobre a distribuição do próximo token, resultando em uma estrutura de quatro quadrantes de vantagem-surpresa e uma propriedade de quase-criticalidade. Motivados por isso, propomos o STARE (Surprisal-guided Token-level Advantage Reweighting for policy Entropy stability - Rebalanceamento de Vantagem em Nível de Token Guiado por Surpresa para Estabilidade de Entropia da Política), que identifica subconjuntos críticos de tokens para entropia por meio de quantis de surpresa intra-lote, reajusta seletivamente suas vantagens efetivas e incorpora uma porta de malha fechada de entropia alvo para regulação estável da entropia. Em escalas de modelo de 1,5B a 32B e três famílias de tarefas (CoT Curto, CoT Longo e Uso de Ferramentas em Múltiplas Etapas), o STARE mantém treinamento RL estável ao longo de milhares de passos, enquanto mantém a entropia da política dentro da faixa alvo. No AIME24 e AIME25, o STARE supera o DAPO e outras linhas de base competitivas em 4%-8% de acurácia média, com tokens de reflexão e comprimento de resposta crescendo em conjunto, indicando um equilíbrio sustentado entre exploração e aproveitamento que desbloqueia ainda mais o potencial de treinamento RL. O código está disponível em https://github.com/hp-luo/STARE.

English

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.