🤖 AI Summary
This work identifies a critical yet previously overlooked issue in reinforcement learning fine-tuning of large language models: training instability and inference degradation are frequently triggered by an extremely small fraction (approximately 0.01%) of rare, spurious tokens. The study demonstrates that these tokens destabilize training by adversely affecting policy gradients and reducing local policy entropy. To address this, the authors propose STAPO, a novel method that analyzes token-level policy gradients to mask gradient updates from spurious tokens and re-normalizes the loss over valid tokens. Unlike prior approaches relying on heuristic regularization, STAPO operates without such constraints and achieves an average improvement of 7.13% across six mathematical reasoning benchmarks on Qwen-1.7B, -8B, and -14B models, significantly outperforming GRPO, 20-Entropy, and JustRL while maintaining more stable entropy throughout training.
📝 Abstract
Reinforcement Learning (RL) has significantly improved large language model reasoning, but existing RL fine-tuning methods rely heavily on heuristic techniques such as entropy regularization and reweighting to maintain stability. In practice, they often experience late-stage performance collapse, leading to degraded reasoning quality and unstable training. We derive that the magnitude of token-wise policy gradients in RL is negatively correlated with token probability and local policy entropy. Building on this result, we prove that training instability is driven by a tiny fraction of tokens, approximately 0.01\%, which we term \emph{spurious tokens}. When such tokens appear in correct responses, they contribute little to the reasoning outcome but inherit the full sequence-level reward, leading to abnormally amplified gradient updates. Motivated by this observation, we propose Spurious-Token-Aware Policy Optimization (STAPO) for large-scale model refining, which selectively masks such updates and renormalizes the loss over valid tokens. Across six mathematical reasoning benchmarks using Qwen 1.7B, 8B, and 14B base models, STAPO consistently demonstrates superior entropy stability and achieves an average performance improvement of 7.13\% over GRPO, 20-Entropy and JustRL.