GRAIL: Gradient-Reweighted Advantages for Reinforcement Learning with Verifiable Rewards
2026-06-03 • Computation and Language
Computation and Language
AI summaryⓘ
The authors studied a way to teach large language models to reason better in math problems by improving how feedback is given during learning. They found that giving the same feedback to every word in an answer can weaken learning because not all words help reach the right solution. To fix this, they created a method called GRAIL that gives more attention to important words based on how much they affect the final answer. Tests showed GRAIL helps models do better than older methods without needing detailed step-by-step supervision.
Reinforcement LearningLarge Language ModelsMathematical ReasoningReward ModelsGradient ActivationAdvantage ReweightingSequence-level AdvantageToken-wise SupervisionGRPOPass@3
Authors
Tej Deep Pala, Vernon Toh, Soujanya Poria
Abstract
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.