Accelerating Monte Carlo Bayesian Inference via Approximating Predictive Uncertainty over the Simplex

1 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 are not conjugate hence marginalizing over model prior or posterior can not be performed analytically , which hinders the practical applicability . For tractability , a simple point estimate such as 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 a 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 , we consider a Bayesian classification model trained on dataset D = { ( yn , xn ) } Nn=1 , where xn ∈ X and yn ∈ Y are the nth input and output , respectively . Let the model posterior p ( θ|D ) be approximated by MC estimate 1S ∑S s=1 δ ( θ−θs ) , and the predictive distribution ( a categorical distribution parameterized by the predicted class probabilities ) is thus approximated by p ( y|x , D ) = ∫ p ( y|x , θ ) p ( θ|D ) dθ ≈ 1 S S∑ s=1 p ( y|x , θs ) . The predictive distribution can be accurately estimated as S → ∞ . However , to perform the computation , we need to maintain a large number of samples , repeatedly evaluate the model for S 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 of uncertainty to avoid making mistakes with a high confidence ( Kendall & Gal , 2017 ; Sünderhauf et al. , 2018 ) . Due to 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 . In this work , aiming at boosting the prediction speed while maintaining a rich characterization of the prediction , we propose to approximate the distribution of class probabilities over the simplex induced by the model posterior p ( θ|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. , 2015 ; Bulò et al. , 2016 ) 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 as 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 et al. , 2016 ; Han et al. , 2015 ) ) 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-mover ’ s 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 a significant speedup ( ∼ 500× ) of Bayes models . The results on Bayes NN show that OPU performs better in misclassification detection and OOD detection than state-ofthe-art works in Bayesian knowledge distillation . It can also be observed that explicit disentangling 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 is universally applicable to all Bayesian classification models . 2 ONE-PASS UNCERTAINTY FRAMEWORK . 2.1 INDUCED DISTRIBUTION OVER SIMPLEX . In this section we 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 , θ ) = Cat ( y|T ( x ; θ ) ) ( 1 ) and a prior distribution p ( θ ) be specified over the parameter space Θ , where T : X ×Θ→ SK−1 is 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 ( θ|D ) is obtained and focus on the computation of the predictive distribution . In what follows , let p ( θ|D ) be approximate or exact ( if available ) model posterior , from which samples could be obtained . From a dual perspective , with a fixed input x , the mapping Tx = T ( x ; · ) is also understood as a data-dependent mapping Tx : Θ→ SK−1 , which could be used to transform the posterior p ( θ|D ) to a conditional distribution ( on x ) over the simplex . That is , for each x , we can define a random variable π = Tx ( θ ) whose distribution is induced by θ ∼ p ( θ|D ) and Tx . This is effectively a welldefined push-forward measure Tx # p ( θ|D ) over the simplex ( see Fig . 1 ) . We term the conditional distribution of π|x , D as an “ induced distribution ” and define px ( π ) = p ( π|x , D ) to keep the notation uncluttered . The induced distribution isolates the dependence between input and output and simplifies the computation of predictive distribution due to the change-of-variable formula . Specifically , given a test point x∗ , the predictive distribution can be alternatively written as p ( y|x∗ , D ) = ∫ Θ Cat ( y|Tx∗ ( θ ) ) p ( θ|D ) dθ = ∫ SK−1 Cat ( y|π ) px∗ ( π ) dπ . ( 2 ) Our key insight is that px contains all information we need for prediction and uncertainty measurement . Hence π is sufficient in the sense that , given π , y is independent of both θ and x . The isolation combines the complexities in both the likelihood and posterior into a single object px , and keeps a simple dependence structure between y and π . Also , the isolation renders the last probabilistic layer y|π nuisance , which can thus be peeled off for prediction . To validate the idea , one can show that the predictive distribution is simply p ( y|x , D ) = Cat ( y|Epxπ ) . The difference between probabilistic structures of original Bayesian model and the isolated version is showed in Fig . 3 , 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 . θs ∼ p ( θ|D ) ; 2 ) “ push-forward ” the samples by Tx , i.e . πs = Tx ( θs ) . The behavior of π can be empirically observed via the particles { πs } Ss=1 . As shown in Fig . 2 , π is able to distinguish different types of input x based on the behavior of its samples . In Fig . 2 ( a ) , the inputs are similar to the training data . The particles gather around a corner of the simplex , indicating a confident prediction . In Fig . 2 ( b ) , 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 . 2 ( c ) , 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 . 2 ( d ) , the particles spread along a line between two corners , indicating the model is confident that the result is not at the other corner . However , using particles for inference is time-consuming as each particle requires one evaluation of the model . This motivates us to use a tractable conditional distribution q ( π|x ) to approximate px . 2.2 AMORTIZED APPROXIMATION . The view of px enables flexible choices of q ( π|x ) , as any distribution defined on RK can be transformed to SK−1 via logistic transformation ( Aitchison , 1982 ) . However , modeling q ( π|x ) locally for every input is not practical , as the design efforts and number of parameters grows linearly with the number of data points . Therefore , we propose to approximate px in two aspects : 1 ) use a single family of distribution q ; 2 ) to generalize to unseen examples , let the parameter of q , αx = α ( x ; φ ) , be a function depending on x and parameterized by a set of global adaptive parameters φ , and thus qx = q ( π|x ) = q ( π|αx ) . The computational cost is amortized by casting the problem of learning a series of conditional distributions to a regression problem . We term px as “ teacher distribution ” and qx as “ student ” . In our method , as seen in graphical representation in Fig . 3 , 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 . 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 , the amortized approximation can be applied to any Bayesian classifier . Note that the approximation framework can be extended to non-parametric model like Gaussian process , where the computational cost of inference is high 1 . In this work , we choose the student model q to be a Dirichlet distribution , q ( π|αx ) = Dir ( π|αx ) , where αx is a function mapping input x to a 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 αx = hx · egx , where hx = h ( x ; φ1 ) and gx = g ( x ; φ2 ) are two neural networks , and the vector output h sums to 1 . Vector output h determines the mean of the Dirichlet ( i.e. , the predicted class probabilities ) , and g 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 ( y ` = 1|x , φ ) = ∫ Cat ( y ` = 1|π ) Dir ( π|αx ) dπ = αx , ` ∑ c αx , c = hx , ` where y ` and αx , ` are the ` -th coordinate of y and αx respectively . Therefore , we call hx as the “ prediction model ” ( PM ) . Similarly , the precision parameter α0 ( determines sharpness ) of the Dirichlet solely depends on gx , α0 = ∑ c αx , c = ∑ c hx , ce gx = egx . Therefore , we term gx as the “ concentration model ” ( CM ) . Based on this property , whether the Dirichlet is flat or not , can be fully characterized by CM . It can be expected that , when approximating particles in Fig . 2 ( a ) and ( b ) , the output value of CM is high , as the samples are concentrated , which means high confidence . CM outputs a low value for particles in Fig . 2 ( c ) , yielding a flat distribution and low confidence .

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