Title: SAMPLING-FREE LEARNING OF BAYESIAN QUANTIZED NEURAL NETWORKS

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

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 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 differentiable lower bound to ELBO, which addresses challenges (2), ( 3) and (4) above. • BQNs leverage efficient tensor operations for probabilistic propagation, further addressing challenge (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; Herná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; Orú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 = {(x n , y n )} N

n=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; Herná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 L n (φ) + R(φ)

where

L n (φ) = E Q [log Pr[y n |x n , θ]] and R(φ) = E Q [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:

EQUATION

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

g n (φ) = E Q [Pr[y n |x n , θ]] = θ Pr[y n |x n , θ]Q(θ; φ)dθ (4a) L n (φ) = E Q [log(Pr[y n |x n , θ])] = θ log (Pr[y n |x n , θ]) Q(θ; φ)dθ (4b)

Although L n (φ) is a function of φ, it is not an explicit function of g n (φ). Consequently, the chain rule in Equation ( 3) on which backpropagation is based is not directly applicable. 

3. PROPOSED LEARNING METHOD FOR BAYESIAN NEURAL NETWORKS:

EQUATION

In Appendix C.1, we proved one feasible L n (φ) which only depends on second last output h (L-1) .

Theorem 3.1 (Alternative Evidence Lower Bound). Define each term L n (φ) in L(φ) as

L n (φ) := E h (L-1) ∼P ; θ (L-1) ∼Q log Pr[y n |h (L-1) , θ (L-1) ] (6) then L n (φ) is a lower bound of L n (φ), i.e. L n (φ) ≤ L n (φ). The equality L n (φ) = L n (φ) holds if h (L-1

) is deterministic given input x and all parameters before the last layer θ ( : L-2) .

Analytic Forms of L n (φ). 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 that L n (φ) 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 L n (φ) involves the last layer only, we omit the superscripts/subsripts of h (L-1) , ψ (L-1) , φ (L-1) , x n , y n , and denote them as h, ψ, φ, x, y . Theorem 3.2 (Analytic Form of L n (φ) for Classification). Let h ∈ R K (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(h c /s)/ K k=1 exp(h k /s). Suppose the logits {h k } K k=1 are pairwise independent (which holds under mean-field approximation) and h k follows a Gaussian distribution

h k ∼ N (µ k , ν k ) (therefore ψ = {µ k , ν k } K k=1 ) and s is a deterministic parameter. Then L n (φ) is further lower bounded as L n (φ) ≥ µc s -log K k=1 exp µ k s + ν k 2s 2

.

The regression case and proofs for both layers are deferred to Appendix C.

4. BAYESIAN QUANTIZED NETWORKS (BQNS):

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

4.1. PROBABILISTIC PROPAGATION AS TENSOR CONTRACTIONS:

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) ∈ Q I , θ (l) ∈ Q K , and h (l+1) ∈ Q J , and their distributions are characterized by P (h (l) ;

EQUATION

Pr[h (l+1) |h (l) , θ (l) ] ∈ R D J ×D I ×D K 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(D IJK ), which is intractable for quantized neural network of any reasonable size. We thus turn to approximations.

4.2. APPROXIMATE PROPAGATION VIA RANK-1 TENSOR CP DECOMPOSITION:

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 ψ

EQUATION

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 k } according to the connectivity pattern in the layer. We denote the set of dependent input units and parameters for h , and define the fan-in number E for the layer as max j I

(l+1) j + M (l+1) j .

Complexity of Approximate Propagation. The approximate propagation reduces the computational complexity from O(D IJK ) to O(JD E ), 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).

4.3. FAST ALGORITHMS FOR APPROXIMATE PROPAGATION:

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(JD E ) 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(JE 2 D log(ED)).

Large Fan-in Layers: Lyapunov Central Limit Theorem. In a typical linear layer, the fan-in E 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 numbers E, 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.

5. EXPERIMENTS:

In this section, we demonstrate the effectiveness of BQNs on the MNIST, Fashion-MNIST, KM-NIST 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. Table 1 : Comparison of performance of BQNs against the baseline E-QNN. Each E-QNN is an ensemble of 10 networks, which are trained individually and but make predictions jointly. We report both NLL (which accounts for prediction uncertainty) and 0-1 test error (which doesn't account for prediction uncertainty). All the numbers are averages over 10 runs with different seeds, the standard deviation are exhibited following the ± sign.

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:

EQUATION

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 (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 over Q(θ) 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.

5.1. ANALYSIS OF RESULTS:

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 The performance gap between our analytical inference and the Monte Carlo Sampling is displayed.

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

Appendix: Sampling-Free Learning of Bayesian Quantized Neural Networks

A RELATED WORK:

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; Herná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) 

EQUATION

assuming g n (φ) 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 L n (φ) 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 } S s=1 are drawn according to the current estimate of Q, i.e. θ s ∼ Q, by which the predictive function g n (φ) and the loss L n (φ) can be approximated by

EQUATION

where f n (θ) = Pr[y n |x n , θ] denotes the predictive function given a specific realization θ of the model parameters. The gradients of L n (φ) can now be approximated as

EQUATION

This approach has multiple drawbacks: (a) Repeated sampling suffers from high variance, besides being computationally expensive in both learning and prediction phases; (b) While g n (φ) is differentiable w.r.t. φ, f n (θ) 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 L n (φ) := E h (L-1) ∼P ; θ (L-1) ∼Q log Pr[y n |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 (Orú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.

B SUPERVISED LEARNING WITH BAYESIAN NEURAL NETWORKS (BNNS):

The problem settings of general Bayesian model and Bayesian neural networks for supervised learning are illustrated in Figures 3a and 3b using graphical models.

... ... ... General Bayesian model Formally, the graphical model in Figure 3a implies the joint distribution of the model parameters θ, the observed dataset D = {(x n , y n )} N n=1 and any unseen data point (x, y) is factorized as follows:

x θ (0) h (1) θ (1) h (2) θ (L-1)

EQUATION

where Pr[x i ]'s and Pr[x] are identically distributed, and so are the conditional distributions Pr[y i |x i , θ]'s and Pr [y|x, θ] . In other words, we assume that (1) the samples (x n , y n )'s (and unseen data point (x, y)) are are identical and independent distributed according to the same data distribution; and (2) x n (or x) and θ together predict the output y n (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:

EQUATION

)

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

EQUATION

)

where 

EQUATION

)

