Original: G ENERALIZED B AYESIAN P OSTERIOR E XPECTATION D ISTILLATION FOR D EEP N EURAL N ETWORKS
Revised: G ENERALIZED B AYESIAN P OSTERIOR E XPECTATION D ISTILLATION FOR D EEP N EURAL N ETWORKS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: In this paper, we present a general framework for distilling expectations with respect to the Bayesian posterior distribution of a deep neural network, signiﬁcantly extending prior work on a method known as “Bayesian Dark Knowledge.”
Revised: In this paper, we present a general framework for distilling expectations with respect to the Bayesian posterior distribution of a deep neural network, signiﬁcantly extending prior work on a method known as “Bayesian Dark Knowledge.”

Original: Our generalized framework applies to the case of classiﬁcation models and takes as input the architecture of a “teacher” network, a general posterior expectation of interest, and the architecture of a “student” network.
Revised: Our generalized framework applies to the case of classiﬁcation models and takes as input the architecture of a “teacher” network, a general posterior expectation of interest, and the architecture of a “student” network.

Original: The distillation method performs an online compression of the selected posterior expectation using iteratively generated
Revised: The distillation method performs an online compression of the selected posterior expectation using iteratively generated

Original: Monte Carlo samples from the parameter posterior of the teacher model.
Revised: Monte Carlo samples from the parameter posterior of the teacher model.

Original: We further consider the problem of optimizing the student model architecture with respect to an accuracy-speed-storage trade-off.
Revised: We further consider the problem of optimizing the student model architecture with respect to an accuracy-speed-storage trade-off.

Original: We present experimental results investigating multiple data sets, distillation targets, teacher model architectures, and approaches to searching for student model architectures.
Revised: We present experimental results investigating multiple data sets, distillation targets, teacher model architectures, and approaches to searching for student model architectures.

Original: 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.
Revised: 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.

Original: Lastly, we show that student architecture search methods can identify student models with signiﬁcantly improved performance.
Revised: Lastly, we show that student architecture search methods can identify student models with signiﬁcantly improved performance.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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;
Revised: 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;

Original: Huang et al., 2016; Devlin et al., 2018).
Revised: Huang et al., 2016; Devlin et al., 2018).

Original: However, existing point estimation-based training methods for these models may result in predictive uncertainties that are not well calibrated, including the occurrence of conﬁdent errors.
Revised: However, existing point estimation-based training methods for these models may result in predictive uncertainties that are not well calibrated, including the occurrence of conﬁdent errors.

Original: It is well-known that Bayesian inference can often provide more robust posterior predictive distributions in the classiﬁcation setting compared to the use of point estimation-based training.
Revised: It is well-known that Bayesian inference can often provide more robust posterior predictive distributions in the classiﬁcation setting compared to the use of point estimation-based training.

Original: However, the integrals required to perform Bayesian inference in neural network models are also well-known to be intractable.
Revised: However, the integrals required to perform Bayesian inference in neural network models are also well-known to be intractable.

Original: 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).
Revised: 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).

Original: To help overcome these problems, Balan et al.
Revised: To help overcome these problems, Balan et al.

Original: (2015) introduced an interesting model training method referred to as Bayesian Dark Knowledge .
Revised: (2015) introduced an interesting model training method referred to as Bayesian Dark Knowledge .

Original: In the classiﬁcation setting, Bayesian Dark
Revised: In the classiﬁcation setting, Bayesian Dark

Original: Knowledge attempts to compress the Bayesian posterior predictive distribution induced by the full parameter posterior of a “teacher” network into a “student” network.
Revised: Knowledge attempts to compress the Bayesian posterior predictive distribution induced by the full parameter posterior of a “teacher” network into a “student” network.

Original: The parameter posterior of the teacher network is represented through a Monte Carlo ensemble of speciﬁc 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.
Revised: The parameter posterior of the teacher network is represented through a Monte Carlo ensemble of speciﬁc 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.

Original: 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.
Revised: 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.

Original: As a result, methods of this type have the potential to be much better suited to learning models for deployment in resource constrained settings.
Revised: As a result, methods of this type have the potential to be much better suited to learning models for deployment in resource constrained settings.

Original: In this paper, we present a Bayesian posterior distillation framework that generalizes the Bayesian
Revised: In this paper, we present a Bayesian posterior distillation framework that generalizes the Bayesian

Original: Dark Knowledge approach in several signiﬁcant directions.
Revised: Dark Knowledge approach in several signiﬁcant directions.

Original:  (1) We generalize the target of distillation in the classiﬁcation case from the posterior predictive distribution to general posterior expectations. This allows us to capture targets like the posterior entropy and posterior marginal variance that are of interest from an uncertainty quantiﬁcation perspective.
Revised: The primary modeling and algorithmic contributions of this work are: (1) we generalize the target of distillation in the classiﬁcation 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.

Original: (2) We generalize the approach used to estimate posterior expectations to include options with reduced variance. (3) We generalize the student architecture from being restricted to match the teacher architecture to being a parameter of  the distillationprocedure and examine two basic approaches for student architecture optimization. In terms ofexperimental evaluation, we (1) consider both  posterior predictive distributions and  posterior entropy as distillation targets, (2) consider both fully connected and convolutional teacher model architectures, (3) investigate exhaustive search and sparsity-inducing regularization with pruning as student architecture optimization methods, and (4) present results on both MNIST and CIFAR10, as well as manipulations ofboth data sets that increase posterior uncertainty.
Revised: 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.

Original: Key results include the ﬁnding that  distilling into a student model that matches the architecture of the teacher, as in Balan et al.
Revised: The key empirical ﬁndings are that (1) distilling into a student model that matches the architecture of the teacher, as in Balan et al.

Original: (2015), can be sub-optimal, and that student architecture optimization methods can identify signiﬁcantly improved student models.
Revised: (2015), can be sub-optimal; and (2) student architecture optimization methods can identify signiﬁcantly improved student models.

Revised: We note that the signiﬁcance 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 quantiﬁcation perspective.

Revised: As noted above, we focus on the case of distilling the expected posterior entropy in addition to the posterior predictive distribution itself.

Revised: 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).

Revised: This distinction is extremely important in determining why predictions are uncertain for a given data case.

Revised: 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.

Original: We begin by presenting background material and related work in Section 2.
Revised: 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.

Original: In Section 3, we present the proposed framework and associated Generalized Posterior Expectation Distillation (GPED) algorithm.
Revised: In Section 3, we present the proposed framework and associated Generalized Posterior Expectation Distillation (GPED) algorithm.

Original: In Section 4, we present experiments and results.
Revised: In Section 4, we present experiments and results.

Original: Additional details regarding data sets and experiments can be found in Appendix A, with supplemental results included in Appendix B.
Revised: Additional details regarding data sets and experiments can be found in Appendix A, with supplemental results included in Appendix B.

Original:  B ACKGROUND AND R ELATED W ORK
Revised:  B ACKGROUND AND R ELATED W ORK

Original: In this section we present background material on Bayesian inference for neural networks, and related work on approximate inference, and model compression and pruning.
Revised: In this section we present background material on Bayesian inference for neural networks, and related work on approximate inference, and model compression and pruning.

Original: 2. B AYESIAN N EURAL N ETWORKS
Revised: 2. B AYESIAN N EURAL N ETWORKS

Original: Let p ( y | x , θ ) represent the probability distribution induced by a deep neural network classiﬁer over classes y ∈ Y = { 1 , .., C } given feature vectors x ∈ R D .
Revised: Let p ( y | x , θ ) represent the probability distribution induced by a deep neural network classiﬁer over classes y ∈ Y = { 1 , .., C } given feature vectors x ∈ R D .

Original: The most common way to ﬁt a model of this type given a data set D = { ( x i , y i ) | 1 ≤ i ≤ N } is to use maximum conditional likelihood estimation, or equivalently, cross entropy loss minimization (or their penalized or regularized variants).
Revised: The most common way to ﬁt a model of this type given a data set D = { ( x i , y i ) | 1 ≤ i ≤ N } is to use maximum conditional likelihood estimation, or equivalently, cross entropy loss minimization (or their penalized or regularized variants).

Original: However, when the volume of labeled data is low, there can be multiple advantages to considering a full Bayesian treatment of the model.
Revised: However, when the volume of labeled data is low, there can be multiple advantages to considering a full Bayesian treatment of the model.

Original: Instead of attempting to ﬁnd the single (locally) optimal parameter set θ ∗ according to a given criterion, Bayesian inference uses Bayes rule to deﬁne 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.
Revised: Instead of attempting to ﬁnd the single (locally) optimal parameter set θ ∗ according to a given criterion, Bayesian inference uses Bayes rule to deﬁne 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.

Original: [Equation]
Revised: [Equation]

Original: 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
Revised: 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

Original: Equation 2.
Revised: Equation 2.

Original: 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.
Revised: 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.

Original: 2. A PPROXIMATE I NFERENCE M ETHODS FOR B AYESIAN N EURAL N ETWORKS
Revised: 2. A PPROXIMATE I NFERENCE M ETHODS FOR B AYESIAN N EURAL N ETWORKS

Original: 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).
Revised: 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).

