Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

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

About

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.

Kasidit Sermsri, Teerapong Panboonyuen• 2026

Related benchmarks

TaskDatasetResultRank
Commonsense ReasoningStrategyQA
Accuracy62.4
208
Commonsense ReasoningCSQA
CSQA Accuracy69.5
195
Symbolic ReasoningLast Letter Concatenation
Accuracy65.7
68
Logical reasoningShuffled Objects
Accuracy92.1
39
ReasoningBBH
BBH Score48.9
39
Symbolic ReasoningLast Letter
Accuracy0.679
31
Showing 6 of 6 rows

Other info

Follow for update