SAMPLING-FREE LEARNING OF BAYESIAN QUANTIZED NEURAL NETWORKS
Year: 2019

Abstract:
Bayesian learning of model parameters in neural networks is important in scenarios where estimates with well-calibrated uncertainty are important. In this paper, we propose Bayesian quantized networks (BQNs), quantized neural networks (QNNs) for which we learn a posterior distribution over their discrete parameters. We provide a set of efficient algorithms for learning and prediction in BQNs without the need to sample from their parameters or activations, which not only allows for differentiable learning in QNNs, but also reduces the variance in gradients. We evaluate BQNs on MNIST, Fashion-MNIST, KMNIST and CIFAR10 image classification datasets. compared against bootstrap ensemble of QNNs (E-QNN). We demonstrate BQNs achieve both lower predictive errors and better-calibrated uncertainties than E-QNN (with less than 20% of the negative log-likelihood).

heading: 1 INTRODUCTION
text: 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 cannot 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).

heading: Contributions:
text: • 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.

heading: 2 BAYESIAN NEURAL NETWORKS
text: 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 .

heading: 2.1 PROBLEM SETTING
text: 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)].

heading: 2.2 THE PREDICTION PROBLEM
text: 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.

heading: 2.3 THE LEARNING PROBLEM
text: 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.

heading: 3 PROPOSED LEARNING METHOD FOR BAYESIAN NEURAL NETWORKS
text: Alternative Evidence Lower Bound. We make learning in BNNs amenable to backpropagation by developing a lower bound Ln(φ) ≤ Ln(φ) such that ∂Ln(φ)/∂φ can be obtained by chain rule (i.e. Ln(φ) is an explicit function of the results from the forward pass.) With Ln(φ) in hand, we can (approximately) find φ? by maximizing the alternative objective via gradient-based method:
φ? = arg max φ L(φ) = arg max φ
( R(φ) +
N∑ n=1 Ln(φ)
) (5)
In Appendix C.1, we proved one feasible Ln(φ) which only depends on second last output h(L−1). Theorem 3.1 (Alternative Evidence Lower Bound). Define each term Ln(φ) in L(φ) as
Ln(φ) := Eh(L−1)∼P ; θ(L−1)∼Q [ log ( Pr[yn|h(L−1),θ(L−1)] )] (6)
then Ln(φ) is a lower bound of Ln(φ), i.e. Ln(φ) ≤ Ln(φ). The equality Ln(φ) = Ln(φ) holds if h(L−1) is deterministic given input x and all parameters before the last layer θ( :L−2).
Analytic Forms of Ln(φ). While the lower bound in Theorem 3.1 applies to BNNs with arbitrary distributions P on hidden variables h, Q on model parameters θ, and any problem setting (e.g. classification or regression), in practice sampling-free probabilistic backpropagation requires thatLn(φ) can be analytically evaluated (or further lower bounded) in terms of φ(L−1) and θ(L−1). This task is nontrivial since it requires redesign of the output layer, i.e. the function of Pr[y|h(L−1),θ(L−1)]. In this paper, we develop two layers for classification and regression tasks, and present the classification case in this section due to space limit. Since Ln(φ) involves the last layer only, we omit the superscripts/subsripts of h(L−1), ψ(L−1), φ(L−1), xn, yn, and denote them as h, ψ, φ, x, y . Theorem 3.2 (Analytic Form of Ln(φ) for Classification). Let h ∈ RK (with K the number of classes) be the pre-activations of a softmax layer (a.k.a. logits), and φ = s ∈ R+ be a scaling factor that adjusts its scale such that Pr[y = c|h, s] = exp(hc/s)/ ∑K k=1 exp(hk/s). Suppose the logits {hk}Kk=1 are pairwise independent (which holds under mean-field approximation) and hk follows a Gaussian distribution hk ∼ N (µk, νk) (therefore ψ = {µk, νk}Kk=1) and s is a deterministic parameter. Then Ln(φ) is further lower bounded as Ln(φ) ≥ µcs − log (∑K k=1 exp ( µk s + νk 2s2 )) .
The regression case and proofs for both layers are deferred to Appendix C.

heading: 4 BAYESIAN QUANTIZED NETWORKS (BQNS)
text: While Section 3 provides a general solution to learning in BNNs, the solution relies on the ability to perform probabilistic propagation efficiently. To address this, we introduce Bayesian quantized
networks (BQNs) — BNNs where both hidden units h(l)’s and model parameters θ(l)’s take discrete values — along with a set of novel algorithms for efficient sampling-free probabilistic propagation in BQNs. For simplicity of exposition, we assume activations and model parameters take values from the same set Q, and denote the degree of quantization as D = |Q|, (e.g. Q = {−1, 1}, D = 2).

heading: 4.1 PROBABILISTIC PROPAGATION AS TENSOR CONTRACTIONS
text: Lemma 4.1 (Probabilistic Propagation in BQNs). After quantization, the iterative step of probabilistic propagation in Equation (1) is computed with a finite sum instead of an integral:
P (h(l+1);ψ(l+1)) = ∑
h(l),θ(l) Pr[h(l+1)|h(l), θ(l)] Q(θ(l);φ(l)) P (h(l);ψ(l)) (7)
and a categorically distributed h(l) results in h(l+1) being categorical as well. The equation holds without any assumption on the operation Pr[h(l+1)|h(l), θ(l)] performed in the neural network.
Notice all distributions in Equation (7) are represented in high-order tensors: Suppose there are I input units, J output units, and K model parameters at the l-th layer, then h(l) ∈ QI , θ(l) ∈ QK , and h(l+1) ∈ QJ , and their distributions are characterized by P (h(l);ψ(l)) ∈ RDI , Q(θ(l);φ(l)) ∈ RDK , P (h(l+1);ψ(l+1)) ∈ RDJ , and Pr[h(l+1)|h(l),θ(l)] ∈ RDJ×DI×DK respectively. Therefore, each step in probabilistic propagation is a tensor contraction of three tensors, which establishes the duality between BQNs and hierarchical tensor networks (Robeva & Seigal, 2017).
Since tensor contractions are differentiable w.r.t. all inputs, BQNs thus circumvent the difficulties in training QNNs (Courbariaux et al., 2015; Rastegari et al., 2016), whose outputs are not differentiable w.r.t. the discrete parameters. This result is not surprising: if we consider learning in QNNs as an integer programming (IP) problem, solving its Bayesian counterpart is equivalent to the approach to relaxing the problem into a continuous optimization problem (Williamson & Shmoys, 2011).
Complexity of Exact Propagation. The computational complexity to evaluate Equation (7) is exponential in the number of random variables O(DIJK), which is intractable for quantized neural network of any reasonable size. We thus turn to approximations.

heading: 4.2 APPROXIMATE PROPAGATION VIA RANK-1 TENSOR CP DECOMPOSITION
text: We propose a principled approximation to reduce the computational complexity in probabilistic propagation in BQNs using tensor CP decomposition, which factors an intractable high-order probability tensor into tractable lower-order factors (Grasedyck et al., 2013). In this paper, we consider the simplest rank-1 tensor CP decomposition, where the joint distributions of P and Q are fully factorized into products of their marginal distributions, thus equivalent to the mean-field approximation (Wainwright et al., 2008). With rank-1 CP decomposition on P (h(l);ψ(l)),∀l ∈ [L], the tensor contraction in (7) reduces to a standard Tucker contraction (Kolda & Bader, 2009)
P (h (l+1) j ;ψ (l+1) j ) ≈ ∑ h(l),θ(l) Pr[h (l+1) j |θ (l), h(l)] ∏ k Q(θ (l) k ;φ (l) k ) ∏ i P (h (l) i ;ψ (l) i ) (8)
where each term of ψ(l)i , φ (l) k parameterizes a single categorical variable. In our implementation, we store the parameters in their log-space, i.e. Q(θ(l)k = Q(d)) = exp(ψ (l) k (d))/ ∑D q=1 exp(φ (l) k (q)).
Fan-in Number E. In a practical model, for the l-th layer, an output unit h(l+1)j only (conditionally) depends on a subset of all input units {h(l)i } and model parameters {θ (h) k } according to the connectivity pattern in the layer. We denote the set of dependent input units and parameters for h(l+1)j as I(l+1)j andM (l+1) j , and define the fan-in number E for the layer as maxj
∣∣∣I(l+1)j ∣∣∣+ ∣∣∣M(l+1)j ∣∣∣. Complexity of Approximate Propagation. The approximate propagation reduces the computational complexity from O(DIJK) to O(JDE), which is linear in the number of output units J if we assume the fan-in number E to be a constant (i.e. E is not proportional to I).

