Uniform-Precision Neural Network Quantization via Neural Channel Expansion

1 INTRODUCTION . Deep neural networks ( DNNs ) have reached human-level performance in a wide range of domains including image processing ( He et al . ( 2016 ) ; Tan & Le ( 2019 ) ) , object detection ( Ren et al . ( 2015 ) ; Liu et al . ( 2016 ) ; Tan et al . ( 2020 ) ) , machine translation ( Wu et al . ( 2016 ) ; Devlin et al . ( 2018 ) ) , and speech recognition ( Zhang et al . ( 2016 ) ; Nassif et al . ( 2019 ) ) . However , tremendous computation and memory costs of these state-of-the-art DNNs make them challenging to deploy on resourceconstrained devices such as mobile phones , edge sensors , and drones . Therefore , several edge hardware accelerators specifically optimized for intensive DNN computation have emerged , including Google ’ s edge TPU ( Google ( 2019 ) ) and NVIDIA ’ s NVDLA ( NVIDIA ( 2019 ) ) . One of the central techniques innovating these edge DNN accelerators is the quantization of deep neural networks ( QDNN ) . QDNN reduces the complexity of DNN computation by quantizing network weights and activations to low-bit precision . Since the area and energy consumption of the multiplyaccumulate ( MAC ) unit can be significantly reduced with the bit-width reduction ( Sze et al . ( 2017 ) ) , thousands of them can be packed in a small area . Therefore , the popular edge DNN accelerators are equipped with densely integrated MAC arrays to boost their performance in compute-intensive operations such as matrix multiplication ( MatMul ) and convolution ( Conv ) . Early studies of QDNN focused on the quantization of weights and activations of MatMul and Conv to the same bit-width ( Hubara et al . ( 2016 ) ; Rastegari et al . ( 2016 ) ; Zhou et al . ( 2016 ) ) . This uniformprecision QDNN gained popularity because it simplifies the dense MAC array design for edge DNN accelerators . However , uniform bit allocation did not account for the properties of individual layers in a network . Sakr & Shanbhag ( 2018 ) showed that the optimal bit-precision varies within a neural network from layer to layer . As a result , uniform-precision quantization may lead to sub-optimal inference accuracy for a given network . Mixed-precision networks address this limitation by optimizing bit-widths at each layer . In this approach , the sensitivity of the layer to the quantization error is either numerically estimated ( Zhou et al . ( 2017 ) ; Dong et al . ( 2019 ) ) or automatically explored under the framework of neural architecture search ( NAS , Wang et al . ( 2019 ) ; Elthakeb et al . ( 2018 ) ) to allocate bit-precision properly . However , mixed-precision representation requires specific variable precision support in hardware , restricting computation units ’ density and power efficiency ( Camus et al . ( 2019 ) ) . Therefore , mixed-precision support imposes a significant barrier for the low-profile edge accelerators with stringent hardware constraints . In this work , we propose a novel NAS based hardware-friendly DNN quantization method that can address the layer-wise heterogeneous sensitivity under uniform-precision quantization . The proposed method explores network structure in terms of the number of channels . Different from the previous work that only includes pruning of the channels in its search space ( Dong & Yang ( 2019 ) ) , we further incorporate the expansion of the channels , thus called neural channel expansion ( NCE ) . During a search of NCE , search parameters associated with different numbers of channels are updated based on each layer ’ s sensitivity to the uniform-precision quantization and the hardware constraints ; the more sensitive to quantization errors , the larger number of channels preferred in that layer . When the preference to the larger number of channels in a layer exceeds a certain threshold , we expand the channels in that layer ’ s search space so that the more number of channels can be explored . Therefore , NCE allows both pruning and expansion of each layer ’ s channels , finding the sweet-spot for the trade-off between the robustness against the quantization error and the hardware cost . We analytically and empirically demonstrate that NCE can adequately facilitate the search to adapt the target model ’ s structure for better quantization accuracy . The experimental results on CIFAR10 and ImageNet show that the network structures adapted from the popular convolutional neural networks ( CNNs ) achieve superior accuracy when the challenging 2-bit quantization is uniformly applied to MatMul and Conv layers . In particular , we achieve the best-to-date accuracy of 74.03/91.63 % ( Top-1/Top-5 ) for NCE-ResNet50 on ImageNet with slightly lower FLOPs and 30 % reduced number of parameters . Our contributions can be summarized as follows : • We propose a new NAS-based quantization algorithm called neural channel expansion ( NCE ) , which is equipped with a simple yet innovative channel expansion mechanism to balance the number of channels across the layers under uniform-precision quantization . • We provide an in-depth analysis of NCE , shedding light on understanding the impact of channel expansion for compensation of quantization errors . • We demonstrate that the proposed method can adapt the structure of target neural networks to significantly improve the quantization accuracy . 2 RELATED WORK . Neural architecture search : The goal of NAS is to find a network architecture that can achieve the best test accuracy . Early studies ( Zoph & Le ( 2016 ) ; Zoph et al . ( 2018 ) ) often employed meta-learners such as reinforcement learning ( RL ) agents to learn the policy for accurate network architectures . However , RL-based approaches may incur prohibitive search costs ( e.g. , thousands of GPU hours ) . As a relaxation , differentiable neural architecture search ( DNAS ) has been proposed ( Liu et al . ( 2018 ) ) , which updates the search parameters and the weights via bi-level optimization . Recent DNAS approaches considered hardware constraints such as latency , the number of parameters , and FLOPs so that the search explored the trade-off between the cross-entropy loss and the hardware constraint loss . This search resulted in the discovery of light-weight models . As an example , Dong & Yang ( 2019 ) explored channel pruning that satisfies the target hardware constraints . In this work , we adopt the successful NAS framework in the domain of QDNN , for which we devise a novel channel expansion search to robustify networks against the quantization errors . Low-precision quantization of deep neural network : QDNN has been actively studied in the literature . Early work on QDNN ( Hubara et al . ( 2016 ) ; Rastegari et al . ( 2016 ) ; Zhou et al . ( 2016 ) ) introduced the concept of a straight-through estimator ( STE ) for the approximation of gradients of the non-differentiable rounding operation . This approximation enabled uniform-precision ( 1- or multi-bit ) quantization during the model training procedure , which fine-tunes the weight parameters towards lower training loss . QDNN techniques have evolved to adaptively find the quantization step size ( Choi et al . ( 2018 ) ; Zhang et al . ( 2018 ) ; Jung et al . ( 2019 ) ; Esser et al . ( 2020 ) ) , which significantly enhanced the accuracy of the uniform-precision quantization . However , this line of research lacks consideration of the heterogeneous quantization sensitivity for individual layers in a network . On the other hand , mixed-precision quantization allows layer-specific bit-precision optimization ; the higher bit-precision is assigned to the more quantization sensitive layers . Zhou et al . ( 2017 ) ; Dong et al . ( 2019 ) numerically estimated the sensitivity via approximating the impact of quantization errors on model prediction accuracy . Wang et al . ( 2019 ) ; Elthakeb et al . ( 2018 ) employed a reinforcement learning framework to learn the bit-allocation policy . Wu et al . ( 2018 ) ; Cai & Vasconcelos ( 2020 ) adopted DNAS with the various bit-precision operators in the search space . However , mixed-precision representation requires specific variable precision support in hardware , restricting computation units ’ density and power efficiency ( Camus et al . ( 2019 ) ) . Therefore , mixed-precision support imposes a significant barrier for the low-profile edge accelerators with stringent hardware constraints . In this work , we exploit NAS to address the layer-wise heterogeneous sensitivity under uniform-precision quantization . Channel expansion for accurate DNN quantization : Researchers have actively studied channel expansion for accurate DNN quantization . Pioneering work by Mishra et al . ( 2018 ) ( WRPN ) demonstrated that an increased number of channels during the training helped regain QDNN accuracy , but this work lack discussion about the detailed mechanism of channel expansion compensating the quantization error . Zhao et al . ( 2019 ) and Park & Choi ( 2019 ) further attempted to split the channels with large magnitude weights in the pre-trained models . This channel splitting reduced the dynamic range of weights to be represented with lower bit-precision ( 6-8-bits ) . Regarding the control over the dynamic range , Meller et al . ( 2019 ) also adjusted the scale factors of the weight parameters after training to balance the dynamic range across the layers . However , these approaches focused on the numerical remedy for quantization of pre-trained models ( with relatively high bit-precision ) . Thus , it not straightforward to extend their work for quantization-aware training , which is necessary for ultra-low bit QDNN . We provide insights with empirical supports that the structure with the channel expanded layers itself matters , reducing the dynamic range of activation during the quantization-aware training . These exciting insights become a pivotal motivation for us to explore channel expansion in the NAS framework . 3 NEURAL CHANNEL EXPANSION . In this section , we explain the detail of our neural channel expansion method . Similar to TAS ( Dong & Yang ( 2019 ) ) , we construct the search space over the number of channels C = { 1 : cout } with the search parameters α ∈ R|C| . Then the output activation is computed as the weighted sum of sampled activations with a different number of channels aligned via channel-wise interpolation ( CWI ) : Ô = ∑ j∈I Softmax ( αj ; { αk } k∈I ) × CWI ( O1 : Cj , max { ckout } k∈I ) , ( 1 ) where output activation Oj:1≤j≤cout = ∑cin k=1Q ( X { k , : , : } ) ∗ Q ( W { j , k , : , : } ) is computed with input activation X and weight parameters W quantized by the quantizer Q , and I is the sampled subset of C. During the search , the search parameters are updated via channel selection based on the trade-off between the cross-entropy loss and the hardware constraint loss ( e.g. , FLOPs ) . In TAS , the number of channels ( |C| ) is fixed , limiting the exploration scope to the pruning . In NCE , we enable channel expansion of individual layers when the search parameter associated with the maximum number of channels exceeds a certain threshold . The intuition is that if one layer is susceptible to the quantization errors , its search parameters are updated toward the preference for a larger number of channels to decrease the cross-entropy loss . With this simple expansion condition , we can expand channels to those layers affected most by the quantization errors and prune channels of the other layers robust to quantization ; therefore , the overall hardware constraints are met . Algorithm 1 summarizes the overall procedure . NCE consists of three phases : warm-up , search , and train . As advocated by Wu et al . ( 2018 ) and Bender et al . ( 2020 ) , we first perform a warm-up of the entire super-net so that all the super-net weight parameters can be reasonably initialized . The search phase consists of the iterative updates of weights ( w ) and the search parameters ( α ) via bi-level optimization . The updated search parameter associated with the maximum number of channels is compared with a threshold ( pre-determined as a hyper-parameter ) to identify if a layer needs a channel expansion in each layer . When a channel expansion happens ( = Expand ) , the additional weight parameters are added to that layer ( and the search parameter is also copied ) , increasing the number of channels . Once the search is done , the candidate model is derived by the `` winner-takes-all '' strategy ; i.e. , for each layer , the number of channels with the largest magnitude search parameter is selected . Algorithm 1 : Neural Channel Expansion Input : Split the training set into two dis-joint sets : Dweight and Darch ( n ( Dweight ) = n ( Darch ) ) Search Parameter : { αl1 , α l 2 , .. , α l n } ∈ Al , { A1 , A2 , .. , AL } ⊂ A , L =number of layer Expand Threshold : T 1 For Warm-up Epoch do 2 Sample batch data Dw from Dweight and network from A ∼ U ( 0 , 1 ) 3 Calculate Lossweight on Dw to update network weights 4 End For 5 For Search Epoch do 6 Sample batch data Dw from Dweight and network from Softmax ( A ) 7 Calculate Lossweight on Dw to update network weights 8 Sample batch data Da from Darch and network from Softmax ( A ) 9 Calculate Lossarch on Da to update A 10 For layer do 11 j ← # Al 12 If Softmax ( αlj ; { αlk } k∈j ) ≥ T do 13 Expand search space ( αlj+1 ) 14 αlj+1 ← αlj # copy search parameter 15 End if 16 End for 17 End for 18 Derive the searched network from A 19 Randomly initialize the searched network and optimize it on the training set

Summary: The authors propose neural channel expansion (NCE), a neural architecture search (NAS) and quantization method. Existing NAS+Q methods typically search for the architecture of the DNN along with the precision at each layer, maximizing accuracy while respecting some kind of hardware constraint. The result is a DNN with mixed-precision, which is challenging for most existing hardware (which only support one or a few precisions). NCE keeps precision the same in each layer, and instead uses the precision sensitivity signal in the NAS to adjust the width of the layer (expand or shrink). The result is uniform-precision, hardware-friendly DNN.