Accelerating Monte Carlo Bayesian Inference via Approximating Predictive Uncertainty over the Simplex
Authors: Yufei Cui
Wuguannan Yao
Qiao Li
Antoni Chan
Chun Jason Xue
Venue:  (2020)

Abstract: Estimating the predictive uncertainty of a Bayesian learning model is critical in various decision-making problems, e.g., reinforcement learning, detecting adversarial attack, self-driving car. As the model posterior is almost always intractable, most efforts were made on finding an accurate approximation the true posterior. Even though a decent estimation of the model posterior is obtained, another approximation is required to compute the predictive distribution over the desired output. A common accurate solution is to use Monte Carlo (MC) integration. However, it needs to maintain a large number of samples, evaluate the model repeatedly and average multiple model outputs. In many real-world cases, this is computationally prohibitive. In this work, assuming that the exact posterior or a decent approximation is obtained, we propose a generic framework to approximate the output probability distribution induced by model posterior with a parameterized model and in an amortized fashion. The aim is to approximate the true uncertainty of a specific Bayesian model, meanwhile alleviating the heavy workload of MC integration at testing time. The proposed method is universally applicable to Bayesian classification models that allow for posterior sampling. Theoretically, we show that the idea of amortization incurs no additional costs on approximation performance. Empirical results validate the strong practical performance of our approach.

Under review as a conference paper at ICLR 2020
ACCELERATING MONTE CARLO BAYESIAN PREDIC- TION VIA APPROXIMATING PREDICTIVE UNCER- TAINTY OVER THE SIMPLEX
Anonymous authors Paper under double-blind review
ABSTRACT
Estimating the predictive uncertainty Bayesian learning model is critical var- ious periisionmmaking problems, e.g., reinforcement learning, detecting adversarial attack, self-driving car. As the model posterior is almost always intractable, most efforts were made on finding an accurate approximation the true posterior. Even though decent estimation of the model posterior is obtained, another approxima- tion is required to compute the predictive distribution over the desired output. A common accurate solution is to use Monte Carlo (MC) integration. However, it needs maintain large number of samples, evaluate the model repeatedly and average multiple model outputs. In many real-world cases, this is computationally prohibitive. In this work, assuming that the exact posterior or a decent approxi- mation is obtained, we propose a generic framework to approximate the output probability distribution induced by model posterior with parameterized model and in amortized fashion. The aim is to approximate the true uncertainty of specific Bayesian model, meanwhile alleviating the heavy workload of MC inte- gration at testing time. The proposed method is universally applicable to Bayesian classification models that allow for posterior sampling. Theoretically, we show that the idea of amortization incurs additional costs approximation performance. Empirical results validate the strong practical performance of our approach.
INTRODUCTION
Bayesian inference is a principled method to estimate the uncertainty of probabilistic models. In most applications, especially in deep learning, the likelihood model and model prior not conjugate hence marginalizing over model prior or posterior cannot be performed analytically, which hinders the practical applicability. For tractability, simple point estimate such maximum a posteriori (MAP) estimate could be used to approximate the full model posterior. The price paid is the loss of model uncertainty due to incomplete characterization of the model posterior. Approximate inference methods, such as Markov chain Monte Carlo and variational inference, enhance the approximate posterior by a better probability distribution while keeping inference tractability. However, even though decent approximation of posterior can be obtained, computation of predictive distribution is usually intractable due to loss of conjugacy, and is of high cost if tractable.
To introduce the problem, {yn, >= where x,, E X and yn E y are the rith input and output, respectively. Let the model posterior p(0[D) be approximated by MC estimate > 5(0 - 0.), and the predictive distribution (a categorical distribution parameterized by the predicted class probabilities) thus approximated by
consider Bayesian classification model trained on dataset D
S
p D) =
f plyix () do ~
1
Σ Epyvix,..). 2=1
S
The predictive distribution can be accurately estimated as S 00. However, to perform the computation, we need to maintain large number samples, repeatedly evaluate the model for times and finally average the model outputs. This problem is critical in many real-world cases. For example, assisted-driving car system requires an accurate measure uncertainty to avoid making mistakes with high confidence (Kendall & Gall 2017;; Sünderhauf et al. 2018). Due the limited computational resources and storage in such system, it's hard to maintain a large number of samples and perform S times evaluation of the Bayes model for the real-time image data.
00.

---

