Title: MIXED PRECISION DNNS: ALL YOU NEED IS A GOOD PARAMETRIZATION

Abstract: Efficient deep neural network (DNN) inference on mobile or embedded devices typically involves quantization of the network parameters and activations. In particular, mixed precision networks achieve better performance than networks with homogeneous bitwidth for the same size constraint. Since choosing the optimal bitwidths is not straight forward, training methods, which can learn them, are desirable. Differentiable quantization with straight-through gradients allows to learn the quantizer's parameters using gradient methods. We show that a suited parametrization of the quantizer is the key to achieve a stable training and a good final performance. Specifically, we propose to parametrize the quantizer with the step size and dynamic range. The bitwidth can then be inferred from them. Other parametrizations, which explicitly use the bitwidth, consistently perform worse. We confirm our findings with experiments on CIFAR-10 and ImageNet and we obtain mixed precision DNNs with learned quantization parameters, achieving state-of-the-art performance.

1. INTRODUCTION:

Quantized DNNs apply quantizers Q : R → {q 1 , ..., q I } to discretize the weights and/or activations of a DNN (Han et al., 2015; Zhou et al., 2017; Li et al., 2016; Liu & Mattina, 2019; Cardinaux et al., 2018; Jain et al., 2019; Bai et al., 2018) . They require considerably less memory and have a lower computational complexity, since discretized values {q 1 , ..., q I } can be stored, multiplied and accumulated efficiently. This is particularly relevant for inference on mobile or embedded devices with limited computational power.

However, gradient based training of quantized DNNs is difficult, as the gradient of a quantization function vanishes almost everywhere, i.e., backpropagation through a quantized DNN almost always returns a zero gradient. Different solutions to this problem have been proposed in the literature: A first possibility is to use DNNs with stochastic weights from a categorical distribution and to optimize the evidence lower bound (ELBO) to obtain an estimate of the posterior distribution of the weights. As proposed in (Jang et al., 2016; Maddison et al., 2016; Louizos et al., 2019) , the categorical distribution can be relaxed to a concrete distribution -a smoothed approximation of the categorical distribution -such that the ELBO becomes differentiable under reparametrization. A second possibility is to use the straight through estimator (STE) (Bengio et al., 2013) . STE allows the gradients to be backpropagated through the quantizers and, thus, the network weights can be adapted with standard gradient descent (Hubara et al., 2016) . Compared to STE based methods, stochastic methods suffer from large gradient variance, which makes training of large quantized DNNs difficult. Therefore, STE based methods are more popular in practice.

More recent research (Jain et al., 2019; Esser et al., 2019; Wang et al., 2018; Elthakeb et al., 2018) focuses on methods which can also learn the optimal quantization parameters, e.g., the stepsize, dynamic range and bitwidth, in parallel to the network weights. This is a promising approach as DNNs with learned quantization parameters almost always outperform DNNs with handcrafted ones.

Recently, and in parallel to our work, (Jain et al., 2019) explored the use of STE to define the gradient with respect to the quantizers's dynamic range. The authors applied a per-tensor quantization and used the dynamic range as an additional trainable parameter also learned with gradient descent. Similarly, (Esser et al., 2019) learned the stepsize using gradient descent. However, neither of them learned the optimal bitwidth of the quantizers. One approach was proposed in (Wang et al., 2018; Elthakeb et al., 2018) . They learn the bitwidth with reinforcement learning, i.e., they learn an optimal bitwidth assignment policy. Their experiments show that a DNN with a learned and heterogeneous bitwidth assignment outperforms quantized DNNs with a homogeneous bitwidth assignment. However, such methods have a high computational complexity as the bitwidth policy must be learned, which involves training many quantized DNNs.

In this paper, we will use the STE approach and show that the quantizer's parameters, including the bitwidth, can be learned with gradient methods if a good parametrization is chosen. Specifically, we show that directly learning the bitwidth is not optimal. Instead, we propose to learn the stepsize and dynamic range. The bitwidth can then be inferred from them. Compared to (Wang et al., 2018; Elthakeb et al., 2018) , our method has the advantage that training quantized DNNs has nearly the same computational complexity as standard float32 training.

The contributions of this paper are:

1. We show that there are three different parametrizations for uniform and power-of-two quantization and that, in both cases, one of them has gradients particularly well suited to train quantized DNNs. The other parametrizations have the problem of yielding gradients with an unbounded gradient norm and coupled components.

2. Using this parametrization, we are able to learn all quantization parameters for DNNs with per-tensor quantization and global memory constraints. We formulate the training as a constrained optimization problem, where the quantized DNN is constrained not to exceed a given overall memory budget, and show how to solve it in a penalty framework.

3. We confirm our findings with experiments on CIFAR-10 and ImageNet. For example, we train a heterogeneously quantized MobileNetV2 on ImageNet requiring a total of only 1.65MB to store the weights and only 0.57MB to store its largest feature map. This is equivalent to a homogenous 4bit quantization of both weights and activations. However, our network learns to allocate the bitwidth heterogeneously in an optimal way. Our MobileNetV2 achieves an error of 30.26% compared to 29.82% for the floating point baseline. This is state-of-the-art for such a heavily quantized MobileNetV2.

We use the following notation throughout this paper: x, x, X and X denote a scalar, a (column) vector, a matrix and a tensor with three or four dimensions, respectively; . and . are the floor and ceiling operators. Finally, δ(.) denotes the Dirac delta function.

2. CHOOSING A QUANTIZATION PARAMETRIZATION:

Let Q(x; θ) be a quantizer with the parameters θ, which maps x ∈ R to discrete values {q 1 , ..., q I }. In this section, we compare different parametrizations of Q(x; θ) for uniform quantization and powerof-two quantization and analyze how well the corresponding straight-through gradient estimates ∂ x Q(x; θ) and ∇ θ Q(x; θ) are suited to optimize the quantizer parameters θ. Our key result is, that the training of quantized DNNs which learns both, the optimal quantized weights and the optimal quantization parameters θ, is very sensitive to the choice of the parametrization of the quantizers. From an optimization point of view, it is best to parametrize the quantizer Q(x; θ) with the stepsize d and the dynamic range q max as it leads to gradients with stable norms. Doing so, we can use standard gradient descent to learn the quantization parameters and do not need to use stochastic or reinforcement based algorithms, which are computationally expensive.

2.1. PARAMETRIZATION AND STRAIGHT THROUGH GRADIENT ESTIMATES:

A symmetric uniform quantizer Q U (x; θ) which maps a real value x ∈ R to one of I = 2k + 1 quantized values q ∈ {-kd, ..., 0, ..., kd} computes 

EQUATION

1: Maximum gradient norm max x ∇ θ Q U (x; θ)

. For "U1" and "U2" the maximum gradient norm can grow exponentially with varying bitwidth b whereas it is bounded for "U3".

