S ELF -D ISTRIBUTION D ISTILLATION : E FFICIENT U NCERTAINTY E STIMATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Deep learning is increasingly being applied in safety-critical domains.
For these scenarios it is important to know the level of uncertainty in a model’s prediction to ensure that appropriate decisions are made by a system.
Deep ensembles are the de-facto standard approach to obtaining various measures of uncertainty.
However, ensembles normally signiﬁcantly increase the resources required in both the training and deployment phases.
Approaches have been developed that typically address the costs in one of these phases.
In this work we propose a novel training approach, self-distribution distillation (S2D), which is able to efﬁciently, both in time and memory, train a single model that can estimate uncertainties in an integrated training phase.
Furthermore it is possible to build ensembles of these models and apply ensemble distillation approaches, hierarchical distribution distillation, in cases where one is less limited by computational resources in the training phase, but still requires efﬁciency in the deployment phase.
Experiments on CIFAR-100 showed that S2D models outperformed standard models and Monte-Carlo dropout.
Additional out-of-distribution detection experiments on LSUN, Tiny ImageNet, SVHN showed that even a standard deep ensemble can be outperformed using S2D based ensembles and novel distilled models.
I NTRODUCTION
Neural networks (NNs) have enjoyed much success in recent years achieving state-of-the-art performance on a large number of tasks within domains such as natural language processing (Vaswani et al., 2017),
speech recognition (Hinton et al., 2012) and computer vision (Krizhevsky et al., 2012).
Unfortunately, despite the prediction performance of NNs they are known to yield poor estimates of the uncertainties in their predictions—in knowing what they do not know (Lakshminarayanan et al., 2017; Guo et al., 2017).
With the increasing application of neural network based systems in performing safety-critical tasks such as biometric identiﬁcation (Schroff et al., 2015), medical diagnosis (De Fauw et al., 2018) or fully autonomous driving (Kendall et al., 2019), it becomes increasingly important to be able to robustly estimate the uncertainty in a model’s prediction.
By having access to accurate measures of prediction uncertainty, a system can act in a more safe and informed manner.
Ensemble methods, and related schemes, have become the standard approach for uncertainty estimation.
Lakshminarayanan et al.
(2017) proposed generating a deep (random-seed) ensemble by training each member model with a different initialisation and stochastic gradient descent (SGD).
Not only does this ensemble perform signiﬁcantly better than a standard trained NN, it also displays better predictive uncertainty estimates.
Although simple to implement, training and deploying an ensemble results in a linear increase in the computational cost.
Alternatively Gal & Ghahramani (2016) introduced the Monte Carlo (dropout) ensemble (MC ensemble) which at test time estimates predictive uncertainty by sampling members of an ensemble using dropout.
Though this approach generally does not perform as well as a deep ensemble (given the same computational power and neglecting memory) (Lakshminarayanan et al., 2017), it is signiﬁcantly cheaper to train as it integrates the ensemble generation method into training.
Despite ensemble generation methods being computationally more expensive, they have an important ability to decompose predictive (total) uncertainty into data and knowledge uncertainty (Depeweg et al., 2018; Gal & Ghahramani, 2016).
Knowledge or epistemic uncertainty refers to the lack of knowledge or ignorance about the most optimal choice of model (parameters) (H¨ullermeier
& Waegeman, 2021).
As additional data is collected, the uncertainty in model parameters should decrease.
This form of uncertainty becomes important whenever the model is tasked with making predictions for out-of-distribution data-points.
For in-distribution inputs, it is expected that the trained model can return reliable predictions.
On the other hand data or aleatoric uncertainty, represents inherent noise in the data being modelled, for example from overlapping classes.
Even if more data is collected, this type of noise is inherent to the process and cannot be avoided or reduced (Malinin & Gales, 2018; Gal & Ghahramani, 2016; Ovadia et al., 2019)
The ability to decompose and distinguish between these sources of uncertainty is important as it allows the cause of uncertainty in the prediction to be known and how it should be used in downstream tasks (Houlsby et al., 2011; Kirsch et al., 2019).
Summary of contributions : In this work we make two important contributions to NN classiﬁer training and uncertainty prediction.
First we introduce self-distribution distillation (S2D), a new general training approach that in an integrated, simultaneous fashion, trains a teacher ensemble and distribution distils the knowledge to a student.
This integrated training allows the user to bypass training a separate expensive teacher ensemble while distribution distillation (Malinin et al., 2020) allows the student to capture the diversity and model a distribution over ensemble member predictions.
Additionally, distribution distillation would give the student the ability to estimate both data and knowledge uncertainty in a single forward pass unlike standard NNs which inherently can not decompose predictive uncertainty, and unlike ensemble methods which can not perform the decomposition in a single pass.
Second, we train an ensemble of these newly introduced models and investigate different distribution distillation techniques giving rise to hierarchical distributions over predictions for uncertainty .
This approach is useful when there are no, or few, computational constraints in the training phase but still require robust uncertainties and efﬁciency in the deployment stage.
B ACKGROUND AND R ELATED W ORK
This section describes two techniques for uncertainty estimation.
First, ensemble methods for predictive uncertainty estimation will be viewed from a Bayesian viewpoint.
Second, a speciﬁc form of distillation for efﬁcient uncertainty estimation will be discussed.
2. E NSEMBLE M ETHODS
From a Bayesian perspective the parameters, θ , of a neural net are treated as random variables with some prior distribution p ( θ ) .
Together with the training data D , this allows the posterior distribution p ( θ |D ) to be derived.
To obtain the predictive distribution over all classes y ∈ Y (for some input x ∗ ) marginalisation over θ is required:
[Equation]
Since ﬁnding the true posterior is intractable a variational approximation p ( θ |D )
≈ q ( θ ) is made (Jordan et al., 1999; Blundell et al., 2015; Graves, 2011; Maddox et al., 2019).
Furthermore, marginalising over all weight values remains intractable leading to a sampling ensemble, approximation method (Gal & Ghahramani, 2016; Lakshminarayanan et al., 2017):
[Equation]
Here, an ensemble generation method is required to obtain the predictive distribution and uncertainty.
Two previously mentioned approaches to generate an ensemble are deep (naive) randomseed and MC-dropout ensemble 1 .
Deep ensembles are based on training M models on the same data but with different initialisations leading to functionally different solutions.
On the other hand, a
MC-dropout ensemble explicitly deﬁnes a variational approximation through the hyper-parameters of dropout (Srivastava et al., 2014) (used during training), allowing for straightforward sampling of model parameters.
Another approach, SWA-Gaussian (Maddox et al., 2019), ﬁnds a Gaussian approximation based on the ﬁrst two moments of stochastic gradient descent iterates.
Unlike the deep ensemble approach, and similar to MC-dropout, this method allows for simple and efﬁcient sampling but suffers from higher memory consumption.
Even a diagonal Gaussian approximation requires twice the memory of a standard network.
There also exists alternative memory and/or compute efﬁcient ensemble approaches such as BatchEnsembles (Wen et al., 2020) and MIMO (Havasi et al., 2021).
While the former approach is parameter efﬁcient it requires multiple forward passes at test time similar to MC ensembles.
The latter avoids this issue by the use of independent subnetworks within a single deep model leading to both efﬁcient training and testing in terms of computational and memory costs.
However, MIMO still requires multiple output heads at test time; this is an issue when scaling to large scale classiﬁcation tasks where the output layer has tens or hundreds of thousands classes.
Given an ensemble, the goal is to estimate and decompose the predictive uncertainty.
First, the entropy of the predictive distribution P ( y | x ∗ , D ) can be seen as a measure of total uncertainty.
Second, this can be decomposed (Depeweg et al., 2018; Kendall & Gal, 2017) as:
[Equation]
where I is mutual information and H represents entropy.
This speciﬁc decomposition allows total uncertainty to be decomposed into separate estimates of knowledge and data uncertainty.
Furthermore, the conditional mutual information can be rephrased as:
[Equation]
For an in-domain sample x ∗ the mutual information should be low as appropriately trained models P ( y | x ∗ , θ ) should be close to the predictive distribution.
High predictive uncertainty will only occur if the input exists in a region of high data uncertainty, for example when an input has signiﬁcant class overlap.
When the input x ∗ is out-of-distribution of the training data, one should expect inconsistent, different, predictions P ( y | x ∗ , θ ) leading to a much higher knowledge uncertainty estimate.
2. E NSEMBLE (D ISTRIBUTION ) D ISTILLATION
Ensemble methods have generally shown superior performance on a range of tasks but suffer from being computationally expensive.
To tackle this issue, a technique called knowledge distillation (KD) and its variants were developed for transferring the knowledge of an ensemble (teacher) into a single (student) model while maintaining good performance (Hinton et al., 2014; Kim & Rush, 2016; Guo et al., 2020).
This is generally achieved by minimising the KL-divergence between the student prediction and the predictive distribution of the teacher ensemble.
In essence, KD trains a new student model to predict the average prediction of its teacher model.
However from the perspective of uncertainty estimation the student model no longer has any information about the diversity of various ensemble member predictions; it was only trained to model the average prediction.
Hence, it is no longer possible to decompose the total uncertainty into different sources, only the total uncertainty can be obtained from the student.
To tackle this issue ensemble distribution distillation (En2D) was developed (Malinin et al., 2020).
Let π signify a categorical distribution, that is π c
= P ( y = ω c
| π ) .
The goal is to directly model the space of categorical predictions { π ( m ) =
f ( x ∗ ; θ ( m ) )
} Mm =1 made by the ensemble.
In work developed by Malinin et al.
(2020) this is done by letting a student model (with weights φ ) predict the parameters of a Dirichlet, which is a continuous distribution over categorical distributions:
[Equation]
The key idea in this concept is that we are not directly interested in the posterior p ( θ |D ) but how predictions π for particular inputs behave when induced by this posterior.
Therefore, it is possible to replace p ( θ |D ) with a trained distribution p ( π | x ∗ , φ ) .
It is now necessary to train the student given the information from the teacher which is straightforwardly done using negative log-likelihood:
[Equation]
A decomposable estimate of total uncertainty is then possible by using conditional mutual information between the class y and prediction π instead of θ Malinin & Gales (2018):
[Equation]
This decomposition has a similar interpretation to eq.
Using a Dirichlet model, these uncertainties can be found using a single forward pass, achieving a much higher level of efﬁciency compared to an ensemble.
Assuming this distillation technique is successful, the distribution distilled student should be able to closely emulate the ensemble and be able to estimate similar high quality uncertainties on both ID and OOD data.
However, ensemble distribution distillation is only applicable and useful when the ensemble members are not overconﬁdent and display diversity in their predictions—there is no need in capturing diversity when there is none.
It is often the case that, for example, convolutional neural networks are over-parameterised, display severe overconﬁdence and can essentially achieve perfect training accuracy restricting the effectiveness of this method (Guo et al., 2017; Seo et al., 2019; Ryabinin et al., 2021).
Furthermore, this method can only be used when an ensemble is available, leading to a high training cost.
S ELF -D ISTRIBUTION D ISTILLATION
In this section we propose self-distribution distillation (S2D) for efﬁcient training and uncertainty estimation, bypassing the need for a separate teacher ensemble.
This combines:
• parameter sharing : allowing the teacher and student to share a common feature extraction base would accelerate training signiﬁcantly, each will branch off and have their own head; • stochastic regularisation : the teacher can generate multiple predictions efﬁciently by forward propagating an input through its head (with a stochastic regulariser) several times, emulating the behaviour of an ensemble; • distribution distillation : while the teacher branch is trained on cross-entropy, the student is taught to predict a distribution over teacher predictions capturing the diversity compactly.
This process is summarised in Fig.
This approach can take many speciﬁc forms with regards to the type of feature extraction module, stochastic regulariser, teacher branch and student modelling choice.
For example, the teacher could entail a much larger branch capturing complex patterns in the data, while the student could consist of a smaller branch used for compressing teacher knowledge into a more efﬁcient form, at test time.
On the other end, training efﬁciency can be achieved by forcing the teacher and student share the same branch parameters.
In this work we choose a particular model conﬁguration that is highly efﬁcient, shown in Fig.
The main functional difference between the teacher and the student branches is the use of logit values, z : for the teacher branch a probability is predicted; whereas the student uses the logits for a Dirichlet distribution.
Furthermore the teacher uses stochastic regularisation techniques (SRTs) in generating multiple teacher predictions, analogous to an ensemble.
In this work multiplicative Gaussian noise (Gaussian dropout) with unit mean and uniformly random standard deviation is used.
This form was chosen due to simplicity of sampling and possible ensemble diversity.
There is a wide range of other choices regarding what SRTs to use, from Bernoulli dropout, additive Gaussian noise to deciding at which teacher branch layers this should be introduced.
Furthermore, since the Dirichlet distribution has bounded ability to represent diverse ensemble predictions (Malinin et al., 2020), simply generating multiple teacher prediction by propagating through the last layer will not limit the ability of the model.
To further improve the memory efﬁciency of the model, a single ﬁnal linear layer shared by both student and teacher branches is used.
This parameter sharing makes the S2D model efﬁcient even when the number of classes is large, and does not use any more parameters compared to a standard model.
Note any NN classiﬁer can be cast into a self-distribution distillation format by inserting stochasticity prior to the ﬁnal linear layer and can easily be combined with many other approaches such as MIMO (Havasi et al., 2021) and SWAG (Maddox et al., 2019).
This choice of integrating ensemble teacher training and distribution distillation into a single entity utilising parameter tying also serves as a regulariser (optimising two objectives using the same set of weights) and allows for inexpensive training.
The only added training cost is from multiple forward passes through the ﬁnal linear layer, a process which can be parallelised.
Additionally, the restricted form of Fig.
2 brings some numerical stability.
As noted by Malinin et al.
(2020), optimising a student to predict a Dirichlet distribution can be unstable when there is a lack of common support between prediction and extremely sharp teacher outputs.
However, note that teacher predictions are closely related to the expected student prediction:
[Equation]
leading to increased common support.
Additionally, stochasticity in the teacher forces the outputs to have some diversity, mildly limiting overconﬁdence.
Now we train the teacher branch using crossentropy, and simultaneously, use the teacher predictions to train the student branch.
Let the weights of this model be denoted by φ and say we have some input-target pair ( x , y ) .
The teacher loss (for a single sample) is then:
[Equation]
where δ is the indicator function.
The student branch could be trained using log-likelihood as in eq.
but it has been found that this approach could be unstable (Fathullah et al., 2021; Ryabinin et al., 2021).
Instead we use the teacher categorical predictions in estimating a proxy teacher Dirichlet ˜
α using maximum log-likelihood.
The resulting student loss is KL-divergence based:
[Equation]
The proxy Dirichlet is estimated using a numerical approach developed by Minka (2000).
The overall training loss becomes L ( φ ) =
L th ( φ ) + µ L st ( φ ) with a small constant µ .
Deep learning models often overﬁt on training data leading to less informative outputs.
To alleviate these issues we integrate temperature scaling in the student branch loss.
While training the teacher branch predictions on cross-entropy we temperature scale the same predictions and use the resulting ones in estimating a proxy teacher Dirichlet.
The student branch will repeatedly be taught to predict a smoother/wider Dirichlet distribution, while the teacher branch’s objective is to maximise the probability of the correct class resulting in a middle ground.
S ELF -D ISTRIBUTION D ISTILLED E NSEMBLE A PPROACHES
If computational resources during the training phase are not constrained it would open up the possibility for self-distribution distilled ensembles and various hierarchical distillation approaches of such models.
First it can be noted that the ensemble generation methods mentioned in previous sections can easily be used with the S2D models in the previous section.
The predictive distribution of such an ensemble would take the following form:
[Equation]
Furthermore, an ensemble of Dirichlet models can be used to estimate similar uncertainty measures as previously described:
[Equation]
This is a generalisation of eq.
(7) since speciﬁc weights φ have been replaced with conditioning on the dataset D .
Computing these uncertainties requires only a few modiﬁcations compared to the standard ensemble in eq.
Next, the most natural step is to transfer the knowledge of an S2D (Dirichlet) ensemble into a single model.
A choice needs to be made regarding the hierarchy of student modelling: should the student predict a categorical 2 , Dirichlet, or a distribution over Dirichlets—hereby given the family name hierarchical distribution distillation (H2D).
Initially we start by training a student model to predict a single Dirichlet identical to eq.
However, since the S2D ensemble provides, for an input x ∗ , a set of Dirichlets { α ( m ) =
f ( x ∗ ; φ [Equation]
[Equation] a modiﬁed distillation criterion is needed:
[Equation]
This KL-divergence based loss also allows the reverse KL criterion to be used (Malinin & Gales, 2019) if desired.
One criticism of this form of model, Dirichlet H2D (H2D-Dir), is that the diversity across ensemble members is lost, similar to the drawback in standard distillation.
Therefore, we seek a distribution over Dirichlets to capture this higher level of diversity.
To model the space of Dirichlets we need to deﬁne a distribution over the parameters.
Here we are faced with a choice: (1) model the parameters α ∈ R K + directly (restricted to the non-negative real space) or (2) apply a transformation to simplify the modelling.
Here a logarithmic transformation z = ln α ∈ R K is applied and a simple distribution over the Dirichlet parameters, a diagonal
Gaussian, to be used (see Appendix C for a justiﬁcation for this modelling choice).
With these building blocks, the goal of H2D is to train a student model, with weights λ , predict the parameters of a diagonal Gaussian ( µ , σ ) (H2D-Gauss):
[Equation]
By sampling from this Gaussian, one can obtain multiple Dirichlet distributions similar to, but cheaper than, an S2D ensemble.
Clearly, the ﬂexibility of such a model can easily be extended by allowing the model to predict a fully speciﬁed covariance, however due to computational tractability only diagonal covariance models are used in this work.
A secondary head is required for such a model, see Fig.
In a similar fashion to previous approaches, this model can be trained using negative log-likelihood or by estimating a proxy teacher Gaussian and use KL-divergence.
In this work we have adopted the proxy approach, see Appendix A.1 for details
E XPERIMENTAL E VALUATION
This section investigates the self-distribution distillation approach on classifying image data.
First, this approach is compared to standard trained models and established ensemble based methods (deep ensembles and MC-dropout) as well as the diagonal version of SWAG (SWAG-Diag) and MIMO.
Second, self-distribution distillation is combined with all above mentioned approaches.
Finally, knowledge distillation is compared to hierarchical distribution distillation of Dirichlet ensembles.
This comparison is based on two sets of experiments.
The ﬁrst set compares the performance of all baselines and proposed models in terms of image classiﬁcation performance and calibration on
CIFAR-100 (Krizhevsky & Hinton, 2009) without (C100) and with (C100+) data augmentation.
The second set of experiments compares the out-of-distribution (OOD) detection performance using the
LSUN (Yu et al., 2015), Tiny ImageNet (CS231N, 2017) and SVHN (Netzer et al., 2011) datasets.
All experiments are based on training the DenseNet-BC (k = 12) with a depth of 100 (Huang et al., 2017).
For ensemble generation methods M = 5 models were sampled (in the case of MC-dropout ensembles and SWAG) or trained (in the case of deep ensembles).
For MIMO we use two output heads ( M = 2 ) due to limited capacity in the chosen model (Havasi et al., 2021).
Note that for this choice of model it was not possible to use ensemble distribution distillation since DenseNet-BC models display high conﬁdence on the training data of CIFAR-100 causing instability in distillation.
All single model training runs were repeated 5 times; mean ± 2 standard deviations are reported.
The experimental setup and additional experiments are described in Appendix A-D.
5. CIFAR-100 C LASSIFICATION P ERFORMANCE E XPERIMENTS
The ﬁrst batch of experiments show the classiﬁcation performance using a range of metrics such as accuracy, negative log-likelihood (NLL) and expected calibration error (ECE), see Table 1.
Perhaps the most noteworthy result is the improvement in all metrics and datasets of a self-distribution distilled model compared to its standard counterpart.
The improvement is more than 2 standard deviations.
A similar picture can be observed for the S2D versions of SWAG-Diag and MC-dropout which, without any notable cost of training and inference, improve upon their equivalent standard counterparts in all metrics notably. Regarding MIMO a small gain can still be observed when switching to the self-distribution distillation framework but this boost is smaller.
Finally for the deep ensemble approach, the S2D version only shows a marginal improvement in accuracy and NLL but a notable increase in ECE.
In fact, it is observed that ensembling standard and S2D models reduces and increases ECE respectively.
This trend is associated with the level of ensemble calibration.
Unlike a standard deep ensemble, the members of the S2D counterpart are close to being calibrated, displaying little to no overconﬁdence.
Ensembling these calibrated models lead to under-conﬁdent average predictions hence, the increased calibration error.
Note, calibration error and negative loglikelihood can easily be reduced post-training by temperature scaling predictions.
The next set of comparisons regard distilled models, the ﬁnal block of Table 1.
As expected they all perform in between the performance of an individual model and deep ensemble.
While standard ensemble distillation (knowledge distillation) was found to achieve better accuracy than other distillation methods, this success was highly dependent on the value of temperature scaling used.
A sub-optimal choice of temperature can drastically reduce performance.
On the other hand, when distilling an S2D ensemble, no additional hyper-parameters are needed.
We observe that while both
[Table]
H2D-Dir and H2D-Gauss obtained a higher NLL they also achieved better calibration than their
S2D ensemble teacher.
Lastly, one can observe that H2D-Dir and H2D-Gauss both outperform the standard SWAG-Diag and MC-dropout ensemble using only a single forward pass.
Although these distilled models involve an expensive training phase (a teacher ensemble is required) they are able to, at test time, achieve much higher computational efﬁciency and estimate and decompose uncertainty.
5. O UT - OF - DISTRIBUTION D ETECTION E XPERIMENTS
The second batch of experiments investigate the out-of-distribution detection performance of models.
The goal is to differentiate between two types of data, negative in-distribution (ID, sampled from the same source as the training data) and positive out-of-distribution (OOD) data.
In all experiments the models were trained on C100.
The ID data was always set to the test set of
C100 and OOD data was the test set of LSUN/TIM/SVHN.
Both LSUN and TIM examples had to be resized or randomly cropped as preprocessing before being fed to the model.
The detection was done using four uncertainty estimates: conﬁdence, total uncertainty (TU), data or aleatoric uncertainty (DU) and knowledge or epistemic uncertainty (KU).
Performance was measured using the threshold independent AUROC (Manning & Sch ¨ utze, 1999) and AUPR (Fawcett, 2006) metrics.
Due to limited space, some LSUN and TIM experiments have been moved to Appendix B.1.
First, there is not a single case in Tables 2 and 3 where an individual model, MIMO, SWAG-Diag or MC-dropout ensemble is able to outperform the detection performance of a single S2D model.
This statement holds for all the analysed uncertainties apart from conﬁdence where both MIMO and SWAG-Diag are insigniﬁcantly better.
When comparing to a deep ensemble, the S2D model is outperformed in many cases.
The general trend is that the ensemble is able to output marginally higher quality conﬁdence and total uncertainty estimates in most datasets, but that S2D sometimes outperforms the ensemble when using data uncertainty (as in Table 3).
Interestingly, the MC ensemble seems to degrade the quality of conﬁdence and total uncertainty when compared to its standard individual counterpart.
However, since a MC-dropout ensemble can estimate data uncertainty, it is able to outperform the standard model overall.
Similarly, the S2D MC ensemble generally has inferior detection performance compared to its single deterministic model equivalent.
The only exception is in detecting SVHN where the ensemble has marginally better data uncertainty estimates.
Regarding SWAG-Diag and MIMO they both gain from being cast into a self-distribution distillation viewpoint drastically increasing their detection performance without additional cost at inference.
Although the S2D deep ensemble, when compared to its vanilla counterpart, wasn’t able to show any noticeable accuracy boost (on CIFAR-100) it does outperform in this detection task.
The only case where the S2D ensemble was not able to outshine the vanilla ensemble is when both use knowl- [Table] [Table] edge uncertainty to detect SVHN examples using the AUPR metric.
Generally, S2D based systems outperform their standard counterparts.
Regarding distillation based approaches, it is observed that knowledge ensemble distillation, EnD, is able to outperform the standard model in all cases except SVHN detection, and in no case is able to reach the deep ensemble performance, which it was distilled from.
On the other hand, both the
H2D-Dir and H2D-Gauss models outperform the distilled model and are able to decompose predictive uncertainty.
Speciﬁcally we discover that H2D-Dir is able to generate the highest quality knowledge uncertainty estimates in almost all cases, and is able to outperform its S2D ensemble teacher using this uncertainty.
The H2D-Gauss model however, was not able to boast similar high quality knowledge uncertainty.
Instead, this model displayed the generally best performing data uncertainty estimates, able to outperform the vanilla deep ensemble in all cases, and the S2D equivalent in all but SVHN detection.
C ONCLUSION
Uncertainty estimation within deep learning is becoming increasingly importance, with deep ensembles being the standard for estimating various sources of uncertainty.
However, ensembles suffer from signiﬁcantly higher computational requirements.
This work proposes self-distribution distillation (S2D), a novel collection of approaches for directly training models able to estimate and decompose predictive uncertainty, without explicitly training an ensemble, and can seamlessly be combined with other approaches.
Additionally, if one is not resource restricted during the training phase, a novel approach, hierarchical distribution distillation (H2D), is described for transferring/distilling the knowledge of S2D style ensembles into a single ﬂexible and robust student model.
It is shown that S2D models are able to outperform standard models and rival MC ensembles on the CIFARtest set.
Additionally, S2D is able to estimate higher quality uncertainty estimates compared to standard models and MC ensembles and in most cases, able to better detect out-of-distribution images from the LSUN, SVHN and TIM datasets.
Combination of S2D with other promising approaches such as MIMO and SWAG also show additional gains in accuracy and detection performance.
S2D is also able to rival the deep ensemble in certain cases even though it only requires a single forward pass.
Furthermore, S2D deep ensembles and H2D derived student models are shown to notably outperform the deep ensemble in almost all detection problems.
These promising results show that the efﬁcient self-distribution and novel hierarchical distribution distillation approaches have the potential to train robust uncertainty estimating models able to outperform deep ensembles.
Future work should further investigate self-distribution distillation in other domains such as natural language processing and speech recognition.
The need for more efﬁcient uncertainty estimation is especially useful for these areas as they often utilise large-scale models.
Furthermore, one could also analyse variations of S2D such as utilising less weight sharing, generating more diverse teacher predictions or changing the student modelling choices.
A E XPERIMENTAL C ONFIGURATION
[Table]
All models were trained on the CIFAR-100 dataset, with and without data augmentation.
The augmentation scheme involves randomly mirroring and shifting images following He et al.
Huang et al.
Remaining datasets were used as out-of-distribution samples in the detection task.
All individual models, and ensemble members were based of off the DenseNet-BC ( k = 12 ,layers) architecture and trained according to Huang et al.
{+SWAG-Diag was obtained by checkpointing the weights of the last 20 epochs with a reduced learning rate of η = 1 . 
+}
× 10 − 4 .
MIMO with two output heads was trained using the same setup as for the standard model.
To keep training costs comparable to (S2D) individual models, no batch or input repetition was used (Havasi et al., 2021).
Similarly all self-distribution distilled equivalents were trained with identical training recipes with the addition of a student loss ( µ = 1 .
28 × 10 − 4 ).
Regarding distilled based models, the EnD baseline was trained using negative log-likelihood using the average temperature scaled prediction of the teacher ensemble, with T ∈ { 1 . 0 , 2 . 0 , 3 . 0 , 4 . 0 , 5 . 0 } .
For the hierarchical distribution distillation approaches the students were ﬁrst initialised with the weights of an S2D model trained for 150 epochs, for increased stability.
Thereafter, each student was trained using the appropriate H2D criteria with a signiﬁcantly reduced learning rate.
H2D-Dir was trained using η = 5 .
× 10 − 5 for an additional 150 epochs.
H2D-Gauss required an initial learning rate of η = 5 .
× 10 − 3 which was reduced by a factor of 2 after 75 and 150 epochs.
It was trained for 170 epochs.
Additionally, uncertainties were computed by generating 50 samples from each Gaussian prediction, since this modelling choice does not result in closed form expressions.
A. P ROXY T ARGET T RAINING
Since the use of negative log-likelihood can be unstable in training S2D and distilling H2D models we utilise proxy targets and KL-divergence.
It has already been mentioned that the proxy target in
S2D follows:
[Equation]
Each categorical prediction will be temperature scaled, with T = 1 . 5 , to mitigate overconﬁdent predictions.
While H2D-Dir does not require any proxy targets, the Gaussian equivalent does.
The proxy diagonal Gaussian, estimated according to maximum log-likelihood, has a closed-form expression:
[Equation]
where v 2 = v (cid:12) v represents an element-wise multiplication.
This is then used in a KL-divergence based loss, training the student with prediction µ , σ according to:
[Equation]
Note however, that the proxy targets are detached from any back gradient propagation calculations.
This is to simulate typical teacher-student knowledge transfer where teacher weights are kept ﬁxed during student training.
B O UT - OF - DISTRIBUTION D ETECTION
This section covers remaining out-of-distribution detection experiments.
First, we cover the LSUN and Tiny ImageNet detection problem for all models considered in section 5.2.
Thereafter, additional experiments will be run on ensembles of various sizes.
This is to investigate if the low quality of knowledge uncertainty estimates is caused by a limited number of ensemble members.
B. T INY I MAGE N ET E XPERIMENTS
Similar to the results section 5.2 the S2D
Deep ensemble and H2D-Gauss outperformed all other models, see Table 6 and 7.
The only exception is the use of conﬁdence on resized TIM with the AUROC metric where the Deep ensemble marginally outperforms the S2D equivalent.
However, unlike previous results, knowledge uncertainty seems to perform on par with or outperform conﬁdence.
The one exception is the MC ensemble.
[Table]
[Table]
EnD
[Table]
[Table]
B. E NSEMBLE S IZE E XPERIMENTS
Knowledge uncertainty was found to have underwhelming performance (especially for MC and
Deep ensembles) and did not show similar trends to prior work (Malinin & Gales, 2018; 2021;
Malinin et al., 2020).
To possibly mitigate this, the ensemble size was increased as a smaller number of models could lead to inaccurate measures of diversity and knowledge uncertainty.
Results are compiled in Tables 8-13.
Performance on the CIFAR-100 test set is shown in Table 8.
Increasing the ensemble size leads to improved accuracy and lower negative log-likelihoods as would be expected.
The MC ensemble also becomes better calibrated.
The Deep ensemble on the other hand has increasing calibration error with the number of members.
This is due to the ensemble prediction becoming under-conﬁdent when averaging over a large number of members.
Out-of-distribution detection performance on LSUN, SVHN and TIM are compiled in Tables 9-13.
Although the MC ensemble enjoys improved accuracy when increased in size, it seems to remain relatively unaffected in terms of OOD detection using any uncertainty metric.
In detecting LSUN using random crops, the performance of KU interestingly deteriorates notably.
Overall this points to MC ensembles’ lacking ability in utilising new information from additional ensemble member draws/samples for better uncertainty estimation.
Regarding the Deep ensemble, it generally improves with increasing size with any metric, however with diminishing returns.
In this case all uncertainties improve with ensemble size, not only knowledge uncertainty.
Therefore it seems that the cause for conﬁdence, total and data outperforming knowledge uncertainty is not due to the ensemble size being limited to ﬁve members.
[Table] [Table]
[Table] [Table] [Table]
[Table] [Table]
C B EHAVIOUR OF U NCERTAINTIES
This section investigates how the uncertainties produced from a vanilla Deep ensemble differ from self-distribution distilled derived systems, and how well hierarchical distribution distillation captures the behaviour of its teacher.
The comparison will be made between the in-domain CIFAR-100 and, out of simplicity, only the out-of-domain SVHN test set.
Figure 4 shows the contrast of various uncertainties between an CIFAR-100 (ID) and SVHN (OOD) test sets.
Clearly, the S2D systems output ID uncertainties in a consistent manner, even matching the [Figure] conceptually different Deep ensemble.
Observe that S2D integrates temperature scaling (smoothing predictions) into the training of models; total and data uncertainties 3 estimated by these models will naturally have larger entropy than Deep ensembles.
While it is expected that the Deep ensemble would have different behaviour on the SVHN OOD set, it is surprising to observe how well H2DGauss aligns with its S2D Deep ensemble teacher.
An individual S2D model was also able to generate closely related total and data uncertainty estimates, but suffers signiﬁcantly in producing consistent knowledge uncertainties.
These results raise the question if a Gaussian student could capture the diversity in a vanilla Deep ensemble by modelling the logits, in a similar fashion to how
H2D-Gauss models its teacher—a possible avenue for future work.
D A DDITIONAL E XPERIMENTS : W IDE R ES N ET
Following the DenseNet-BC experiments in section 5 we repeated them with a different architecture.
In this section we focus on a signiﬁcantly larger WideResNet (Zagoruyko & Komodakis, 2016) model with a depth of 28 and a widening factor of 10.
The standard and S2D models were both trained as described in Zagoruyko & Komodakis (2016), with the S2D speciﬁc parameters being the same as previously described.
The only difference is that teacher predictions were generated using multiplicative Gaussian noise with a ﬁxed standard deviation of 0 . 10 .
The H2D-Gauss model was also trained in a different manner.
First, it was initialised from an S2D model trained for 150 epochs.
Thereafter it was trained for an additional 80 epochs with a starting learning rate of η = 2 × 10 − 3 which was reduced by a factor of 4 after 60 epochs.
For this section,
EnD and H2D-Dir were not investigated.
Table 14 shows test set performance.
Unlike previous experiments, S2D was not able to outperform an individual model by more than two standard deviations, in this case achieving around one standard deviation improvement in accuracy.
Interestingly, the MC approach has worse accuracy for both the standard and S2D case, however this could be due to the small number of drawn samples ( M = 5 ).
Furthermore, both Deep ensembles signiﬁcantly outperform their individual equivalents with the
S2D version being slightly better in all measured performance metrics.
The notable result in this table is the high performance of H2D-Gauss, able to outperform the Deep ensemble in C100 and achieve near ensemble performance in C100+.
In the OOD detection task we observe that both versions of the MC ensemble struggle to outperform their individual counterparts.
There also seems to be a disparity in performance when comparing resize and random cropped LSUN and TIM.
With random crops, all S2D systems notably outperform [Table] their standard counterparts.
In this case both S2D Individual and H2D-Gauss were able to outperform the Deep ensemble using any uncertainty metric.
In the other case of resizing LSUN and TIM images and in SVHN the detection performance difference is smaller but the S2D
Deep ensemble still remains the best model with both H2D-Gauss and Deep ensemble performing similarly
[Table]
[Table]
[Table]
[Table]
[Table]
[Table] [Table] [Table] [Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]