S AMPLING -F REE L EARNING OF B AYESIAN Q UANTIZED N EURAL N ETWORKS
A BSTRACT
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 efﬁcient 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 CIFAR image classiﬁcation 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).
I NTRODUCTION
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 efﬁciency 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 efﬁcient 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-ﬁeld approximation) that speeds up the forward pass in BQNs.
• BQNs provide a tunable trade-off between computational resource and model complexity: using a reﬁned 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.
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-ﬁeld 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-ﬁeld approximation in BQNs.
Related Works.
In Appendix A, we survey different approaches for training Bayesian neural networks including sampling-free assumed density ﬁltering (Minka, 2001; Soudry et al., 2014;
Hern ´ andez-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 ´ us, 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 efﬁcient learning and prediction for BQNs.
B AYESIAN N EURAL N ETWORKS
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. P ROBLEM S ETTING
Given a dataset D = { ( x n , y n ) } 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 ﬁrst solve the learning problem to ﬁnd 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 ( θ ; φ )
= (cid:81)
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. T HE P REDICTION P ROBLEM
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
[Equation]
This iterative process to compute the predictive distributions layer-by-layer sequentially is known as probabilistic propagation (Soudry et al., 2014; Hern ´ andez-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 difﬁcult to ﬁnd 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. T HE L EARNING P ROBLEM
Objective Function.
A standard approach to ﬁnding a good approximation Q ( θ ; φ ) is variational inference , which ﬁnds φ (cid:63) 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 ( φ ) .
[Equation]
where L n ( φ ) =
E Q [log Pr [ y n | x n , θ ]] and R ( φ ) =
E Q [log [Equation] Pr [ θ ] [Equation]
+ H [Equation]
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 ﬁrst-order methods, typically stochastic gradient descent/ascent .
Notice that (1) For classiﬁcation, 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 difﬁculty 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):
[Equation]
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.
P ROPOSED L EARNING M ETHOD FOR B AYESIAN N EURAL N ETWORKS
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)
ﬁnd φ (cid:63) by maximizing the alternative objective via gradient-based method:
[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 ) .
Deﬁne each term L n ( φ ) in L ( φ ) as
[Equation]
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. classiﬁcation 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 classiﬁcation and regression tasks, and present the classiﬁcation 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 Classiﬁcation ) .
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 ) / (cid:80)
Kk
=1 exp( h k /s ) .
Suppose the logits { h k } Kk =1 are pairwise independent (which holds under mean-ﬁeld approximation) and h k follows a Gaussian distribution h k ∼ N ( µ k , ν k ) (therefore ψ = { µ k , ν k } Kk =1 ) and s is a deterministic parameter.
Then L n ( φ ) is further lower bounded as L n ( φ ) ≥
µ c s − log (cid:16)(cid:80)
Kk =1 exp (cid:0) µ k s + ν k 2 s 2 (cid:1)(cid:17) .
The regression case and proofs for both layers are deferred to Appendix C.
B AYESIAN Q UANTIZED N ETWORKS (BQN S )
While Section 3 provides a general solution to learning in BNNs, the solution relies on the ability to perform probabilistic propagation efﬁciently.
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 efﬁcient 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 ).
P ROBABILISTIC P ROPAGATION AS T ENSOR C ONTRACTIONS
Lemma 4.1 ( Probabilistic Propagation in BQNs ) .
After quantization, the iterative step of probabilistic propagation in Equation (1) is computed with a ﬁnite sum instead of an integral:
[Equation]
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 ) ; ψ ( l ) )
∈
R D I , Q ( θ ( l ) ; φ ( l ) )
∈ R D K , P ( h ( l +1) ; ψ ( l +1) )
∈ R D J , and 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 difﬁculties 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. A PPROXIMATE P ROPAGATION VIA R ANK -1 T ENSOR CP D ECOMPOSITION
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-ﬁeld 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
[Equation]
j
[Equation]
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 )) /
(cid:80) Dq =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 and M ( l +1) j , and deﬁne the fan-in number E for the layer as max j (cid:12)(cid:12)(cid:12) I ( l +1) j (cid:12)(cid:12)(cid:12) + (cid:12)(cid:12)(cid:12) M ( l +1) j (cid:12)(cid:12)(cid:12) .
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. F AST A LGORITHMS FOR A PPROXIMATE P ROPAGATION
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 inefﬁcient 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 sufﬁciently 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 .
E XPERIMENTS
In this section, we demonstrate the effectiveness of BQNs on the MNIST, Fashion-MNIST, KMNIST and CIFAR10 classiﬁcation 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 ﬂipping 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]
(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 coefﬁcient λ 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 [Figure]
are displayed. 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 :
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.
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-ﬁeld approximation and (b) the tightness of the our proposed alternative ELBO.
MAP estimation is similar to MC sampling, except that only one set of model parameters θ (cid:63) is obtained θ (cid:63) =
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. A NALYSIS OF R ESULTS
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.
Compared to E-QNNs, our BQNs have signiﬁcantly lower NLL and smaller predictive error (except for
[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-ﬁeld 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 inﬁnite samples).
Therefore we can evaluate the approximations in probabilistic propagation, namely mean-ﬁeld 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 justiﬁes the approximations we use in BQNs.
To further decouple these two factors of mean-ﬁeld approximation and relaxation of the original
ELBO , we vary the regularization coefﬁcient λ in the learning objective.
(1) For λ =
(where the prior term is removed), the models are forced to become deterministic during training.
Since the deterministic models do not have mean-ﬁeld approximation in the forward pass, the gap between analytic inference and MC-sampling reﬂects the tightness of our alternative ELBO.
(2) As λ increases, the gaps increases slightly as well, which shows that the mean-ﬁeld approximation becomes slightly less accurate with higher learned uncertainty in the model.
[Table]
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 inﬁnite 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
[Table] 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
[Table] we compare the models by both approaches to their counterparts (a single QNN for MAP [Table] and E-QNN for MC sampling) trained from scratch as in Courbariaux et al.
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.
C ONCLUSION
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 efﬁcient 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.
Appendix: Sampling-Free Learning of Bayesian Quantized Neural Networks
A
R ELATED W ORK
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´andez-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 (Hern ´ andez-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 n ( φ ))
/∂φ , and EBP Soudry et al.
(2014) derives a similar rule for Bernoulli parameters in binary classiﬁcation.
Notice that ADF is compatible to backpropagation:
[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 difﬁcult 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 } Ss =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 speciﬁc 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.,
The partial derivatives ∂θ s /∂φ are difﬁcult 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
[Equation]
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 difﬁcult 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´us, 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 S UPERVISED L EARNING WITH B AYESIAN NEURAL NETWORKS (BNN S )
The problem settings of general Bayesian model and Bayesian neural networks for supervised learning are illustrated in Figures 3a and 3b using graphical models .
[Equation] ... x θ (0) h (1) [Figure]
[Figure] [Figure] ( L
[Figure]
1) y [Figure] 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:
[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:
[Equation]
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 ﬁnding a good approximation Q ( θ ; φ ) is variational inference , which ﬁnds φ (cid:63) such that the KL-divergence KL ( Q ( θ ; φ ) ||
Pr [ θ |D ]) of Q ( θ ; φ ) from Pr [ θ |D ] is minimized (or alternatively the negative KL-divergence is maximized.)
[Equation]
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):
[Equation]
[Equation]
where (1) L n ( φ ) is the expected log-likelihood , which reﬂects 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 satisﬁes 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 − 1 l =0 (where each φ ( l ) parameterizes θ ( l ) in the l -th layer) such that:
[Equation]
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 ) :
[Equation]
where we use short-hand notations h (: l ) and θ (: l ) to represent the sets of previous layers { h ( k ) } lk
= 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:
[Equation]
C A LTERNATIVE E VIDENCE L OWER B OUND AND ITS A NALYTIC F ORMS C. A LTERNATIVE E VIDENCE L OWER B OUND (P ROOF FOR T HEOREM 3.1)
The steps to prove the inequality (6) almost follow the ones for probabilistic propagation above:
[Equation]
[Equation]
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)
[Equation]
L n ( φ ) can be simpliﬁed as:
[Equation]
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:
[Equation]
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 +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.
C. S OFTMAX L AYER FOR C LASSIFICATION P ROBLEM
In this part, we ﬁrst 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 Classiﬁcation ) .
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 ) / (cid:80)
Kk
=1 exp( h k /s ) .
Suppose the logits { h k } Kk =1 are pairwise independent (which holds under mean-ﬁeld approximation) and h k follows a Gaussian distribution h k ∼ N ( µ k , ν k ) (therefore ψ = { µ k , ν k } Kk =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]
(cid:32) K (cid:89) k
=
Pr [ h k | x ] (cid:33) dh (42)
= µ c s − (cid:90) h log (cid:32) K (cid:88) k
= exp (cid:18)
h k s (cid:19)(cid:33) (cid:32) K (cid:89) k
=
Pr [ h k | x ] (cid:33) dh (43) ≥ µ c s − log
cid:32)(cid:90)
h
K (cid:88) k
= exp (cid:18) h k s (cid:19) (cid:32) K (cid:89) k
=
Pr [ h k | x ] (cid:33) dh (cid:33) (44) = µ c s − log (cid:32) K (cid:88) k
= (cid:90) h k exp (cid:18)
h k s (cid:19)
Pr [ h k | x ] dh k (cid:33) (45) = µ c s − log (cid:32) K (cid:88) k
= (cid:90) h k exp (cid:18)
h k s (cid:19) · 1 √ 2 πν k exp (cid:18) − ( h k − µ k )
ν k (cid:19) dh
k (cid:33) (46) = µ c s − log (cid:32) K (cid:88) k
= exp (cid:16) µ k s + ν k 2 s(cid:17)(cid:33)
=
ˆ L ( φ ) (47) 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 ﬁrst need to obtain the derivatives w.r.t.
ψ ( L − 1) =
{ µ k , ν k } K k =1 .
[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-ﬁeld assumption
[Equation]
where we denote the softmax function as (cid:96) c ( h ) =
exp( h c /s ) /
[ (cid:80) k exp( h k /s )] .
Unfortunately
[Equation]
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 sk } Ss = independently and compute the prediction:
[Equation] [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
[Equation]
First [Equation] we expand the function (cid:96) c
( h ) by Taylor’s series at the point µ (up to the second order):
[Equation] [Equation]
Before we derive the forms of these derivatives, we ﬁrst show the terms of odd orders do not contribute to the expectation.
For example, if (cid:96) c ( h ) is approximated by its ﬁrst 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 deﬁnition of µ k ’s).
Therefore, the ﬁrst-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 (cid:96) c
( h ) .
Speciﬁcally, the ﬁrst-order derivatives are obtained from the deﬁnition of (cid:96) c
( h ) .
[Equation]
and subsequently the second-order derivatives from the ﬁrst ones:
[Equation]
with these derivatives we can compute the second-order approximation as
[Equation]
The equation above can be further written in vector form as:
[Equation]
C. G AUSSIAN O UTPUT L AYER FOR R EGRESSION P ROBLEM
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 difﬁculty 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 deﬁne the predictive distribution over output y as
[Equation]
Suppose the hidden units { h k } Kk =1 are pairwise independent (which holds under mean-ﬁeld approximation), and each h i has mean µ i and variance
ν i , then L n ( φ ) takes an analytic form:
[Equation]
where ( w ◦ 2 )
i =
w 2 i and µ = [ µ 1 , · · · , µ I ] (cid:62)
∈ R I and ν =
[ ν 1 , · · · , ν I ] (cid:62) ∈ 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 ﬁrst term can be further simpliﬁed with notations of µ and ν : (cid:88)
[Equation]
where w ◦ 2 denotes element-wise square, i.e. w ◦ 2 =
[ w 2 1 , · · · , w 2 I ] (cid:62) .
Derivatives of L n ( φ ) in Equation (68)
It is not difﬁcult 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 ν :
[Equation]
Furthermore, the parameters { w, s } can be updated along with other parameters with their gradients:
[Equation]
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 ] = (cid:88)
[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 (cid:62) h + (cid:15) , where (cid:15) ∼ N (0 , s ) , its mean and variance can be easily computed as
[Equation]
Furthermore, if we denote the (normalized) skewness and kurtosis of h i as γ i and κ i :
[Equation]
Then the (normalized) skewness and kurtosis of the prediction y are also easily computed with the vectors of γ =
[ γ 1 , · · · , γ I ] (cid:62) ∈ R I and κ = [ κ 1 , · · · , κ I ] ∈ R I .
[Equation]
D P ROBABILISTIC P ROPAGATION IN B AYESIAN Q UANTIZED N ETWORKS
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 speciﬁc range of fan-in numbers E .
D. S MALL F AN - IN L AYERS : D IRECT T ENSOR C ONTRACTION
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 ( l +1) i is a transformation (parameterized by θ ( l ) i ) of its corresponding input h ( l ) i , i.e.
[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.
M EDIUM F AN - IN L AYERS : D ISCRETE F OURIER T RANSFORM
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 = max i ∈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
[Equation]
For both cases, the predictive distribution of h ( l +1)
j can be computed as
[Equation]
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 u i take values in { b i , b i + 1 , . . .
, B i } between integers
b i and B i , then the summation v = (cid:80)
Ei =1 u i takes values between b and B , where b = (cid:80)
Ei =1 b i and B = (cid:80)
Ei =1 B i .
Let C v , C u i be the discrete Fourier transforms of P v , P u i respectively, i.e.
[Equation]
Then C v ( f ) is the element-wise product of all Fourier transforms C u
i ( f ) , i.e.
[Equation]
Proof.
We only prove the theorem for two discrete random variable, and the extension to multiple variables can be proved using induction.
Now consider u 1 ∈ [ b 1 , B 1 ] , u 2 ∈
[ b 2 , B 2 ] and their sum v = u
1 + u 2 ∈ [ b, B ] , where b = b 1 + b 2 and B = B 1 + B 2 .
Denote the probability vectors of u 1 , u 2 and v as P 1 ∈ (cid:52) B 1 − b 1 , P 2 ∈ (cid:52) B 2 − b 2 and P ∈ (cid:52) B − b respectively, then the entries in P are computed with P 1 and P 2 by standard convolution as follows:
[Equation]
The relation above is usually denoted as P = P 1 ∗ P 2 , where ∗ is the symbol for convolution.
Now deﬁne the characteristic functions C , C 1 , and C 2 as the discrete Fourier transform (DFT) of the probability vectors P , P 1 and P 2 respectively:
[Equation]
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, C 1 , C 2 ∈ C R , and we denote the (functional) mappings as C = F ( P ) and
C i
=
F i ( P i ) .
Given a characteristic function, its original probability vector can be recovered by inverse discrete Fourier transform (IDFT) :
[Equation]
93b) which we denote the inverse mapping as P = F − 1 ( C ) and P i = F − 1 i ( C i ) .
Now we plug the convolution in Equation (91) into the characteristic function C ( f ) in (92a) and rearrange accordingly:
[Equation]
[Equation]
(Let u 2 = v − u 1 ) =
B 1 (cid:88) u
1 = b
B 2 (cid:88) u
= b
P 1 ( u 1 ) P 2 ( u 2 ) exp (cid:18)
[Equation]
(cid:19) (95) (Since b = b 1 + b 2 ) =
(cid:34) B 1 (cid:88) u 1 = b
[Equation]
[Equation] (cid:18) [Equation] (cid:19)(cid:35)
[Equation]
The equation above can therefore be written as C = C 1 ◦ C 2 , 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 efﬁciently as follows
[Equation]
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 log R )
=
O ( E 2 D log( ED ))
(since R = O ( ED ) ), compared to O ( D E ) 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 /∂P i } Ii =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. L ARGE F AN - IN L AYERS : L YAPUNOV C ENTRAL L IMIT A PPROXIMATION
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 ) ∈
R I to h ( l +1) ∈
R J as
[Equation]
where u ( l ) ji = θ ( l ) ji · h ( l ) i and v ( l +1)
j = (cid:80) i ∈I ( j ) u ( l ) ji .
The key difﬁculty 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 v j = σ (˜ v j ) =
σ ( (cid:80) 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 sufﬁciently large |I
( j ) | , the distribution of ˜ v j converges to a Gaussian distribution N (˜ µ j , ˜ ν j ) with mean and variance as ˜ µ j =
I (cid:88) i = m ji µ i (103a) ˜ ν j =
I (cid:88) i
= m
2 ji ν i + v ji µ 2 i + v ji ν i (103b) 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 / (cid:112) ˜
[Equation]
where Φ is the culminative probability function of a standard Gaussian distribution N (0 , 1) .
Proof.
The proof follows directly from the deﬁnitions of mean and variance:
[Equation] [Equation]
For fully-connected layers, these two equations can be concisely written in matrix forms:
[Equation]
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.
[Equation]
Notice that these equations do not take into account the fact that V implicitly deﬁned with M (i.e. v ji is deﬁned 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
[Equation]
[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 S UPPLEMENTARY M ATERIAL FOR E XPERIMENTS E. N ETWORK A RCHIECTURES
(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 ﬁrst 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. M ORE R ESULTS FOR M ULTI - LAYER P ERCEPTRON (MLP)
[Figure]
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]
The performance gap between our analytical inference and the Monte Carlo Sampling is displayed
[Figure]
E. R EGRESSION ON B OSTON H OUSING D ATASET
In this part, we evaluate our proposed BQN on Boston housing dataset, a regression benchmark widely used in testing Bayesian neural networks (Hern´andez-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 Hern´andez-Lobato & Adams (2015); Ghosh et al.
; 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.
[Table]