Original: In VI, an auxiliary distribution q φ ( θ ) is deﬁned to approximate the true parameter posterior p ( θ | D , θ 0 ) .
Revised: In VI, an auxiliary distribution q φ ( θ ) is deﬁned to approximate the true parameter posterior p ( θ | D , θ 0 ) .

Original: The variational parameters φ are selected to minimize the
Revised: The variational parameters φ are selected to minimize the

Original: Kullback-Leibler (KL) divergence between q φ ( θ ) and p ( θ | D , θ 0 ) .
Revised: Kullback-Leibler (KL) divergence between q φ ( θ ) and p ( θ | D , θ 0 ) .

Original: Hinton & Van Camp (1993)
Revised: Hinton & Van Camp (1993)

Original: ﬁrst studied applying VI to neural networks.
Revised: ﬁrst studied applying VI to neural networks.

Original: Graves (2011) later presented a method based on stochastic
Revised: Graves (2011) later presented a method based on stochastic

Original: VI with improved scalability.
Revised: VI with improved scalability.

Original: In the closely related family of expectation propagation (EP) methods (Minka, 2001), Soudry et al.
Revised: In the closely related family of expectation propagation (EP) methods (Minka, 2001), Soudry et al.

Original: (2014) present an online EP algorithm for neural networks with the ﬂexibility of representing both continuous and discrete weights.
Revised: (2014) present an online EP algorithm for neural networks with the ﬂexibility of representing both continuous and discrete weights.

Original: Hern´andez-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 ﬁltering (ADF) algorithm that, like VI and EP, generally relies on simpliﬁed posterior densities.
Revised: Hern´andez-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 ﬁltering (ADF) algorithm that, like VI and EP, generally relies on simpliﬁed posterior densities.

Original: The main drawback of VB, EP, and ADF is that they all typically result in biased posterior estimates for complex posterior distributions.
Revised: The main drawback of VB, EP, and ADF is that they all typically result in biased posterior estimates for complex posterior distributions.

Original: 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.
Revised: 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.

Original: MCMC methods allow for drawing a correlated sequence of samples θ t ∼ p ( θ | D , θ 0 ) from the parameter posterior.
Revised: MCMC methods allow for drawing a correlated sequence of samples θ t ∼ p ( θ | D , θ 0 ) from the parameter posterior.

Original: These samples can then be used to approximate the posterior predictive distribution as a Monte Carlo average as shown in Equation 3.
Revised: These samples can then be used to approximate the posterior predictive distribution as a Monte Carlo average as shown in Equation 3.

Original: [Equation]
Revised: [Equation]

Original: Neal (1996) addressed the problem of Bayesian inference in neural networks using Hamiltonian
Revised: Neal (1996) addressed the problem of Bayesian inference in neural networks using Hamiltonian

Original: Monte Carlo (HMC) to provide a set of posterior samples.
Revised: Monte Carlo (HMC) to provide a set of posterior samples.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: In particular, the method uses SGLD to approximate the posterior distribution using a set of posterior parameter samples.
Revised: In particular, the method uses SGLD to approximate the posterior distribution using a set of posterior parameter samples.

Original: These samples can be thought of as an ensemble of neural network models with identical architectures, but different parameter values.
Revised: These samples can be thought of as an ensemble of neural network models with identical architectures, but different parameter values.

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Wang et al.
Revised: Wang et al.

Original: (2018) and Henning et al.
Revised: (2018) and Henning et al.

Original: both propose methods for learning to generate samples that mimic those produced by SGLD.
Revised: both propose methods for learning to generate samples that mimic those produced by SGLD.

Original: 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.
Revised: 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.

Original: This is again a potentially costly operation and is exactly the computation that Bayesian Dark Knowledge addresses.
Revised: This is again a potentially costly operation and is exactly the computation that Bayesian Dark Knowledge addresses.

Original: 2. M ODEL C OMPRESSION AND P RUNING
Revised: 2. M ODEL C OMPRESSION AND P RUNING

Original: 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.
Revised: 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.

Original: In this work, we are particularly concerned with the issue of enabling test-time speed-storage-accuracy trade-offs.
Revised: In this work, we are particularly concerned with the issue of enabling test-time speed-storage-accuracy trade-offs.

Original: The relevant background material includes methods for network compression and pruning.
Revised: The relevant background material includes methods for network compression and pruning.

Original: Previous work has shown that overparameterised deep learning models tend to show much better learnability.
Revised: Previous work has shown that overparameterised deep learning models tend to show much better learnability.

Original: Further, it has also been shown that such overparameterised models rarely use their full capacity and can often be pruned back substatially without signiﬁcant loss of generality (Hassibi et al., 1993; LeCun et al., 1989;
Revised: Further, it has also been shown that such overparameterised models rarely use their full capacity and can often be pruned back substatially without signiﬁcant loss of generality (Hassibi et al., 1993; LeCun et al., 1989;

Original: Luo et al., 2017; Louizos et al., 2017; Frankle & Carbin, 2018;
Revised: Luo et al., 2017; Louizos et al., 2017; Frankle & Carbin, 2018;

Original: Han et al., 2015; Zhang & Ou, 2018).
Revised: Han et al., 2015; Zhang & Ou, 2018).

Original: Hassibi et al.
Revised: Hassibi et al.

Original: ; LeCun et al.
Revised: ; LeCun et al.

Original: (1989) use the second order derivatives of the objective function to guide pruning network connections.
Revised: (1989) use the second order derivatives of the objective function to guide pruning network connections.

Original: More recently, Han et al.
Revised: More recently, Han et al.

Original: (2015) introduced a weight magnitude-based technique for pruning connections in deep neural networks using simple thresholding.
Revised: (2015) introduced a weight magnitude-based technique for pruning connections in deep neural networks using simple thresholding.

Original: Guo et al.
Revised: Guo et al.

Original: ; Jin et al.
Revised: ; Jin et al.

Original: ; Han et al.
Revised: ; Han et al.

Original: (2016) introduce thresholding methods which also support restoration of connections.
Revised: (2016) introduce thresholding methods which also support restoration of connections.

Original: A related line of work includes pruning neurons/channels/ﬁlters instead of individual weights.
Revised: A related line of work includes pruning neurons/channels/ﬁlters instead of individual weights.

Original: Pruning these components explicitly reduces the number of computations by making the networks smaller.
Revised: Pruning these components explicitly reduces the number of computations by making the networks smaller.

Original: Group LASSO-based methods have the advantage of turning the pruning problem into a continuous optimization problem with a sparsity-inducing regularizer.
Revised: Group LASSO-based methods have the advantage of turning the pruning problem into a continuous optimization problem with a sparsity-inducing regularizer.

Original: Zhang & Ou (2018); Alvarez & Salzmann (2016); Wen et al. (2016); He et al.
Revised: Zhang & Ou (2018); Alvarez & Salzmann (2016); Wen et al. (2016); He et al.

Original: (2017) are some examples that use Group LASSO regularization at their core.
Revised: (2017) are some examples that use Group LASSO regularization at their core.

Original: Similarly Louizos et al.
Revised: Similarly Louizos et al.

Original: (2017) use hierarchical priors to prune neurons instead of weights.
Revised: (2017) use hierarchical priors to prune neurons instead of weights.

Original: 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.
Revised: 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.

Original: This makes it easier to realize the resulting computational savings, even on platforms that do not directly support sparse matrix operations.
Revised: This makes it easier to realize the resulting computational savings, even on platforms that do not directly support sparse matrix operations.

Original:  P ROPOSED F RAMEWORK
Revised:  P ROPOSED F RAMEWORK

Original: In this section, we describe our proposed framework for distilling general Bayesian posterior expectations for neural network classiﬁcation models and discuss methods for enabling test-time speed-storage-accuracy trade-offs for ﬂexible deployment of the resulting models.
Revised: In this section, we describe our proposed framework for distilling general Bayesian posterior expectations for neural network classiﬁcation models and discuss methods for enabling test-time speed-storage-accuracy trade-offs for ﬂexible deployment of the resulting models.

Original: 3. G ENERALIZED P OSTERIOR E XPECTATIONS
Revised: 3. G ENERALIZED P OSTERIOR E XPECTATIONS

Original: There are many possible inferences of interest given a Bayesian parameter posterior P ( θ | D , θ 0 ) .
Revised: There are many possible inferences of interest given a Bayesian parameter posterior P ( θ | D , θ 0 ) .

Original: We consider the general case of inferences that take the form of posterior expectations as shown in
Revised: We consider the general case of inferences that take the form of posterior expectations as shown in

Original: Equation 4 where g ( y, x , θ ) is an arbitrary function of y , x and θ .
Revised: Equation 4 where g ( y, x , θ ) is an arbitrary function of y , x and θ .

Original: [Equation]
Revised: [Equation]

Original: Important examples of functions g ( y, x , θ ) include g ( y, x , θ )
Revised: Important examples of functions g ( y, x , θ ) include g ( y, x , θ )

Original: = p ( y | x , θ ) , which results in a posterior expectation yielding the posterior predictive distribution p ( y | x , D , θ 0 ) , as used in Bayesian
Revised: = p ( y | x , θ ) , which results in a posterior expectation yielding the posterior predictive distribution p ( y | x , D , θ 0 ) , as used in Bayesian

