Technical Summary: Combating Supervision Fidelity Decay in On-Policy Distillation

Problem: Supervision Fidelity Decay (SFD)

The paper identifies a critical bottleneck in On-Policy Distillation (OPD), a paradigm where a student model learns from its own generated trajectories using token-level feedback from a teacher model. While OPD has been effective for transferring reasoning capabilities, the authors observe that its performance degrades significantly as the generation length increases (e.g., performance plateaus around 16k tokens and declines at 39k tokens).

The authors attribute this to Supervision Fidelity Decay (SFD). As the student generates longer sequences, its outputs progressively drift from the teacher's training distribution (Out-of-Distribution or OOD). This drift causes the teacher's next-token distribution to become less confident and less discriminative.

• Mechanism: In standard reverse-KL distillation, the learning signal (advantage) relies on the difference between the student's log-probability and the teacher's log-probability. As the teacher's distribution becomes diffuse (high entropy) due to SFD, the variance in the teacher's log-probabilities collapses.
• Consequence: The teacher's discriminative contribution to the gradient vanishes. The learning signal effectively reduces to a "student-only" signal that reinforces the student's existing modes without correction. This creates a self-reinforcing feedback loop where drift compounds, shrinking the effective learning horizon and establishing a ceiling on reasoning length.

Methodology: Lookahead Group Reward (LGR)

To address SFD, the authors propose Lookahead Group Reward (LGR), a method that actively maintains supervision quality by optimizing for future teacher confidence rather than current local matching.

Core Insight

Even when the teacher's local supervision at position $t$ has collapsed (i.e., the distribution is near-uniform), different token choices at $t$ can lead to significantly different teacher confidence levels at the next position ($t+1$). The paper proves theoretically that one-step-ahead discriminability can survive even when local supervision fails. By selecting tokens that maximize the teacher's confidence at the subsequent step, the student can slow the rate of drift.

Algorithm Design

1. Confidence Reward: At each position $t$, the method evaluates the teacher's peak next-token probability ($\max_v \pi_T(v|x_{<t}, x_t)$) for the student's top-$K$ candidate tokens. This value serves as a reward signal indicating how well a candidate preserves the teacher's ability to provide high-quality supervision in the next step.
2. Group Normalization: To handle variance in absolute confidence levels across different tasks and positions, the raw rewards are normalized within the group of top-$K$ candidates (similar to GRPO). This retains the relative ranking of candidates while removing absolute scale noise.
3. Entropy-Triggered Activation: Computing the lookahead reward for all tokens is computationally expensive. The authors observe that SFD and the need for lookahead guidance are most critical at high-entropy positions (where the student is uncertain). Therefore, the confidence reward is only computed for positions where the student's generation entropy exceeds a threshold $\tau$. This reduces computational overhead by 4–7 times.
4. Tree-Attention Mechanism: To efficiently evaluate $K$ candidates for multiple positions without re-running the teacher's prefill for every candidate, the authors employ a tree-attention mechanism. This constructs an extended sequence where the main branch is the student's trajectory and candidate tokens are appended as branches. A sparse attention mask allows the teacher to process all candidates in a single forward pass, achieving a theoretical speedup of up to 1000$\times$ compared to a naive implementation.

The final loss function combines the standard reverse-KL advantage with the lookahead confidence reward, weighted by a hyperparameter $\gamma$.

Key Contributions

1. Identification of SFD: The paper formally characterizes Supervision Fidelity Decay as a fundamental failure mode of on-policy reverse-KL distillation, proving that it collapses the gradient into a self-reinforcing drift loop that limits reasoning length.
2. Lookahead Group Reward (LGR): The authors propose a principled remedy that looks one step ahead. Instead of trying to match a degraded local distribution, LGR selects tokens that maximize future teacher confidence, effectively steering the student toward states where supervision remains reliable.
3. Efficient Implementation: The design of an entropy-triggered tree-attention mechanism makes the computationally intensive lookahead process practical, enabling training on long contexts (up to 39k tokens) without prohibitive costs.

Experimental Results

The method was evaluated on six benchmarks covering mathematical reasoning (AIME-24, AIME-25, AIME-26, HMMT-25, HMMT-26) and code generation (LiveCodeBench v6).

• Performance Gains: For a 7B student model, LGR improved the mean@8 score by 2.57 points over standard OPD (9k tokens) on average.
• Long-Context Scaling: The benefits of LGR scale with generation length. On the AIME-26 benchmark at 39k tokens, LGR achieved a gain of +4.92 points over OPD. In contrast, standard OPD performance degraded significantly at this length.
• Comparison to Baselines: LGR outperformed other distillation methods including JSD, REOPOLD, and standard OPD variants. Notably, methods that only stabilized training (like REOPOLD) without addressing the supervision quality degradation failed to match LGR's performance.
• Ablation Studies:
  • Max-p vs. Entropy: Using the teacher's max-probability (Max-p) as the confidence metric outperformed using entropy or perplexity, aligning with the theoretical finding that Max-p is a direct proxy for proximity to the teacher's in-distribution manifold.
  • Weight Sensitivity: The confidence reward weight $\gamma$ showed optimal performance in the range of 1.0 to 1.5. Excessively high weights led to reward hacking (repetitive sequences).

Significance and Claims

The paper claims that LGR provides a structurally different solution to the problem of long-horizon reasoning in distillation. Unlike strategies that passively avoid weak-supervision regions (e.g., truncation) or merely stabilize optimization, LGR actively steers the student toward trajectories where the teacher maintains high supervision fidelity.

The authors emphasize that their work addresses the root cause of performance decay in long reasoning chains: the compounding effect of student drift on teacher confidence. By operationalizing the insight that "higher next-position confidence implies a more discriminative future signal," LGR breaks the vicious cycle of supervision decay, allowing models to scale effectively to much longer reasoning chains (up to 39k tokens) than previously possible with standard on-policy distillation.