Reinforcement Learning with Promising Tokens for Large Language Models

TL;DR

This paper introduces RLPT, a reinforcement learning framework that improves large language model training by focusing on promising tokens, reducing action space, and enhancing efficiency and stability.

Contribution

RLPT leverages semantic priors to dynamically select promising tokens, decoupling decision-making from token generation, and improves training stability and efficiency.

Findings

RLPT reduces gradient variance and stabilizes training.

RLPT outperforms standard RL baselines on reasoning tasks.

Effective across different model sizes and RL algorithms.

Abstract

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 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).