Original: Dark Knowledge; g ( y, x , θ ) =
Revised: Dark Knowledge; g ( y, x , θ ) =

Original: (cid:80) Cy (cid:48)
Revised: (cid:80) Cy (cid:48)

Original: =1 p ( y (cid:48)
Revised: =1 p ( y (cid:48)

Original: | x , θ ) log p ( y (cid:48) | x , θ ) , which yields the posterior predictive entropy H ( y | x , D , θ 0 ) 1 ; and g ( y, x , θ ) =
Revised: | x , θ ) log p ( y (cid:48) | x , θ ) , which yields the posterior predictive entropy H ( y | x , D , θ 0 ) 1 ; and g ( y, x , θ ) =

Original: p ( y | x , θ )(1 − p ( y | x , θ )) , which results in the posterior marginal variance σ 2 ( y | x , D , θ 0 ) .
Revised: p ( y | x , θ )(1 − p ( y | x , θ )) , which results in the posterior marginal variance σ 2 ( y | x , D , θ 0 ) .

Original: While the posterior predictive distribution p ( y | x , D , θ 0 ) is certainly the most important posterior inference from a predictive standpoint, the entropy and variance are also important from the perspective of uncertainty quantiﬁcation.
Revised: While the posterior predictive distribution p ( y | x , D , θ 0 ) is certainly the most important posterior inference from a predictive standpoint, the entropy and variance are also important from the perspective of uncertainty quantiﬁcation.

Original: 3. G ENERALIZED P OSTERIOR E XPECTATION D ISTILLATION
Revised: 3. G ENERALIZED P OSTERIOR E XPECTATION D ISTILLATION

Original: Our goal is to learn to approximate posterior expectations E p
Revised: Our goal is to learn to approximate posterior expectations E p

Original: ( θ | D ,θ 0 )
Revised: ( θ | D ,θ 0 )

Original: [ g ( y, x , θ )] under a given teacher model architecture using a given student model architecture.
Revised: [ g ( y, x , θ )] under a given teacher model architecture using a given student model architecture.

Original: The method that we propose takes as input the teacher model p ( y | x , θ ) , the prior p ( θ | θ 0 ) , a labeled data set D , an unlabeled data set
Revised: The method that we propose takes as input the teacher model p ( y | x , θ ) , the prior p ( θ | θ 0 ) , a labeled data set D , an unlabeled data set

Original: D (cid:48) , the function g ( y, x , θ ) , a student model f
Revised: D (cid:48) , the function g ( y, x , θ ) , a student model f

Original: ( y, x | φ ) , an online expectation estimator, and a loss function (cid:96) ( · , · ) that measures the error of the approximation given by the student model f
Revised: ( y, x | φ ) , an online expectation estimator, and a loss function (cid:96) ( · , · ) that measures the error of the approximation given by the student model f

Original: ( y, x | φ ) .
Revised: ( y, x | φ ) .

Original: Similar to Balan et al.
Revised: Similar to Balan et al.

Original: , we propose an online distillation method based on the use of the SGLD sampler.
Revised: , we propose an online distillation method based on the use of the SGLD sampler.

Original: We describe all of the components of the framework in the sections below, and provide a complete description of the resulting method in Algorithm 1.
Revised: We describe all of the components of the framework in the sections below, and provide a complete description of the resulting method in Algorithm 1.

Original: SGLD Sampler: We deﬁne the prior distribution over the parameters p
Revised: SGLD Sampler: We deﬁne the prior distribution over the parameters p

Original: ( θ | θ 0 ) to be a spherical
Revised: ( θ | θ 0 ) to be a spherical

Original: Gaussian distribution centered at µ = 0 with precision τ (we thus have θ 

Revised: Gaussian distribution centered at µ = 0 with precision τ (we thus have θ

Original: =
Revised: =

Original: [ µ, τ ] ).
Revised: [ µ, τ ] ).

Original: We deﬁne S to be a minibatch of size M drawn from D .
Revised: We deﬁne S to be a minibatch of size M drawn from D .

Original: θ t denotes the parameter set sampled for the teacher model at sampling iteration t , while η t denotes the step size for the teacher model at iteration t .
Revised: θ t denotes the parameter set sampled for the teacher model at sampling iteration t , while η t denotes the step size for the teacher model at iteration t .

Original: The
Revised: The

Original: Langevin noise is denoted by z t ∼ N (0 , η t I ) .
Revised: Langevin noise is denoted by z t ∼ N (0 , η t I ) .

Original: The sampling update for SGLD is given by:
Revised: The sampling update for SGLD is given by:

Original: [Equation]
Revised: [Equation]

Original: Distillation Procedure: For the distillation learning procedure, we make use of a secondary unlabeled data set D (cid:48)
Revised: Distillation Procedure: For the distillation learning procedure, we make use of a secondary unlabeled data set D (cid:48)

Original: = { x i | 1 ≤ i ≤ N (cid:48) } .
Revised: = { x i | 1 ≤ i ≤ N (cid:48) } .

Original: This data set could use feature vectors from the primary data set D , or a larger data set.
Revised: This data set could use feature vectors from the primary data set D , or a larger data set.

Original: We note that due to autocorrelation in the sampled teacher model parameters θ t , we may not want to run a distillation update for every Monte Carlo sample drawn.
Revised: We note that due to autocorrelation in the sampled teacher model parameters θ t , we may not want to run a distillation update for every Monte Carlo sample drawn.

Original: We thus use two different iteration indices: t for SGLD iterations and s for distillation iterations.
Revised: We thus use two different iteration indices: t for SGLD iterations and s for distillation iterations.

Original: On every distillation step s , we sample a minibatch S (cid:48) from D (cid:48) of size M (cid:48) .
Revised: On every distillation step s , we sample a minibatch S (cid:48) from D (cid:48) of size M (cid:48) .

Original: For every data case i in
Revised: For every data case i in

Original: S (cid:48) , we update an estimate ˆ g yis of the posterior expectation using the most recent parameter sample θ t , obtaining an updated estimate ˆ g yis +

Revised: S (cid:48) , we update an estimate ˆ g yis of the posterior expectation using the most recent parameter sample θ t , obtaining an updated estimate ˆ g yis +

Original: ≈ E p ( θ | D ,θ 0 )
Revised: ≈ E p ( θ | D ,θ 0 )

Original: [ g ( y, x , θ )]
Revised: [ g ( y, x , θ )]

Original: (we discuss update schemes in the next section).
Revised: (we discuss update schemes in the next section).

Original: Next, we use the minibatch of examples S (cid:48) to update the student model.
Revised: Next, we use the minibatch of examples S (cid:48) to update the student model.

Original: To do so, we take a step in the gradient direction of the regularized empirical risk of the student model as shown below where α s is the student model learning rate, R ( φ ) is the regularizer, and λ is the regularization hyper-parameter.
Revised: To do so, we take a step in the gradient direction of the regularized empirical risk of the student model as shown below where α s is the student model learning rate, R ( φ ) is the regularizer, and λ is the regularization hyper-parameter.

Original: We next discuss the estimation of the expectation targets ˆ g yis .
Revised: We next discuss the estimation of the expectation targets ˆ g yis .

Original: [Equation]
Revised: [Equation]

Original: Expectation Estimation:
Revised: Expectation Estimation:

Original: Given an explicit collection of posterior samples θ 1 , ..., θ s , the standard
Revised: Given an explicit collection of posterior samples θ 1 , ..., θ s , the standard

Original: Monte Carlo estimate of E p ( θ | D ,θ 0 )
Revised: Monte Carlo estimate of E p ( θ | D ,θ 0 )

Original: [ g ( y, x , θ )] is simply ˆ g yis
Revised: [ g ( y, x , θ )] is simply ˆ g yis

Original: = 1 s (cid:80) sj =1 g
Revised: = 1 s (cid:80) sj =1 g

Original: ( y, x i , θ j ) .
Revised: ( y, x i , θ j ) .

Original: However, this estimator requires retaining the sequence of samples θ 1 , ..., θ s , which may not be feasible in terms of storage cost.
Revised: However, this estimator requires retaining the sequence of samples θ 1 , ..., θ s , which may not be feasible in terms of storage cost.

Original: Instead, we consider the application of an online update function.
Revised: Instead, we consider the application of an online update function.

Original: We deﬁne m is to be the count of the number of times data case i has been sampled up to and including distillation iteration s .
Revised: We deﬁne m is to be the count of the number of times data case i has been sampled up to and including distillation iteration s .

Original: An online update function U (ˆ g yis , θ t , m is ) 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 an updated estimate of the expectation ˆ g yis +1 .
Revised: An online update function U (ˆ g yis , θ t , m is ) 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 an updated estimate of the expectation ˆ g yis +1 .

Original: Below, we deﬁne two different versions of the function.
Revised: Below, we deﬁne two different versions of the function.

Original: U s (ˆ g yis , θ t , m is ) , updates ˆ g yis using the current sample only, while U o
Revised: U s (ˆ g yis , θ t , m is ) , updates ˆ g yis using the current sample only, while U o