where the parameter vector θ = [d, q max , b] T consists of the step size d ∈ R, the maximum value q max ∈ R and the number of bits b ∈ N, b ≥ 2 used to encode the quantized values q.

When training quantized DNNs, we want to optimize Q U (x; θ) with respect to the input x and the quantization parameters θ, meaning that we need the gradients ∇ x Q U (x; θ) and ∇ θ Q(x; θ). A common problem is, that the exact gradients are not useful for training. For example,

∂ x Q U (x; θ) = d 2 b-1 -2 k=-2 b-1 +1 δ x -d k + 1 2

vanishes almost everywhere. A solution is to define the derivative using STE (Bengio et al., 2013) , which ignores the floor operation in (1). This leads to

EQUATION

which is non-zero in the interesting region |x| ≤ q max and which turned out to be very useful to train quantized DNNs in practice (Yin et al., 2019) . In this work, we follow this idea and define the gradients ∇ x Q(x; θ) and ∇ θ Q(x; θ), using STE whenever we need to differentiate a floor operation. We refer to this as differentiable quantization (DQ).

An important observation from ( 1) is that the parameters θ = [d, q max , b] T of a quantizer depend on each other, i.e., q max = (2 b-1 -1)d. This means, that we can choose from three equivalent parametrizations of Q U (x; θ): Case "U1" with θ = [b, d] T , case "U2" with θ = [b, q max ] T and case "U3" with θ = [d, q max ] T . Interestingly, they differ in their gradients: Case U1: Parametrization with respect to θ = [b, d] T , using q max = q max (b, d) gives

∇ θ Q U (x; θ) = ∂ b Q U (x; θ) ∂ d Q U (x; θ) =        0 1 d (Q U (x; θ) -x) |x| ≤ (2 b-1 -1)d 2 b-1 log(2)d 2 b-1 -1 sign(x) |x| > (2 b-1 -1)d (3a)

Case U2: Parametrization with respect to θ = [b, q max ] T , using d = d(b, q max ) gives

∇ θ Q U (x; θ) = ∂ b Q U (x; θ) ∂ qmax Q U (x; θ) =          -2 b-1 log 2 2 b-1 -1 1 qmax (Q U (x; θ) -x) |x| ≤ q max 0 sign(x) |x| > q max (3b)

Case U3: Parametrization with respect to θ = [d, q max ] T , using b = b(d, q max ) gives

∇ θ Q U (x; θ) = ∂ d Q U (x; θ) ∂ qmax Q U (x; θ) =        1 d 0 (Q U (x; θ) -x) |x| ≤ q max 0 sign(x) |x| > q max (3c)

Fig. 1 shows the maximum gradient norm max x ∇ θ Q U (x; θ) for the three parametrizations "U1" to "U3". For the parametrizations "U1" and "U2", max x ∇ θ Q U (x; θ) can grow exponentially with varying bitwidth b as

∂ d Q U (x; θ) ∈ [-2 b-1 -1, 2 b-1 -1] for "U1" and ∂ b Q U (x; θ) ∈ -qmax 2 b-1 -1 log 2, qmax 2 b-1 -1 log 2

for "U2". This is not desirable when training quantized DNNs, because it will lead to large changes of the gradient norm and forces us to use small learning rates to

q max 2 b-1 -1 x ∂ b Q U (x) ∂ d Q U (x) (a) Case U1 q max 1 x ∂ b Q U (x) ∂ q max Q U (x) (b) Case U2 q max 1 x ∂ d Q U (x) ∂ q max Q U (x) (c) Case U3 q max 1 x ∂ x Q U (x) (d) Input derivative

Figure 2 : Partial derivatives of Q U (x; θ) with respect to the input and the quantization parameters d, q max and b. Partial derivatives are coupled for "U1" and "U2" but are decoupled for "U3". avoid divergence. However, parametrization "U3" does not suffer from such an unbounded gradient norm as both partial derivatives 2 shows the gradients for the parametrization "U1" to "U3". For parametrization "U3", the partial derivatives in ∇ θ Q U (x; θ) are decoupled, i.e., ∇ θ Q U (x; θ) is a unit vector, which either points only in the direction of d if |x| ≤ q max or only in the direction of q max , if |x| > q max . We will show in Sec. 2.3 that this implies a diagonal Hessian, which results in a better convergence behavior of gradient descent. In contrast, both parametrizations "U1" and "U2" have partial derivatives that are coupled. In summary, this implies that parametrization "U3" is the best DQ parametrization.

∂ d Q U (x; θ) ∈ [-1 2 , 1 2 ] and ∂ qmax Q U (x; θ) ∈ {-1, 1} are bounded. Fig.

Similar considerations can be made for the power-of-two quantization Q P (x; θ), which maps a real-valued number x ∈ R to a quantized value q ∈ {±2 k : k ∈ Z} by

EQUATION

where q min and q max are the minimum and maximum absolute values of the quantizer for a bitwidth of b bit. Power-of-two quantization is an especially interesting scheme for DNN quantization, since a multiplication of quantized values can be implemented as an addition of the exponents. Using the STE for the floor operation, the derivative ∂ x Q P (x; θ) is given by

EQUATION

)

The power-of-two quantization has the three parameters [b, q min , q max ] =: θ, which depend on each other with the relationship q max = 2 2 b-1 -1 q min . Therefore, we have again three different parametrizations with θ = [b, q min ], θ = [b, q max ] or θ = [q min , q max ], respectively. Similarly to the uniform case, one parametrization (θ = [q min , q max ]) leads to a gradient of a very simple form

EQUATION

which has again a bounded gradient magnitude and independent components and is, hence, best suited for first order gradient based optimization.

2.2. CONSTRAINTS ON θ:

In practice, for an efficient hardware implementation, we need to ensure that the quantization parameters only take specific discrete values: for uniform quantization, only integer values are allowed for the bitwidth b, and the stepsize d must be a power-of-two, see e.g. (Jain et al., 2019) ; for power-of-two quantization, the bitwidth must be an integer, and the minimum and maximum absolute values q min and q max must be powers-of-two.

We fulfill these constraints by rounding the parameters in the forward pass to the closest integer or power-of-two value. In the backward pass we update the original float values, i.e., we used again the STE to propagate the gradients.

2.3. EXPERIMENTAL COMPARISON OF DQ PARAMETRIZATIONS:

In the following we compare the parametrizations using two experiments. b, q min (P2) q min , qmax (P3) (b) Power-of-two quantization Figure 3 : MSE for quantizing Gaussian data x ∼ N (0, 1) with uniform and power-of-two quantization. Parametrizations "U3" and "P3" converge to the lowest MSE without any oscillations.

1) Quantization of Gaussian data:

In our first experiment we use DQ to learn the optimal quantization parameters θ * which minimize the mean squared error (MSE) E 1 2 (Q(x; θ) -x) 2 with gradient descent and compare the convergence speed for three possible parametrizations of a uniform and power-of-two quantizer. We choose this example as the gradient