where (1) L n (φ) is the expected log-likelihood, which reflects the predictive performance of the Bayesian model on the data point (x n , y n ); 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: l) ). The approximate posterior Q(θ; φ) over the model parameters θ are partitioned into L disjoint and statistically independent layers {θ (l) } L-1 l=0 (where each φ (l) parameterizes θ (l) in the l-th layer) such that:

Assumption B.1 (Independence of Model Parameters θ (

EQUATION

)

Assumption B.2 (Markovianity of Hidden Units h (l) ). The hidden variables h = {h (l) } L l=0 satisfy the Markov property that h (l+1) depends on the input x only through its previous layer h (l) :

EQUATION

)

where we use short-hand notations h ( : l) and θ ( : l) to represent the sets of previous layers {h (k) } l k=0 and {θ (k) } l k=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)

EQUATION

C ALTERNATIVE EVIDENCE LOWER BOUND AND ITS ANALYTIC FORMS:

C.1 ALTERNATIVE EVIDENCE LOWER BOUND (PROOF FOR THEOREM 3.1)

The steps to prove the inequality (6) almost follow the ones for probabilistic propagation above:

L n (φ) = E Q [log(Pr[y n |x n , θ])] (29) = θ log (Pr[y n |x n , θ]) Q(θ; φ)dθ (30) = θ log h (L-1) Pr[y n , h (L-1) |x n , θ]dh (L-1) Q(θ; φ)dθ (31) = θ log h (L-1) Pr[y n |h (L-1) , θ (L-1) ]Pr[h (L-1) |x n , θ (0:L-2) ]dh (L-1) Q(θ; φ)dθ (32) ≥ θ h (L-1) log Pr[y n |h (L-1) , θ (L-1) ] Pr[h (L-1) |x n , θ (0:L-1) ]dh (L-1) Q(θ; φ)dθ (33) = h (L-1) ,θ (L-1) log Pr[y n |h (L-1) , θ (L-1) ] Q(θ (L-1) ; φ (L-1) ) θ (0:L-2)

