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

TL;DR

This paper addresses the issue of low-probability tokens dominating gradients in RL training of LLMs, proposing methods to balance token influence and improve reasoning performance.

Contribution

It introduces Advantage Reweighting and Lopti, two techniques to mitigate low-probability token dominance in RL training of LLMs, enhancing learning efficiency.

Findings

Up to 46.2% improvement in reasoning tasks

Effective attenuation of low-probability token gradients

Enhanced balanced token updates in RL training

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.