PaperlayerView exact revision on arXiv

Source-linked AI summary

GRAIL: Gradient-Reweighted Advantages for Reinforcement Learning with Verifiable Rewards

Tej Deep Pala, Vernon Toh, Soujanya Poria

arXiv:2606.04889v1cs.CL

TL;DR

GRPO-style methods assign the same sequence-level advantage to every token, limiting credit assignment across reasoning traces. GRAIL uses gradient-activation saliency to reweight token advantages and consistently outperforms GRPO across five models, improving average accuracy by 3.60% and Pass@3 by 3.05%.

  • Problem

    GRPO-style methods broadcast one sequence-level advantage to all tokens, assuming equal contribution and limiting fine-grained credit assignment across reasoning traces.

  • Method

    GRAIL computes gradient-activation saliency for each token and uses it to reweight the sequence-level advantage in the policy-gradient objective.

  • Results

    Across five models and six mathematical reasoning benchmarks, GRAIL consistently outperformed GRPO, improving average accuracy by 3.60% and Pass@3 by 3.05%.

  • Takeaways & Limitations

    GRAIL provides a self-contained approach to fine-grained reasoning alignment without process-level supervision in evaluated mathematical problem-solving.

  • Takeaways & Limitations

    GRAIL’s first-order influence approximation misses higher-order token interactions, while more rigorous saliency estimators would require computationally prohibitive multiple forward and backward passes.

Abstract

from arXiv · show

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.

1 Introduction

GRPO-style methods assign the same sequence-level advantage to every token, diluting updates by treating flawed reasoning and filler words like critical logical steps. GRAIL addresses this with gradient-activation saliency for intrinsic token-wise advantage reweighting and outperforms GRPO in evaluations across five models.

  • Motivation: GRPO-style methods broadcast one sequence-level advantage across all tokens, assuming equal contribution to the final outcome.This coarse credit assignment dilutes the gradient signal across reasoning traces.
  • Motivation: Flawed reasoning steps and filler words can receive updates as strong as critical logical steps under uniform advantage distribution.The resulting mismatch motivates focusing updates with an intrinsic saliency signal.
  • GRAIL: GRAIL computes gradient-activation saliency for each token and uses it to reweight the sequence-level advantage in the policy-gradient loss.This intrinsic token-wise reweighting gives stronger updates to salient tokens and weaker updates to incidental filler words.
  • Results: 3.60% improvement in average accuracy is reported for GRAIL over the GRPO baseline across evaluations of five Qwen3, R1-distilled, and OctoThinker models.The paper reports that GRAIL consistently outperforms GRPO.

2 Related Work

Related work contrasts coarse outcome-level rewards and externally supervised process rewards with GRAIL’s intrinsic gradient-based token weighting. It also situates GRAIL among saliency methods and recent token-level advantage redistribution approaches.

  • Reward Models: Outcome reward models provide only a binary completion-level signal, whereas process reward models assign rewards to intermediate reasoning steps for denser error-localizing supervision.PRMs may be trained with human labels or automated verification.
  • GRAIL: GRAIL derives token-level weights directly from intrinsic training gradients, enabling fine-grained credit assignment without external-supervision overhead.This distinguishes GRAIL from approaches relying on separately trained or externally supervised process reward models.
  • Saliency Methods: GRAIL’s element-wise gradient × activation estimator approximates first-order Taylor sensitivity and is cheaper than path-integral methods such as Integrated Gradients.Gradient-based saliency uses output gradients with respect to inputs as local sensitivity signals.
  • Saliency Methods: Prior language-model saliency work typically identifies influential tokens after training or traces factual knowledge, while attention weights are not reliable proxies for causal effect.Using saliency during training remains less explored.
  • Token-Level Advantage Redistribution: Recent methods redistribute sequence-level advantages using token entropy, implicit online process rewards, or token-influence gating to improve credit assignment.Examples include entropy-based reweighting, PRIME’s outcome-only implicit PRM, and OAR’s bi-level gating mechanism.

3 Methodology

