Sparse Knowledge Distillation: A Mathematical Framework for Probability-Domain Temperature Scaling and Multi-Stage Compression
Aaron R. Flouro, Shawn P. Chadwick
TL;DR
This work reframes knowledge distillation as an operator-level problem operating in probability space, anchored by probability-domain softening axioms. It shows that multiple non-unique operator families satisfy the axioms, enabling logit-free KD especially under partial access and privacy constraints. The authors prove operator-agnostic bias--variance bounds, a homotopy-based view of multi-stage compression with $O(1/n)$ convergence, and a classification of equivalence classes of softened-target operators. The framework provides practical guidance for black-box teacher distillation and privacy-preserving compression while offering universal guarantees across operator choices. Together, these results illuminate when and why sparse students can outperform dense teachers and how iterative pruning achieves reliable compression in a probability-domain setting.
Abstract
We develop a unified theoretical framework for sparse knowledge distillation based on probability-domain softening operators. While the equivalence $p^{1/T} \propto \mathrm{softmax}(z/T)$ is well known, our contribution is an operator-level analytical framework built on this foundation rather than the equivalence itself. The framework comprises four core components: (i) operator-agnostic bias--variance decompositions that characterize when sparse students outperform dense teachers, (ii) a homotopy path formalization of multi-stage pruning in function space explaining why iterative compression succeeds where one-shot pruning fails, (iii) convergence guarantees establishing $O(1/n)$ rates for $n$-stage distillation with explicit parameter dependence, and (iv) equivalence class characterizations identifying distinct probability-domain operators that yield identical student models under capacity constraints. We introduce an axiomatic definition of probability-domain softening operators based on ranking preservation, continuity, entropy monotonicity, identity, and boundary behavior, and show that multiple non-equivalent operator families satisfy these axioms. All learning-theoretic guarantees are shown to hold uniformly across this operator class, independent of implementation details. These results provide theoretical grounding for black-box teacher distillation, partial-access settings such as top-$k$ truncation and text-only outputs, and privacy-preserving model compression.