Original: (ˆ g yis , θ t , m is ) performs an online update equivalent to a full Monte Carlo average.
Revised: (ˆ g yis , θ t , m is ) performs an online update equivalent to a full Monte Carlo average.

Original: [Equation]
Revised: [Equation]

Original: We note that both update functions provide unbiased estimates of E p ( θ | D ,θ 0 )
Revised: We note that both update functions provide unbiased estimates of E p ( θ | D ,θ 0 )

Original: [ g ( y, x , θ )] after a suitable burn-in time B .
Revised: [ g ( y, x , θ )] after a suitable burn-in time B .

Original: The online update U o () will generally result in much lower variance in the estimated values of ˆ g yis , but it comes at the cost of needing to explicitly maintain the expectation estimates ˆ g yis across learning iterations, increasing the storage cost of the algorithm.
Revised: The online update U o () will generally result in much lower variance in the estimated values of ˆ g yis , but it comes at the cost of needing to explicitly maintain the expectation estimates ˆ g yis across learning iterations, increasing the storage cost of the algorithm.

Revised: It is worthwhile noting that the extra storage and computation cost required by U o grows linearly in the size of the training set for the student.

Original: By contrast, the fully stochastic update is memoryless in terms of past expectation estimates, so the estimated expectations ˆ g yis do not need to be retained across iterations resulting in a space savings.
Revised: By contrast, the fully stochastic update is memoryless in terms of past expectation estimates, so the estimated expectations ˆ g yis do not need to be retained across iterations resulting in a space savings.

Original: General Algorithm and Special Cases: We show a complete description of the proposed method in Algorithm 1.
Revised: General Algorithm and Special Cases: We show a complete description of the proposed method in Algorithm 1.

Original: The algorithm takes as input the teacher model p ( y | x , θ ) , the parameters of the prior
Revised: The algorithm takes as input the teacher model p ( y | x , θ ) , the parameters of the prior

Original: P ( θ | θ 0 ) , a labeled data set D , an unlabeled data set D (cid:48) , the function g ( y, x , θ ) , the student model f ( y, x | φ ) , an online expectation estimator U (ˆ g yis , θ t , m is ) , a loss function (cid:96) ( · , · ) that measures the error of the approximation given by f ( y, x | φ ) , a regularization function R () and regularization hyper-parameter λ , minibatch sizes M and M (cid:48) , the thinning interval parameter H , the SGLD burn-in time parameter B and step size schedules for the step sizes η t and α s .
Revised: P ( θ | θ 0 ) , a labeled data set D , an unlabeled data set D (cid:48) , the function g ( y, x , θ ) , the student model f ( y, x | φ ) , an online expectation estimator U (ˆ g yis , θ t , m is ) , a loss function (cid:96) ( · , · ) that measures the error of the approximation given by f ( y, x | φ ) , a regularization function R () and regularization hyper-parameter λ , minibatch sizes M and M (cid:48) , the thinning interval parameter H , the SGLD burn-in time parameter B and step size schedules for the step sizes η t and α s .

Original: We note that the original Bayesian Dark Knowledge method is recoverable as a special case of this framework via the the choices g ( y, x , θ )
Revised: We note that the original Bayesian Dark Knowledge method is recoverable as a special case of this framework via the the choices g ( y, x , θ )

Original: = p ( y | x , θ ) , (cid:96) ( p, q ) = − p log( q ) , U
Revised: = p ( y | x , θ ) , (cid:96) ( p, q ) = − p log( q ) , U

Original: = U s and p ( y | x , θ ) =
Revised: = U s and p ( y | x , θ ) =

Original: f ( y, x , φ ) (e.g., the architecture of the student is selected to match that of the teacher).
Revised: f ( y, x , φ ) (e.g., the architecture of the student is selected to match that of the teacher).

Original: The original approach also uses a distillation data set D (cid:48) obtained from D by adding randomly generated noise to instances from D on each distillation iteration, taking advantage of the fact that the choice U{+, D (cid:48) , p ( y | x , θ ) , θ 0 , g, f, U, (cid:96), R, M, M (cid:48) , H, B, λ, { η t } Tt =1 , { α s } Ss =1 ) 2: Initialize s = 0 , φ 0 , θ 0 , ˆ g yi 0 = 0 , m i 
 = 0 , η3: for t = 0 to T do 4:+}
Revised: The original approach also uses a distillation data set D (cid:48) obtained from D by adding randomly 1: procedure GPED( D {+, D (cid:48) , p ( y | x , θ ) , θ 0 , g, f, U, (cid:96), R, M, M (cid:48) , H, B, λ, { η t } Tt =1 , { α s } Ss =1 ) 2: Initialize s = 0 , φ 0 , θ 0 , ˆ g yi 0 = 0 , m i 
 = 0 , η3: for t = 0 to T do 4:+}

Original: = U s means that no aspect of the algorithm scales with |  D (cid:48)  | .
Revised: Sample S from D with | S | = M 5: θ t +1 ← θ t + η t 2 (cid:0) ∇ θ log p ( θ | θ 0 ) + NM (cid:80) i ∈ S ∇ θ log p ( y i | x i , θ t ) (cid:1) + z t 6: if mod ( t, H ) = 0 and t > B then 7: Sample S (cid:48) from D (cid:48) with | S (cid:48) | =

Revised: M (cid:48) 8: for i ∈ S (cid:48) do 9:

Revised: ˆ g yis +1 ← U (ˆ g yis , θ t , m is ) 10: m is +1 ← m is +11: end for 12: φ s +1 ← φ s + α s (cid:16) N (cid:48)

Revised: M (cid:48) (cid:80) i ∈ S (cid:48) (cid:80)

Revised: y ∈ Y ∇ φ (cid:96) (cid:0)

Revised: ˆ g yis +1 , f

Revised: ( y, x i | φ s ) (cid:1) +

Revised: λ ∇ φ R ( φ s ) (cid:17) 13: s ← s +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

Revised: = U s means that no aspect of the algorithm scales with | D (cid:48) | .

Original: Our general framework allows for other trade-offs, including reducing the variance in the estimates of ˆ g yis at the cost of additional storage in proportion to | D (cid:48) | .
Revised: Our general framework allows for other trade-offs, including reducing the variance in the estimates of ˆ g yis at the cost of additional storage in proportion to | D (cid:48) | .

Original: We also note that note that the loss function (cid:96) ( p, q ) = − p log( q ) and the choice g ( y, x , θ )
Revised: We also note that note that the loss function (cid:96) ( p, q ) = − p log( q ) and the choice g ( y, x , θ )

Original: = p ( y | x , θ ) are somewhat of a special case when used together as even when the full stochastic expectation update U s is used, the resulting distillation parameter gradient is unbiased.
Revised: = p ( y | x , θ ) are somewhat of a special case when used together as even when the full stochastic expectation update U s is used, the resulting distillation parameter gradient is unbiased.

Original: To distill posterior entropy, we set g ( y, x , θ ) =
Revised: To distill posterior entropy, we set g ( y, x , θ ) =

Original: (cid:80)
Revised: (cid:80)

Original: y ∈
Revised: y ∈

Original: Y p ( y | x , θ ) log p ( y | x , θ ) , U =
Revised: Y p ( y | x , θ ) log p ( y | x , θ ) , U =

Original: U o and (cid:96) ( h, h (cid:48) )
Revised: U o and (cid:96) ( h, h (cid:48) )

Original: =
Revised: =

Original: | h − h (cid:48) | .
Revised: | h − h (cid:48) | .

Original: 3. M ODEL C OMPRESSION AND P RUNING
Revised: 3. M ODEL C OMPRESSION AND P RUNING

Original: One of the primary motivations for the original Bayesian Dark Knowledge approach is that it provides an approximate inference framework that results in signiﬁcant computational and storage savings at test time.
Revised: One of the primary motivations for the original Bayesian Dark Knowledge approach is that it provides an approximate inference framework that results in signiﬁcant computational and storage savings at test time.

Original: However, a drawback of the original approach is that the architecture of the student is

Original: 1: procedure GPED( D , D (cid:48) , p ( y | x , θ ) , θ 0 , g, f, U, (cid:96), R, M, M (cid:48) , H, B, λ, { η t } Tt =1 , { α s } Ss =1 ) 2:

Original: [-Initialize s = 0 , φ 0 , θ 0 , ˆ g yi 0 = 0 , m i 
-]
Revised: [-Initialize s = 0 , φ 0 , θ 0 , ˆ g yi 0 = 0 , m i 
-]

Original: = 0 , η3: for t = 0 to T do 4:

Original: Sample S from D with | S | =

Original: M 5: θ t +1 ← θ t + η t 2 (cid:0) ∇ θ log p ( θ | θ 0 )

Original: + NM (cid:80) i ∈ S ∇ θ log p ( y i | x i , θ t ) (cid:1) + z t 6: if mod ( t, H ) = 0 and t > B then 7:

Original: Sample S (cid:48) from D (cid:48) with | S (cid:48) | =

Original: M (cid:48) 8: for i ∈ S (cid:48) do 9:

Original: ˆ g yis +1 ← U (ˆ g yis , θ t , m is ) 10: m is +1 ← m is +11: end for 12: φ s +1 ← φ s + α s (cid:16) N (cid:48)

