🤖 machine learning

Localizing Credit at the Divergence: Path-Conditioned Self-Distillation for LLM Reasoning

This paper proposes Hindsight Self-Distillation (HSD), a method that improves token-level credit assignment in long reasoning traces by conditioning a teacher model on successful peer rollouts to provide dense, path-specific guidance at critical divergence points, thereby outperforming existing reinforcement learning and distillation baselines on math and code benchmarks.

Original authors: Yu Li, Shu Hong, Tian Lan

Published 2026-06-16

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yu Li, Shu Hong, Tian Lan

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 teaching a student to solve a complex math problem. You give them a question, and they write out a long, step-by-step reasoning process on a piece of paper. At the very end, they write down a final answer.

The Problem: The "Silent Teacher"
In traditional AI training (specifically a method called Reinforcement Learning), the teacher only looks at the final answer.

If the answer is wrong, the teacher says, "Bad job," and gives the student a single "thumbs down" for the entire essay.
The student doesn't know where they went wrong. Did they mess up the first step? The middle? Or just the very last calculation?
Because the teacher only sees the final result, they are "silent" during the long, confusing middle parts of the reasoning. They can't point to the specific sentence where the logic broke.

This is especially bad for tasks where the answer is very short (like a single number). If the answer is just "42," the teacher has no way of knowing which of the 500 words the student wrote led to that number, or which word caused the mistake.

The Old Fix: "Answer-Conditioned" Distillation
Researchers tried a smarter approach called "Self-Distillation." Here, the AI plays two roles:

The Student: Tries to solve the problem with no help.
The Teacher: Tries to solve the problem while knowing the correct final answer.

The idea was that if the Teacher knows the answer is "42," they can guide the Student better. But the paper found a flaw: If the answer is just a short number, the Teacher still doesn't know how to get there. The Teacher stays silent during the middle steps because the final answer doesn't give enough clues about the path. It's like telling a hiker, "The summit is at the top of the mountain," but not showing them the trail. The hiker still gets lost in the woods.

The New Solution: "Hindsight Self-Distillation" (HSD)
The authors propose a clever new trick called Hindsight Self-Distillation. Instead of just looking at the final answer, the Teacher looks at a successful peer.

Here is how it works in the classroom:

The teacher asks 8 students (the AI generates 8 different attempts) to solve the same problem.
7 students fail. 1 student succeeds.
The Teacher takes the successful student's full essay and shows it to the failing students.
The Teacher says: "Look at Student #4. They got the right answer. You two (Student #1 and #2) were following the exact same path as Student #4 for the first 50 words. But then, at word 51, you took a wrong turn. That is where you need to change your thinking."

Why This Works (The "Divergence" Moment)
The magic happens at the divergence point—the exact moment where the failing student's path splits away from the successful student's path.

Before the split: The Teacher is quiet because both students were doing the same thing.
At the split: The Teacher screams, "Stop! Look at the successful path! You went left, but you should have gone right!"
After the split: The Teacher's guidance fades because the failing student is now on a totally different track.

This gives the AI a precise "credit score" for every single word. It tells the AI exactly which word caused the failure, rather than just saying "the whole essay was bad."

The Results
The paper tested this on two powerful AI models (Qwen3-8B and Qwen3-32B) using math and coding problems.

The Test: They used difficult benchmarks like AIME (math competitions) where answers are short numbers.
The Outcome: HSD beat every other method, including the standard "thumbs up/thumbs down" methods and the old "answer-only" teacher methods.
The Big Win: The improvement was biggest on the hardest, shortest-answer tasks. This proves that giving the AI a "successful peer's path" to compare against is much better than just giving it the final answer.

In Summary
Instead of just grading the final exam, this method lets the AI learn by comparing its own failed attempts against a successful attempt made by its "sibling" in the same training session. It highlights the exact moment the logic went off the rails, allowing the AI to learn much faster and more accurately.

Technical Summary: Localizing Credit at the Divergence

1. Problem Statement

The paper addresses a fundamental bottleneck in Reinforcement Learning from Verifiable Rewards (RLVR) for Large Language Models (LLMs): token-level credit assignment.

In standard RLVR pipelines (e.g., GRPO), a group of rollouts is sampled, scored by a verifier, and updated based on a scalar advantage. This scalar is identical for every token in a trajectory. Consequently, while the model learns that a trajectory failed, it cannot identify where the failure originated. A long reasoning trace might contain a correct setup, a single erroneous algebraic step, and a final answer derived from that error; the scalar reward treats the entire sequence uniformly, failing to penalize the specific point of divergence.

On-policy self-distillation attempts to solve this by using the model as both a student (conditioned on the prompt) and a teacher (conditioned on the prompt plus privileged context, typically the ground-truth answer $a^\star$ ). The Kullback-Leibler (KL) divergence between the teacher and student distributions provides a dense, per-token gradient.

However, the authors identify a critical flaw in the standard choice of privileged context ( $c = a^\star$ ): The Endpoint Context Bottleneck.

On tasks with terse answers (e.g., AIME, where the answer is a 3-digit number), the ground-truth answer constrains only the final token span.
It provides no information about the intermediate reasoning path required to reach that answer.
Empirically, answer-conditioned teachers produce a "silence" (near-zero predictive shift) across the majority of the reasoning chain, with signals appearing only near the final answer. This leaves the model without guidance at the exact positions where path-level correction is most needed.

