🤖 machine learning

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

The paper proposes **SHEAR**, a method that improves reinforcement learning for reasoning tasks by using the Wasserstein distance between hidden-state distributions of correct and incorrect rollouts to perform fine-grained, token-level credit assignment without requiring additional reward models or step-level annotations.

Original authors: Xinzhu Chen, Wei He, Huichuan Fan, Wenzhe Niu, Zhongxiang Sun, Xuanru Wang, Jiuchong Gao, Jinghua Hao, Renqing He, Weijie Yu

Published 2026-04-28

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Xinzhu Chen, Wei He, Huichuan Fan, Wenzhe Niu, Zhongxiang Sun, Xuanru Wang, Jiuchong Gao, Jinghua Hao, Renqing He, Weijie Yu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a teacher grading a long, complex math exam. Most AI models today are like a teacher who only looks at the final answer. If the student gets the answer right, the teacher gives them an A+ for the whole page. If they get it wrong, the teacher gives them an F for the whole page.

The problem? The student might have done 90% of the work perfectly but made one tiny calculation error at the very end. Or, they might have guessed the right answer by pure luck despite having a totally broken logical process. Giving the same "reward" to every single step is unfair and makes it hard for the student (the AI) to learn exactly where they went wrong.

This paper introduces a new method called SHEAR to solve this "Credit Assignment" problem.

The Core Idea: The "Vibe Check" of Hidden States

Instead of hiring a second, expensive teacher (called a Process Reward Model) to watch every single step, the researchers discovered something brilliant: The AI’s own "brain waves" (hidden states) already know when it's starting to lose the plot.

Think of it like a professional chef cooking a complex meal.

When the chef is following a perfect recipe, their movements are rhythmic, confident, and consistent.
The moment they realize they’ve added too much salt or burned the garlic, their "vibe" changes. Their movements become frantic, hesitant, or disorganized.

The researchers found that when an AI is on a "correct" path, its internal mathematical patterns (hidden states) look a certain way. When it starts to make a mistake, those patterns suddenly "diverge"—they drift away from the "correct" pattern and start looking like the patterns seen in failed attempts.

The Tool: The Wasserstein Distance (The "Difference Meter")

To measure this "vibe shift," they use a mathematical tool called the Wasserstein Distance.

Imagine you have two piles of sand. One pile represents the "brain waves" of a successful reasoning path, and the other represents a failed one. The Wasserstein Distance is essentially the minimum amount of work required to move the sand from one pile to match the shape of the other.

If the piles are almost identical, the distance is low (the "vibe" is the same).
If the piles look completely different, the distance is high (the "vibe" has shifted).

The Method: SHEAR (The Smart Grader)

SHEAR uses this "Difference Meter" to grade the AI's work with surgical precision. Here is how it works in three steps:

Scanning: It slides a window across the AI's reasoning chain, looking at small chunks (spans) of its "brain waves" at a time.
Comparing: It compares those chunks to the "brain waves" of successful attempts.
Reweighting:
- If a chunk of reasoning looks just like a successful path, it gets a standard grade.
- If a chunk looks wildly different from the successful paths (meaning the AI likely just made a mistake), SHEAR amplifies the signal. It tells the AI: "Hey! This specific part was the turning point! Pay extra attention here!"

Why is this a big deal?

It’s Self-Supervised: It doesn't need humans to sit there and label every single step of a math problem. It learns purely by comparing its own "good" thoughts to its "bad" thoughts.
It’s Efficient: It doesn't require training a whole new "teacher" model. It just uses the brain waves the AI is already producing.
It Actually Works: In their tests, this method helped AI get much better at complex math and coding, often performing as well as (or better than) much more expensive, human-supervised methods.

In short: SHEAR teaches AI to listen to its own "gut feeling" to figure out exactly which step caused a mistake, making it a much faster and more precise learner.

Technical Summary: Hidden States Know Where Reasoning Diverges

