Source-linked AI summary
Purified OPSD: On-Policy Self-Distillation Without Losing How to Think
Zhanming Shen, Jintao Tong, Shaotian Yan, Chen Shen, Hao Chen, Wentao Ye, Xiaomeng Hu, Rui Miao, Haobo Wang, Junbo Zhao, Gang Chen, Jieping Ye
TL;DR
OPSD can fail on long-CoT models because reference-induced supervision dominates the transferable reasoning signal and destabilizes reflective behavior. The paper removes that component with a reference-only teacher, converts the residual into a PMI target, and reports consistent improvements over both the base model and standard OPSD across four models and two datasets.
Problem
Standard OPSD yields marginal or negative results on long-CoT models while destabilizing the reflective reasoning capability they rely on.
Method
A reference-only teacher isolates the non-transferable component, and PMI transforms the residual question-conditioned signal into a target distribution for student distillation.
Results
Across four long-CoT models and two datasets, OPSD-PMI consistently improves over both the base model and standard OPSD while preserving reflective reasoning throughout training.
Takeaways & Limitations
Filtering reference-induced shortcuts provides useful supervision for long-CoT models without losing their natural epistemic behavior.
Abstract
from arXiv · showhide
On-policy self-distillation (OPSD) has emerged as a promising paradigm for improving LLM reasoning, where a privileged teacher with access to reference solutions provides token-level supervision on the student's own generated trajectories. However, we find that OPSD consistently fails on long chain-of-thought (long-CoT) reasoning models, yielding at best marginal gains while destabilizing the reflective reasoning capability these models depend on. Through a novel decomposition of the teacher's supervision signal, we identify the root cause: the teacher's supervision is dominated by a reference-induced component that drives rote memorization of reference-specific shortcuts, while the question-conditioned, inference-transferable component is ignored or actively opposed. Based on this diagnosis, we propose a two-step solution. First, we construct a reference-only teacher (the same model conditioned on the reference without the question) to isolate the non-transferable component of the supervision signal; the residual after subtracting this component captures the question-conditioned, inference-transferable correction. Second, we use pointwise mutual information (PMI) as the mechanism to transform this residual into a well-formed PMI target distribution that the student can directly distill from, filtering out the reference-induced shortcut. Experiments on four long-CoT models across two datasets demonstrate consistent improvements over both the base model and standard OPSD, while preserving the models' natural epistemic behavior throughout training.
1 Introduction
Long-CoT models expose a consistent weakness in standard OPSD: gains are marginal or transient, performance and reflective reasoning can degrade, and the teacher’s update is dominated by reference-specific supervision. The paper proposes removing that shortcut with a reference-only teacher and converting the residual into a distillable PMI target.
- Long-CoT reasoning improves complex-task performance, making its distillation into smaller models a central challenge.
- OPSD provides at best marginal, short-lived gains on long-CoT models and may degrade performance.
- During OPSD, epistemic markers show pathological, model-dependent fluctuations associated with destabilized reflective reasoning.
- The privileged teacher may drive rote memorization of reference-specific reasoning paths instead of transferable corrections.
- A reference-only teacher decomposes supervision into reference-induced and question-conditioned components, with the former dominating update direction and magnitude.
- The proposed residual removes the reference-specific shortcut and is transformed with PMI into a target distribution for direct student distillation.
2 Revisiting On-Policy Self-Distillation for Long-CoT Models
The paper revisits OPSD for long-CoT models and finds that its privileged-reference signal correlates with performance and epistemic instability. A reference-only teacher exposes the non-transferable component, while the residual represents question-conditioned supervision that can be purified for learning.
- 2.1 Preliminaries: OPSD trains an on-policy student from token-level evaluations by a privileged teacher sharing the base model but conditioned on a reference solution.
- 2.1 Preliminaries: PMI measures the association between two events beyond what independence would predict.
- 2.2 OPSD Degrades Long-CoT Model Performance: Across four long-CoT models on Math-CoT-20K, standard OPSD yields marginal, short-lived gains and may degrade performance.
- 2.4 Diagnosing the Root Cause: Rote Memorization of Privileged Information: A reference-only teacher isolates supervision available from the reference alone, separating non-transferable reference-induced information from potentially useful question-conditioned information.
- 2.4 Diagnosing the Root Cause: Rote Memorization of Privileged Information: The reference-induced component dominates update direction and magnitude, while the inference-transferable component partially cancels it.
- 2.4 Diagnosing the Root Cause: Rote Memorization of Privileged Information: As reference memorization saturates, the inference-transferable signal partially recovers, but accuracy continues declining after early reasoning damage.
- 2.4 Diagnosing the Root Cause: Rote Memorization of Privileged Information: The residual ∆it = log πT −log πref is a conditional PMI-style quantity that captures question-conditioned corrections but requires transformation into a distribution.
3 Method
The method removes the reference-induced shortcut from OPSD supervision and converts the remaining question-conditioned correction into a PMI target anchored to the clean base distribution. It stabilizes this target before distilling it into the student, preserving the base reasoning prior while adding transferable corrections.
- Signal decomposition: The reference-induced component is dominated by the question-conditioned residual, motivating its removal from OPSD supervision.The residual is interpreted as the inference-transferable correction, but it is initially only a log-probability difference rather than a valid distribution.
- PMI target construction: PMI converts the question-conditioned residual into a well-formed target distribution.The residual is computed as log πT(v) − log πref(v), then anchored to the clean base distribution.
- PMI target construction: β controls correction strength: β = 1 applies the full correction, while larger β produces a target closer to the base distribution.The target starts from a reference-free base distribution and adjusts it using the question-conditioned residual.
- Optimality: The PMI target is the closed-form optimum of a KL-regularized objective that maximizes inference-transferable reward while staying close to the clean base distribution.The clean base distribution serves as the reference policy, and the residual serves as the implicit reward.
- Stabilized implementation: The practical procedure centers and softly clips the residual, then adds it to base-model log-probabilities and normalizes the result.Centering removes global logit shifts, while soft clipping bounds extreme PMI values; the experiments use c = 10 and β = 1 unless otherwise specified.
- Training procedure: Training uses generalized Jensen-Shannon divergence between the student distribution and the stabilized PMI target, with two additional forward passes and no extra trainable parameters.The added passes probe the reference-only teacher and base distribution, increasing wall-clock training time by less than 10% in the reported implementation.
4 Experiments
Across four long-CoT models and two datasets, OPSD-PMI consistently improves over both the base model and OPSD-Standard while maintaining stable training dynamics and reflective reasoning behavior. Ablations further show robustness to the soft-clipping threshold and correction strength.
- Main results: OPSD-PMI improves over both the base model and OPSD-Standard across all four models and both training datasets.OPSD-Standard provides negligible or negative gains, whereas OPSD-PMI achieves gains on every model-dataset combination.
- Main results: OPSD-Standard degrades three of four models on DASD-10K and three of four models on Math-CoT-20K, while OPSD-PMI remains consistently superior.The direct gap between the two methods is large and consistent across models and datasets.
- Training dynamics: OPSD-PMI remains stable across training and above the baseline, whereas OPSD-Standard peaks briefly before steadily declining.The smaller checkpoint variance reduces reliance on careful early stopping.
- Epistemic marker analysis: OPSD-PMI preserves epistemic behavior: Qwen3-8B and R1-Distill-7B stay near baseline marker counts while OPSD-Standard collapses or explodes.OPSD-Standard changes marker distributions substantially, whereas OPSD-PMI closely preserves the base distribution.
- Ablation studies: All tested soft-clipping thresholds improve over baseline and show similar trajectories, indicating robustness to c.The c ablation compares c ∈ {5, 10, 20} with β = 1 fixed.
- Ablation studies: All tested correction strengths improve over baseline and OPSD-Standard, although β values produce different trajectories and checkpoint-specific peaks.β = 1 is used in the main experiments for simplicity and to avoid additional tuning.
5 Related Work
Prior work established long-CoT reasoning and on-policy self-distillation as important approaches for improving and transferring LLM reasoning. This paper situates its contribution as identifying and addressing a failure mode specific to applying OPSD to long-CoT models.
- Long chain-of-thought reasoning: Long-CoT reasoning extends chain-of-thought prompting with extended reflective processes for complex reasoning tasks.These processes include self-correction, backtracking, and uncertainty externalization.
- On-policy self-distillation: The paper identifies and solves a fundamental failure mode of on-policy self-distillation specific to long-CoT reasoning models.The failure mode concerns applying OPSD to models whose reasoning depends on reflective processes.
- On-policy self-distillation: OPSD has students generate their own trajectories and receive token-level feedback from privileged teachers conditioned on reference solutions.This directly addresses distribution mismatch associated with off-policy distillation.
6 Conclusion
The paper attributes OPSD's failure on long-CoT models to reference-induced supervision that overwhelms transferable reasoning, then proposes a PMI-based purification procedure. Across four long-CoT models and two datasets, the method improves over the base model and standard OPSD while preserving reflective reasoning.
- Diagnosis: Reference-induced supervision dominates OPSD's training signal, suppressing question-conditioned corrections and driving rote memorization instead of transferable reasoning improvement.This accounts for both performance degradation and epistemic-marker destabilization.
- Solution: The method constructs a reference-only teacher, isolates the non-transferable component, and uses PMI to produce a purified target distribution for distillation.The residual question-conditioned signal is anchored to the base distribution within the standard OPSD framework.
- Results: Across four long-CoT models and two training datasets, the approach consistently improves over both the base model and standard OPSD while preserving reflective reasoning throughout training.The ablation studies also support robustness to key hyperparameters.