GateKD: Confidence-Gated Closed-Loop Distillation for Robust Reasoning

📄 arXiv: 2605.13136v1 📥 PDF

作者: Kasidit Sermsri, Teerapong Panboonyuen

分类: cs.CL

发布日期: 2026-05-13

备注: 16 pages


💡 一句话要点

提出GateKD:一种置信度门控的闭环蒸馏框架,用于提升鲁棒推理能力。

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 知识蒸馏 大型语言模型 鲁棒推理 置信度门控 闭环反馈 多步推理 模型压缩

📋 核心要点

  1. 现有推理蒸馏方法忽略了教师模型推理过程中的噪声和不确定性,导致学生模型学习到错误的中间推理。
  2. GateKD通过引入置信度门控机制,动态地评估教师模型的可靠性,并有选择地进行知识转移,从而减少幻觉和噪声的影响。
  3. 实验表明,GateKD在常识、逻辑和符号推理任务上显著优于现有方法,尤其是在逻辑和符号推理方面提升明显。

📝 摘要(中文)

将大型语言模型(LLMs)的多步推理能力提炼到紧凑的学生模型中仍然具有挑战性,这主要是由于噪声推理、幻觉监督和静态的师生交互。现有的推理蒸馏方法,包括基于导师的方法,主要以开环方式运行,隐式地假设教师具有一致的可靠性,从而传播错误的中间推理。我们提出了GateKD,一个置信度门控的闭环蒸馏框架,通过将教师视为动态的门卫而不是静态的预言机,从而实现鲁棒的推理转移。GateKD引入了三个互补的机制:(i)置信度门控的软监督,选择性地提炼可靠的预测信号;(ii)门控隐藏状态演化,仅当教师置信度高时才对齐中间表示;(iii)可靠性过滤的注意力蒸馏,在抑制噪声模式的同时保留稳定的推理结构。这些组件共同形成一个闭环反馈,其中教师置信度不断调节蒸馏过程,减少幻觉转移并稳定学生推理。在常识、逻辑和符号推理基准上的大量实验,使用不同大小的T5和Flan-T5骨干网络,表明GateKD始终优于强大的开环蒸馏基线。值得注意的是,GateKD在逻辑和符号推理方面产生了显著的收益,在低资源蒸馏设置下保持鲁棒性,并且当任何门控组件被移除时,性能会明显下降。我们的结果表明,置信度门控的闭环监督对于构建可靠且可扩展的小型推理模型至关重要。

🔬 方法详解

问题定义:现有推理蒸馏方法,特别是开环方法,假设教师模型是完全可靠的,忽略了教师模型在多步推理过程中可能产生的错误或不确定性。这导致学生模型学习到教师模型产生的噪声推理过程,降低了推理的准确性和鲁棒性。现有方法缺乏动态评估和过滤教师模型输出的能力,无法有效应对教师模型的幻觉问题。

核心思路:GateKD的核心思路是将教师模型视为一个动态的“门卫”,而不是一个静态的“预言机”。通过引入置信度评估机制,GateKD能够动态地判断教师模型在每个推理步骤中的可靠性。只有当教师模型的置信度足够高时,才将相应的知识转移给学生模型。这种选择性的知识转移能够有效地减少噪声和幻觉的影响,提高学生模型的推理能力。

技术框架:GateKD包含三个主要模块:(1) 置信度门控的软监督:根据教师模型的置信度,选择性地将教师模型的预测结果作为软标签来指导学生模型的训练。(2) 门控隐藏状态演化:仅当教师模型的置信度较高时,才对齐教师模型和学生模型的中间隐藏状态,从而避免将教师模型中的错误信息传递给学生模型。(3) 可靠性过滤的注意力蒸馏:通过过滤掉教师模型中不稳定的注意力模式,保留稳定的推理结构,从而提高学生模型的推理能力。这三个模块形成一个闭环反馈系统,教师模型的置信度不断调节蒸馏过程。

关键创新:GateKD的关键创新在于引入了置信度门控机制,将教师模型的可靠性纳入到蒸馏过程中。与传统的开环蒸馏方法不同,GateKD能够动态地评估教师模型的可靠性,并有选择地进行知识转移。这种闭环反馈机制能够有效地减少噪声和幻觉的影响,提高学生模型的推理能力。

关键设计:GateKD使用教师模型的预测概率作为置信度指标。对于置信度门控的软监督,使用一个阈值来过滤掉置信度较低的预测结果。对于门控隐藏状态演化,使用一个门控函数来控制隐藏状态的对齐程度,门控函数的输入是教师模型的置信度。对于可靠性过滤的注意力蒸馏,使用一个稳定性指标来评估注意力模式的稳定性,并过滤掉不稳定的注意力模式。损失函数包括预测损失、隐藏状态对齐损失和注意力蒸馏损失,这些损失函数都受到教师模型置信度的调节。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,GateKD在常识、逻辑和符号推理任务上均优于现有的开环蒸馏方法。例如,在逻辑推理任务上,GateKD相比于基线方法提升了5%以上。此外,GateKD在低资源蒸馏设置下也表现出良好的鲁棒性,表明其能够有效地利用有限的数据进行知识迁移。消融实验表明,移除任何一个门控组件都会导致性能下降,验证了各个组件的有效性。

🎯 应用场景

GateKD可应用于各种需要多步推理的场景,例如问答系统、对话系统、知识图谱推理等。通过将大型语言模型的推理能力迁移到小型模型,可以降低计算成本和部署难度,使得这些应用能够在资源受限的环境中运行。GateKD还有助于提高模型的鲁棒性和可靠性,减少幻觉问题,从而提升用户体验。

📄 摘要(原文)

Distilling multi-step reasoning abilities from large language models (LLMs) into compact student models remains challenging due to noisy rationales, hallucinated supervision, and static teacher-student interactions. Existing reasoning distillation methods, including mentor-based approaches, predominantly operate in an open-loop manner, implicitly assuming uniform teacher reliability and consequently propagating erroneous intermediate reasoning. We propose GateKD, a confidence-gated closed-loop distillation framework that enables robust reasoning transfer by treating the teacher as a dynamic gatekeeper rather than a static oracle. GateKD introduces three complementary mechanisms: (i) confidence-gated soft supervision that selectively distills reliable predictive signals, (ii) gated hidden-state evolution that aligns intermediate representations only when teacher confidence is high, and (iii) reliability-filtered attention distillation that preserves stable reasoning structures while suppressing noisy patterns. These components jointly form a closed feedback loop in which teacher confidence continuously modulates the distillation process, reducing hallucination transfer and stabilizing student reasoning. Extensive experiments across commonsense, logical, and symbolic reasoning benchmarks, using T5 and Flan-T5 backbones of varying sizes, demonstrate that GateKD consistently outperforms strong open-loop distillation baselines. Notably, GateKD yields substantial gains in logical and symbolic reasoning, remains robust under low-resource distillation settings, and shows clear performance degradation when any gating component is removed. Our results highlight that confidence-gated closed-loop supervision is critical for building reliable and scalable small reasoning models.