Reinforcement Learning with Promising Tokens for Large Language Models
This addresses a bottleneck in RL for LLMs by reducing action space noise, offering an incremental improvement for researchers and practitioners in AI alignment and optimization.
The paper tackles the problem of reinforcement learning for large language models being inefficient due to the large vocabulary action space, and introduces RLPT, which focuses on a subset of promising tokens to improve training stability and sample efficiency, achieving better performance on math, coding, and telecom reasoning tasks across model sizes and algorithms.
Reinforcement learning (RL) has emerged as a key paradigm for aligning and optimizing large language models (LLMs). Standard approaches treat the LLM as the policy and apply RL directly over the full vocabulary space. However, this formulation includes the massive tail of contextually irrelevant tokens in the action space, which could distract the policy from focusing on decision-making among the truly reasonable tokens. In this work, we verify that valid reasoning paths could inherently concentrate within a low-rank subspace. Based on this insight, we introduce Reinforcement Learning with Promising Tokens (RLPT), a framework that mitigates the action space issue by decoupling strategic decision-making from token generation. Specifically, RLPT leverages the semantic priors of the base model to identify a dynamic set of \emph{promising tokens} and constrains policy optimization exclusively to this refined subset via masking. Theoretical analysis and empirical results demonstrate that RLPT effectively reduces gradient variance, stabilizes the training process, and improves sample efficiency. Experiment results on math, coding, and telecom reasoning show that RLPT outperforms standard RL baselines and integrates effectively across various model sizes (4B and 8B) and RL algorithms (GRPO and DAPO).