Uncertainty-based knowledge distillation for Bayesian deep neural network compression

Deep learning models have been widely employed across various fields. In real-world scenarios, especially safety-critical applications, quantifying uncertainty is as crucial as achieving high accuracy. To address this concern, Bayesian deep neural networks (BDNNs) emerged to estimate two different types of uncertainty: Aleatoric and Epistemic. Nevertheless, implementing a BDNN on resource-constrained devices poses challenges due to the substantial computational and storage costs imposed by approximation inference techniques. Thus, efficient compression methods should be utilized. We propose an uncertainty-based knowledge distillation method to compress BDNNs. Knowledge distillation is a model compression technique that involves transferring knowledge from a complex network, known as the teacher network, to a simpler one, referred to as the student network. Our method incorporates uncertainty into knowledge distillation to address situations where inappropriate teacher supervision undermines compression performance. We utilize the Epistemic uncertainty of teacher predictions to tailor supervision for each sample individually to take into account teacher's limited knowledge. Additionally, we adjust the temperature parameter of the distillation process for each sample based on the Aleatoric uncertainty of the teacher predictions, ensuring that the student receives appropriate supervision even in the presence of ambiguous data. As a result, the proposed method enables the Bayesian student network to be trained under both appropriate supervision of the Bayesian teacher network and ground truth labels. We evaluated our method on the CIFAR-10, CIFAR-100, and RAF-DB datasets, demonstrating notable improvements in accuracy over state-of-the-art knowledge distillation-based methods. Furthermore, the robustness of our approach was assessed through testing weakly trained teacher networks and the analysis of blurred and low-resolution data, which have high uncertainty. Experimental results show that the proposed method outperformed existing methods.