Source-linked AI summary

Outcome-Grounded Advantage Reshaping for Fine-Grained Credit Assignment in Mathematical Reasoning

Ziheng Li, Liu Kang, Feng Xiao, Luxi Xing, Qingyi Si, Zhuoran Li, Weikang Gong, Deqing Yang, Yanghua Xiao, Hongcheng Guo

arXiv:2601.07408v2cs.CLcs.LG

TL;DR

Standard GRPO assigns one group-level advantage uniformly across a response, limiting fine-grained credit assignment for reasoning tokens. OAR reshapes advantages according to token influence on the model’s outcome distribution, using perturbation-based OAR-P or gradient-based OAR-G with conservative mass-preserving reallocation. Across mathematical reasoning benchmarks, OAR consistently improves over strong GRPO baselines; OAR-P is the upper bound while OAR-G retains most gains at lower cost.

  • Problem

    GRPO broadcasts a sequence-level advantage across tokens even though reasoning tokens contribute unevenly to correctness, and entropy is not necessarily outcome-relevant.

  • Method

    OAR estimates token influence from shifts in the model’s final-answer distribution, using counterfactual perturbations in OAR-P, gradient sensitivity in OAR-G, and conservative mass-preserving advantage reshaping.

  • Results

    OAR consistently improves over strong GRPO baselines on mathematical reasoning benchmarks, with OAR-P setting the performance upper bound and OAR-G retaining most gains at minimal computational cost.

  • Takeaways & Limitations

    Outcome-grounded token attribution provides a fine-grained credit-assignment alternative for critic-free GRPO, with OAR-G offering a scalable efficiency–fidelity trade-off.

  • Takeaways & Limitations

    The answer-distribution surrogate can be imperfect when distributional changes are weakly coupled with verifier acceptance, and the evaluation focuses on mathematical reasoning with verifiable final answers.

Abstract

from arXiv · show

Group Relative Policy Optimization (GRPO) has emerged as a promising critic-free reinforcement learning paradigm for reasoning tasks. However, standard GRPO employs a coarse-grained credit assignment mechanism that propagates group-level rewards uniformly to to every token in a sequence, neglecting the varying contribution of individual reasoning steps. We address this limitation by introducing Outcome-grounded Advantage Reshaping (OAR), a fine-grained credit assignment mechanism that redistributes advantages based on how much each token influences the model's final answer. We instantiate OAR via two complementary strategies: (1) OAR-P, which estimates outcome sensitivity through counterfactual token perturbations, serving as a high-fidelity attribution signal; (2) OAR-G, which uses an input-gradient sensitivity proxy to approximate the influence signal with a single backward pass. These importance signals are integrated with a conservative Bi-Level advantage reshaping scheme that suppresses low-impact tokens and boosts pivotal ones while preserving the overall advantage mass. Empirical results on extensive mathematical reasoning benchmarks demonstrate that while OAR-P sets the performance upper bound, OAR-G achieves comparable gains with negligible computational overhead, both significantly outperforming a strong GRPO baseline, pushing the boundaries of critic-free LLM reasoning.

1 Introduction

GRPO’s sequence-level credit assignment broadcasts one advantage across every token, despite uneven token contributions to reasoning outcomes. OAR addresses this mismatch by estimating outcome-sensitive token influence and conservatively reallocating advantages.

  • Motivation: GRPO broadcasts a single group-normalized advantage to every token, overlooking that only some reasoning tokens determine correctness.This coarse assignment can increase gradient variance and accelerate entropy collapse.
  • Motivation: Entropy-based token weighting may misallocate credit because uncertainty does not necessarily reflect outcome-relevant importance.High entropy can occur for stylistic or lexically ambiguous tokens with little effect on correctness.
  • OAR: OAR scores tokens by how much perturbing them shifts the model’s final-answer distribution, using that distribution as a practical outcome surrogate.The surrogate avoids directly attributing sparse, non-differentiable verifier rewards to individual tokens.
  • OAR: OAR-P uses counterfactual perturbations for high-fidelity attribution, whereas OAR-G uses an efficient gradient-based proxy for scalable online training.The two strategies trade attribution fidelity against computational efficiency.
  • Advantage reshaping: The Bi-Level mechanism suppresses low-impact tokens, boosts pivotal tokens, and renormalizes weights to preserve overall advantage mass.This design targets sharper learning signals without changing the effective total advantage scale.
  • Results: OAR consistently outperforms strong GRPO baselines, with OAR-P setting the performance upper bound and OAR-G retaining most gains at minimal computational cost.The reported results come from mathematical reasoning benchmarks.

