Purified OPSD: On-Policy Self-Distillation Without Losing How to Think
For researchers training large language models on reasoning tasks, this work addresses a critical failure mode of OPSD on long-CoT models, offering a principled fix that improves performance without destabilizing reasoning.
On-policy self-distillation (OPSD) fails on long chain-of-thought reasoning models due to a reference-induced component that drives rote memorization. The proposed method, Purified OPSD, isolates the transferable supervision signal using a reference-only teacher and pointwise mutual information, achieving consistent improvements (e.g., +2-5% accuracy) across four models and two datasets while preserving reflective reasoning.
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.