Pr[h (L-1) |x n , θ (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[y n |h (L-1) , θ (L-1) ] Q(θ (L-1) )Pr[h (L-1) |x n ]dh (L-1) dθ (L-1) (35) = E h (L-1) ∼P ; θ (L-1) ∼Q log Pr[y n |h (L-1) , θ (L-1) ] = L n (φ) (36)

where the key is the Jensen's inequality

E Q [log(•)] ≥ log (E Q [•]

) in Equation ( 33). Notice that if θ (L-1) is not random variable (typical for an output layer), L n (φ) can be simplified as:

L n (φ) = h (L-1)

log Pr[y n |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 ∂L n (φ)/∂φ (L-1) can be obtained by differentiating over Equation (37), while other gradients ∂L n (φ)/φ ( : L-2) further obtained by chain rule:

∂L n (φ) ∂φ ( : L-2) = ∂L n (φ) ∂ψ (L-1) • ∂ψ (L-1) ∂φ ( : L-2) (38)

which requires us to compute ∂L n (φ)/∂ψ (L-1) and ∂ψ (L-1) /∂φ ( : L-2) . While ∂L n (φ)/∂ψ (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) and φ (l) , the partial derivatives ∂ψ (L-1) /∂φ ( : L-2) can be obtained by iterative chain rule.

ψ (l+1) = g (l) (ψ (l) , φ (l) ) is differentiable w.r.t. ψ (

C.2 SOFTMAX LAYER FOR CLASSIFICATION PROBLEM:

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 L n (φ) for Classification). Let h ∈ R K (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(h c /s)/ K k=1 exp(h k /s). Suppose the logits {h k } K k=1 are pairwise independent (which holds under mean-field approximation) and

h k follows a Gaussian distribution h k ∼ N (µ k , ν k ) (therefore ψ = {µ k , ν k } K k=1

) and s is a deterministic parameter. Then L n (φ) can be further upper bound by the following analytic form:

EQUATION

Proof. The lower bound follows by plugging Pr[y|h, s] and Pr[h k |x] into Equation (6).

EQUATION

where the last equation follows

EQUATION

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 of L n (φ) 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.

EQUATION

Furthermore, the scale s can be (optionally) updated along with other parameters using the gradient

EQUATION

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:

EQUATION

(Mean-field assumption) = h1 • • • h K c (h) K k=1 Pr[h k |x] dh 1 • • • dh k (54)

where we denote the softmax function as

c (h) = exp(h c /s)/[ k exp(h k /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 h k we draw S samples {h s k } S s=1 independently and compute the prediction:

EQUATION

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

EQUATION

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

EQUATION

where the second term is zero by the symmetry of Pr[h k |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).

EQUATION

∂ c ∂h k (h) = - 1 s • exp (h c /s) • exp (h k /s) K k=1 exp (h k /s) 2 = - 1 s c (h) k (h), ∀k = c (60b)

and subsequently the second-order derivatives from the first ones:

EQUATION

∂ 2 c ∂h 2 k (h) = - 1 s ∂ c ∂h c (h) k (h) + c (h) ∂ k ∂h c (h) = 1 s 2 -c (h) k (h) + 2 2 c (h) k (h) , ∀k = c (61b)

with these derivatives we can compute the second-order approximation as

Pr[y = c|x] ≈ h1,•,h K   c (µ) + 1 2 K i=1 K j=1 ∂ 2 c ∂µ i µ j (µ)(h i -µ i )(h j -µ j )   K k=1 Pr[h k |x] dh 1 • • • dh K (62) = c (µ) + 1 2 ∂ 2 c ∂µ 2 c (µ) hc (h c -µ c ) 2 Pr[h c |x]dh c + 1 2 k =c ∂ 2 c ∂µ 2 k (µ) h k (h k -µ k ) 2 Pr[h k |x]dh k (63) = c (µ) + 1 2s 2 c (µ) -3 2 c (µ) + 2 3 c (µ) ν c + 1 2s 2 k =c -c (µ) k (µ) + 2 2 c (µ) k (µ) ν k (64) = c (µ) + 1 2s 2 c (µ) -2 2 c (µ) ν c - K k=1 k (µ)ν k (65)

The equation above can be further written in vector form as:

EQUATION

C.3 GAUSSIAN OUTPUT LAYER FOR REGRESSION PROBLEM:

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 L n (φ) for Regression). Let h ∈ R I be the output of last hidden layer (with I the number of hidden units), and φ = (w, s) ∈ R I × R + be the parameters that define the predictive distribution over output y as

EQUATION

)

Suppose the hidden units {h k } K k=1 are pairwise independent (which holds under mean-field approximation), and each h i has mean µ i and variance ν i , then L n (φ) takes an analytic form:

L n (φ) = - (y -w µ) 2 + (w •2 ) ν 2s - log(2πs) 2 (68)

where

(w •2 ) i = w 2 i and µ = [µ 1 , • • • , µ I ] ∈ R I and ν = [ν 1 , • • • , ν I ] ∈ R I

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

EQUATION

where the long summation in the first term can be further simplified with notations of µ and ν:

EQUATION

=y 2 -2y

EQUATION

)

where w •2 denotes element-wise square, i.e.

w •2 = [w 2 1 , • • • , w 2 I ] .

Derivatives of L n (φ) in Equation (68) It is not difficult to show that the gradient of L n (φ) can be backpropagated through the last layer. by computing derivatives of L n (φ) w.r.t. µ and ν:

∂L n (φ) ∂µ = - (y -w µ)w s (79a) ∂L n (φ) ∂ν = - w •2 2s (79b)

Furthermore, the parameters {w, s} can be updated along with other parameters with their gradients:

∂L n (φ) ∂w = - (y -w µ)µ + (w • ν) s (80a) ∂L n (φ) ∂s = - 1 2s + (y -w µ) 2 + (w •2 ) ν 2s 2 (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

EQUATION

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 h i as γ i and κ i :

γ i = E (h i -µ i ) 3 |x /ν 3/2 i = hi (h i -µ i ) 3 Pr[h i |x]/ν 3/2 i (83a) κ i = E (h i -µ i ) 4 |x /ν 2 i = hi (h i -µ i ) 4 Pr[h i |x]/ν 2 i (83b)

Then the (normalized) skewness and kurtosis of the prediction y are also easily computed with the vectors of γ

EQUATION

D PROBABILISTIC PROPAGATION IN BAYESIAN QUANTIZED NETWORKS:

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.

D.1 SMALL FAN-IN LAYERS: DIRECT TENSOR CONTRACTION:

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

EQUATION

Depth-wise Layers In a depth-wise layer, each output unit h

EQUATION

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(JD 2 ) since E <= 2.

D.2 MEDIUM FAN-IN LAYERS: DISCRETE FOURIER TRANSFORM:

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

= max i∈I(j) h

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

EQUATION

Prob: P (h

EQUATION

where P (h

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. 

EQUATION

Then C v (f ) is the element-wise product of all Fourier transforms C ui (f ), i.e.

EQUATION

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/∂P i } I i=1 . For this purpose, we break the computation in Equation (99) into three steps, and compute the derivative for each of these steps.

EQUATION

where in (100b) we use C/C i to denote element-wise division. Since P i lies into real domain, we need to project the gradients back to real number in (100c). Putting all steps together:

EQUATION

D.3 LARGE FAN-IN LAYERS: LYAPUNOV CENTRAL LIMIT APPROXIMATION:

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

EQUATION

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 } I i=1 , i.e. addition of a large number of random variables. Theorem D.2 (Fast summation via Lyapunov Central Limit Theorem). Let v j = σ(ṽ j ) = σ( i∈I(j) θ ji u i ) be an activation of a linear layer followed by nonlinearity σ. Suppose both inputs {u i } i∈I and parameters {θ ji } i∈I(j) have bounded variance, then for sufficiently large |I(j)|, the distribution of ṽj converges to a Gaussian distribution N (μ j , νj ) with mean and variance as

EQUATION

)

where

m ji = E [θ ji ], v ji = V [θ ji ] and µ i = E [u i ], ν i = V [u i ].

And if the nonlinear transform σ is a sign function, each activation v j follows a Bernoulli distribution P (v j = 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:

EQUATION

Published as a conference paper at ICLR 2020

EQUATION

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. v ji is defined upon m ji ). Therefore, we adjust the backpropagation rule for the probabilities: denote

Q ji (d) = Q(θ ji = Q(d); φ (l)

ji ), then the backpropagation rule can be written in matrix form as ∂L ∂Q(d)

EQUATION

Lastly, we derive the backpropagation rules for sign activations. Let p j denote the probability that the hidden unit v j is activated as p j = Pr[v j = 1|x], ∂L/p j relates to {∂L/∂ μj , ∂L/∂ νj } as:

EQUATION

E SUPPLEMENTARY MATERIAL FOR EXPERIMENTS E.1 NETWORK ARCHIECTURES

(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. The performance gap between our analytical inference and the Monte Carlo Sampling is displayed.

E.3 REGRESSION ON BOSTON HOUSING DATASET:

In this part, we evaluate our proposed BQN on Boston housing dataset, a regression benchmark widely used in testing Bayesian neural networks (Hernández-Lobato & Adams, 2015; Ghosh et al., 2016) and Probabilistic neural networks (Wang et al., 2016) 2016), and the one for NPN is from (Wang et al., 2016) .

:

CONCLUSIONWe 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.

Figures and Tables:

on a subset of all input units {h (l) i } and model parameters {θ

Figure 2: Illustration of mean-field approximation and tightness of alternative ELBO on a CNN. The performance gap between our analytical inference and the Monte Carlo Sampling is displayed.

Graphical depiction of the problem setting in Bayesian neural networks.

Graphical model depiction of a Bayesian neural network as hierarchical model, where predicting y from x can be performed iteratively through the hidden variables h (l) 's.

Figure 3: Graphical models.

Fast summation via discrete Fourier transform). Suppose u i take values in {b i , b i + 1, . . . , B i } between integers b i and B i , then the summation v = E i=1 u i takes values between b and B, where b = E i=1 b i and B

Figure 4: Comparison of the predictive performance of our BQNs against the E-QNN as well as the non-quantized BNN trained by SGVB on a MLP. Negative log-likelihood (NLL) which accounts for uncertainty and 0-1 test error which doesn't account for uncertainty are displayed.

Figure 5: Illustration of mean-field approximation and tightness of alternative ELBO on a MLP. The performance gap between our analytical inference and the Monte Carlo Sampling is displayed.

Alternative Evidence Lower Bound. We make learning in BNNs amenable to backpropagation by developing a lower bound L n (φ) ≤ L n (φ) such that ∂L n (φ)/∂φ can be obtained by chain rule (i.e. L n (φ) is an explicit function of the results from the forward pass.) With L n (φ) in hand, we can (approximately) find φ by maximizing the alternative objective via gradient-based method:

Table content: <table/>

Deterministic model compression through direct training of QNN(Courbariaux et al., 2016) v.s. MAP estimation in our proposed BQN. All the numbers are averages over 10 runs with different seeds, the standard deviation are exhibited following the ± sign.

Table content: <table><tr><td>Methods</td><td colspan="2">MNIST NLL(10 -3 )</td><td>% Err.</td><td colspan="2">KMNIST NLL(10 -3 )</td><td>% Err.</td><td>Fashion-MNIST NLL(10 -3 ) % Err.</td></tr><tr><td>QNN on MLP</td><td>522.4±42.2</td><td colspan="2">4.14±0.25</td><td>2019.1±281.2</td><td colspan="2">19.56±1.97</td><td>2427.1±193.5</td><td>15.67±1.19</td></tr><tr><td>MAP of BQN on MLP</td><td>137.60±4.40</td><td colspan="2">3.69±0.09</td><td>464.60±12.80</td><td colspan="2">14.79±0.21</td><td>461.30±13.40</td><td>12.89±0.17</td></tr><tr><td>QNN on CNN</td><td>497.4±139.5</td><td/><td>1.08±0.2</td><td>4734.5±1697.2</td><td colspan="2">14.2±2.29</td><td>1878.3±223.8</td><td>3.88±0.33</td></tr><tr><td>MAP of BQN on CNN</td><td>30.3±1.6</td><td colspan="2">0.92±0.07</td><td>293.6±4.4</td><td colspan="2">10.82±0.37</td><td>179.1±4.4</td><td>5.00±0.11</td></tr></table>

Bayesian Model compression through direct training of Ensemble-QNN vs a Monte-Carlo sampling on our proposed BQN. Each ensemble consists of 5 quantized neural networks, and for fair comparison we use 5 samples for Monte-Carlo evaluation. All the numbers are averages over 10 runs with different seeds, the standard deviation are exhibited following the ± sign.

Table content: <table/>

and PBP(Herná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(g

Table content: <table/>

. 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 Herná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 Table4. The results show that our BQN achieves lower RMSE compared to other models trained in a probabilistic/Bayesian way. Performance of different networks in terms of RMSE. The numbers for BQN are averages over 10 runs with different seeds, the standard deviation are exhibited following the ± sign. The results for PBP, EBP are fromGhosh et al. (

Table content: <table><tr><td>Dataset</td><td>BQN</td><td colspan="3">PBP (Ghosh et al., 2016) EBP (Soudry et al., 2014) NPN (Wang et al., 2016)</td></tr><tr><td colspan="2">Boston 2.04 ± 0.07</td><td>2.79 ± 0.16</td><td>3.14 ± 0.93</td><td>2.57± NA</td></tr></table>

Peer review (iclr2020):
(1) Contributions: The paper proposes an algorithm for training quantized Bayesian neural networks (BQN). The method optimizes a lower bound to the ELBO that can be analytically computed for BQNs. The paper shows that propagation in BQNs is a sequence of tensor contractions, for which there exist fast and efficient approximate algorithms. BQNs show good performance on set of image classification benchmarks (MNIST, KMNIST, Fashion MNIST). (2) Originality: I am not aware of any prior works that attempts to do inference in quantized networks. The idea of BQNs and propagation as tensor contractions is interesting and novel. It is a cool result that in quantized networks, one can exactly calculate the distribution of the activations (h^(l)) and therefore get an accurate estimate of the ELBO. It is also interesting that this algorithm can be implemented as a series of tensor contractions. This leads to a fast and efficient algorithm for training. Regarding the novelty in the lower bound to the ELBO, variational Bayes makes the same approximation by sampling the layer activations (h^(1) .. h^(L-1)). In VB, when we sample the layer activations, we approximate the logarithm of the expectation by the logarithm of a Monte Carlo sample, which is a lower bound due to Jensen’s inequality. The only difference here is that this approximation is only done in the final hidden layer instead of every layer. (2) Clarity: I want to point out some issues with the writing. Overall, I believe that the number of equations and theorems hinders the readability of the paper. It is very confusing to try and navigate the various notations and assess the significance of each result. I would recommend including a few key equations and theorems in the paper and leaving the rest to the appendix. The ideas and concepts in the paper could be easily communicated in writing. I do not think this paper should be 9.5 pages. The work could be comfortably explained in the standard 8 pages. Namely, in Section 3.1 and 3.2, it is unnecessary to provide analytic gradients for the expressions (in the age of automatic differentiation). Section 4.3 should be left to the appendix. I found this section very confusing and it does not seem to be important in the narrative. (4) Significance: I think this work is significant, although I felt that the experiments section could have included more regression benchmarks as well as more difficult datasets such as CIFAR10. (5) Impact: This work would certainly be interesting to researchers and anyone working on quantized neural networks. Overall assessment: I think this paper presents interesting and novel ideas, but it has shortcomings in the presentation.