heading: 4.3 FAST ALGORITHMS FOR APPROXIMATE PROPAGATION
text: Different types of network layers have different fan-in numbers E, and for those layers with E greater than a small constant, Equation (8) is inefficient since the complexity grows exponential in E. Therefore in this part, we devise fast(er) algorithms to further lower the complexity.
Small Fan-in Layers: Direct Tensor Contraction. If E is small, we implement the approximate propagation through tensor contraction in Equation (8). The computational complexity is O(JDE) as discussed previously. See Appendix D.1 for a detailed discussion.
Medium Fan-in Layers: Discrete Fourier Transform. If E is medium, we implement approximate propagation through fast Fourier transform since summation of discrete random variables is equivalent to convolution between their probability mass function. See Appendix D.2 for details. With the fast Fourier transform, the computational complexity is reduced to O(JE2D log(ED)).
Large Fan-in Layers: Lyapunov Central Limit Theorem. In a typical linear layer, the fan-inE is large, and a super-quadratic algorithm using fast Fourier transform is still computational expensive. Therefore, we derive a faster algorithm based on the Lyapunov central limit theorem (See App D.3) With CLT, the computational complexity is further reduced to O(JED).
Remarks: Depending on the fan-in numbersE, we adopt CLT for linear layers with sufficiently large E such as fully connected layers and convolutional layers; DFT for those with medium E such as average pooling layers and depth-wise layers; and direct tensor contraction for those with small E such as shortcut layers and nonlinear layers.

heading: 5 EXPERIMENTS
text: In this section, we demonstrate the effectiveness of BQNs on the MNIST, Fashion-MNIST, KMNIST and CIFAR10 classification datasets. We evaluate our BQNs with both multi-layer perceptron (MLP) and convolutional neural network (CNN) models. In training, each image is augmented by a random shift within 2 pixels (with an additional random flipping for CIFAR10), and no augmentation is used in test. In the experiments, we consider a class of quantized neural networks, with both binary weights and activations (i.e. Q = {−1, 1}) with sign activations σ(·) = sign(·). For BQNs, the distribution parameters φ are initialized by Xavier’s uniform initializer, and all models are trained by ADAM optimizer (Kingma & Ba, 2014) for 100 epochs (and 300 epochs for CIFAR10) with batch size 100 and initial learning rate 10−2, which decays by 0.98 per epoch.
Training Objective of BQNs. To allow for customized level of uncertainty in the learned Bayesian models, we introduce a regularization coefficient λ in the alternative ELBO proposed in Equation (5) (i.e. a lower bound of the likelihood), and train the BQNs by maximizing the following objective:
L(φ) = N∑ n=1 Ln(φ) + λR(φ) = λ
( 1/λ
N∑ n=1 Ln(φ) +R(φ)
) (9)
where λ controls the uncertainty level, i.e. the importance weight of the prior over the training set.
Baselines. (1) We compare our BQN against the baseline – Bootstrap ensemble of quantized neural networks (E-QNN). Each member in the ensemble is trained in a non-Bayesian way (Courbariaux et al., 2016), and jointly make the prediction by averaging over the logits from all members. Note
that Courbariaux et al. (2016) is chosen over other QNN training methods as the baseline since it trains QNN from random initialization, thus a fair comparison to our approach. Details are discussed in Appendix A. (2) To exhibit the effectiveness of our BQN, we further compare against continuousvalued Bayesian neural network (abbreviated as BNN) with Gaussian parameters. The model is trained with stochastic gradient variational Bayes (SGVB) augmented by local re-parameterization trick (Shridhar et al., 2019). Since the BNN allows for continuous parameters (different from BQN with quantized parameters), the predictive error is expected to be lower than BQN.
Evaluation of BQNs. While 0-1 test error is a popular metric to measure the predictive performance, it is too coarse a metric to assess the uncertainty in decision making (for example it does not account for how badly the wrong predictions are). Therefore, we will mainly use the negative log-likelihood (NLL) to measure the predictive performance in the experiments.
Once a BQN is trained (i.e. an approximate posterior Q(θ) is learned), we consider three modes to evaluate the behavior of the model: (1) analytic inference (AI), (2) Monte Carlo (MC) sampling and (3) Maximum a Posterior (MAP) estimation:
1. In analytic inference (AI, i.e. our proposed method), we analytically integrate overQ(θ) to obtain the predictive distribution as in the training phase. Notice that the exact NLL is not accessible with probabilistic propagation (which is why we propose an alternative ELBO in Equation (5)), we will report an upper bound of the NLL in this mode.
2. In MC sampling, S sets of model parameters are drawn independently from the posterior posterior θs ∼ Q(θ),∀s ∈ [S], and the forward propagation is performed as in (non-Bayesian) quantized neural network for each set θs, followed by an average over the model outputs. The difference between analytic inference and MC sampling will be used to evaluate (a) the effect of mean-field approximation and (b) the tightness of the our proposed alternative ELBO.
3. MAP estimation is similar to MC sampling, except that only one set of model parameters θ? is obtained θ? = arg maxθ Q(θ). We will exhibit our model’s ability to compress a Bayesian neural network by comparing MAP estimation of our BQN with non-Bayesian QNN.

heading: 5.1 ANALYSIS OF RESULTS
text: Expressive Power and Uncertainty Calibration in BQNs. We report the performance via all evaluations of our BQN models against the Ensemble-QNN in Table 1 and Figure 1. (1) Compared to E-QNNs, our BQNs have significantly lower NLL and smaller predictive error (except for
Fashion-MNIST with architecture CNN). (2) As we can observe in Figure 1, BQNs impressively achieve comparable NLL to continuous-valued BNN, with slightly higher test error. As our model parameters only take values {−1, 1}, small degradation in predictive accuracy is expected. Evaluations of Mean-field Approximation and Tightness of the Alternative ELBO. If analytic inference (by probabilistic propagation) were computed exactly, the evaluation metrics would have been equal to the ones with MC sampling (with infinite samples). Therefore we can evaluate the approximations in probabilistic propagation, namely mean-field approximation in Equation (8) and relaxation of the original ELBO in Equation (5), by measuring the gap between analytic inference and MC sampling. As shown in Figure 2, such gaps are small for all scenarios, which justifies the approximations we use in BQNs.
To further decouple these two factors of mean-field approximation and relaxation of the original ELBO, we vary the regularization coefficient λ in the learning objective. (1) For λ = 0 (where the prior term is removed), the models are forced to become deterministic during training. Since the deterministic models do not have mean-field approximation in the forward pass, the gap between analytic inference and MC-sampling reflects the tightness of our alternative ELBO. (2) As λ increases, the gaps increases slightly as well, which shows that the mean-field approximation becomes slightly less accurate with higher learned uncertainty in the model.
Compression of Neural Networks via BQNs. One advantage of BQNs over continuous-valued BNNs is that deterministic QNNs can be obtained for free, since a BQN can be interpreted as an ensemble of infinite QNNs (each of which is a realization of posterior distribution). (1) One simple approach is to set the model parameters to their MAP estimates, which compresses a given BQN to 1/64 of its original size (and has the same number of bits as a single QNN). (2) MC sampling can be
interpreted as another approach to compress a BQN, which reduces the original size to its S/64 (with the same number of bits as an ensemble of S QNNs). In Tables 2 and 3, we compare the models by both approaches to their counterparts (a single QNN for MAP, and E-QNN for MC sampling) trained from scratch as in Courbariaux et al. (2016). For both approaches, our compressed models outperform their counterparts (in NLL) . We attribute this to two factors: (a) QNNs are not trained in a Bayesian way, therefore the uncertainty is not well calibrated; and (b) Non-differentiable QNNs are unstable to train. Our compression approaches via BQNs simultaneously solve both problems.

