Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Reinforcement Learning with Promising Tokens for Large Language Models

Jing-Cheng Pang, Liang Lu, Xian Tang, Kun Jiang, Sijie Wu, Kai Zhang, Xubin Li

TL;DR

RLPT addresses the large action space in RL for LLMs by constraining policy optimization to a dynamic set of promising tokens identified from the base model's priors. This token-masking approach reduces gradient variance and stabilizes training, yielding better sample efficiency while preserving linguistic coherence. Empirical results across mathematical reasoning, coding, and telecom tasks show consistent gains over strong baselines and across model sizes, with theoretical justification for variance reduction. The method offers a practical, model-agnostic pathway to more efficient RL for LLMs and invites future work on adaptive, context-driven token selection.

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

Reinforcement Learning with Promising Tokens for Large Language Models

TL;DR

RLPT addresses the large action space in RL for LLMs by constraining policy optimization to a dynamic set of promising tokens identified from the base model's priors. This token-masking approach reduces gradient variance and stabilizes training, yielding better sample efficiency while preserving linguistic coherence. Empirical results across mathematical reasoning, coding, and telecom tasks show consistent gains over strong baselines and across model sizes, with theoretical justification for variance reduction. The method offers a practical, model-agnostic pathway to more efficient RL for LLMs and invites future work on adaptive, context-driven token selection.

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 \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).
Paper Structure (32 sections, 1 theorem, 13 equations, 7 figures, 7 tables)

This paper contains 32 sections, 1 theorem, 13 equations, 7 figures, 7 tables.

Table of Contents

  1. Introduction
  2. Preliminary
  3. Analyzing the Potential for Decision Making over Promising Tokens
  4. Analysis of Ground-truth Solution
  5. Analysis of Model-generated Solution
  6. Method
  7. Promising Token Construction
  8. Policy Rollout and Optimization with Promising Token Masking
  9. Policy Rollout.
  10. Policy Optimization.
  11. Integration with RL Algorithms
  12. Theoretical Justification
  13. Related Work
  14. Reinforcement Learning for LLM Training
  15. Decoding-time Action Pruning Strategies
  16. ...and 17 more sections

Key Result

Proposition 4.1

Assuming the advantage $A_t$ is bounded, optimizing the policy over the constrained space $\mathcal{P}_t$ strictly reduces the variance of the gradient estimator associated with the tail tokens $\mathcal{T}$, compared to optimization over the full vocabulary $\mathcal{V}$.

Figures (7)

  • Figure 1: An illustration of decision making with promising tokens. At each step, the policy selects tokens solely from a high-likelihood subset, enabling it to focus on strategic decision-making.
  • Figure 2: The overall framework of RLPT method.
  • Figure 3: Training curves on Math-17k dataset.
  • Figure 4: Performance of RLPT method on different sizes of models.
  • Figure 5: Gradient norm curves of GRPO and GRPO+RLPT during the training process.
  • ...and 2 more figures

Theorems & Definitions (3)

  • Definition 3.1: Promising Tokens
  • Proposition 4.1: Variance Reduction
  • proof