2. Methodology: Hindsight Self-Distillation (HSD)

The authors propose Hindsight Self-Distillation (HSD), a method that shifts the privileged context from an endpoint (the answer) to a path (a successful peer trajectory).

Core Mechanism

Instead of conditioning the teacher solely on the ground-truth answer $a^\star$ , HSD conditions the teacher on a successful peer rollout ( $o^+$ ) drawn from the same training group, alongside the answer.

Input: For a failing rollout $o_i$ , if the group contains a successful rollout $o^+$ , the teacher is conditioned on $c_i = (a^\star, o^+)$ .
Fallback: If no peer in the group succeeds ( $S_i = \emptyset$ ), the method falls back to the standard endpoint context $c_i = a^\star$ .
Cost: This requires no additional sampling. The successful peer is already generated as part of the standard group-relative sampling process (e.g., GRPO's group of $G$ rollouts).

Theoretical Justification

Ideal Credit: The ideal per-token credit should favor tokens that increase the probability of eventual verifier success.
Path-Level Information: A successful peer $o^+$ is an exact sample from the success-conditioned policy $\pi^+$ . By providing $o^+$ as context, the teacher gains access to the full successful continuation of the reasoning path.
Localization: When a failing rollout $o_i$ $o_{i}$ and a successful peer $o^+$ $o^{+}$ share a prefix up to a divergence point $\tau$ $τ$ :
- Pre-divergence ( $t < \tau$ ): The teacher and student see the same prefix. The credit is near zero (the "silence floor").
- At divergence ( $t = \tau$ ): The teacher sees the successful continuation $o^+_{\tau}$ , while the student generates the failing token. The KL divergence spikes, creating a strong, localized gradient signal exactly at the point of error.
- Post-divergence ( $t > \tau$ ): The signal decays as the student's prefix drifts further from the demonstration.

3. Key Contributions

Diagnosis of Endpoint Context Failure: The paper demonstrates that answer-conditioned self-distillation fails on terse-answer benchmarks because the answer imposes an information ceiling ( $\log V_k$ ) that is insufficient to constrain intermediate reasoning paths. The teacher remains "silent" on the reasoning steps that matter most.
Hindsight Self-Distillation (HSD): A minimal, one-line modification to the training loop that conditions the teacher on a successful peer trajectory. This transforms the credit assignment landscape, localizing the gradient signal to the divergence position $\tau$ .
Empirical Validation: Controlled comparisons on Qwen3-8B and Qwen3-32B across math (MATH, GSM8K, AIME) and code (APPS, LiveCodeBench) benchmarks.
- HSD achieves state-of-the-art results against GRPO variants and on-policy distillation baselines (OPSD, SDPO, RLSD).
- Gains are most pronounced on terse-answer tasks (e.g., AIME), where the endpoint context is weakest. For instance, HSD exceeds the GRPO baseline by >5 points on AIME-24 and >7 points on AIME-24 for the 32B model.
- The empirical credit profile confirms the theory: near-zero values on shared prefixes and a sharp spike at the divergence point.

4. Results and Analysis

Performance: HSD outperforms all baselines on math and code benchmarks at both model scales. The improvement over the strongest non-HSD baseline (a hybrid of GRPO and OPSD) is consistent, suggesting that path-based supervision adds value beyond linear combinations of sparse rewards and endpoint distillation.
Credit Localization: Analysis of the per-token credit profile shows that HSD concentrates roughly 55% of its credit mass within ±8 tokens of the divergence point $\tau$ . In contrast, baseline methods (OPSD) spread credit uniformly or remain in the silence band.
Coverage: The effectiveness of HSD depends on the existence of a successful peer. The probability of having a successful peer peaks when the group success rate $p$ is moderate (around 0.25 for group size 8). This aligns with the observation that HSD yields the largest gains on AIME (where success rates are low-to-moderate) compared to HumanEval+ (where success rates are high, leaving fewer failing rollouts to supervise).
Compute Cost: HSD incurs a modest overhead (approx. 25% more wall-clock time per step than OPSD) due to the longer teacher context, but no additional rollouts are required.

5. Significance and Claims

The paper claims that the primary limitation of current on-policy distillation is not the method of distillation, but the choice of privileged context. By moving from an endpoint context (answer) to a path context (successful peer), HSD unlocks the ability to localize credit assignment without requiring external reward models or additional sampling.

The authors position HSD as a simple, effective modification that resolves the "silence" of teachers on terse-answer tasks. The work suggests that the "hindsight" of a successful peer within the same training group is a sufficient and powerful signal for guiding reasoning models, effectively turning the group-relative sampling mechanism into a source of dense, path-level supervision.

Limitations Acknowledged:

HSD relies on the existence of a successful peer; if the entire group fails, it falls back to the weaker endpoint signal.
The method is evaluated on math and code with programmatic verifiers; its applicability to noisy rewards, multi-turn interactions, or instruction-tuned starting points is not yet established.
The "hard tail" of difficulty (where success rates are extremely low) may suffer from low coverage, a potential area for future work involving buffering past successes.