Introduction
Deep Neural Networks (DNNs) are often criticized for being overconfident, frequently making incorrect predictions with high certainty. Deep ensembles have emerged as the gold standard for uncertainty quantification by training multiple models with different initializations and using their predictive variance as a proxy for uncertainty. However, ensembles are computationally expensive, requiring significant memory and compute power, which makes them impractical for edge devices like smartphones or drones.
Standard Knowledge Distillation (KD) attempts to compress these ensembles into a single "student" model. However, traditional KD typically collapses the ensemble's variance into a single point estimate, losing critical uncertainty information. This paper introduces Gaussian Distillation, a method that leverages a Deep Latent Factor (DLF) model to compress an entire ensemble into a single student distribution, preserving reliable uncertainty at a fraction of the cost.
Methodology: The Deep Latent Factor Model
The core innovation of this work is treating ensemble members as independent realizations of a stochastic process rather than static weights. The authors propose modeling the teacher ensemble using a Deep Latent Factor (DLF) architecture.
In this framework, the student model outputs the parameters of a Gaussian distribution—specifically the mean (\mu) and a covariance structure modeled through factor loadings (\phi). This approach allows the model to capture complex relationships between outputs without the parameter explosion associated with full covariance matrices.
To ensure training stability and performance, the authors implement two key strategies:
- Expectation-Maximization (EM) Algorithm: Used to estimate the Maximum Likelihood of the student DNN, providing a more stable training objective than standard gradient descent on complex distributions.
- MMD Initialization: Maximum Mean Discrepancy (MMD) is employed to ensure that the latent factors follow a standard Gaussian distribution, providing a robust starting point for the EM algorithm.
Experimental Evaluation
The proposed method was evaluated across diverse tasks, including regression (UCI benchmarks), image classification (CIFAR-10/100, Tiny ImageNet), and Natural Language Processing (GLUE/SuperGLUE using RoBERTa).
Key Results
Gaussian Distillation consistently outperformed existing baselines such as Hydra, Batch Ensembles (BE), and Dirichlet-based distillation methods.
As shown in the results for image classification, the DLF student model not only provides better uncertainty quantification (lower NLL and ECE) but also frequently improves top-line accuracy. For instance, on CIFAR-100, the DLF model achieved 79.68% accuracy, significantly higher than other ensemble distillation techniques.
In regression tasks, the efficiency of the DLF model was particularly evident. On the Boston housing dataset, the student model required only ~1,800 parameters, a massive reduction from the ~90,000 parameters used by the original teacher ensemble, while maintaining competitive performance.
Conclusion and Future Directions
Gaussian Distillation offers a robust and efficient path for deploying reliable AI on resource-constrained devices. By viewing ensembles as a distribution of functions and employing a latent factor model, the authors successfully distilled the "wisdom of crowds" into a single, lightweight student model.
Future work includes extending this framework to Bayesian DNNs to provide a more formal probabilistic grounding. Additionally, integrating distillation with Low-Rank Adaptation (LoRA) for large language models and exploring online Bayesian learning are promising avenues for further research.