通过动态Token选择实现分布对齐自蒸馏的鲁棒推理

# 通过动态词元选择实现分布对齐自蒸馏的鲁棒推理  
来源: https://arxiv.org/html/2606.00628  
Ruiqi Zhang1,2,\*, Lingxiang Wang1,2, Hainan Zhang1,2, Zhiming Zheng1,2  
1北京航空航天大学未来区块链与隐私计算北京高级创新中心  
2北京航空航天大学人工智能学院  
通讯作者: zhanghainan@buaa\.edu\.cn (https://arxiv.org/html/2606.00628v1/mailto:[email protected])  

###### 摘要  
自蒸馏通过将参考答案重写为更符合模型自身分布的训练数据来提高学习效率。然而，参考答案也引入了强烈的风格偏差，导致生成模型模仿表面形式而非学习有用的推理模式。我们观察到，重写数据中包含大量高困惑度（PPL）词元，这些词元来自两个不同的来源：有益的知识增强型逻辑修正，以及由参考模仿引起的有害风格偏移。同等对待所有此类词元可能会破坏基础模型的原始分布，并降低性能，特别是在困难的推理任务上。为了解决这个问题，我们提出了分布对齐自蒸馏（DASD），它使用一个答案感知的参考模型生成候选词元，并根据基础模型的置信度动态过滤它们。DASD 保留了编码有用逻辑知识的词元，同时抑制了分布不对齐的风格噪声。在数学、代码和常识推理基准上的实验表明，DASD 始终优于竞争基线，减少了高 PPL 词元，并提高了跨不同难度任务的鲁棒性。

# 通过动态词元选择实现分布对齐自蒸馏的鲁棒推理  
Ruiqi Zhang1,2,\*, Lingxiang Wang1,2, Hainan Zhang1,2, Zhiming Zheng1,2  
1北京航空航天大学未来区块链与隐私计算北京高级创新中心  
2北京航空航天大学人工智能学院  
通讯作者: zhanghainan@buaa\.edu\.cn (https://arxiv.org/html/2606.00628v1/mailto:[email protected])  

## 1 引言  
灾难性遗忘 [Luo et al. (2025) (https://arxiv.org/html/2606.00628#bib.bib1); Li et al. (2024) (https://arxiv.org/html/2606.00628#bib.bib2)] 仍然是大语言模型（LLM）后训练中的核心挑战。当下游数据与预训练分布差异很大时，直接微调可能会扭曲模型的原始参数空间。自蒸馏 [Yang et al. (2024) (https://arxiv.org/html/2606.00628#bib.bib3)] 通过将下游示例转换为更符合模型自身生成分布的训练样本来缓解这一问题，从而在减少分布偏移的同时提高训练效率。

参见图注  
图 1: 不同难度级别下平均高 PPL 词元（PPL > 2.5）比例与模型性能之间的相关性。x 轴表示低（级别 1 & 2）、中（级别 3 & 4）和高（级别 5）。

然而，自蒸馏并不总是分布对齐的。在许多下游任务中，尤其是困难的推理问题，生成的解决方案被限制为必须保留真实答案或参考推理过程。当基础模型缺乏足够的推理能力时，这个过程可能导致生成的数据继承参考特定的推理模式、模板和风格表达，而非反映基础模型的本机分布。如图 1 (https://arxiv.org/html/2606.00628#S1.F1) 所示，我们在 MATH 基准 [Hendrycks et al. (2021) (https://arxiv.org/html/2606.00628#bib.bib14)] 上的分层难度分析揭示，随着问题难度的增加，与基础模型不一致的词元比例急剧上升，而微调的效果则稳步下降。这表明自蒸馏可能提供了内容正确但分布不对齐的监督，削弱了知识传递并增加了遗忘的风险。

进一步的词元级分析表明，这种不一致来源于两个性质不同的来源，如图 4 (https://arxiv.org/html/2606.00628#S3.F4) 所示。一些高困惑度（PPL）词元代表有益的逻辑修正：这些推理步骤在基础模型下不太可能出现，但对于解决问题是必要的。另一些则反映有害的风格偏移：从参考答案中继承的冗余表面形式、推理模板或答案特定表达，它们贡献的有用知识很少。朴素的自我蒸馏对所有生成的词元进行统一优化，因此无法区分有用的修正和嘈杂的风格偏差。

现有的词元选择策略只能部分解决这个问题。Masked-PPL [Wu et al. (2026) (https://arxiv.org/html/2606.00628#bib.bib4)] 和 ProFit [Liu et al. (2026) (https://arxiv.org/html/2606.00628#bib.bib5)] 方法从训练中移除高 PPL 词元，这抑制了风格噪声但同时也丢弃了有价值的推理修正。Hint-decoding [Zhang et al. (2026a) (https://arxiv.org/html/2606.00628#bib.bib6)] 方法使用参考模型的不确定性来结合基础和参考分布，但答案条件化的参考模型可能过于自信，导致生成仍然偏向参考解决方案。这些限制要求一种更具选择性的机制，既能保留有用的推理偏差，又能过滤掉分布不对齐的噪声。

在本文中，我们提出 DASD，一种基于置信度的动态词元选择方法，用于鲁棒的分布对齐自蒸馏。DASD 构建参考模型的候选词元以保持答案正确性，同时使用基础模型的置信度来确定每个词元在分布上是否可接受。基础模型本地熟悉的词元被优先选择以维持分布对齐，并且一个强制回退机制保留了必不可少的推理词元，即使它们在基础模型下不太可能出现。通过这种方式，DASD 保留了有益的逻辑修正，同时抑制了有害的风格偏移。在数学推理、代码生成和常识推理基准上的实验表明，DASD 始终优于竞争性的自我蒸馏基线。它减少了高 PPL 风格偏差，提高了跨难度级别的鲁棒性，并在后训练期间更好地保留了预训练知识¹¹¹https://anonymous\.4open\.science/r/emnlp\-SD\-6AB2\.。

我们总结主要贡献如下：
- • 我们揭示了自我蒸馏对参考答案的严重依赖，并验证了任务难度和分布偏移如何影响性能。
- • 我们提出了 DASD，一个分布对齐的自蒸馏框架，具有置信度感知的动态词元选择，能在知识注入和分布保持之间取得平衡。
- • 我们在多种推理基准上评估 DASD，并表明它在很大程度上优于传统的自我蒸馏基线，尤其是在困难的推理任务上。

## 2 相关工作  
##### 分布一致的数据选择  
已经有大量工作探索使用合成自生成数据训练大语言模型 [Wang et al. (2023) (https://arxiv.org/html/2606.00628#bib.bib7)]。在合成数据构建中，训练样本与目标模型之间的分布对齐被广泛认为是关键，因为模型在熟悉的数据上学习更高效，同时可以减轻预训练知识的侵蚀 [Ren et al. (2024) (https://arxiv.org/html/2606.00628#bib.bib9)]。为了保持模型生成的风格一致性，SCAR [Li et al. (2025) (https://arxiv.org/html/2606.00628#bib.bib10)] 通过风格感知排序过滤高质量指令数据，用更少的训练样本稳定模型性能。除了标准困惑度，自对齐 PPL [Ren et al. (2025) (https://arxiv.org/html/2606.00628#bib.bib11)] 度量被提出用于衡量生成内容与模型固有推理模式之间的一致性，从而促进更好的数据选择。像 Mask-PPL [Wu et al. (2026) (https://arxiv.org/html/2606.00628#bib.bib4)] 和 ProFit [Liu et al. (2026) (https://arxiv.org/html/2606.00628#bib.bib5)] 这样的方法直接丢弃高困惑度词元以保留分布内内容，但它们遭受知识注入不足，限制了模型获取新推理能力。

##### 自蒸馏  
在此基础上，使用模型生成数据的自训练已被广泛研究。自蒸馏 [Yang et al. (2024) (https://arxiv.org/html/2606.00628#bib.bib3)] 利用合成数据的固有分布一致性，有效地缓解了微调过程中的分布偏移。Beyond Human Data [Singh et al. (2023) (https://arxiv.org/html/2606.00628#bib.bib8)] 从模型中生成样本，通过二元反馈过滤，对正确样本进行微调，并重复这个过程。在代码生成任务中，SSD [Zhang et al. (2026b) (https://arxiv.org/html/2606.00628#bib.bib12)] 通过平衡格式范式和逻辑推理稳定地提高了模型性能，即使对自生成数据没有严格的正确性过滤。为了确保额外的知识注入，STaR [Zelikman et al. (2024) (https://arxiv.org/html/2606.00628#bib.bib13)] 在模型出错时提示其通过提供正确答案来重新生成答案。Self-Distillation 工作引入了带有推理路径的参考答案来辅助模型生成，同时确保答案正确性和分布接近性。随后的 Hint-Decoding [Zhang et al. (2026a) (https://arxiv.org/html/2606.00628#bib.bib6)] 研究试图在词元级别区分分布内和分布外内容：它通过熵融合基础和参考模型的输出，然而输出容易受到参考答案的干扰，导致置信度膨胀，无法准确区分风格词元和逻辑词元。

## 3 动机  
自蒸馏采用参考答案作为外部指导来确保生成内容的正确性。然而，这种机制导致模型机械地模仿参考特定的推理路径，导致生成内容偏离其原始分布。一个典型现象是自生成过程中广泛出现高困惑度词元。为了揭示其不利影响，我们在 MATH 数据集 [Hendrycks et al. (2021) (https://arxiv.org/html/2606.00628#bib.bib14)] 上进行了难度分层实验。结果表明，任务难度越高，相应高 PPL 词元的比例越高，模型性能下降越远。

见图注  
图 2: MATH 基准中不同难度答案的高 PPL 率分布，横坐标为 5 个难度级别，纵坐标为比例。

### 3.1 蒸馏数据中的高 PPL 词元  
我们在 MATH 数据集上进行了难度分层分析。我们采用 Llama3.2-3B-Instruct 作为基础模型 Mθ\\mathcal{M}_{\\theta}，并构建一个参考模型 Mref\\mathcal{M}_{\\text{ref}} 来生成基于黄金答案的推理路径。我们从基础模型计算词元级困惑度，并经验性地将困惑度 > 2.5 的词元视为高 PPL 词元，表示与基础模型的分布不匹配。然后我们评估了跨难度级别的高 PPL 词元比例。图 2 (https://arxiv.org/html/2606.00628#S3.F2) 展示了不同难度下的段级 PPL 分布。简单问题以低 PPL 内容为主，而超高 PPL（≥\\geq 15）段随着难度增加而持续增长。对于 1 级简单样本，PPL 低于 9 的序列占大多数，超高 PPL 内容仅占 11.89%。相比之下，5 级困难样本的低 PPL 比例急剧下降，超高 PPL 段上升到 30.52%。这些结果表明，在参考引导的生成下，困难推理会引发严重的分布偏移。图 1 (https://arxiv.org/html/2606.00628#S1.F1) 进一步说明了高 PPL 比例与下游性能之间的相关性。随着高 PPL 词元的增加，模型性能逐渐下降。在简单数据上微调带来稳定收益，而困难级别的训练则使整体准确率下降 7%。这表明分布偏移损害了知识学习并破坏了模型的固有表示。朴素的自我蒸馏盲目拟合偏移的噪声词元，这成为复杂推理的关键限制。

见图注  
图 3: 不同难度级别下 Mθ\\mathcal{M}_{\\theta} 和 Mref\\mathcal{M}_{\\text{ref}} 之间的词元级对齐。所有指标均在从 Mref\\mathcal{M}_{\\text{ref}} 采样的完整序列上计算，包括成对 top-1 词元一致性、Mθ\\mathcal{M}_{\\theta} 的 top-1 词元在 Mref\\mathcal{M}_{\\text{ref}} 的 top-3/top-6 中的覆盖率，以及 Mθ\\mathcal{M}_{\\theta} 的 top-3 词元在 Mref\\mathcal{M}_{\\text{ref}} 的 top-3/top-6 中的覆盖率。

### 3.2 分布偏移现象  
为了分析参考模型和基础模型之间的预测差距和分布偏移，我们进行了跨五个难度级别的词元级对齐实验。我们从参考模型 Mref\\mathcal{M}_{\\text{ref}} 采样完整序列，并计算与基础模型 Mθ\\mathcal{M}_{\\theta} 的词元级重叠率。我们测量了五个统计量：两个模型一致的 top-1 词元、Mθ\\mathcal{M}_{\\theta} 的 top-1 词元在 Mref\\mathcal{M}_{\\text{ref}} 的 top-3 和 top-6 中被覆盖的情况，以及 Mθ\\mathcal{M}_{\\theta} 的 top-3 词元在 Mref\\mathcal{M}_{\\text{ref}} 的 top-3 和 top-6 中被覆盖的情况。如图 3 (https://arxiv.org/html/2606.00628#S3.F3) 所示，所有对齐指标随着推理难度的增加呈现稳定下降趋势。两个模型预测的相同 top 词元比例从 93.11% 下降到 90.79%，这表明高复杂度推理扩大了模型最自信预测的分歧。尽管整体词元重叠保持在较高水平，但基础模型的最优词元在参考 top-3 候选中的覆盖率从 99.17% 逐渐下降到 97.78%。类似地，更广泛的 base top-3 候选集的重叠程度在困难样本上也呈现缓慢但持续的下降。这些经验观察为我们的 DASD 方法提供了坚实的动机。一方面，极高的候选覆盖率证明基础模型的大多数原生高质量词元自然包含在参考候选空间中。在由 Mref\\mathcal{M}_{\\text{ref}} 约束的正确性边界内保留原始生成风格是可行的。另一方面，困难任务上不可忽视的预测差距验证了细粒度过滤的必要性。简单地遵循参考模型将不可避免地引入偏移词元，破坏基础模型的固有分布。因此，我们的方法不是僵硬的模仿或直接的词元替换，而是通过置信度校准和候选筛选动态平衡双模型输出，从而在保持复杂数学推理的逻辑正确性的同时减轻分布偏移。

见图注  
图 4: MATH 数据集上的一个生成案例。问题: 原始问题; 参考: 标准解决方案; 自: 基础模型输出; 自蒸馏: 自蒸馏结果; DASD: 我们方法的输出。黄色词元表示对参考答案的机械模仿，这些偏离