Relaxed Quantization for Discretized Neural Networks

1 INTRODUCTION . Neural networks excel in a variety of large scale problems due to their highly flexible parametric nature . However , deploying big models on resource constrained devices , such as mobile phones , drones or IoT devices is still challenging because they require a large amount of power , memory and computation . Neural network compression is a means to tackle this issue and has therefore become an important research topic . Neural network compression can be , roughly , divided into two not mutually exclusive categories : pruning and quantization . While pruning ( LeCun et al. , 1990 ; Han et al. , 2015 ) aims to make the model “ smaller ” by altering the architecture , quantization aims to reduce the precision of the arithmetic operations in the network . In this paper we focus on the latter . Most network quantization methods either simulate or enforce discretization of the network during training , e.g . via rounding of the weights and activations . Although seemingly straighforward , the discontinuity of the discretization makes the gradient-based optimization infeasible . The reason is that there is no gradient of the loss with respect to the parameters . A workaround to the discontinuity are the “ pseudo-gradients ” according to the straight-through estimator ( Bengio et al. , 2013 ) , which have been successfully used for training low-bit width architectures at e.g . Hubara et al . ( 2016 ) ; Zhu et al . ( 2016 ) . The purpose of this work is to introduce a novel quantization procedure , Relaxed Quantization ( RQ ) . RQ can bypass the non-differentiability of the quantization operation during training by smoothing it appropriately . The contributions of this paper are four-fold : First , we show how to make the set of quantization targets part of the training process such that we can optimize them with gradient descent . Second , we introduce a way to discretize the network by converting distributions over the weights and activations to categorical distributions over the quantization grid . Third , we show that we can obtain a “ smooth ” quantization procedure by replacing the categorical distributions with ∗Work done while interning at Qualcomm AI Research . concrete ( Maddison et al. , 2016 ; Jang et al. , 2016 ) equivalents . Finally we show that stochastic rounding ( Gupta et al. , 2015 ) , one of the most popular quantization techniques , can be seen as a special case of the proposed framework . We present the details of our approach in Section 2 , discuss related work in Section 3 and experimentally validate it in Section 4 . Finally we conclude and provide fruitful directions for future research in Section 5 . 2 RELAXED QUANTIZATION FOR DISCRETIZING NEURAL NETWORKS . The central element for the discretization of weights and activations of a neural network is a quantizer q ( · ) . The quantizer receives a ( usually ) continous signal as input and discretizes it to a countable set of values . This process is inherently lossy and non-invertible : given the output of the quantizer , it is impossible to determine the exact value of the input . One of the simplest quantizers is the rounding function : q ( x ) = α ⌊ x α + 1 2 ⌋ , where α corresponds to the step size of the quantizer . With α = 1 , the quantizer rounds x to its nearest integer number . Unfortunately , we can not simply apply the rounding quantizer to discretize the weights and activations of a neural network . Because of the quantizers ’ lossy and non-invertible nature , important information might be destroyed and lead to a decrease in accuracy . To this end , it is preferable to train the neural network while simulating the effects of quantization during the training procedure . This encourages the weights and activations to be robust to quantization and therefore decreases the performance gap between a full-precision neural network and its discretized version . However , the aforementioned rounding process is non-differentiable . As a result , we can not directly optimize the discretized network with stochastic gradient descent , the workhorse of neural network optimization . In this work , we posit a “ smooth ” quantizer as a possible way for enabling gradient based optimization . 2.1 LEARNING ( FIXED POINT ) QUANTIZERS VIA GRADIENT DESCENT . The proposed quantizer comprises four elements : a vocabulary , its noise model and the resulting discretization procedure , as well as a final relaxation step to enable gradient based optimization . The first element of the quantizer is the vocabulary : it is the set of ( countable ) output values that the quantizer can produce . In our case , this vocabulary has an inherent structure , as it is a grid of ordered scalars . For fixed point quantization the grid G is defined as G = [ −2b−1 , . . . , 0 , . . . , 2b−1 − 1 ] , ( 1 ) where b is the number of available bits that allow for K = 2b possible integer values . By construction this grid of values is agnostic to the input signal x and hence suboptimal ; to allow for the grid to adapt to x we introduce two free parameters , a scale α and an offset β . This leads to a learnable grid via Ĝ = αG + β that can adapt to the range and location of the input signal . The second element of the quantizer is the assumption about the input noise ; it determines how probable it is for a specific value of the input signal to move to each grid point . Adding noise to x will result in a quantizer that is , on average , a smooth function of its input . In essense , this is an application of variational optimization ( Staines & Barber , 2012 ) to the non-differentiable rounding function , which enables us to do gradient based optimization . We model this form of noise as acting additively to the input signal x and being governed by a distribution p ( ) . This process induces a distribution p ( x̃ ) where x̃ = x+ . In the next step of the quantization procedure , we discretize p ( x̃ ) according to the quantization grid Ĝ ; this neccesitates the evaluation of the cumulative distribution function ( CDF ) . For this reason , we will assume that the noise is distributed according to a zero mean logistic distribution with a standard deviation σ , i.e . L ( 0 , σ ) , hence leading to p ( x̃ ) = L ( x , σ ) . The CDF of the logistic distribution is the sigmoid function which is easy to evaluate and backpropagate through . Using Gaussian distributions proved to be less effective in preliminary experiments . Other distributions are conceivable and we will briefly discuss the choice of a uniform distribution in Section 2.3 . The third element is , given the aforementioned assumptions , how the quantizer determines an appropriate assignment for each realization of the input signal x . Due to the stochastic nature of x̃ , a deterministic round-to-nearest operation will result in a stochastic quantizer for x. Quantizing x in this manner corresponds to discretizing p ( x̃ ) onto Ĝ and then sampling grid points gi from it . More specifically , we construct a categorical distribution over the grid by adopting intervals of width equal to α centered at each of the grid points . The probability of selecting that particular grid point will now be equal to the probability of x̃ falling inside those intervals : p ( x̂ = gi|x , σ ) = P ( x̃ ≤ ( gi + α/2 ) ) − P ( x̃ < ( gi − α/2 ) ) ) ( 2 ) = Sigmoid ( ( gi + α/2− x ) /σ ) − Sigmoid ( ( gi − α/2− x ) /σ ) , ( 3 ) where x̂ corresponds to the quantized variable , P ( · ) corresponds to the CDF and the step from Equation 2 to Equation 3 is due to the logistic noise assumption . A visualization of the aforementioned process can be seen in Figure 1 . For the first and last grid point we will assume that they reside within ( g0 − α/2 , g0 + α/2 ] and ( gK − α/2 , gK + α/2 ] respectively . Under this assumption we will have to truncate p ( x̃ ) such that it only has support within ( g0 − α/2 , gK + α/2 ] . Fortunately this is easy to do , as it corresponds to just a simple modification of the CDF : P ( x̃ ≤ c|x̃ ∈ ( g0 − α/2 , gK + α/2 ] ) = P ( x̃ ≤ c ) − P ( x̃ < ( g0 − α/2 ) ) P ( x̃ ≤ ( gK + α/2 ) ) − P ( x̃ < ( g0 − α/2 ) ) . ( 4 ) Armed with this categorical distribution over the grid , the quantizer proceeds to assign a specific grid value to x̂ by drawing a random sample . This procedure emulates quantization noise , which prevents the model from fitting the data . This noise can be reduced in two ways : by clustering the weights and activations around the points of the grid and by reducing the logistic noise σ . As σ → 0 , the CDF converges towards the step function , prohibiting gradient flow . On the other hand , if is too high , the optimization procedure is very noisy , prohibiting convergence . For this reason , during optimization we initialize σ in a sensible range , such that L ( x , σ ) covers a significant portion of the grid . Please confer Appendix A for details . We then let σ be freely optimized via gradient descent such that the loss is minimized . Both effects reduce the gap between the function that the neural network computes during training time vs. test time . We illustrate this in Figure 2 . The fourth element of the procedure is the relaxation of the non-differentiable categorical distribution sampling . While we can use an unbiased gradient estimator via REINFORCE ( Williams , 1992 ) , we opt for a continuous relaxation due to high variances with REINFORCE . This is achieved by replacing the categorical distribution with a concrete distribution ( Maddison et al. , 2016 ; Jang et al. , 2016 ) . This relaxation procedure corresponds to adopting a “ smooth ” categorical distribution that can be seen as a “ noisy ” softmax . Let πi be the categorical probability of sampling grid point i , i.e . πi = p ( x̂ = gi ) ; the “ smoothed ” quantized value x̂ can be obtained via : ui ∼ Gumbel ( 0 , 1 ) , zi = exp ( ( log πi + ui ) /λ ) ∑ j exp ( ( log πj + uj ) /λ ) , x̂ = K∑ i=1 zigi , ( 5 ) where zi is the random sample from the concrete distribution and λ is a temperature parameter that controls the degree of approximation , since as λ→ 0 the concrete distribution becomes a categorical . We have thus defined a fully differentiable “ soft ” quantization procedure that allows for stochastic gradients for both the quantizer parameters α , β , σ as well as the input signal x ( e.g . the weights or the activations of a neural network ) . We refer to this algorithm as Relaxed Quantization ( RQ ) . We summarize its forward pass as performed during training in Algorithm 1 . It is also worthwhile to notice that if there were no noise at the input x then the categorical distribution would have non-zero mass only at a single value , thus prohibiting gradient based optimization for x and σ . One drawback of this approach is that the smoothed quantized values defined in Equation 5 do not have to coincide with grid points , as z is not a one-hot vector . Instead , these values can lie anywhere between the smallest and largest grid point , something which is impossible with e.g . stochastic rounding ( Gupta et al. , 2015 ) . In order to make sure that only grid-points are sampled , we propose an alternative algorithm RQ ST in which we use the variant of the straight-through ( ST ) estimator proposed in Jang et al . ( 2016 ) . Here we sample the actual categorical distribution during the forward pass but assume a sample from the concrete distribution for the backward pass . While this gradient estimator is obviously biased , in practice it works as the “ gradients ” seem to point towards a valid direction . This effect was also recently studied at Yin et al . ( 2019 ) . We perform experiments with both variants . After convergence , we can obtain a “ hard ” quantization procedure , i.e . select points from the grid , at test time by either reverting to a categorical distribution ( instead of the continuous surrogate ) or by rounding to the nearest grid point . In this paper we chose the latter as it is more aligned with the low-resource environments in which quantized models will be deployed . Furthermore , with this goal in mind , we employ two quantization grids with their own learnable scalar α , σ ( and potentially β ) parameters for each layer ; one for the weights and one for the activations .

Summary: This paper introduces a method for learning neural networks with quantized weights and activations. The main idea is to stochastically – rather than deterministically – quantize values, and to replace the resulting categorical distribution over quantized values with a continuous relaxation (the "concrete distribution" or "Gumbel-Softax distribution"; Maddison et al., 2016; Jang et al., 2016). Good empirical performance is demonstrated for LeNet-5 applied to MNIST, VGG applied to CIFAR-10, and MobileNet and ResNet-18 applied to ImageNet.