2 Preliminaries

GRPO is a critic-free PPO-style method that estimates advantages by normalizing rewards within groups and applies each resulting scalar across the sampled response. Existing entropy-based shaping modifies this allocation but can misrepresent outcome importance and alter update scale.

  • GRPO: GRPO samples multiple completions, normalizes their rewards within the group, and uses the resulting advantage for policy optimization.The method avoids training a learned value function.
  • GRPO: The same scalar advantage is applied to all token-level log-probability terms in each sampled response.This creates sequence-level rather than token-specific credit assignment.
  • GRPO: GRPO then optimizes a PPO-style clipped surrogate objective using importance ratios.The clipped objective is the policy-optimization component described in the preliminaries.
  • Intrinsic signals: Entropy-based variants redistribute advantages with intrinsic token-uncertainty signals and entropy-dependent bonuses.The shaping rule includes gradient stopping for stability.
  • Intrinsic signals: Intrinsic entropy proxies may emphasize uncertain tokens that have little effect on correctness, while additive shaping can change sequence-dependent update scales.Renormalization or retuning is needed to address the altered advantage mass.

3 Method

OAR adds token-level, outcome-grounded credit assignment to GRPO through perturbation- or gradient-based importance scores, then conservatively reallocates sequence advantages while preserving total advantage mass.

  • OAR framework: OAR equips GRPO with token-level credit assignment through OAR-P perturbation attribution and OAR-G gradient-based approximation.
  • OAR-P: OAR-P masks each reasoning token, recomputes the final prediction distribution, and measures the resulting KL divergence as token importance.Large divergence indicates an outcome-critical token, whereas near-zero divergence indicates syntactic filler; computing scores for length L requires L additional forward passes.
  • OAR-G: OAR-G injects Gaussian noise into token embeddings and uses Gradient × Input sensitivity to approximate outcome effects with a single backward pass.This replaces O(L) counterfactual forward passes and is computationally negligible compared with generation.
  • Importance normalization: Raw importance scores are log-transformed and min-max normalized within each sequence before advantage reshaping.
  • Bi-Level advantage reshaping: Bi-level reshaping suppresses low-importance tokens, boosts high-importance tokens, and sum-preserving renormalization keeps total advantage mass unchanged.The mechanism redistributes credit across tokens rather than globally rescaling advantages.
  • Training objective: By integrating outcome-grounded credit into the surrogate loss, OAR encourages refinement of critical reasoning steps while remaining robust to high-variance noise in long mathematical traces.

4 Experiments

Experiments evaluate OAR across mathematical reasoning benchmarks, backbones, baselines, and training dynamics. OAR improves reasoning performance and optimization behavior, with results also extending to a smaller model.

  • Baselines: The comparison includes vanilla GRPO, random credit, entropy-based weighting, and KTAE token-level advantage corrections.
  • Reasoning performance: OAR-P improves Qwen2.5-7B AIME25 P@1 from 9.4 to 11.8 and the overall average from 51.3 to 53.7.
  • Reasoning performance: OAR-P raises the Qwen2.5-Math-7B overall average from 57.3 to 59.5.
  • Training dynamics: Across training runs, OAR improves optimization stability, reaches higher rewards with fewer steps, and avoids premature entropy collapse.
  • Scaling to different sizes: On Qwen2.5-Math-1.5B, both OAR-G and OAR-P consistently improve over baselines.

5 Analysis

The analysis tests whether OAR importance scores identify causally important tokens and measures the computational cost of its two attribution strategies. OAR-P aligns most closely with the counterfactual oracle, while OAR-G offers a substantially cheaper alternative.

  • Causal-token identification: The counterfactual Oracle labels a token as causal when replacing it flips final-answer correctness; otherwise it is non-causal noise.
  • Alignment visualization: OAR-P assigns high importance to outcome-critical symbols and downweights stylistic or redundant tokens, qualitatively matching the Oracle mask.
  • Quantitative correlation: Across top-K% percentiles, OAR-P and OAR-G achieve higher Oracle causal-token recall than entropy weighting, with OAR-P consistently best.
  • Computational trade-off: Normalized training cost is 1.0× for GRPO, 1.4× for OAR-G, and 4.2× for batched OAR-P; naive serial OAR-P costs 18.7×.

