V ARIATIONAL N ETWORK Q UANTIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
In this paper, the preparation of a neural network for pruning and few-bit quantization is formulated as a variational inference problem.
To this end a quantizing prior that leads to a multi-modal, sparse posterior distribution over weights is introduced and a differentiable Kullback-Leibler divergence approximation for this prior is derived.
After training with Variational Network Quantization, weights can be replaced by deterministic quantization values with small to negligible loss of task accuracy (including pruning by setting weights to 0 ).
The method does not require ﬁne-tuning after quantization.
Results are shown for ternary quantization on LeNet-5 (MNIST) and DenseNet-121 (CIFAR-10).
I NTRODUCTION
Parameters of a trained neural network commonly exhibit high degrees of redundancy (Denil et al., 2013) which implies an over-parametrization of the network.
Network compression methods implicitly or explicitly aim at the systematic reduction of redundancy in neural network models while at the same time retaining a high level of task accuracy.
Besides architectural approaches, such as SqueezeNet (Iandola et al., 2016) or MobileNets (Howard et al., 2017), many compression methods perform some form of pruning or quantization .
Pruning is the removal of irrelevant units (weights, neurons or convolutional ﬁlters) (LeCun et al., 1990).
Relevance of weights is often determined by the absolute value (“magnitude based pruning” (Han et al., 2016; 2017; Guo et al., 2016)), but more sophisticated methods have been known for decades, e.g. based on second-order derivatives (Optimal Brain Damage (LeCun et al., 1990) and Optimal Brain Surgeon Hassibi & Stork (1993)) or ARD (automatic relevance determination, a Bayesian framework for determining the relevance of weights, (MacKay, 1995; Neal, 1995; Karaletsos & R ¨ atsch, 2015)).
Quantization is the reduction of the bit-precision of weights, activations or even gradients, which is particularly desirable from a hardware perspective (Sze et al., 2017).
Methods range from ﬁxed bit-width computation (e.g. 12-bit ﬁxed point) to aggressive quantization such as binarization of weights and activations (Courbariaux et al., 2016; Rastegari et al., 2016;
Zhou et al., 2016; Hubara et al., 2016).
Few-bit quantization ( 2 to 6 bits) is often performed by k-means clustering of trained weights with subsequent ﬁne-tuning of the cluster centers (Han et al., 2016).
Pruning and quantization methods have been shown to work well in conjunction (Han et al., 2016).
In so-called “ternary” networks, weights can have one out of three possible values (negative, zero or positive) which also allows for simultaneous pruning and few-bit quantization (Li et al., 2016; Zhu et al., 2016).
This work is closely related to some recent Bayesian methods for network compression (Ullrich et al., 2017;
Molchanov et al., 2017; Louizos et al., 2017; Neklyudov et al., 2017) that learn a posterior distribution over network weights under a sparsity-inducing prior.
The posterior distribution over network parameters allows identifying redundancies through three means: weights with (1) an expected value very close to zero and (2) weights with a large variance can be pruned as they do not contribute much to the overall computation.
(3) the posterior variance over non-pruned parameters can be used to determine the required bit-precision (quantization noise can be made as large as implied by the posterior uncertainty).
Additionally, Bayesian inference over modelparameters is known to automatically reduce parameter redundancy by penalizing overly complex models MacKay (2003).
In this paper we present Variational Network Quantization (VNQ), a Bayesian network compression method for simultaneous pruning and few-bit quantization of weights.
We extend previous Bayesian pruning methods by introducing a multi-modal quantizing prior that penalizes weights of low variance unless they lie close to one of the target values for quantization.
As a result, weights are either drawn to one of the quantization target values or they are assigned large variance values—see Fig.
After training, our method yields a Bayesian neural network with a multi-modal posterior over weights (typically with one mode ﬁxed at 0 ), which is the basis for subsequent pruning and quantization.
Additionally, posterior uncertainties can also be interesting for network introspection and analysis, as well as for obtaining uncertainty estimates over network predictions Gal & Ghahramani (2015); Gal (2016); Depeweg et al.
After pruning and hard quantization, and without the need for additional ﬁne-tuning, our method yields a deterministic feed-forward neural network with heavily quantized weights.
Our method is applicable to pre-trained networks but can also be used for training from scratch.
Target values for quantization can either be manually ﬁxed or they can be learned during training.
We demonstrate our method for the case of ternary quantization on
LeNet-5 (MNIST) and DenseNet-121 (CIFAR-10).
[Figure] [Figure]
P RELIMINARIES
Our method extends recent work that uses a (variational) Bayesian objective for neural network pruning (Molchanov et al., 2017). In this section we ﬁrst motivate such an approach by discussing that the objectives of compression (in the minimum-description-length sense) and Bayesian inference are well-aligned. We then brieﬂy review the core ingredients that are combined in Sparse Variational
Dropout (Molchanov et al., 2017).
The ﬁnal idea (and also the starting point of our method) is to learn dropout noise levels per weight and prune weights with large dropout noise. Learning dropout noise per weight can be done by interpreting dropout training as variational inference of an approximate weight-posterior under a sparsity inducing prior - this is known as Variational Dropout which is described in more detail below, after a brief introduction to modern approximate posterior inference in Bayesian neural networks by optimizing the evidence lower bound via stochastic gradient ascent and reparameterization tricks.
2. W HY B AYES FOR COMPRESSION ?
Bayesian inference over model parameters automatically penalizes overly complex parametric models, leading to an automatic regularization effect (Gr¨unwald, 2007; Graves, 2011) (see Molchanov et al.
(2017), where the authors show that Sparse Variational Dropout (Sparse VD) successfully prevents a network from ﬁtting unstructured data, that is a random labeling).
The automatic regularization is based on the objective of maximizing model evidence, also know as marginal likelihood.
A very complex model might have a particular parameter setting that achieves extremely good likelihood given the data, however, since the model evidence is obtained via marginalizing parameters, overly complex models are penalized for having many parameter settings with poor likelihood.
This effect is also known as “Bayesian Occams Razor” in Bayesian model selection (MacKay, 2003; Genewein & Braun, 2014). The argument can be extended to variational Bayesian inference (with some caveats) via the equivalence of the variational Bayesian objective and the Minimum description length (MDL) principle Rissanen (1978); Gr¨unwald (2007); Graves (2011);
Louizos et al.
The evidence lower bound (ELBO), which is maximized in variational inference, is the sum of two terms: one, the average message length required to transmit outputs (labels) to a receiver that knows the inputs and the posterior over model parameters and two, the average message length to transmit the posterior parameters to a receiver that knows the prior over parameters:
[Equation]
Maximizing the ELBO minimizes the total message length: max L ELBO =
min
L E + L C , leading to an optimal trade-off between short description length of the data and the model (thus minimizing the sum of error cost L E and model complexity cost L C ).
Interestingly, MDL dictates the use of stochastic models since they are in general “more compressible” compared to deterministic models: high posterior uncertainty over parameters is rewarded by the entropy term in L C —higher uncertainty allows the quantization noise to be higher, thus requiring lower bit-precision for a parameter.
Variational Bayesian inference can also be formally related to the information-theoretic framework for lossy compression, rate-distortion theory, (Cover & Thomas, 2006; Tishby et al., 2000; Genewein et al., 2015).
The only difference is that rate-distortion requires the use of the optimal prior , which is the marginal over posteriors Hoffman & Johnson (2016); Tomczak & Welling (2017); Hoffman et al. (2017) - providing an interesting connection to empirical Bayes where the prior is learned from the data.
2. V ARIATIONAL B AYES AND REPARAMETERIZATION
Let D be a dataset of N pairs ( x n , y n ) Nn =1 and p ( y | x, w ) be a parameterized model that predicts outputs y given inputs x and parameters w .
A Bayesian neural network models a (posterior) distribution over parameters w instead of just a point-estimate.
The posterior is given by Bayes’ rule: p ( w |D )
= p ( D| w ) p ( w ) /p
( D ) , where p ( w ) is the prior over parameters.
Computation of the true posterior is in general intractable.
Common approaches to approximate inference in neural networks are for instance: MCMC methods pioneered in (Neal, 1995) and later reﬁned e.g. via stochastic gradient Langevin dynamics (Welling & Teh, 2011), or variational approximations to the true posterior (Graves, 2011), Bayes by Backprop (Blundell et al., 2015), Expectation Backpropagation (Soudry et al., 2014), Probabilistic Backpropagation (Hern ´ andez-Lobato & Adams, 2015).
In the latter methods the true posterior is approximated by a parameterized distribution q φ ( w ) .
Variational parameters φ are optimized by minimizing the Kullback-Leibler (KL) divergence from the true to the approximate posterior D KL ( q φ ( w ) || p ( w |D )) .
Since computation of the true posterior is intractable, minimizing this KL divergence is approximately performed by maximizing the so-called “evidence lower bound” (ELBO) or “negative variational free energy” (Kingma & Welling, 2014):
[Equation]
where we have used the Reparameterization Trick 2 (Kingma & Welling, 2014) in Eq.
(2) to get an unbiased, differentiable, minibatch-based Monte Carlo estimator of the expected log likelihood
L D ( φ ) .
A mini-batch of data is denoted by (˜ x m , ˜ y m )
Mm =1 .
Additionally, and in line with similar work (Molchanov et al., 2017; Louizos et al., 2017; Neklyudov et al., 2017), we use the Local
Reparameterization Trick (Kingma et al., 2015) to further reduce variance of the stochastic ELBO gradient estimator, which locally marginalizes weights at each layer and instead samples directly from the distribution over pre-activations (which can be computed analytically).
See Appendix A. for more details on the Local reparameterization.
Commonly, the prior p ( w ) and the parametric form of the posterior q φ ( w ) are chosen such that the KL divergence term can be computed analytically (e.g. a fully factorized Gaussian prior and posterior, known as the mean-ﬁeld approximation).
Due to the particular choice of prior in our work, a closed-form expression for the KL divergence cannot be obtained but instead we use a differentiable approximation (see Sec. 3.3).
2. V ARIATIONAL INFERENCE VIA D ROPOUT TRAINING
Dropout (Srivastava et al., 2014) is a method originally introduced for regularization of neural networks, where activations are stochastically dropped (i.e. set to zero) with a certain probability p during training.
It was shown that dropout, i.e. multiplicative noise on inputs, is equivalent to having noisy weights and vice versa (Wang & Manning, 2013; Kingma et al., 2015).
Multiplicative
Gaussian noise ξ ij ∼ N (1 , α = p 1 − p ) on a weight w ij induces a Gaussian distribution
[Equation]
(3) with (cid:15) ij ∼ N (0 , 1) .
In standard (Gaussian) dropout training, the dropout rates α (or p to be precise) are ﬁxed and the expected log likelihood L D ( φ ) (ﬁrst term in Eq.
(1)) is maximized with respect to the means θ .
Kingma et al.
(2015) show that Gaussian dropout training is mathematically equivalent to maximizing the ELBO (both terms in Eq.
(1)), under a prior p ( w )
and ﬁxed α where the KL term does not depend on θ :
[Equation]
where the dependencies on α and θ of the terms in Eq.
(1) have been made explicit.
The only prior that meets this requirement is the scale invariant log-uniform prior:
[Equation]
Using this interpretation, it becomes straightforward to learn individual dropout-rates α ij per weight, by including α ij into the set of variational parameters φ
= ( θ, α ) .
This procedure was introduced in (Kingma et al., 2015) under the name “Variational Dropout”.
With the choice of a log-uniform prior (Eq.
(5)) and a factorized Gaussian approximate posterior q φ ( w ij ) =
N ( θ ij , α ij θ 2 ij ) (Eq. (3))
the KL term in Eq.
(1) is not analytically tractable, but the authors of Kingma et al.
(2015) present an approximation
[Equation]
see original publication for numerical values of c 1 , c 2 , c 3 .
Note that due to the mean-ﬁeld approximation, where the posterior over all weights factorizes into a product over individual weights q φ ( w ) =
(cid:81) q φ ( w ij ) , the KL divergence factorizes into a sum of individual KL divergences
[Equation]
2. P RUNING UNITS WITH LARGE DROPOUT RATES
Learning dropout rates is interesting for network compression since neurons or weights with very high dropout rates p
→ 1 can very likely be pruned without loss in accuracy.
However, as the authors of Sparse Variational Dropout (sparse VD) (Molchanov et al., 2017) report, the approximation in Eq.
(6) is only accurate for α ≤ 1 (corresponding to p ≤ 0 . 5 ).
For this reason, the original variational dropout paper restricted α to values smaller or equal to 1 , which are unsuitable for pruning.
Molchanov et al.
(2017) propose an improved approximation, which is very accurate on the full range of log α :
[Equation]
with k 1 = 0 .
63576 , k 2 = 1 .
87320 and k 3 = 1 . 48695 and S
denoting the sigmoid function.
Additionally, the authors propose to use an additive, instead of a multiplicative noise reparameterization, which signiﬁcantly reduces variance in the gradient ∂ L SGVB ∂θ ij for large α ij .
To achieve this, the multiplicative noise term is replaced with an exactly equivalent additive noise term σ ij (cid:15) ij with σ 2 ij = α ij θ 2 ij and the set of variational parameters becomes φ
= ( θ, σ ) :
[Equation]
After Sparse VD training, pruning is performed by thresholding α ij = σ 2 ij θ 2 ij .
In Molchanov et al.
(2017) a threshold of log α = 3 is used, which roughly corresponds to p > 0 .
Pruning weights that lie above a threshold of T α leads to
[Equation]
(9) which means effectively that weights with large variance but also weights of lower variance and a mean θ ij close to zero are pruned.
A visualization of the pruning threshold can be seen in Fig 1 (the “central funnel”, i.e. the area marked by the red dotted lines for a threshold for T α = 2 ).
Sparse
VD training can be performed from random initialization or with pre-trained networks by initializing the means θ ij accordingly.
In Bayesian Compression (Louizos et al., 2017) and Structured
Bayesian Pruning (Neklyudov et al., 2017), Sparse VD has been extended to include group-sparsity constraints, which allows for pruning of whole neurons or convolutional ﬁlters (via learning their corresponding dropout rates).
2. S PARSITY INDUCING PRIORS
For pruning weights based on their (learned) dropout rate, it is desirable to have high dropout rates for most weights.
Perhaps surprisingly, Variational Dropout already implicitly introduces such a “high dropout rate constraint” via the implicit prior distribution over weights.
The prior p ( w ) can be used to induce sparsity into the posterior by having high density at zero and heavy tails.
There is a well known family of such distributions: scale-mixtures of normals (Andrews & Mallows, 1974; Louizos et al., 2017; Ingraham & Marks, 2017):
[Equation]
where the scales of w are random variables.
A well-known example is the spike-and-slab prior (Mitchell & Beauchamp, 1988), which has a delta-spike at zero and a slab over the real line.
Gal & Ghahramani (2015); Kingma et al.
(2015) show how Dropout training implies a spike-and-slab prior over weights.
The log uniform prior used in Sparse VD (Eq.
(5)) can also be derived as a marginalized scale-mixture of normals
[Equation]
also known as the normal-Jeffreys prior (Figueiredo, 2002).
Louizos et al.
(2017) discuss how the log-uniform prior can be seen as a continuous relaxation of the spike-and-slab prior and how the alternative formulation through the normal-Jeffreys distribution can be used to couple the scales of weights that belong together and thus learn dropout rates for whole neurons or convolutional ﬁlters, which is the basis for Bayesian Compression (Louizos et al., 2017) and Structured Bayesian Pruning (Neklyudov et al., 2017).
V ARIATIONAL N ETWORK Q UANTIZATION
We formulate the preparation of a neural network for a post-training quantization step as a variational inference problem.
To this end, we introduce a multi-modal, quantizing prior and train by maximizing the ELBO (Eq.
under a mean-ﬁeld approximation of the posterior (i.e. a fully factorized Gaussian).
The goal of our algorithm is to achieve soft quantization , that is learning a posterior distribution such that the accuracy-loss introduced by post-training quantization is small.
Our variational posterior approximation and training procedure is similar to Kingma et al.
and
Molchanov et al.
(2017) with the crucial difference of using a quantizing prior that drives weights towards the target values for quantization.
3. A QUANTIZING PRIOR
The log uniform prior (Eq.
(5)) can be viewed as a continuous relaxation of the spike-and-slab prior with a spike at location
(Louizos et al., 2017).
We use this insight to formulate a quantizing prior, a continuous relaxation of a “multi-spike-and-slab” prior which has multiple spikes at locations
c k , k ∈ { 1 ..K } .
Each spike location corresponds to one target value for subsequent quantization.
The quantizing prior allows weights of low variance only at the locations of the quantization target values c k .
The effect of using such a quantizing prior during Variational Network Quantization is shown in Fig.
After training, most weights of low variance are distributed very closely around the quantization target values c k and can thus be replaced by the corresponding value without signiﬁcant loss in accuracy.
We typically ﬁx one of the quantization targets to zero, e.g. c 2 = 0 , which allows pruning weights. Additionally, weights with a large variance can also be pruned. Both kinds of pruning can be achieved with an α ij threshold (see Eq.
as in sparse Variational Dropout (Molchanov et al., 2017).
Following the interpretation of the log uniform prior p ( w ij ) as a marginal over the scale-hyperparameter z ij , we extend Eq.
(10) with a hyper-prior over locations
[Equation]
with p ( z ij ) ∝ | z ij | − 1 .
The location prior p m ( m ij ) is a mixture of weighted delta distributions located at the quantization values c k .
Marginalizing over m yields the quantizing prior p
[Equation]
ij [Equation] In our experiments we use K = 3 , a k = 1 /K ∀ k and c 2 = 0 unless indicated otherwise.
3. P OST -T RAINING Q UANTIZATION
Equation (9) implies that using a threshold on α ij as a pruning criterion is equivalent to pruning weights whose value does not differ signiﬁcantly from zero:
[Equation]
To be precise, T α speciﬁes the width of a scaled standard-deviation band ± σ ij / √ T α around the mean θ ij .
If the value zero lies within this band, the weight is assigned the value 0 .
For instance, a pruning threshold which implies p ≥ 0 .
95 corresponds to a variance band of approximately σ ij / 4 .
An equivalent interpretation is that a weight is pruned if the likelihood for the value 0 under the approximate posterior exceeds the threshold given by the standard-deviation band (Eq.
[Equation] [Equation]
Extending this argument for pruning weights to a quantization setting, we design a post-training quantization scheme that assigns each weight the quantized value c k with the highest likelihood under the approximate posterior.
Since variational posteriors over weights are Gaussian, this translates into minimizing the squared distance between the mean θ ij and the quantized values c k :
[Equation]
Additionally, the pruning rate can be increased by ﬁrst assigning a hard 0 to all weights that exceed the pruning threshold T α (see Eq.
(9)) before performing the assignment to quantization levels as described above.
3. KL DIVERGENCE APPROXIMATION
Under the quantizing prior (Eq.
the KL divergence from the prior D KL ( q φ ( w ) || p ( w )) to the mean-ﬁeld posterior is analytically intractable.
Similar to Kingma et al.
; Molchanov et al.
we use a differentiable approximation F KL ( θ, σ, c ) 3 , composed of a small number of differentiable functions to keep the computational effort low during training.
We now present the approximation for a reference codebook c =
[ − r, 0 , r ] , r = 0 .
2 , however later we show how the approximation can be used for arbitrary ternary, symmetric codebooks as well.
The basis of our approximation is the approximation F KL , LU introduced by Molchanov et al.
(2017) for the KL divergence from a log uniform prior to a Gaussian posterior (see Eq.
(7)) which is centered around zero.
We observe that a weighted mixture of shifted versions of F KL , LU can be used to approximate the KL divergence for our multi-modal quantizing prior (Eq.
(which is composed of shifted versions of the log uniform prior).
In a nutshell, we shift one version of F KL to each codebook entry
c k
and then use θ -dependent Gaussian windowing functions Ω( θ ) to mix the shifted approximations (see more details in the Appendix A.3).
The approximation for the KL divergence from our multimodal quantizing prior to a Gaussian posterior is given as
[Equation]
with
[Equation]
We use τ = 0 .
075 in our experiments.
Illustrations of the approximation, including a comparison against the ground-truth computed via Monte Carlo sampling are shown in Fig.
Over the range of θ - and σ -values relevant to our method, the maximum absolute deviation from the ground-truth is 1 . 07 nats.
See Fig.
4 in the Appendix for a more detailed quantitative evaluation of our approximation.
[Figure]
but ground-truth and approximation align for large values of θ .
This KL approximation in Eq.
(16), developed for the reference codebook c r = [ − r, 0 , r ] , can be reused for any symmetric ternary codebook c a =
[ − a, 0 , a ] , a ∈ R + , since c a can be represented with the reference codebook and a positive scaling factor s , c a = sc r , s = a/r .
As derived in the Appendix (A.4), this re-scaling translates into a multiplicative re-scaling of the variational parameters θ and σ .
The KL divergence from a prior based on the codebook c a to the posterior q φ ( w ) is thus given by D KL ( q φ ( w )
||
p c a ( w ))
≈ F KL ( θ/s, σ/s, c r ) .
This result allows learning the quantization level a during training as well.
E XPERIMENTS
In our experiments we train with VNQ and then ﬁrst prune via thresholding log α ij ≥ log T α = 2 .
Remaining weights are then quantized by minimizing the squared distance to the quantization values c k (see Sec. 3.2).
We use warm-up (Sønderby et al., 2016), that is we multiply the KL divergence term (Eq.
(2)) with a factor β , where β
= 0 during the ﬁrst few epochs and then linearly ramps up to β = 1 .
To improve stability of VNQ training we ensure through clipping that log σ 2 ij ∈ ( − 10 , 1) and θ ij ∈ (0 . 3679 σ ± − a, ) (which corresponds to a shifted log α threshold of 2 , that is we clip θ ij if it lies left of the − a funnel or right of the + a funnel, compare Fig. 1).
This leads to a clipping-boundary that depends on trainable parameters. To avoid weights getting stuck at these boundaries we use gradient-stopping, that is we apply the gradient to a so-called “shadow weight” and use the clipped weight-value only for the forward pass.
Without this procedure our method still works, but accuracies are a bit worse, particularly on CIFAR-10.
When learning codebook values a during training, we use a lower learning rate for adjusting the codebook, otherwise we observe a tendency for codebook values to collapse in early stages of training (a similar observation was made by Ullrich et al.
Additionally, we ensure a ≥ 0 .
by clipping.
4. L E N ET -5 ON MNIST
We demonstrate our method with LeNet-5 4 (LeCun et al., 1998) on the MNIST handwritten digits dataset.
Images are pre-processed by subtracting the mean and dividing by the standard-deviation over the training set.
For the pre-trained network we run 5 epochs on a randomly initialized network (Glorot initialization, Adam optimizer), which leads to a validation accuracy of 99 . 2% .
We initialize means θ with the pre-trained weights and variances with log σ 2 = − 8 .
The warm-up factor β is linearly increased from 0 to 1 during the ﬁrst 15 epochs.
VNQ training runs for a total of 195 epochs with a batch-size of 128 , the learning rate is linearly decreased from 0 .
001 to 0 and the learning rate for adjusting the codebook parameter a uses a learning rate that is 100 times lower. We initialize with a = 0 .
Results are shown in Table 1, a visualization of the distribution over weights after VNQ training is shown in Fig.
We ﬁnd that VNQ training sufﬁciently prepares a network for pruning and quantization with negligible loss in accuracy and without requiring subsequent ﬁne-tuning. Training from scratch yields similar performance compared to initializing with a pre-trained network, with a slightly higher pruning rate.
Compared to pruning methods that do not consider few-bit quantization in their objective, we achieve signiﬁcantly lower pruning rates.
This is an interesting observation since our method is based on a similar objective (e.g. compared to Sparse VD) but with the addition of forcing nonpruned weights to tightly cluster around the quantization levels.
Few-bit quantization severely limits network capacity.
Perhaps this capacity limitation must be countered by pruning fewer weights.
Our pruning rates are roughly in line with other papers on ternary quantization, e.g. Zhu et al.
(2016), who report sparsity levels between 30% and 50% with their ternary quantization method.
Note that a direct comparison between pruning, quantizing and ternarizing methods is difﬁcult and depends on many factors such that a fair computation of the compression rate that does not implicitly favor certain methods is hardly possible within the scope of this paper.
For instance, compression rates for pruning methods are typically reported under the assumption of a CSC storage format which would not fully account for the compression potential of a sparse ternary matrix.
We thus chose not to report any measures for compression rates, however for the methods listed in Table 1, they can easily be found in the literature. [Table]
4. D ENSE N ET -121 ON CIFAR-
Our second experiment uses a modern DenseNet-121 (Huang et al., 2017)
( k =12, with bottleneck) on CIFAR-10 (Krizhevsky & Hinton, 2009).
The training procedure is identical to the procedure on
MNIST with the following exceptions: we use a batch-size of 64 samples, the warm-up weight β of the KL term is 0 for the ﬁrst 5 epochs and is then linearly ramped up from 0 to 1 over the nextepochs, the learning rate of 0 .
005 is kept constant for the ﬁrst 50 epochs and then linearly decreased to a value of 0 . 003 when training stops after 150 epochs.
We pre-train a deterministic DenseNet (reaching validation accuracy of 93 .
to initialize VNQ training.
The codebook parameter for non-zero values a is initialized with the maximum absolute value over pre-trained weights per layer.
Results are shown in Table 2.
A visualization of the distribution over weights after VNQ training is shown in the Appendix Fig. [Table] [Table]
VNQ + P&Q (w/o 1 )
[Table]
[Table] VNQ + P&Q [Table] We generally observe lower levels of sparsity for DenseNet, compared to LeNet.
This might be due to the fact that DenseNet already has an optimized architecture which removed a lot of redundant parameters from the start.
In line with previous publications we generally observed that the ﬁrst and last layer of the network are most sensitive to pruning and quantization.
However, in contrast to many other methods that do not quantize these layers (e.g. Zhu et al.
, we ﬁnd that after sufﬁcient training, the complete network can be pruned and quantized with very little additional loss in accuracy (see Table 2). Inspecting the weight scatter-plot for the ﬁrst and last layer (Figure 3, topleft and bottom-right panel) it can be seen that some weights did not settle on one of the prior modes (the “funnels”) after VNQ training, particularly the ﬁrst layer has a few such weights with very low variance. It is likely that quantizing these weights causes the additional loss in accuracy that we observe when quantizing the whole network. Without gradient stopping (i.e. applying gradients to a shadow weight at the trainable clipping boundary) we have observed that pruning and quantizing the ﬁrst layer leads to a more pronounced drop in accuracy (about 3% compared to a network where the ﬁrst layer is kept with full precision, not shown in results).
R ELATED WORK
Our method is an extension of Sparse VD (Molchanov et al., 2017), originally used for network pruning.
In contrast we use a quantizing prior, leading to a multi-modal posterior suitable for fewbit quantization and pruning.
Bayesian Compression and Structured Bayesian Pruning (Louizos et al., 2017; Neklyudov et al., 2017) extend Sparse VD to prune whole neurons or ﬁlters via groupsparsity constraints.
Additionally, in Bayesian Compression the required bit-precision per layer is determined via the posterior variance.
In contrast to our method, Bayesian Compression does not explicitly enforce clustering of weights during training and thus requires bit-widths in the range between 5 and 18 bits.
Extending our method to include group-constraints for pruning is an interesting direction for future work.
Another Bayesian method for simultaneous network quantization and pruning is soft weight-sharing (SWS) (Ullrich et al., 2017), which uses a Gaussian mixture model prior (and a KL term without trainable parameters such that the KL term reduces to the prior entropy).
SWS acts like a probabilistic version of k-means clustering with the advantage of automatic collapse of unnecessary mixture components.
Similar to learning the codebooks in our method, soft weight-sharing learns the prior from the data, a technique known as empirical Bayes (see also ARD (Karaletsos & R ¨ atsch, 2015)).
We cannot directly compare against soft weight-sharing since the authors do not report results on ternary networks.
Gal et al.
(2017) learn dropout rates by using a continuous relaxation of dropout’s discrete masks (via the concrete distribution).
The authors learn layer-wise dropout rates, which dos not allow for dropout-rate-based pruning.
We have experimented with using the concrete distribution for learning codebooks for quantization with promising early results but we have generally observed lower pruning rates or lower accuracy compared to VNQ.
A non-probabilistic state-of-the-art method for network ternarization is Trained Ternary
Quantization (Zhu et al., 2016) which uses full-precision shadow weights during training, but quantized forward passes.
Additionally it learns a (non-symmetric) scaling per layer for the non-zero quantization values, similar to our learned quantization level a .
While the method achieves impressive accuracy, the sparsity and thus pruning rates are rather low (between 50% and 70% sparsity) and the ﬁrst and last layer need to be kept with full precision.
D ISCUSSION
A potential shortcoming of our method is the KL divergence approximation (Sec. 3.3).
While the approximation is reasonably good on the relevant range of θ - and σ -values, there is still room for improvement which could have the beneﬁt that weights are drawn even more tightly onto the quantization levels, resulting in lower accuracy loss after quantization and pruning.
Since our functional approximation to the KL divergence only needs to be computed once and an arbitrary amount of ground-truth data can be produced, it should be possible to improve upon the approximation presented here at least by some brute-force function approximation, e.g. a neural network, polynomial or kernel regression.
The main difﬁculty is that the resulting approximation must be differentiable and must not introduce signiﬁcant computational overhead since the approximation is evaluated once for each network parameter in each gradient step. We have also experimented with a naive
Monte-Carlo approximation of the KL divergence term.
This has the disadvantage that local reparameterization (where pre-activations are sampled directly) can no longer be used, since weight samples are required for the MC approximation.
To keep computational complexity comparable, we used a single sample for the MC approximation.
In our LeNet-5 on MNIST experiment the MC approximation achieves comparable accuracy with higher pruning rates compared to our functional KL approximation. However, with DenseNet-121 on CIFAR-10 and the MC approximation validation accuracy plunges catastrophically after pruning and quantization. See Section A.3 in the Appendix for more details. Compared to similar methods that only consider network pruning, our pruning rates are signiﬁcantly lower.
This does not seem to be a particular problem of our method since other papers on network ternarization report similar or even lower sparsity levels (Li et al.
(2016) roughly achieve between 30% and 50% sparsity).
The reason for this might be that heavily quantized networks have much lower capacity compared to full-precision networks.
This limited capacity might require that the network compensates by effectively using more weights such that the pruning rates become signiﬁcantly lower.
Similar trends have also been observed with binary networks, where drops in accuracy could be prevented by increasing the number of neurons (with binary weights) per layer.
Principled experiments to test the trade-off between low bit-precision and sparsity rates would be an interesting direction for future work.
One starting point could be to test our method with more quantization levels (e.g. 5 , 7 or 9 ) and investigate how this affects the pruning rate
A. L OCAL R EPARAMETERIZATION
We follow Sparse VD (Molchanov et al., 2017) and use the Local reparameterization trick (Kingma et al., 2015) and Additive Noise Reparmetrization to optimize the stochastic gradient variational lower bound L SGVB (Eq.
We optimize posterior means and log-variances ( θ, log σ 2 ) and in some experiments additionally the codebook level a .
We apply Variational Network Quantization to fully connected and convolutional layers.
Denoting inputs to a layer with A M × I outputs of a layer with B M × O and using local reparameterization we get:
[Equation]
for a fully connected layer.
Similarly activations for a convolutional layer are computed as follows
[Equation]
where ( · ) 2 denotes an element-wise operation, ∗ is the convolution operation and vec( · ) denotes reshaping of a matrix/tensor into a vector.
A. KL APPROXIMATION FOR QUANTIZING PRIOR
Under the quantizing prior (Eq.
the KL divergence between the mean-ﬁeld posterior and prior
D KL ( q φ ( w ij ) || p ( w ij )) is analytically intractable.
Molchanov et al.
(2017) presented an approximation for the KL divergence under a (zero-centered) log uniform prior (Eq.
Since our quantizing prior is essentially a composition of shifted log uniform priors, we construct a composition of the approximation given by Molchanov et al.
(2017), shown in Eq.
The original approximation can be utilized to calculate a KL divergence approximation (up to an additive constant ˜ C ) from a shifted log-uniform prior p ( w ij )
∝ 1 | w ij − r | to a Gaussian posterior q φ ( w ij ) by transferring the shift to the posterior parameter θ
[Equation]
For small posterior variances σ 2 ij ( σ ij (cid:28) r ) and means near the quantization levels (i.e, | θ ij | ≈ r ), the KL divergence is dominated by the mixture prior component located at the respective quantization level r .
For these values of θ and σ , the KL divergence can be approximated by shifting the approximation F LU , KL ( θ, σ ) to the quantization level r , i.e. F LU , KL ( θ ± r, σ ) .
For small σ and values of θ near zero or far away from any quantization level, as well as for large values of σ and arbitrary θ , the KL divergence can be approximated by the original non-shifted approximation
F LU , KL ( θ, σ ) .
Based on these observations we construct our KL approximation by properly mixing shifted versions of F LU , KL ( θ ± r, σ ) .
We use Gaussian window functions Ω( θ ± r ) to perform this weighting (to ensure differentiability).
The remaining θ domain is covered by an approximation located at zero and weighted such that this approximation is dominant near zero and far away from the quantization levels, which is achieved by introducing the constraint that all window functions sum up to one on the full θ domain.
See Fig.
2 for a visual representation of shifted approximations and their respective window functions.
A.3. A PPROXIMATION QUALITY
We evaluate the quality of our KL approximation (Eq.
by comparing against a ground-truth
Monte Carlo approximation on a dense grid over the full range of relevant θ and σ values.
Results of this comparison are shown in Fig.
Alternatively to the functional KL approximation, one could also use a naive Monte Carlo approximation directly. This has the disadvantage that local reparameterization can no longer be used, since actual samples of the weights must be drawn. To assess the quality of our functional KL approximation, we also compare against experiments where we use a naive MC approximation of the KL divergence term, where we only use a single sample for approximating the expectation to keep computational complexity comparable to our original method. Note that the “ground-truth” MC approximation used before to evaluate KL approximation quality uses many more samples which would be prohibitively expensive during training.
To test for the effect of local reparameterization in isolation we also show results for our functional KL approximation without using local reparameterization.
The results in Table 3 show that the naive MC approximation of the KL term leads to slightly lower validation error on MNIST (LeNet-5) (with higher pruning rates) but on CIFAR-10 (DenseNet) the validation error of the network trained with the naive MC approximation catastrophically increases after pruning and quantizing the network.
Except for removing local reparameterization or plugging in the naive MC approximation, experiments are run as described in Section 4.
Setting
[Table]
LeNet
[Table]
MNIST
[Table]
no local reparam,
[Table]
approx
[Table]
no local reparam, naive MC approx
[Table]
DenseNet-121 on CIFARlocal reparam
[Table]
func.
KL
[Table]
[Table]
Inspecting the distribution over weights after training with the naive MC approximation for the KL divergence, shown in Figure 5 for LeNet-5 and in Figure 6 for DenseNet-121, reveals that weightmeans tend to be more dispersed and weight-variances tend to be generally lower than when training with our functional KL approximation (compare Figure 1 for LeNet-5 and Figure 3 for DenseNet121).
We speculate that the combined effects of missing local reparameterization and single-sample
MC approximation lead to more noisy gradients.
[Figure]
A. R EUSING THE KL APPROXIMATION FOR ARBITRARY CODEBOOKS
We show that the KL approximation (Eq.
(16)), developed for a ﬁxed reference codebook, can be reused for arbitrary codebooks as long as codebook learning is restricted to learning a multiplicative scaling factor.
Without loss of generality we consider the case of ternary, symmetric codebooks
[Equation]
(19) where r ∈ R + is the quantization level value and p c r denotes a sparsity-inducing, quantizing prior over weights (sparsity is induced because one of the codebook entries is ﬁxed to 0 ).
We denote c r as the reference codebook for which we design the KL approximation D KL ( q φ ( w ) || p c r ) =
F KL ( θ, σ, c r ) (Eq. (16)).
This approximation can be reused for any symmetric ternary codebook c a =
[ − a, 0 , a ] with quantization level a ∈ R + .
The latter can be seen by representing c a with the reference codebook and a positive scaling factor s > 0 as c a = sc r , s = a/r .
This re-scaling translates into a multiplicative re-scaling of the variational parameters θ and σ .
To see this, consider the prior p c a , based on codebook c a :
[Equation]
The KL divergence from a prior based on the codebook c a to the posterior q φ ( w ) is given by
[Equation]
Since q θ ( sz ) is Gaussian, the scaling s can be transfered into the variational parameters φ
=
( θ, [Equation]
[Equation]
with ˆ φ = ( θ s , σ s ) .
Inserting into Eq.
(21) yields:
[Equation]
Thus, D KL ( q φ ( w ) || p c a ( w )) =
D KL ( q ˆ φ ( w ) ||
p c r ( w ))
+ C ≈ F KL ( θ/s, σ/s, c r ) , where F KL is given by Eq.
This means that the KL approximation can be used for arbitrary ternary, symmetric codebooks of the form c a =
[ − a, 0 , a ] = sc r because the scaling factor s translates into a re-scaling of the variational parameters ˆ φ = ( θs , σs ) .