Generalized Bayesian Posterior Expectation Distillation for Deep Neural Networks

1 INTRODUCTION . Deep learning models have shown promising results in the areas including computer vision , natural language processing , speech recognition , and more ( Krizhevsky et al. , 2012 ; Graves et al. , 2013a ; b ; Huang et al. , 2016 ; Devlin et al. , 2018 ) . However , existing point estimation-based training methods for these models may result in predictive uncertainties that are not well calibrated , including the occurrence of confident errors . It is well-known that Bayesian inference can often provide more robust posterior predictive distributions in the classification setting compared to the use of point estimation-based training . However , the integrals required to perform Bayesian inference in neural network models are also well-known to be intractable . Monte Carlo methods provide one solution to representing neural network parameter posteriors as ensembles of networks , but this can require large amounts of both storage and compute time ( Neal , 1996 ; Welling & Teh , 2011 ) . To help overcome these problems , Balan et al . ( 2015 ) introduced an interesting model training method referred to as Bayesian Dark Knowledge . In the classification setting , Bayesian Dark Knowledge attempts to compress the Bayesian posterior predictive distribution induced by the full parameter posterior of a “ teacher ” network into a “ student ” network . The parameter posterior of the teacher network is represented through a Monte Carlo ensemble of specific instances of the teacher network ( the teacher ensemble ) , and the analytically intractable posterior predictive distributions are approximated as Monte Carlo averages over the output of the networks in the teacher ensemble . The major advantage of this approach is that the computational complexity of prediction at test time is drastically reduced compared to computing Monte Carlo averages over a large ensemble of networks . As a result , methods of this type have the potential to be much better suited to learning models for deployment in resource constrained settings . In this paper , we present a Bayesian posterior distillation framework that generalizes the Bayesian Dark Knowledge approach in several significant directions . The primary modeling and algorithmic contributions of this work are : ( 1 ) we generalize the target of distillation in the classification case from the posterior predictive distribution to general posterior expectations ; ( 2 ) we generalize the student architecture from being restricted to match the teacher architecture to being a free choice in the distillation procedure . The primary empirical contributions of this work are ( 1 ) evaluating the distillation of both the posterior predictive distribution and expected posterior entropy across a range of models and data sets including manipulations of data sets that increase posterior uncertainty ; and ( 2 ) evaluating the impact of the student model architecture on distillation performance including the investigation of sparsity-inducing regularization and pruning for student model architecture optimization . The key empirical findings are that ( 1 ) distilling into a student model that matches the architecture of the teacher , as in Balan et al . ( 2015 ) , can be sub-optimal ; and ( 2 ) student architecture optimization methods can identify significantly improved student models . We note that the significance of generalizing distillation to arbitrary posterior expectations is that it allows us to capture a wider range of useful statistics of the posterior that are of interest from an uncertainty quantification perspective . As noted above , we focus on the case of distilling the expected posterior entropy in addition to the posterior predictive distribution itself . When combined with the entropy of the posterior predictive distribution , the expected posterior entropy enables disentangling model uncertainty ( epistemic uncertainty ) from fundamental uncertainty due to class overlap ( aleatoric uncertainty ) . This distinction is extremely important in determining why predictions are uncertain for a given data case . Indeed , the difference between these two terms is the basis for the Bayesian active learning by disagreement ( BALD ) score used in active learning , which samples instances with the goal of minimizing model uncertainty ( Houlsby et al. , 2011 ) . The remainder of this paper is organized as follows . In the next section , we begin by presenting background material and related work in Section 2 . In Section 3 , we present the proposed framework and associated Generalized Posterior Expectation Distillation ( GPED ) algorithm . In Section 4 , we present experiments and results . Additional details regarding data sets and experiments can be found in Appendix A , with supplemental results included in Appendix B . 2 BACKGROUND AND RELATED WORK . In this section we present background material on Bayesian inference for neural networks , and related work on approximate inference , and model compression and pruning . 2.1 BAYESIAN NEURAL NETWORKS . Let p ( y|x , θ ) represent the probability distribution induced by a deep neural network classifier over classes y ∈ Y = { 1 , .. , C } given feature vectors x ∈ RD . The most common way to fit a model of this type given a data set D = { ( xi , yi ) |1 ≤ i ≤ N } is to use maximum conditional likelihood estimation , or equivalently , cross entropy loss minimization ( or their penalized or regularized variants ) . However , when the volume of labeled data is low , there can be multiple advantages to considering a full Bayesian treatment of the model . Instead of attempting to find the single ( locally ) optimal parameter set θ∗ according to a given criterion , Bayesian inference uses Bayes rule to define the posterior distribution p ( θ|D , θ0 ) over the unknown parameters θ given a prior distribution P ( θ|θ0 ) with prior parameters θ0 as seen in Equation 1. p ( θ|D , θ0 ) = p ( D|θ ) p ( θ|θ 0 ) ∫ p ( D|θ ) p ( θ|θ0 ) dθ ( 1 ) p ( y|x , D , θ0 ) = ∫ p ( y|x , θ ) p ( θ|D , θ0 ) dθ = Ep ( θ|D , θ0 ) [ p ( y|x , θ ) ] ( 2 ) For prediction problems in machine learning , the quantity of interest is typically not the parameter posterior itself , but the posterior predictive distribution p ( y|x , D , θ0 ) obtained from it as seen in Equation 2 . The primary problem with applying Bayesian inference to neural network models is that the distributions p ( θ|D , θ0 ) and p ( y|x , D , θ0 ) are not available in closed form , so approximations are required , which we discuss in the next section . 2.2 APPROXIMATE INFERENCE METHODS FOR BAYESIAN NEURAL NETWORKS . Most Bayesian inference approximations studied in the machine learning literature are based on variational inference ( VI ) ( Jordan et al. , 1999 ) or Markov Chain Monte Carlo ( MCMC ) methods ( Neal , 1996 ; Welling & Teh , 2011 ) . In VI , an auxiliary distribution qφ ( θ ) is defined to approximate the true parameter posterior p ( θ|D , θ0 ) . The variational parameters φ are selected to minimize the Kullback-Leibler ( KL ) divergence between qφ ( θ ) and p ( θ|D , θ0 ) . Hinton & Van Camp ( 1993 ) first studied applying VI to neural networks . Graves ( 2011 ) later presented a method based on stochastic VI with improved scalability . In the closely related family of expectation propagation ( EP ) methods ( Minka , 2001 ) , Soudry et al . ( 2014 ) present an online EP algorithm for neural networks with the flexibility of representing both continuous and discrete weights . Hernández-Lobato & Adams ( 2015 ) present the probabilistic backpropagation ( PBP ) algorithm for approximate Bayesian learning of neural network models , which is an example of an assumed density filtering ( ADF ) algorithm that , like VI and EP , generally relies on simplified posterior densities . The main drawback of VB , EP , and ADF is that they all typically result in biased posterior estimates for complex posterior distributions . MCMC methods provide an alternative family of sampling-based posterior approximations that are unbiased , but are often computationally more expensive to use at training time . MCMC methods allow for drawing a correlated sequence of samples θt ∼ p ( θ|D , θ0 ) from the parameter posterior . These samples can then be used to approximate the posterior predictive distribution as a Monte Carlo average as shown in Equation 3. p ( y|x , D , θ0 ) ≈ 1 T T∑ t=1 p ( y|x , θt ) ; θt ∼ p ( θ|D , θ0 ) ( 3 ) Neal ( 1996 ) addressed the problem of Bayesian inference in neural networks using Hamiltonian Monte Carlo ( HMC ) to provide a set of posterior samples . A bottleneck with this method is that it uses the full dataset when computing the gradient needed by HMC , which is problematic for larger data sets . While this scalability problem has largely been solved by more recent methods such as stochastic gradient Langevin dynamics ( SGLD ) ( Welling & Teh , 2011 ) , the problem of needing to compute over a large set of samples when making predictions at test or deployment time remains . Bayesian Dark Knowledge ( Balan et al. , 2015 ) is precisely aimed at reducing the test-time computational complexity of Monte Carlo-based approximations for neural networks . In particular , the method uses SGLD to approximate the posterior distribution using a set of posterior parameter samples . These samples can be thought of as an ensemble of neural network models with identical architectures , but different parameter values . This posterior ensemble is used as the “ teacher ” in a distillation process that trains a single “ student ” model to match the teacher ensemble ’ s posterior predictive distribution ( Hinton et al. , 2015 ) . The major advantage of this approach is that it can drastically reduce the test time computational complexity of posterior predictive inference relative to using a Monte Carlo average computed using many samples . Finally , we note that with the advent of Generative Adversarial Networks ( Goodfellow et al. , 2014 ) , there has also been work on generative models for approximating posterior sampling . Wang et al . ( 2018 ) and Henning et al . ( 2018 ) both propose methods for learning to generate samples that mimic those produced by SGLD . However , while these approaches may provide a speed-up relative to running SGLD itself , the resulting samples must still be used in a Monte Carlo average to compute a posterior predictive distribution in the case of Bayesian neural networks . This is again a potentially costly operation and is exactly the computation that Bayesian Dark Knowledge addresses . 2.3 MODEL COMPRESSION AND PRUNING . As noted above , the problem that Bayesian Dark Knowledge attempts to solve is reducing the test-time computational complexity of using a Monte-Carlo posterior to make predictions . In this work , we are particularly concerned with the issue of enabling test-time speed-storage-accuracy trade-offs . The relevant background material includes methods for network compression and pruning . Previous work has shown that overparameterised deep learning models tend to show much better learnability . Further , it has also been shown that such overparameterised models rarely use their full capacity and can often be pruned back substatially without significant loss of generality ( Hassibi et al. , 1993 ; LeCun et al. , 1989 ; Luo et al. , 2017 ; Louizos et al. , 2017 ; Frankle & Carbin , 2018 ; Han et al. , 2015 ; Zhang & Ou , 2018 ) . Hassibi et al . ( 1993 ) ; LeCun et al . ( 1989 ) use the second order derivatives of the objective function to guide pruning network connections . More recently , Han et al . ( 2015 ) introduced a weight magnitude-based technique for pruning connections in deep neural networks using simple thresholding . Guo et al . ( 2016 ) ; Jin et al . ( 2016 ) ; Han et al . ( 2016 ) introduce thresholding methods which also support restoration of connections . A related line of work includes pruning neurons/channels/filters instead of individual weights . Pruning these components explicitly reduces the number of computations by making the networks smaller . Group LASSO-based methods have the advantage of turning the pruning problem into a continuous optimization problem with a sparsity-inducing regularizer . Zhang & Ou ( 2018 ) ; Alvarez & Salzmann ( 2016 ) ; Wen et al . ( 2016 ) ; He et al . ( 2017 ) are some examples that use Group LASSO regularization at their core . Similarly Louizos et al . ( 2017 ) use hierarchical priors to prune neurons instead of weights . An advantage of these methods over ones which induce connection-based sparsity is that these methods directly produce smaller networks after pruning ( e.g. , fewer units or channels ) as opposed to networks with sparse weight matrices . This makes it easier to realize the resulting computational savings , even on platforms that do not directly support sparse matrix operations .

Summary: The paper introduces a general framework for distilling expectations of the Bayesian posterior distribution of a deep neural network, aiming to extend the original Bayesian Dark Knowledge approach [1]. More concretely, the generalized framework takes as input a teacher network, a general posterior expectation of interest, a student network, and thus performs an online compression of the selected posterior expectation using iteratively generated Monte Carlo samples from the parameter posterior of the teacher model. The proposed framework is applied to the case of classification models and empirical results demonstrate that distilling into a student model with an architecture that matches the teacher, as is done in Bayesian Dark Knowledge, can lead to sub-optimal performance. It is also shown that student architecture search methods can identify student models with significantly improved speed-storage-accuracy trade-offs.