GRAIL addresses noisy sequence-level credit assignment by reweighting the policy-gradient advantage with token-specific gradient-activation saliency. It derives bounded weights from answer-loss sensitivity at input embeddings, without requiring an external process reward model.

  • Token-wise advantage reweighting: GRAIL reweights the sequence-level advantage with a token-specific saliency weight that proxies each token’s sensitivity to the final answer.This distinguishes filler or flawed intermediary tokens from logical inferences that drive a positive outcome.
  • Gradient-activation saliency: Gradient-activation saliency uses input embeddings as differentiable leaves so answer-loss gradients propagate through causal self-attention to intermediate reasoning tokens.The resulting signal is non-zero when attention routes information from a reasoning position into the final answer states.
  • Gradient-activation saliency: The raw saliency score is computed from the L2-norm of the element-wise product between each token’s embedding and its answer-loss gradient.This provides a first-order local sensitivity signal computed by the current model during training.
  • Weight normalisation and stability: Log-space standardisation and bounded clipping stabilise heavy-tailed saliency weights and prevent extreme tokens from dominating policy-gradient updates.The final weights are controlled by wmean and σw, then clipped to [wmin, wmax].
  • Intrinsic supervision: GRAIL does not require external heuristics or an independently trained process reward model because its token sensitivity scores are computed intrinsically during training.The saliency weights are derived from the current model’s answer-loss gradients.
  • Token-type corrections: Final answer tokens are assigned wmax, delimiters are downweighted to avoid reinforcing formatting, and post-answer positions are reset to wmean.These corrections preserve emphasis on final correctness, reasoning over formatting, and confidence in when to stop.

4 Experimental Setup

The experiments compare GRAIL with GRPO under matched training conditions across five models from three families, using DeepMath-103K and six mathematical reasoning benchmarks. Evaluation samples three generations per problem at temperature 0.6 and reports pass rate and average accuracy.

  • Models and Dataset: Five models from three families are trained on DeepMath-103K to assess the generalisability of GRAIL’s saliency mechanism.The models are Qwen3-4B, Qwen3-8B, Deepseek-R1-Distill-Llama-8B, OctoThinker-3B-Short, and OctoThinker-8B-Short.
  • Controlled Comparison: GRAIL and GRPO use identical base models, datasets, and core training infrastructure, isolating the effect of the loss reweighting formulation.Training hyperparameters are provided in Appendix A.
  • Evaluation Benchmarks: The evaluation covers six benchmarks spanning varying difficulty: Math500, AIME 2024, AMC 2023, MinervaMath, CollegeMath, and OlympiadBench.These benchmarks measure generalisation of mathematical problem-solving capabilities.
  • Inference Protocol: For each problem, inference samples three distinct generations at temperature 0.6, with performance reported as pass rate and average accuracy across generations.The same sampling strategy is applied across all evaluation datasets.

5 Results

Across five model architectures and six mathematical reasoning benchmarks, GRAIL consistently improves over GRPO by reweighting gradient signals toward influential tokens. Results also show statistically significant gains, superiority over OAR-G, and sensitivity to saliency placement and rollout weighting.

  • Overall performance: 3.60% average accuracy and 3.05% Pass@3 improvements over GRPO establish GRAIL’s consistent gains across the evaluation suite.Results were averaged across five model architectures and six mathematical reasoning benchmarks.
  • Statistical significance: 1.3 × 10^-8 p-value confirms statistically significant accuracy improvements for GRAIL over GRPO across models and benchmarks.The Wilcoxon signed-rank test reported W = 459.5.
  • Comparison with OAR-G: 6.29% absolute average-accuracy gain on Qwen3-8B shows GRAIL outperforming OAR-G by 62.04% to 55.75%.On Qwen3-4B, GRAIL’s average-accuracy improvement over OAR-G was 1.47%.
  • Comparison with OAR-G: 15.56% AIME 2024 accuracy improvement on Qwen3-8B demonstrates GRAIL’s strongest advantage over OAR on complex multi-step derivations.The corresponding improvement on Qwen3-4B was 6.67%.
  • Saliency-leaf ablation: 58.15% average accuracy from input-embedding saliency falls to 53.41% at the penultimate layer, showing that deeper leaf placement degrades performance.On AIME24, accuracy declined from 36.67% to 24.44% when the backward pass was truncated.
  • Rollout weighting: 58.86% peak accuracy on Qwen3-4B and 62.04% on Qwen3-8B occur when weights are applied only to wrong rollouts, while correct-only weighting underperforms.Wrong-rollout weighting reaches 47.78% on AIME24 for Qwen3-8B, suggesting sharper sensitivity to logical errors.

6 Conclusion

GRAIL addresses uniform token-level credit assignment in GRPO by reweighting the loss with token-specific importance weights derived from gradient saliency. The method provides a self-contained approach to fine-grained reasoning alignment and may extend beyond mathematical problem-solving to other verifiable-reinforcement-learning domains.

  • Conclusion: GRAIL addresses GRPO’s uniform credit assignment by deriving token-specific importance weights for loss reweighting.The method is designed to distinguish critical logical inferences from flawed intermediate derivations.
  • Conclusion: GRAIL provides a self-contained approach to improving fine-grained reasoning alignment in large language models.It was evaluated on mathematical problem-solving.
  • Conclusion: GRAIL’s intrinsic gradient-weighting mechanism may extend beyond mathematics to other domains requiring verifiable trajectories.The weighting signal is grounded in the model’s own gradients, providing a simple path for improving verifiable reinforcement learning.