6 Ablations

Ablations show that OAR’s credit concentration and performance depend on conservative gating and on how the outcome is represented. Moderate sparsification performs best, while answer-span mean logits provide the strongest outcome definition.

  • Bi-level gating: Larger gating thresholds generally produce sparser credit allocations than vanilla GRPO.ESS ratio measures weight uniformity, while top-10% advantage mass measures concentration in the highest-weighted tokens.
  • Bi-level gating: Moderate thresholding performs best: boost-only gating can collapse, suppress-only gating is stable but inefficient, and balanced gating achieves the best accuracy.
  • Outcome definitions: OAR uses an outcome probe to summarize final-answer prediction, and different probes can produce different token rankings and credit reallocations.
  • Outcome definitions: The ablation compares last-token logits, answer-span mean logits, and answer-span joint likelihood as outcome operators.
  • Outcome-definition results: Answer-span mean logits achieve the highest Oracle top-K% recall and best downstream training performance, followed by last-token logits, while answer-span joint likelihood performs worst.Distribution-based operators can capture outcome-distribution shifts that scalar likelihood may miss.

7 Related Work

Related work situates OAR within reinforcement learning with verifiable rewards and fine-grained credit assignment. The paper distinguishes its outcome-grounded attribution from prior model-internal token-credit methods.

  • Reinforcement learning with verifiable rewards: RLVR uses automatically checkable outcome signals, while GRPO stabilizes large-scale training by group-normalizing verifier rewards.
  • Fine-grained credit assignment: Prior fine-grained credit methods use attention dynamics, response-level confidence, or token-level certainty and correctness associations to redistribute advantages.

8 Conclusion

The paper introduces OAR to improve GRPO credit assignment in long-horizon reasoning by reallocating advantages according to token-level outcome contribution. OAR-P uses counterfactual attribution, while OAR-G provides a lightweight gradient-based proxy, and experiments show consistent gains over strong GRPO baselines.

  • OAR improves credit assignment in GRPO for long-horizon reasoning by redistributing sequence-level advantages according to outcome contribution.
  • OAR-P uses counterfactual attribution, whereas OAR-G uses a lightweight gradient-based proxy to estimate token importance.
  • Experiments on mathematical reasoning benchmarks show consistent gains over strong GRPO baselines.

Limitaions

The paper identifies limitations in OAR’s outcome-distribution proxy, masking intervention, evaluation scope, and reproducibility of related baselines.

  • OAR’s answer-distribution proxy can be imperfect when distributional changes are weakly coupled with verifier acceptance.
  • OAR-P’s masking intervention may introduce distribution shift compared with in-distribution edits.
  • Experiments focus on mathematical reasoning with verifiable final answers, leaving open-ended and interactive settings for future extension.
  • Several closely related methods are not yet fully reproducible, limiting the breadth of current comparisons.

A.1 Theoretical Analysis

The analysis explains why uniform GRPO advantages create noisy token updates and why OAR’s outcome-grounded weighting can suppress nuisance-token noise while concentrating credit on outcome-relevant steps. It also defines outcome influence, concentration metrics, and qualitative comparisons for evaluating these effects.

  • Additive proxy shaping can create sample-dependent update scales because its gradient magnitude depends on sequence length and high-proxy-token density.
  • The outcome probe is the model’s predictive distribution at an outcome-bearing position, and token influence is measured by the KL divergence between factual and masked distributions.
  • If token influence is approximately zero, masking it cannot substantially change the probability mass of any verifier-defined correct outcome set.
  • Standard GRPO broadcasts one trajectory-level advantage to every token, injecting variance when nuisance tokens did not influence the reward.
  • For a counterfactually non-influential token, OAR suppresses gradient-variance contribution relative to standard GRPO.
  • As ε →0, OAR’s gradient noise at a nuisance-token position vanishes, focusing optimization on tokens that causally affect the outcome.
  • Figure 8 compares OAR-G with an oracle causal mask and an entropy-based baseline on a representative reasoning trace.
  • The ESS ratio measures weight uniformity, while top-10% advantage mass measures how much absolute advantage is carried by the highest-weighted subset of valid tokens.
Loading 2601.07408v2…