Diagnosing and Mitigating Thinking Collapse in On-Policy Self-Distillation

2026-07-12

AI Summary

This work addresses “reasoning collapse” in on-policy self-distillation for complex reasoning tasks, where a model’s innate intermediate reasoning capabilities degrade during training. The study formally defines and quantifies this phenomenon, identifying its root cause as excessive suppression of student gradients by the teacher at high-entropy decision points. To mitigate this, the authors propose OPS-D, an adaptive dual-perspective framework featuring an asymmetric pointwise divergence gating mechanism that dynamically anchors distillation targets—preserving error-correction benefits while safeguarding native reasoning pathways. Integrating entropy-aware gradient masking, token-level target analysis, and frozen base-model priors, the method achieves an average accuracy gain of 4.1% across multiple mathematical reasoning benchmarks, substantially alleviating reasoning collapse and demonstrating robust generalization across model scales and post-training paradigms.

Abstract

On-Policy Self-Distillation (OPSD) has emerged as a crucial paradigm for enhancing and aligning Large Language Models (LLMs). However, in complex reasoning tasks, OPSD paradoxically degrades downstream performance. In this paper, we systematically investigate this pathology and identify a severe optimization trap we define as \textbf{Thinking Collapse} -- a sharp decline in the model's native intermediate reasoning behavior, measured by epistemic-token density (ET per 1k). Through entropy-based gradient masking and token-level target analysis, we show that this collapse is triggered by aggressive teacher gradients at high-student-entropy decision forks, where student epistemic tokens are frequently suppressed into teacher non-epistemic targets and are highly concentrated in high pointwise student-teacher divergence regions. To resolve this optimization pathology, we propose \textbf{Adaptive Dual-Perspective OPSD (AD-OPSD)}, a robust control framework that dynamically moderates the self-distillation objective. AD-OPSD selectively anchors high-suppression-risk sandboxed tokens to a reference prior derived from the frozen base model via an asymmetrical pointwise divergence gate, preserving native thinking capacity while retaining OPSD's error-correcting power. Extensive experiments across competitive mathematical benchmarks show that AD-OPSD improves over standard OPSD by up to \textbf{+4.1\%} absolute average accuracy across diverse model scales and datasets. Further analysis demonstrates that AD-OPSD mitigates thinking collapse and generalizes robustly to different post-training paradigms.

Problem

• On-Policy Self-Distillation
• Thinking Collapse
• Large Language Models
• Reasoning Degradation
• Optimization Trap

Innovation

• Thinking Collapse
• On-Policy Self-Distillation
• Epistemic-Token Density
• Adaptive Dual-Perspective OPSD
• Gradient Masking