Title: SMART TERNARY QUANTIZATION

Abstract: Neural network models are resource hungry. Low bit quantization such as binary and ternary quantization is a common approach to alleviate this resource requirements. Ternary quantization provides a more flexible model and often beats binary quantization in terms of accuracy, but doubles memory and increases computation cost. Mixed quantization depth models, on another hand, allows a trade-off between accuracy and memory footprint. In such models, quantization depth is often chosen manually (which is a tiring task), or is tuned using a separate optimization routine (which requires training a quantized network multiple times). Here, we propose Smart Ternary Quantization (STQ) in which we modify the quantization depth directly through an adaptive regularization function, so that we train a model only once. This method jumps between binary and ternary quantization while training. We show its application on image classification.

1. INTRODUCTION:

Deep Neural Networks (DNN) models have achieved tremendous attraction because of their success on a wide variety of tasks including computer vision, automatic speech recognition, natural language processing, and reinforcement learning (Goodfellow et al., 2016) . More specifically, in computer vision DNN have led to a series of breakthrough for image classification (Krizhevsky et al., 2017) , (Simonyan & Zisserman, 2014) , (Szegedy et al., 2015) , and object detection (Redmon et al., 2015) , (Liu et al., 2015) , (Ren et al., 2015) . DNN models are computationally intensive and require large memory to store the model parameters. Computation and storage resource requirement becomes an impediment to deploy such models in many edge devices due to lack of memory, computation power, battery, etc. This motivated the researchers to develop compression techniques to reduce the cost for such models.

Recently, several techniques have been introduced in the literature to solve the storage and computational limitations of edge devices. Among them, quantization methods focus on representing the weights of a neural network in lower precision than the usual 32-bits float representation, saving on the memory footprint of the model. Binary quantization (Courbariaux et al., 2015) , (Hubara et al., 2016) , (Rastegari et al., 2016) , (Zhou et al., 2016) , (Lin et al., 2017) represent weights with 1 bit precision and ternary quantization (Lin et al., 2015) , (Li & Liu, 2016) , (Zhu et al., 2016) with 2 bits precision. While the latter frameworks lead to significant memory reduction compared to their full precision counterpart, they are constrained to quantize the model with 1 bit or 2 bits, on demand. We relax this constraint, and present Smart Ternary Quantization (STQ) that allows mixing 1 bit and 2 bits layers while training the network. Consequently, this approach automatically quantizes weights into binary or ternary depending upon a trainable control parameter. We show that this approach leads to mixed bit precision models that beats ternary networks both in terms of accuracy and memory consumption. Here we only focus on quantizing layers because it is easier to implement layer-wise quantization at inference time after training. However, this method can be adapted for mixed precision training of sub-network, block, filter, or weight easily. To the best of our knowledge this is the first attempt to design a single training algorithm for low-bit mixed precision training.

2. RELATED WORK:

There are two main components in DNN's, namely, weight and activation. These two components are usually computed in full precision, i.e. floating point 32-bits. This work focuses on quantizing the weights of the network, i.e. generalizing BinaryConnect (BC) of Courbariaux et al. (2015) and Ternary Weight Network (TWN) of Li & Liu (2016) towards automatic 1 or 2 bits mixed-precision using a single training algorithm.

In BC the real value weights w are binarized to w b ∈ {-1, +1} during the forward pass. To map a full precision weight to a binary weight, the deterministic sign function is used,

EQUATION

The derivative of the sign function is zero on R \ {0}. During back propagation, this cancels out the gradient of the loss with respect to the weights after the sign function. Therefore, those weights cannot get updated. Ternary Weight Network (TWN) (Li & Liu, 2016 ) is a neural network with weights constrained to {-1, 0, +1}. The weight resolution is reduced from 32 bits to 2 bits, replacing full precision weights with ternary weights. TWN aims to fill the gap between full precision and binary precision weight.

Compared to binary weight networks, ternary weight networks have stronger expressive capability.

As pointed out in (Li & Liu, 2016) , for a 3×3 weight filter in a convolutional neural network, there is 2 3×3 = 512 possible variation with binary precision and 3 3×3 = 19683 with ternary precision.

