M IXED P RECISION DNN S : A LL YOU NEED IS A GOOD PARAMETRIZATION
A BSTRACT
Efﬁcient 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 ﬁnal performance.
Speciﬁcally, 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 conﬁrm our ﬁndings with experiments on CIFAR-10 and ImageNet and we obtain mixed precision
DNNs with learned quantization parameters, achieving state-of-the-art performance.
I NTRODUCTION
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 efﬁciently.
This is particularly relevant for inference on mobile or embedded devices with limited computational power.
However, gradient based training of quantized DNNs is difﬁcult, 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 ﬁrst 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 difﬁcult.
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 deﬁne 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.
Speciﬁcally, 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 ﬂoat32 training.
The contributions of this paper are:
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.
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.
We conﬁrm our ﬁndings with experiments on CIFAR-10 and ImageNet.
For example, we train a heterogeneously quantized MobileNetV2 on ImageNet requiring a total of only 1 .
65 MB to store the weights and only 0 .
57 MB to store its largest feature map.
This is equivalent to a homogenous 4 bit 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 ﬂoating 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; (cid:98) .
(cid:99) and (cid:100) .
(cid:101) are the ﬂoor and ceiling operators.
Finally, δ ( . ) denotes the Dirac delta function.
C HOOSING 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. P ARAMETRIZATION AND STRAIGHT THROUGH GRADIENT ESTIMATES
A symmetric uniform quantizer Q U ( x ; θ ) which maps a real value x ∈ R to one of I = 2 k +quantized values q ∈ {− kd, ..., 0 , ..., kd } computes
[Equation]
[Figure]
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 ; θ ) =
[Equation] vanishes almost everywhere.
A solution is to deﬁne the derivative using STE (Bengio et al., 2013), which ignores the ﬂoor 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 deﬁne the gradients ∇ x Q ( x ; θ ) and ∇ θ Q ( x ; θ ) , using STE whenever we need to differentiate a ﬂoor 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 U
Parametrization with respect to θ =
[ b, d
T
, using q max
=
q
max ( b, d ) gives
[Equation]
Case U2: Parametrization with respect to θ =
[ b, q max ] T using d = d ( b, q max ) gives
[Equation]
Case U3: Parametrization with respect to θ =
[ d, q max ] T , using b = b ( d, q max ) gives
[Equation]
Fig. 1 shows the maximum gradient norm max x (cid:107) ∇ θ Q U ( x ; θ ) (cid:107) for the three parametrizations “U1” to “U3”. For the parametrizations “U1” and “U2”, max x (cid:107) ∇ θ Q U ( x ; θ ) (cid:107) 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 ; θ ) ∈ (cid:104) −
[Equation]
log 2 , q max 2 b
[Equation]
1 − 1 log 2 (cid:105) 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 −x ∂ b Q U ( x ) ∂ d Q U ( x ) (a) Case U q max x ∂ b Q U ( x ) ∂ q max Q U ( x ) (b) Case U q max x ∂ d Q U ( x ) ∂ q max Q U ( x ) (c) Case U q max x ∂ x Q U ( x ) (d) Input derivative avoid divergence.
However, parametrization “U3” does not suffer such an unbounded gradient norm as both partial derivatives ∂ d Q U ( x ; θ ) ∈ [ − 12 , 12 ] and ∂ are bounded.
Fig.
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.
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 ﬂoor 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 ﬁrst order gradient based optimization.
2. C ONSTRAINTS ON θ
In practice, for an efﬁcient hardware implementation, we need to ensure that the quantization parameters only take speciﬁc 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 fulﬁll 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 ﬂoat values, i.e., we used again the
STE to propagate the gradients.
2. E XPERIMENTAL COMPARISON OF DQ PARAMETRIZATIONS
In the following we compare the parametrizations using two experiments.
[Figure]
1) Quantization of Gaussian data
In our ﬁrst experiment we use DQ to learn the optimal quantization parameters θ ∗ which minimize the mean squared error (MSE) E (cid:2) 12 ( Q ( x ; θ ) − x ) 2 (cid:3) 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 (cid:2) ( Q ( x ; θ ) − x ) 2 (cid:3) ∈
R 2 × 2 of the MSE:
H = E (cid:2) ∇ θ Q ( x ; θ ) ∇ θ Q ( x ; θ ) T
+ ( Q ( x ; θ ) − x ) ∇ θ ∇ Tθ Q ( x ; θ
[Equation]
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 (cid:2) ∇ θ Q ( x ; θ ) ∇ θ Q ( x ; θ ) T (cid:3) 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, efﬁcient.
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 artiﬁcially 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) CIFARIn 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 ﬂoat 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 ﬂips 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 ﬂoating point network and the initial stepsize for a layer is chosen to be d l = 2 (cid:98) log 2 (max | W l | / (2 b − 1 − 1))
(cid:99) .
The remaining quantization parameters are chosen such that we start from an initial bitwidth of b
=bit.
This is a reasonable upper limit for b , as in practice no performance degradation can be observed for b > 4 bit.
Even simple ofﬂine algorithms like min/max quantization result in networks with good accuracies.
We deﬁne 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 [Figure] [Figure] [Figure]
to the best local optimum.
Furthermore, we observe the smallest oscillation of the validation error for this parametrization.
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.
T RAINING QUANTIZED DNN S 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 .
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 matrixvector 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:
Total memory S w ( θ w 1 , ..., θ wL ) =
[Equation]
Ll =1 S wl ( θ wl ) to store all weights: We use the constraint
[Equation]
as well as required memory to store the [Table] 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 wl .
Total activation
[Equation]
Ll
[Equation]
l [Equation]
θ x l ) to store all feature maps: We use the constraint
[Equation]
8b
[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 ˆ S x ( θ
x [Equation] [Equation]
, θ xL ) = max l =1 ,... L S xl to store the largest feature map: We use the constraint
[Equation]
to ensure that the maximum activation size ˆ S x does not exceed a given limit ˆ S 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 θ xl , θ wl .
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 fulﬁlls the memory constraints.
The probability to fulﬁll 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 fulﬁlled the constraints at the end of training.
E XPERIMENTS
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 ) =
(cid:6) log 2 (cid:0) q max d + 1 (cid:1) + 1 (cid:7) and b ( θ P )
= (cid:108) log 2 (cid:16) log 2 (cid:16) q max q min (cid:17) + 1 (cid:17)
+ 1 (cid:109) .
All quantized networks use a pre-trained ﬂoat 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
[Table]
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 ﬁrst and/or last layer.
In our experiments [Table] 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 ﬂoat32 activations.
We start with the most restrictive quantization scheme with ﬁxed q max and b = 2 bit (“Fixed”).
Then, we allow the model to learn q max while b = 2 bit remains ﬁxed 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 70 KB (“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 ﬁrst one constrains the total activation memory S x while the second constrains the maximum activation memory ˆ S x such that both have the same size as a homogeneously quantized model with 4 bit 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 ﬁxed quantization scheme loses more than 6 % accuracy while our quantization scheme loses less than 0 .
5 % compared to a ﬂoat32 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 ﬁrst and last layers, meaning that their network is much bigger.
Speciﬁcally, 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 efﬁcient inference, effectively allowing to efﬁciently 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 efﬁcient since it is comparable to the training of unquantized network.
Speciﬁcally, 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.
[Table]
[Figure]
[Table]
[Figure]
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).
C ONCLUSIONS
In this paper we discussed differentiable quantization and its application to the training of compact
DNNs with memory constraints.
In order to fulﬁll memory constraints, we introduced penalty functions during training and used stochastic gradient descent to ﬁnd 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 CKNOWLEDGEMENTS
We would like to thank Masato Ishii for many helpful comments during the preparation of this manuscript.
A D ERIVATION 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 ﬂoor function, i.e., we assume
[Equation]
A. D ERIVATIVES 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 = 2 k + 1 quantized values q ∈ {− kd, ..., 0 , ..., kd } by computing
[Equation]
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. C ASE U1: P ARAMETRIZATION WITH RESPECT TO b AND d
For the parametrization with respect to the bitwidth b and steps size d , (12) is given by
[Equation]
and the derivatives are given by
[Equation]
A.1. C ASE U2: P ARAMETRIZATION 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 =
[Equation]
[Equation]
∂b = | x | and β 2 =
[Equation]
|
[Equation]
|
[Equation]
max =
−| x | , if we use the straight-through gradient estimate for the ﬂoor function.
[Figure]
[Equation]
[Figure]
max
[Figure]
q max q min q max x
Q P
[Equation]
x )
[Figure]
Eq.
(12) gives the quantization with respect to the step size d and maximum value
q max .
The derivatives are
[Equation]
17b
[Equation]
A. D ERIVATIVES OF THE POWER - OF - TWO QUANTIZER
Power-of-two quantization Q P ( x ; θ ) maps a real-valued number x ∈ R to a quantized value q ∈ {± 2 k : k ∈ Z } by [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 ﬂoor 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
[Equation]
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 ﬁrst order gradient based optimization
[Equation]
A.2. C ASE P1: P ARAMETRIZATION WITH RESPECT TO b AND q MAX
For the parametrization with θ =
[ b, q max ] , (18) is given by
[Equation]
[Figure]
and the partial derivatives are
[Equation]
A.2. C ASE P2: P ARAMETRIZATION WITH RESPECT TO b AND q MIN
For the parametrization with θ =
[ b, q min ] , (18) is given by
[Equation] [Equation]
and the partial derivatives are
[Equation]
(24b)
A.2. C ASE P3: P ARAMETRIZATION 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
[Equation] [Equation]
(25a)
[Equation]

[Equation]
x | ≤
[Equation]
min 0 q min
[Equation]
x | ≤
[Equation]
max | x | > q max .
(25b)
[Figure] [Figure]
A. V ISUALIZATION OF THE ERROR SURFACE FOR THE QUANTIZATION OF G AUSSIAN DATA
In Sec.
2.3 of the paper, we compared the three different parametrizations of the uniform quantizer at the example of optimal quantization of Gaussian data.
To get a better understanding of Fig 3, 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 (cid:104) ( x − Q ( x ; θ )) 2 (cid:105) .
We use three different parametrizations, adapt the quantizer’s parameters with gradient descent and compare the convergence speed as well as the ﬁnal quantization error.
As an input, we generate 10samples 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 (cid:47) 2 − 13 , q max = 4 , since we allow a maximal bitwidth of 16 bit and the largest sample magnitude in our dataset is max { x 1 , ..., x N } (cid:47) 4 .
In each of the cases U1-U3, the error surface is composed of steep ridges and large ﬂat 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 ﬂat 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 ﬂat 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.shows the evolution of the training and validation error if we start from a random or a pre-trained ﬂoat 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
∂q min can take very large (absolute) values.
A. F URTHER 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-
6 · 1010 −10 −
Iteration
T r a i n i ng e rr o r (a)Randominitialization   6 · 1010 −
Iteration
V a li d a ti on e rr o r b , d (U1) b , q max (U2) d , q max (U3)   6 · 1010 −10 −
Iteration
T r a i n i ng e rr o r (b)Pre-trainedinitialization   6 · 1010 −
Iteration
V a li d a ti on e rr o r
[Table]
B I MPLEMENTATION 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 ﬁnished.
B. U NIFORM QUANTIZATION B.1. C ASE U1: P ARAMETRIZATION WITH RESPECT TO b AND d
[Figure]
[Figure] [Figure] [Figure]
[Figure]
_init, n
[Figure]
_init
[Figure]
[Figure]
[Figure]
fix_parameters=
False):"""Parametric version of ‘fixed_point_quantize‘ where thebitwidth ‘b‘ and stepsize ‘d‘ are learnable parameters.
Returns:~nnabla.
Variable: N-D array. """def clip_scalar(v, min_value, max_value):return F.minimum_scalar(F.maximum_scalar(v, min_value), max_value) def broadcast_scalar(v, shape):return F.broadcast(F.reshape(v, (1,)
* len(shape), inplace=False), shape=shape) def quantize_pow2(v):return 2 ** F.round(F.log(v) / np.log(2.))
n = get_parameter_or_create("n", (),
ConstantInitializer(n_init),need_grad=True,as_need_grad=not fix_parameters)d = get_parameter_or_create("d", (),
ConstantInitializer(d_init),need_grad=
True,as_need_grad=not fix_parameters)
# ensure that bitwidth is in specified range and an integern =
F.round(clip_scalar(n, n_min, n_max))if sign:n = n - # ensure that stepsize is in specified range and a power of twod = quantize_pow2(clip_scalar(d, d_min, d_max))
# ensure that dynamic range is in specified rangexmax =
d * (2 ** n - 1) # compute min/max value that we can representif sign:xmin
= -xmaxelse
xmin = nn.Variable((1,),
need_grad=False)xmin.d = 0.
# broadcast variables to correct sized = broadcast_scalar(d, shape=
x.shape)xmin = broadcast
[Figure]
scalar(xmin, shape=
x.shape)xmax = broadcast_scalar(xmax, shape=x.shape) # apply fixed-point quantizationreturn d * F.round(F.clip_by_value(x, xmin, xmax) / d)
B.1.
ASE U2: P ARAMETRIZATION WITH RESPECT TO b AND q MAX
1 def parametric_fixed_point_quantize_b_xmax(x, sign,n_init, n_min, n_max,xmax_init, xmax_min, xmax_max,fix_parameters=False):"""Parametric version of ‘fixed_point_quantize‘ where thebitwidth ‘b‘ and dynamic range ‘xmax‘ are learnable parameters.
Returns:~nnabla.
Variable: N-D array.
"""def clip_scalar(v, min_value, max_value):return F.minimum_scalar(F.maximum_scalar(v, min_value), max_value) def broadcast_scalar(v, shape):return F.broadcast(F.reshape(v, (1,)
* len(shape), inplace=False), shape=shape) def quantize_pow2(v):return 2 ** F.round(F.log(v) / np.log(2.))
n = get_parameter_or_create("n", (),
ConstantInitializer(n_init),need_grad=True,as_need_grad=not fix_parameters)xmax = get_parameter_or_create("xmax", (),
ConstantInitializer(xmax_init),need_grad=
True,as_need_grad=not fix_parameters)
# ensure that bitwidth is in specified range and an integern =
F.round(clip_scalar(n, n_min, n_max))if sign:n = n - # ensure that dynamic range is in specified range
[Figure]
=
[Figure]
[Figure]
(xmax, xmax_min, xmax_max) # compute step size from dynamic range and make sure that it is a pow2d = quantize_pow2(xmax / (2 ** n - 1))
# compute min/max value that we can representif sign:
xmin =
-xmaxelse:
xmin = nn.Variable((1,),
need_grad=False)xmin.d = 0.
# broadcast variables to correct sized = broadcast_scalar(d, shape=
x.shape)xmin = broadcast_scalar(xmin, shape=x
[Figure]
shape)xmax = broadcast_scalar(xmax, shape=x.shape) # apply fixed-point quantizationreturn d * F.round(F.clip_by_value(x, xmin, xmax) / d)
B.1. C ASE U3: P ARAMETRIZATION WITH RESPECT TO d AND q MAX
1 def parametric_fixed_point_quantize_d_xmax(x, sign,d_init, d_min, d_max,xmax_init, xmax_min, xmax_max,fix_parameters=False):"""Parametric version of ‘fixed_point_quantize‘ where thestepsize ‘d‘ and dynamic range ‘xmax‘ are learnable parameters.
Returns:~nnabla.
Variable: N-D array.
"""def clip_scalar(v, min_value, max_value):return F.minimum_scalar(F.maximum_scalar(v, min_value), max_value) def broadcast_scalar(v, shape):return F.broadcast(F.reshape(v, (1,)
* len(shape), inplace=False), shape=shape) def quantize_pow2(v):return 2 ** F.round(F.log(v) / np.log(2.))
d = get_parameter_or_create("d", (),
ConstantInitializer(d_init),need_grad=
True,as_need_grad=not fix_parameters)xmax = get_parameter_or_create("xmax", (),
ConstantInitializer(xmax_init),need_grad=
True,as_need_grad=not fix_parameters)
# ensure that stepsize is in specified range and a power of twod = quantize_pow2(clip_scalar(d, d_min, d_max))
# ensure that dynamic range is in specified rangexmax = clip_scalar(xmax, xmax_min, xmax_max) # compute min/max value that we can representif sign
[Figure] [Figure]
[Figure]
B. P OWER - OF - TWO QUANTIZATION B.2. C ASE P1: P ARAMETRIZATION WITH RESPECT TO b AND q MAX
[Figure] [Figure]
def clip_scalar(v
[Figure]
min_value
[Figure]
max_value
[Figure]
:return F
[Figure]
minimum_scalar(F
[Figure]
maximum_scalar(v, min_value), max_value) def
[Figure]
scalar(v
[Figure]
shape
[Figure]
:return F
[Figure]
broadcast(F
[Figure]
reshape(v
[Figure] [Figure]
* len(shape), inplace=False), shape
[Figure]
shape)
def quantize
[Figure]
pow2(
[Figure]
:return 2 ** F
[Figure]
(F
[Figure]
log(F
[Figure]
abs(v
[Figure]
np
[Figure]
[Figure] [Figure] [Figure]
[Figure]
ConstantInitializer
[Figure]
[Figure] [Figure]
ConstantInitializer
[Figure]
as
[Figure]
fix_parameters
[Figure]
# ensure that bitwidth is in specified range and an integern =
[Figure]
(clip_scalar(n
[Figure]
n_min
[Figure]
n
[Figure]
max)
[Figure]
if sign:n = n - 1if with_zero:n = n - # ensure that dynamic
[Table]
[Table] [Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
B.2. C ASE P2: P ARAMETRIZATION WITH RESPECT TO b AND q MIN
1 def parametric_pow2_quantize_b_xmin(x, sign, with_zero,n_init, n_min, n_max,xmin_init, xmin_min, xmin_max,fix_parameters=False):"""Parametric version of ‘pow2_quantize‘ where thebitwidth ‘n‘ and the smallest value ‘xmin‘ are learnable parameters.
Returns:~nnabla.
Variable: N-D array.
"""def clip_scalar(v, min_value, max_value):return F.minimum_scalar(F.maximum_scalar(v, min_value), max_value) def broadcast_scalar(v, shape):return F.broadcast(F.reshape(v, (1,) * len(shape), inplace=False), shape=shape) def quantize_pow2(v):return 2 ** F.round(F.log(F.abs(v)) / np.log(2.))
n = get_parameter_or_create("n", (),
ConstantInitializer(n_init),need_grad=True,as_need_grad=not fix_parameters)xmin = get_parameter_or_create("xmin", (),
ConstantInitializer(xmin_init),need_grad=
True,as_need_grad=not fix_parameters)
# ensure that bitwidth is in specified range and an integern =
F.round(clip_scalar(n, n_min, n_max))if sign:n = n - 1if with_zero:n = n - # ensure that minimum dynamic range is in specified range and a power-of-twoxmin = quantize_pow2(clip_scalar(xmin, xmin_min, xmin_max)) # compute min/max value that we can representxmax = xmin * (2 ** ((2 ** n) - 1)) # broadcast variables to correct sizexmin = broadcast
[Table]
scalar(xmin, shape=x.shape)
[Table]
broadcast
[Table]
scalar(
[Table]
, shape=x.shape)
[Table]
if
[Table]
if not
[Table]
:x
[Table]
(x)
[Table]
[Table] [Table]
if with
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
else:
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
False #
[Table]
[Table]
# take care of values outside of dynamic rangereturn
[Table]
1 def parametric_pow2_quantize_xmin_xmax(x, sign, with_zero,xmin_init, xmin_min, xmin_max,xmax_init, xmax_min, xmax_max,fix_parameters=False):"""Parametric version of ‘pow2_quantize‘ where themin value ‘xmin‘ and max value ‘xmax‘ are learnable parameters.
Returns:~nnabla.
Variable: N-D array.
"""def clip_scalar(v, min_value, max_value):return F.minimum_scalar(F.maximum_scalar(v, min_value), max_value) def broadcast_scalar(v, shape):return F.broadcast(F.reshape(v, (1,)
* len(shape), inplace=False), shape=shape) def quantize_pow2(v):return 2.
** F.round(F.log(F.abs(v)) / np.log(2.)
xmin = get_parameter_or_create("xmin", (),
ConstantInitializer(xmin_init),need_grad=
True,as_need_grad=not fix_parameters)xmax = get_parameter_or_create("xmax", (),
ConstantInitializer(xmax_init),need_grad=
True,as_need_grad=not fix_parameters)
# ensure that minimum dynamic range is in specified range and a power-of-twoxmin = quantize_pow2(clip_scalar(xmin, xmin_min, xmin_max))
# ensure that minimum dynamic range is in specified range and a power-of-twoxmax = quantize_pow2(clip_scalar(xmax, xmax_min, xmax_max))
# broadcast variables to correct sizexmin = broadcast_scalar(xmin, shape=
x.shape)xmax = broadcast_scalar(xmax, shape=x.shape) # if unsigned, then quantize all negative values to zeroif not sign:x = F.relu(x) # compute absolute value/sign of inputax =
F.abs(x)sx = F.sign(x) if with_zero:# prune smallest elements (in magnitude) to zero if they are smaller# than ‘x_min / \sqrt(2)‘x_threshold = xmin / np.sqrt(2) idx1 = F.greater_equal(ax, x_threshold) *
[Table]
.less(
[Table]
, xmin)idx2 = F.greater_equal(
[Table]
, xmin)
* F.less(
[Table]
, xmax)idx3 = F.greater_equal(
[Table]
{+, xmax)else:idx1 = F.less(ax, xmin)idx
+}
= F.greater_equal(ax, xmin) *
{+F.less(ax, xmax)idx
+}
= F.greater_equal(ax, xmax) # do not backpropagate gradient through indicesidx1.need_grad =
Falseidx2.need_grad = Falseidx3.need_grad =
False # do not backpropagate gradient through signsx.need_grad =
False # take care of values outside of dynamic rangereturn sx * (xmin * idx1 + quantize_pow2(ax) * idx2 + xmax * idx3)