Hidden States Know Where Reasoning Diverges: Credit Assignment via Span-Level Wasserstein Distance

• URL: http://arxiv.org/abs/2604.23318v1
• Date: Sat, 25 Apr 2026 14:11:23 GMT
• Title: Hidden States Know Where Reasoning Diverges: Credit Assignment via Span-Level Wasserstein Distance
• Abstract summary: We show that hidden-state distributions contain a useful signal for local reasoning quality.<n>Within each GRPO group, the Wasserstein distance between span-level hidden state distributions increases around regions where their local reasoning quality diverges.<n>We propose textbfSpan-level textbfHidden state textbfEd textbfAdvantage textbfReweighting.
• Score: 8.66471442661456
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Group Relative Policy Optimization (GRPO) performs coarse-grained credit assignment in reinforcement learning with verifiable rewards (RLVR) by assigning the same advantage to all tokens in a rollout. Process reward models can provide finer-grained supervision, but they require step-level annotation or additional reward modeling. We show that hidden-state distributions contain a useful signal for local reasoning quality that can be extracted using only outcome-level correctness labels available in RLVR. Specifically, within each GRPO group, the Wasserstein distance between span-level hidden state distributions of correct and incorrect rollouts increases around regions where their local reasoning quality diverges. This association holds both across examples and within individual trajectories, suggesting that hidden-state distributional divergence can serve as a self-supervision signal for fine-grained credit assignment. We formalize this observation with a separation theorem showing that, under mild structural assumptions, post-divergence spans have larger Wasserstein distances than pre-divergence spans whenever the population-level distributional gap exceeds finite-sample noise. Motivated by this result, we propose \textbf{S}pan-level \textbf{H}idden state \textbf{E}nabled \textbf{A}dvantage \textbf{R}eweighting (SHEAR), which modifies GRPO by using span-level Wasserstein distances to scale token-level advantages, amplifying updates on tokens whose hidden states are more separated from the opposing group. The method requires no additional model and only minimal changes to the training pipeline. Experiments on five mathematical reasoning benchmarks and five code generation benchmarks show improvements over standard GRPO and strong performance relative to supervised process reward models, while requiring no additional annotation or reward model training.

Related papers

• When Correct Solutions Repeat: Rarity-Aware Credit Redistribution for GRPO [25.979415525596007]
  Reinforcement learning with verifiable rewards (RLVR)<n>We formalize this behavior as multiplicity-induced structure-level credit concentration.<n> Cue-GRPO instantiates this rule with- out auxiliary-model inference.
  arXiv  Detail & Related papers  (2026-08-04T11:02:26Z)
• ExTra: Exploratory Trajectory Optimization for Language Model Reinforcement Learning [75.26064451955199]
  ExTra is a framework that extracts exploration signals from the model's own rollouts.<n>It improves Qwen3-1.7B over GRPO by about +5 points on pass@1 and +7 points on pass@16.
  arXiv  Detail & Related papers  (2026-06-23T15:51:39Z)
• TRACE: A Unified Rollout Budget Allocation Framework for Efficient Agentic Reinforcement Learning [52.67756371231985]
  Reinforcement learning with verifiable rewards (RLVR) is a promising approach for enhancing reasoning and agentic behavior in large language models.<n>We introduce Tree Rollout Allocation for Contrastive Exploration (TRACE), a unified rollout allocation framework.<n>Technically, TRACE allocates rollout budget to both prompt roots and intermediate prefixes that are most likely to yield mixed terminal rewards.
  arXiv  Detail & Related papers  (2026-06-09T17:16:03Z)
• GRASP: Geometry-aware Residual Alignment for Scalable Pretraining Data Attribution [5.770893169582546]
  This work reframes attribution as subset-level counterfactual utility prediction.<n>We introduce GRASP, an interaction-aware surrogate.<n>It decisively outperforms existing scalable baselines.
  arXiv  Detail & Related papers  (2026-06-05T04:17:50Z)
• ARCA: Adapter-Residual Credit Assignment When Token Signals Degenerate [0.0]
  Token-level credit assignment for language-model reinforcement learning is usually formulated as if the policy were fully trainable.<n>We propose a lightweight alternative that derives token salience from the adapter's own hidden-state residual.<n>In a compact MATH/Qwen3-1.7B GRPO sweep, ARCA exhibits the predicted non-degenerate middle-regime credit distribution under matched rollout budgets.
  arXiv  Detail & Related papers  (2026-05-29T18:42:06Z)
• GEAR: Granularity-Adaptive Advantage Reweighting for LLM Agents via Self-Distillation [33.370957547486775]
  Granularity-adaptivE Advantage Reweighting reshapes the trajectory-level GRPO advantage using token- and segment-level signals.<n>GEAR consistently outperforms standard GRPO, self-distillation-only baselines, and token- or turn-level credit-assignment methods.
  arXiv  Detail & Related papers  (2026-05-12T09:38:38Z)
• Relative Score Policy Optimization for Diffusion Language Models [29.344961499429257]
  Diffusion large language models (dLLMs) offer a promising route to parallel and efficient text generation.<n>The lack of tractable sequence-level log-ratios forces existing methods to rely on high-variance ELBO-based approximations.<n>We propose textbfRelative textbfScore textbfPolicy textbfOptimization (RSPO) to calibrate noisy likelihood estimates in dLLMs.
  arXiv  Detail & Related papers  (2026-05-11T08:58:40Z)
