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

arXiv: https://arxiv.org/abs/2604.23318 · HTML: https://mlsanity.com/paper/2604.23318/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Reinforcement learning for reasoning, one of 17 directions in LLM reasoning.
- Builds on 31 of its 36 references that are in the map.
- First posted on arXiv 25 Apr 2026.

## Placed in

- LLM reasoning / [Reinforcement learning for reasoning](https://mlsanity.com/lr/reinforcement-learning-for-reasoning/)
- LLM alignment / RLHF / preference optimisation / [Process reward models for reasoning](https://mlsanity.com/alignment/process-reward-models-for-reasoning/)

## Builds on (most cited first, 30 of 31)

- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [LiveCodeBench: Holistic and Contamination Free Evaluation of Large Language Models for Code](https://mlsanity.com/paper/2403.07974/) (2024)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code Generation](https://mlsanity.com/paper/2305.01210/) (2023)
- [Understanding R1-Zero-Like Training: A Critical Perspective](https://mlsanity.com/paper/2503.20783/) (2025)
- [OlympiadBench: A Challenging Benchmark for Promoting AGI with Olympiad-Level Bilingual Multimodal Scientific Problems](https://mlsanity.com/paper/2402.14008/) (2024)
- [Beyond the 80/20 Rule: High-Entropy Minority Tokens Drive Effective Reinforcement Learning for LLM Reasoning](https://mlsanity.com/paper/2506.01939/) (2025)
- [Process Reinforcement through Implicit Rewards](https://mlsanity.com/paper/2502.01456/) (2025)
- [Self-Distilled Reasoner: On-Policy Self-Distillation for Large Language Models](https://mlsanity.com/paper/2601.18734/) (2026)
- [The Lessons of Developing Process Reward Models in Mathematical Reasoning](https://mlsanity.com/paper/2501.07301/) (2025)
- [Learning to Reason under Off-Policy Guidance](https://mlsanity.com/paper/2504.14945/) (2025)
- [Reasoning with Exploration: An Entropy Perspective](https://mlsanity.com/paper/2506.14758/) (2025)
- [Learning to Reason without External Rewards](https://mlsanity.com/paper/2505.19590/) (2025)
- [Learning beyond Teacher: Generalized On-Policy Distillation with Reward Extrapolation](https://mlsanity.com/paper/2602.12125/) (2026)
- [Critique-GRPO: Advancing LLM Reasoning with Natural Language and Numerical Feedback](https://mlsanity.com/paper/2506.03106/) (2025)
- [Stop Summation: Min-Form Credit Assignment Is All Process Reward Model Needs for Reasoning](https://mlsanity.com/paper/2504.15275/) (2025)
- [The Latent Space: Foundation, Evolution, Mechanism, Ability, and Outlook](https://mlsanity.com/paper/2604.02029/) (2026)
- [EDGE-GRPO: Entropy-Driven GRPO with Guided Error Correction for Advantage Diversity](https://mlsanity.com/paper/2507.21848/) (2025)
- [DCPO: Dynamic Clipping Policy Optimization](https://mlsanity.com/paper/2509.02333/) (2025)
- [GHPO: Adaptive Guidance for Stable and Efficient LLM Reinforcement Learning](https://mlsanity.com/paper/2507.10628/) (2025)
- [Hybrid Latent Reasoning via Reinforcement Learning](https://mlsanity.com/paper/2505.18454/) (2025)
- [LaDiR: Latent Diffusion Enhances LLMs for Text Reasoning](https://mlsanity.com/paper/2510.04573/) (2025)
- [Latent Thinking Optimization: Your Latent Reasoning Language Model Secretly Encodes Reward Signals in Its Latent Thoughts](https://mlsanity.com/paper/2509.26314/) (2025)
- [Outcome-Grounded Advantage Reshaping for Fine-Grained Credit Assignment in Mathematical Reasoning](https://mlsanity.com/paper/2601.07408/) (2026)
- [SSPO: Subsentence-level Policy Optimization](https://mlsanity.com/paper/2511.04256/) (2025)

## Built on by (most cited first, 1 of 1)

- [Spend Teacher Tokens Where They Matter: Success-Referenced On-Policy Distillation](https://mlsanity.com/paper/2610.02678/) (2026)

## Abstract (arXiv)

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 Span-level Hidden state Enabled Advantage Reweighting (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.