How Is Uncertainty Propagated in Knowledge Distillation?
Ziyao Cui, Jian Pei
TL;DR
This work analyzes how uncertainty propagates through knowledge distillation across linear models, neural networks, and large language models, revealing that standard single-response distillation suppresses intra-student uncertainty while leaving inter-student variability substantial. The authors formalize uncertainty into inter- and intra-student components, derive exact results for linear regression, and validate patterns in neural nets and LLMs, demonstrating that teacher output noise drives inter-student variance and that single-sample supervision underestimates teacher variance in LLMs. They propose two variance-aware strategies—averaging multiple teacher responses and inverse-variance weighting—that provably reduce target noise and yield minimum-variance estimators in simple settings and robust improvements in practice, including reduced hallucination in LLM distillation. The results recast distillation as an uncertainty transformation and show variance-aware distillation can produce more stable students that better reflect teacher uncertainty, with practical implications for safety-critical and high-stakes deployments. Conceptually, this work provides a principled framework and actionable techniques to balance fidelity to the teacher with faithful uncertainty preservation across distillation scales.
Abstract
Knowledge distillation transfers behavior from a teacher to a student model, but the process is inherently stochastic: teacher outputs, student training, and student inference can all be random. Collapsing these uncertainties to a single point estimate can distort what is learned. We systematically study how uncertainty propagates through knowledge distillation across three representative model classes--linear regression, feed-forward neural networks, and large language models (LLMs)--and propose simple corrections. We distinguish inter-student uncertainty (variance across independently distilled students) from intra-student uncertainty (variance of a single student's predictive distribution), showing that standard single-response knowledge distillation suppresses intra-student variance while leaving substantial inter-student variability. To address these mismatches, we introduce two variance-aware strategies: averaging multiple teacher responses, which reduces noise at rate $O(1/k)$, and variance-weighting, which combines teacher and student estimates via inverse-variance weighting to yield a minimum-variance estimator. We provide formal guarantees in linear regression, validate the methods in neural networks, and demonstrate empirical gains in LLM distillation, including reduced systematic noise and hallucination. These results reframe knowledge distillation as an uncertainty transformation and show that variance-aware distillation produces more stable students that better reflect teacher uncertainty.