• Too Correct to Learn: Reinforcement Learning on Saturated Reasoning Data [55.84428098924793]
  We propose Constrained Uniform Top-K Sampling (CUTS), a parameter-free decoding strategy enforcing structure-preserving exploration.<n>We integrate this into Mixed-CUTS, a training framework synergizing exploitative and exploratory rollouts to amplify intra-group advantage variance.<n> Notably, Mixed-CUTS improves Pass@1 accuracy on the challenging AIME25 benchmark by up to 15.1% over standard GRPO.
  arXiv  Detail & Related papers  (2026-04-20T16:43:28Z)
• Balanced Aggregation: Understanding and Fixing Aggregation Bias in GRPO [70.38763678943648]
  Reinforcement learning with verifiable rewards (RLVR) has become a central paradigm for improving reasoning and code generation in large language models.<n>Standard GRPO uses sequence aggregation, while recent work has advocated token aggregation as a better alternative.<n>We show that these two rules induce different optimization biases; token aggregation introduces sign-length coupling, while sequence aggregation implicitly downweights longer responses.
  arXiv  Detail & Related papers  (2026-04-14T09:48:46Z)
• When Right Meets Wrong: Bilateral Context Conditioning with Reward-Confidence Correction for GRPO [18.988527161000203]
  Group Relative Policy Optimization (GRPO) has emerged as an effective method for training reasoning models.<n>We present a contrastive reformulation of GRPO, showing that the GRPO objective implicitly maximizes the margin between the policy ratios of correct and incorrect samples.<n>We propose Bilateral Context Conditioning (BICC), a mechanism that allows the model to cross-reference successful and failed reasoning traces.
  arXiv  Detail & Related papers  (2026-03-13T16:25:02Z)
• Improving Search Agent with One Line of Code [68.58667107354253]
  Tool-based Agentic Reinforcement Learning (TARL) has emerged as a promising paradigm for training search agents to interact with external tools.<n>We propose textbfSearch textbfAgent textbfPolicy textbfOptimization (textbfSAPO), which stabilizes training via a conditional token-level KL constraint.
  arXiv  Detail & Related papers  (2026-03-10T04:07:39Z)
• SimGR: Escaping the Pitfalls of Generative Decoding in LLM-based Recommendation [68.00727783181289]
  A core objective in recommender systems is to accurately model the distribution of user preferences over items to enable personalized recommendations.<n>We observe that existing methods inevitably introduce systematic bias when estimating item-level preference distributions.<n>We propose textbfSimply textbfGenerative textbfRecommendation (textbfSimGR), a framework that directly models item-level preference distributions in a shared latent space.
  arXiv  Detail & Related papers  (2026-02-08T07:26:52Z)
• Owen-Shapley Policy Optimization (OSPO): A Principled RL Algorithm for Generative Search LLMs [12.75200353208858]
  Owen-Shapley Policy Optimization (OSPO) is a framework that redistributes sequence-level advantages based on tokens' marginal contributions to outcomes.<n>Unlike value-model-based methods requiring additional computation, OSPO employs potential-based reward shaping via Shapley-Owen attributions to assign segment-level credit.<n> Experiments on Amazon ESCI and H&M Fashion datasets show consistent gains over baselines.
  arXiv  Detail & Related papers  (2026-01-13T10:17:46Z)
• Beyond Majority Voting: Towards Fine-grained and More Reliable Reward Signal for Test-Time Reinforcement Learning [12.354777054071379]
  Test-time reinforcement learning mitigates reliance on annotated data by using majority voting results as pseudo-labels.<n>This voting strategy often induces confirmation bias and suffers from sparse rewards, limiting the overall performance.<n>We propose subgroup-specific step-wise confidence-weighted pseudo-label estimation (SCOPE) to address these issues.
  arXiv  Detail & Related papers  (2025-12-17T07:21:54Z)
• Rewarding the Unlikely: Lifting GRPO Beyond Distribution Sharpening [36.81125165911328]
  Reinforcement learning is emerging as a primary driver for improving language model reasoning capabilities.<n>We investigate whether current reinforcement learning algorithms merely sharpen the base model's distribution around problems it can already solve.<n>We show that unlikeliness reward mitigates rank bias and improves pass@$N$ across a large range of $N$ in both synthetic and real theorem proving settings.
  arXiv  Detail & Related papers  (2025-06-03T01:15:15Z)
• Submodular Reinforcement Learning [38.40138241424851]
  In reinforcement learning (RL), rewards of states are typically considered additive, and following the Markov assumption, they are $textitindependent$ states visited previously. In many important applications, such as coverage control, experiment design and informative path planning, rewards naturally have diminishing returns, i.e., their value decreases in light of similar states visited previously. We propose $textitsubmodular RL$ (SubRL), a paradigm which seeks to optimize more general, non-additive (and history-dependent) rewards modelled via submodular set functions which capture diminishing returns
  arXiv  Detail & Related papers  (2023-07-25T09:46:02Z)
• Implicit Distributional Reinforcement Learning [61.166030238490634]
  implicit distributional actor-critic (IDAC) built on two deep generator networks (DGNs) Semi-implicit actor (SIA) powered by a flexible policy distribution. We observe IDAC outperforms state-of-the-art algorithms on representative OpenAI Gym environments.
  arXiv  Detail & Related papers  (2020-07-13T02:52:18Z)
• Nested-Wasserstein Self-Imitation Learning for Sequence Generation [158.19606942252284]
  We propose the concept of nested-Wasserstein distance for distributional semantic matching. A novel nested-Wasserstein self-imitation learning framework is developed, encouraging the model to exploit historical high-rewarded sequences.
  arXiv  Detail & Related papers  (2020-01-20T02:19:13Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.