heading: 6 CONCLUSION
text: We present a sampling-free, backpropagation-compatible, variational-inference-based approach for learning Bayesian quantized neural networks (BQNs). We develop a suite of algorithms for efficient inference in BQNs such that our approach scales to large problems. We evaluate our BQNs by Monte-Carlo sampling, which proves that our approach is able to learn a proper posterior distribution on QNNs. Furthermore, we show that our approach can also be used to learn (ensemble) QNNs by taking maximum a posterior (or sampling from) the posterior distribution.

heading: A RELATED WORK
text: Probabilistic Neural Networks and Bayesian Neural Networks These models consider weights to be random variables and aim to learn their distributions. To further distinguish two families of such models, we call a model Bayesian neural network if the distributions are learned using a prior-posterior framework (i.e. via Bayesian inference) (Soudry et al., 2014; Hernández-Lobato & Adams, 2015; Ghosh et al., 2016; Graves, 2011; Blundell et al., 2015; Shridhar et al., 2019), and otherwise probabilistic neural network (Wang et al., 2016; Shekhovtsov & Flach, 2018; Gast & Roth, 2018). In particular, our work is closely related to natural-parameters networks (NPN) (Wang et al., 2016), which consider both weights and activations to be random variables from exponential family. Since categorical distribution (over quantized values) belongs to exponential family, our BQN can be interpreted as categorical NPN, but we learn the distributions via Bayesian inference.
For Bayesian neural networks, various types of approaches have been proposed to learn the posterior distribution over model parameters.
(1) Sampling-free Assumed Density Filtering (ADF), including EBP (Soudry et al., 2014) and PBP (Hernández-Lobato & Adams, 2015), is an online algorithm which (approximately) updates the posterior distribution by Bayes’ rule for each observation. If the model parameters θ are Gaussian distributed, Minka (2001) shows that the Bayes’ rule can be computed in analytic form based on ∂ log(gn(φ))/∂φ, and EBP Soudry et al. (2014) derives a similar rule for Bernoulli parameters in binary classification. Notice that ADF is compatible to backpropagation:
∂log(gn(φ))
∂φ =
1 gn(φ) · ∂gn(φ) ∂φ (10)
assuming gn(φ) can be (approximately) computed by sampling-free probabilistic propagation as in Section 2. However, this approach has two major limitations: (a) the Bayes’ rule needed to be derived case by case, and analytic rule for most common cases are not known yet. (b) it is not compatible to modern optimization methods (such as SGD or ADAM) as the optimization is solved analytically for each data point, therefore difficult to cope with large-scale models.
(2) Sampling-based Variational inference (SVI), formulates an optimization problem and solves it approximately via stochastic gradient descent (SGD). The most popular method among all is, Stochastic Gradient Variational Bayes (SGVB), which approximates Ln(φ) by the average of multiple samples (Graves, 2011; Blundell et al., 2015; Shridhar et al., 2019). Before each step of learning or prediction, a number of independent samples of the model parameters {θs}Ss=1 are drawn according to the current estimate of Q, i.e. θs ∼ Q, by which the predictive function gn(φ) and the loss Ln(φ) can be approximated by
gn(φ) ≈ 1
S S∑ s=1 Pr[yn|xn, θs] = 1 S S∑ s=1 fn(θs) (11a)
Ln(φ) ≈ 1
S S∑ s=1 log (Pr[yn|xn, θs]) = 1 S S∑ s=1 log (fn(θs)) (11b)
where fn(θ) = Pr[yn|xn, θ] denotes the predictive function given a specific realization θ of the model parameters. The gradients of Ln(φ) can now be approximated as
∂Ln(φ) ∂φ ≈ 1 S S∑ s=1 ∂Ln(φ) ∂fn(θs) · ∂fn(θs) ∂θs · ∂θs ∂φ
(12)
This approach has multiple drawbacks: (a) Repeated sampling suffers from high variance, besides being computationally expensive in both learning and prediction phases; (b) While gn(φ) is differentiable w.r.t. φ, fn(θ) may not be differentiable w.r.t. θ. One such example is quantized neural networks, whose backpropagation is approximated by straight through estimator (Bengio et al.,
2013). (3) The partial derivatives ∂θs/∂φ are difficult to compute with complicated reparameterization tricks (Maddison et al., 2016; Jang et al., 2016).
(3) Deterministic Variational inference (DVI) Our approach is most similar to Wu et al. (2018), which observes that if the underlying model is deterministic, i.e. Pr[h(l+1)|h(l),θ(l)] is a dirac function
Ln(φ) := Eh(L−1)∼P ; θ(L−1)∼Q [ log ( Pr[yn|h(L−1),θ(L−1)] )] (13)
Our approach considers a wider scope of problem settings, where the model could be stochastic, i.e. Pr[h(l+1)|h(l),θ(l)] is an arbitrary function. Furthermore, Wu et al. (2018) considers the case that all parameters θ are Gaussian distributed, whose sampling-free probabilistic propagation requires complicated approximation (Shekhovtsov & Flach, 2018).
Quantized Neural Networks These models can be categorized into two classes: (1) Partially quantized networks, where only weights are discretized (Han et al., 2015; Zhu et al., 2016); (2) Fully quantized networks, where both weights and hidden units are quantized (Courbariaux et al., 2015; Kim & Smaragdis, 2016; Zhou et al., 2016; Rastegari et al., 2016; Hubara et al., 2017). While both classes provide compact size, low-precision neural network models, fully quantized networks further enjoy fast computation provided by specialized bit-wise operations. In general, quantized neural networks are difficult to train due to their non-differentiability. Gradient descent by backpropagation is approximated by either straight-through estimators (Bengio et al., 2013) or probabilistic methods (Esser et al., 2015; Shayer et al., 2017; Peters & Welling, 2018). Unlike these papers, we focus on Bayesian learning of fully quantized networks in this paper. Optimization of quantized neural networks typically requires dedicated loss function, learning scheduling and initialization. For example, Peters & Welling (2018) considers pre-training of a continuous-valued neural network as the initialization. Since our approach considers learning from scratch (with an uniform initialization), the performance could be inferior to prior works in terms of absolute accuracy.
Tensor Networks and Tensorial Neural Networks Tensor networks (TNs) are widely used in numerical analysis (Grasedyck et al., 2013), quantum physiscs (Orús, 2014), and recently machine learning (Cichocki et al., 2016; 2017) to model interactions among multi-dimensional random objects. Various tensorial neural networks (TNNs) (Su et al., 2018; Newman et al., 2018) have been proposed that reduce the size of neural networks by replacing the linear layers with TNs. Recently, (Robeva & Seigal, 2017) points out the duality between probabilistic graphical models (PGMs) and TNs. I.e. there exists a bijection between PGMs and TNs. Our paper advances this line of thinking by connecting hierarchical Bayesian models (e.g. Bayesian neural networks) and hierarchical TNs.

