Generalized Bayesian Posterior Expectation Distillation for Deep Neural Networks
Authors: Meet P. Vadera
Benjamin M. Marlin
Venue:  (2020)
Keywords: Bayesian Neural Networks
Distillation

Abstract: In this paper, we present a general framework for distilling expectations with respect to the Bayesian posterior distribution of a deep neural network, significantly extending prior work on a method known as ``Bayesian Dark Knowledge."  Our generalized framework applies to the case of classification models and takes as input the architecture of a ``teacher" network, a general posterior expectation of interest, and the architecture of a ``student" network. The distillation method performs an online compression of the selected posterior expectation using iteratively generated Monte Carlo samples from the parameter posterior of the teacher model. We further consider the problem of optimizing the student model architecture with respect to an accuracy-speed-storage trade-off. We present experimental results investigating multiple data sets, distillation targets,  teacher model architectures, and approaches to searching for student model architectures. We establish the key result 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. Lastly, we show that student architecture search methods can identify student models with significantly improved performance. 

Under review as a conference paper at ICLR 2020
GENERALIZED BAYESIAN POSTERIOR EXPECTATION DISTILLATION FOR DEEP NEURAL NETWORKS
Anonymous authors Paper under double-blind review
ABSTRACT
In this paper, we present general framework for distilling expectations with respect to the Bayesian posterior distribution of a deep neural network, significantly extending prior work a method known "Bayesian Dark Knowledge." Our generalized framework applies to the case of classification models and takes as input the architecture of a "teacher" network, a general posterior expectation of interest, and the architecture of a "student" network. The distillation method performs an online compression of the selected posterior expectation using iteratively generated Monte Carlo samples from the parameter posterior of the teacher model. We further consider the problem of optimizing the student model architecture with respect to accuracy-speed-storee trade-off. We present experimental results investigating multiple data sets, distillation targets, teacher model architectures, and approaches searching for student model architectures. We establish the key result 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. Lastly, we show that student architecture search methods can identify student models with significantly improved performance.
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 al.] [2013a b) Huang et al. 2016 Devlin et al.) 2018). However, existing point estimation- training methods for these models may result predictive uncertainties that are not well calibrated, including the occurrence of confident errors.
is well-known that Bayesian inference can often provide more robust posterior predictive distribu- tions 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 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 "teacher" network into "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 large ensemble of networks.
drastically reduced compared to computing Monte Carlo averages 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 general posterior expectations; (2) we generalize the

---

Under review as a conference paper at ICLR 2020
student architecture from being restricted to match the teacher architecture being 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 pparsity-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 al. (2015), can be sub-optimal; and (2) student architecture optimization methods can identify significanlly improved student models.
We note that the significance of generalizing distillation to arbitrary posterior expectations that it allows us to capture 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 (alcatoric uncertainty). This distinction is extremely important in determining why predictions uncertain for 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 Bayesian inference for neural networks, and related work on approximate inference, and model compression and pruning.
2.1 BAYESIAN NEURAL NETWORKS
Let xx, 0) represent the probability distribution induced by a deep neural network classifier over classes E y = (1....C C'} given feature vectors x E RD. The most common way to fit model of this type given a data set estimation, 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 8, according to a given criterion. Bayesian inference uses Bayes rule to define the posterior distribution p( p(0)2 (D, (d) over the unknown parameters 0 given prior distribution P(0 (*) with prior parameters 8 as seen in Equation]
{(x, y1 S S N) is to use maximum conditional likelihood
(( ((D())((000)))))
p(0)D,00)=
[ [p(D (8)))))
/ plulx 0 (0 (0 = brvlxx 0)]
p(y1x, D,0")
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, 8*) obtained from as seen in Equation 2] The primary problem with applying Bayesian inference to neural network models is that the distributions p(0)D, (o) and p/y/x, D. (%) are not available in closed form, so approximations required, which we discuss in the next section.
2

---

Under review as a conference paper at ICLR 2020
2.2
APPROXIMATE INFERENCE METHODS FOR BAYESIAN NEURAL NETWORKS
Most Bayesian inference approximations studied in the machine learning literature are based on variational inference (VD) (Jordan et al. 1999) or Markov Chain Monte Carlo (MCMC) methods (Neal 1996; Welling & Teh 2011) In VI, an auxiliary distribution (0) is defined to approximate the true parameter posterior p 89). The variational parameters & are selected to minimize the Kullback-Leibler (KI.) divergence between (0( and p(0|D, 80)) Hinton & Van Camp (993) first studied applying VI to neural networks. Graves (2011) later presenied a method based on stochastic VI with improved scalability. In the closely related family of expectation propagation (EP) methods (Minka (00T). Soudry et al. (2014) present an online EP algorithm for neural networks with the flexibility of representing both continuous and discrete weights. Homnnndez-Lobatte & 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 VT and EP, generally relies on simplified posterior densities.
The main drawback of VB, EP, and ADF that they all typically result in biased posterior estimates for complex posterior distributions. MCMC methods provide an alternative family of campling-based posterior approximations that are unbiased, but are often computationally more expensive to use at training time. MCMC methods allow for drawing correlated sequence of samples 8, ~ p(0 8²) from the parameter posterior. These samples can then be used to approximate the posterior predictive distribution Monte Carlo average as shown in Equation 3
plu|x, 0"
ply|x, 0(); 0₂ 0~~ P()
T
Neal (1996) addressed the problem of Bayesian inference in neural networks using Hamiltonian Monte Carlo (HMC) provide 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 & TTh 201T), the problem of needing to compute over 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 compu- tational 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 distillation process that trains a single "student" model to match the teacher ensemble's '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 Generative Adversarial Networks (Goodfellow et al., 2014), there has also been work 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 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 reducing the test-time computational complexity of using a Monte-Catte posterior to make predictions. In this work, we particularly concerned with the issue of enabling test-time speed-storage-accuss relevant background material includes methods for network compression and pruning.
trade-offs. The
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
let al. 1993 LeCun et al. 1989 Luo al. 2017
Louizos et al. 2017 Frankle & Carbin 2018
3

---

Under review as a conference paper at ICLR 2020
Han et al. 2013: 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, Hann et al. 2013) introduced a weight magnitude- -asded 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/chammels/ileer 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
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 casier to realize the resulting computational savings, even on platforms that do not directly support sparse matrix operations.
3 PROPOSED FRAMEWORK
In this section, we describe our proposed framework for distilling general Bayesian posterior ex- pectations for neural network classification models and discuss methods for enabling test-time speed-storaae--cccraayy trade-offs for flexible deployment of the resulting models.
3.1 GENERALIZED POSTERIOR EXPECTATIONS
There are many possible inferences of interest given a Bayesian parameter posterior P(0[D, 0%). We consider the general case of inferences that take the form of posterior expectations as shown in Equation 4 where g(y, x. 0) is arbitrary function of y. x and 0.
10 0)] / g(y. x, 0) (D, )d0
(4)
Important examples of functions gly, x, 0) include g(v. x, 0) = p(yl x, 0), which results in pos- terior expectation yielding the posterior predictive distribution p(y) D, 89), as used in Bayesian Dark Knowledge; gyy x, 0) ES - p(o'l 0) Mogg(y' (/'l/x |x, 0). which yields the posterior predictive entropy H x, D, 80 I and g x, = p(y/x, (1 p(y)x; x, 0)). which results the posterior marginal variance o² (y D, O)) While the posterior predictive distribution ply|x, D, 0²) is cer- tainly the most important posterior inference from predictive standpoint, the entropy and variance are also important from the perspective of uncertainty quantification.
3.2 GENERALIZED POSTERIOR EXPECTATION DISTILLATION
Our goal is to learn to approximate posterior expectations eo x, 0)] under given teacher model architecture using a given student model architecture. The method that we propose takes as input the teacher model p(v) x. 0), the prior p(0)0%), labeled data set D, an unlabeled data set D', the function ( x, 0), a student model [y, x (d),, an online expectation estimator, and loss function ((-..) that measures the error of the approximation given by the student model f(y x o). Similar to Balan et al. (015), we propose an online distillation method based on the use of the SGLD sampler. We describe all of the components of the framework in the sections below, and provide complete description of the resulting method in Algorithm
We define the prior distribution over the parameters p(0]00) be spherical Gaussian distribution centered at 0 with precision (we thus have 0 = [1,, r)). We define S to be minibatch of size M drawn from D. O, denotes the parameter set sampled for the teacher model at sampling iteration f, while denotes the step size for the teacher model at iteration t. The Langevin noise is denoted by 2 A (0, ncl). The sampling update for SGLD given by:
SGLD Sampler:
'Note that the posterior predictive entropy represents the average entropy integrated over the parameter posterior. It is not equal to the entropy of the posterior predictive distribution p(y|x, D, () in general.
4

