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

Abstract: On-Policy Self-Distillation (OPSD) has emerged as a crucial paradigm for enhancing and aligning 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.

• The paper demonstrates that on-policy self-distillation can induce a localized collapse of reasoning tokens, leading to reduced accuracy in mathematical tasks.
• It employs entropy-based gradient masking and token-level KL divergence analysis to identify high-uncertainty regions where teacher suppression harms epistemic token density.
• The proposed AD-OPSD method dynamically interpolates between teacher and base model targets, recovering thinking density and delivering up to a +6.7% accuracy gain.

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

Introduction and Motivation

On-Policy Self-Distillation (OPSD) has become a prevalent fine-tuning mechanism for aligning and improving LLMs on complex reasoning tasks. While OPSD provides dense, token-level supervision by aligning the student policy with the ground-truth-conditioned teacher along the student’s sampled trajectories, empirical evidence shows that standard OPSD can degrade downstream mathematical reasoning performance relative to the pre-trained baseline. This paper introduces and systematically diagnoses the phenomenon of Thinking Collapse: a sharp reduction of the model's intermediate reasoning behavior, as measured by epistemic-token density, under OPSD.

Empirical Identification of Thinking Collapse

The authors define thinking density as the number of epistemic tokens (ETs) per 1000 generated tokens—where ETs cover natural language markers for planning, uncertainty, and verification. Through controlled experiments with Qwen3-1.7B on AIME25, it is observed that OPSD with both thinking teacher and student (T/T) triggers a catastrophic drop in thinking density—from 10.3 (base) to 7.9 ET per 1k—and a concomitant accuracy drop of 4.2%.

Figure 1: AIME25 accuracy of Qwen3-1.7B degrades sharply as thinking density collapses under T/T on-policy distillation.

Furthermore, disabling student thinking during OPSD (NT/T) yields stronger accuracy and preserved thinking density, suggesting that the observed accuracy improvement is closely linked to the model’s ability to maintain exploratory reasoning tokens. However, NT/T creates confounds due to cross-mode distillation. Therefore, subsequent analysis focuses on the in-mode (T/T) OPSD setup.

Mechanistic Diagnosis via Gradient Masking and Token-Level Analysis

To localize the cause of thinking collapse, the paper deploys a three-pronged diagnostic strategy:

1. Entropy-Based Gradient Masking: By selectively masking teacher gradients based on the student’s tokenwise entropy, it is shown that suppressive teacher signals are spatially localized at high-entropy (i.e., high-uncertainty) tokens. Masking just the top 20% entropy tokens substantially restores thinking density but at a cost to accuracy—unveiling an optimization deadlock where binary masking cannot avoid either collapse or loss of corrective signals.
2. Microscopic Discrepancy Analysis: At initialization, the majority (66.3%) of ET-to-Non-ET suppression events occur within the top 10% highest-entropy tokens; essentially, native student epistemic behaviors are highly concentrated at uncertain decision forks, precisely where teacher targets diverge and suppression risk is high.
3. Pointwise KL Divergence Mapping: The local pointwise KL divergence between student and teacher token probabilities ($C_i = P_s(y_i) [\log P_s(y_i) - \log P_t(y_i)]$) robustly identifies regions where epistemic tokens cluster and suppression risk is maximal.

Figure 2: Entropy-based masking demonstrates that targeting high student-entropy tokens can almost entirely recover native thinking density; most suppression events are concentrated in the highest-entropy percentiles, confirming the sharp locality of the pathology.

Figure 3: Epistemic tokens under active teacher suppression ($P_s > P_t$) are extremely concentrated in regions with high pointwise KL; this provides a precise real-time gate for adaptive interventions.

Adaptive Dual-Perspective OPSD (AD-OPSD): Framework and Implementation

To break the accuracy–reasoning preservation trade-off, the authors introduce Adaptive Dual-Perspective OPSD (AD-OPSD). This framework dynamically interpolates the distillation target between the teacher and the frozen base model (the epistemic anchor) through a soft, asymmetrical gating mechanism:

• Dual-Perspective Target: For each token, the target is computed as a convex combination of teacher and base model distributions, with the weight determined by the unreliability index $U_i$.
• Adaptive Gating: $U_i$ is determined by (a) restricting gating to high-entropy ("sandboxed") regions, and (b) applying a sigmoid over the pointwise KL divergence, but only where $P_s > P_t$.
• Continuous Shielding: This design ensures only those high-risk regions suffer teacher suppression get protected, allowing standard factual correction elsewhere.

This method is compared to a static blend (SD-OPSD, $U_i=0.5$ constant in the sandbox) and several alternative gating baselines. The paper implements all variants using LoRA parameter-efficient adapters with frozen teacher and base checkpoints for stabilization.

Experimental Results and Ablations

Main Results

AD-OPSD systematically and consistently outperforms standard OPSD across three model families—Qwen3-1.7B, Qwen3-4B, and DeepSeek-R1-1.5B—on four competitive mathematical benchmarks.

• For Qwen3-1.7B, AD-OPSD delivers +4.1% absolute average accuracy gain over standard OPSD.
• On difficult tasks like AIME 2025, the accuracy jump can reach +6.7%.

Detailed Analysis of Reasoning Behavior

AD-OPSD recovers a significant portion of collapsed thinking density (e.g., from 7.9→8.6 on AIME 2025 with Qwen3-1.7B), with accuracy rising in parallel. Static blending is effective but underperforms compared to the dynamic pointwise KL-based gating.

Robustness, Generalization, and Hyperparameter Sensitivity

• Non-Thinking Paradigms: Even when both teacher and student are non-thinking, AD-OPSD improves accuracy for both non-thinking and thinking inference, demonstrating generalization to latent reasoning alignment.
• Context Length: The relative gain from AD-OPSD is most pronounced in shorter contexts, aligning with recent observations that early reasoning tokens concentrate the critical difference. The margin narrows for longer contexts due to diminishing base-teacher divergence.
• Sandbox Size Ablation: Optimal results are obtained with a 20–30% high-entropy sandbox; beyond 40%, accuracy tapers due to dilution of corrective teacher signals.
• Gating Function Ablation: Sigmoid pointwise KL gating is the most robust; alternatives such as global entropy or plain log-ratio underperform or are unstable.

Figure 4: Downstream accuracy and relative gains of AD-OPSD hold for increasing maximum context lengths; sandbox size controls the trade-off between preserving native reasoning and maintaining teacher correction strength.

Theoretical and Practical Implications

The analysis rigorously demonstrates that the core failure mode of on-policy self-distillation in reasoning domains is a highly localized, entropy-driven suppression of exploratory student behavior by the teacher. This contradicts claims that reasoning post-training failure is a global train-test mismatch, and instead points to the precise locality of the collapse at high-entropy decision forks. By anchoring the distillation objective to the model’s unsuppressed prior selectively in these regions, AD-OPSD achieves substantial and robust accuracy gains without discarding the benefits of dense OPSD correction.

From a practical perspective, AD-OPSD is straightforward to implement as a post-training fix and is compatible with Parameter-Efficient Fine-Tuning protocols. It neither requires modification of the underlying LLM architecture nor incurs any inference time overhead.

Limitations and Prospects for Future Research

The experiments are conducted with moderate model scales (up to Qwen3-4B) and on competitive mathematics benchmarks. Extension to ultra-large-scale LLMs and to domains outside symbolic mathematics—such as code synthesis, logical reasoning, or open-ended multi-hop question answering—represents a compelling direction. The approach of localized dual-perspective anchoring based on dynamic risk indicators (e.g., pointwise KL) potentially generalizes as a recipe for robust post-training in other compositional reasoning tasks.

Conclusion

Diagnosing and mitigating reasoning degradation in on-policy self-distillation, this paper shows that teacher-induced suppression of epistemic tokens at high-entropy forks is the central cause of thinking collapse. By employing adaptive, token-level, dual-perspective anchoring governed by pointwise student-teacher divergence, AD-OPSD preserves the model’s native reasoning capacity and yields strong accuracy improvements over standard post-training methods, setting a new benchmark in controllable and robust reasoning alignment (2607.10805).