∇ θ Q(x; θ) = E [(Q(x; θ) -x)∇ θ Q(x; θ)] is just a scaled version of ∇ θ Q(x; θ), i.e.

, the gradient direction depends directly on the parametrization of Q(x; θ) and thus the effects of changing the parametrization can be observed.

It is interesting to study the Hessian:

H = ∇ θ ∇ T θ E (Q(x; θ) -x) 2 ∈ R 2×2 of the MSE: H = E ∇ θ Q(x; θ)∇ θ Q(x; θ) T + (Q(x; θ) -x)∇ θ ∇ T θ Q(x; θ) ≈ E ∇ θ Q(x; θ)∇ θ Q(x; θ) T .

Note that we use the outer-product approximation (Bishop, 2006) in order to simplify our considerations. From this equation it is apparent that the Hessian will be diagonal for the case U3 as

∇ θ Q(x; θ)∇ θ Q(x; θ) T only contains an element in either (1, 1) or (2, 2) and, therefore, E ∇ θ Q(x; θ)∇ θ Q(x; θ) T is a diagonal matrix.

From this, we can see that gradient descent with an individual learning rate for each parameter is equivalent to Newton's method and, therefore, efficient. In general this will not be the case for U1 and U2.

We conduct an experiment, using ADAM to optimize the mean squared quantization error on artificially generated data, which is generated by drawing 10 4 samples from N (0, 1). Please note that the same example was studied in (Jain et al., 2019) . The results in Fig. 3 clearly show that the parametrizations "U3" and "P3" are best suited to optimize the uniform and power-of-two quantization parameters, respectively. Indeed, these quantizers converge without oscillation to the lowest MSE. It is interesting to see, that even adaptive gradient methods like ADAM can not solve the scaling issue described in Sec. 2.1. In the Appendix A.4 we give further empirical evidence to support this claim and compare the different parametrizations for the training of a quantized ResNet-20 on CIFAR-10 using ADAM. Note that all cases use the same learning rate. For the interested reader, a more detailed visualization of the error surfaces over the quantization parameters can be found in Appendix A.3.

2) CIFAR-10:

In our second experiment we train a ResNet-20 (He et al., 2016) with quantized parameters and activations on CIFAR-10 (Krizhevsky & Hinton, 2009) using the same settings as proposed by (He et al., 2016) . Fig. 4 shows the evolution of the training and validation error during training for the case of uniform quantization. The plots for power-of-two quantization can be found in the appendix (Fig. 10 ). We initialize this network from random parameters or from a pre-trained float network. The quantized DNNs are trained for 160 epochs, using SGD with momentum 0.9 and a learning rate schedule starting with 0.01 and reducing it by a factor of 10 after 80 and 120 epochs, respectively. We use random flips and crops for data augmentation. Each epoch takes about 2.5 min on a single GTX 1080 Ti.

In case of randomly initialized weights, we use an initial stepsize d l = 2 -3 for the quantization of weights and activations. Otherwise, we initialize the weights using a pre-trained floating point network and the initial stepsize for a layer is chosen to be d l = 2 log 2 (max |W l |/(2 b-1 -1)) . The remaining quantization parameters are chosen such that we start from an initial bitwidth of b = 4 bit. This is a reasonable upper limit for b, as in practice no performance degradation can be observed for b > 4bit. Even simple offline algorithms like min/max quantization result in networks with good accuracies. We define no memory constraints during training, i.e., the network can learn to use a large number of bits to quantize weights and activations of each layer. From Fig. 4 , we again observe that the parametrization θ = [d, q max ] T is best suited to train a uniformly quantized DNN as it converges (validation error with "random"/"float net" initialization) to the best local optimum. Furthermore, we observe the smallest oscillation of the validation error for this parametrization.

Quantization Float32 Uniform quantization Power-of-two quantization θ = [b, d] T θ = [b, qmax] T θ = [d, qmax] T θ = [b, qmax] T θ = [b, qmin] T θ = [

Table 1 compares the best validation error for all parametrizations of the uniform and power-of-two quantizations. We trained networks either with quantized weights and full precision activations or with both being quantized. In case of activation quantization with power-of-two, we use one bit to explicitly represent the value x = 0. This is advantageous as the ReLU nonlinearity will map many activations to this value. We can observe that training the quantized DNN with the optimal parametrization of DQ, i.e., using either θ = [d, q max ] T or θ = [q min , q max ] T results in a network with the lowest validation error. This result again supports our theoretical considerations from Sec. 2.1.

3. TRAINING QUANTIZED DNNS WITH MEMORY CONSTRAINTS:

We now discuss how to train quantized DNNs with memory constraints. Such constraints appear in many applications when the network inference is performed on an embedded device with limited computational power and memory resources.

A quantized DNN consists of layers which compute

EQUATION

where f l (•) denotes the nonlinear activation function of layer l and Q(•; θ) is a per-tensor quantization with parameters θ applied separately to the input and output tensors X l-1 ∈ I l and X l ∈ I l , and also to both the weight tensors W l ∈ P l and the bias vector c l ∈ R M l .1 For a fully connected layer,

I l-1 = R M l-1 , I l = R M l are vectors, P l = R M l ×M l-1 are matrices and A * B is a matrix- vector product. In case of a convolutional layer, I l-1 = R M l-1 ×N l-1 ×N l-1 , I l = R M l ×N l ×N l , P l = R M l ×M l-1 ×K l ×K l are tensors and A * B is a set of M l-1 M l 2D convolutions,

where the convolution is performed on square-sized feature maps of size N l-1 × N l-1 using square-sized kernels of size K l × K l .

DNNs with quantized weights and activations have a smaller memory footprint and are also computationally cheaper to evaluate since Q(α; θ) • Q(β; θ) for α, β ∈ R requires only an integer multiplication for the case of uniform quantization or an integer addition of the exponents for powerof-two quantization. Furthermore, Q(α; θ) + Q(β; θ) for α, β ∈ R only requires an integer addition. Table 2 compares the computational complexity and the memory footprint of layers which apply uniform or power-of-two quantization to weights and activations.

We consider the following memory characteristics of the DNN, constraining them during training: 1. Total memory S w (θ w 1 , ..., θ w L ) = L l=1 S w l (θ w l ) to store all weights: We use the constraint as well as required memory to store the weights S w l and activations S x l of fully connected and convolutional layers.

EQUATION

Layer Quantization C mul l C add l S w l S x l Fully connected uniform M l M l-1 M l M l-1 M l (M l-1 + 1)b w l M l b x l pow-2 0 2M l M l-1 Convolutional uniform M l M l-1 N 2 l K 2 l M l M l-1 N 2 l K 2 l M l (M l-1 K 2 l + 1)b w l M l N 2 l b x l pow-2 0 2M l M l-1 N 2 l K 2 l

to ensure that the total weight memory requirement S w (θ w 1 , ..., θ w L ) is smaller than a certain maximum weight memory size S w 0 . Table 2 gives S w l (θ w l ) for the case of fully connected and convolutional layers. Each layer's memory requirement S w l (θ w l ) depends on the bitwidth b w l : reducing S w l (θ w l ) will reduce the bitwidth b w l . 2. Total activation memory S x (θ x 1 , ...,

θ x L ) = L l=1 S x l (θ x l )