---

Under review as a conference paper at ICLR 2020
0t+1 t 0: + "t V logg)(()))))
N
Vo log (vil,, + 2t
(5)
M
FES
Distillation Procedure: For the distillation learning procedure, we make use of secondary unlabeled data set D' = [x, < i S N*.. This data set could use feature vectors from the primary data set D, larger data set. We note that due autocorrelation in the sampled teacher model parameters Qz.
we may not want to run a distillation update for every Monte Carlo sample drawn. We thus use two different iteration indices: for SGLD iterations and for distillation iterations.
On every distillation step s, we sample a minibatch S' from D' of size M'. For every data case in S', we update an estimate Oyis of the posterior expectation using the most recent parameter sample O.. obtaining an updated estimate 0yis+11 R Epte/D.e 80) (g x, 0)] (we discuss update schemes in the next section). Next, we use the minibatch of examples S' to update the student model. To do so, we take step in the gradient direction of the regularized empirical risk of the student model as shown below where a2 is the student model learning rate, R()) is the regularizer, and is the regularization hyper-parameter. We next discuss the estimation of the expectation targets Byis.
N' ΣΣ VES'
Dst1 t- +
V
pell
AVo
M
Given an explicit collection of posterior samples o1...., Gs, the standard Monte Carlo estimate of E [g(y, x, 8)] is simply 9yse = S}=19h (yy x,, 8,). However, this estimator requires retaining the sequence samples 01, Og. which may not be feasible in terms of storage cost. Instead, consider the application of online update function. We define mis be the count of the number of times data case has been sampled up to and including distillation iteration
Expectation Estimation:
An online update function U (ggiis 0z, mis) takes as input the current estimate of the expectation, the current sample of the model parameters, and the number of times data case i has been sampled, and produces updated estimate of the expectation gyss+1. Below, define two different versions of the function. Ug (ggie 6,, mis), updates gyis using the current sample only, while (0lii f2 mia) performs an online update equivalent to a full Monte Carlo average.
U, (igiis G2, m4a) = g(y. 0,)
U U Oz, mis) Mi+++ (mis Iyis + g(v, Xi, 0,))
We note that both update functions provide unbiased estimates of EpcorD. [g x, 0)] after suitable burn-in time B. The online update Uo() will generally result in much lower variance in the estimated values of Syiis but it comes at the cost of needing to explicitly maintain the expectation estimates Gyis across learning iterations, increasing the storage cost of the algorithm. It is worthwhile noting that the extra storage and computation cost required by U, grows linearly in the size of the training set for the student. By contrast, the fully stochastic update memoryless terms of past expectation estimates, so the estimated expectations Pyis do not need to be retained across iterations resulting in space savings.
General Algorithm and Special Cases:
We show complete description of the proposed method in Algorithm The algorithm takes input the teacher model p(y) x, 8), the parameters of the prior ((0))"), labeled data set DD an unlabeled data set D', the function g 0), the student model
[(, xlØ), an online expectation estimator U U[gges Ot, m,s), a loss function ((...) that measures the
error of the approximation given by f(y, x|o). a regularization function R() and regularization
syper-parameter X, minibatch sizes and M', the thinning interval parameter H, the SGLD burn-in time parameter B and step size schedules for the step sizes ni and a,,
We note that the original Bayesian Dark Knowledge method is recoverable as a special case of -p log(g). U U, and p(y|x, f) = fyy Φ) (e.g.. the architecture of the student is selected to match that of the teacher). The original approach also uses distillation data set D' obtained from D by adding randomly
this framework via the the choices g(y. x, 0) p(y/xx 0). ((p,g)
5

---

Under review as a conference paper at ICLR 2020
Algorithm 1 Generalized Posterior Expectation Distillation
1: procedure GPED(D, D', D''p(ulx, 0), 0% 9, f, e, R. M, H, B, A, (W {as)=z))
2: Initialize 0, cu. 8o. Pyau = 0, m4. 0.7o
3: for t = 0 to T do
4: Sample S from D with |||= M
5: 0 0 Iu (V log -00 M
Exes V Vo + zt
6:
if mod (t,, H) - 0 and > B then
7:
Sample S' from D' with |S M'
8:
for i S' do
9:
gyse+1 t (Âgie 8, mis)
mis+1 t Mis +
11:
end for
12:
Φ.41
6s + 0ss
N' M
Lies EyEy Vol J(..x.o)) AV +V (V)(()))
8 t s+1
14:
end if
15:
end for
16:
end procedure
generated noise to instances from D on each distillation iteration, taking advantage of the fact that the choice U, means that no aspect of the algorithm scales with |
Our general framework allows for other trade-offs, including reducing the variance in the esti- mates of Tyis the cost of additional storage in proportion [D'|. We also note that note that the loss function {p, q) --plog log(g) and the choice g(y, x, 0) p(y|x, 0) are somewhat of a special case when used together even when the full stochastic expectation update U, is used, the resulting distillation parameter gradient is unbiased. To distill posterior entropy, we set g(y, x, 0) = eeyyp(U)x pog x, 0), U. and ((h,1/) = -//.
3.3 MODEL COMPRESSION AND PRUNING
One of the primary motivations for the original Bayesian Dark Knowledge approach is that it provides an approximate inference framework that results in significant computational and storage savings at test time. However, drawback of the original approach is that the architecture of the student is chosen to match that of the teacher. As we will show Section 4, this will sometimes result in student network that has too little capacity to represent particular posterior expectation accurately. On the other hand, if we plan to deploy the student model in low resource compute environment, the teacher architecture may not meet the specified computational constraints. In either case, need a general approach for selecting an architecture for the student model.
To begin to explore this problem, we consider to basic approaches to choosing student model archi- tectures that enable trading off test time inference speed and storage for accuracy. A helpful aspect of the distillation process relative teacher model is available as a starting point. As a first approach, we consider wrapping the proposed GPED algorithm with an explicit search over set of student models that are "close" to the teacher. Specifically, we consider a search space obtained by starting from the teacher model and applying width multiplier to the width of every fully connected layer and kernel multiplier to the number of kernels in every convolutional layer. While this search requires exponential time in the number of layers, provides a baseline for evaluating other methods.
de novo architecture search problem is that the architecture of the
As an alternative approach with better computational complexity, we leverage the regularization function R(Ø) included in the GPED framework to prune large initial network using group 4 regularization (Zhang& & Oul 2018 Wen et al. 2016). To apply this approach, we first must partition the parameters the parameter vector & across K groups Gt. The form of the regularizer is R(o) - _k (Ejes, 03) o/) As is well-established in the literature, this regularizer causes all parameters in a group go to zero simultancously when they not needed in model. To use for model pruning for a unit in a fully connected layer, we collect all of that unit's inputs into a group. Similarly, collect all of the incoming weights for particular channel in convolution layer

---

Under review as a conference paper at ICLR 2020
together into a group. If all incoming weights associated with unit or a channel have magnitude below a small threshold e, we can explicitly remove them from the model, obtaining a more compact architecture. We also fine-tune our models after pruning.
Finally, we note that any number of weight compressing, pruning, and architecture search methods could be combined with the GPED framework. Our goal is not to exhaustively compare such methods, but rather to demonstrate that GPED is sensitive to the choice of student model to highlight the need for additional research on the problem of selecting student model architectures.
4 EXPERIMENTS AND RESULTS
In this section, we present experiments and results evaluating the proposed approach using multiple data sets, posterior expectations, teacher model architectures, student model architectures and basic architecture search methods. We begin by providing an overview of the experimental protocols used.
4.1 EXPERIMENTAL PROTOCOLS
Data Sets: We use the MNIST (LeCun 1998) and CIFARIO (Kizhevsky et al. 2009) data sets as base data sets in our experiments. In the case of MNIST, posterior predictive uncertainty very low, so we introduce two different modifications to explore the impact of uncertainty on distillation performance. The first modification is simply to subsample the data. The second modification is to introduce occlusions into the data set using randomly positioned square masks of different sizes, resulting masking rates from 0% to 86.2% For CIFARIO, we only use sub-sampling. Full details for both data sets and the manipulations applied can be found in Appendix
Models: We evaluate a total of three teacher models in this work: a three-layer fully connected network (FCNN) for MNIST matching the architecture used by Balan et al. (2013), four-layer convolutional network for MNIST, and a five-layer convolutional network for CIFARIO. Full details of the teacher model architectures model architectures, we use the teacher model architectures as base models and search over a space of layer width multipliers Ki and K₂ that can be used expand sets of layers in the teacher models.
given in Appendix [A.2] For exhaustive search for student
A full description of the search space of student models can be found in Appendix A.2
We consider distilling both the posterior predictive distribution and the posterior entropy, as described in the previous section. For the posterior predictive distribution, we use the stochastic expectation estimator U, while for entropy we used the full online update U.- We allow B 1000 burn-in iterations and total of T 10 training iterations. The prior hyper-parameters learning rate schedules and other parameters vary by data set distillation target and are fully described in Appendix A.2
Distillation Procedures:
4.2 EXPERIMENT 1: DISTILLING POSTERIOR EXPECTATIONS
For this experiment, we use the MNIST and CIFARIO datasets without any subsampling or masking. For each dataset and model, we consider separately distilling the posterior predictive distribution and the posterior entropy. We fix the architecture of the student to match that of the teacher. To evaluate the performance while distilling the posterior predictive distribution, we use the negative log- oogliikllhood (NLL) of the model on the test set. For evaluating the performance of distilling posterior entropy, we use the mean absolute difference between the teacher ensemble's entropy estimate and the student model output on the test set. The results are given in Table [T First, we note that the FCNN NLL results on MNIST closely replicate the results in Balan et al. (2015). as expected. We also note that the error in the entropy is low for both the FCNN and CNN architectures on MNIST. However, the student model fails to match the NLL of the teacher on CIFAR10 10 and the entropy MAE is also relatively high. In Experiment 2, will investigate the effect of increasing uncertainty on models applied to both data sets, while Experiment we will search for student model architectures that improve performance.
7

---

Under review as a conference paper at ICLR 2020
Model Dataset
Teacher NLLL
Student NLL
MAB (Entropy)
FCNN MNIST
0.052
0.082
0.016
CNN MNIST
0.022
0.053
0.016
CNN CIFAR10
0.671
0.932
0.245
Table 1: Results of posterior distillation when the student architecture is fixed match the teacher architecture and the base data sets are used with no sub-sampling or occlusion.
2
NLL(S) (NL(T)
MAE Entropy (Test)
NLL(T)
0.2
Masking Rate
0.0
0.13
0.29
10000
20000
30000 60000
10000
20000
30000
60000
0.0
10000 20000 30000 60000
0.51
0.8
No. of training samples
No. training samples
No. of training samples
(a)
(b)
(c)
NLL(S) -ULIIT)
MAE Entropy (Test)
0.75
0.4
0.2
0.50
0.2
0.1
0.25
0.0
0.0
10000
20000
50000
10000
20000
50000
10000 No.
20000
50000
No. of training samples
No. of training samples
training samples
(d)
(c)
(f)
Figure 1: Top Row: Distillation performance using CNNs MNIST while varying data set size and masking rate. (a) Test negative log likelihood of the teacher posterior predictive distribution. (b) Difference in test negative log likelihood between student and teacher posterior predictive distribution estimates. (c) Difference between teacher and student posterior entropy estimates on test data set. Bottom Row: Distillation performance using CNNs CIFAR10 while varying data size. (d) Test negative log likelihood of the teacher posterior predictive distribution. (e) Difference in test negative log likelihood between student and teacher posterior predictive distribution estimates. (f) Difference between teacher and student posterior entropy estimates on test data set. In the plots above, S denotes the student and T denotes the teacher.
4.3 EXPERIMENT 2: ROBUSTNESS TO UNCERTAINTY
This experiment builds Experiment by exploring methods for increasing posterior uncertainty MNIST (sub-sampling and masking) and CIFAR 10 (sub-sampling). We consider the cross product of four sub-sampling rates and six masking rates for MNIST and three sub-sampling rates for CIFAR 10. We consider the posterior predictive distribution and posterior entropy distillation targets. For the posterior predictive distribution we report the negative log likelihood (NLL) of the teacher, and the NLL gap between the teacher and student. For entropy, we report the mean absolute error between the teacher ensemble and the student. All metrics are evaluated on held-out test data. We also restrict the experiment to the case where the student architecture matches the teacher architecture, mirroring the Bayesian Dark Knowledge approach. In Figure [ we show the results for the convolutional models on MNIST and CIFAR 10 respectively. The FCNN results are similar to the CNN results on MNIST and are shown Figure f in Appendix B In Appendix we also provide performance comparison between the Uo and U, estimators while distilling posterior expectations.
As expected, the the NLL of the teacher decreases as the data set size decreases. We observe that changing the number of training samples has similar effect on NLL gap for both CIFARIO and MNIST. More specifically, for any fixed masking rate of MNIST (and zero masking rate for CIFAR10),
that the NLL difference between the student and teacher decreases with increasing training data. However, for MNIST we can see that the teacher NLL increases much more rapidly function of the masking rate. Moreover, the gap between the teacher and student peaks
8

---

Under review as a conference paper at ICLR 2020
Student
Teacher
Individual Student
Student
Teacher
Individual Student
0.5
0.5
NLL (Test)
0.4
NLL (Test)
0.4
0.3
0.3
0.2
0.2
0.0
0.5
1.0
1.5
2.0
2.5
3.0
3.5
4.0
0.0
0.5
1.0
1.5
2.0
2.5
3.0
3.5
4.0
FLOPS (x100)
FLOPS (x10')
(a)
(b)
Teacher
Individual Student
Teacher
Individual Student
0.5
0.5
NLL (Tex)
0.4
NLL (Test
0.4
0.3
0.3
0.2
0.2
0.05
0.25
0.05
0.15
0.25
Num, Parameters (x10*)
Num. Parameters (x109)
(c)
(d)
Figure 2: ALL-Storage-Computtaiie tradeoff while using CNNs on MNIST with masking rate 29%. (a,b) Test negative log likelihood of posterior predictive distribution FLOPS found using exhaustive search and group 4 4 with pruning. (c,d) Test negative log likelihood of posterior predictive distribution The optimal student model for this configuration is obtained with group / pruning. It has approximately 6.6x the number of parameters and 6.4x the FLOPS of the base student model.
storage found using exhaustive search and group L with pruning.
for moderate values of the masking rate. This fact is explained through the observation that when the masking rate is low, posterior uncertainty is low, and distillation is relatively easy. On the other hand, when the masking rate high, the teacher essentially outputs the uniform distribution for every example, which is very easy for the student represent. As result, the moderate values of the masking rate result in the hardest distillation problem and thus the largest performance gap. For varying masking rates, MNIST. However, the gap for entropy prediction increases as a function data set size for CIFARIO. Finally, would expect, the performance of distillation using the U, estimator is almost always better than that of the U, estimator (refer Appendix B
exactly the same trend for the gap in posterior entropy predictions
The key finding of this experiment is simply that the quality of the approximations provided by the student model varies as a function of properties of the underlying data set. Indeed, restricting the student architecture match the teacher sometimes result in significant performance gaps. In the next experiment, we address the problem of searching for improved student model architectures.
4.4 EXPERIMENT 3: TOWARDS STUDENT MODEL ARCHITECTURE SEARCH
In this experiment, we compare the exhaustive search the group / (group lasso) regularizer combined with pruning. For the pruming approach, we start with the largest student model considered under exhaustive search, and prune back from there using different regularization parameters leading to different student model architectures. We present results in terms of performance versus computation time (estimated FL.OPS), as well as performance vs storage cost (estimated in number of parameters). As performance measures for the posterior predictive distribution, accuracy and negative log likelihood. For entropy, we use mean absolute error. In all cases results are reported test data. We consider both fully connected and convolutional models.
consider
Figure 2 shows results for negative the log likelihood (NLL) of the convolutional model on MNIST with masking rate 29% and 60,000 training samples. We select this setting illustrative of a difficult case for posterior predictive distribution distillation. We plot NLL vs FLOPS and NLL vs storage for all points encountered in each search. The solid blue line indicates the Pareto frontier.

---

Under review as a conference paper at ICLR 2020
First, we note that the baseline student model (with architecture matching the teacher) from Experi- ment on MNIST achieves an NLL of 0.469 at approximately 0.48 x 10 FLOPs and 0.03 x 10 parameters on this configuration of the data set. We can that both methods for selecting student architectures provide highly significant improvement over the baseline student architectures. On MNIST, the NLL is reduced to 0.30. Further, we can also that the group & f2 approach is able obtain much better NLL at the same computation and storage cost relative to the exhaustive search method. Lastly, the group L l2 method able to obtain models on MNIST at less than 50% the computational cost needed by the baseline model with only a small loss in performance. Results for other models and distillation targets show similar trends and are presented in Appendix B Additional experimental details are given in Appendix A.2
In summary, the key finding of this experiment that the capacity of the student model has a significant impact on the performance of the distillation procedure, and methods for optimizing the student architecture are needed achieve a desired speed-storan pereddstorage-accuraa trade-off.
5 CONCLUSIONS & FUTURE DIRECTIONS
We have presented a framework for distilling expectations with respect to the Bayesian posterior distribution of deep neural network that generalizes the Bayesian Dark Knowledge approach in several significant directions. Our results show that the performance of posterior distillation can be highly sensitive to the architecture of the student model, but that basic architecture search methods can help to identify student model architectures with improved speed-storage-accuraa There are many directions for future work including considering the distillation of broader class of posterior statistics including percentiles, assessing and developing more advanced student model architecture search methods, and applying the framework to larger state-of-the-ar models.
trade-offs.
REFERENCES
Jose M Alvarez and Mathieu Salzmann. Learning the number of neurons in deep networks. In Advances in Neural Information Processing Systems, pp. 2270-2278, 2016.
Anoop Korattikara Balan, Vivek Rathod, Kevin P Murphy, and Max Welling. Bayesian dark knowledge. In Advances Neural Information Processing Systems, pp. 3438-3446, 2015.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In NAACL-HLT, 2018
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Training pruned neural networks. ArXiv, habs/1803.03635, 2018
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances neural informa- tion processing systems. pp. 1672-2680. 2014.
Alex Graves. Practical variational inference for neural networks. Advances in neural information processing systems, pp. 2348 22448,23666-11 2011.
Alex Graves, Navdeep Jaitly, and Abdel-rahman Mohamed. Hybrid speech recognition with deep bidirectional lstm. In Automatic Speech Recognition and Understanding (ASRU), 2013 IEEE Workshop on, pp. 273-278. IEEE, 2013a.
Alex Graves, Abdel-rahman Mohamed, and Geoffrey Hinton. Speech recognition with deep recurrent neural networks. In Acoustics, speech and signal processing (icassp), 2013 ieee international conference on, pp. 6645-6649. IEFF, 2013b.
Yiwen Guo, Anbang Yao, and Yurong Chen. Dynamic network surgery for efficient dnns. In NIPS, 2016.
Song Han, Jeff Pool, John Tran, and William J. Dally. Learning both weights and connections for efficient neural networks. In NIPS, 2015.
10

---

Under review as a conference paper at ICLR 2020
Song Han, Jeff Pool, Sharan Narang, Huizi Mao, Shijian Tang, Erich Elsen, Bryan Catanzaro, John Tran, and William J. Dally. Dsd: Regularizing deep neural networks with Sense-sparse-dense training flow. ArXiv, abs/1607,,4381, 2016.
Babak Hassibi, David G. Stork, and Gregory J. Wolff. Optimal brain surgeon and general network pruning. IEEE International Conference on Neural Networks, pp. 293-299 vol.1, 1993.
Yihui He, Xiangyu Zhang, and Jian Sun. Channel pruning for accelerating very deep neural networks. 2017 IEEE International Conference on Computer Vision (ICCV), pp. 1398-1406, 2017.
Christian Henning, Johannes von Oswald, João Sacramento, Simone Carlo Surace, Jean- Pasaal Pfister, and Benjamin Grewe. Approximating the predictive distribution via udversarially-train hypernetworks. In Bayesian Deep Learning Workshop, NeurlPS (Spotlight) 2018, 2018.
José Miguel Hernnndez-L.ae Lobato and Ryan Adams. Probabilistic backpropagation for scalable learning of bayesian neural networks. In International Conference Machine Learning, pp. 1861-1866 2015.
Geoffrey Hinton and Drew Van Camp. Keeping neural networks simple by minimizing the description length of the weights. In Proc. of the 6th Ann. ACM Conf. on Computational Learning Theory. Citeseer, 1993.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge neural network. arXiv preprint arXiv:1503.02531, 2015.
Neil Houlsby, Ferenc Huszár, Zoubin Ghahramani, and Máté Lengyel. Bayesian active learning for classification and preference learning. arXiv preprint arXiv: 2011.
Gao Huang, Zhuang Liu, Laurens van der Maaten, and Kilian Q. Weinberger. Densely connected convolutional networks. 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 2261-2261 2016.
Xiaojie Jin, Xiao-Tong Tong Yuan, Jiashi Feng, and Shuicheng Yan. Training skinny deep neural networks with iterative hard thresholding methods. ArXiv, abs/1607.00423,, 2016.
Michael Jordan, Zoubin Ghahramani, Tommi Jaakkola, and Lawrence Saul. An introduction to variational methods for graphical models. Machine learning, 37(2):183-233, 1999.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. Technical report, Citeseer, 2009.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolu- tional neural networks. In Advances in neural information processing systems, pp. 1097-1105, 2012.
Yann LeCun. The mnist database of handwritten digits. http://yann. lecun. com/exdb/mnist// 1998.
Yann LeCun, John S. Denker, and Sara A. Solla. Optimal brain damage. In NIPS, 1989.
Christos Louizos, Karen Ullrich, and Max Welling. Bayesian compression for deep learning. ArXiv, abs/1705.08665, 2017.
Jian-Hao Luo, Jianxin Wu, and Weiyao Lin. Thinet: A filter level pruning method for deep neural network compression. 2017 IEEE International Conference on Computer Vision (ICCV), pp. 5068-5076, 2017.
Thomas P Minka. Expectation propagation for approximate bayesian inference. In Proceedings of the Seventeenth conference on Uncertainty in artificial intelligence, pp. 362-369. Morgan Kaufmann Publishers Inc., 2001.
Radford M. Neal. Bayesian Learning for Neural Networks. Springer- Verlag, Berlin, Heidelberg, 1996. ISBN 0387947248.
11

---

Under review as a conference paper at ICLR 2020
Daniel Soudry, Itay Hubara, and Ron Meir. Expectation backpropagation: Parameter-free free training of multilayer neural networks with continuous or discrete weights. In Advances in Neural Information Processing Systems, pp. 963-971, 2014.
Kuan-Chieh Wang, Paul Vicol, James Lucas, Li Gu, Roger Grosse, and Richard Zemel. Adversarial distillation of bayesian neural network posteriors. arXiv preprint arXiv: 10317, 2018.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-I1), pp. 681-688,
2011.
Wei Wen, Chunpeng Wu, Yandan Wang, Yiran Chen, and Hai Li. Learning structured sparsity in deep neural networks. Advances in neural information processing systems, pp. 2074-2082, 2016.
Yichi Zhang and Zhijian Ou. Learning sparse structured ensembles with stochastic gradient meme sampling and network pruning. 2018 IEEE 28th International Workshop on Machine Learning for Signal Processing (MLSP), pp- 1-6, 2018.
A DATASETS AND MODEL DETAILS
A.1 DATASETS
As noted earlier in the paper. the original empirical investigation of Bayesian Dark Knowledge for classification focused on the MNIST data set (LeCun 1998). However, the models fit to the MNIST data set have very low posterior uncertainty and we argue that it is thus a poor benchmark for assessing the performance of posterior distillation methods. In this section, we investigate two orthogonal modifications of the standard MNIST data set to increase uncertainty: reducing the training set size and masking regions of the input images. Our goal is to produce a range benchmark problems with varying posterior predictive uncertainty. We also use the CIFAR10 10 data set (Krizhevsky et al.) 2009) in our experiments and employ the same subsampling technique.
MNIST: The full MNIST dataset consists of 60,000 training images and 10,000 test images, each of size 28 x 28, distributed among 10 classes classes[Leuun (1998). ). As a first manipulation, we consider sub-sampling the labeled training data to include 10,000, 20,000, 30,000 all 60,000 data cases in the primary data set DD when performing posterior sampling for the teacher model. Importantly, we use all 60,000 unlabeled training cases in the distillation data set D'. This allows de-couple the impact of reduced labeled training data on posterior predictive distributions from the effect of the amount of unlabeled data available for distillation.
As second manipulation, we generate images with occlusions by randomly masking out parts of each available training and test image. For generating such images, we randomly choose square m X m region (mask) and set the value for pixels in that region to O. Thus, the masking rate for 28 MNIST image corresponding to the mask of size m is given by x8 28' We illustrate original and masked data in Figure 3 between 0% and 86.2%.
We consider range of square masks resulting in masking rates
5041
/
9213
1 4
3 5
3 6 1
7
3
(a) Original images
(b) Processed images
Figure 3: Example MNIST data after masking with 14.
CIFARIO: The full CIFARIO dataset consists of 50,000 training images and 10,000 test images, each of size 32 x 32 pixels. We sub-sample the data into primary training sets D containing 10,000,
12

---

Under review as a conference paper at ICLR 2020
20,000, and 50,000 images. As with MNIST, the sub-sampling is limited training the teacher model only and we utilize all the 50,000 unlabeled training images in the distillation data set D'.
A.2 MODELS
To demonstrate the generalizability of our methods to a range of model architectures, we run our experiments with both connecco,,,,,ed, and convolutional neural networks. We note that our goal in this work is not to evaluate the GPED framework on state-of-the-art architectures, but rather to provide illustrative results and establish methodology for assessing the impact of several factors including the level of uncertainty and the architecture of the student model.
Teacher Models: We begin by defining the architectures used for the teacher model as follows:
1. FCNN (MNIST): We use a 3-layer fully connected neural network. The architecture used is:
This matches the architecture used by Balan
ct al. (2015).
2. CNN (MNIST): For CNN, use two consecutive sets of 2D convolution and max- pooling layers, followed by two fully-connected layers. The architecture used is: In- put(1, (88,28))-onn(nnn)) kernel==10, kernel size=4, stride=1) MaxPool(kernel size=2)- Conv(num 2 kernel _ize=4, stride=1) MaxPPoo((eermelssize2) FC (80) - FC (output).
3. CNN (CIFAR10): Similar to the CNN architecture used for MNIST, we use two consec- utive sets of 2D convolution and max-pooling layers followed by fully-connected layers. Conv(num kernels=16, kernel kemnelssize=5) MaxPool(kernel size=2) Conv(num kernels=32, sizeel MaxPool(kemel 2ize==) FC((00) FC (50) FC (output).
In the architectures mentioned above, the "output" size will change depending on the expectation that we're distilling. For classification, the output size will be 10 for both datasets, while for the case of entropy. it will be 1. We use ReLU non- non-lincassties everywhere between the hidden layers. For the final output layer, softmax used for classification. In the case of entropy, we use an exponential activiation to ensure positivity.
The student models used in our experiments use the above mentioned architectures as the base architecture. For explicitly searching the space of the student models, we use set of width multipliers starting from the teacher architecture. The space of student architectures corresponding to each teacher model defined earlier is given below. The width multiplier values of Ki and K2 are determined differently for each of the experiments, and thus will be mentioned in later sections.
Student Models:
1. FCNN (MNIST): Input7744 K1)-FC(400 FC)uuC((t
2. CNN (MNIST): Input(1, (onp)))CCooommmmkeeeeeeee ]0 K1J, kerneLsize=4,, stride=1) MaxPool(kernel size=2) Conv(num kernels= KiJ. kernel size=4, stride=1) Max- Pool(kernel size) FC ([80- K2)) FC (output).
3. CNN (CIFAR10): Input(3,, MaxPPoo(kernel_size-- Cony(num_kemeee kemnels= Pool(kernel.size=2) FC ([200 K2]) FC ([50- K2)) - FC (output).
kernels= 16 Kj), kernel size=5) K1]. kernel_sizess Max-
Model and Distillation Hyper-Parameterss We run the distillation procedure using the following hyperparameters: fixed teacher learning rate 17: 4 x 10- for models on MNIST and m 2 x 10 for models on CIFAR10, teacher prior precision - 10, initial student learning rate a, 10 10-3 student dropout rate 0.5 for fully-connected models on MNIST' (and zero otherwise), burn-in iterations B = 1000, thinning interval H 100 for distilling predictive means and H = 10 for distilling entropy values, and total training iterations T 10 For training the student model, the Adam algorithm (instead of plain steepest descent as indicated Algorithm 1) and set learning schedule for the student such that it halves its learning rate every 200 epochs for models on MNIST, and every 400 epochs for models on CIFAR10. Also, note that we only apply the regularization function R().) while doing Group f pruning. Otherwise, we dropout indicated before.
for Group / pruning experiments: For experiments involving group e / regularizer, the regularization strength values À are chosen from log-scale ranging from 10-8
13

---

Under review as a conference paper at ICLR 2020
NLL/Student) NL[[Teache)
0.20
NL.Teacher
0.15
Masking ftate
0.0
0.03
1.0
0.13
0.29
0.5
0.05
0.51
0.8
0.0
0.00
10000
20000
30000
60000
10000
20000
30000
60000
training samples
No. training samples
(a)
(b)
0.10
0.08
Masking Rate
0.06
0.03
0.13
0.04
0.29
0.02
0.51
No. of training samples
(c)
Figure 4: Distillation performance using Fully-Connected Networks MNIST while varying data set size and masking rate. (a) Test negative log likelihood of the teacher posterior predictive distribution. (b) Difference in negative log likelihood between teacher and student posterior predictive distribution estimates. (c) Difference between teacher and student posterior entropy estimates on test data set.
When using Group l//l2 regularizer, do not dropout for the student model. The number of fine-tuning epochs for models on MNIST and CIFAR 100 are 600 and 800 respectively. At the start of fine-tuning, also reinitialize the student learning rate o₂ 10-1 for fully-connected models and ar 10-3 for convolutional models. The magnitude threshold for pruning is 10-3
B SUPPLEMENTAL EXPERIMENTS AND RESULTS
Supplemental Results for Experiment 2: Robustness to Uncertainty In Figure 4 we demonstrate the results of Experiment (Section 4.3), on fully-connected networks for MNIST. Additionally, in Tables 2 we provide a performance comparison between U, and U, estimators while distilling posterior expectations for all model-data set combinations. We follow the same experimental configurations as in Experiment 2.
Supplemental Results for Experiment 3: Towards Student Model Architecture Search The additional results from running Experiment (Section 4.4) on different combinations of model type, dataset, and performance metries have been given Figures 12
14

---

Under review as a conference paper at ICLR 2020
Num. training samples
Masking rate
NLL (Teacher)
NLL (Student, U+)
NLL (Student, U.)
MAE (Entropy, U.)
MAE (Entropy, U.)
0
0.048
0.214
0.218
0.025
0.030
0.03
0.069
0,274
0.274
0.033
0.038
0.13
0.161
0.509
0.509
0.058
0.069
10000
0.29
0.394
0.902
0.907
0.115
0.129
0.51
1.099
1.615
1.630
0.194
0.170
0.8
2.298
2.301
2.301
0.016
0.019
0
0.034
0.126
0.126
0.020
0.021
0,03
0,054
0.180
0.181
0.026
0.030
0.13
0.123
0.342
0.344
0.053
0.066
20000
0.29
0.326
0.684
0.697
0.104
0.122
0.51
1,050
1.369
1.378
0.145
0.150
0.8
2.298
2.300
2.299
0.016
0.020
0
0.028
0.084
0.086
0.017
0.019
0.03
0.044
0.132
0.134
0.024
0.027
0.13
0.106
0.292
0.294
0.051
0.061
30000
0.29
0.300
0.620
0.618
0.106
0.120
0.51
1.044
1.307
1.308
0.130
0.141
0.8
2.296
2.297
2.296
0.017
0.021
0
0.022
0.053
0.053
0.016
0.017
0.03
0.035
0.088
0.090
0.025
0.026
0.13
0.090
0.219
0.221
0.049
0.058
60000
0.29
0.267
0.463
0.472
0.108
0.120
0.51
1.024
1.184
1.187
0.118
0.127
0.8
2.297
2.297
2.297
0.020
0.023
Table 2: Performance comparison between U, and U, estimators for convolutional neural network on MNIST. The NLL results correspond to the case of distilling the posterior predictive distribution while the MAE entropy results correspond to the case of distilling the expectation of predictive entropy.
15

---

Under review as a conference paper at ICLR 2020
Num. training samples
Masking rate
NLL (Teacher)
NLL (Student, U,)
NLL (Student, U,)
MAE (Entropy, U.)
MAE
(Entropy, U.)
0
0.137
0.184
0.243
0.013
0.018
0.03
0.180
0.233
0.300
0.018
0.023
0.13
0.312
0.389
0.483
0.031
0.040
10000
0.29
0,556
0.637
0.760
0.059
0.089
0.51
1.183
1.229
1.371
0.111
0.135
0.8
2.103
2.111
2.129
0.023
0.019
0
0.089
0.115
0.161
0.011
0.015
0.03
0.131
0.165
0.220
0.014
0.021
0.13
0,230
0.280
0.366
0.024
0.042
20000
0.29
0,452
0,510
0.607
0.049
0.104
0.51
1,080
1.120
1.215
0.094
0.112
0.8
2.104
2.108
2.117
0.019
0.021
0
0,071
0.083
0.124
0.011
0.014
0.03
0.107
0.129
0.180
0.015
0.021
0.13
0,201
0.243
0.314
0.028
0,052
30000
0.29
0,414
0.459
0.555
0.062
0.105
0.51
1.044
1.082
1.172
0.091
0.105
0.8
2.089
2.092
2.101
0.022
0.023
0
0.052
0.054
0.082
0.016
0.020
0.03
0.081
0.094
0.133
0.023
0.034
0.13
0.155
0.186
0.240
0.043
0.068
60000
0.29
0.360
0.398
0.471
0.086
0.109
0.51
1.010
1.033
1.107
0.106
0.099
0.8
2.088
2.089
2.094
0.021
0.022
Table 3: Performance comparison between U, and U, estimators for fully-connected network MNIST, The NLL results correspond to the case of distilling the posterior predictive distribution while the MAE on entropy results correspond to the case of distilling the expectation of predictive entropy.
Num. training samples
NLL
NLL (Student, U.)
NLL (Student, U.)
MAE (Entropy, U.)
MAE (Entropy, U,)
(Teacher)
10000
0.912
1.372
1.391
0.144
0.192
20000
0.798
1.184
1.179
0.210
0.231
50000
0.671
0.924
0.932
0.245
0.290
Table 4: Performance comparison between U, and U, estimators for convolutional neural network on CIFARIO. The NLL results correspond to the case of distilling the posterior predictive distribution while the MAE on entropy results correspond to the case of distilling the expectation of predictive entropy.
16

---

Under review as a conference paper at ICLR 2020
0.94
Student Accuracy (Test)
0.94
Student Accuracy (Test)
Teacher Accuracy (Test)
Teacher Accuracy (Test)
Accuracy (Test)
0.92
Individual Student
Accuracy (Test)
Individual Student
0.90
0.90
0.88
0.88
0.84
0.84
0.5
1.0
1.5
2.0
2.5
3.0
3.5
4.0
0.0
0.5
1.0
1.5
2.0
2.5
3.0
3.5
4.0
FLOPS (x109)
FLOPS (x10%)
(h)
0.94
Student Accuracy (Test)
0.94
Student Accuracy (Test)
Teacher Accuracy (Test)
Teacher Accuracy (Test)
Accuracy (Test)
Individual Student
Accuracy (Test)
0.92
Individual Student
0.90
0.90
0.88
0.88
0.86
0.86
0.84
0.84
0.00
0.05
0.10
0.15
0.20
0.25
0.30
0.00
0.05
0.10
0.15
0.20
0.25
0.30
Num. Parameters (x10f)
Num. Parameters (x10*)
(c)
(d)
Figure 5: Saarrryyy-toonge-Compaaaaii tradeoff while using CNNs on MNIST with masking rate 29%. (a) Test accuracy using posterior predictive distribution FLOPS found using exhaustive search. (b) Test accuracy using posterior predictive distribution vs FLOPS found using group f with pruning. (c) Test accuracy using posterior predictive distribution storage found using exhaustive search. (d) Test accuracy using posterior predictive distribution vs storage found using group f with pruning. The optimal student model for this configuration is obtained with group l pruning. It has approximately 6.6x the number of parameters and 6.4x the FLOPS of the base student model.
17

---

Under review as a conference paper at ICLR 2020
MAE- Entropy Output (Test)
0.110
MAE MAE-Em E-EErropy Output (Test)
0.110
MAE- Entropy Output (Test)
MAE- Entropy Output (Test)
Individual Student
Individual Student
0,100
0.100
0.095
0.095
0.090
0.090
0.085
0.085
0.0
1.5
2.0
2.5
3.0
3.5
4.0
0.0
0.5
1.0
2.5
3.0
3.5
4.0
(x10%)
FLOPS (x10*)
(a)
(b)
MAE-ER Entropy Output (Test)
0.110
IAE-Entropy Output (Test)
0.110
MAE Entropy Output (Test)
0.105
0.105
Individual Student
0.100
0.100
0.095
0.095
0.090
0.090
0.085
0.085
MAE: Entropy Output (Test)
Individual Student
0.080
0.060
0.00
0.05
0.10
0.15
0.20
0.25
0.30
0.00
0.05
0.10
0.15
0.20
0.25
0.30
Parameters (x106)
Parameters (x109)
(c)
(d)
Figure 6: Entropy Error-Storage www--Sttote-Co-ooee Computation tradeoff while using CNNs on MNIST with masking rate 29%. (a) Test mean absolute error for posterior entropy vs FLOPS found using exhaustive search. (b) Test mean absolute error for posterior entropy vs FLOPS found using group 4 with pruning. (c) Test mean absolute error for posterior entropy vs storage found using exhaustive search. (d) Test mean absolute error for posterior entropy vs storage found using group E with pruning. The optimal student model for this configuration is obtained with group é /22 pruning. It has approximately 1.8x the number of parameters and 4.3x x the FLOPS of the base student model.
18

---

Under review as a conference paper at ICLR 2020
0.90
0.90
Accuracy (Test))
0.88
Accuracy (Test)
0.88
0.86
0.86
0.84
0.84
StudentAccuracy Accuracy (Test)
Student
Student Accuracy (Test)
0.82
Teacher Accuracy (Test)
Teacher Accuracy (Test)
Individual Student
Individual Student
6
10
FLOPS (x10*)
FLOPS (x10*)
(a)
(b)
0.90
0.90
Accuracy (Test)
0.88
Accuracy (Test)
0.88
0.86
0.86
0.84
0.84
Student Accuracy (Test)
Student Accuracy (Test)
0.82
Teacher Accuracy (Test)
0.82
Teacher Accuracy (Test)
Individual Student
Individual Student
0.80
0.80
6
Num. Parameters (x10f)
Num. Parameters (x10*)
(c)
(d)
Figure 7: Sccuracy-Storage-Cnee Citorage-Computattionn tradeoff while using Fully-connected networks MNIST with masking rate 29%. (a) Test accuracy using posterior predictive distribution vs FLOPS found using exhaustive search. (b) Test accuracy using posterior predictive distribution vs FLOPS found using group f with pruming. (c) Test accuracy using posterior predictive distribution vs storage found using exhaustive search. (d) Test accuracy using posterior predictive distribution vs storage found using group f1/22 with pruning. The optimal student model for this configuration is obtained with group / pruning. has approximately 9.9x the number parameters and 10x the FLOPS of the base student model.
19

---

Under review as a conference paper at ICLR 2020
Student NLL (Test)
Student NLL (Test)
Teacher NLL (Test)
Teacher NLL (Test)
0.45
Individual Student
0.45
Individual Student
NLL N (Test)
NLL (Test)
0.40
0.40
0.35
0.35
0.30
0.30
FLOPS (x10%)
FLOPS (x10F)
(a)
(b)
0.50
0.50
Student NLL (Test)
Student NLL (Test)
Teacher NLL (Test)
Teacher NLL (Test)
0.45
Individual Student
0.45
Individual Student
NLL (Test)
NLL (Test)
0.40
0.40
0.30
Num. Parameters (x10%)
Num. Parameters (x106)
(c)
(d)
Figure 8: NLL-Storage- AiL-Storage-Computatioo tradeoff while using Fully-connected networks on MNIST with masking rate 29%. (a) Test negative log likelihood of posterior predictive distribution vs FLOPS found using exhaustive search. (b) Test negative log likelihood of posterior predictive distribution vs FLOPS found using group & L22 with pruming. (c) Test negative log likelihood of posterior predictive distribution storage found using exhaustive search. (d) Test negative log likelihood of posterior predictive distribution vs storage found using group f with pruning. The optimal student model for this configuration is obtained with group f pruning. has approximately 9.9 the number of parameters and 10x the FLOPS of the base student model.
20

---

Under review as a conference paper at ICLR 2020
MAE-Entropy Output (Test)
MAE-Entropy Output (Test)
0.14
Entropy Output (Test)
0.14
Entropy Output (Test)
Individual Student
Individual Student
0.12
0.12
0.10
0.10
0.08
0.08
0.06
0.06
6
10
0
8
FLOPS (x10*)
FLOPS (x10*)
(a)
(b)
MAE-Entropy Output (Test)
MAE-Entropy Output (Test)
0.14
MAE- Entropy Output (Test)
0.14
MAE- Entropy Output (Test)
Individual Student
Individual Student
0.12
0.12
0.10
0.10
0.08
0.08
0.06
0.06
2
1
6
Num. Parameters (x10f)
Num. Parameters (x10*)
(c)
(d)
Figure 9: Entropy Proo--Soorae-Compuaane tradeoff while using Fully-connected networks MNIST with masking rate 29%. (a) Test mean absolute eπоо for posterior entropy vs FLOPS found using exhaustive search. (b) Test mean absolute error for posterior entropy vs FLOPS found using group / with pruning. (c) Test mean absolute erгоr for posterior entropy vs storage found using exhaustive search. (d) Test mean absolute error for posterior entropy vs storage found using group f with pruning. The optimal student model for this configuration is obtained with group 4 6/6 pruning. has approximately 4.2x the number of parameters and 4.2x the FLOPS of base student model.
21

---

Under review as a conference paper at ICLR 2020
1.2
1.2
Student (Test)
1.1
1.1
Teacher (Test)
Individual Student
NLL (Test)
NLT (Test)
1.0
1.0
0.9
0.9
Student NLL (Test)
Teacher NLL (Test)
0.8
Individual Student
0.0
10.0
12.5
15.0
17.5
0.0
2.5
5.0
10.0
12.5
15.0
17.5
FLOPS (x109)
FLOPS (x10%)
(a)
(b)
1.2
1.2
Student NLL (Test)
Student NLL (Test)
Teacher NLL (Test)
1.1
Teacher NLL (Test)
Individual Student
Individual Student
NLL (Test)
NLL (Test)
1.0
1.0
0.9
0.8
0.8
0.7
0.7
Num. Parameters (x10%)
Num. Parameters (x10*)
(c)
(d)
Figure 10: MLL-Storage-Compuatton tradeoff while using CNNs on CIFARIO with training set size of 20,000 samples. (a) Test negative log likelihood of posterior predictive distribution FL.OPS found using exhaustive search.
(b) Test negative log likelihood of posterior predictive distribution vs FLOPS found using group & / with pruning. (c) Test negative log likelihood of posterior predictive distribution vs storage found using exhaustive search. (d) Test negative log likelihood of posterior predictive distribution vs storage found using group & f with pruning. The optimal student model for this configuration is obtained with group L pruning. It has approximately 4.7x x the number of parameters and 5.2x the FLOPS of the base student model.
22

---

Under review as a conference paper at ICLR 2020
0.750
0,750
0.725
0.725
Accuracy (Test)
Accuracy (Test)
0.700
0.700
Student Accuracy (Test)
Teacher Accuracy (Test)
Individual Student
Student Accuracy (Test)
0.625
0.625
Teacher Accuracy
Individual Student
0.600
0.600
12.5
7.5
15.0
FLOPS (x10f)
FLOPS (x109)
(a)
(b)
0.750
0.750
0.725
Accuracy (Test)
Accuracy (Test)
0.700
0.700
Student Accuracy (Test)
0.675
Teacher Accuracy (Test)
0.675
Individual Student
0.650
0.650
Student Accuracy (Test)
0.625
Accuracy (Test)
Individual Student
0.600
0.600
1.00 1.25
1.50
Num. Parameters (x10f)
Num. Parameters (x10*)
(c)
(d)
(ceuracy-Storage Chorage-Cooppuationnn tradeoff while using CNNs CIFARIO with sub- sampling training data to 20,000 samples. (a) Test accuracy using posterior predictive distribution vs FLOPS found using exhaustive search. (b) Test accuracy using posterior predictive distribution vs FLOPS found using group L with pruning. (c) Test accuracy using posterior predictive distribution storage found using exhaustive search. (d) Test accuracy using posterior predictive distribution vs storage found using group & f//Ez with pruning. The optimal student model for this configuration is obtained with group / pruning. has approximately 5.4x the number of parameters and 5.6x the FLOPPS of the base student model.
Figure 11:
23

---

Under review as a conference paper at ICLR 2020
MAE-Entropy Output (Test)
0.250
Entropy Output (Test) Individual Student
MAE-Entropy Output (Test)
0.20
MAE- Entropy Output (Test)
0.225
0.18
Individual Student
0.200
0.16
0.14
0.12
0.125
0.100
0.10
5.0
12.5
0.0
2.5
7.5
10.0
12.5
15.0
17.5
FLOPS (x10f)
FLOPS (x10f)
(a)
(b)
MAE- MAE-Entroon Entropy Output (Test)
0.250
MAE- Entropy Output (Test)
MAE- VAE-Entropy Output (Test)
0.20
MAE Entropy Output (Test)
0.225
Individual Student
0.18
Individual Student
0.200
0.16
0.175
0.14
0.12
0.125
0.100
0.10
0.00
0.25
0.50
0.75
1.00 1.25 1.50 1.75
2.00
0.00
0.25
0.50
0.75
1.00
1.25 1.50 1.75 2.00
Num. Parameters (x10f)
Num. Parameters (x10*)
(c)
(d)
Figure 12: Entropy Sioo--Storge-Compuaaaie tradeoff while using CNNs on CIFARI0 with sub- sampling training data to 20,000 samples.
(a) Test mean absolute error for posterior entropy vs (b) Test mean absolute error for posterior entropy FLOPP found using group f1/l2 with pruning. (c) Test mean absolute error for posterior entropy vs storage found using exhaustive search. (d) Test mean absolute error for posterior entropy storage found using group / with pruning. The optimal student model for this configuration obtained with group / pruning. It has approximately 1.6x the number of parameters and 2.8x the FLOPS of the base student model.
FLOPS found using exhaustive search.
24


Reviews:
[{"review_id": "r1xgFEvWqH", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThe 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.\n\nStrengths:\nOverall, the paper is well written and the relationship to previous works is well described. I personally like the Bayesian Dark Knowledge approach, which combines SGLD and knowledge distillation or dark knowledge, and very happy to see its generalization. Unlike the previous work, it is clearly shown that restricting the student architecture to match the teacher can sometimes lead to a significant performance drop, which provides a basis for guiding future developments.\n\nWeaknesses:\n- I think it is a valuable contribution, but my major concern is that the authors only conduct experiments for the classification task, whereas the original Bayesian Dark Knowledge approach also deals with the regression task and shows some interesting results (see Sect. 3.2 and 3.3 in Ref. [1]). I would recommend the authors to extend the experimental evaluation and provide some insight on how to extend the proposed framework to cover the regression task.\n- On page 5, the choice of loss function does not seem to be discussed. I would like the authors to clarify why cross entropy loss is replaced with l(h, h’)=|h-h’| in the classification case.\n- The size of some figures appears too small, for example Fig. 1 and Fig. 2, which may hinder readability.\n\nAt the moment, I recommend a weak reject as the main weakness is the experimental evaluation, but I could be open to increasing my score if my concerns are addressed.\n\nReferences:\n[1] Anoop Korattikara Balan, Vivek Rathod, Kevin P Murphy, and Max Welling. Bayesian dark knowledge. In Advances in Neural Information Processing Systems, pp. 3438–3446, 2015.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "SyeK3VXTYS", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors consider the problem of distilling expectations with respect to Bayesian neural network (BNN) posteriors. These expectations rely on Monte Carlo integration and owing to the large number of BNN parameters can be computationally expensive and memory intensive to compute, motivating the need for distillation.  \n\nI recommend a weak accept for the paper. The authors generalize previous work on distilling posterior predictives by allowing for the computation of posterior expectations beyond posterior predictive distributions, proposing alternate low variance MC estimators, and using an amortization network whose architecture need not be identical to the original BNNs architecture. \n\nWhile the extensions individually are incremental and not particularly exciting, taken together, I believe, they do address a gap in the existing literature. The experiments successfully demonstrate a) when naive distillation fails and b) the proposed extensions help alleviate some of the observed issues. The paper would likely be an useful resource for practitioners in the area. \n\nMinor:\n+ Us vs Uo estimators: It would be interesting to more clearly see what the additional storage (and computation) of Uo is buying us. How much worse are the posterior predictive entropies if Uo is switched with Us? And do the posterior predictive estimates improve if Uo is used inlace of Us? \n\n+ In the paragraph following equation 4, the posterior marginal variance expression implicitly assumes that p(y|x, \\theta) is a Categorical distribution. This should be clarified. The expression doesn’t generally hold, for example if p(y | x, \\theta) is a Gaussian.\n\n+ Figures 1 and 2 are too small and difficult to parse. I would recommend moving some of these to the supplement. \n\n+ It would be good to explicitly point out how much larger is the best (one with the smallest teacher student gap) l1/l2 regularized model compared to the base student model. I realize this is hiding in Figure 2 somewhere, but is not obvious.  ", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}, {"review_id": "HJeLOilLtr", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Contributions:\n\nThe paper considers the distillation of a Bayesian neural network as presented in [Balan et al. 2015]\n\nThe main contribution of the paper is the extension of [Balan et al. 2015] to apply to general posterior expectations instead of being restricted to predictions. \n\nA second contribution of the paper is the finding that restricting the architecture of the student network to coincide with the teacher can lead to suboptimal performance and this can be mitigated by expanding the student's architecture using architecture search.\n\nOriginality/Significance:\n\nI want to discuss the result regarding the generalization of the posterior expectation (section 3.1). To my knowledge this is novel, however, I am failing to see the importance of this result. The paper mentions two cases as motivations: to calculate the entropy and the variance of the marginal. The  problem with these two examples is that it is unclear why these are important and they are not used in the experiments anywhere. Their use should be motivated and the performance of the distillation should be properly evaluated in the experiments.\n\nThe result regarding architecture search is interesting, but it should be expanded and explained more to be a main contribution.\n\nClarity:\n\nThe paper generally understandable, and well written, but it could be better organized. It should expand on the motivation and the architecture search, since these are key components of the paper. The figures are not legible. Even fully zoomed in, they are difficult to read.\n\nOverall assessment:\n\nThe paper has some interesting ideas, but it lacks motivation and significant results.\n\n_______________________________________________________________________\n\nResponse to the rebuttal:\n\nThank you for the detailed reply.\n\n> A primary motivation for generalising posterior expectations is to help quantify model uncertainty. Indeed, the expectation of the predictive entropy is an important quantity that is distinct from the entropy of the posterior predictive distribution. For instance, the difference between these two quantities is exactly the BALD score used in active learning [1]. ...\n\nBALD would be problematic to use with this framework due to computational costs. The main benefit of distillation is the reduced computational cost at inference time. But training itself is still expensive. In BALD, the bulk of the computational cost is fitting the model after each new observation. Distillation does not provide a speedup here.\n\nIf the method indeed works well in an active learning setting, it would be interesting to see experiments showcasing this result.\n\n> - The result regarding architecture search is interesting, but it should be expanded and explained more to be a main contribution.\n\nI was hoping for more discussion/guidance on finding the right architecture or perhaps an algorithm that efficiently optimises the architecture. But I understand that this is more of a future work so I am not holding this against the paper.\n\nI realise that my initial assessment was rather short so I decided to increase my rating and lower my confidence score. I think the paper is borderline, but I am slightly leaning towards rejection due to the insufficient motivation.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I did not assess the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors consider distilling posterior expectations for Bayesian neural networks. While reviewers found the material interesting, and the responses thoughtful, there were questions about the practical utility of the work. Evaluations of classification favour NLL (and typically do not show accuracy), and regression (which was considered in the original Bayesian Dark Knowledge paper) is not considered. In general, it is difficult to assess and interpret how the approach is working, and in what application regime it would be a gold standard, e.g., with respect to downstream tasks. The authors are encouraged to continue with this work, taking reviewer comments into account in a final version.
Decision: Reject