Token-Regulated Group Relative Policy Optimization for Stable Reinforcement Learning in Large Language Models

2025-10-29

AI Summary

In GRPO, excessively large gradients for low-probability tokens destabilize training and diminish the learning contribution of high-probability tokens. To address this, we propose Token-Level Weighted GRPO (TLW-GRPO), which introduces a token-level dynamic weighting mechanism within the GRPO framework—where weights are positively correlated with predicted token probabilities. This explicitly suppresses gradient amplification for low-probability tokens and strengthens the influence of high-confidence predictions on policy updates. TLW-GRPO integrates verifiable reward reinforcement learning without requiring auxiliary models or human annotations. Empirical evaluation across logical reasoning, mathematical problem solving, and agent-based tasks demonstrates that TLW-GRPO consistently outperforms baseline methods, significantly improving both training stability and generalization performance. These results substantiate the critical role of token-level gradient balancing in enhancing the reasoning capabilities of large language models.

Technology Category

• Natural Language Processing: (Large) Language Models
• Machine Learning: Reinforcement Learning
• Multiagent Systems: Multiagent Learning

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

Reinforcement learning with verifiable rewards (RLVR) has emerged as a powerful approach for strengthening the reasoning capabilities of large language models (LLMs). Among existing algorithms, Group Relative Policy Optimization (GRPO) has demonstrated strong performance, yet it suffers from a critical issue: low-probability tokens disproportionately dominate gradient updates due to their inherently large gradient magnitudes. This imbalance leads to unstable training and suppresses the contribution of high-probability tokens that are more reliable for learning. In this work, we introduce Token-Regulated Group Relative Policy Optimization (TR-GRPO), a simple yet effective extension of GRPO that assigns token-level weights positively correlated with the model's predicted probability. By downweighting low-probability tokens and emphasizing high-probability ones, TR-GRPO mitigates gradient over-amplification while preserving informative learning signals. Extensive experiments demonstrate that TR-GRPO consistently outperforms GRPO across RLVR tasks, including logic, math, and agentic reasoning, highlighting the importance of regulating token contributions during RL training and establishing TR-GRPO as a robust framework for enhancing LLM reasoning.

Problem

• Addresses unstable training from low-probability token gradients
• Mitigates gradient imbalance suppressing high-probability token contributions
• Enhances reinforcement learning stability for large language models

Innovation

• Token-level weights regulate gradient updates
• Downweight low-probability tokens to stabilize training
• Emphasize high-probability tokens for reliable learning