heading: B SUPERVISED LEARNING WITH BAYESIAN NEURAL NETWORKS (BNNS)
text: The problem settings of general Bayesian model and Bayesian neural networks for supervised learning are illustrated in Figures 3a and 3b using graphical models.
1
General Bayesian model Formally, the graphical model in Figure 3a implies the joint distribution of the model parameters θ, the observed dataset D = {(xn, yn)}Nn=1 and any unseen data point
13
(x,y) is factorized as follows:
Pr[x,y,D,θ] = (Pr[y|x,θ]Pr[x] (Pr[D|θ]))Pr[θ] (14)
= ( Pr[y|x,θ]Pr[x] ( N∏ n=1 Pr[yn|xn,θ]Pr[xn] )) Pr[θ] (15)
where Pr[xi]’s and Pr[x] are identically distributed, and so are the conditional distributions Pr[yi|xi,θ]’s and Pr[y|x,θ]. In other words, we assume that (1) the samples (xn, yn)’s (and unseen data point (x,y)) are are identical and independent distributed according to the same data distribution; and (2) xn (or x) and θ together predict the output yn (or y) according to the same conditional distribution. Notice that the factorization above also implies the following equations:
Pr[y|x,D,θ] = Pr[y|x,θ] (16a) Pr[θ|x,D] = Pr[θ|D] (16b)
With these implications, the posterior predictive distribution Pr[y|x,D] can now expanded as: Pr[y|x,D] = ∫ θ Pr[y|x, θ,D]Pr[θ|x,D]dθ = ∫ θ
Pr[y|x, θ]Pr[θ|D]︸ ︷︷ ︸ ≈Q(θ; φ) dθ (17)
where we approximate the posterior distribution Pr[θ|D] by a parameterized distribution Q(θ;φ).
Variational Learning The reason we are learning an approximate posterior Q and not the exact distribution Pr[θ|D] is that for complex models the latter is intractable to compute. The exact posterior Pr[θ|D] generally does not take the form of Q(θ;φ) even if its prior Pr[θ] does. A standard approach to finding a good approximation Q(θ;φ) is variational inference, which finds φ? such that the KL-divergence KL(Q(θ;φ)||Pr[θ|D]) of Q(θ;φ) from Pr[θ|D] is minimized (or alternatively the negative KL-divergence is maximized.)
φ? = arg max φ (−KL(Q(θ;φ)||Pr[θ|D])) (18)
= arg max φ ( − ∫ θ Q(θ;φ) log ( Q(θ;φ) Pr[θ|D] ) dθ ) (19)
where Pr[θ|D] is obtained via standard Bayes’ rule, i.e. Pr[θ|D] = Pr[D|θ]Pr[θ]/Pr[D]. Now we are able to decompose the maximization objective into two terms by plugging the rule into (19):
L(φ) = − ∫ θ Q(θ;φ) log ( Q(θ;φ) · Pr[D] Pr[θ]Pr[D|θ] ) dθ (20)
= N∑ n=1 ∫ θ log (Pr[yn|xn, θ])Q(θ;φ)dθ + ∫ θ Q(θ;φ) log ( Q(θ;φ) Pr[θ] ) dθ + const. (21)
= N∑ n=1 EQ [log (Pr[yn|xn, θ])]︸ ︷︷ ︸ Ln(φ) +KL(Q(θ;φ)||Pr[θ])︸ ︷︷ ︸ R(φ) − log (Pr[D])︸ ︷︷ ︸ const.
(22)
where (1) Ln(φ) is the expected log-likelihood, which reflects the predictive performance of the Bayesian model on the data point (xn, yn); and (2) R(φ) is the KL-divergence between Q(θ;φ) and its prior Pr[θ], which reduces to entropy H(Q) if the prior of θ follows a uniform distribution.
Hierarchical Bayesian Model A Bayesian neural network can be considered as a hierarchical Bayesian model depicted in Figure 3b, which further satisfies the following two assumptions:
Assumption B.1 (Independence of Model Parameters θ(l)). The approximate posterior Q(θ;φ) over the model parameters θ are partitioned into L disjoint and statistically independent layers {θ(l)}L−1l=0 (where each φ(l) parameterizes θ(l) in the l-th layer) such that:
Q(θ;φ) = L−1∏ l=0 Q(θ(l);φ(l)) (23)
Assumption B.2 (Markovianity of Hidden Units h(l)). The hidden variables h = {h(l)}Ll=0 satisfy the Markov property that h(l+1) depends on the input x only through its previous layer h(l):
Pr[h(l+1)|h( : l),θ( : l)] = Pr[h(l+1)|h(l),θ(l)] (24)
where we use short-hand notations h( : l) and θ( : l) to represent the sets of previous layers {h(k)}lk=0 and {θ(k)}lk=0. For consistency, we denote h(0) = x and h(L) = y.
Proof of probabilistic prorogation Based on the two assumptions above, we provide a proof for probabilistic propagation in Equation (1) as follows:
P(h(l+1); ψ(l+1))︷ ︸︸ ︷ Pr[h(l+1)|x] = ∫ θ( : l) Pr[h(l+1)|x, θ( : l)] Q(θ( : l);φ( : l)) dθ( : l) (25)
= ∫ θ( : l) (∫ h(l) Pr[h(l+1)|h(l), θ(l)]Pr[h(l)|x, θ( : l−1)]dh(l) ) Q(θ( : l);φ( : l)) dθ( : l) (26)
=
∫ h(l),θ(l)
Pr[h(l+1)|h(l), θ(l)]Q(θ(l);φ(l))(∫ θ( : l−1) Pr[h(l)|x, θ( : l−1)]Q(θ( : l−1);φ( : l−1))dθ( : l−1) ) dh(l)dθ(l)
(27)
= ∫ 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) (28)

heading: C ALTERNATIVE EVIDENCE LOWER BOUND AND ITS ANALYTIC FORMS

heading: C.1 ALTERNATIVE EVIDENCE LOWER BOUND (PROOF FOR THEOREM 3.1)
text: The steps to prove the inequality (6) almost follow the ones for probabilistic propagation above:
Ln(φ) = EQ [log(Pr[yn|xn,θ])] (29)
= ∫ θ log (Pr[yn|xn, θ])Q(θ;φ)dθ (30)
= ∫ θ log (∫ h(L−1) Pr[yn, h (L−1)|xn, θ]dh(L−1) ) Q(θ;φ)dθ (31)
= ∫ θ log (∫ h(L−1) Pr[yn|h(L−1), θ(L−1)]Pr[h(L−1)|xn, θ(0:L−2)]dh(L−1) ) Q(θ;φ)dθ (32)
≥ ∫ θ (∫ h(L−1) log ( Pr[yn|h(L−1), θ(L−1)] ) Pr[h(L−1)|xn, θ(0:L−1)]dh(L−1) ) Q(θ;φ)dθ (33)
=
∫ h(L−1),θ(L−1) log ( Pr[yn|h(L−1), θ(L−1)] ) Q(θ(L−1);φ(L−1))(∫
θ(0:L−2) Pr[h(L−1)|xn, θ(0:L−2)]Q(θ(0:L−2);φ(0:L−2))dθ(0:L−2)
) dh(L−1)dθ(L−1)
(34)
= ∫ h(L−1),θ(L−1) log ( Pr[yn|h(L−1), θ(L−1)] ) Q(θ(L−1))Pr[h(L−1)|xn]dh(L−1)dθ(L−1) (35)
= Eh(L−1)∼P ; θ(L−1)∼Q [ log ( Pr[yn|h(L−1), θ(L−1)] )] = Ln(φ) (36)
where the key is the Jensen’s inequality EQ [log(·)] ≥ log (EQ [·]) in Equation (33). Notice that if θ(L−1) is not random variable (typical for an output layer), Ln(φ) can be simplified as:
Ln(φ) = ∫ h(L−1) log ( Pr[yn|h(L−1);φ(L−1)] ) P (h(L−1);ψ(L−1))dh(L−1) (37)
where we write Pr[h(L−1)|x] in its parameterized form P (h(L−1);ψ(L−1)). Now, the gradient ∂Ln(φ)/∂φ(L−1) can be obtained by differentiating over Equation (37), while other gradients ∂Ln(φ)/φ( :L−2) further obtained by chain rule:
∂Ln(φ) ∂φ( :L−2) = ∂Ln(φ) ∂ψ(L−1) · ∂ψ (L−1) ∂φ( :L−2) (38)
which requires us to compute ∂Ln(φ)/∂ψ(L−1) and ∂ψ(L−1)/∂φ( :L−2). While ∂Ln(φ)/∂ψ(L−1) can be derived from Equation (37), ∂ψ(L−1)/∂φ( :L−2) can be obtained by backpropagating outputs of the (L − 2)th layer obtained from probabilistic propagation in Equation (1). In other words: since P (h(L−1);ψ(L−1)) is an intermediate step of the forward pass, ψ(L−1) is a function of all parameters from previous layers φ( :L−2), and if each step ψ(l+1) = g(l)(ψ(l), φ(l)) is differentiable w.r.t. ψ(l) and φ(l), the partial derivatives ∂ψ(L−1)/∂φ( :L−2) can be obtained by iterative chain rule.