Title: Hidden States Know Where Reasoning Diverges: Credit Assignment via Span-Level Wasserstein Distance
Core Concept: Using the distributional divergence of a model's own hidden states to perform fine-grained credit assignment in Reinforcement Learning with Verifiable Rewards (RLVR).

1. The Problem: Coarse-Grained Credit Assignment

In Reinforcement Learning for reasoning (e.g., math and code), algorithms like Group Relative Policy Optimization (GRPO) typically use outcome-level rewards. This creates a "credit assignment" problem: if a reasoning chain is 500 tokens long and contains one logical error at token 200, the entire trajectory is penalized.

Current solutions to this involve Process Reward Models (PRMs), which provide step-level supervision. However, PRMs are problematic because:

Cost: They require expensive human annotation or auxiliary model training.
Miscalibration: As the policy model evolves during RL, a static PRM may become misaligned with the model's current reasoning patterns.

2. Methodology: SHEAR (Span-level Hidden state Enabled Advantage Reweighting)

The authors propose a self-supervised approach that extracts "process-level" signals directly from the model's internal representations without any external supervision.

A. The Empirical Observation:
The authors discovered that the Wasserstein distance (a metric measuring the distance between probability distributions) between the hidden states of "correct" rollouts and "incorrect" rollouts increases significantly at the exact points where reasoning quality diverges. This divergence tracks local reasoning success even within a single trajectory.

B. The SHEAR Algorithm:

Span Decomposition: Instead of looking at single tokens, the method partitions rollouts into overlapping spans (using a sliding window). This captures local context.
Wasserstein Distance Calculation: For each span, the algorithm calculates the minimum Wasserstein distance to the "opposing group" (e.g., if evaluating an incorrect rollout, it compares its spans to the spans of correct rollouts in the same group).
Max-Pooling Aggregation: The span-level distances are projected back to individual tokens using max-pooling, ensuring that a token receives the highest discrepancy signal from any span covering it.
Advantage Reweighting: The standard GRPO advantage is scaled by these token-level weights ( $\omega_t$ $ω_{t}$ ).
- Incorrect rollouts: Tokens in high-divergence regions (where the error occurred) receive larger negative updates.
- Correct rollouts: Tokens in high-divergence regions (where the model was "most distinct" from failure modes) receive larger positive updates.

C. Theoretical Foundation:
The authors provide a Separation Theorem proving that, under mild assumptions, the expected Wasserstein distance for post-error spans is strictly larger than for pre-error spans, provided the distributional gap exceeds finite-sample noise.

3. Key Contributions

Self-Supervised Signal: Demonstrated that hidden-state distributions carry a "latent" process reward signal that can be extracted using only outcome labels.
SHEAR Method: A plug-and-play modification to GRPO that requires no additional models or human annotations.
Theoretical Guarantee: Provided a mathematical justification for why Wasserstein distance is a valid proxy for local reasoning quality.
Efficiency: Proved that the computational overhead is minimal (under 16% increase in training time).

4. Experimental Results

The method was tested on multiple backbones (Qwen2.5, Llama 3.1) across math and code benchmarks.

Mathematical Reasoning: SHEAR consistently outperformed standard GRPO and, significantly, outperformed PRM-based methods (like PURE). It achieved higher accuracy on difficult benchmarks like AIME and AMC.
Code Generation: SHEAR showed consistent improvements over GRPO, particularly on complex, long-form tasks like LiveCodeBench.
Robustness: Ablation studies showed that using "distribution-aware" metrics (Wasserstein, MMD, Chamfer) is critical; using simple "mean-shift" metrics (Cosine distance) actually harms performance.

5. Significance

This paper shifts the paradigm of credit assignment in LLM reasoning. It suggests that the model already "knows" when it is making a mistake through its internal representations. By mathematically surfacing these moments of divergence, we can guide the model to learn more precisely from its successes and failures, achieving the benefits of expensive Process Reward Models through purely self-supervised, representation-based methods.