Original: M (cid:48) (cid:80) i ∈ S (cid:48) (cid:80)

Original: y ∈ Y ∇ φ (cid:96) (cid:0) ˆ g yis +1 , f ( y, x i | φ s ) (cid:1) + λ ∇ φ R ( φ s ) (cid:17) 13: s ← s +14: end if 15: end for 16: end procedure chosen to match that of the teacher.
Revised: However, a drawback of the original approach is that the architecture of the student is chosen to match that of the teacher.

Original: As we will show in Section 4, this will sometimes result in a student network that has too little capacity to represent a particular posterior expectation accurately.
Revised: As we will show in Section 4, this will sometimes result in a student network that has too little capacity to represent a particular posterior expectation accurately.

Original: On the other hand, if we plan to deploy the student model in a low resource compute environment, the teacher architecture may not meet the speciﬁed computational constraints.
Revised: On the other hand, if we plan to deploy the student model in a low resource compute environment, the teacher architecture may not meet the speciﬁed computational constraints.

Original: In either case, we need a general approach for selecting an architecture for the student model.
Revised: In either case, we need a general approach for selecting an architecture for the student model.

Original: To begin to explore this problem, we consider to basic approaches to choosing student model architectures that enable trading off test time inference speed and storage for accuracy.
Revised: To begin to explore this problem, we consider to basic approaches to choosing student model architectures that enable trading off test time inference speed and storage for accuracy.

Original: A helpful aspect of the distillation process relative to a de novo architecture search problem is that the architecture of the teacher model is available as a starting point.
Revised: A helpful aspect of the distillation process relative to a de novo architecture search problem is that the architecture of the teacher model is available as a starting point.

Original: As a ﬁrst approach, we consider wrapping the proposed
Revised: As a ﬁrst approach, we consider wrapping the proposed

Original: GPED algorithm with an explicit search over a set of student models that are “close” to the teacher.
Revised: GPED algorithm with an explicit search over a set of student models that are “close” to the teacher.

Original: Speciﬁcally, we consider a search space obtained by starting from the teacher model and applying a width multiplier to the width of every fully connected layer and a kernel multiplier to the number of kernels in every convolutional layer.
Revised: Speciﬁcally, we consider a search space obtained by starting from the teacher model and applying a width multiplier to the width of every fully connected layer and a kernel multiplier to the number of kernels in every convolutional layer.

Original: While this search requires exponential time in the number of layers, it provides a baseline for evaluating other methods.
Revised: While this search requires exponential time in the number of layers, it provides a baseline for evaluating other methods.

Original: As an alternative approach with better computational complexity, we leverage the regularization function R ( φ ) included in the GPED framework to prune a large initial network using group (cid:96) 1 /(cid:96)regularization
Revised: As an alternative approach with better computational complexity, we leverage the regularization function R ( φ ) included in the GPED framework to prune a large initial network using group (cid:96) 1 /(cid:96)regularization

Original: (Zhang & Ou, 2018; Wen et al., 2016).
Revised: (Zhang & Ou, 2018; Wen et al., 2016).

Original: To apply this approach, we ﬁrst must partition the parameters in the parameter vector φ across K groups G k .
Revised: To apply this approach, we ﬁrst must partition the parameters in the parameter vector φ across K groups G k .

Original: The form of the regularizer is
Revised: The form of the regularizer is

Original: R ( φ ) =
Revised: R ( φ ) =

Original: (cid:80)
Revised: (cid:80)

Original: Kk =1 (cid:0) (cid:80) j ∈ G k
Revised: Kk =1 (cid:0) (cid:80) j ∈ G k

Original: φ 2 j (cid:1) 1 / 2 .
Revised: φ 2 j (cid:1) 1 / 2 .

Original: As is well-established in the literature, this regularizer causes all parameters in a group to go to zero simultaneously when they are not needed in a model.
Revised: As is well-established in the literature, this regularizer causes all parameters in a group to go to zero simultaneously when they are not needed in a model.

Original: To use it for model pruning for a unit in a fully connected layer, we collect all of that unit’s inputs into a group.
Revised: To use it for model pruning for a unit in a fully connected layer, we collect all of that unit’s inputs into a group.

Original: Similarly, we collect all of the incoming weights for a particular channel in a convolution layer together into a group.
Revised: Similarly, we collect all of the incoming weights for a particular channel in a convolution layer together into a group.

Original: If all incoming weights associated with a unit or a channel have magnitude below a small threshold (cid:15) , we can explicitly remove them from the model, obtaining a more compact architecture.
Revised: If all incoming weights associated with a unit or a channel have magnitude below a small threshold (cid:15) , we can explicitly remove them from the model, obtaining a more compact architecture.

Original: We also ﬁne-tune our models after pruning.
Revised: We also ﬁne-tune our models after pruning.

Original: Finally, we note that any number of weight compressing, pruning, and architecture search methods could be combined with the GPED framework.
Revised: Finally, we note that any number of weight compressing, pruning, and architecture search methods could be combined with the GPED framework.

Original: 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.
Revised: 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.

Original:  E XPERIMENTS AND R ESULTS
Revised:  E XPERIMENTS AND R ESULTS

Original: 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.
Revised: 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.

Original: We begin by providing an overview of the experimental protocols used.
Revised: We begin by providing an overview of the experimental protocols used.

Original: [Table]

Original: 4. E XPERIMENTAL P ROTOCOLS
Revised: 4. E XPERIMENTAL P ROTOCOLS

Original: Data Sets: We use the MNIST (LeCun, 1998) and CIFAR10 (Krizhevsky et al., 2009) data sets as base data sets in our experiments.
Revised: Data Sets: We use the MNIST (LeCun, 1998) and CIFAR10 (Krizhevsky et al., 2009) data sets as base data sets in our experiments.

Original: In the case of MNIST, posterior predictive uncertainty is very low, so we introduce two different modiﬁcations to explore the impact of uncertainty on distillation performance.
Revised: In the case of MNIST, posterior predictive uncertainty is very low, so we introduce two different modiﬁcations to explore the impact of uncertainty on distillation performance.

Original: The ﬁrst modiﬁcation is simply to subsample the data.
Revised: The ﬁrst modiﬁcation is simply to subsample the data.

Original: The second modiﬁcation is to introduce occlusions into the data set using randomly positioned square masks of different sizes, resulting in masking rates from 0% to 86 .
Revised: The second modiﬁcation is to introduce occlusions into the data set using randomly positioned square masks of different sizes, resulting in masking rates from 0% to 86 .

Original: For CIFAR10, we only use sub-sampling.
Revised: For CIFAR10, we only use sub-sampling.

Original: Full details for both data sets and the manipulations applied can be found in Appendix A.1.
Revised: Full details for both data sets and the manipulations applied can be found in Appendix A.1.

Original: 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.
Revised: 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.

Original: (2015), a four-layer convolutional network for MNIST, and a four-layer convolutional network for CIFAR10.
Revised: (2015), a four-layer convolutional network for MNIST, and a ﬁve-layer convolutional network for CIFAR10.

Original: Full details of the teacher model architectures are given in Appendix A.2.
Revised: Full details of the teacher model architectures are given in Appendix A.2.

Original: For exhaustive search for student model architectures, we use the teacher model architectures as base models and search over a space of layer width multipliers K 1 and K 2 that can be used to expand sets of layers in the teacher models.
Revised: For exhaustive search for student model architectures, we use the teacher model architectures as base models and search over a space of layer width multipliers K 1 and K 2 that can be used to expand sets of layers in the teacher models.

Original: A full description of the search space of student models can be found in Appendix A.2.
Revised: A full description of the search space of student models can be found in Appendix A.2.

Original: Distillation Procedures: We consider distilling both the posterior predictive distribution and the posterior entropy, as described in the previous section.
Revised: Distillation Procedures: We consider distilling both the posterior predictive distribution and the posterior entropy, as described in the previous section.

Original: For the posterior predictive distribution, we use the stochastic expectation estimator U s while for entropy we used the full online update
Revised: For the posterior predictive distribution, we use the stochastic expectation estimator U s while for entropy we used the full online update

Original: U o .
Revised: U o .

Original: We allow B = 1000 burn-in iterations and total of T = 10 6 training iterations.
Revised: We allow B = 1000 burn-in iterations and total of T = 10 6 training iterations.

Original: The prior hyper-parameters, learning rate schedules and other parameters vary by data set or distillation target and are fully described in Appendix A.2.
Revised: The prior hyper-parameters, learning rate schedules and other parameters vary by data set or distillation target and are fully described in Appendix A.2.

Original: 4. E XPERIMENT 1: D ISTILLING P OSTERIOR E XPECTATIONS
Revised: 4. E XPERIMENT 1: D ISTILLING P OSTERIOR E XPECTATIONS

Original: For this experiment, we use the MNIST and CIFAR10 datasets without any subsampling or masking.
Revised: For this experiment, we use the MNIST and CIFAR10 datasets without any subsampling or masking.

Original: For each dataset and model, we consider separately distilling the posterior predictive distribution and the posterior entropy.
Revised: For each dataset and model, we consider separately distilling the posterior predictive distribution and the posterior entropy.

Original: We ﬁx the architecture of the student to match that of the teacher.
Revised: We ﬁx the architecture of the student to match that of the teacher.