heading: C.2 SOFTMAX LAYER FOR CLASSIFICATION PROBLEM
text: In this part, we first prove the alternative evidence lower bound (ELBO) for Bayesian neural networks with softmax function as their last layers. Subsequently, we derive the corresponding backpropagation rule for the softmax layer. Finally, we show a method based on Taylor’s expansion to approximately evaluate a softmax layer without Monte Carlo sampling.
Theorem C.1 (Analytic Form of Ln(φ) for Classification). Let h ∈ RK (with K the number of classes) be the pre-activations of a softmax layer (a.k.a. logits), and φ = s ∈ R+ be a scaling factor that adjusts its scale such that Pr[y = c|h, s] = exp(hc/s)/ ∑K k=1 exp(hk/s). Suppose the logits {hk}Kk=1 are pairwise independent (which holds under mean-field approximation) and hk follows a Gaussian distribution hk ∼ N (µk, νk) (therefore ψ = {µk, νk}Kk=1) and s is a deterministic parameter. Then Ln(φ) can be further upper bound by the following analytic form:
Ln(φ) ≥ µc s − log ( K∑ k=1 exp (µk s + νk 2s2 )) , L̂(φ) (39)
Proof. The lower bound follows by plugging Pr[y|h, s] and Pr[hk|x] into Equation (6). Ln(φ) = ∫ h log (Pr[yn = c|h; s])Pr[h|x]dh (40)
= ∫ h ( hc s − log ( K∑ k=1 exp ( hk s )))( K∏ k=1 Pr[hk|x] ) dh (41)
= 1
s ∫ hc hcPr[hc|xn]dhc − ∫ h log ( K∑ k=1 exp ( hk s ))( K∏ k=1 Pr[hk|x] ) dh (42)
= µc s − ∫ h log ( K∑ k=1 exp ( hk s ))( K∏ k=1 Pr[hk|x] ) dh (43)
≥ µc s − log (∫ h K∑ k=1 exp ( hk s )( K∏ k=1 Pr[hk|x] ) dh ) (44)
= µc s − log ( K∑ k=1 ∫ hk exp ( hk s ) Pr[hk|x]dhk ) (45)
= µc s − log ( K∑ k=1 ∫ hk exp ( hk s ) · 1√ 2πνk exp ( − (hk − µk) 2 2νk ) dhk ) (46)
= µc s − log ( K∑ k=1 exp (µk s + νk 2s2 )) = L̂(φ) (47)
where the last equation follows∫ hk exp ( hk s ) · 1√ 2πνk exp ( − (hk − µk) 2 2νk ) dhk (48)
= ∫ hk 1√ 2πνk exp ( −h 2 k − 2(µk + νk/s)hk + µ2k 2νk ) dhk (49)
= ∫ hk 1√ 2πνk exp ( − (hk − (µk + νk)) 2 2νk ) dhk︸ ︷︷ ︸ · exp (µk s + νk 2s2 ) (50)
where the under-braced term is unity since it takes the form of Gaussian distribution.
From Equation (43) to (44), we use the Jensen’s inequality to achieve a lower bound for integral of log-sum-exp. The bound can be tighten with advanced techniques in Khan (2012).
Derivatives ofLn(φ) in (39) To use probabilistic backpropagation to obtain the gradients w.r.t. the parameters from previous layers, we first need to obtain the derivatives w.r.t. ψ(L−1) = {µk, νk}Kk=1.
∂L̂n(φ) ∂µk = −1 s
( exp ( µk/s+ νk/2s 2 )∑K
k=1 exp (µk/s+ νk/2s 2) − 1[k = c]
) (51a)
∂L̂n(φ) ∂νk = − 1 2s2
( exp ( µk/s+ νk/2s 2 )∑K
k=1 exp (µk/s+ νk/2s 2)
) (51b)
Furthermore, the scale s can be (optionally) updated along with other parameters using the gradient
∂L̂n(φ) ∂s = −µc s2 +
∑K k=1 ( µk/s 2 + νk/s 3 ) exp ( µk/s+ νk/2s 2 )∑K
k=1 exp (µk/s+ νk/2s 2)
(52)
Prediction with Softmax Layer Once we learn the parameters for the Bayesian neural network, in principle we can compute the predictive distribution of y by evaluating the following equation:
Pr[y = c|x] = ∫ h Pr[y = c|h, s]Pr[h|x]dh = ∫ h `c(h)Pr[h|x]dh (53)
(Mean-field assumption) = ∫ h1 · · · ∫ hK `c(h) ( K∏ k=1 Pr[hk|x] ) dh1 · · · dhk (54)
where we denote the softmax function as `c(h) = exp(hc/s)/[ ∑ k exp(hk/s)]. Unfortunately, the equation above can not be computed in closed form. The most straight-forward work-around is to approximate the integral by Monte Carlo sampling: for each hk we draw S samples {hsk}Ss=1 independently and compute the prediction:
Pr[y = c|x] ≈ 1 S S∑ s=1 `c(h s), ∀c ∈ [K] (55)
Despite its conceptual simplicity, Monte Carlo method suffers from expensive computation and high variance in estimation. Instead, we propose an economical estimate based on Taylor’s expansion. First, we expand the function `c(h) by Taylor’s series at the point µ (up to the second order):
`c(h) = `c(µ) + [ ∂`c ∂h (µ) ]> (h− µ) + 1 2 (h− µ)> [ ∂2`c ∂h2 (µ) ] (h− µ) +O ( ‖h− c‖3 ) (56)
= `c(µ) + K∑ k=1 [ ∂`c ∂hk (µ) ] (hk − µk) + K∑ i=1 K∑ j=1 [ ∂2`c ∂hihj (µ) ] (hi − µi)(hj − µj) +O ( ‖h− µ‖3 ) (57)
Before we derive the forms of these derivatives, we first show the terms of odd orders do not contribute to the expectation. For example, if `c(h) is approximated by its first two terms (i.e. a linear function), Equation (54) can be written as
Pr[y = c|x] ≈ ∫ h1 · · · ∫ hK ( `c(µ) + K∑ k=1 [ ∂`c ∂hk (µ) ] (hk − µk) )( K∏ k=1 Pr[hk|x] ) dh1 · · · dhk
(58)
= `c(µ) + K∑ k=1 [ ∂`c ∂hk (µ) ](∫ hk (hk − µk)Pr[hk|x]dhk ) = `c(µ) (59)
where the second term is zero by the symmetry of Pr[hk|x] around µk (or simply the definition of µk’s). Therefore, the first-order approximation results exactly in a (deterministic) softmax function of the mean vector µ. In order to incorporate the variance into the approximation, we will need to derive the exact forms of the derivatives of `c(h). Specifically, the first-order derivatives are obtained from the definition of `c(h).
∂`c ∂hc (h) = 1 s · exp (hc/s)− exp (2hc/s)(∑K k=1 exp (hk/s) )2 = 1s (`c(h)− `2c(h)) (60a)
∂`c ∂hk (h) = −1 s · exp (hc/s) · exp (hk/s)(∑K k=1 exp (hk/s) )2 = −1s `c(h)`k(h), ∀k 6= c (60b)
and subsequently the second-order derivatives from the first ones:
∂2`c ∂h2c (h) = 1 s ( ∂`c ∂hc (h)− 2`c(h) ∂`c ∂hc (h) ) = 1 s2 ( `c(h)− 3`2c(h) + 2`3c(h) ) (61a)
∂2`c ∂h2k (h) = −1 s ( ∂`c ∂hc (h)`k(h) + `c(h) ∂`k ∂hc (h) ) = 1 s2 ( −`c(h)`k(h) + 2`2c(h)`k(h) ) , ∀k 6= c
(61b)
with these derivatives we can compute the second-order approximation as Pr[y = c|x] ≈ ∫ h1,·,hK `c(µ) + 1 2 K∑ i=1 K∑ j=1 ∂2`c ∂µiµj (µ)(hi − µi)(hj − µj) ( K∏ k=1 Pr[hk|x] ) dh1 · · · dhK
(62)
= `c(µ) + 1
2 ∂2`c ∂µ2c (µ) ∫ hc (hc − µc)2Pr[hc|x]dhc + 1 2 ∑ k 6=c ∂2`c ∂µ2k (µ) ∫ hk (hk − µk)2Pr[hk|x]dhk
(63)
= `c(µ) + 1 2s2 ( `c(µ)− 3`2c(µ) + 2`3c(µ) ) νc + 1 2s2 ∑ k 6=c ( −`c(µ)`k(µ) + 2`2c(µ)`k(µ) ) νk
(64)
= `c(µ) + 1 2s2 ( `c(µ)− 2`2c(µ) )( νc − K∑ k=1 `k(µ)νk ) (65)
The equation above can be further written in vector form as:
Pr[y|x] ≈ `(µ) + 1 2s2
( `(µ)− `(µ)◦2 ) ◦ ( ν − `(µ)>ν ) (66)

heading: C.3 GAUSSIAN OUTPUT LAYER FOR REGRESSION PROBLEM
text: In this part, we develop an alternative evidence lower bound (ELBO) for Bayesian neural networks with Gaussian output layers, and derive the corresponding gradients for backpropagation. Despite the difficulty to obtain an analytical predictive distribution for the output, we show that its central moments can be easily computed given the learned parameters.
Theorem C.2 (Analytic Form of Ln(φ) for Regression). Let h ∈ RI be the output of last hidden layer (with I the number of hidden units), and φ = (w, s) ∈ RI ×R+ be the parameters that define the predictive distribution over output y as
Pr[y|h;w, s] = 1√ 2πs exp
( − (y − w >h)2
2s
) (67)
Suppose the hidden units {hk}Kk=1 are pairwise independent (which holds under mean-field approximation), and each hi has mean µi and variance νi, then Ln(φ) takes an analytic form:
Ln(φ) = − (y − w>µ)2 + (w◦2)>ν
2s − log(2πs) 2 (68)
where (w◦2)i = w2i and µ = [µ1, · · · , µI ]> ∈ RI and ν = [ν1, · · · , νI ]> ∈ RI are vectors of mean and variance of the hidden units h.
Proof. The Equation (68) is obtained by plugging Pr[y|h;w, s] into Equation (6).
Ln(φ) = ∑ h1 · · · ∑ hI log (Pr[y|h1, · · · , hI ;w, s]) ( I∏ i=1 Pr[hi|xn] ) (69)
= − ∑ h1 · · · ∑ hI
 ( y − ∑I i=1 wihi )2 2s + log(2πs) 2 ( I∏ i=1 Pr[hi|xn] ) (70)
= − 1 2s ∑ h1 · · · ∑ hI
( y −
I∑ i=1 wihi )2( I∏ i=1 Pr[hi|xn] ) − log(2πs) 2 (71)
where the long summation in the first term can be further simplified with notations of µ and ν:
∑ h1 · · · ∑ hI
( y −
I∑ i=1 wihi )2( I∏ i=1 Pr[hi|xn] ) (72)
= ∑ h1 · · · ∑ hI y2 − 2y I∑ i=1 wihi + I∑ i=1 w2i h 2 i + I∑ j=1 ∑ k 6=j wjwkhjhk ( I∏ i=1 Pr[hi|xn] ) (73)
=y2 − 2y I∑ i=1 wi (∑ hi hiPr[hi|x] ) + I∑ i=1 w2i (∑ hi h2iPr[hi|xn] )
+ I∑ j=1 ∑ k 6=j wjwk ∑ hj hjPr[hj |xn] (∑ hk hkPr[hk|xn] ) (74)
=y2 − 2y I∑ i=1 wiµi + I∑ i=1 w2i (µ 2 i + νi) + I∑ j=1 ∑ k 6=j wjwkµjµk (75) =y2 − 2y I∑ i=1 wiµi + I∑ i=1 w2i νi +  I∑ j=1 wjµj ( I∑ k=1 wkµk ) (76)
=y2 − 2y w>µ+ (w◦2)>ν + ( w>µ )2 (77)
=(y − w>µ)2 + (w◦2)>ν (78)
where w◦2 denotes element-wise square, i.e. w◦2 = [w21, · · · , w2I ]>.
Derivatives of Ln(φ) in Equation (68) It is not difficult to show that the gradient of Ln(φ) can be backpropagated through the last layer. by computing derivatives of Ln(φ) w.r.t. µ and ν:
∂Ln(φ) ∂µ = − (y − w >µ)w s
(79a)
∂Ln(φ) ∂ν = −w ◦2 2s (79b)
Furthermore, the parameters {w, s} can be updated along with other parameters with their gradients:
∂Ln(φ) ∂w = − (y − w >µ)µ+ (w ◦ ν) s (80a)
∂Ln(φ) ∂s = − 1 2s + (y − w>µ)2 + (w◦2)>ν 2s2 (80b)
Prediction with Gaussian Layer Once we determine the parameters for the last layer, in principle we can compute the predictive distribution Pr[y|x] for the output y given the input x according to
Pr[y|x] = ∑ h Pr[y|h;w, s]Pr[h|x] = ∑ h1 · · · ∑ hI Pr[y|h;w, s] ( I∏ i=1 Pr[hi|x] )
= ∑ h1 · · · ∑ hI 1√ 2πs exp
− ( y − ∑I i=1 wihi )2 2s ( I∏ i=1 Pr[hi|x] ) (81)
Unfortunately, exact computation of the equation above for arbitrary output value y is intractable in general. However, the central moments of the predictive distribution Pr[y|x] are easily evaluated. Consider we interpret the prediction as y = w>h + , where ∼ N (0, s), its mean and variance can be easily computed as
E [y|x] = w>E [h] = w>µ (82a) V [y|x] = (w◦2)>V [h] + V [ ] = (w◦2)>ν + s (82b)
Furthermore, if we denote the (normalized) skewness and kurtosis of hi as γi and κi: γi = E [ (hi − µi)3|x ] /ν 3/2 i = ∑ hi (hi − µi)3Pr[hi|x]/ν3/2i (83a)
κi = E [ (hi − µi)4|x ] /ν2i = ∑ hi (hi − µi)4Pr[hi|x]/ν2i (83b)
Then the (normalized) skewness and kurtosis of the prediction y are also easily computed with the vectors of γ = [γ1, · · · , γI ]> ∈ RI and κ = [κ1, · · · , κI ] ∈ RI .
γ[y|x] = E [ (y − w>µ)3|x ] V [y|x]3/2 = (w◦3)>(γ ◦ ν◦3/2) [(w◦2)>ν + s]3/2
(84a)
κ[y|x] = E [ (y − w>µ)4|x ] V [y|x]2 = (w◦4)>(κ ◦ ν◦2) + s(w◦2)>ν [(w◦2)>ν + s]2 (84b)

heading: D PROBABILISTIC PROPAGATION IN BAYESIAN QUANTIZED NETWORKS
text: In this section, we present fast(er) algorithms for sampling-free probabilistic propagation (i.e. evaluating Equation (8)). According to Section 4, we divide this section into three parts, each part for a specific range of fan-in numbers E.

heading: D.1 SMALL FAN-IN LAYERS: DIRECT TENSOR CONTRACTION
text: If E is small, tensor contraction in Equation (8) is immediately applicable. Representative layers of small E are shortcut layer (a.k.a. skip-connection) and what we name as depth-wise layers.
Shortcut Layer With a skip connection, the output h(l+1) is an addition of two previous layers h(l) and h(m). Therefore and the distribution of h(l+1) can be directly computed as
P (h (l+1) i ;ψ (l+1) i ) = ∑ h (l) i ,h (m) i δ[h (l+1) i = h (l) i + h (m) i ] P (h (l) i ;ψ (l) i ) P (h (m) i ;ψ (m) i ) (85)
Depth-wise Layers In a depth-wise layer, each output unit h(l+1)i is a transformation (parameterized by θ(l)i ) of its corresponding input h (l) i , i.e.
P (h (l+1) i ;ψ (l+1) i ) = ∑ h (l) i ,θ (l) i Pr[h (l+1) i |h (l) i , θ (l) i ] Q(θ (l) i ;φ (m) i ) P (h (l) i ;ψ (l) i ) (86)
Depth-wise layers include dropout layers (where θ(l) are dropout rates), nonlinear layers (where θ(l) are threshold values) or element-wise product layers (where θ(l) are the weights). For both shortcut and depth-wise layers, the time complexity is O(JD2) since E <= 2.

heading: D.2 MEDIUM FAN-IN LAYERS: DISCRETE FOURIER TRANSFORM
text: In neural networks, representative layers with medium fan-in number E are pooling layers, where each output unit depends on a medium number of input units. Typically, the special structure of pooling layers allows for faster algorithm than computing Equation (8) directly.
Max and Probabilistic Pooling For each output, (1) a max pooling layer picks the maximum value from corresponding inputs, i.e. h(l+1)j = maxi∈I(j) h (l) i , while (2) a probabilistic pooling layer selects the value the inputs following a categorical distribution, i.e. Pr[h(l+1)j = h (l) i ] = θi. For both cases, the predictive distribution of h(l+1)j can be computed as
Max: P (h(l+1)j ≤ q) = ∏ i∈I(j) P (h (l) i ≤ q) (87)
Prob: P (h(l+1)j = q) = ∑ i∈I(j) θiP (h (l) i = q) (88)
where P (h(l)i ≤ q) is the culminative mass function of P . Complexities for both layers are O(ID).
Average Pooling and Depth-wise Convolutional Layer Both layers require additions of a medium number of inputs. We prove a convolution theorem for discrete random variables and show that discrete Fourier transform (DFT) (with fast Fourier transform (FFT)) can accelerate the additive computation. We also derive its backpropagation rule for compatibility of gradient-based learning.
Theorem D.1 (Fast summation via discrete Fourier transform). Suppose ui take values in {bi, bi + 1, . . . , Bi} between integers bi and Bi, then the summation v = ∑E i=1 ui takes values
between b and B, where b = ∑E i=1 bi and B = ∑E i=1Bi. Let C
v , Cui be the discrete Fourier transforms of P v , Pui respectively, i.e.
Cv(f) = B∑ v=b P v(v) exp(−j2π(v − b)f/(B − b+ 1)) (89a)
Cui(f) = Bi∑ ui=bi Pui(ui) exp(−j2π(ui − bi)f/(Bi − bi + 1)) (89b)
Then Cv(f) is the element-wise product of all Fourier transforms Cui(f), i.e.
Cv(f) = E∏ i=1 Cui(f),∀f (90)
Proof. We only prove the theorem for two discrete random variable, and the extension to multiple variables can be proved using induction. Now consider u1 ∈ [b1, B1], u2 ∈ [b2, B2] and their sum v = u1 + u2 ∈ [b, B], where b = b1 + b2 and B = B1 + B2. Denote the probability vectors of u1, u2 and v as P1 ∈ 4B1−b1 , P2 ∈ 4B2−b2 and P ∈ 4B−b respectively, then the entries in P are computed with P1 and P2 by standard convolution as follows:
P (v) = B1∑ u1=b1 P1(u1)P2(v − u1) = B2∑ u2=b2 P1(v − u2)P2(u2), ∀v ∈ {b, · · · , B} (91)
The relation above is usually denoted as P = P1 ∗ P2, where ∗ is the symbol for convolution. Now define the characteristic functions C, C1, and C2 as the discrete Fourier transform (DFT) of the probability vectors P , P1 and P2 respectively:
C(f) = B∑ v=b P (v) exp ( −j2π R (v − b)f ) , f ∈ [R] (92a)
Ci(f) = Bi∑ ui=bi Pi(ui) exp ( −j2π R (ui − bi)f ) , f ∈ [R] (92b)
where R controls the resolution of the Fourier transform (typically chosen as R = B − b + 1, i.e. the range of possible values). In this case, the characteristic functions are complex vectors of same length R, i.e. C,C1, C2 ∈ CR, and we denote the (functional) mappings as C = F(P ) and Ci = Fi(Pi). Given a characteristic function, its original probability vector can be recovered by inverse discrete Fourier transform (IDFT):
P (v) = 1
R R−1∑ f=0 C(f) exp ( j 2π R (v − b)f ) , ∀v ∈ {b, · · · , B} (93a)
Pi(ui) = 1
R R−1∑ f=0 Ci(f) exp ( j 2π R (ui − bi)f ) , ∀ui ∈ {bi, · · · , Bi} (93b)
which we denote the inverse mapping as P = F−1(C) and Pi = F−1i (Ci). Now we plug the convolution in Equation (91) into the characteristic function C(f) in (92a) and rearrange accordingly:
C(f) = B∑ v=b
( B1∑
u1=b1
P1(u1)P2(v − u1) ) exp ( −j2π
R (v − b)f
) (94)
(Let u2 = v − u1) = B1∑
u1=b1 B2∑ u2=b2 P1(u1)P2(u2) exp ( −j2π R (u1 + u2 − b)f ) (95)
(Since b = b1 + b2) =
[ B1∑
u1=b1
P1(u1) exp
( −j2π
R (u1 − b1)f )] [
B2∑ u2=b2 P2(u2) exp ( −j2π R (u2 − b2)f )] (96) = C1(f) · C2(f) (97)
The equation above can therefore be written asC = C1◦C2, where we use ◦ to denote element-wise product. Thus, we have shown summation of discrete random variables corresponds to element-wise product of their characteristic functions.
With the theorem, addition of E discrete random variables can be computed efficiently as follows
P v = Pu1 ∗ Pu2 ∗ · · · ∗ PuE (98) = F−1 (F (Pu1) ◦ F (Pu2) ◦ · · · ◦ F (PuE )) (99)
where F denotes the Fourier transforms in Equations (93a) and (93b). If FFT is used in computing all DFT, the computational complexity of Equation (99) is O(ER logR) = O(E2D log(ED)) (since R = O(ED)), compared to O(DE) with direct tensor contraction.
Backpropagation When fast Fourier transform is used to accelerate additions in Bayesian quantized network, we need to derive the corresponding backpropagation rule, i.e. equations that relate ∂L/∂P to {∂L/∂Pi}Ii=1. For this purpose, we break the computation in Equation (99) into three steps, and compute the derivative for each of these steps.
Ci = Fi(Pi) =⇒ ∂L ∂Pi = R · F−1i ( ∂L ∂Ci ) (100a)
C = C1 ◦ · · · ◦ CI =⇒ ∂L ∂Ci = C Ci ◦ ∂L ∂C
(100b)
P = F−1(C) =⇒ ∂L ∂C = R−1 · F ( ∂L ∂P ) (100c)
where in (100b) we use C/Ci to denote element-wise division. Since Pi lies into real domain, we need to project the gradients back to real number in (100c). Putting all steps together:
∂L ∂Pi
= < { F−1i ( C
Ci ◦ F ( ∂L ∂P ))} ,∀i ∈ [I] (101)

heading: D.3 LARGE FAN-IN LAYERS: LYAPUNOV CENTRAL LIMIT APPROXIMATION
text: In this part, we show that Lyapunov central limit approximation (Lyapunov CLT) accelerates probabilistic propagation in linear layers. For simplicity, we consider fully-connected layer in the derivations, but the results can be easily extended to types of convolutional layers. We conclude this part by deriving the corresponding backpropagation rules for the algorithm.
Linear Layers Linear layers (followed by a nonlinear transformations σ(·)) are the most important building blocks in neural networks, which include fully-connected and convolutional layers. A linear layer is parameterized by a set of vectors θ(l)’s, and maps h(l) ∈ RI to h(l+1) ∈ RJ as
h (l+1) j = σ  ∑ i∈I(j) θ (l) ji · h (l) i  = σ  ∑ i∈I(j) u (l) ji  = σ (v(l+1)j ) (102) where u(l)ji = θ (l) ji ·h (l) i and v (l+1) j = ∑ i∈I(j) u (l) ji . The key difficulty here is to compute the distribution of v(l+1)j from the ones of {u (l) ji }Ii=1, i.e. addition of a large number of random variables. Theorem D.2 (Fast summation via Lyapunov Central Limit Theorem). Let vj = σ(ṽj) = σ( ∑ i∈I(j) θjiui) be an activation of a linear layer followed by nonlinearity σ. Suppose both inputs {ui}i∈I and parameters {θji}i∈I(j) have bounded variance, then for sufficiently large |I(j)|, the distribution of ṽj converges to a Gaussian distribution N (µ̃j , ν̃j) with mean and variance as
µ̃j = I∑ i=1 mjiµi (103a)
ν̃j = I∑ i=1 m2jiνi + vjiµ 2 i + vjiνi (103b)
where mji = E [θji], vji = V [θji] and µi = E [ui], νi = V [ui]. And if the nonlinear transform σ is a sign function, each activation vj follows a Bernoulli distribution P (vj = 1) = Φ(µ̃j/ √ ν̃j), where Φ is the culminative probability function of a standard Gaussian distribution N (0, 1).
Proof. The proof follows directly from the definitions of mean and variance:
µ̃j = E [ I∑ i=1 θji hi ] = I∑ i=1 E [θji hi] (104)
= I∑ i=1 E [θji]E [hi] = I∑ i=1 mjiµi (105)
ν̃j = V [ I∑ i=1 θji hi ] = I∑ i=1 V [θji hi] (106)
= I∑ i=1 ( E [ θ2ji ] E [ h2i ] − E [ θ2ji ] E [ h2i ])
(107)
= I∑ i=1 [( m2ji + vi ) ( µ2i + νji ) −m2jiµ2i ] (108)
= I∑ i=1 ( m2jiνi + vjiµ 2 i + vjiνi ) (109)
For fully-connected layers, these two equations can be concisely written in matrix forms:
µ̃ = Mµ (110a) ν̃ = ( M◦2 ) ν + V ( µ◦2 + ν ) (110b)
where M◦2 and µ◦2 are element-wise square of M and µ respectively.
Backpropagation With matrix forms, the backpropagation rules that relate ∂L/∂ψ̃(l+1) = {∂L/∂µ̃, ∂L/∂ν̃} to ∂L/∂φ(l) = {∂L/∂M, ∂L/∂V } and ∂L/∂ψ(l) = {∂L/∂µ, ∂L/∂ν} can be derived with matrix calculus.
∂L ∂M = ( ∂L ∂µ̃ ) µ> + 2M ◦ [( ∂L ∂ν̃ ) ν> ]
(111a)
∂L ∂V = ( ∂L ∂ν̃ )( µ◦2 )>
(111b)
∂L ∂µ = M> ( ∂L ∂µ̃ ) + 2µ ◦ [ V > ( ∂L ∂ν̃ )] (111c)
∂L ∂ν
= ( M◦2 )>(∂L ∂ν̃ ) (111d)
Notice that these equations do not take into account the fact that V implicitly defined with M (i.e. vji is defined uponmji). Therefore, we adjust the backpropagation rule for the probabilities: denote Qji(d) = Q(θji = Q(d);φ(l)ji ), then the backpropagation rule can be written in matrix form as
∂L ∂Q(d) = ( ∂L ∂M + ∂L ∂V · ∂V ∂M ) ∂M ∂Q(d) + ∂L ∂V · ∂ν ∂Q(d)
(112)
= Q(d) · ∂L ∂M + 2(Q(d)−M) ◦ ( ∂L ∂V ) (113)
Lastly, we derive the backpropagation rules for sign activations. Let pj denote the probability that the hidden unit vj is activated as pj = Pr[vj = 1|x], ∂L/pj relates to {∂L/∂µ̃j , ∂L/∂ν̃j} as:
∂pj ∂µ̃j = N ( µ̃j√ ν̃j ) (114a)
∂pj ∂ν̃j = − 1 2ν̃
3/2 j
· N ( µ̃j√ ν̃j ) (114b)

heading: E SUPPLEMENTARY MATERIAL FOR EXPERIMENTS

heading: E.1 NETWORK ARCHIECTURES
text: (1) For MNIST, Fashion-MNIST and KMNIST, we evaluate our models on both MLP and CNN. For MLP, we use a 3-layers network with 512 units in the first layer and 256 units in the second; and
for CNN, we use a 4-layers network with two 5× 5 convolutional layers with 64 channels followed by 2× 2 average pooling, and two fully-connected layers with 1024 hidden units. (2) For CIFAR10, we evaluate our models on a smaller version of VGG (Peters & Welling, 2018), which consists of 6 convolutional layers and 2 fully-connected layers: 2 x 128C3 – MP2 – 2 x 256C3 – MP2 – 2 x 512C3 – MP2 – 1024FC – SM10.
E.2 MORE RESULTS FOR MULTI-LAYER PERCEPTRON (MLP)

heading: E.3 REGRESSION ON BOSTON HOUSING DATASET
text: In this part, we evaluate our proposed BQN on Boston housing dataset, a regression benchmark widely used in testing Bayesian neural networks (Hernández-Lobato & Adams, 2015; Ghosh et al., 2016) and Probabilistic neural networks (Wang et al., 2016). The dataset consists of 456 training and 50 test samples, each sample has 13 features as input and a scalar (housing) price as output. Following Hernández-Lobato & Adams (2015); Ghosh et al. (2016); Wang et al. (2016), we train a two-layers network with 50 hidden units, and report the performance in terms of root mean square error (RMSE) in Table 4. The results show that our BQN achieves lower RMSE compared to other models trained in a probabilistic/Bayesian way.

Review summaries:
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.