Sampling-Free Learning of Bayesian Quantized Neural Networks

1 INTRODUCTION . A Bayesian approach to deep learning considers the network ’ s parameters to be random variables and seeks to infer their posterior distribution given the training data . Models trained this way , called Bayesian neural networks ( BNNs ) ( Wang & Yeung , 2016 ) , in principle have well-calibrated uncertainties when they make predictions , which is important in scenarios such as active learning and reinforcement learning ( Gal , 2016 ) . Furthermore , the posterior distribution over the model parameters provides valuable information for evaluation and compression of neural networks . There are three main challenges in using BNNs : ( 1 ) Intractable posterior : Computing and storing the exact posterior distribution over the network weights is intractable due to the complexity and high-dimensionality of deep networks . ( 2 ) Prediction : Performing a forward pass ( a.k.a . as probabilistic propagation ) in a BNN to compute a prediction for an input can not be performed exactly , since the distribution of hidden activations at each layer is intractable to compute . ( 3 ) Learning : The classic evidence lower bound ( ELBO ) learning objective for training BNNs is not amenable to backpropagation as the ELBO is not an explicit function of the output of probabilistic propagation . These challenges are typically addressed either by making simplifying assumptions about the distributions of the parameters and activations , or by using sampling-based approaches , which are expensive and unreliable ( likely to overestimate the uncertainties in predictions ) . Our goal is to propose a sampling-free method which uses probabilistic propagation to deterministically learn BNNs . A seemingly unrelated area of deep learning research is that of quantized neural networks ( QNNs ) , which offer advantages of computational and memory efficiency compared to continuous-valued models . QNNs , like BNNs , face challenges in training , though for different reasons : ( 4.1 ) The non- differentiable activation function is not amenable to backpropagation . ( 4.2 ) Gradient updates cease to be meaningful , since the model parameters in QNNs are coarsely quantized . In this work , we combine the ideas of BNNs and QNNs in a novel way that addresses the aforementioned challenges ( 1 ) ( 2 ) ( 3 ) ( 4 ) in training both models . We propose Bayesian quantized networks ( BQNs ) , models that ( like QNNs ) have quantized parameters and activations over which they learn ( like BNNs ) categorical posterior distributions . BQNs have several appealing properties : • BQNs solve challenge ( 1 ) due to their use of categorical distributions for their model parameters . • BQNs can be trained via sampling-free backpropagation and stochastic gradient ascent of a dif- ferentiable lower bound to ELBO , which addresses challenges ( 2 ) , ( 3 ) and ( 4 ) above . • BQNs leverage efficient tensor operations for probabilistic propagation , further addressing chal- lenge ( 2 ) . We show the equivalence between probabilistic propagation in BQNs and tensor contractions ( Kolda & Bader , 2009 ) , and introduce a rank-1 CP tensor decomposition ( mean-field approximation ) that speeds up the forward pass in BQNs . • BQNs provide a tunable trade-off between computational resource and model complexity : using a refined quantization allows for more complex distribution at the cost of more computation . • Sampling from a learned BQN provides an alternative way to obtain deterministic QNNs . In our experiments , we demonstrate the expressive power of BQNs . We show that BQNs trained using our sampling-free method have much better-calibrated uncertainty compared with the stateof-the-art Bootstrap ensemble of quantized neural networks ( E-QNN ) trained by Courbariaux et al . ( 2016 ) . More impressively , our trained BQNs achieve comparable log-likelihood against Gaussian Bayesian neural network ( BNN ) trained with stochastic gradient variational Bayes ( SGVB ) ( Shridhar et al. , 2019 ) ( the performance of Gaussian BNNs are expected to be better than BQNs since they allows for continuous random variables ) . We further verify that BQNs can be easily used to compress ( Bayesian ) neural networks and obtain determinstic QNNs . Finally , we evaluate the effect of mean-field approximation in BQN , by comparing with its Monte-Carlo realizations , where no approximation is used . We show that our sampling-free probabilistic propagation achieves similar accuracy and log-likelihood — justifying the use of mean-field approximation in BQNs . Related Works . In Appendix A , we survey different approaches for training Bayesian neural networks including sampling-free assumed density filtering ( Minka , 2001 ; Soudry et al. , 2014 ; Hernández-Lobato & Adams , 2015 ; Ghosh et al. , 2016 ) , sampling-based variational inference ( Graves , 2011 ; Blundell et al. , 2015 ; Shridhar et al. , 2019 ) , as well as sampling-free variational inference ( Wu et al. , 2018 ) , probabilistic neural networks ( Wang et al. , 2016 ; Shekhovtsov & Flach , 2018 ; Gast & Roth , 2018 ) , quantized neural network ( Han et al. , 2015 ; Courbariaux et al. , 2015 ; Zhu et al. , 2016 ; Kim & Smaragdis , 2016 ; Zhou et al. , 2016 ; Rastegari et al. , 2016 ; Hubara et al. , 2017 ; Esser et al. , 2015 ; Peters & Welling , 2018 ; Shayer et al. , 2017 ) , and tensor networks and tensorial neural networks ( Grasedyck et al. , 2013 ; Orús , 2014 ; Cichocki et al. , 2016 ; 2017 ; Su et al. , 2018 ; Newman et al. , 2018 ; Robeva & Seigal , 2017 ) . Contributions : . • We propose an alternative evidence lower bound ( ELBO ) for Bayesian neural networks such that optimization of the variational objective is compatible with the backpropagation algorithm . • We introduce Bayesian quantized networks ( BQNs ) , establish a duality between BQNs and hierarchical tensor networks , and show prediction a BQN is equivalent to a series of tensor contractions . • We derive a sampling-free approach for both learning and inference in BQNs using probabilistic propagation ( analytical inference ) , achieving better-calibrated uncertainty for the learned models . • We develop a set of fast algorithms to enable efficient learning and prediction for BQNs . 2 BAYESIAN NEURAL NETWORKS . Notation . We use bold letters such as θ to denote random variables , and non-bold letters such as θ to denote their realizations . We abbreviate Pr [ θ = θ ] as Pr [ θ ] and use bold letters in an equation if the equality holds for arbitrary realizations . For example , Pr [ x , y ] = Pr [ y|x ] Pr [ x ] means Pr [ x = x , y = y ] = Pr [ y = y|x = x ] Pr [ x = x ] , ∀x ∈ X , y ∈ Y . 2.1 PROBLEM SETTING . Given a dataset D = { ( xn , yn ) } Nn=1 of N data points , we aim to learn a neural network with model parameters θ that predict the output y ∈ Y based on the input x ∈ X . ( 1 ) We first solve the learning problem to find an approximate posterior distribution Q ( θ ; φ ) over θ with parameters φ such that Q ( θ ; φ ) ≈ Pr [ θ|D ] . ( 2 ) We then solve the prediction problem to compute the predictive distribution Pr [ y|x , D ] for arbitrary input x = x given Q ( θ ; φ ) . For notational simplicity , we will omit the conditioning on D and write Pr [ y|x , D ] as Pr [ y|x ] in what follows . In order to address the prediction and learning problems in BNNs , we analyze these models in their general form of probabilistic graphical models ( shown in Figure 3b in Appendix B ) . Let h ( l ) , θ ( l ) and h ( l+1 ) denote the inputs , model parameters , and ( hidden ) outputs of the l-th layer respectively . We assume that θ ( l ) ’ s are layer-wise independent , i.e . Q ( θ ; φ ) = ∏L−1 l=0 Q ( θ ( l ) ; φ ( l ) ) , and h ( l ) follow the Markovian property , i.e . Pr [ h ( l+1 ) |h ( : l ) , θ ( : l ) ] = Pr [ h ( l+1 ) |h ( l ) , θ ( l ) ] . 2.2 THE PREDICTION PROBLEM . Computing the predictive distribution Pr [ y|x , D ] with a BNN requires marginalizing over the random variable θ . The hierarchical structure of BNNs allows this marginalization to be performed in multiple steps sequentially . In Appendix B , we show that the predictive distribution of h ( l+1 ) given input x = x can be obtained from its preceding layer h ( l ) by Pr [ h ( l+1 ) |x ] ︸ ︷︷ ︸ P ( h ( l+1 ) ; ψ ( l+1 ) ) = ∫ h ( l ) , θ ( l ) Pr [ h ( l+1 ) |h ( l ) , θ ( l ) ] Q ( θ ( l ) ; φ ( l ) ) Pr [ h ( l ) |x ] ︸ ︷︷ ︸ P ( h ( l ) ; ψ ( l ) ) dh ( l ) dθ ( l ) ( 1 ) This iterative process to compute the predictive distributions layer-by-layer sequentially is known as probabilistic propagation ( Soudry et al. , 2014 ; Hernández-Lobato & Adams , 2015 ; Ghosh et al. , 2016 ) . With this approach , we need to explicitly compute and store each intermediate result Pr [ h ( l ) |x ] in its parameterized form P ( h ( l ) ; ψ ( l ) ) ( the conditioning on x is hidden in ψ ( l ) , i.e . ψ ( l ) is a function of x ) . Therefore , probabilistic propagation is a deterministic process that computes ψ ( l+1 ) as a function of ψ ( l ) and φ ( l ) , which we denote as ψ ( l+1 ) = g ( l ) ( ψ ( l ) , φ ( l ) ) . Challenge in Sampling-Free Probabilistic Propagation . If the hidden variables h ( l ) ’ s are continuous , Equation ( 1 ) generally can not be evaluated in closed form as it is difficult to find a family of parameterized distributions P for h ( l ) such that h ( l+1 ) remains in P under the operations of a neural network layer . Therefore most existing methods consider approximations at each layer of probabilistic propagation . In Section 4 , we will show that this issue can be ( partly ) addressed if we consider the h ( l ) ’ s to be discrete random variables , as in a BQN . 2.3 THE LEARNING PROBLEM . Objective Function . A standard approach to finding a good approximation Q ( θ ; φ ) is variational inference , which finds φ ? such that the KL-divergence KL ( Q ( θ ; φ ) ||Pr [ θ|D ] ) from Q ( θ ; φ ) to Pr [ θ|D ] is minimized . In Appendix B , we prove that to minimizing the KL-divergence is equivalent to maximizing an objective function known as the evidence lower bound ( ELBO ) , denoted as L ( φ ) . max φ L ( φ ) = −KL ( Q ( θ ; φ ) ||Pr [ θ|D ] ) = N∑ n=1 Ln ( φ ) +R ( φ ) where Ln ( φ ) = EQ [ logPr [ yn|xn , θ ] ] andR ( φ ) = EQ [ log ( Pr [ θ ] ) ] +H ( Q ) ( 2 ) Probabilistic Backpropagation . Optimization in neural networks heavily relies on the gradientbased methods , where the partial derivatives ∂L ( φ ) /∂φ of the objective L ( φ ) w.r.t . the parameters φ are obtained by backpropagation . Formally , if the output produced by a neural network is given by a ( sub- ) differentiable function g ( φ ) , and the objective L ( g ( φ ) ) is an explicit function of g ( φ ) ( and not just an explicit function of φ ) , then the partial derivatives can be computed by chain rule : ∂L ( g ( φ ) ) /∂φ = ∂L ( g ( φ ) ) /∂g ( φ ) · ∂g ( φ ) /∂φ ( 3 ) The learning problem can then be ( approximately ) solved by first-order methods , typically stochastic gradient descent/ascent . Notice that ( 1 ) For classification , the function g ( φ ) returns the probabilities after the softmax function , not the categorical label ; ( 2 ) An additional regularizer R ( φ ) on the parameters will not cause difficulty in backpropagation , given ∂R ( φ ) /∂φ is easily computed . Challenge in Sampling-Free Probabilistic Backpropagation . Learning BNNs is not amenable to standard backpropagation because the ELBO objective function L ( φ ) in ( 4b ) is not an explicit ( i.e . implicit ) function of the predictive distribution g ( φ ) in ( 4a ) : gn ( φ ) = EQ [ Pr [ yn|xn , θ ] ] = ∫ θ Pr [ yn|xn , θ ] Q ( θ ; φ ) dθ ( 4a ) Ln ( φ ) = EQ [ log ( Pr [ yn|xn , θ ] ) ] = ∫ θ log ( Pr [ yn|xn , θ ] ) Q ( θ ; φ ) dθ ( 4b ) Although Ln ( φ ) is a function of φ , it is not an explicit function of gn ( φ ) . Consequently , the chain rule in Equation ( 3 ) on which backpropagation is based is not directly applicable .

Summary: This paper proposed Bayesian Quantized Neural Networks (BQN), which is a rough Bayesian treatment for quantized neural networks (QNN). The advantages the BQN over QNN is that it provides better uncertainty estimates as well as accuracy. The disadvantage is that BQN loses the memory efficient property of QNN, that is, BQN in general still needs to store real-value parameters, which parameterize the binary weights. The authors compare BQN to EQNN and show that BQN outperforms the latter.