Functional Ensemble Distillation
Bayesian models have many desirable properties, most notable is their ability
to generalize from limited data and to properly estimate the uncertainty in
their predictions. However, these benefits come at a steep computational cost
as Bayesian inference, in most cases, is computationally intractable. One
popular approach to alleviate this problem is using a Monte-Carlo estimation
with an ensemble of models sampled from the posterior. However, this approach
still comes at a significant computational cost, as one needs to store and run
multiple models at test time. In this work, we investigate how to best distill
an ensemble's predictions using an efficient model. First, we argue that
current approaches that simply return distribution over predictions cannot
compute important properties, such as the covariance between predictions, which
can be valuable for further processing. Second, in many limited data settings,
all ensemble members achieve nearly zero training loss, namely, they produce
near-identical predictions on the training set which results in sub-optimal
distilled models. To address both problems, we propose a novel and general
distillation approach, named Functional Ensemble Distillation (FED), and we
investigate how to best distill an ensemble in this setting. We find that
learning the distilled model via a simple augmentation scheme in the form of
mixup augmentation significantly boosts the performance. We evaluated our
method on several tasks and showed that it achieves superior results in both
accuracy and uncertainty estimation compared to current approaches.