Under review as a conference paper at ICLR 2020
In this work, aiming at boosting the prediction speed while maintaining rich characterization of the prediction, we propose to approximate the distribution of class probabilities over the simplex induced by the model posterior p(0 D), in an amortized fashion. This naturally diverts the heavy-load MC integration process from testing period to approximation period. Different from the previous work in Bayesian knowledge distillation (Balan et al., 2013 Bull et al. (0T6) that only focuses on the output categorical distribution (a point on simplex), the induced distribution over simplex provides: 1) rich knowledge includes prediction confidence for identifying out-of-domain (OOD) data (see empirical examples in Fig. 2 ; 2) the possibility to use more expressive distributions as the approximate model.
We term the Bayes classifier as "Bayes teacher" and the approximate distribution as "student", due to the analogy with teacher-student learning. A Dirichlet distribution is used as the student due to its expressiveness, conjugacy to categorical distribution and its efficient reparameterization for training. We propose to explicitly disentangle the parameters of the student into a prediction model (PM) and concentration model (CM), which capture class probability and sharpness of Dirichlet respectively. The CM output can directly be used as a measure for detecting OOD data. We term our approximation method One-Pass Uncertainty (OPU) as it simplifies real-world evaluation of Bayesian models by computing the predictive distribution with only one model evaluation. Note that, OPU allows choosing various types of student model (e.g., compressed neural network (Lin et al. 2017 Hubara eet al. 2016 Han et al. 20155) for further speedup on specific platforms with no extra design efforts.
As the amortized approximation of induced distributions is unexplored in the literature, we consider and compare several choices of probability distance measure: forward KL, earth- distance (EMD) and maximum mean discrepancy (MMD). We theoretically analyze the performance gap incurred by the amortized approximation and show that, under MMD, besides model loss due to restriction of student distribution, the amortized approximation does not introduce additional loss.
Empirical evaluations show significant speedup (~~ 500x) of Bayes models. The results on Bayes NN show that OPU performs better in misclassification detection and OOD detection than state-of- the-art works in Bayesian knowledge distillation. It can also be observed that explicit disentang taninnnginnng of mean and concentration helps improve the performance. The comparisons of different probability measures validate the theoretical analysis. We also conduct empirical evaluations and comparisons on Bayes logistic regression and Gaussian process, to show OPU universally applicable to all Bayesian classification models.
2
ONE-PASS UNCERTAINTY FRAMEWORK
2.1
INDUCED DISTRIBUTION OVER SIMPLEX
In this section present our OPU framework for a generic Bayesian parametric classifier, e.g- Bayesian logistic regression (BLR) or Bayesian neural networks (NN). Let the Bayesian classifier be specified with categorical likelihood and a parametric function, i.e.,
p(y|x, 0) Cat(yIT(x;;)) 0))
and prior distribution p(0) be specified over the parameter space θ, where T a parametric function from input space to simplex, e.g., a neural network with softmax output layer, and K is the number of classes. In this paper, we assume the posterior p(0 D) is obtained and focus on the computation of the predictive distribution. In what follows, let ((( be approximate or exact (if available) model posterior, from which samples could be obtained.
- SK is
From dual perspective, with fixed input x, the mapping Tx T(x; is also understood as data-dependent mapping Tx e > Sk-1, which could be used to transform the posterior p(0 DD) to conditional distribution (on x) over the simplex. That is. for each we can define random variable π = Tx(0) whose distribution is induced by 0 ~ (D) and Tx. This is effectively a well- defined purw-forward measure #p0 D) over the simplex (see Fig. 11. We term the conditional distribution of π xx D as an "induced distribution" " and define px (T) = p(T] x, D) to keep the notation uncluttered.
The induced distribution isolates the dependence between input and output and simplifies the compu- tation of predictive distribution due to the change-of- formula. Specifically, given a test point x", the predictive distribution can be alternatively written
p(y|x", (D) = / Cat(y] (()))( do
/ -SX-1 Cat(y) (T (TT) dr.
(2)
2

---

Under review as a conference paper at ICLR 2020
2
1
Teacher
Softmax
Margin.
Categorical ply x. D)
X
Tx
(8)
(b)
2
0
p
(D)
Simplex
Approximation
Student
Categorical Class probabilitt
Prediction
(c)
X
g(T x
Scalar g
Concentration
Figure 2: An empirical example of MC estimate of the induced distribution over simplex. The classi- fier (MCDP) is trained on real-world digits images of {0, 1,2] (corresponding to left, right, top cor- of simplex). The 1st and 3rd indicate the input while the 2nd and 4th rows indicate MC estimate of π over a 2-simplex.
Confidence
An intuitive example of the proposed framework. Only one imput is consider this example. "Margin." " indicates marginalization over π.
Figure I:
Our key insight is that Px contains all information we need for prediction and uncertainty measurement. Hence π is sufficient the sense that, given æ, y is independent of both 0 and x. The isolation combines the complexities in both the likelihood and posterior into single object px, and keeps simple dependence structure between y and æ. Also, the isolation renders the last probabilistic layer y Ir nuisance, which can thus be pecled off for prediction. To validate the idea, can show that the predictive distribution is simply p(y) x, D) = Cat(y E,, ).. The difference between probabilistic structures of original Bayesian model and the isolated version is showed in Fig. Appendix.
Although the density function Px is sometimes hard to compute, especially for complicated Tx (like NN), its samples can be obtained via 1) sampling the posterior. i.e. 0₂ p(0) D); 2) "push-forward" the samples by Tx. i.e. n. Tx (02). )- The behavior of π can be empirically observed via the particles is able to distinguish different types of input based the behavior of its samples. In Fig. 2 a). the inputs are similar to the training data. The particles gather around corner of the simplex, indicating a confident prediction. In Fig. 2b) the inputs are digit images but are relatively hard to predict. Therefore, the particles gathers around the center, indicating the model is certain that the input is on decision boundary. In Fig.2jo (e)). the inputs are images outside the domain of training data, the particles spread over the simplex, which means the model has a high uncertainty about input, indicating out-of-domain (OOD). In Fig. dd) the particles spread along a line between two corners, indicating the model is confident that the result is not at the other corner.
(T, ] 1-1 As shown in Fig.
However, using particles for inference is time-consuming as each particle requires one evaluation of the model. This motivates us to use tractable conditional distribution q g( (x) to approximate Px-
2.2 AMORTIZED APPROXIMATION
The view of Ps enables flexible choices of q(< x), as any distribution defined R-K can be transformed to SK-1 via logistic transformation (Aitchison 1982). However, modeling g( locally for every input not practical, as the design efforts and number parameters grows linearly with the number of data points. Therefore, we propose to approximate Px in two aspects: I) use a single family of distribution 9: to generalize to unseen examples, let the parameter of q. ax axx: Φ), be a function depending on x and parameterized by set of global adaptive parameters Φb, and thus gx = q x) = q(<)| (xx). The computational cost is amortized by casting the problem of learning series of conditional distributions to regression problem.
We term Px as "teacher distribution" and (x as "student". representation in Fig. Appendix, by proper approximation, the stochasticity and knowledge in node of the teacher is transferred into node π of the student model, such that the full predictive uncertainty is maintained. In terms of computation, the approximation requires sampling only in the training stage. While in the testing stage, to obtain the predictive distribution, only one evaluation is required. Thus, we denote the framework as the One-Pass Uncertainty (OPU) model.
our method, as seen in graphical
The above benefits do not introduce any concession on generalizability. Since OPU is based on approximating the distribution of the output class probabilities which is common for all classifiers,
3

---

Under review as a conference paper at ICLR 2020
the amortized approximation can be applied to any Bayesian classifier. Note that the approximation framework can be extended to non pon-maaaettric model like Gaussian process, where the computational cost of inference is high |
In this work, we choose the student model to be a Dirichlet distribution, ( = Dir(sr/as)) where Ox is function mapping input
Dirichlet parameter. The reasons for choosing Dirichlet is the tractability: the Dirichlet is the conjugate prior to the Categorical, and thus enables tractable integration of 2 given the parameters. To better disentangle the uncertainty measures, we use the design Oxx hx exx, where hx h(x; Φi) and O g(;; Φ2)) are two neural networks, and the vector output sums to 1. Vector output h determines the mean of the Dirichlet (i.e., the predicted class probabilities). and determines the concentration of the Dirichlet (i.e., the prediction confidence). To see this, the posterior of the class labels is the Dirichlet mean, p(ye = 1|x, 3) = Cat(ye = 1 Dir()i((()))) dr E..,. ax.c lxxt where ye and 0x,, are the ε-th coordinate of y and ax respectively. Therefore, we call h as the "prediction model" (PM). Similarly, the precision parameter ao (determines sharpness) of the Dirichlet solely depends on 9x. ag = Ee axxcc = -x Therefore, we term Ox as the "concentration model" (CM). Based on this property, whether the Dirichlet flat or not, can be fully characterized by CM. It can be expected that, when approximating particles in Fig. a and (b), the output value of CM is high, as the samples are concentrated, which means high confidence. CM outputs low value for particles in Fig. [[c), yielding a flat distribution and low confidence.
lxxceex e0x,
2.3 LEARNING
With probability distance or divergence, we define the approximation loss L(x, ax) =((pe; (x)
The student model is trained to minimize the aggregated objective
min aeF
E pxtx ((x,
where F is some hypothesis space and p(x) is some distribution over X. In practice we take p(x) PD (x)( with D' = [xmm }MM a dataset containing features only. As the amortized approximation of induced distributions in Bayesian classifiers is unexplored in the literature, we consider and compare several choices of including KL divergence, earth mover's distance (EMD) and maximum mean discrepancy (MMD). The corresponding derivation of the training objectives and training algorithms are in the appendix.
Forward KL divergence Minimizing aggregated reverse KL divergence is not tractable in our scenario as the density function px is not available in general. This difficulty is avoided by using forward KL, in which the intractable density function is only involved in the irrelevant entropy term. It is equivalent to using cross-entropy as a local loss, i.c. L(x, a) = -Ep In By plugging in particle estimation Px = E 5(T (x,)))) the training objective becomes min. -Ep (x) 5 Σ. In (x((( which is equivalent to an "amortized" MLE problem with particles providing the estimation of sufficient statistics of (x- Due to the zero-avoiding nature, forward KI. tends over-estimate the support of Px- This leads under-confidence approximation ("flat" approximate distribution) and hence might deteriorate the quality of uncertainty measurements. This is expected be more serious when Px multi-modal.
EMD It known that EMD provides much weaker topology than other probability distance measures Peyré et al. (2019). In the application where data is supported on strictly lower-dimensional manifolds, EMD provides more stable gradient than KL divergence (Arjovsky et al., 2017 Tolstikhin
et al. 2017). An example of particles on fow-dimensional manifold is shown in Fig. 2
Specifically, the KR dual representation of EMD Viillani 2008) in our problem is given by E Ex (T)) T)], where || |l. denotes the Lips-
R T (())) chitz semi-norm and v is known as the critic (or the discriminator (Arjovsky et al. 2017 As the
Wi xx) = suPpWlass
induced distribution is conditioned on x, a local critic vix should be defined for each x. Following Arjovsky et al. (2017), intractable supremum is solved by parameterizing vix (h( wx). To avoid training [D]] local critics, we propose to let u be the global weight, and let ta depends on x, i.e., x/x= v/( (T; w. x) (see Fig. Appendix). The final aggregated training objective becomes, min max
E (xpr(x) [E)(o P) IV (Tx); x)] - Egg (TT;; w. x]]] + AR(w),
(4)
'The details of extracting samples from various Bayes parametric models and non-parametric model (GP classifier) are summarized in the appendix.
4

---

Under review as a conference paper at ICLR 2020
where is the introduced global parameter for th. (w) is the imposed gradient penalty (Gulrajani iet al., (211)) over w to enforce the Lipschitz constraint. Solving the minimax problem requires the supremum to be attained for each under the Lipschitz constraint. Specifically, in every optimization step, v is trained to generate a critic for each x that matches the exact EMD. Practically, this needs a high-capacity critic and the required capacity increases with the number of classes K.
MMD Let H4 be a reproducing kernel Hilbert space (RKHS) defined by positive-definite kernel k, the MMD between y>x and (x can be written as
MMDa(P.d) =
sup ME 11 [x - E,g [+)(
(5)
Compared with EMD, the advantage of MMD is that there is no need to train an NN as the critic that maximizes Eq. 5 With kernel trick, MMD can be readily estimated in closed-form with its empirical version under finite sample (Sec. Cp.
Compared with KL divergence, MMD is valid statistical metric. Due to the symmetry property, the approximation is expected to be neither mean-seeking nor mode-seeking, Therefore, MMD is not expected to have an under-confidence issue.
Reparameterization Note that optimization under both EMD and MMD requires gradient of the expectation of critic via sampling from q. which contains parameters. obtain efficient gradient estimator and reduce variance, we use the reparameterization trick (specifically, implicit reparameter- ization trick). For details, see Sec. Ci Appendix.
2.4 AMORTIZATION GAP
To better understand nature of the proposed approximation, we consider the "unamortized" version of the approximation as an intermediate stage, which involves fitting separate approximations to each Px. To demonstrate the idea, we leverage to MMD due to its nice property. For fixed px. the optimal point-wise approximation within family Q is defined as
arg min MMDx Kx
(6)
4EQ
Then we have the following lemma:
Lemma 1. Let PI) be the space of probability measures over the simplex. equipped with MMD metric defined by a universal kernel. If satisfies Assumption [ then the map x + Px is continuous. Further, if C P() is closed convex model space, the projection Tx is unique and the map X Tx is also continuous.
Further, if we assume the model space is parameterized and identifiable in MMD, i.c. Q= {g(TT(a)
E A} and MMD+(q(T) a)), (a)))') 0 if and only if || ax| || 0, we may obtain continuity in parameter space. The continuity of optimal parameters implics there exists a continuous function a" x E X' + ax (x) E A, which serves as the essential target of amortized goal.
To analyze how amortization affects the approximation, we define the local amortization gap as
(xx MMD (gx, Px) MMD Px).
Then it holds that
0 Axx) MMDx Rx),
where the lower bound is because 77 is the projection and the upper bound is due triangle inequality. Then our goal, minimizing aggregated MMD loss minimizing aggregated amortization gap E)(x) A(x), up to an irrelevant additive constant.
(4x, Px), is essentially equivalent to
One can see that ax is the global minimizer of aggregated amortization gap and hence the global minimizer of the aggregated MMD loss. Intuitively, J is of enough capacity, then the infimum (over F) of aggregated amortization gap could be O, i.c. we can push the amortization gap arbitrarily small. If the optimal solution is covered by hypothesis space, additional cost is introduced by amortizing the approximation. the global optimum is reached, OPU approximation exactly matches the point-wise minimizer, i.e. MMD (gx, (x) = O, due to uniqueness of projection.
other words, model loss due to using restrictive Q will dominate. Further,
5

---

Under review as a conference paper at ICLR 2020
3 RELATED WORK
In this section, related works are reviewed and compared with OPU in terms of methodology. In PoppactApprox (Snelson & Ghahramani, 2005), a parametric model composed of small subset of "best samples" selected from the original MC samples used to approximate the full predictive Categorical distribution. Extending the approximation to Bayes NN, BDK (Balan et al. 2015) proposes to use NN to approximate the predictive Categorical distribution of a Bayes NN trained by stochastic gradient Langevin dynamics (SGLD). Specifically, the teacher network generates samples via SGLD and KL between the two distributions is minimized in an online fashion. However, the disadvantage is that data uncertainty, model uncertainty and distributional uncertainty are all entangled in the class probabilities because categorical distributions are of limited expressiveness.
Different from these previous works that only approximate the class probabilities, OPU approximates the induced distribution of class probabilities, which contains richer information including both class probabilities and prediction confidence (e.g., the three types of uncertainty observable via the samples in Fig. 2 We choose a Dirichlet distribution as the student model, and explicitly disentangle the mean and concentration to fully capture the thee types of uncertainty. We also explore other probability distance measures (EMD and MMD), showing that KL yields degenerate prediction performance.
Using a Dirichlet to estimate uncertainty has also been explored by Deep Prior Network (DPN) (Malinin & Gales) 2018), where parameterized Dirichlet is used in Bayesian model to characterize the "distributional uncertainty", i.e., to tell if the data in the training domain or not. However, DPN adds stochastic layer in the Bayes model, rather than approximating a well-trained Bayes teacher. Due to the intractable inference, DPN uses an MAP estimate of the model posterior which incurs a loss of uncertainty. To compensate for the lost characterization of uncertainty, DPN uses hand-crafted training goal that explicitly requires OOD examples (which are typically unavailable in real-world applications). In contrast to DPN, our OPU is able to: 1) extract predictive uncertainty from any Bayesian classification model according to the practical requirements; 2) choose various types for student model (e-g.: quantized neural network) to enable fast prediction; 3) use only in- domain data in training to get a good uncertainty measure (see Sec. 4. Note that none of these properties can be achieved by DPN.
4 EXPERIMENTS
4.1 EXPERIMENTAL SETUP
Models and Tasks In this section, present the experimental results using Bayesian NN (BNN) For each model, we choose a few Bayesian methods as teachers, and approximate them with our OPU, We also compare with state-of-the-art approximations that are proposed for the specific types of methods. For BNN, we use MCDP (Gal & Ghahramani 2016) and SGLD (Welling & Teh ) as the teacher, and BDK and DPN for comparisons. The methods for obtaining posterior samples from the Bayes teachers are in Appendix B. For each type of model, in-domain misclassification (MisC) detection, off-oin (OOD) input detection, prediction performance and prediction time are presented.
Baselines and Uncertainty Measures For the uncertainty measures of teachers, we take MCDP and KL as a example. For each testing data x, a Dirichlet is fit on particles under KL, to get an optimal 2x whose differential entropy (D) is an uncertainty measure. For the uncertainty in prediction, to avoid the model loss, entropy (E) and maximum probability (P) of the averaged particle are directly adopted as uncertainty measures. This gives MCDP-KL model as baseline. The other baselines MCDP-EMD and MCDP-MMD are also obtained similarly. Note that they share the same E and P as the sample mean is the same. The of oz is expected to be the best uncertainty measure that OPU can approach theoretically (when the amortization loss is zero, see Sec. 2.4). For fairness, the same set of posterior particles is used for obtaining D. The basclines for SGLD are obtained in a similar way. Students use categorical entropy (for BDK) or D (for DPN) and P of the output distributions. For OPU, we consider E and P of the prediction model, and the scalar output of the concentration model (C) as the measures.
Data and Evaluation Metrics For a fair comparison, we let D' = D in the approximation of OPU. The in do-aoniin dataset is split to training data and testing data, i.e., Din = D' = (Dir DDe), ), which is used for training models, evaluating prediction and MisC detection. The OOD dataset Dood and Dre are used for OOD detection. To assess the performance, we use accuracy, time, Area under the ROC
"hhe setup and results BLR and GP are available in the appendix.
6

