SGD-Based Knowledge Distillation with Bayesian Teachers: Theory and Guidelines

TL;DR

The paper addresses why knowledge distillation benefits from probabilistic teacher outputs by formulating a Bayesian, SGD-centered view of KD. It develops a rigorous analysis distinguishing two regimes: exact Bayes class probabilities and noisy probability estimates, showing variance reduction and interpolation advantages when supervision aligns with the true Bayes posterior. The authors advocate using Bayesian deep learning as KD teachers and demonstrate through theory and experiments (e.g., CIFAR-100) that Bayesian teachers lead to higher accuracy and more stable convergence, even in noisy settings or with limited data. This work thus links teacher calibration to optimization dynamics, providing concrete guidelines for constructing Bayesian KD pipelines with practical impact on model compression and generalization. Overall, the study offers a principled framework connecting Bayesian calibration, stochastic optimization, and distillation efficacy, with demonstrated gains in both convergence and generalization.

Abstract

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.

Figures (7)

• Figure 1: Learning curves, test accuracies, and average generalization error gaps in the learning curves when training with one-hot labels, when supervised with the true bcp, when supervised with true bcp corrupted with different noise levels, and supervision by combinations of one-hot labels and noisy bcp adjusted based on several $\lambda$ values. Complete experimental details are provided in Appendix \ref{['app:toy_experiment']}.
• Figure 2: Test accuracy stability over the last 50 training epochs with standard one-hot training, deterministic teachers, and Bayesian teachers.
• Figure 3: Effect of noise level $\epsilon$ on student performance. Each plot shows results for students trained with noisy bcp. Top: average generalization error (left) and variability in generalization error (right). Bottom: average test accuracy (left) and variability in accuracy (right). Each plot includes a fit proportional to $\dfrac{1}{1+\varepsilon}$, alongside the result achieved by a student trained from One-hot labels, and a student trained with the true bcp.
• Figure 4: Correlations between performance and stability metrics across different noise levels $\epsilon$. Each point corresponds to a student trained with noisy bcp, with baselines for one-hot labels and true bcp also included. Top left: higher accuracy is strongly associated with lower generalization error. Top right: models with noisier generalization error curves also display noisier accuracy curves. Bottom left: higher accuracy coincides with reduced variability in accuracy. Bottom right: lower generalization error coincides with reduced generalization error variability.
• Figure 5: The accuracy achieved by students (averaged over 5 runs) trained from both deterministic and Bayesian teachers in kd with different teacher and student temperatures applied. The teacher-student pairs displayed are ResNet-50 $\rightarrow$ ResNet-18 (left), ResNet-50$\to$WRN-40-2 (middle), and WRN-40-2$\to$WRN-16-2 (right).

Theorems & Definitions (20)

• Definition 1: Interpolation
• Definition 2: Gradient noise
• Proposition 1
• Proposition 2
• Lemma 1
• Theorem 1
• Theorem 2
• Theorem 3
• Theorem 4
• Proposition 3