Original: To evaluate the performance while distilling the posterior predictive distribution, we use the negative log-likelihood (NLL) of the model on the test set.
Revised: To evaluate the performance while distilling the posterior predictive distribution, we use the negative log-likelihood (NLL) of the model on the test set.

Original: 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.
Revised: 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.

Original: The results are given in Table 1.
Revised: The results are given in Table 1.

Original: First, we note that the FCNN NLL results on MNIST closely replicate the results in Balan et al.
Revised: First, we note that the FCNN NLL results on MNIST closely replicate the results in Balan et al.

Original: (2015), as expected.
Revised: (2015), as expected.

Original: We also note that the error in the entropy is low for both the FCNN and CNN architectures on MNIST.
Revised: We also note that the error in the entropy is low for both the FCNN and CNN architectures on MNIST.

Original: However, the student model fails to match the NLL of the teacher on CIFAR10 and the entropy MAE is also relatively high.
Revised: However, the student model fails to match the NLL of the teacher on CIFAR10 and the entropy MAE is also relatively high.

Original: In Experiment 2, we will investigate the effect of increasing uncertainty on models applied to both data sets, while in Experiment 3 we will search for student model architectures that improve performance. 
Revised: In Experiment 2, we will investigate the effect of increasing uncertainty on models applied to both data sets, while in Experiment 3 we will search for student model architectures that improve performance. [Figure] [Figure]

Original: 4. E XPERIMENT 2: R OBUSTNESS TO U NCERTAINTY
Revised: 4. E XPERIMENT 2: R OBUSTNESS TO U NCERTAINTY

Original: This experiment builds on Experiment 1 by exploring methods for increasing posterior uncertainty on
Revised: This experiment builds on Experiment 1 by exploring methods for increasing posterior uncertainty on

Original: MNIST (sub-sampling and masking) and CIFAR10 (sub-sampling).
Revised: MNIST (sub-sampling and masking) and CIFAR10 (sub-sampling).

Original: We consider the cross product of four sub-sampling rates and six masking rates for MNIST and three sub-sampling rates for CIFAR10.
Revised: We consider the cross product of four sub-sampling rates and six masking rates for MNIST and three sub-sampling rates for CIFAR10.

Original: We consider the posterior predictive distribution and posterior entropy distillation targets.
Revised: We consider the posterior predictive distribution and posterior entropy distillation targets.

Original: For the posterior predictive distribution we report the negative log likelihood (NLL) of the teacher, and the
Revised: For the posterior predictive distribution we report the negative log likelihood (NLL) of the teacher, and the

Original: NLL gap between the teacher and student.
Revised: NLL gap between the teacher and student.

Original: For entropy, we report the mean absolute error between the teacher ensemble and the student.
Revised: For entropy, we report the mean absolute error between the teacher ensemble and the student.

Original: All metrics are evaluated on held-out test data.
Revised: All metrics are evaluated on held-out test data.

Original: We also restrict the[Figure] experiment to the case where the student architecture matches the teacher architecture, mirroring the
Revised: We also restrict the experiment to the case where the student architecture matches the teacher architecture, mirroring the

Original: Bayesian Dark Knowledge approach.
Revised: Bayesian Dark Knowledge approach.

Original: In Figure 1, we show the results for the convolutional models on MNIST and CIFAR10 respectively.
Revised: In Figure 1, we show the results for the convolutional models on MNIST and CIFAR10 respectively.

Original: The FCNN results are similar to the CNN results on MNIST and are shown in Figure 4 in Appendix B.
Revised: The FCNN results are similar to the CNN results on MNIST and are shown in Figure 4 in Appendix B.

Revised: In Appendix B, we also provide a performance comparison between the U o

Revised: and U s estimators while distilling posterior expectations.

Original: As expected, the the NLL of the teacher decreases as the data set size decreases.
Revised: As expected, the the NLL of the teacher decreases as the data set size decreases.

Original: We observe that changing the number of training samples has a similar effect on NLL gap for both CIFAR and MNIST.
Revised: We observe that changing the number of training samples has a similar effect on NLL gap for both CIFAR and MNIST.

Original: More speciﬁcally, for any ﬁxed masking rate of MNIST (and zero masking rate for
Revised: More speciﬁcally, for any ﬁxed masking rate of MNIST (and zero masking rate for

Original: CIFAR10), we can see that the NLL difference between the student and teacher decreases with increasing training data.
Revised: CIFAR10), we can see that the NLL difference between the student and teacher decreases with increasing training data.

Original: However, for MNIST we can see that the teacher NLL increases much more rapidly as a function of the masking rate.
Revised: However, for MNIST we can see that the teacher NLL increases much more rapidly as a function of the masking rate.

Original: Moreover, the gap between the teacher and student peaks  for moderate values of the masking rate.
Revised: Moreover, the gap between the teacher and student peaks [Figure] The optimal student model for this conﬁguration is obtained with group (cid:96) 2 pruning.

Revised: It has approximately 6 .