---

Under review as a conference paper at ICLR 2020
Table 1: Results BNN- MNIST. The doubt quote means "same above".
MisC detection
Omniglot
SEMEION
Acc.
Test
Model
AUROC
AUPR
AUROC
AUPR
AUROC
AUPR
(%)
time(s)
MCDP-KL
97.3 (E)
43.0 (B)
99.4 (D)
99.7 (D)
86.8(E)
53.8 (P)
97.9
210.6
MCDP-EMD
97.3 (E)
43.0 (E)
99.6 (D)
99.9 (D)
86.8(E)
53.8 (P)
97.9
MCDP-MMD
97.3 (E)
43.0 (E)
99.7 (D)
(D)
90.1 (D)
71.2 (D)
97.9
OPU-MCDP-KL
94.2 (E)
37.7 (B)
100 (C)
(C) (C)
(C)
96.2
0.443
OPU-MCDP-EMII
95.3 (P)
43.7 (P)
100 (C)
100 (C)
(C) (C)
96,1
OPU-MCDP-NM -MD
97.2 (P)
40.9
100 (C)
100 (C)
99.8 (C) 98.6
97.9
SGLD-KL KL
(P)
46.2 (E)
99.2 (E)
99.6 (E)
89.3 (E)
46)8
98.4
233.5
SGLD-EMD
(E)
46.2
99.4 (D)
99.7 (D)
89.9 (D)
47.1 (D)
98.4
SGLD-MME
97.9 (E)
46.2 (E)
99.2 (E)
99.6 (E)
89.3 (E)
46.8 (E)
98.4
OPU-SGID-KI.
(F)
46.7 (R)
100 (C)
100 (C)
99.5 (C)
98.4 (C)
98.2
0.443
OPU-SGLD-EMII
93.7 (P)
44.4 (E)
100 (C)
100 (C)
98.9 (C)
96.2 (C)
98.0
OPU-SGLD-MMDI
97.2 (P)
44.6 (E)
100 (C)
100 (C)
99.1 (C)
(C.
98.1
BDK-SGLD
85.9
46.6 (E)
46.1 (E)
41.7 (B)
35.3
(P)
92.1
0.441
BDK-MCDP
86.9 (E)
41.1 (E)
47.5 (P)
44.1 (P)
43.3 (P)
47.2 (P)
92.4
ADK-DIR-SGLJ
89.9
40.0 (E)
95.4 (E)
96.4 (E)
74.7
38.3 (E)
94.1
DPN
99.0 (E)
43.6 (E)
100 (E)
100 (E)
99.7 (E)
98.6 (E)
99.4
Table 2: Results on BNN Balanced EMNIST.
MisC detection
Omniglot
Acc.
Model
AUROC
AUPR
AUROC
AUPR
(%)
MCDP-KL
89.7 (P)
46.8 (P)
99.7 (E)
99.7 (E)
88.8
MCDP-MMD
89.7 (P)
46.8 (P)
99.9 (D)
99.9 (D)
88.8
OPU-MCDP-KL
84.8 (P)
40.6 (P)
96.2 (E) /67.5 (C)
96.5 (E) /63.7 (C)
87.9
OPU-MCDP-MMD
89.8 (P)
49.6 (P)
100.0 (C)
100 (C)
88.4
(AUROC) and PR (AUPR), following the baseline in (Hendrycks & Gimpel 2017) (The detailed calculation process is shown appendix D.3j. Time evaluated on the whole De To save space, we only present the best performing uncertainty measure (E, P or C) for each task and method. (Full results on undistilled MCDP/SGLD are shown in appendix (D33) We use the MXNet implementation of BDK and GPflow implementation of GP, and the remaining models are implemented in Pytorch. All experiments run on desktop with an i7-8700 CPU and an RTX-2080 Ti GPU. The experiments for Bayesian logistic regression follows the same setup Gaussian process.
4.2 BAYESIAN NEURAL NETWORK
The experiments for Bayesian neural network use MNIST and balanced EMNIST datasets as Din
and use Omniglot and SEMEION dataset as Dood as in Malinin & Gales 2018)) Given MCDP and SGLD as teacher models. the baselines are obtained as illustrated in Sec. 4.1 The rest models are: OPU approximating MCDP (OPU-MDDP), OPU-SGLD, BDK-SGLD, BDK-Dir-SGLD and DPN. For BDK-Dir-SGLDD, we replace the Categorical distribution in BDK by Dirichlet without disentangling the mean and concentration, then train it with the same MC ensemble as OPU-SGID. This is to show the benefits of explicit disentanglement.
The NN architecture used by these models is an MLP with size 784-400-4000 10, ReLU activations. and softmax outputs, following Balan et al. (2015). For CM. we use an MLP with size 784-400-400-
MCDP trained by SGD with hyper-parameters: dropout-rate of 0.5, learning rate x 10-1, mini-batch size of 256, number of iterations 10 The parameters of critic are shown in Appendix. For MMD, we use a summation of RBF kernel and polynomial kernel. OPU-MCDP is trained by Adam with syper-paramedt training of SGLD and BDK follows Balan et al. (2015). Then OPU-SGLD is trained with the same follows
number of iterations 100, learning rate for student 10-3 The
hyperparameters as OPU-MCDP. Results of DPN are from Malinin & Gales (2018). The results presented in Table [
Computation time. OPU offers 5500 speedup compared to the original MCDP/SGLD, as OPU only evaluate the model twice (PM and CM in the student network) while MCDP/SGLD evaluates for S times. This confirms our idea of accelerating Bayesian prediction by diverting the sampling process from the test period to the approximation period. Note that the time cost of MCDP/SGLD increases with more posterior samples involved. BDK is slightly faster than OPU because it runs one network while OPU runs both PM and CM. Experimental results on original MCDP/SGLD with fewer samples are given in appendix D.3 to show performance degradation with fewer samples.
7

---

Under review as a conference paper at ICLR 2020
OPU BDK. In some tasks especially OOD detection, the measure of concentration outperforms the baseline. This is because the explicit disentanglement of mean and concentration helps "targeted" knowledge distillation, forward KI.), we observe that OPU-SGID--K is significantly better in OOD detection tasks and AUROC in MisC detection. BDK shows slight advantage in AUPR in MisC detection task. This is because the knowledge distillation only happens between two categorical variables in BDK, which only helps capturing prediction information. In contrast, OPU framework first extracts all information in BNN with the induced distribution, then transfers the knowledge to more expressive distribution with a small loss guaranteed (Sec. (Sec.2.4). Adding a Dirichlet distribution to BDK (BDK-DIL DIR-SGLD) helps improving the performance OOD detection. However, on SEMEION, which is expected to be harder as it is more similar to MNIST, there is a large performance difference from OPU. This further validates the necessity of explicit disentanglement of mean and concentration.
shown in Sec. 2.2 By comparing OPU-SGLD-KR -LLLLLLLLLL and BDK (trained by
OPU DPN. Our OPU model (without OOD data in training) has comparable performance to DPN (which uses a hand-crafted goal and OOD data in training). Another reason that DPN performs slightly better is that DPN uses VGG-6 (4 Convolutional layer and FC layer), which much stronger model than the 2-layer MLP model that other models use.
KL vs EMD VS MMD. With MCDP, MCDP-MMD gives the best performance of differential entropy baseline (D). This is because the samples of MCDP are relatively spread out over the simplex and might be multi- modal.. The probability distance is fully captured by MMD under such case. EMD is expected to perform well as there are lot of samples π. residing on a low-dimensional manifold. However, the performance seems to degenerated due to the limited capacity of the hyper-network and the difficulty to train the minimax problem. KI presents the worst performance as expected because it is likely to be under confident with samples of MCDP. With SGLD, the performance of KI. is the best except for AUROC of MisC detection. This is because the samples T, of SGLD over the simplex are much denser and are typically unimodal.
KL vs MMD (EMNIST). The experiments are conducted on EMNIST whose number of classes is large. We choose Omniglot as the OOD detection dataset as it is also contains handwritten characters, which are harder than SEMEION for model trained on EMNIST. A CNN with structure similar with LeNet is used for MCDP and OPU (20 and 50 output channels in two convolutional layers). The baseline approach achieves classification accuracy of 88.8%. The samples of MCDP Bayes NN are expected to be even more dispersive and multi-modal than MCDP trained with MNIST. OPU trained with EMD failed to converge, possibly because the capacity of the hypernetwork was not enough. Therefore, we do not recommend to use EMD for amortized approximation of predictive uncertainty. unless more scalable estimator of EMD can be provided. As shown in Table the performance gap between OPU-MCDP- KI. and the baseline is larger because the multi-modality is severe - the entropy of sample mean a better measure for OOD detection than the concentration. This further validates that OPU trained by KL suffers from an cover the support of all samples, making the student distribution more dispersive. This inaccurate estimate of concentration affects the estimation of prediction results (mean) in turn, thus the accuracy lower and MisC detection performance degenerated. By contrast, MMD consistently provides an approximation that has similar performance with the teacher, that echoes the analysis in Sec. 2.4 Amortization Gap. The details of experiments on the approximation error/gap can be found in appendix D33 It is shown that the amortization gap of using OPU is small, indicating the function space is enough to cover the essential target. It also shows the approximation error is mainly from the model error of using Dirichlet distribution, which is small as well, see appendix D33 for details. 5 DISCUSSION
issue. Specifically, KL forces OPU to
The idea of "transferring" the randomness from model posterior to simple-structure distribution at the output can be generalized other problems where real-time evaluation of uncertainty is critical, e.g., object segmentation. This allows interesting designs of structured output distributions.
The choice of student distribution for OPU is flexible. Although we verify that Dirichlet suitable for the tested classification tasks, there are still chances that Dirichlet incurs great model error. Note that, in our framework, using a Dirichlet for the student modeling choice, similar to assuming Gaussian posteriors for variational approximations, The OPU framework is general and any student distribution can be adopted, e.g., generalized Dirichlet, mixture of Dirichlets. To do this requires: 1) a suitable parametrization of the model that can capture uncertainty: 2) deriving/computing the approximation loss (KL, MMD, EMD); 3) reparameterization trick of expectations for efficient gradient estimation. The algorithms of KI., EMD and MMD as well the analysis still apply.
8

