Purified OPSD: On-Policy Self-Distillation Without Losing How to Think

Abstract: 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.

• The paper finds that standard on-policy self-distillation degrades question-conditioned reasoning in long chain-of-thought (CoT) models.
• It introduces a PMI-anchored two-step correction leveraging a reference-only teacher to isolate transferable signals.
• Experiments show that Purified OPSD improves accuracy and preserves reflective epistemic markers in various reasoning benchmarks.

Purified OPSD: On-Policy Self-Distillation Without Losing How to Think

Introduction

This work systematically diagnoses and addresses a core failure mode of on-policy self-distillation (OPSD) when applied to long chain-of-thought (long-CoT) reasoning models. While OPSD has been successful in aligning compact LMs with reasoning capabilities exhibited in larger teacher models, it exhibits divergent and sometimes pathological behavior on long-CoT models, degrading the reflective, epistemic processes crucial for robust generalization and deliberative reasoning. The analysis leads to a new, theoretically-grounded distillation paradigm, Purified OPSD, informed by pointwise mutual information (PMI), that rectifies these failings without introducing significant computational cost.

Analysis of OPSD Failure in Long-CoT Reasoning

Long-CoT models are distinguished not only by their ability to produce extended, structured reasoning traces but also by their use of epistemic markers—tokens denoting uncertainty and self-reflection (e.g., "Wait", "Let me think", "Perhaps"). These markers serve as signals of generalizable reasoning strategies rather than rote pattern matching. However, empirical analysis reveals that standard OPSD either suppresses these markers or escalates degenerate repetition, fundamentally disrupting the epistemic behavior that underpins generalization.

Quantitative decomposition of the teacher’s update signal demonstrates that the supervision provided by the privileged teacher in OPSD is dominated by a reference-induced component—the portion of the gradient field aligned with the reference answer—but not necessarily with the question or the correct inferential trajectory. This reference-specific shortcut not only fails to support question-conditioned, inference-transferable improvements but actively displaces them, causing the student to overfit to privileged solution heuristics and lose general problem-solving skills. Cosine similarity analysis confirms that the alignment of the teacher’s update with the inference-transferable component is negligible or negative throughout most of training, confirming the dominance of pathology-inducing reference-based signals.

Purified OPSD: PMI-Anchored Distillation Target Construction

To eliminate reference-specific memorization and preserve general-purpose reasoning, the authors introduce a two-step corrective scheme:

1. Reference-Only Teacher: The first step is to isolate the non-transferable, reference-induced component by constructing a reference-only teacher—i.e., the base model conditioned on the solution reference but not the question—thereby making explicit the component of the gradient that can never transfer to deployment conditions where ground-truth references are unavailable.
2. PMI Transformation: The residual between the full privileged teacher update and the reference-only teacher update is a per-token log-probability difference, reflecting the question-conditioned correction. This is interpreted as a conditional PMI signal, which is then combined with the clean base model distribution (conditioned only on the question, never the reference) to form a purified distillation target:

$$P_{\text{PMI}}(v) \propto P_0(v) \exp\left(\beta \; (\log T_T(v) - \log T_{\text{ref}}(v))\right)$$

where $P_0$ is the base model distribution, $T_T$ the teacher output, $T_{\text{ref}}$ the reference-only teacher, $\beta$ a correction strength parameter, and $v$ a vocabulary token.

A numerically robust, stabilized PMI target is constructed by centering and softly clipping extreme residuals. This target is then used in place of the raw teacher in the JSD distillation objective.

Experimental Results

Extensive evaluation across four long-CoT models (Qwen3-8B, Qwen3-4B, R1-Distill-7B, OLMo-7B) and two challenging datasets (DASD-10K, Math-CoT-20K) demonstrates that:

• Standard OPSD yields negligible or negative gains, often destabilizing reflective reasoning and degrading accuracy across multiple mathematical reasoning benchmarks.
• Purified OPSD with the PMI target consistently improves accuracy relative both to the base model and OPSD-Standard, robustly across all model-dataset combinations.
• Reflective behaviors, as measured by the distribution and magnitude of epistemic markers, are preserved by Purified OPSD but severely disrupted by standard OPSD.
• The PMI-based approach is robust to reasonable choices of key hyperparameters (clipping threshold and correction strength), and its performance is not sensitive to checkpoint selection, reducing operational complexity.

Theoretical Implications

The decomposition approach establishes that naive use of privileged reference solutions for token-level alignment can be counterproductive in reflective, deliberate reasoning regimes, in contrast to success stories in shorter, less epistemically rich tasks. The PMI grounding situates the method within the literature on KL-regularized policy improvement, as the purified distillation target emerges as an information-theoretically optimal solution to the constrained reward-maximization problem where transferability, not privileged alignment, is the objective.

Practical Implications and Future Directions

Practically, Purified OPSD offers a simple drop-in replacement for OPSD that avoids the catastrophic collapse of epistemic reasoning in long-CoT distillation. The computational cost is minimal, requiring only additional frozen-model forward passes for PMI computation.

Potential avenues for further investigation include extending this decomposition/Purification approach to other domains with privileged supervision (e.g., multimodal LMs, RL-from-human-feedback with intermediate signals), exploration of more nuanced residual decomposition techniques, and understanding the impact on sample efficiency and implicit representation learning of reasoning features.

Conclusion

This study identifies and solves a core pathology of on-policy self-distillation for long-CoT models: the dominance of non-transferable, reference-induced gradients that corrupt reflective, generalizable reasoning. By leveraging a reference-only teacher and PMI-based target construction, Purified OPSD ensures that only question-conditioned, inference-transferable corrections are distilled, yielding superior reasoning accuracy and stable epistemic behavior across diverse models and benchmarks. The results highlight the necessity of carefully separating transferable and non-transferable supervision in privileged distillation for robust AI reasoning transfer.