Original: {+6 × the number of parameters and 6 . 
 × the FLOPS of the base student model. 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.
Revised: {+6 × the number of parameters and 6 . 
 × the FLOPS of the base student model. 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.

Original: On the other hand, when the masking rate is high, the teacher essentially outputs the uniform distribution for every example, which is very easy for the student to represent.
Revised: On the other hand, when the masking rate is high, the teacher essentially outputs the uniform distribution for every example, which is very easy for the student to represent.

Original: As a result, the moderate values of the masking rate result in the hardest distillation problem and thus the largest performance gap.
Revised: As a result, the moderate values of the masking rate result in the hardest distillation problem and thus the largest performance gap.

Original: For varying masking rates, we see exactly the same trend for the gap in posterior entropy predictions on
Revised: For varying masking rates, we see exactly the same trend for the gap in posterior entropy predictions on

Original: MNIST.
Revised: MNIST.

Original: However, the gap for entropy prediction increases as a function of data set size for CIFAR10.
Revised: However, the gap for entropy prediction increases as a function of data set size for CIFAR10.

Revised: Finally, as we would expect, the performance of distillation using the U o estimator is almost always better than that of the U s estimator (refer Appendix B).

Original: The key ﬁnding 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.
Revised: The key ﬁnding 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.

Original: Indeed, restricting the student architecture to match the teacher can sometimes result in signiﬁcant performance gaps.
Revised: Indeed, restricting the student architecture to match the teacher can sometimes result in signiﬁcant performance gaps.

Original: In the next experiment, we address the problem of searching for improved student model architectures.
Revised: In the next experiment, we address the problem of searching for improved student model architectures.

Original: 4. E XPERIMENT 3: T OWARDS S TUDENT M ODEL A RCHITECTURE S EARCH
Revised: 4. E XPERIMENT 3: T OWARDS S TUDENT M ODEL A RCHITECTURE S EARCH

Original: In this experiment, we compare the exhaustive search to the group (cid:96) 1 /(cid:96) 2 (group lasso) regularizer combined with pruning.
Revised: In this experiment, we compare the exhaustive search to the group (cid:96) 1 /(cid:96) 2 (group lasso) regularizer combined with pruning.

Original: For the pruning 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.
Revised: For the pruning 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.

Original: We present results in terms of performance versus

Original: [Figure]

Original: The optimal student model for this conﬁguration is obtained with group (cid:96)

Original: (cid:96) 2 pruning.

Original: [-It has approximately 6 . 6 × the number of parameters and 6 . 
 × the FLOPS-] of the base student model. computation time (estimated in FLOPS), as well as performance vs storage cost (estimated in number of parameters).
Revised: [-It has approximately 6 . 6 × the number of parameters and 6 . 
 × the FLOPS-]We present results in terms of performance versus computation time (estimated in FLOPS), as well as performance vs storage cost (estimated in number of parameters).

Original: As performance measures for the posterior predictive distribution, we consider accuracy and negative log likelihood.
Revised: As performance measures for the posterior predictive distribution, we consider accuracy and negative log likelihood.

Original: For entropy, we use mean absolute error.
Revised: For entropy, we use mean absolute error.

Original: In all cases results are reported on test data.
Revised: In all cases results are reported on test data.

Original: We consider both fully connected and convolutional models.
Revised: We consider both fully connected and convolutional models.

Original: 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.
Revised: 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.

Original: We select this setting as illustrative of a difﬁcult case for posterior predictive distribution distillation.
Revised: We select this setting as illustrative of a difﬁcult case for posterior predictive distribution distillation.

Original: We plot NLL vs FLOPS and NLL vs storage for all points encountered in each search.
Revised: We plot NLL vs FLOPS and NLL vs storage for all points encountered in each search

Original: The solid blue line indicates the Pareto frontier.

Original: First, we note that the baseline student model (with architecture matching the teacher) from Experiment 2 on MNIST achieves an NLL of 0 .
Revised: First, we note that the baseline student model (with architecture matching the teacher) from Experiment 2 on MNIST achieves an NLL of 0 .

Original: 469 at approximately 0 .
Revised: 469 at approximately 0 .

Original: 48 × 10 6 FLOPs and 0 .
Revised: 48 × 10 6 FLOPs and 0 .

Original: 03 × 10parameters on this conﬁguration of the data set.
Revised: 03 × 10parameters on this conﬁguration of the data set.

Original: We can see that both methods for selecting student architectures provide a highly signiﬁcant improvement over the baseline student architectures.
Revised: We can see that both methods for selecting student architectures provide a highly signiﬁcant improvement over the baseline student architectures.

Original: On
Revised: On

Original: MNIST, the NLL is reduced to 0 . 30 .
Revised: MNIST, the NLL is reduced to 0 . 30 .

Original: Further, we can also see that the group (cid:96) 1 /(cid:96) 2 approach is able to obtain much better NLL at the same computation and storage cost relative to the exhaustive search method.
Revised: Further, we can also see that the group (cid:96) 1 /(cid:96) 2 approach is able to obtain much better NLL at the same computation and storage cost relative to the exhaustive search method.

Original: Lastly, the group (cid:96) 1 /(cid:96) 2 method is 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.
Revised: Lastly, the group (cid:96) 1 /(cid:96) 2 method is 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.

Original: 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.
Revised: 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.

Original: In summary, the key ﬁnding of this experiment is that the capacity of the student model has a signiﬁcant impact on the performance of the distillation procedure, and methods for optimizing the student architecture are needed to achieve a desired speed-storage-accuracy trade-off.
Revised: In summary, the key ﬁnding of this experiment is that the capacity of the student model has a signiﬁcant impact on the performance of the distillation procedure, and methods for optimizing the student architecture are needed to achieve a desired speed-storage-accuracy trade-off.

Original:  C ONCLUSIONS & F UTURE D IRECTIONS
Revised:  C ONCLUSIONS & F UTURE D IRECTIONS

Original: We have presented a framework for distilling expectations with respect to the Bayesian posterior distribution of a deep neural network that generalizes the Bayesian Dark Knowledge approach in several signiﬁcant directions.
Revised: We have presented a framework for distilling expectations with respect to the Bayesian posterior distribution of a deep neural network that generalizes the Bayesian Dark Knowledge approach in several signiﬁcant directions.

Original: 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-accuracy trade-offs.
Revised: 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-accuracy trade-offs.

Original: There are many directions for future work including considering the distillation of a 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-art models
Revised: There are many directions for future work including considering the distillation of a 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-art models.

Original: A D ATASETS AND M ODEL D ETAILS A. D ATASETS
Revised: A D ATASETS AND M ODEL D ETAILS A. D ATASETS

Original: As noted earlier in the paper the original empirical investigation of Bayesian Dark Knowledge for classiﬁcation focused on the MNIST data set (LeCun, 1998).
Revised: As noted earlier in the paper the original empirical investigation of Bayesian Dark Knowledge for classiﬁcation focused on the MNIST data set (LeCun, 1998).

Original: However, the models ﬁt 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.
Revised: However, the models ﬁt 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.

Original: In this section we investigate two orthogonal modiﬁcations of the standard MNIST data set to increase uncertainty: reducing the training set size[Figure] and masking regions of the input images.
Revised: In this section, we investigate two orthogonal modiﬁcations of the standard MNIST data set to increase uncertainty: reducing the training set size and masking regions of the input images.

Original: Our goal is to produce a range of benchmark problems with varying posterior predictive uncertainty.
Revised: Our goal is to produce a range of benchmark problems with varying posterior predictive uncertainty.

Original: We also use the CIFAR10 data set (Krizhevsky et al., 2009) in our experiments and employ the same subsampling technique.
Revised: We also use the CIFAR10 data set (Krizhevsky et al., 2009) in our experiments and employ the same subsampling technique.

Original: MNIST:
Revised: MNIST:

Original: The full MNIST dataset consists of 60,000 training images and 10,000 test images, each of size 28 × 28 , distributed among 10 classes LeCun (1998).
Revised: The full MNIST dataset consists of 60,000 training images and 10,000 test images, each of size 28 × 28 , distributed among 10 classes LeCun (1998).

Original: As a ﬁrst manipulation, we consider sub-sampling the labeled training data to include 10,000, 20,000, 30,000 or all 60,000 data cases in the primary data set D when performing posterior sampling for the teacher model.
Revised: As a ﬁrst manipulation, we consider sub-sampling the labeled training data to include 10,000, 20,000, 30,000 or all 60,000 data cases in the primary data set D when performing posterior sampling for the teacher model.

Original: Importantly, we use all 60,000 unlabeled training cases in the distillation data set D (cid:48) .
Revised: Importantly, we use all 60,000 unlabeled training cases in the distillation data set D (cid:48) .

Original: This allows us 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.
Revised: This allows us 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.

Original: As a second manipulation, we generate images with occlusions by randomly masking out parts of each available training and test image.
Revised: As a second manipulation, we generate images with occlusions by randomly masking out parts of each available training and test image.

Original: For generating such images, we randomly choose a square m × m region (mask) and set the value for pixels in that region to 0.
Revised: For generating such images, we randomly choose a square m × m region (mask) and set the value for pixels in that region to 0.

Original: Thus, the masking rate for a 28 × 28 MNIST image corresponding to the mask of size m × m is given by r = m × m 28 × 28 .
Revised: Thus, the masking rate for a 28 × 28 MNIST image corresponding to the mask of size m × m is given by r = m × m 28 × 28 .

Original: We illustrate original and masked data in Figure 3.
Revised: We illustrate original and masked data in Figure 3.

Original: We consider a range of square masks resulting in masking rates between 0% and 86.2%.
Revised: We consider a range of square masks resulting in masking rates between 0% and 86.2%.

Original: CIFAR10:
Revised: CIFAR10:

Original: The full CIFAR10 dataset consists of 50,000 training images and 10,000 test images, each of size 32 × 32 pixels.
Revised: The full CIFAR10 dataset consists of 50,000 training images and 10,000 test images, each of size 32 × 32 pixels.

Original: We sub-sample the data into a primary training sets D containing 10,000, 20,000, and 50,000 images.
Revised: We sub-sample the data into a primary training sets D containing 10,000, 20,000, and 50,000 images.

Original: As with MNIST, the sub-sampling is limited to training the teacher model only and we utilize all the 50,000 unlabeled training images in the distillation data set D (cid:48) .
Revised: As with MNIST, the sub-sampling is limited to training the teacher model only and we utilize all the 50,000 unlabeled training images in the distillation data set D (cid:48) .

Original: A. M ODELS
Revised: A. M ODELS

Original: To demonstrate the generalizability of our methods to a range of model architectures, we run our experiments with both fully-connected, and convolutional neural networks.
Revised: To demonstrate the generalizability of our methods to a range of model architectures, we run our experiments with both fully-connected, and convolutional neural networks.

Original: 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.
Revised: 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.

Original: Teacher Models: We begin by deﬁning the architectures used for the teacher model as follows:
Revised: Teacher Models: We begin by deﬁning the architectures used for the teacher model as follows:

Original: FCNN (MNIST) :
Revised: FCNN (MNIST) :

Original: We use a 3-layer fully connected neural network.
Revised: We use a 3-layer fully connected neural network.

Original: The architecture used is: Input(784)-FC(400)-FC(400)-FC(output).
Revised: The architecture used is: Input(784)-FC(400)-FC(400)-FC(output).

Original: This matches the architecture used by Balan et al.
Revised: This matches the architecture used by Balan et al.

Original: CNN (MNIST) :
Revised: CNN (MNIST) :

Original: For a CNN, we use two consecutive sets of 2D convolution and maxpooling layers, followed by two fully-connected layers.
Revised: For a CNN, we use two consecutive sets of 2D convolution and maxpooling layers, followed by two fully-connected layers.

Original: The architecture used is: Input(1, (28,28))-Conv(num kernels=10, kernel size=4, stride=1) - MaxPool(kernel size=2) - Conv(num kernels=20, kernel size=4, stride=1) - MaxPool(kernel size=2) - FC (80) - FC (output).
Revised: The architecture used is: Input(1, (28,28))-Conv(num kernels=10, kernel size=4, stride=1) - MaxPool(kernel size=2) - Conv(num kernels=20, kernel size=4, stride=1) - MaxPool(kernel size=2) - FC (80) - FC (output).

Original: CNN (CIFAR10) :
Revised: CNN (CIFAR10) :

Original: Similar to the CNN architecture used for MNIST, we use two consecutive sets of 2D convolution and max-pooling layers followed by fully-connected layers.
Revised: Similar to the CNN architecture used for MNIST, we use two consecutive sets of 2D convolution and max-pooling layers followed by fully-connected layers.

Original: Conv(num kernels=16, kernel size=5) - MaxPool(kernel size=2) - Conv(num kernels=32, kernel size=5) - MaxPool(kernel size=2) - FC(200) - FC (50) - FC (output).
Revised: Conv(num kernels=16, kernel size=5) - MaxPool(kernel size=2) - Conv(num kernels=32, kernel size=5) - MaxPool(kernel size=2) - FC(200) - FC (50) - FC (output).

Original: In the architectures mentioned above, the “output” size will change depending on the expectation that we’re distilling.
Revised: In the architectures mentioned above, the “output” size will change depending on the expectation that we’re distilling.

Original: For classiﬁcation, the output size will be 10 for both datasets, while for the case of entropy, it will be 1 .
Revised: For classiﬁcation, the output size will be 10 for both datasets, while for the case of entropy, it will be 1 .

Original: We use ReLU non-linearities everywhere between the hidden layers.
Revised: We use ReLU non-linearities everywhere between the hidden layers.

Original: For the ﬁnal output layer, softmax is used for classiﬁcation.
Revised: For the ﬁnal output layer, softmax is used for classiﬁcation.

Original: In the case of entropy, we use an exponential activiation to ensure positivity.
Revised: In the case of entropy, we use an exponential activiation to ensure positivity.

Original: Student Models: The student models used in our experiments use the above mentioned architectures as the base architecture.
Revised: Student Models: The student models used in our experiments use the above mentioned architectures as the base architecture.

Original: For explicitly searching the space of the student models, we use a set of width multipliers starting from the teacher architecture.
Revised: For explicitly searching the space of the student models, we use a set of width multipliers starting from the teacher architecture.

Original: The space of student architectures corresponding to each teacher model deﬁned earlier is given below.
Revised: The space of student architectures corresponding to each teacher model deﬁned earlier is given below.

Original: The width multiplier values of K 1 and K 2 are determined differently for each of the experiments, and thus will be mentioned in later sections.
Revised: The width multiplier values of K 1 and K 2 are determined differently for each of the experiments, and thus will be mentioned in later sections.

Original: FCNN (MNIST) :
Revised: FCNN (MNIST) :

Original: Input(784)-FC( 400 · K 1 )-FC( 400 · K 2 )-FC(output).
Revised: Input(784)-FC( 400 · K 1 )-FC( 400 · K 2 )-FC(output).

Original: CNN (MNIST) : Input(1, (28,28))-Conv(num kernels= (cid:98) 10 · K 1 (cid:99) , kernel size=4, stride=1) - MaxPool(kernel size=2) - Conv(num kernels= (cid:98) 20 · K 1 (cid:99) , kernel size=4, stride=1) - MaxPool(kernel size=2) - FC ( (cid:98) 80 · K 2 (cid:99) ) - FC (output).
Revised: CNN (MNIST) : Input(1, (28,28))-Conv(num kernels= (cid:98) 10 · K 1 (cid:99) , kernel size=4, stride=1) - MaxPool(kernel size=2) - Conv(num kernels= (cid:98) 20 · K 1 (cid:99) , kernel size=4, stride=1) - MaxPool(kernel size=2) - FC ( (cid:98) 80 · K 2 (cid:99) ) - FC (output).

Original: CNN (CIFAR10) : Input(3, (32,32))-Conv(num kernels= (cid:98) 16 · K 1 (cid:99) , kernel size=5) - MaxPool(kernel size=2) - Conv(num kernels= (cid:98) 16 · K 1 (cid:99) , kernel size=5) - MaxPool(kernel size=2) - FC ( (cid:98) 200 · K 2 (cid:99) ) - FC ( (cid:98) 
Revised: CNN (CIFAR10) : Input(3, (32,32))-Conv(num kernels= (cid:98) 16 · K 1 (cid:99) , kernel size=5) - MaxPool(kernel size=2) - Conv(num kernels= (cid:98) 16 · K 1 (cid:99) , kernel size=5) - MaxPool(kernel size=2) - FC ( (cid:98) 200 · K 2 (cid:99) ) - FC ( (cid:98) 50 · K 2 (cid:99) ) -

Original: 50 · K 2 (cid:99) ) - FC (output).
Revised: FC (output).

Original: Model and Distillation Hyper-Parameters : We run the distillation procedure using the following hyperparameters: ﬁxed teacher learning rate η t = 4 × 10 − 6 for models on MNIST and η t = 2 × 10 −for models on CIFAR10, teacher prior precision τ = 10 , initial student learning rate
Revised: Model and Distillation Hyper-Parameters : We run the distillation procedure using the following hyperparameters: ﬁxed teacher learning rate η t = 4 × 10 − 6 for models on MNIST and η t = 2 × 10 −for models on CIFAR10, teacher prior precision τ = 10 , initial student learning rate

Original: α s = 10 − 3 , student dropout rate p = 0 .
Revised: α s = 10 − 3 , student dropout rate p = 0 .

Original: 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 6 .
Revised: 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 6 .

Original: For training the student model, we use the Adam algorithm (instead of plain steepest descent as indicated in Algorithm 1) and set a 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.
Revised: For training the student model, we use the Adam algorithm (instead of plain steepest descent as indicated in Algorithm 1) and set a 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.

Original: Also, note that we only apply the regularization function R ( φ s ) while doing Group (cid:96) 1 /(cid:96)
Revised: Also, note that we only apply the regularization function R ( φ s ) while doing Group (cid:96) 1 /(cid:96)

Original: 2 pruning.
Revised: 2 pruning.

Original: Otherwise, we use dropout as indicated before.
Revised: Otherwise, we use dropout as indicated before.

Original: Hyper-parameters for Group (cid:96) 1 /(cid:96)
Revised: Hyper-parameters for Group (cid:96) 1 /(cid:96)

Original: 2 pruning experiments :
Revised: 2 pruning experiments :

Original: For experiments involving group (cid:96) 1 /(cid:96)regularizer, the regularization strength values λ are chosen from a log-scale ranging from 10 − 8 to 10 − 3 .
Revised: For experiments involving group (cid:96) 1 /(cid:96)regularizer, the regularization strength values λ are chosen from a log-scale ranging from 10 − 8 to 10 − 3 .

Revised: [Figure]

Original: When using Group (cid:96) 1 /(cid:96) 2 regularizer, we do not use dropout for the student model.
Revised: When using Group (cid:96) 1 /(cid:96) 2 regularizer, we do not use dropout for the student model.

Original: The number of ﬁne-tuning epochs for models on MNIST and CIFAR100 are 600 and 800 respectively.
Revised: The number of ﬁne-tuning epochs for models on MNIST and CIFAR100 are 600 and 800 respectively.

Original: At the start of ﬁne-tuning, we also reinitialize the student learning rate α t = 10 − 4 for fully-connected models and α t = 10 − 3 for convolutional models.
Revised: At the start of ﬁne-tuning, we also reinitialize the student learning rate α t = 10 − 4 for fully-connected models and α t = 10 − 3 for convolutional models.

Original: The magnitude threshold for pruning is (cid:15) =
Revised: The magnitude threshold for pruning is (cid:15) =

Original: 10 − 3 .
Revised: 10 − 3 .

Original: B S UPPLEMENTAL E XPERIMENTS AND R ESULTS
Revised: B S UPPLEMENTAL E XPERIMENTS AND R ESULTS

Original: Supplemental Results for Experiment 2: Robustness to Uncertainty In Figure 4, we demonstrate the results of Experiment 2 (Section 4.3), on fully-connected networks for MNIST. 
Revised: Supplemental Results for Experiment 2: Robustness to Uncertainty In Figure 4, we demonstrate the results of Experiment 2 (Section 4.3), on fully-connected networks for MNIST. Additionally, in

Revised: Tables [2-4], we provide a performance comparison between U o

Revised: and U s estimators while distilling posterior expectations for all model-data set combinations.

Revised: We follow the same experimental conﬁgurations as in Experiment 2.

Original: Supplemental Results for Experiment 3:
Revised: Supplemental Results for Experiment 3:

Original: Towards Student Model Architecture Search
Revised: Towards Student Model Architecture Search

Original: The additional results from running Experiment 3 (Section 4.4) on different combinations of model type, dataset, and performance metrics have been given in Figures[5 - 12]. 
Revised: The additional results from running Experiment 3 (Section 4.4) on different combinations of model type, dataset, and performance metrics have been given in Figures[5 - 12]. [Table] [Table]

Original: [Figure]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: 2 with pruning.
Revised: 2 with pruning.

Original: The optimal student model for this conﬁguration is obtained with group (cid:96) 1 /(cid:96)
Revised: The optimal student model for this conﬁguration is obtained with group (cid:96) 1 /(cid:96)

Original: 2 pruning.
Revised: 2 pruning.

Original: It has approximately 6 .
Revised: It has approximately 6 .

Original: 6 × the number of parameters and 6 . 

Revised: 6 × the number of parameters and 6 .

Original: × the FLOPS of the base student model.
Revised: × the FLOPS of the base student model.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: of the base student model.
Revised: of the base student model.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: It has approximately 4 .
Revised: It has approximately 4 .

Original: 2 × the number of parameters and 4 .
Revised: 2 × the number of parameters and 4 .

Original: 2 × the FLOPS of the base student model.
Revised: 2 × the FLOPS of the base student model.

Original: [Figure]
Revised: [Figure]

Original: conﬁguration is obtained with group (cid:96) 1 /(cid:96)
Revised: conﬁguration is obtained with group (cid:96) 1 /(cid:96)

Original: 2 pruning.
Revised: 2 pruning.

Original: It has approximately 4 .
Revised: It has approximately 4 .

Original: 7 × the number of parameters and 5 .
Revised: 7 × the number of parameters and 5 .

Original: 2 × the FLOPS of the base student model.
Revised: 2 × the FLOPS of the base student model.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]