---

Under review as a conference paper at ICLR 2020
REFERENCES
John Aitchison. The statistical analysis of compositional data. Journal of the Royal Statistical Society: Series (Methodological), 44(2):139-161) 1982.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein GAN. arXiv preprint
arXiv:1701.07877 077770770777. 2017.
Anoop Korattikara Balan, Vivek Rathod, Kevin Murphy, and Max Welling. Bayesian dark knowledge. In Advances in Neural Information Processing Systems, pp. 3438-3446 2015.
Samuel Rota Bulò, Lorenzo Porzi, and Peter Kontschieder. Dropout distillation. In International Conference Machine Learning, pp. 99-107, 2016.
Mikhail Figurnov, Shakir Mohamed, and Andriy Mnih. Implicit reparameterization gradients. In Advances in Neural Information Processing Systems, pp. 441-452, 2018.
Yarin Gal and Zoubin Ghahramani. Dropout a Bayesian approximation: Representing model uncertainty in deep learning. In International conference on machine learning, pp. 1050-1059, 2016.
Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville. Improved training of Wasserstein GANs. In Advances in Neural Information Processing Systems, pp. 5767-5777,
2017. 2017.
Song Han, Huizi Mao, and William Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. arXiv preprint arXiv:1510.000 2015.
Dan Hendrycks and Kevin Gimpel. A baseline for detecting misclassified and out-of-distribution ex- amples in neural networks. Proceedings of International Conference on Learning Representations, 2017.
James Hensman, Alexander G Matthews, Maurizio Filippone, and Zoubin Ghahramani. MCMC for variationally sparse Gaussian processes. In Advances in Neural Information Processing Systems, pp. 1648-1656, 2015a.
James Hensman, Alexander G de G Matthews, and Zoubin Ghahramani. Scalable variational gaussian process classification. In Proceedings of the Eighteenth International Conference on Artificial Intelligence and Statistics, 2015b.
Itay Hubara, Matthicu Courbariaux, Daniel Soudry, Ran EI- Yaniv, and Yoshua Bengio. Binarized neural networks. In Advances in neural information processing systems, pp- 4107-4115, 2016.
Alex Kendall and Yarin Gal. What uncertainties we need bayesian deep learning for computer vision? In Advances in neural information processing systems, pp- 5574-5584, 2017.
Xiaofan Lin, Cong Zhao, and Wei Pan. Towards accurate binary convolutional neural network. In Advances in Neural Information Processing Systems, pp. 345-353, 2017.
Andrey Malinin and Mark Gales. Predictive uncertainty estimation via prior networks. In Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Bianchi and R. Garnett (eds.), Advances in Neural Information Processing Systems 31, pp. 7047-7058. Curran Associates, Inc., 2018
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport. Foundations and Trends() in Machine Learning, 11(5-6)3355-60) 2019.
Nicholas Polson, James G Scott, and Jesse Windle. Bayesian inference for logistic models using Pólya-Gamma latent variables. Journal of the American statistical Association, 108(504): 1339-1349, 2013.
Edward Snelson and Zoubin Ghahramani. Compact approximations to Bayesian predictive distri- butions. Proceedings of the 22nd international conference on Machine learning, pp. 240-847. ACM, 2005.

