GRAIL: Gradient-Reweighted Advantages for Reinforcement Learning with Verifiable Rewards
Abstract
Gradient-Reweighted Advantage (GRAIL) improves mathematical reasoning in LLMs by reweighting token-wise advantages based on gradient-activation saliency, outperforming GRPO in accuracy and Pass@3 metrics.
Reinforcement learning with verifiable rewards (e.g. GRPO) is now a common way to improve mathematical reasoning in Large Language Models (LLMs). However, current methods usually broadcast one sequence-level advantage to all tokens, or use costly process reward models (PRMs) for step-level supervision. Uniform advantage distribution assumes that all tokens contribute equally to the final reward. This dilutes the gradient signal, since flawed reasoning steps and filler words are updated as strongly as valid logical inferences. To address this, we introduce Gradient-Reweighted Advantage (GRAIL), an intrinsic token-wise advantage reweighting method. GRAIL uses gradient-activation saliency to place more weight on tokens that are more locally sensitive to the final answer. Evaluations across five models from the Qwen3, R1-distilled and OctoThinker families show that GRAIL consistently outperforms GRPO. GRAIL achieved an average improvement of 3.60% in accuracy and 3.05% in Pass@3, demonstrating that fine-grained reasoning alignment can be achieved without process-level supervision.
Community
We introduce Gradient-Reweighted Advantage (GRAIL), an intrinsic token-wise advantage reweighting method. GRAIL uses gradient-activation saliency to place more weight on tokens that are more locally
sensitive to the final answer. Evaluations across five models from the Qwen3, R1-distilled and
OctoThinker families show that GRAIL consistently outperforms GRPO. GRAIL achieved an average
improvement of 3.60% in accuracy and 3.05% in Pass@3.
Models citing this paper 0
No model linking this paper
Datasets citing this paper 0
No dataset linking this paper
Spaces citing this paper 0
No Space linking this paper
Collections including this paper 0
No Collection including this paper