English

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

Computation and Language 2025-05-20 v1 Artificial Intelligence Machine Learning

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.

Keywords

Cite

@article{arxiv.2505.12929,
  title  = {Do Not Let Low-Probability Tokens Over-Dominate in RL for LLMs},
  author = {Zhihe Yang and Xufang Luo and Zilong Wang and Dongqi Han and Zhiyuan He and Dongsheng Li and Yunjian Xu},
  journal= {arXiv preprint arXiv:2505.12929},
  year   = {2025}
}

Comments

24 pages, 12 figures