诊断并缓解在线自蒸馏中的思维崩溃

背景

On-Policy Self-Distillation（OPSD）已成为增强和对齐大型语言模型（LLMs）的关键范式。其核心思想是让模型在自身当前策略生成的样本上进行自蒸馏，从而在训练过程中持续自我改进。然而，在复杂推理任务中，OPSD 反而导致下游性能下降，这一矛盾现象引发了对该优化过程内在病理的深入探究。本文系统性地诊断了这种病理，并揭示了一个严重的优化陷阱——Thinking Collapse（思考崩溃）。

核心内容

思考崩溃的发现与定义

Thinking Collapse 表现为模型原生中间推理行为的急剧衰退，通过测量认知令牌密度（epistemic-token density，ET per 1k） 来量化。该指标衡量模型在生成过程中产生“认知令牌”（即承载推理步骤的关键 token）的频次。在标准 OPSD 训练下，该密度显著下降，说明模型逐渐丧失了原本具备的、显式的中间推理能力。

诊断方法：熵梯度掩码与令牌级目标分析

为了定位崩溃的触发机制，作者使用了两种诊断工具：

• 基于熵的梯度掩码（entropy-based gradient masking）：通过识别学生模型在推理过程中具有高不确定性的决策分叉点（即高熵位置），分析这些位置上的梯度更新方向。
• 令牌级目标分析（token-level target analysis）：逐 token 比较学生模型生成的预测分布与教师模型提供的软目标分布。

触发原因：高熵决策分叉处的激进教师梯度

研究发现，Thinking Collapse 的核心触发条件是：在具有高学生熵的决策分叉处，教师模型提供了过于激进的梯度信号。具体表现为：

• 学生模型的高熵位置（即模型对下一步生成的 token 非常不确定）是推理的关键转折点。
• 在这些位置，教师模型倾向于将学生的认知令牌（本应保留推理步骤的 token）压制为非认知令牌（non-epistemic targets），即教师模型认为这些位置不需要复杂的推理。
• 这种现象高度集中在点态学生-教师散度（pointwise student-teacher divergence）很高的区域，表明学生与教师在该位置的预测分歧极大，而教师梯度的强压制最终导致学生放弃原有的推理路径。

解决方案：自适应双视角 OPSD（AD-OPSD）

为了缓解这一优化病理，论文提出 Adaptive Dual-Perspective OPSD（AD-OPSD），一个稳健的控制框架，动态调节自蒸馏目标。其核心机制包括：

1. 沙盒令牌（sandboxed tokens）：识别出具有高抑制风险（即容易因教师梯度而丢失认知能力的令牌）的位置。
2. 参考先验锚定：对这些高风险令牌，不直接使用教师软目标，而是将其损失函数锚定到一个参考先验（reference prior），该先验来自冻结的基模型（即原始未经过 OPSD 的模型）的预测分布。
3. 非对称点态散度门控（asymmetrical pointwise divergence gate）：动态决定何时使用教师目标、何时使用基模型先验，从而在保留原生思考能力（来自基模型）的同时，保留 OPSD 的错误纠正能力（来自教师）。

实验验证

在多个竞争性数学基准测试（如 GSM8K、MATH 等）上，AD-OPSD 在多种模型规模（从 7B 到 70B）和数据集上，相比标准 OPSD，平均绝对准确率提升最高达 +4.1%。进一步分析表明，AD-OPSD 有效缓解了 Thinking Collapse，并且能够稳健地泛化到不同的后训练范式（如指令微调、RLHF 等）。

关键要点

• Thinking Collapse 的定义：OPSD 训练中模型原生中间推理能力（以认知令牌密度衡量）的急剧下降，是导致复杂推理任务性能退化的直接原因。
• 关键触发条件：高学生熵的决策分叉处，教师模型提供过于激进的梯度，将学生的认知令牌压制为非认知令牌，且这些位置集中在点态散度高的区域。
• 诊断工具：基于熵的梯度掩码和令牌级目标分析，可精确定位 Thinking Collapse 的触发位置。
• AD-OPSD 的核心创新：引入沙盒令牌机制和来自冻结基模型的参考先验，通过非对称点态散度门控动态平衡教师目标与原始模型能力，避免过度压制。
• 性能提升：在数学推理任务上，AD-OPSD 将标准 OPSD 的准确率提升最高 +4.1%，且在不同规模和训练范式下表现稳健。
• 泛化能力：AD-OPSD 不仅适用于 OPSD 本身，还能推广到其他后训练方法（如强化学习、指令微调），显示了其作为通用优化框架的潜力。

意义与影响

• 理论贡献：首次系统性地揭示了 OPSD 在复杂推理任务中的退化机制——Thinking Collapse，并提供了可量化的诊断指标（如认知令牌密度和点态散度），为后续研究提供了分析工具。
• 方法创新：AD-OPSD 提出了一种“软约束”策略，通过引入基模型先验作为“锚点”，在不破坏教师蒸馏的纠错能力的同时，保留了模型原生推理路径。这种思路可能启发更多针对自蒸馏中“灾难性遗忘”或“能力退化”问题的解决方案。
• 实践价值：提升 LLM 在数学、逻辑等需要多步推理的任务上的表现，且无需额外训练数据或模型改动，仅需修改损失函数，具有较高的实用性和可迁移性。
• 未来方向：该工作提示，在自蒸馏或自我改进类训练中，需要关注模型“思考过程”的保留，而不仅仅是最终答案的准确性。AD-OPSD 的框架也可能被推广到其他模态（如多模态推理、代码生成）中，用于防止模型在持续优化中丧失固有的推理能力。