Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Do Not Let Low-Probability Tokens Over-Dominate in RL for LLMs

Zhihe Yang, Xufang Luo, Zilong Wang, Dongqi Han, Zhiyuan He, Dongsheng Li, Yunjian Xu

TL;DR

The paper identifies a gradient-dominance bias in RL training for LLMs where low-probability tokens disproportionately drive updates, hindering learning for high-probability tokens. It provides a theoretical bound showing token-gradient magnitudes scale with $(1-\pi_\theta(o))$ and introduces two practical mitigations: Advantage Reweighting and Lopti, which suppress low-prob gradients and re-order updates to favor high-prob tokens. Empirical results across K&K Logic Puzzles and math datasets demonstrate substantial performance gains (up to $46.2\%$ on challenging tasks) when combined, with ablations highlighting the importance of high-prob tokens, update order, and careful hyperparameter tuning. The work advances RL for LLMs by addressing gradient interference directly, enabling more balanced and efficient learning and broader applicability to policy-gradient methods.

Abstract

Reinforcement learning (RL) has become a cornerstone for enhancing the reasoning capabilities of large language models (LLMs), with recent innovations such as Group Relative Policy Optimization (GRPO) demonstrating exceptional effectiveness. In this study, we identify a critical yet underexplored issue in RL training: low-probability tokens disproportionately influence model updates due to their large gradient magnitudes. This dominance hinders the effective learning of high-probability tokens, whose gradients are essential for LLMs' performance but are substantially suppressed. To mitigate this interference, we propose two novel methods: Advantage Reweighting and Low-Probability Token Isolation (Lopti), both of which effectively attenuate gradients from low-probability tokens while emphasizing parameter updates driven by high-probability tokens. Our approaches promote balanced updates across tokens with varying probabilities, thereby enhancing the efficiency of RL training. Experimental results demonstrate that they substantially improve the performance of GRPO-trained LLMs, achieving up to a 46.2% improvement in K&K Logic Puzzle reasoning tasks. Our implementation is available at https://github.com/zhyang2226/AR-Lopti.

Do Not Let Low-Probability Tokens Over-Dominate in RL for LLMs

TL;DR

The paper identifies a gradient-dominance bias in RL training for LLMs where low-probability tokens disproportionately drive updates, hindering learning for high-probability tokens. It provides a theoretical bound showing token-gradient magnitudes scale with and introduces two practical mitigations: Advantage Reweighting and Lopti, which suppress low-prob gradients and re-order updates to favor high-prob tokens. Empirical results across K&K Logic Puzzles and math datasets demonstrate substantial performance gains (up to on challenging tasks) when combined, with ablations highlighting the importance of high-prob tokens, update order, and careful hyperparameter tuning. The work advances RL for LLMs by addressing gradient interference directly, enabling more balanced and efficient learning and broader applicability to policy-gradient methods.

Abstract

Reinforcement learning (RL) has become a cornerstone for enhancing the reasoning capabilities of large language models (LLMs), with recent innovations such as Group Relative Policy Optimization (GRPO) demonstrating exceptional effectiveness. In this study, we identify a critical yet underexplored issue in RL training: low-probability tokens disproportionately influence model updates due to their large gradient magnitudes. This dominance hinders the effective learning of high-probability tokens, whose gradients are essential for LLMs' performance but are substantially suppressed. To mitigate this interference, we propose two novel methods: Advantage Reweighting and Low-Probability Token Isolation (Lopti), both of which effectively attenuate gradients from low-probability tokens while emphasizing parameter updates driven by high-probability tokens. Our approaches promote balanced updates across tokens with varying probabilities, thereby enhancing the efficiency of RL training. Experimental results demonstrate that they substantially improve the performance of GRPO-trained LLMs, achieving up to a 46.2% improvement in K&K Logic Puzzle reasoning tasks. Our implementation is available at https://github.com/zhyang2226/AR-Lopti.
Paper Structure (32 sections, 1 theorem, 21 equations, 12 figures, 7 tables, 1 algorithm)

This paper contains 32 sections, 1 theorem, 21 equations, 12 figures, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminary
  4. Large Language Models.
  5. Group Relative Policy Optimization.
  6. Methodology
  7. Explanation on Low-Probability Tokens' Dominance
  8. Mitigating the Over-Dominance of Low-Probability Tokens
  9. Adverse Effect of the Dominance.
  10. Advantage Reweighting.
  11. Low-Probability Tokens Isolation (Lopti).
  12. Experimental Results
  13. Experiments on K&K Logic Puzzles
  14. Experiments on Math-related Datasets
  15. Ablation Studies
  16. ...and 17 more sections

Key Result

Proposition 4.2

Under Assumption assumption4.1, let $\delta_\ell(o_{i,t}):= \nabla_{a_\ell} J_{GRPO}(o_{i,t})$ denote the gradient of the GRPO objective with respect to activation $a_\ell$ at any layer for a single token $o_{i,t}$. Let $\Vert \cdot \Vert$ denote the spectral norm, and define the vocabulary size as

Figures (12)

  • Figure 1: Experimental analysis on the K&K Logic Puzzle dataset during GRPO training of Qwen2.5-7B-Instruct-1M. Tokens are divided into four groups based on probability quartiles. (a) Token probability distribution and (b) corresponding advantages. (c) Token probability changes after updates (using SGD with lr=1e-3) and (d) gradient norms for each probability group. Effects of selective updates: (e) Probability changes when only tokens in the lowest quartile (probability < 0.25) are updated, and (f) when only tokens in the highest quartile (probability > 0.75) are updated. To ensure clarity, the top 1% of outlier samples in the violin plots for token probability changes are excluded. Results are averaged over 10 randomly sampled batches.
  • Figure 2: Diagram of Proposition \ref{['proposition4.2']}.
  • Figure 3: The proportion of positive tokens updated in the correct direction for different updating methods, under the same experimental settings as in Figure \ref{['fig:motivation']}.
  • Figure 4: Experimental results on the K&K Logic Puzzles benchmark. For Advantage Reweight, $\alpha=0.3$, and for Lopti, $\eta=0.5$. The reward curve during training (left) is truncated to exclude the first epoch and smoothed with an exponential moving average (coefficient: 0.95). The evaluation accuracy on the test set (right) are averaged over the last three checkpoints to mitigate randomness.
  • Figure 5: (a) The relationship between the frequency of six categories of inference-related words and the corresponding sample rewards for Qwen-2.5-7B-Instruct-1M trained with naive GRPO. The Pearson correlation coefficient ($r$) and Spearman rank correlation coefficient ($\rho$) are annotated. (b) A comparison of the frequency of the six categories of words across the starting point (Qwen-2.5-7B-Instruct-1M), naive GRPO, and GRPO enhanced with Advantage Reweighting and/or Lopti.
  • ...and 7 more figures

Theorems & Definitions (1)

  • Proposition 4.2