---

Under review as a conference paper at ICLR 2020
Niko Sünderhauf, Oliver Brock, Walter Scheirer, Raia Hadsell, Dieter Fox, Jürgen Leitner, Ben Uperoft. Pieter Abbeel, Wolfram Burgard, Michael Milford, et al. The limits and potentials of deep learning for robotics. The International Journal of Roboties Research, 37(4-5):405 420, 2018.
Ilya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto-encoders. arXiv preprint arXiv: 2017.
Cédric Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
Max Welling and Yee Whye Teh. Bayesian learning via stochastic gradient Langevin dynamics. In Proceedings of the 28nd international conference on Machine learning, pp. 681-688. ACM, 2011.
A APPENDIX: PROOF
Assumption 1. T y be a map between finite dimensional vector spaces. We say TT satisfies Assumption for distribution p (:;0) Lipschitz and the Lipschitz constant Le satisfies Eg < +00.
Proof. We simply show both maps are Lipschitz continuous with MMD metric P(S Let x, E X and Hy such that |,lll, I 1. For x + Px. the "push-forwar" definition of px leads
MMD (x)) = sup Bp(pT)((Tx()) (xx
lill11
< sup
(x(0))))(TT( - (Txx
|
sup
|| d4 (0) (x)
lells1
K CBp(ODD (0)
& LEp -
where dę = is a kernel-based distance. The constant C emits when bounding de with Fuclidean norm and the existence of such constant due to topology equivalence in fiite-dimenniioal space. For x 2x we notice that 0% is the projection of Px onto Q (effectively the projection of kernel mean embeddings in H4)) By projection theorem in Hilbert space, the projection map P> + 9x is non-expansive, i.e.
MMD Rx) & MMD (Px,
which leads to Lipschitz property of x + C
B APPENDIX: PLOTS.
The graph representations of original view, isolated view of Bayes teacher and the student shown in Fig.
C APPENDIX: ALGORITHMS.
For presenting the algorithms, we slightly change the notation. We let a(x) a(x, Φ), h(x) h(x, Φ1), g(x) g(x Φ2), where Φ = (Φ1, ( Φi and Φ2 are the parameters of prediction model and concentration model g. respectively.
KL. With the training objective
min
E p (x) 5
Σ In (Tx(xxx))xx
10

---

Under review as a conference paper at ICLR 2020
N
0
π
y
y
(a)
(b)
(c)
Figure 3: Graphical representation of probabilistic structure the Bayes teacher (left and middle) and the student (right). Dashed edges denote deterministic dependence, box nodes are deterministic and circle nodes are stochastic. The left and middle graphs correspond to the LHS and RHS of D. respectively.
The student model is updated by doing $ = of - Vo, (x)]axx;))))(<';; (1))) and Φ : Φb2 --' Ee, Vez In (x(""; 0.) fa)); ;2)))) alternately, where 7 is the learning rate at iteration and x" is the input at this iteration.
EMD. With the following training objective,
min max
Epp (x) [Ep(ord) (Tx w, x)] - E (,; + XR(w),
Algorithm 1: OPU Training Algorithm with EMD
Input: Posterior samples: [0,)5-1- OPU training data DD; Gradient penalty coefficient A; Number of training iterations: Tstu and Twit-
while & not converge do
Sample (6) pD (x)
x* Update Approximation
for iter in .Ttu do
Sample (T 1o ((T(x(()) Φ))
((1) Ds v ((TT;;,
Φ1 Adam( V VΦ, (())
Sample (T } ~ ((T((())
( (122))
Es' (iT w. x()
Φ2 Adam $ 2())
end
/* Update Critic
for iter in 1... Twit do
Sample ST 1->-1) ()T
Compute R(v)
(1)(v)
ES S
m=1 (T 0,) x(( Es- (TTT;,,x))) + XR(w)
v + Adam (O)
end
end
MMD. The kernel mean embedding of Px and Ix are given by Up = Epg [A((TT .)] and Hg E.g (k(x,,)) )) The training objective is then,
min 1|1p Hyll
(10)
11

---

Under review as a conference paper at ICLR 2020
Crilic #i [5,
π
Concatenate
Feat
Feat
NNs
NN layers
Scalar
Output
Figure 4: A brief example of the critic in EMD.
Algorithm 2: OPU Training Algorithm with MMD
Input:
Posterior samples: (0.)a-1: OPU training data D'; Gradient penalty coefficient A.
while & converge do Sample
x()) PD (x) Sample
{Tg, 15'
='=1 ((T(x(())
Φ), get (T.. [T((x00)))==
2 ("
Σ 5
Empr mm Tm + 11 Em k p. ppm))
Σ m,m== n= (TTg.m. p.
S'
57(5
Φ1 Adam Vo, C(0)
Sample {T. 15' ~ (TT(((()) Φ)
(6) -1(--1) Emeen k(T Tgmm) + 11 Empin k(Tppm ppm)
S' 2 Em mm,== k(TTg.m Tp..)
('
Φ - Adam V (())
end
Reparameterization:: To obtain efficient gradient estimator and reduce variance, we reparame- terize the Dirichlet by an equivalent product of K independent Gamma distributions. If क PG (x) = I Gam(@x)oxs, then π = (Σx मх)
Dir(T (x). By Thm. 3 in (Arjovsky al. 2017 in each L(dlxx, the supremum attained at wt E L and the gradient is V [( [m xx) [ (T)]. Then as noted by (Figurnov et al. 2018), the gradient VoL(ol ())x can be implicitly computed without knowing the inverse of standardization function (e.g., CDF). Specifically, by Eq. in (Figurnov et al. 2018). V [ ) = EpG(#)x:0) < V where the
first term V V (T) is computed via the chain rule and the second term V is obtained by solving
local diagonal linear system. Refer to (Figurnov al. 2018 and references therein for details.
D APPENDIX: EXPERIMENTS
For the critic in EMD used in MNIST experiments, NNI is 784-256 MLP. NN2 a 10-256 MLP and NN3 is 512-256-1 MLP (see Fig. D. For the critic in EMD used in MNIST experiments, NN1 the same as convolutional layers used in Ba, NN2 is a 10-256 MLP and NN3 a 512-256-1 MLP (see Fig. 4).
D. BAYESIAN LOGISTIC REGRESSION
We test two models in this experiment: Polya Gamma (PG), CompactApprox approximating PG (CA-PG) and OPU approximating PG (OPU-PG). We draw 500 posterior samples from PG and train OPU with the following hyperparameters: number of epochs 100. learning rate for student. CA-PG is trained by first drawing 5000 samples from PG then evaluating the model with 50 randomly selected samples from them (same setup CoppactApprox). The random selection is repeated for 10 times
12

---

Under review as a conference paper at ICLR 2020
Table 3: Results Bayesian logistic regression models.
MisC detection
OOD detection
Acc.
Time
Data
Model
AUROC
AUPR
AUROC
AUPR
(%)
(sec.)
PG
60.0 (Eat)
24.2 (Ent)
87.0 (Ent)
76.7 (Ent)
64,4
1.01 1 0.001
Pima
CA-PG
58.2 (Ent)
24.2 (Ent)
80.1 (Ent)
74.5 (Ent)
62.3
0.18 + 0.001
OPU-PG
59.7 (Ent)
25.6 (Ent)
100.0 (CM)
100.0 (CM)
64.4
0.01=000022
PG
83.9 (Ent)
24.3 (Ent)
54.6 (Ent)
53.5 (Ent)
92.4
Spam
CA-PG
64.1 (Ent)
24.2 (Ent)
71.5 (Ent)
67.5 (Ent)
85,4
0.36 1 0.001
OPU-PG
83.9 (Ent)
23.8 (Ent)
99.7 (CM)
99.3 (CM)
92.4
0.01=00002
Table 4: Results on Gaussian process classification models.
MisC detection
OOD detection
Ace.
Time
Data
Model
AUROC
AUPR
AUROC
AUPR
(%)
(second)
SGPMC
65.3 (F)
(E)
96.7 (F)
94.9 (F)
79.3
0.003
Pima
SVGP
64.3 (E)
43.2 (E)
96.0 (E)
91.1 (F)
77.1
0.004
OPU-SGPMC
65.7 (E)
(E)
100.0 (C)
100.0 (C)
79.2
0.010
SGPMC
86.7 (E)
(E)
97.6 98)
92.4
0.056
Spam
SVGP
86.2 (E)
33.3 (E)
99.2 (E)
98)5
92.1
0.032
OPU-SGPMC
86.5 (E)
39,5 (E)
100.0 (C)
100.0 (C)
92.0
0.011
and we pick the best group of samples. As the results for the three metrics are similar, we only show the results trained with KL divergence.
The results are shown in Table 3 OPU-PG maintains similar performance with the original PG on prediction accuracy and MisC detection. Meanwhile, OPU outperforms CA-PG on prediction accuracy, MisC and OOD detection. For OPU, CM outperforms other uncertainty measures at OOD detection, which indicates it captures the distributional uncertainty well. OPU performs better than the PG Bayes teacher. The reason might be that parametric model is Iearned to approximate the ensemble of discrete samples, which could produce a smoother output distribution (regularization), leading to better performance. OPU also achieves ~ 100-600x speedup from the original PG.
D.2 GAUSSIAN PROCESS
For GP, we use SGPMC (Hensman et al. 2015a) as the teacher, and SVGP (Hensman et al. 2015b
for comparison.
This experiment uses Pima and Spambase datasets as Din, Pima a medical dataset with 769 data points and 9 dimensions. Spambase is text dataset with 4601 data points and 57 dimensions for identifying spam email. We generate the same number data points from a zero-mean multivariate Gaussian distribution for Dood For each dataset, 10% of data points uniformly selected into the testing set D We normalize the data by features with L.2 norm. where K is the number of classes. There are models tested: SGPMC, OPU approximating SGPMC (OPU-SGPMC) and SVGP. SGPMC and SVGP are trained with 1/ [Dim data points randomly selected from Din as inducing points. Then 500 samples over functions of Diu are generated from SGPMC. As the results for the three metrics are similar, we only show the results trained with KL divergence.
The results are presented in Table 4 On MisC detection and prediction accuracy, OPU has similar performance to SGPMC, which indicates the effectiveness of approximating prediction results with ensemble of samples in the nonparametric family. With the same number of inducing points, SVGP performs slightly worse than SGPMC, because it incurs a two-fold approximation. Measuring uncertainty with CM in OPU outperforms other measures and models, which indicates the sharpness of the nogima distribution can be captured via the CM the designed Dirichlet.
SGPMC and SVGP are faster than OPU on Pima, but are slower than OPU on the larger Spambase. This is due to the static latency for setting up the GPU for OPU. which becomes the main time cost when the dataset is small (as in Pima). For the two GP methods, the computation time depends the number of inducing points and the number of dimensions. Therefore, as the dataset becomes larger (Spambase), the computation time increases.
13

---

Under review as a conference paper at ICLR 2020
Table 5: Results error/gap EMNIST.
Avg. Awg.Apprxx.Frr
Frr.Model..rr.
Avg.Amt,
Avg.Amt.
0.0650
0.0601
0.0049
0.0053
D.3 MORE EXPERIMENTS ON BAYESIAN NEURAL NETWORK
In this section, we present more experimental results/details on Bayesian NN.
Computation of AUROC/AUPR: First, we show how AUROC/AUPR is calculated. We use entropy (E) and max probability (P) of the particle mean, and differential entropy (D) of locally fitted student distribution as the measures for the teacher. We use entropy (E) and max probability (P) of the particle mean, and the scalar output of student concentration model (C) as the measures for the OPU student. For OOD detection, say 9x, has higher E (or lower P or lower C) than 9x2+ then 9x. is more likely to be an OOD data. The value of E, P and C (D for the teacher) are computed for all testing data points (both in domain data and out-of-domaii data), which rank the data points based on the numerical magnitude. The ROC curve is plotted by setting threshold on each magnitude and compute the True Positive Rate and False Positive Rate at each threshold. The PR curve is plotted similar by computing the Precision and Recall. Then the area under two curves (AUROC and AUPR) can be obtained. For misclassification detection, the similar calculation process is applied.
Error/gap: Second, we use the four types of error/gap under MMD defined in Sec. 2.4 to show the effectiveness of amortized approximation and the suitableness of using Dirichlet family. Specifically, the following measures are used:
Averaged total approximation error (Avg.Approx.Err.) : The averaged MMD be- tween teacher's particles (for each x) and the predicted Dirichlet by OPU, i.c.,
E (MM((xx:Px..)
Averaged total approximation error (Avg.Model.Em.): The averaged MMD between teacher's 's particles (for each x) and locally fitted Dirichlet, i..., N E MMD(Px;: Tx. ). Averaged local amortization gap lAvg.Amt.Gaan The difference between Avg.Approx.Eni Err and Avg.Model.Errs i.c., Axx), defined in Eq. 7 Averaged local amortization error Avv..Amt.Erg We define this type of error to be the averaged MMD between locally fitted Dirichlet (for each x) and the predicted Dirichlet by
OPU, i..., is given by Eq. 8
1
5 MMD( gx,) The relation between Avg.Amt.Cr Gap and Avg.Amt.Enm
N
We show the numerical results EMNIST dataset by Table can be observed that the Avg.Amt. Err. 0.0053 is low and it bounds the Avg.Amt.Gap (xx)=Avg.Approx.f Err. Mod.Mo............ which is consistent with Eq. g The approximation error is mainly determined by the model error, which turns out to be acceptably small. This is consistent with the analysis and also shows the effectiveness and suitableness of using Dirichlet family.
Cifar10 results: Third, we show the experimental results of OPU approximating Bayesian NN on 10 dataset in Table 6 The NN model for teacher and student is standard VGG19 with the output dimension of concentration model to be 1. The MCDP teacher model is trained with SGD optimizer and cyclical learning rate policy. The base learning rate is initialized to be 0.01 and linearly increased to 10x of learning rate, then decreases back to the base learning rate. The base learning rate is then scaled by half. The number of MCDP particles is 700. The OPU students are trained with KL/EMD/MMD using the corresponding al algorithm for 200 epochs.
As the task is harder than MNIST, particles tends to be more spread over the corners. Under such case, it can be observed that the performance of EMD and MMD is much better than that of KL.
Teachers with fewer samples: Fourth, the performance of original MCDP/SGLD with fewer samples is shown in Table 7 We note that using few particles would affect the performance of Bayesian classifier, especially the performance on OOD detection. In real-world cases like automatic driving where the robustness is critical, it is not worth to trade safety for speed. Therefore, OPU solves this issue by providing accurate estimation of predictive uncertainty with short evaluation time. Besides the speedup, OPU approximation provides full distribution to characterize the predictive
14

---

Under review as a conference paper at ICLR 2020
Table 6: Results BNN Cifar10.
MisC detection
OOD detection LSUN
Acc.
Model
AUROC
AUPR
AUROC
AUPR
(%)
MCDP-KI
92.2 (P)
(P)
(E)
88.7 (E)
92.4
MCDP-EMD
92.2 (P)
(P)
91.4 (D)
89.1 (D)
92.4
MCDP-MMD
92)
(P)
91.0 (D)
89.3 (D)
92.4
OPU-MCDP-KL
(P)
(P)
86.1 (E)
(E)
89,9
OPU-MCDP-EMD
91.8 (B)
46.9 (P)
93.5 (C)
(2)
91.8
OPU-MCDP-MMI
91.3 (B)
46.6
92.9 (C)
91.7 (C)
91.8
Table 7: Results on MNIST MCDP/SGLDD with fewer samples.
Model
MisC detection
Omniglot AUROC
SEMEION
Acc.
Time
AUROC
AUROC
(%)
MCDP
96.1 (-1.2)
98.5
(-4.1)
(-3.0)
132.1
SGLD
97.1 (-0.9)
91.2 (-8.0)
82.5 (-6.8)
98.0 (-0.4)
141.9
distribution, which is not available with particle approximation. This allows for better uncertainty measures such as differential entropy.
Un-distilled teachers: Finally, we show the detailed un-distilled MCDP/SGLD, whose elements might be covered by differential entropy of locally fitted Dirichlet in Table I The detailed results are shown in Table g
E APPENDIX: SAMPLING
In this section, we illustrate the details for extracting samples from Bayesian logistic regression, Bayesian neural network and Gaussian process. Under some contexts, we use X and Y to collectively denote the inputs and outputs respectively for previous 2)
E.1 BAYESIAN LOGISTIC REGRESSION
The Polya-Gamma (PG) scheme (Polson et al. 2013) is a data augmentation strategy that allows for a closed-form Gibbs sampler. In binary classification, i.c., Yn y {1. 0). let be the regression coefficients with a Gaussian conditional conjugate prior p(0) No A A-1). The PG Gibbs sampler composed of the following two conditionals,
Wn 0, Xx
PG x N(m,, Vw),
(1I)
0,w,
(12)
where wa is the augmenting data corresponding to the nihh data point. The posterior conditional variance and mean are given by V_ = =x
and m V V ( respectively.
Given this formulation we will be able to collect samples from the posterior after burn-in period. The posterior samples { together with dataset D', are used train our approximation with goal defined in Eq. 10
As an MCMC method, the PG augmentation scheme offers accurate samples. Other alternative methods such as local variational approximation can be employed, which are much faster but sacrifice accuracy.
E.2 MONTE CARLO DROPOUT
A neural network with dropout applied before every weight layer shown to be approximation probabilistic deep Gaussian process (GP) Gal & Ghahramani 2016). Let (()) to be the approximate
distribution the GP posterior. Here, { HO and O is parameter matrix of dimensions K; x K for NN layer i. In this approximation, g()) can be defined through direct modification:
θ M,diag((a
(13)
15

---

Under review as a conference paper at ICLR 2020
Table 8: Results MNIST undistilled MCDP/SGLD.
MisC detection
Omniglot
SEMEION
Acc.
Model
AUROC
AUPR
AUROC
AUPR
AUROC
AUPR
(%)
MCDP
97.3 (E)
43.0 (E)
(E) 99.2 (E)
98.8 (P)
86.8 (E)
53.8 (P)
97.9
SGLD
97.9 (E)
46.2 (E)
99.6 (E)
89.3 (E)
46.8 (E)
98.4
where 2t.j Bern(p;) for [T] and [K,-1]. given some prior dropout probabilities p; and matrices Q are treated as variational parameters. The binary variable žj. = 0 indicates that unit j in layer being dropped out an input layer i. The predictive mean of this approximation is given by R8 [|y|X| R E. f(x,, ...;:..., ,..,.,,)))) which hence referred to MC dropout (MCDP).
We use OPU to approximate the uncertainty induced by g(0). A sample in the MC ensemble is given by 0. i. ) M,, We set Φ { Z to the mean of O., i.e.,
$ = Σ ==1 i M, 1 Σ diag([zz,,o
MCDP provides a simple way of approximating Bayesian inference through dropout sampling. However, it still introduces a variational approximation to exact Bayesian posterior. Therefore, we further include more accurate way to generate MC ensemble stochastic gradient Langevin dynamics (SGLD).
E.3 VANILLA SGLD
SGLD enables mini-batch MC sampling from the posterior via adding noise step to SGD (Welling Teh 201T). We choose the "vanilla" version of SGLD our approximation. Specifically, we start training of fo(-) from ( In each epoch with mini-batch size B,
(++11) (()) + e, e logg( (D) + m
(14)
0 + c Vlogp( (
+ Σ log p(yslx,, (()) + me.
(15)
where is the size of a mini-batch and T/t ~ (O, 2e,1). After SGLD converges at step T, the samples 0, = ((T++) is collected by running the training process for another S iterations. We use averaged samples as parameter of h, i.e. d1 = 5 - 0, and train OPU by Eq. 10 with this ensemble.
E.4 MONTE CARLO GAUSSIAN PROCESS
We apply the OPU framework to the GP framework, which demonstrate its use on parametric classifier. Let data D be split into as input matrix X and output matrix Y. We consider GP prior over the space of functions, i.e., d GP(O, (C). where I. is positive definite kernel controlling the prior belief on smoothness. Existing techniques allow us to compute ((p)) which approximates p(u(D) and is comparable to previous (()) under a parametric model. In a classification task, the posterior p(p1) (D) can be sampled via MCMC (Hensman et al. 2015a) or approximated, c.g., via variational approximation (Hensman et al.) 2015b). Let μ4 be a shorthand for μ* and is then defined as S(1="). If Gaussian variational approximation is used, the marginal posterior p( [|x="xx, |x, D) induces registic-normal distribution for π |x. D). .. In our approximation, { from p(u")) |x, D). The optimization goal is the same as that in the NN case with a different target distribution defined
obtain samples
S >
f( D)
8( S(n)))
(16)
16


Reviews:
[{"review_id": "Skg2gUHcFr", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Overall I liked several results presented in this paper. The findings in Figure 2 gives clear illustration on how Bayesian classification models distinguish between in-distribution difficult-to-classify data and out-of-distribution data, namely uncertain predicted mean and large predicted variance. Though I believe this eventually depends on what kind of \"kernel\"s are used to correlate data points in the prior, throughout the paper I assume meaningful \"kernel\"s are used (for Bayesian NNs this is rooted in the inductive bias of neural networks). \n\nAnother result that I liked is in experiments we can clearly see the advantage of considering the bayesian predictive distribution over a single predictive mean. As demonstrated by BDK-SGLD vs. BDK-DIR-SGLD. \n\nThe proposed idea is a simple and meaningful improvement over previous works. Though the contribution is quite limited, the authors present it with great clarity, which I appreciated. However, I do find some discussions in the paper unnecessary and would expect for more technical contributions. For example, I didn't see the argument for the whole section discussing amortization gap. Everything seems straightforward given the hypothesis F is of enough capacity, which obviously does not hold in practice.\n\nMany other concerns are summarized below:\n* What if the teacher predictive distribution is far unlike a Dirichlet? How much is the discrepancy between teacher and student predictive distribution? Theoretical or empirical evidence is needed for this modeling choice.\n* One importance advantage of Bayesian classification models is that they can capture the covariance between predictions  of different data points. By amortization this advantage no longer exists.\n* In the paper the authors keep mentioning that the method can be applied to GPs but I don't see experiments or algorithms for it?\n* The concentration model is parameterized using an exponential activation, how does this activation affect the performance?\n* The distilling process is done on a held-out dataset. Which may not be wanted because an advantage of Bayesian classification models (eg. GPs) is that all hyperparameters can be automatically selected by marginal likelihoods and don't need a held-out validation set.\n* MMD/wasserstein distances are cool but they require also samples from the student, which adds more variance to the distillation process.\n* The experiment setup is extremely unclear to me. What is \"uncertainty measures\", are they used as metrics for detecting out-of-distribution data, how are AUROC/AUPR calculated using the uncertainty measures? I can guess the meaning but the paper should be more clear about this.\n* I found most numbers convincing except that sometimes BDK-SGLD outperforms BDK-DIR-SGLD, if I understand it right, the predicted mean of BDK-DIR-SGLD should be as good as BDK-SGLD?\n\nMinor:\n* On page 4, above Eq. (4) there is a broken figure link.\n* On Page 7, \"To save space, we only present the best performing uncertainty measure (E, P or C)\". What is \"C\" here?", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "SJx1BMq5FB", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Thank the authors for your detailed rebuttal. I agree with the authors that the proposed method acts as a useful tool for \"real-time evaluation of induced predictive uncertainty\", and the experiments also validate that the method indeed achieves comparable performance with smaller computations. But for now, I am inclined to not change my score.\n\n###################\n\n\nBayesian models maintain the posterior distribution for predictions, which might bring up big computational costs of multiple forwards or big memory costs of multiple particles. To resolve the computational and memory issues at predictions, this paper proposes to distill Bayesian models into an amortized prediction model, avoiding the original multiple forwards. Specifically, in classification, they distill the predictive probabilities into an amortized Dirichlet distribution. They evaluated different distillation metrics, including KL divergence, Earth moving distance, and Maximum mean discrepancy.  Empirically, they evaluate the proposed method over out-of-distribution detection. They demonstrate that their method achieves comparable performance with much speedup.\n\nStrengths, \n1, This paper is well-written and the ideas are well-presented. They evaluated the proposed method over different Bayesian models (MCDP & SGLD) as well different metrics (KL, EMD, MMD), and demonstrate the effectiveness of their method. Overall, this paper is very comprehensive.\n2, As evaluated and validated in the experiments, the proposed method vastly reduces the inference time at test phase. \n\nWeakness,\n1, The paper kind of lacks of novelty. Basically the proposed method distills a Bayesian models into an amortized Dirichlet distribution, which is straightforward. \n2, The baselines such as MCDP-KL, MCDP-EMD are strange, it is wired why you would distill the predictive distribution of a single point to a Dirichlet distribution. And I think it is probably unfair, as distilling the single-point distribution to the Dirichlet under KL, EMD, MMD might require large amount of particles, which they don't have.\n3, Related to (2), more baselines should be compared with to better demonstrate the method's effectiveness. 1) performance of the un-distilled MCDP and SGLD models. 2) BDK and DPN for the MCDP models. 3) MCDP and SGLD with fewer particles. The paper claims to achieve 500x speed up, while I reckon the performance of MCDP and SGLD won't deteriorate a lot if you use only fewer particles. \n4, It would be interesting to see experiments other than out-of-distribution detection, such as calibration. \n\nMinor Issues,\n1, The paper has several un-complied references, such as above eq(4) and appendix D.\n2, The \\Tau(x | \\theta) in Figure 1 is a typo.\n3, Assumption 1 should be put forward to the main articles for comprehensiveness of Lemma 1.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory."}, "extra_text": {}}, {"review_id": "SJgNjT9W9B", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper studies the problem of avoiding Monte Carlo (MC) estimate for the predictive distribution during the test for Bayesian methods. MC estimate will incur multiple passes where the number of passes depends on the number of samples and therefore the cost can be huge. The authors propose One-Pass Uncertainty (OPU) methods to approximate the predictive distribution through distillation. Experiments on Bayesian neural networks are conducted to demonstrate the proposed method.\n\nQuality:\nThe proposed method appears to be technically sound. The view of approximating the predictive distribution over simplex is interesting and may inspire future studies under this formulation. Although the restriction of the student distribution to be tractable seems to limit the design of the student model significantly. And this restrictive distribution family may cause large amortization error, as suggested by Lemma 1 in the paper.\n\nThe experiments are well-conducted, and the proposed method is well-evaluated.\n\nSignificance:\nThis paper studies an important problem in Bayesian machine learning and the proposed method can be combined with many Bayesian methods to reduce the computational cost during the test. \n\nOriginality:\nAs far as I know, the method is novel. The related work is adequately cited. \n\nClarity:\nThis paper is well-written and easy to follow. \n\n", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes to speed up Bayesian deep learning at test time by training a student network to approximate the BNN's output distribution. The idea is certainly a reasonable thing to try, and the writing is mostly good (though as some reviewers point out, certain sections might not be necessary). The idea is fairly obvious, though, so the question is whether the experimental results are impressive enough by themselves to justify acceptance. The method is able to get close to the performance achieved by Monte Carlo estimators with much lower cost, although there is a nontrivial drop in accuracy. This is probably worth paying if it achieves 500x computation reduction as claimed in the paper, though the practical gains are probably much smaller since Monte Carlo methods are rarely used with 500 samples. Overall, this seems a bit below the bar for ICLR.

Decision: Reject