SGD-Based Knowledge Distillation with Bayesian Teachers: Theory and Guidelines
This work provides theoretical insights and practical guidelines for improving Knowledge Distillation, particularly in applications requiring stable and accurate model compression, though it is incremental in nature.
The paper tackles the theoretical understanding of Knowledge Distillation by analyzing SGD convergence with Bayesian teachers, showing that using exact Bayes Class Probabilities reduces variance and improves bounds compared to one-hot supervision. It demonstrates experimentally that students distilled from Bayesian teachers achieve up to +4.27% higher accuracy and up to 30% more stable convergence.
Knowledge Distillation (KD) is a central paradigm for transferring knowledge from a large teacher network to a typically smaller student model, often by leveraging soft probabilistic outputs. While KD has shown strong empirical success in numerous applications, its theoretical underpinnings remain only partially understood. In this work, we adopt a Bayesian perspective on KD to rigorously analyze the convergence behavior of students trained with Stochastic Gradient Descent (SGD). We study two regimes: $(i)$ when the teacher provides the exact Bayes Class Probabilities (BCPs); and $(ii)$ supervision with noisy approximations of the BCPs. Our analysis shows that learning from BCPs yields variance reduction and removes neighborhood terms in the convergence bounds compared to one-hot supervision. We further characterize how the level of noise affects generalization and accuracy. Motivated by these insights, we advocate the use of Bayesian deep learning models, which typically provide improved estimates of the BCPs, as teachers in KD. Consistent with our analysis, we experimentally demonstrate that students distilled from Bayesian teachers not only achieve higher accuracies (up to +4.27%), but also exhibit more stable convergence (up to 30% less noise), compared to students distilled from deterministic teachers.