Li & Liu (2016) find the closest ternary weights matrix W t to its corresponding real value weight matrix W using μ, Ŵt = arg min

EQUATION

The ternary weight W t is achieved by applying a symmetric threshold ∆ 2016) define a weight dependant threshold ∆ and a scaling factor µ that approximately solves (3). TWN is trained using stochastic gradient descent. Similar to BC and BWN schemes; ternary-value weights are only used for the forward pass and back propagation, but not for the parameter updates. At inference, the scaling factor can be folded with the input X X W ≈ X (µW t ) = (µX) W t , (5) where denotes the convolution. Zhu et al. (2016) proposed a more general ternary method which reduces the precision of weights in neural network to ternary values. However, they quantize the weights to asymmetric values {-µ 1 , 0, +µ 2 } using two full-precision scaling coefficients µ 1 and µ 2 for each layer of neural network. While the method achieve better accuracy as opposed to TWN, its hardware implementation becomes a challenge, because there are two unequal full precision scaling factors to deal with.

W t =    +1 w ij > ∆, 0 |w ij | ≤ ∆, -1 w ij < -∆. (4) Li & Liu (

Our method provides a compromise between BC and TWN and trains weights with a single trainable scaling factor µ. Weights jumps between ternary {-µ, 0, +µ} and binary {-µ, +µ}. This provides a single algorithm for 1 or 2 bits mixed precision.

2.2. REGULARIZATION:

Regularization term is the key to prevent over-fitting problem and to obtain robust generalization for unseen data. Standard regularization functions, such as L 2 or L 1 encourage weights to be concentrated about the origin. However, in case of binary network with binary valued weights, it is more appropriate to have a regularization function to encourage the weights about µ × {-1, +1}, with a scaling factor µ > 0 such as

EQUATION

proposed in Belbahri et al. (2019) .

A straightforward generalization for ternary quantization is

EQUATION

)

Regularizer ( 6) encourages weights about {-µ, +µ}, and (7) about {-µ, 0, +µ}. The two functions are depicted in Figure 1 . These regularization functions are only useful when the quantization depth is set before training start. We propose a more flexible version to smoothly move between these two functions using a shape parameter. 

3. SMART TERNARY QUANTIZATION:

Here we propose an adaptive regularization function that switches between binary regularization of (6) and ternary regularization of ( 7)

EQUATION

in which µ is a trainable scaling factor, and β ∈ ( π 4 , π 2 ) controls the transition between ( 6) and ( 7). As a special case β → π 2 converges to the binary regularizer (6) and β → π 4 coincide with the ternary regularizer (7), depicted in Figure 2 . A large value of tan(β) repels estimated weights away from zero thus yielding binary quantization, and small tan(β) values encourage zero weights. The shape parameter β controls the quantization depth. Quantization is done per layer, therefore we let β very per layer. We recommend to regularize β about π 2 i.e. preferring binary quantization apriori

EQUATION

in which γ controls the proportion of binary to ternary layers.

For a single filter W the regularization function is a sum over its elements

EQUATION

)

Large values of γ encourage binary layers. In each layer, weights are pushed to binary or ternary values, depending on the trained value of the corresponding β. A generalization of (9) towards L p norms of Belbahri et al. ( 2019) is also possible. However, here we only focus on regularizer constructed using the L 1 norm as the accuracy did not change significantly by using L p norm with different values of p.

The introduced regularization function is added to the empirical loss function L(.). The objective function defined on weights W, scaling factors µ, and quantization depths β is optimized using back propagation

EQUATION

where k indexes the channel, and l indexes the layer. One may use a different regularization constant λ l for each layer to keep the impact of the regularization term balanced across layers, indeed different layers may involve different number of parameters. We set λ l = λ #W l where λ is a constant, and #W l is the number of weights in layer l. We propose to use the same threshold-based function of Li & Liu (2016) (3), but with a fixed threshold ∆ l per layer l. Note that Li & Liu (2016) propose a weight-dependant threshold. We let the possibility for the weights to only accumulate about {-µ, +µ} and not about 0, depending on β.

