Ensemble Distribution Distillation

1 INTRODUCTION . Neural Networks ( NNs ) have emerged as the state-of-the-art approach to a variety of machine learning tasks ( LeCun et al. , 2015 ) in domains such as computer vision ( Girshick , 2015 ; Simonyan & Zisserman , 2015 ; Villegas et al. , 2017 ) , natural language processing ( Mikolov et al. , 2013b ; a ; 2010 ) , speech recognition ( Hinton et al. , 2012 ; Hannun et al. , 2014 ) and bio-informatics ( Caruana et al. , 2015 ; Alipanahi et al. , 2015 ) . Despite impressive supervised learning performance , NNs tend to make over-confident predictions ( Lakshminarayanan et al. , 2017 ) and , until recently , have been unable to provide measures of uncertainty in their predictions . As NNs are increasingly being applied to safety-critical tasks such as medical diagnosis ( De Fauw et al. , 2018 ) , biometric identification ( Schroff et al. , 2015 ) and self driving cars , estimating uncertainty in model ’ s predictions is crucial , as it enables the safety of an AI system ( Amodei et al. , 2016 ) to be improved by acting on the predictions in an informed manner . Ensembles of NNs are known to yield increased accuracy over a single model ( Murphy , 2012 ) , allow useful measures of uncertainty to be derived ( Lakshminarayanan et al. , 2017 ) , and also provide defense against adversarial attacks ( Smith & Gal , 2018 ) . There is both a range of Bayesian MonteCarlo approaches ( Gal & Ghahramani , 2016 ; Welling & Teh , 2011 ; Garipov et al. , 2018 ; Maddox et al. , 2019 ) , as well as non-Bayesian approaches , such as random-initialization ( Lakshminarayanan et al. , 2017 ) and bagging ( Murphy , 2012 ; Osband et al. , 2016 ) , to generating ensembles . Crucially , ensemble approaches allow total uncertainty in predictions to be decomposed into knowledge uncertainty and data uncertainty . Data uncertainty is the irreducible uncertainty in predictions which arises due to the complexity , multi-modality and noise in the data . Knowledge uncertainty , also known as epistemic ∗Equal Contribution uncertainty ( Gal , 2016 ) or distributional uncertainty ( Malinin & Gales , 2018 ) , is uncertainty due to a lack of understanding or knowledge on the part of the model regarding the current input for which the model is making a prediction . This form of uncertainty arises when the test input x∗ comes either from a different distribution than the one that generated the training data or from an in-domain region which is sparsely covered by the training data . Mismatch between the test and training distributions is also known as a dataset shift ( Quiñonero-Candela , 2009 ) , and is a situation which often arises for real world problems . Distinguishing between sources of uncertainty is important , as in certain machine learning applications it may be necessary to know not only whether the model is uncertain , but also why . For instance , in active learning , additional training data should be collected from regions with high knowledge uncertainty , but not data uncertainty . A fundamental limitation of ensembles is that the computational cost of training and , more importantly , inference can be many times greater than that of a single model . One solution to speed up inference is to distill an ensemble of models into a single network to yield the mean predictions of the ensemble ( Hinton et al. , 2015 ; Korattikara Balan et al. , 2015 ) . However , this collapses an ensemble of conditional distributions over classes into a single point-estimate conditional distribution over classes . As a result , information about the diversity of the ensemble is lost . This prevents measures of knowledge uncertainty , such as mutual information ( Malinin & Gales , 2018 ; Depeweg et al. , 2017a ) , from being estimated . In this work , we investigate the explicit modelling of the distribution over the ensemble predictions , rather than just the mean , with a single model . This problem — referred to as Ensemble Distribution Distillation ( EnD2 ) — yields a method that preserves both the distributional information and improved classification performance of an ensemble within a single neural network model . It is important to highlight that Ensemble Distribution Distillation is a novel task which , to our knowledge , has not been previously investigated . Here , the goal is to extract as much information as possible from an ensemble of models and retain it within a single , possibly simpler , model . As an initial solution to this problem , this paper makes use of a recently introduced class of models , known as Prior Networks ( Malinin & Gales , 2018 ; 2019 ) , which explicitly model a conditional distribution over categorical distributions by parameterizing a Dirichlet distribution . Within the context of EnD2 this effectively allows a single model to emulate the complete ensemble . The contributions of this work are as follows . Firstly , we define the task of Ensemble Distribution Distillation ( EnD2 ) as a new challenge for machine learning research . Secondly , we propose and evaluate a solution to this problem using Prior Networks . EnD2 is initially investigated on artificial data , which allows the behaviour of the models to be visualized . It is shown that distributiondistilled models are able to distinguish between data uncertainty and knowledge uncertainty . Finally , EnD2 is evaluated on CIFAR-10 , CIFAR-100 and TinyImageNet datasets , where it is shown that EnD2 yields models which approach the classification performance of the original ensemble and outperform standard DNNs and regular Ensemble Distillation ( EnD ) models on the tasks of identifying misclassifications and out-of-distribution ( OOD ) samples . 2 ENSEMBLES . In this work , a Bayesian viewpoint on ensembles is adopted , as it provides a particularly elegant probabilistic framework , which allows knowledge uncertainty to be linked to Bayesian model uncertainty . However , it is also possible to construct ensembles using a range of non-Bayesian approaches . For example , it is possible to explicitly construct an ensemble of M models by training on the same data with different random seeds ( Lakshminarayanan et al. , 2017 ) and/or different model architectures . Alternatively , it is possible to generate ensembles via Bootstrap methods ( Murphy , 2012 ; Osband et al. , 2016 ) in which each model is trained on a re-sampled version of the training data . The essence of Bayesian methods is to treat the model parameters θ as random variables and place a prior distribution p ( θ ) over them to compute a posterior distribution p ( θ|D ) via Bayes ’ rule : p ( θ|D ) = p ( D|θ ) p ( θ ) p ( D ) ∝ p ( D|θ ) p ( θ ) ( 1 ) Here , model uncertainty is captured in the posterior distribution p ( θ|D ) . Consider an ensemble of models { P ( y|x∗ , θ ( m ) ) } Mm=1 sampled from the posterior : { P ( y|x∗ , θ ( m ) ) } M m=1 → { P ( y|π ( m ) ) } M m=1 , π ( m ) = f ( x∗ ; θ ( m ) ) , θ ( m ) ∼ p ( θ|D ) ( 2 ) where π are the parameters of a categorical distribution [ P ( y = ω1 ) , · · · , P ( y = ωK ) ] T and x∗ is a test input . The expected predictive distribution , or predictive posterior , for a test input x∗ is obtained by taking the expectation with respect to the model posterior : P ( y|x∗ , D ) = Ep ( θ|D ) [ P ( y|x∗ , θ ) ] ( 3 ) Each of the models P ( y|x∗ , θ ( m ) ) yields a different estimate of data uncertainty . Uncertainty in predictions due to model uncertainty is expressed as the level of spread , or ‘ disagreement ’ , of an ensemble sampled from the posterior . The aim is to craft a posterior p ( θ|D ) , via an appropriate choice of prior p ( θ ) , which yields an ensemble that exhibits the set of behaviours described in figure 1 . Specifically , for an in-domain test input x∗ , the ensemble should produce a consistent set of predictions with little spread , as described in figure 1a and figure 1b . In other words , the models should agree in their estimates of data uncertainty . On the other hand , for inputs which are different from the training data , the models in the ensemble should ‘ disagree ’ and produce a diverse set of predictions , as shown in figure 1c . Ideally , the models should yield increasingly diverse predictions as input x∗ moves further away from the training data . If an input is completely unlike the training data , then the level of disagreement should be significant . Hence , the measures of model uncertainty will capture knowledge uncertainty given an appropriate choice of prior . Given an ensemble { P ( y|x∗ , θ ( m ) ) } M m=1 which exhibits the desired set of behaviours , the entropy of the expected distribution P ( y|x∗ , D ) can be used as a measure of total uncertainty in the prediction . Uncertainty in predictions due to knowledge uncertainty can be assessed via measures of the spread , or ‘ disagreement ’ , of the ensemble such as Mutual Information : MI [ y , θ|x∗ , D ] ︸ ︷︷ ︸ Knowledge Uncertainty = H [ Ep ( θ|D ) [ P ( y|x∗ , θ ) ] ] ︸ ︷︷ ︸ Total Uncertainty − Ep ( θ|D ) [ H [ P ( y|x∗ , θ ) ] ] ︸ ︷︷ ︸ Expected Data Uncertainty ( 4 ) This formulation of mutual information allows the total uncertainty to be decomposed into knowledge uncertainty and expected data uncertainty ( Depeweg et al. , 2017a ; b ) . The entropy of the predictive posterior , or total uncertainty , will be high whenever the model is uncertain - both in regions of severe class overlap and out-of-domain . However , the difference of the entropy of the predictive posterior and the expected entropy of the individual models will be non-zero only if the models disagree . For example , in regions of class overlap , each member of the ensemble will yield a high entropy distribution ( figure 1b ) - the entropy of the predictive posterior and the expected entropy will be similar and mutual information will be low . In this situation total uncertainty is dominated by data uncertainty . On the other hand , for out-of-domain inputs the ensemble yields diverse distributions over classes such that the predictive posterior is near uniform ( figure 1c ) , while the expected entropy of each model may be much lower . In this region of input space the models ’ understanding of data is low and , therefore , knowledge uncertainty is high . 3 ENSEMBLE DISTRIBUTION DISTILLATION . Previous work ( Hinton et al. , 2015 ; Korattikara Balan et al. , 2015 ; Wong & Gales , 2017 ; Wang et al. , 2018 ; Papamakarios , 2015 ; Buciluǎ et al. , 2006 ) has investigated distilling a single large network into a smaller one and an ensemble of networks into a single neural network . In general , distillation is done by minimizing the KL-divergence between the model and the expected predictive distribution of an ensemble : L ( φ , Dens ) = Ep̂ ( x ) [ KL [ Ep̂ ( θ|D ) [ P ( y|x ; θ ) ] || P ( y|x ; φ ) ] ] ] ( 5 ) This approach essentially aims to train a single model that captures the mean of an ensemble , allowing the model to achieve a higher classification performance at a far lower computational cost . The use of such models for uncertainty estimation was investigated in ( Li & Hoiem , 2019 ; Englesson & Azizpour , 2019 ) . However , the limitation of this approach with regards to uncertainty estimation is that the information about the diversity of the ensemble is lost . As a result , it is no longer possible to decompose total uncertainty into knowledge uncertainty and data uncertainty via mutual information as in equation 4 . In this work we propose the task of Ensemble Distribution Distillation , where the goal is to capture not only the mean of the ensemble , but also its diversity . In this section , we outline an initial solution to this task . An ensemble can be viewed as a set of samples from an implicit distribution of output distributions : { P ( y|x∗ , θ ( m ) ) } M m=1 → { P ( y|π ( m ) ) } M m=1 , π ( m ) ∼ p ( π|x∗ , D ) ( 6 ) Recently , a new class of models was proposed , called Prior Networks ( Malinin & Gales , 2018 ; 2019 ) , which explicitly parameterize a conditional distribution over output distributions p ( π|x∗ ; φ̂ ) using a single neural network parameterized by a point estimate of the model parameters φ̂ . Thus , a Prior Network is able to effectively emulate an ensemble , and therefore yield the same measures of uncertainty . A Prior Network p ( π|x∗ ; φ̂ ) models a distribution over categorical output distributions by parameterizing the Dirichlet distribution . p ( π|x ; φ̂ ) = Dir ( π|α̂ ) , α̂ = f ( x ; φ̂ ) , α̂c > 0 , α̂0 = K∑ c=1 α̂c ( 7 ) The distribution is parameterized by its concentration parametersα , which can be obtained by placing an exponential function at the output of a Prior Network : α̂c = eẑc , where ẑ are the logits predicted by the model . While a Prior Network could , in general , parameterize arbitrary distributions over categorical distributions , the Dirichlet is chosen due to its tractable analytic properties , which allow closed form expressions for all measures of uncertainty to be obtained . However , it is important to note that the Dirichlet distribution may be too limited to fully capture the behaviour of an ensemble and other distributions may need to be considered . In this work we consider how an ensemble , which is a set of samples from an implicit distribution over distributions , can be distribution distilled into an explicit distribution over distributions modelled using a single Prior Network model , ie : { P ( y|x ; θ ( m ) ) } M m=1 → p ( π|x ; φ̂ ) . This is accomplished in several steps . Firstly , a transfer dataset Dens = { x ( i ) , π ( i,1 : M ) } Ni=1 ∼ p̂ ( x , π ) is composed of the inputs xi from the original training set D = { x ( i ) , y ( i ) } Ni=1 and the categorical distributions { π ( i,1 : M ) } Ni=1 derived from the ensemble for each input . Secondly , given this transfer set , the model p ( π|x ; φ ) is trained by minimizing the negative log-likelihood of each categorical distribution π ( im ) : L ( φ , Dens ) = − Ep̂ ( x ) [ Ep̂ ( π|x ) [ ln p ( π|x ; φ ) ] ] = − 1 N N∑ i=1 [ ln Γ ( α̂ ( i ) 0 ) − K∑ c=1 ln Γ ( α̂ ( i ) c ) + 1 M M∑ m=1 K∑ c=1 ( α̂ ( i ) c − 1 ) lnπ ( im ) c ] ( 8 ) Thus , Ensemble Distribution Distillation with Prior Networks is a straightforward application of maximum-likelihood estimation . Given a distribution-distilled Prior Network , the predictive distribution is given by the expected categorical distribution π̂ under the Dirichlet prior : P ( y = ωc|x∗ ; φ̂ ) = Ep ( π|x∗ ; φ̂ ) [ P ( y = ωc|π ) ] = π̂c = α̂c∑K k=1 α̂k = eẑc∑K k=1 e ẑk ( 9 ) Separable measures of uncertainty can be obtained by considering the mutual information between the prediction y and the parameters of π of the categorical : MI [ y , π|x∗ ; φ̂ ] ︸ ︷︷ ︸ Knowledge Uncertainty = H [ Ep ( π|x∗ ; φ̂ ) [ P ( y|π ) ] ] ︸ ︷︷ ︸ Total Uncertainty − Ep ( π|x∗ ; φ̂ ) [ H [ P ( y|π ) ] ] ︸ ︷︷ ︸ Expected Data Uncertainty ( 10 ) Similar to equation 4 , this expression allows total uncertainty , given by the entropy of the expected distribution , to be decomposed into data uncertainty and knowledge uncertainty ( Malinin & Gales , 2018 ) . If Ensemble Distribution Distillation is successful , then the measures of uncertainty derivable from a distribution-distilled model should be identical to those derived from the original ensemble .

Summary: Ensembles of probabilistic models (e.g. for classification) provide measures of both data uncertainty (i.e. how uncertain each model is on average) and model uncertainty (i.e. how much the models disagree in their predictions). When naively distilling an ensemble to a single model, the ability to decompose total uncertainty into data uncertainty and model uncertainty is lost. This paper describes a method for ensemble distillation that retains both types of uncertainty in the case of classification. The idea is to train a prior network (i.e. a conditional Dirichlet distribution) to model the distribution of categorical probabilities within the ensemble. Both total uncertainty and data uncertainty can then be computed analytically from the prior network.