Self-Distribution Distillation: Efficient Uncertainty Estimation

1 INTRODUCTION . 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 identification ( 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 significantly 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 significantly 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 ) ( Hüllermeier & 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 can not 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 classifier 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 efficiency in the deployment stage . 2 BACKGROUND AND RELATED WORK . This section describes two techniques for uncertainty estimation . First , ensemble methods for predictive uncertainty estimation will be viewed from a Bayesian viewpoint . Second , a specific form of distillation for efficient uncertainty estimation will be discussed . 2.1 ENSEMBLE METHODS . From a Bayesian perspective the parameters , θ , of a neural net are treated as random variables with some prior distribution p ( θ ) . Together with the training dataD , 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 : P ( y|x∗ , D ) = Ep ( θ|D ) [ P ( y|x∗ , θ ) ] ( 1 ) Since finding 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 ) : P ( y|x∗ , D ) ≈ 1 M M∑ m=1 P ( y|x∗ , θ ( m ) ) , θ ( m ) ∼ q ( θ ) ( 2 ) 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 ensemble1 . 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 defines 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 ) , finds a Gaussian approximation based on the first two moments of stochastic gradient descent iterates . Unlike the deep ensemble approach , and similar to MC-dropout , this method allows for simple and efficient sampling but suffers from higher memory consumption . Even a diagonal Gaussian approximation 1In-depth comparisons of ensemble methods were conducted in Ovadia et al . ( 2019 ) ; Ashukha et al . ( 2020 ) requires twice the memory of a standard network . There also exists alternative memory and/or compute efficient ensemble approaches such as BatchEnsembles ( Wen et al. , 2020 ) and MIMO ( Havasi et al. , 2021 ) . While the former approach is parameter efficient 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 efficient 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 classification 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 : H [ P ( y|x∗ , D ) ] ︸ ︷︷ ︸ Total Uncertainty = I [ y , θ|x∗ , D ] ︸ ︷︷ ︸ Knowledge Uncertainty + Ep ( θ|D ) [ H [ P ( y|x∗ , θ ) ] ] ︸ ︷︷ ︸ Data Uncertainty ( 3 ) where I is mutual information and H represents entropy . This specific decomposition allows total uncertainty to be decomposed into separate estimates of knowledge and data uncertainty . Furthermore , the conditional mutual information can be rephrased as : I [ y , θ|x∗ , D ] = Ep ( θ|D ) [ KL ( P ( y|x∗ , θ ) ∥∥ P ( y|x∗ , D ) ) ] ( 4 ) 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 significant 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.2 ENSEMBLE ( DISTRIBUTION ) DISTILLATION . 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 : p ( π|x∗ , φ ) = Dir ( π ; α ) , α = f ( x∗ ; φ ) ( 5 ) 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 : L ( φ ) = − 1 M M∑ m=1 ln Dir ( π ( m ) ; α ) ( 6 ) 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 ) : H [ P ( y|x∗ , φ ) ] ︸ ︷︷ ︸ Total Uncertainty = I [ y , π|x∗ , φ ] ︸ ︷︷ ︸ Knowledge Uncertainty + Ep ( π|x∗ , φ ) [ H [ P ( y|π ) ] ] ︸ ︷︷ ︸ Data Uncertainty ( 7 ) This decomposition has a similar interpretation to eq . ( 3 ) . Using a Dirichlet model , these uncertainties can be found using a single forward pass , achieving a much higher level of efficiency 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 overconfident 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 overconfidence 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 .

Summary: This paper contributes to neural network classifier training and uncertainty prediction.  It proposes a self-distribution distillation method that can train a single model to estimate uncertainties in an integrated training phase.  Also, it is flexible to be extended to build ensembles of models in the training phase and efficiently deployed in the test phase. Experiments are done on both classification and out-of-distribution detection tasks to show effectiveness.