One may set ∆ l to have the same balanced weights in {-µ, 0, +µ} at initialization for all layers and let the weights evolve during training. Formally, if σ l is the standard deviation of the initial Gaussian weights in layer l, we propose ∆ l = 0.2 × σ l . The probability that a single weight lies in the range [-∆ l , ∆ l ] is ≈ 0.16. All the weights falling in this range will be quantized as zeros after applying the threshold function.

Weights are naturally pushed to binary or ternary values depending on β l during training. Eventually, a threshold δ close to π 2 ≈ 1.57 defines the final quantization depth for each layer.

Final quantization depth of layer l : Binary β l ≥ δ, Ternary β l < δ

4. EXPERIMENT:

We run experiments on two simple image classification tasks MNIST (LeCun et al., 1998) and CIFAR10 (Krizhevsky & Hinton, 2009) datasets. We compare our method (STQ) with BinaryConnect (BC) of Courbariaux et al. (2015) , Binary Weight Networks (BWN) of Rastegari et al. (2016) , Ternary Weights Network (TWN) of Li & Liu (2016) and also with a Full Precision network (FP).

The quality of the compression is measured only in terms of memory, it is difficult to compare mixed precision models, with binary and ternary, in terms of consumed energy. Assume that n l is the quantization depth for the layer l and #W l is the number of weights in layer l, therefore the compression ratio is

L l=1 #W l ×32 L l=1 #W l ×n l

. The compression ratio for a binary network is 32, for a ternary network 16, and our approach falls in between.

STQ network generalizes binary and ternary regularization in a single regularization function. First we show how to control the proportion of binary and ternary layers using γ in (8). Figure 3 clarifies the effect of γ on the weight distribution. When γ is large, β is encouraged towards π 2 which corresponds to binary quantization. Consequently, weights are pushed about {-µ, +µ} and 0 is removed from the trained values, see Figure 3a . On the contrary, when γ is small, β tends to π 4 and the weights started including 0 in their values, see Figure 3b . 

4.1. MNIST:

MNIST is an image classification benchmark dataset with 28 × 28 gray-scales images representing digits ranging from 0 to 9. The dataset is split into 60k training images and 10k testing images. We used the LeNet-5 (LeCun et al., 1998) architecture consisting of 5 layers, 2 convolution followed by maxpooling, stacked with two fully connected layers and a softmax layer at the end. We train the network for 60 epochs using Adam optimizer. We used the initial learning rate of 0.01, but divided by 10 in epoch 15 and 30 to stabilize training. The batch size is set to 64 with L 2 weight decay constant 10 -4 only for BC, BWC and TWN. The full precision LeNet-5 is trained with no regularization as it provided a superior accuracy. STQ is trained with λ = 0.1 and γ = 1 × 10 -2 and the effective regularization constant is divided by the number of weights in each layer to compensate for the layer size. The best validation accuracy for each method is reported in Table 1 , as well as the quantization depth, and the overall compression ratio. We observe that STQ network quantized the first two convolutional layers in 1 bit, and the last fully-connected layers in 2 bits. The accuracy improvement and the compression ratio is marginal for simple task and simple architectures. The effect of smart training becomes more visible for more complex tasks with more layers. We train the network for 150 epochs, using Adam optimizer with the initial learning rate 0.001 divided by 10 at epochs 40 and 80. The batch size is set to 64 with L 2 weight decay constant 10 -4 , moreover λ = 0.1, γ = 10 -3 for STQ. The best validation accuracy for each method is reported in Table 2 . STQ beats pure 2 bits network TWN, even in terms of accuracy. It recommends three 1 bit layers for VGG-7 and seven 1 bit layers for VGG-16. The compression ratio is significantly higher than a ternary network. The weight distribution of each layers are depicted in Figure 4 for the VGG-7 architecture. Weights are pushed to {-µ, +µ} or {-µ, 0, +µ} depending on the shape parameter β. 

EQUATION

5. CONCLUSION:

Smart ternary quantization (STQ) is a training method to build a 1 and 2 bits mixed quantized layers. Depth optimization requires training network multiple times which is costly, specially if the network is complex. This approach successfully combines quantization with different depths, while training the network only once. We tried layer-wise quantization, because it is more suitable for mixed-precision inference implementation. However, subnetwork, block, filter, or weight mixed quantization is feasible using a similar algorithm.

STQ makes manual tuning of quantization depth unnecessary. It allows to improve the memory consumption, by automatically quantizing some layers with smaller precision. This method sometimes even outperforms pure ternary networks in terms of accuracy thank to a formal regularization function that shapes trained weights towards mixed-precision.

It is well-known that some layers are more resilient to aggressive quantization. Our introduced methodology trains network similar to pure ternary but give an insight about which layers can be simplified further by going to binary quantization depth.

:



Figures and Tables:

Figure 1: Binary and ternary regularizers; R 1 encourages binary weights, with minimums at {-µ, +µ}, and R 2 encourages ternary weights, with minimums at {-µ, 0, +µ}.

Figure 2: Adaptive regularization function. When β → π 2 the regularization function switches from ternary to binary.

Figure 3: Effect of γ on the weights distribution of a layer while training.

Figure 4: Layer wise weights distribution in VGG-7 for STQ. The weights are pushed to binary when the shape parameter β is close to π 2 ≈ 1.57.

Smart ternary quantization (STQ) compared with Binary Connect (BC), Binary Weight Network (BWN), Terneary Weight Network (TWN), and Full Precision (FP) on MNIST dataset. image classification benchmark that contains 32 × 32 RGB images from ten classes. The dataset is split into 50k training images and 10k testing images. All are normalized using mean = (0.4914, 0.4822, 0.4465) and std = (0.247, 0.243, 0.261). For the training session, we pad the sides of the images with 4 pixels, then a 32 × 32 crop is sampled, and flipped horizontally at random.We use two VGG-like architectures, i) VGG-7 architecture defined inLi & Liu (2016) in which we apply batch normalization after each layer and use ReLU activations, ii) a standard VGG-16 architecture. Deep architectures are sensitive to early layer quantization. As commonly practiced, we did not quantize the first and the last layers in VGG-16 for all competing methods.

Table content: <table><tr><td/><td>Quantization depth</td><td>Compression</td><td>Accuracy</td></tr><tr><td/><td>per layer (-bits)</td><td>ratio</td><td>(top-1)</td></tr><tr><td>BC</td><td>1-1-1-1-1</td><td>32</td><td>99.35</td></tr><tr><td>BWN</td><td>1-1-1-1-1</td><td>32</td><td>99.32</td></tr><tr><td>TWN</td><td>2-2-2-2-2</td><td>16</td><td>99.38</td></tr><tr><td>STQ</td><td>1-1-2-2-2</td><td>16.3</td><td>99.37</td></tr><tr><td>FP</td><td/><td>1</td><td>99.44</td></tr><tr><td>4.2 CIFAR10</td><td/><td/><td/></tr><tr><td>CIFAR10 is an</td><td/><td/><td/></tr></table>

Smart ternary quantization (STQ) compared with Binary Connect (BC), Binary Weight Network (BWN), Terneary Weight Network (TWN), and Full Precision (FP) on CIFAR10 dataset.

Table content: <table><tr><td>Architecture Method</td><td>per layer (-bits) Quantization depth</td><td>ratio Compression</td><td>Accuracy</td></tr></table>

Peer review (iclr2020):
The Paper talks about the Smart Ternary Quantization method that improves the quantization over binary and ternary quantizations by specifying an adaptive quantization. The proposed regularization function is covered in detail and the results are evaluated on MNIST and CIFAR10 datasets The authors can improve the submission by 1. evaluating more modern networks with bigger datasets, as opposed to the ones demonstrated. 2. describing gamma (eq 10) , it wasn*t clear why that parameter was introduced (in addition to the beta) and it*s significance, what was more confusing is coverage for this instead of beta in the experimental setup 3. why the LR needed to be modified for the described method