Trust the uncertain teacher: distilling dark knowledge via calibrated uncertainty

TL;DR

Calibrated Uncertainty Distillation is proposed, a framework designed to make dark knowledge more faithfully accessible by directly shaping the teacher's predictive distribution before transfer, allowing students to benefit from both confident signals on easy cases and structured uncertainty on hard ones.

Abstract

The core of knowledge distillation lies in transferring the teacher's rich 'dark knowledge'-subtle probabilistic patterns that reveal how classes are related and the distribution of uncertainties. While this idea is well established, teachers trained with conventional cross-entropy often fail to preserve such signals. Their distributions collapse into sharp, overconfident peaks that appear decisive but are in fact brittle, offering little beyond the hard label or subtly hindering representation-level transfer. This overconfidence is especially problematic in high-cardinality tasks, where the nuances among many plausible classes matter most for guiding a compact student. Moreover, such brittle targets reduce robustness under distribution shift, leaving students vulnerable to miscalibration in real-world conditions. To address this limitation, we revisit distillation from a distributional perspective and propose Calibrated Uncertainty Distillation (CUD), a framework designed to make dark knowledge more faithfully accessible. Instead of uncritically adopting the teacher's overconfidence, CUD encourages teachers to reveal uncertainty where it is informative and guides students to learn from targets that are calibrated rather than sharpened certainty. By directly shaping the teacher's predictive distribution before transfer, our approach balances accuracy and calibration, allowing students to benefit from both confident signals on easy cases and structured uncertainty on hard ones. Across diverse benchmarks, CUD yields students that are not only more accurate, but also more calibrated under shift and more reliable on ambiguous, long-tail inputs.

Figures (4)

• Figure 1: Illustration of the main framework. Left: Difficulty-aware Uncertainty Shaping (DUS) adjusts weights based on the correctness of the model’s predictions, reducing overconfidence and encouraging a more calibrated probability distribution. Right: Wrong-mass Clipping (W-Clip) modifies the probability assigned to incorrect classes when the model makes a wrong prediction, providing a more stable and reliable distribution for the student.
• Figure 2: Effect of calibration on uncertainty and distributions.
• Figure 3: Effect of DUS on predictive calibration. Across datasets, DUS reduces overconfident predictions by reshaping the teacher distribution, leading to a controlled underconfident shift. Although this increases ECE due to lower confidence than accuracy, the behavior is safer than overconfidence and yields more reliable uncertainty estimates for downstream OOD detection and risk-sensitive settings.
• Figure 4: Learning curves illustrating the convergence and generalization behavior of our method. Our approach demonstrates smoother optimization and stronger validation performance throughout training.

Theorems & Definitions (1)

• Theorem 1: Uniqueness of Projection