to store all feature maps: We use the constraint

EQUATION

to ensure an upper limit on the total activation memory size S x 0 . Table 2 gives S x l (θ x l ) for the case of fully connected and convolutional layers. Such a constraint is important if we use pipelining for accelerated inference, i.e., if we evaluate multiple layers with several consecutive inputs in parallel. This can, e.g., be the case for FPGA implementations (Guo et al., 2017) .

3. Maximum activation memory Ŝx (θ x 1 , ..., θ x L ) = max l=1,...,L S x l to store the largest feature map: We use the constraint

EQUATION

to ensure that the maximum activation size Ŝx does not exceed a given limit Ŝx 0 . This constraint is relevant for DNN implementations where layers are processed sequentially.

To train the quantized DNN with memory constraints, we need to solve the optimization problem

EQUATION

where J(X L , Y) is the loss function for yielding the DNN output X L although the ground truth is Y. Eq. ( 9) learns the weights W l , c l as well as the quantization parameters θ x l , θ w l . In order to use simple stochastic gradient descent solvers, we use the penalty method (Bertsekas, 2014) to convert (9) into the unconstrained optimization problem

EQUATION

)

where λ j ∈ R + are individual weightings for the penalty terms. Hence, training with weight and activation size constraints requires choosing two penalty weightings λ j , one for (8a) and one for either (8b) or (8c).

Note, that the optimization problem (10) does not necessarily give a quantized DNN which fulfills the memory constraints. The probability to fulfill the constraint g j depends on the choice of λ j . In particular, this probability increases with larger λ j . However, choosing a too large λ j will yield a penalty term that dominates over the network loss decreasing the network performance. In our experiments, we choose λ j such that the initial loss and the penalty term have approximately the same magnitude. Using this simple heuristic, we optained quantized DNNs that reached a high accuracy and at the same time fulfilled the constraints at the end of training.

4. EXPERIMENTS:

In the following, we will use the best parametrizations for uniform and power-of-two DQ, i.e., θ U = [d, q max ] T and θ P = [q min , q max ] T , that we found in Sec. 2. Both parametrizations do not directly depend on the bitwidth b. Therefore, we compute it by using b(θ U ) = log 2 qmax d + 1 + 1 and b(θ P ) = log 2 log 2 qmax qmin + 1 + 1 . All quantized networks use a pre-trained float32 network for initialization and all quantizers are initialized as described in Sec. 2.3. For our experiments on CIFAR-10, we use the same training setup as described in Sec. 2.3. For the experiments on ImageNet, we train the quantized DNNs for 50 epochs, using SGD with momentum 0.9 and a learning rate schedule starting with 0.01 and reducing it by a factor of 10 after 16 and 32 epochs, respectively. Please note that we quantize all layers opposed to other papers which use a higher precision for the first and/or last layer.

In our experiments, we noticed that the performance of DQ is not sensitive to the choice of λ j in (10). For the CIFAR-10 experiments, we use λ = 0.1 for both constraints (for sizes in kB). For the ImageNet experiments, we kept the same regularization level by scaling λ j with the square of the size ratio between the ImageNet model and the CIFAR-10 model. We scale with the square-ratio as the constraints in (10) are squared penalty terms.

First, in Table 3 /top, we train a ResNet-20 on CIFAR-10 with quantized weights and float32 activations. We start with the most restrictive quantization scheme with fixed q max and b = 2bit ("Fixed").

Then, we allow the model to learn q max while b = 2bit remains fixed as was done in (Jain et al., 2019) ("TQT"). Finally, we learn both q max and b with the constraint that the weight size is at most 70KB ("Ours"), which is just 4.5kB larger that the previous 2Bit networks. This allows the model to allocate more than two bits to some layers. From Table 3 /top, we observe that the error is smallest when we learn all quantization parameters.

In Table 3 /bottom, weights and activations are quantized. For activation quantization, we consider two cases as discussed in Sec. 3. The first one constrains the total activation memory S x while the second constrains the maximum activation memory Ŝx such that both have the same size as a homogeneously quantized model with 4bit activations. Again, we observe that the error is smallest when we learn all quantization parameters.

We also use DQ to train quantized ResNet-18 (He et al., 2016) and MobileNetV2 (Sandler et al., 2018) on ImageNet (Deng et al., 2009) with 4bit uniform weights and activations or equivalent-sized networks with learned quantization parameters. This is quite aggressive and, thus, a fixed quantization scheme loses more than 6% accuracy while our quantization scheme loses less than 0.5% compared to a float32 precision network.

Our results compare favorably to other recent quantization approaches. To our knowledge, the best result for a 4bit ResNet-18 was reported by (Esser et al., 2019) (29.91% error). This is very close to our performance (29.92% error). Importantly, (Esser et al., 2019) did not quantize the first and last layers, meaning that their network is much bigger. Specifically, compared to our quantized ResNet-18, their model with high precision input and output layers requires 37% more memory to store the weights. Moreover, (Esser et al., 2019) learns stepsizes which are not restricted to powers-of-two.

As explained in Sec. 2.2, uniform quantization with power-of-two stepsize leads to more efficient inference, effectively allowing to efficiently compute any multiplication with an integer multiplication and bit-shift. To our knowledge only (Wang et al., 2018) reported results of MobileNetV2 quantized to 4bit. They keep the baseline performance constraining the network to the same size as the 4bit network. However, they do not quantize the activations in this case. In addition, DQ training is efficient since it is comparable to the training of unquantized network. Specifically, one epoch on ImageNet takes 37min for MobileNetV2 and 18min for ResNet-18 on four Nvidia Tesla V100.

Fig. 5 shows the weight bitwidth assignment over layers. We observe that small bitwidths are used for layers with many parameters, i.e., pointwise convolutions and fully connected layers. However, the resulting bitwidth assignments are complex, meaning that there is no simple heuristic. Therefore, it is important to learn the optimal bitwidth assignment. Figure 5 : Weight bitwidth assignment over layers for ResNet-18 and MobileNetV2 on ImageNet with weights constrained to a maximum size of 5.57MB. Our method has learned a heterogeneous bitwidth distribution, which gives a better performance than a homogeneous one (see Table 4 ).

5. CONCLUSIONS:

In this paper we discussed differentiable quantization and its application to the training of compact DNNs with memory constraints. In order to fulfill memory constraints, we introduced penalty functions during training and used stochastic gradient descent to find the optimal weights as well as the optimal quantization values in a joint fashion. We showed that there are several possible parametrizations of the quantization function. In particular, learning the bitwidth directly is not optimal; therefore, we proposed to parametrize the quantizer with the stepsize and dynamic range instead. The bitwidth can then be inferred from them. This approach is competitive to other recent quantization methods while it does not require to retrain the network multiple times in contrast to reinforcement learning approaches (Wang et al., 2018; Elthakeb et al., 2018) .

A DERIVATION OF THE GRADIENTS FOR DIFFERENTIABLE QUANTIZATION (DQ):

In the following sections, we will give the derivatives ∂ ∂x Q(x; θ) and gradients ∇ θ Q(x; θ) for the uniform and the power-of-two quantizers. The results are summarized in Sec. 2.

We use the straight-through gradient estimate whenever we need to differentiate a non-differentiable floor function, i.e., we assume

EQUATION

A.1 DERIVATIVES OF THE UNIFORM QUANTIZER Fig. 6 (a) shows a symmetric uniform quantizer Q U (x; θ) which maps a real value x ∈ R to one of I = 2k + 1 quantized values q ∈ {-kd, ..., 0, ..., kd} by computing

q = Q U (x; θ) = sign(x) d |x| d + 1 2 |x| ≤ q max q max |x| > q max (12)

using the parameters θ = [d, q max , b] T where d ∈ R is the stepsize, q max ∈ R is the maximum value and b ∈ N is the number of bits that we use to encode the quantized values q. The elements of θ are dependent as there is the relationship q max = (2 b-1 -1)d.

A.1.1 CASE U1: PARAMETRIZATION WITH RESPECT TO b AND d:

For the parametrization with respect to the bitwidth b and steps size d, (12) is given by

q = Q U (x; b, d) = sign(x)d |x| d + 1 2 |x| ≤ (2 b-1 -1)d 2 b-1 -1 |x| > (2 b-1 -1)d (13)

and the derivatives are given by

EQUATION

)

EQUATION

A.1.2 CASE U2: PARAMETRIZATION WITH RESPECT TO b AND q MAX For the parametrization with respect to the bitwidth b and maximum value q max , ( 12) is given by

EQUATION

)

and the derivatives are given by

EQUATION

)

where

β 1 = qmax 2 b-1 log 2 ∂ |x| 2 b-1 -1 qmax + 1 2 ∂b = |x| and β 2 = q 2 max 2 b-1 -1 ∂ |x| 2 b-1 -1 qmax + 1 2 ∂qmax

= -|x|, if we use the straight-through gradient estimate for the floor function. Eq. ( 12) gives the quantization with respect to the step size d and maximum value q max . The derivatives are

d q max d q max x Q U (x) (a) Uniform quantization q min q max q min q max x Q P (x) (b) Power-of-two quantization

EQUATION

)

EQUATION

where q min and q max are the minimum and maximum (absolute) values of the quantizer for a bitwidth of b bits. Fig. 6b shows the quantization curve for this quantization scheme.

Using the STE for the floor operation, the derivative ∂ x Q P (x; θ) is given by

EQUATION

The power-of-two quantization has the three parameters θ = [b, q min , q max ], which are dependent on each other, i.e., q max = 2 2 b-1 -1 q min . Therefore, we have again three different parametrizations with θ = [b, q min ], θ = [b, q max ] or θ = [q min , q max ], respectively. The resulting partial derivatives for each parametrization are shown in Fig. 7 and summarized in the following sections. Similar to the uniform case, one parametrization (θ = [q min , q max ]) leads to a gradient with the nice form

EQUATION

which has a bounded gradient magnitude and independent components and is, hence, well suited for first order gradient based optimization.

A.2.1 CASE P1: PARAMETRIZATION WITH RESPECT TO b AND q MAX For the parametrization with θ = [b, q max ], (18) is given by and the partial derivatives are

EQUATION

EQUATION

EQUATION

A.2.2 CASE P2: PARAMETRIZATION WITH RESPECT TO b AND q MIN For the parametrization with θ = [b, q min ], (18) is given by

Q P (x; b, q min ) = sign(x)      q min |x| ≤ q min 2 0.5+log 2 |x| q min < |x| ≤ 2 2 b-1 -1 q min 2 2 b-1 -1 q min |x| > 2 2 b-1 -1 q min (23)

and the partial derivatives are

EQUATION

EQUATION

A.2.3 CASE P3: PARAMETRIZATION WITH RESPECT TO q MIN AND q MAX Eq. ( 18) gives the parametrization of Q(x; θ) with respect to the minimum value q min and maximum value q max . The derivatives are ∂Q P (x; q min , q max ) ∂q min = sign(x)

EQUATION

∂Q P (x; q min , q max ) ∂q max = sign(x) we show how the error surfaces look like for this example problem. The experimental setup is the same as in Sec. 2.3, i.e., we use DQ to learn the optimal quantization parameters of a uniform and a power-of-two quantizer, which minimize the expected quantization error min θ E (x -Q(x; θ)) 2 .

EQUATION

We use three different parametrizations, adapt the quantizer's parameters with gradient descent and compare the convergence speed as well as the final quantization error. As an input, we generate 10 4 samples from N (0, 1).

Fig. 8 shows the corresponding error surfaces for the three different parametrizations of the uniform quantization. The red curve shows the path through the parameter space taken by gradient descent in order to optimize the MSE, starting with the initial values b = 2, d = q max = 1. The optimum θ * is located at b = 16, d 2 -13 , q max = 4, since we allow a maximal bitwidth of 16bit and the largest sample magnitude in our dataset is max{x 1 , ..., x N } 4. In each of the cases U1-U3, the error surface is composed of steep ridges and large flat regions. The steep ridges force us to use small learning rates to avoid divergence. For cases U1 and U2, the optimal θ * can not be reached. However, for U3, θ * lies at the border of a flat region and can be easily reached. Furthermore, case U3 shows a much faster and more stable convergence without oscillation, since the gradient magnitudes are bounded and the error surface has fewer steep ridges where gradient descent starts oscillating.

Fig. 9 shows the corresponding error surfaces for the three different parametrizations of the power-oftwo quantization. Again, the optimum θ * is not attained for two parametrizations, namely P1 and P2, as θ * is surrounded by a large, mostly flat region. For these two cases, gradient descent tends to oscillate at steep ridges and tends to be unstable. However, gradient descent converges to a point close to θ * for parametrization P3, where θ = [q min , q max ].

Finally, we also did a comparison of the different power-of-two quantizations on CIFAR-10. Fig. 10 shows the evolution of the training and validation error if we start from a random or a pre-trained float network initialization. We can observe that θ = [q min , q max ] has the best convergence behavior and thus also results in the smallest validation error (cf. Table 1 ). The unstable behavior of P2 is expected as the derivative ∂Q P ∂qmin can take very large (absolute) values.

A.4 FURTHER EXPERIMENTS WITH ADAM:

Finally, we did an experiment to verify that the parametrization is important, even if adaptive gradient descent methods like ADAM are used for optimization. Table 5 gives the results for a ResNet-20 trained on CIFAR-10. We observe, that again U3 and P3 are the best parametrizations. 

B IMPLEMENTATION DETAILS FOR DIFFERENTIABLE QUANTIZATION:

The following code gives our differentiable quantizer implementation in NNabla (Sony) . The source code for reproducing our results will be published after the review process has been finished. 

:

In this paper, we use "weights" to refer to W and c.

Figures and Tables:

Case U3: θ = [d, qmax] T

Figure1: Maximum gradient norm max x ∇ θ Q U (x; θ) . For "U1" and "U2" the maximum gradient norm can grow exponentially with varying bitwidth b whereas it is bounded for "U3".

Figure 4: ResNet-20 with uniformly quantized weights and activations.

Figure 6: Examples of uniform quantizer Q U (x) and power-of-two quantizer Q P (x) for b = 3 bits A.1.3 CASE U3: PARAMETRIZATION WITH RESPECT TO d AND q MAX

Figure 7: Derivatives for the three different parametrizations of Q P (x; θ)

Figure 8: MSE surfaces for uniform quantization. Only U3 reaches the optimum θ * .

Figure 10: ResNet-20 with power-of-two quantized weights and activations.

Comparison of different DQ parametrizations for ResNet-20 on CIFAR-10.

Table content: <table/>

qmin, qmax] T

Table content: <table><tr><td colspan="2">Weights Weights+Activations</td><td colspan="2">8.50%/7.29%</td><td>17.8%/8.18% 8.80%/7.44% 8.50%/7.32% 11.70%/7.90% 53.07%/23.01% 28.9%/9.03% 9.43%/7.74% 9.23%/7.40% 22.91%/11.68% diverging/35.68% 15.10%/9.86% 10.61%/7.56%</td></tr><tr><td>0</td><td>2</td><td>4</td><td>6 •10 4</td></tr></table>

Number of multiplications C mul

Table content: <table/>

Homogeneous vs. heterogeneous quantization of ResNet-20 on CIFAR-10.

Table content: <table><tr><td/><td>Bitwidth</td><td>q max</td><td>Size</td><td colspan="2">Uniform quant. Power-of-two quant.</td></tr><tr><td/><td colspan="4">Weight/Activ. Weight/Activ. Weight/Activ.(max)/Activ.(sum) Validation error</td><td>Validation error</td></tr><tr><td>Baseline</td><td>32bit/32bit</td><td>-</td><td>1048KB/64KB/736KB</td><td/><td>7.29%</td></tr><tr><td>Fixed</td><td>2bit/32bit</td><td>fixed/-</td><td>65.5KB/64KB/736KB</td><td>10.81%</td><td>8.99%</td></tr><tr><td>TQT (Jain et al., 2019)</td><td>2bit/32bit</td><td>learned/ -</td><td>65.5KB/64KB/736KB</td><td>9.47%</td><td>8.79%</td></tr><tr><td>Ours (w/ constr. (8a))</td><td>learned/32bit</td><td>learned/-</td><td>70KB/64KB/736KB</td><td>8.59%</td><td>8.53%</td></tr><tr><td>Fixed</td><td>2bit/4bit</td><td>fixed/fixed</td><td>65.5KB/8KB/92KB</td><td>11.30%</td><td>11.62%</td></tr><tr><td>TQT (Jain et al., 2019)</td><td>2bit/4bit</td><td>learned/learned</td><td>65.5KB/8KB/92KB</td><td>9.62%</td><td>11.29%</td></tr><tr><td colspan="3">Ours (w/ constr. (8a) and (8b)) learned/learned learned/learned</td><td>70KB/ -/92KB</td><td>9.38%</td><td>11.29%</td></tr><tr><td colspan="3">Ours (w/ constr. (8a) and (8c)) learned/learned learned/learned</td><td>70KB/8KB/ -</td><td>8.58%</td><td>11.23%</td></tr></table>

Homogeneous vs. heterogeneous quantization of MobileNetV2 and ResNet-18 on ImageNet.

Table content: <table><tr><td/><td/><td/><td/><td/><td/><td colspan="2">MobileNetV2</td><td>ResNet-18</td></tr><tr><td/><td/><td>Bitwidth</td><td>q max</td><td/><td/><td>Size</td><td>Validation</td><td>Size</td><td>Validation</td></tr><tr><td/><td/><td colspan="2">Weight/Activ. Weight/Activ.</td><td colspan="4">Weight/Activ(max)</td><td>Error</td><td>Weight/Activ(max)</td><td>Error</td></tr><tr><td colspan="2">Baseline</td><td>32bit/32bit</td><td>-</td><td colspan="4">13.23MB/4.59MB 29.82%</td><td>44.56MB/3.04MB 29.72%</td></tr><tr><td>Fixed</td><td/><td>4bit/4bit</td><td>fixed/fixed</td><td colspan="4">1.65MB/0.57MB</td><td>36.27%</td><td>5.57MB/0.38MB</td><td>34.15%</td></tr><tr><td colspan="2">TQT (Jain et al., 2019)</td><td>4bit/4bit</td><td>learned/learned</td><td colspan="4">1.65MB/0.57MB</td><td>32.21%</td><td>5.57MB/0.38MB</td><td>30.49%</td></tr><tr><td colspan="4">Ours (w/ constr. (8a) and (8c)) learned/learned learned/learned</td><td colspan="4">1.55MB/0.57MB</td><td>30.26%</td><td>5.40MB/0.38MB</td><td>29.92%</td></tr><tr><td colspan="2">Ours (w/o constr.)</td><td colspan="2">learned/learned learned/learned</td><td colspan="4">3.14MB/1.58MB</td><td>29.41%</td><td>10.50MB/1.05MB 29.34%</td></tr><tr><td/><td/><td>ResNet-18</td><td/><td/><td/><td/><td>MobileNetV2</td></tr><tr><td/><td>10 bit</td><td/><td/><td/><td/><td>10 bit</td></tr><tr><td/><td>8 bit</td><td/><td/><td/><td/><td>8 bit</td></tr><tr><td>Bitwidth</td><td>4 bit 6 bit</td><td/><td/><td/><td>Bitwidth</td><td>4 bit 6 bit</td></tr><tr><td/><td>2 bit</td><td/><td/><td/><td/><td>2 bit</td></tr><tr><td/><td/><td>Layer</td><td/><td/><td/><td/><td>Layer</td></tr><tr><td/><td/><td/><td>Convolution</td><td>Affine</td><td colspan="2">Depthwise</td><td>Pointwise</td></tr></table>

Error rate of ResNet-20 on CIFAR-10 using different quantization parametrizations. Training is done either by SGD with momentum or ADAM.