Limitations

GRAIL’s gradient-activation product provides an effective intrinsic learning signal but relies on a first-order approximation of token influence. Consequently, it captures linear sensitivity while omitting higher-order and counterfactual interactions, although Integrated Gradients could offer a more rigorous estimator.

  • Limitations: GRAIL’s gradient-activation product relies on a first-order approximation of influence.It captures the linear sensitivity of the final answer to each token’s input representation.
  • Limitations: The approximation omits higher-order interactions, including tokens that become critical only in combination.It also does not account for counterfactual contributions.
  • Limitations: Integrated Gradients could provide a more theoretically rigorous saliency estimator.The passage presents it as an advanced alternative to GRAIL’s current saliency method.

A Hyperparameters · A.1 Training Hyperparameters

The training setup fine-tunes all models with Hugging Face’s open-source trl library, uses the Dr. GRPO loss formulation, and sets the KL-divergence penalty β to 0. Table 5 provides the complete hyperparameter list for the training runs.

  • A Hyperparameters: Table 5 lists the hyperparameters used to train Qwen3 and OctoThinker models.The table covers training through the Hugging Face trl library.
  • A.1 Training Hyperparameters: The hyperparameter summary covers training runs across Qwen3 and OctoThinker model families.Table 5 is specifically described as containing the hyperparameters used for these models.
  • A.1 Training Hyperparameters: All models are fine-tuned using Hugging Face’s open-source trl library.The implementation uses the trl library from Hugging Face.
  • A.1 Training Hyperparameters: The training objective is configured with the Dr. GRPO loss formulation.The configured loss is identified as dr_grpo.
  • A.1 Training Hyperparameters: β = 0 for the KL-divergence penalty.The GRPO configuration sets the KL divergence penalty to zero.
  • A Hyperparameters: The training setup uses GRPO as its objective framework.The passage specifies the GRPO objective before identifying its Dr. GRPO loss formulation.
  • A.1 Training Hyperparameters: The paper directs readers to Table 5 for a comprehensive list of training hyperparameters.The passage states that the full list is provided there.

A.2 Evaluation Hyperparameters · A.3 Prompt

Evaluation uses a consistent sampling, answer-extraction, and mathematical-verification protocol across six benchmarks, with rollout parameters specified in Table 6. A standardized zero-shot prompt elicits step-by-step reasoning, requires boxed final answers, and supports both automated extraction and GRAIL’s positional corrections.

  • A.2 Evaluation Hyperparameters: Across all six mathematical reasoning benchmarks, evaluation maintains a consistent sampling and verification protocol for fair comparisons.Each problem receives multiple candidate solutions.
  • A.2 Evaluation Hyperparameters: Final mathematical answers are extracted from reasoning traces by parsing the contents of \boxed{}.The extracted predictions are evaluated against dataset ground truth.
  • A.2 Evaluation Hyperparameters: Math-verify evaluates mathematical equivalence between each extracted prediction and the dataset’s ground truth.This verification procedure is part of the benchmark evaluation protocol.
  • A.2 Evaluation Hyperparameters: Table 6 details the sampling parameters used to generate evaluation rollouts.The table is identified as covering sampling hyperparameters for the benchmark evaluation phase.
  • A.3 Prompt: A standardized zero-shot template formats query q during both training rollout generation and final evaluation.The template is designed to elicit structured, step-by-step reasoning and reliably formatted final answers.
  • A.3 Prompt: The prompt asks the model to reason step by step and place its final answer within \boxed{}.This explicit delimiter enables reliable automated extraction.
  • A.3 Prompt: The boxed-answer constraint makes reasoning traces compatible with math-verify extraction and provides a delimiter for GRAIL’s positional weight corrections.The same structure is used to support both verification and the GRAIL mechanism.

B Computational Efficiency and Training Time

GRAIL increases training cost because its intrinsic token-level saliency requires an additional backward pass before policy updates. Across evaluated models, this produces approximately 50%–60% average computational overhead over GRPO under the reported training setup.

  • Source of overhead: An additional backward pass through the attention layers to the input embeddings computes GRAIL’s intrinsic token-level saliency scores before the policy update.This extra pass is the source of GRAIL’s computational overhead.
  • Training overhead: 50%–60% average computational overhead is introduced by GRAIL over standard GRPO across the evaluated models.Wall-clock training times were compared on a single compute node with 4 × NVIDIA H200 GPUs.
  • Evaluation setup: All models in the wall-clock comparison were trained for 200 steps on 4 × NVIDIA H200 GPUs.The comparison covered standard GRPO and GRAIL.
Loading 2606.04889v1…