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

2026-05-13

AI Summary

This work addresses the unreliability of knowledge transfer when distilling multi-step reasoning capabilities from large language models to smaller student models, a challenge exacerbated by noisy reasoning chains, hallucinated supervision signals, and static teacher–student interactions. To overcome these limitations, the authors propose GateKD, a novel confidence-gated closed-loop distillation framework. GateKD dynamically routes reliable signals from the teacher based on prediction confidence, aligns high-confidence intermediate representations, and filters unstable attention patterns. The approach integrates confidence-gated soft supervision, gated hidden state evolution, and reliability-aware attention distillation into a cohesive closed-loop system. Experimental results demonstrate that GateKD significantly outperforms existing open-loop distillation methods across commonsense, logical, and symbolic reasoning benchmarks, with particularly strong gains in logical/symbolic tasks and low-resource settings. Ablation studies further confirm the contribution of each component.

Abstract

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.

Problem

• reasoning distillation
• hallucination
• teacher-student interaction
• robust reasoning
• noisy rationales

Innovation

• confidence-gated distillation
• closed-loop distillation
• reasoning transfer
• hallucination suppression
• dynamic teacher-student interaction