Table content: <table><tr><td>Parametrization</td><td>SGD momentum</td><td>ADAM</td></tr><tr><td>U1</td><td>11.74</td><td>7.61</td></tr><tr><td>U2</td><td>7.44</td><td>7.85</td></tr><tr><td>U3</td><td>7.32</td><td>7.36</td></tr><tr><td>P1</td><td>15.35</td><td>7.54</td></tr><tr><td>P2</td><td>7.74</td><td>7.79</td></tr><tr><td>P3</td><td>7.40</td><td>7.40</td></tr></table>

.2.3 CASE P3: PARAMETRIZATION WITH RESPECT TO q MIN AND q MAX

Table content: <table><tr><td colspan="2">B.1 UNIFORM QUANTIZATION B.1.1 CASE U1: PARAMETRIZATION WITH RESPECT TO b AND d 1 def parametric_fixed_point_quantize_d_b(x, sign, 2 n_init, n_min, n_max, 3 d_init, d_min, d_max, 4 fix_parameters=False): 5 """Parametric version of 'fixed_point_quantize' where the 6 bitwidth 'b' and stepsize 'd' are learnable parameters. 7 8 Returns: 9 ~nnabla.Variable: N-D array. 10 """ 11 def clip_scalar(v, min_value, max_value): 12 return F.minimum_scalar(F.maximum_scalar(v, min_value), max_value) 13 14 def broadcast_scalar(v, shape): 15 return F.broadcast(F.reshape(v, (1,) * len(shape), inplace=False), shape=shape) 16 17 def quantize_pow2(v): 18 return 2 ** F.round(F.log(v) / np.log(2.)) 19 20 n = get_parameter_or_create("n", (), 21 ConstantInitializer(n_init), 22 need_grad=True, 23 as_need_grad=not fix_parameters) 24 d = get_parameter_or_create("d", (), 25 ConstantInitializer(d_init), 26 need_grad=True, 27 as_need_grad=not fix_parameters) 28 29 # ensure that bitwidth is in specified range and an integer 30 n = F.round(clip_scalar(n, n_min, n_max)) 31 if sign: 32 n = n -1 33 34 # ensure that stepsize is in specified range and a power of two 35 d = quantize_pow2(clip_scalar(d, d_min, d_max)) 36 37 # ensure that dynamic range is in specified range 38 xmax = d * (2 ** n -1) 39 40 # compute min/max value that we can represent 41 if sign: 42 xmin = -xmax B.1.2 CASE U2: PARAMETRIZATION WITH RESPECT TO b AND q MAX 1 def parametric_fixed_point_quantize_b_xmax(x, sign, 2 n_init, n_min, n_max, 3 xmax_init, xmax_min, xmax_max, 4 fix_parameters=False): 5 """Parametric version of 'fixed_point_quantize' where the 6 bitwidth 'b' and dynamic range 'xmax' are learnable parameters. 7 8 Returns: 9 ~nnabla.Variable: N-D array. 10 """ 11 def clip_scalar(v, min_value, max_value): 12 return F.minimum_scalar(F.maximum_scalar(v, min_value), max_value) 13 14 def broadcast_scalar(v, shape): 15 return F.broadcast(F.reshape(v, (1,) * len(shape), inplace=False), shape=shape) 16 17 def quantize_pow2(v): 18 return 2 ** F.round(F.log(v) / np.log(2.)) 19 20 n = get_parameter_or_create("n", (), 21 ConstantInitializer(n_init), 22 need_grad=True, 23 as_need_grad=not fix_parameters) 24 xmax = get_parameter_or_create("xmax", (), 25 ConstantInitializer(xmax_init), 26 need_grad=True, 27 as_need_grad=not fix_parameters) 28 29 # ensure that bitwidth is in specified range and an integer 30 n = F.round(clip_scalar(n, n_min, n_max)) 31 if sign: 32 n = n -1 33 34 # ensure that dynamic range is in specified range 35 xmax = clip_scalar(xmax, xmax_min, xmax_max) 36 37 # compute step size from dynamic range and make sure that it is a pow2 38 d = quantize_pow2(xmax / (2 ** n -1)) 39 40 # compute min/max value that we can represent 41 if sign: 42 xmin = -xmax 43 else: 44 xmin = nn.Variable((1,), need_grad=False) 45 xmin.d = 0. B.1.3 CASE U3: PARAMETRIZATION WITH RESPECT TO d AND q MAX 1 def parametric_fixed_point_quantize_d_xmax(x, sign, 2 d_init, d_min, d_max, 3 xmax_init, xmax_min, xmax_max, 4 fix_parameters=False): 5 """Parametric version of 'fixed_point_quantize' where the 6 stepsize 'd' and dynamic range 'xmax' are learnable parameters. 7 8 Returns: 9 ~nnabla.Variable: N-D array. 10 """ 11 def clip_scalar(v, min_value, max_value): 12 return F.minimum_scalar(F.maximum_scalar(v, min_value), max_value) 13 14 def broadcast_scalar(v, shape): 15 return F.broadcast(F.reshape(v, (1,) * len(shape), inplace=False), shape=shape) 16 17 def quantize_pow2(v): 18 return 2 ** F.round(F.log(v) / np.log(2.)) 19 20 d = get_parameter_or_create("d", (), 21 ConstantInitializer(d_init), 22 need_grad=True, 23 as_need_grad=not fix_parameters) 24 xmax = get_parameter_or_create("xmax", (), 25 ConstantInitializer(xmax_init), 26 need_grad=True, 27 as_need_grad=not fix_parameters) 28 29 # ensure that stepsize is in specified range and a power of two 30 d = quantize_pow2(clip_scalar(d, d_min, d_max)) 31 32 # ensure that dynamic range is in specified range 33 xmax = clip_scalar(xmax, xmax_min, xmax_max) 34 35 # compute min/max value that we can represent 36 if sign: 37 xmin = -xmax 38 else: 39 xmin = nn.Variable((1,), need_grad=False) 40 xmin.d = 0. 41 42 # broadcast variables to correct size 43 d = broadcast_scalar(d, shape=x.shape) 44 xmin = broadcast_scalar(xmin, shape=x.shape) 45 xmax = broadcast_scalar(xmax, shape=x.shape) B.2 POWER-OF-TWO QUANTIZATION B.2.1 CASE P1: PARAMETRIZATION WITH RESPECT TO b AND q MAX 1 def parametric_pow2_quantize_b_xmax(x, sign, with_zero, 2 n_init, n_min, n_max, 3 xmax_init, xmax_min, xmax_max, 4 fix_parameters=False): 5 """Parametric version of 'pow2_quantize' where the 6 bitwidth 'n' and range 'xmax' are learnable parameters. 7 8 Returns: 9 ~nnabla.Variable: N-D array. 10 """ 11 def clip_scalar(v, min_value, max_value): 12 return F.minimum_scalar(F.maximum_scalar(v, min_value), max_value) 13 14 def broadcast_scalar(v, shape): 15 return F.broadcast(F.reshape(v, (1,) * len(shape), inplace=False), shape=shape) 16 17 def quantize_pow2(v): 18 return 2 ** F.round(F.log(F.abs(v)) / np.log(2.)) 19 20 n = get_parameter_or_create("n", (), 21 ConstantInitializer(n_init), 22 need_grad=True, 23 as_need_grad=not fix_parameters) 24 xmax = get_parameter_or_create("xmax", (), 25 ConstantInitializer(xmax_init), 26 need_grad=True, 27 as_need_grad=not fix_parameters) 28 29 # ensure that bitwidth is in specified range and an integer 30 n = F.round(clip_scalar(n, n_min, n_max)) 31 if sign: 32 n = n -1 33 if with_zero: 34 n = n -1 35 36 # ensure that dynamic range is in specified range and an integer 37 xmax = quantize_pow2(clip_scalar(xmax, xmax_min, xmax_max)) 38 39 # compute min value that we can represent 40 xmin = (2 ** (-(2 ** n) + 1)) * xmax 41 42 # broadcast variables to correct size B.2.2 CASE P2: PARAMETRIZATION WITH RESPECT TO b AND q MIN 1 def parametric_pow2_quantize_b_xmin(x, sign, with_zero, 2 n_init, n_min, n_max, 3 xmin_init, xmin_min, xmin_max, 4 fix_parameters=False): 5 """Parametric version of 'pow2_quantize' where the 6 bitwidth 'n' and the smallest value 'xmin' are learnable parameters. 7 8 Returns: 9 ~nnabla.Variable: N-D array. 10 """ 11 def clip_scalar(v, min_value, max_value): 12 return F.minimum_scalar(F.maximum_scalar(v, min_value), max_value) 13 14 def broadcast_scalar(v, shape): 15 return F.broadcast(F.reshape(v, (1,) * len(shape), inplace=False), shape=shape) 16 17 def quantize_pow2(v): 18 return 2 ** F.round(F.log(F.abs(v)) / np.log(2.)) 19 20 n = get_parameter_or_create("n", (), 21 ConstantInitializer(n_init), 22 need_grad=True, 23 as_need_grad=not fix_parameters) 24 xmin = get_parameter_or_create("xmin", (), 25 ConstantInitializer(xmin_init), 26 need_grad=True, 27 as_need_grad=not fix_parameters) 28 29 # ensure that bitwidth is in specified range and an integer 30 n = F.round(clip_scalar(n, n_min, n_max)) 31 if sign: 32 n = n -1 33 if with_zero: 34 n = n -1 35 36 # ensure that minimum dynamic range is in specified range and a power-of-two 37 xmin = quantize_pow2(clip_scalar(xmin, xmin_min, xmin_max)) 38 39 # compute min/max value that we can represent 40 xmax = xmin * (2 ** ((2 ** n) -1)) 41 42 2 xmin_init, xmin_min, xmin_max, 3 xmax_init, xmax_min, xmax_max, 4 fix_parameters=False): 5 """Parametric version of 'pow2_quantize' where the 6 min value 'xmin' and max value 'xmax' are learnable parameters. 7 8 Returns: 9 ~nnabla.Variable: N-D array. 10 """ 11 def clip_scalar(v, min_value, max_value): 12 return F.minimum_scalar(F.maximum_scalar(v, min_value), max_value) 13 14 def broadcast_scalar(v, shape): 15 return F.broadcast(F.reshape(v, (1,) * len(shape), inplace=False), shape=shape) 16 17 def quantize_pow2(v): 18 return 2. ** F.round(F.log(F.abs(v)) / np.log(2.)) 19 20 xmin = get_parameter_or_create("xmin", (), 21 ConstantInitializer(xmin_init), 22 need_grad=True, 23 as_need_grad=not fix_parameters) 24 xmax = get_parameter_or_create("xmax", (), 25 ConstantInitializer(xmax_init), 26 need_grad=True, 27 as_need_grad=not fix_parameters) 28 29 # ensure that minimum dynamic range is in specified range and a power-of-two 30 xmin = quantize_pow2(clip_scalar(xmin, xmin_min, xmin_max)) 31 32 # ensure that minimum dynamic range is in specified range and a power-of-two 33 xmax = quantize_pow2(clip_scalar(xmax, xmax_min, xmax_max)) 34 35 # broadcast variables to correct size 36 xmin = broadcast_scalar(xmin, shape=x.shape) 37 xmax = broadcast_scalar(xmax, shape=x.shape) 38 39 # if unsigned, then quantize all negative values to zero # broadcast variables to correct size 1 def parametric_pow2_quantize_xmin_xmax(x, sign, with_zero, 40 if not sign:</td></tr><tr><td>43 46 46</td><td>else:</td></tr><tr><td>44 47 47</td><td>xmin = nn.Variable((1,), need_grad=False) # broadcast variables to correct size # apply fixed-point quantization</td></tr><tr><td>45 48 48</td><td>xmin.d = 0. d = broadcast_scalar(d, shape=x.shape) return d * F.round(F.clip_by_value(x, xmin, xmax) / d)</td></tr><tr><td>46 49</td><td>xmin = broadcast_scalar(xmin, shape=x.shape)</td></tr><tr><td>47 50</td><td># broadcast variables to correct size xmax = broadcast_scalar(xmax, shape=x.shape)</td></tr><tr><td>48 51</td><td>d = broadcast_scalar(d, shape=x.shape)</td></tr><tr><td>49 52</td><td>xmin = broadcast_scalar(xmin, shape=x.shape) # apply fixed-point quantization</td></tr><tr><td>50 53</td><td>xmax = broadcast_scalar(xmax, shape=x.shape) return d * F.round(F.clip_by_value(x, xmin, xmax) / d)</td></tr><tr><td>51</td><td/></tr><tr><td>52</td><td># apply fixed-point quantization</td></tr><tr><td>53</td><td>return d * F.round(F.clip_by_value(x, xmin, xmax) / d)</td></tr></table>

Peer review (iclr2020):
The work studies differentiable quantization of deep neural networks with straight-through gradient (Bengio et. al., 2013). The authors find that a proper parametrization of the quantizer is critical to stable training and good quantization performance and demonstrated their findings to obtain mixed precision DNNs on two datasets, i.e., CIFAR-10 and Imagenet. The paper is clearly written and easy to follow. The idea proposed is fairly straight-forward. Although the argument the authors used to support the finding is not very rigorous, the finding itself is still worth noting. One of the arguments that the authors used to support the specific form of parametrization is that it leads to diagonal Hessian. From optimization perspective, what matters is the condition number, i.e., max/min of the eigenvalues of the Hessian. Could this explain the small difference between the three different parametrization forms with uniform quantization and the big difference for power-of-two quantization? The penalty method used to address the memory constraints will not necessarily lead to solutions that satisfy the constraints. The authors noted that the algorithm is not sensitive to the choice of the penalty parameters. Have the authors tried to tackle problems of hard memory constraints?