解决大模型强化学习中低概率词主导更新的问题
Do Not Let Low-Probability Tokens Over-Dominate in RL for LLMs
- 通过重加权优势和隔离低概率词来抑制其梯度影响
- 使高概率词的梯度得到充分更新，提升模型性能
- 特别适合改进基于GRPO的推理类大模型训练
强化学习（RL）已成为提升大语言模型（LLM）推理能力的核心方法，近期提出的组相对策略优化（GRPO）展现出卓越效果。本研究发现，RL训练中一个关键但未被充分关注的问题：低概率词因梯度幅值过大而过度主导模型更新，从而抑制了对性能至关重要的高概率词的梯度。为此，我们提出两种新方法——优势重加权与低概率词隔离（Lopti），有效减弱低概率词的梯度影响，同时强化高概率词驱动的参数更新。该策略促进不同概率词间的平衡更新，显著提升RL训练效率。实验表明，应用所提方法后，GRPO训练的LLM在K&K逻辑谜题任务上性能最高提升46.2%。代码已开源：https://github.com/zhyang2226/AR-Lopti。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。