BatchQuant: Quantized-for-all Architecture Search with Robust Quantizer

1 Introduction . In order to deploy deep learning models on resource-constrained edge devices , careful model optimization , including pruning and quantization is required . While existing works have demonstrated the effectiveness of model optimization techniques in speeding up model inference [ 1 , 2 , 3 , 4 ] , model optimization increases human labor by introducing extra hyperparameters . Consequently , automated methods such as neural architecture search ( NAS ) [ 5 , 6 , 7 , 8 , 9 , 10 , 11 , 12 , 13 , 14 , 15 , 16 ] and automated quantization policy search [ 17 , 18 , 19 ] have emerged to alleviate the human bandwidth required for obtaining compact models with good performance . In this paper , we focus on finding the best of both worlds—the best combination of architecture and mixed-precision quantization policy . However , combining two complex search spaces is inherently challenging , not to mention that quantization usually requires a lengthy quantization-aware training ( QAT ) procedure to recover performance . Thus , previous methods tend to employ proxies to estimate the performance of an architecture and quantization policy combination . For example , APQ [ 20 ] performs QAT on each of 5000 architecture and quantization policy combinations for 0.2 GPU hours ∗Corresponding author : Haoping Bai haoping_bai @ apple.com . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . and uses the sampled combinations to train an accuracy predictor . BATS [ 21 ] searches for cell structures that are repeatedly stacked to form the target architecture . Proxy-based methods necessitate both careful treatments to ensure reliable ranking of quantized architectures and a time-consuming retraining procedure when the target quantized architecture is identified , rendering proxy-based approaches impractical for frequently changing deployment scenarios . To avoid the lengthy retraining process , NAS methods that train a single-shot weight-sharing supernet [ 12 , 22 ] are ideal . However , many previous works have shown evidence that QAT of mixedprecision supernets can easily become highly unstable [ 23 , 21 ] . As a result , existing single-shot quantized architecture search methods usually limit the size of the combined search space . For example , SPOS [ 12 ] sacrifices architecture search space size to only allow channel search and requires retraining to recover performance . OQA [ 23 ] limits its quantization policy search space to fixed-precision quantization policies and trains a separate set of weights for each bitwidth . To successfully train the mixed-precision supernet , we propose BatchQuant ( BQ ) . Analogous to batch normalization , BQ leverages batch statistics to adapt to the shifting activation distribution as a result of quantized subnet selection , offering better robustness to outliers than quantizer with vanilla running min/max based scale estimation , and better flexibility than a learnable quantizer that only learns a fixed set of parameters . Without limiting the architecture search space , our joint architecture and quantization policy search space contains over 1076 possible quantized subnets , providing much more flexibility than previous search spaces ( e.g . The OQA search space has 1020 possible quantized subnets ) . While our approach and OQA both follow the supernet training strategy introduced in [ 22 ] , our weight-sharing supernet takes only 190 epochs to train despite the complex search space , significantly less than the 495 epochs required by OQA . We further leverage the NSGA-II algorithm to produce a Pareto set of quantized architectures that densely covers varying complexity constraints , eliminating the marginal cost of adapting architecture to new deployment scenarios . The contributions of the paper are • To the best of our knowledge , we present the first result to train one-shot weight-sharing supernet to support subnets with arbitrary mixed-precision quantization policy without retraining . • We propose BatchQuant , an activation quantizer formulation for stable mixed-precision supernet training . The general formulation of BatchQuant allows easy adaptation of new scale estimators . • Compared with existing methods , our method , QFA , takes a shorter time to train a significantly more complex supernet with over 1076 possible quantized subnets and discovers quantized subnets at SOTA efficiency with no marginal search cost for new deployment scenarios . 2 Related Work . 2.1 Mixed Precision Quantization . Different layers of a network have different redundancy and representation power , each layer will react to quantization differently and may achieve varying levels of efficiency gain on hardware . In fact , many fixed-bitwidth quantization methods are inherently mixed-precision by making design decisions to leave the first convolution layer , batch-norm layers [ 24 ] , squeeze and excitation layers [ 25 ] , and the last fully connected layer at full precision / int8 precision [ 3 , 2 , 19 , 12 , 23 ] . As a result , mixedprecision quantization methods [ 26 , 18 , 27 , 17 ] emerge in the place of fixed precision quantization and are receiving increasing attention from hardware manufacturers [ 17 ] . In contrast to mixed-precision quantization methods that optimize for a single quantization policy on a single architecture , QFA allows the discovery of the most suitable quantization policy for arbitrary subnet architectures . Adaptive Quantization is also closely related to our work . Methods include Adabits [ 28 ] , AnyPrecision DNN [ 29 ] , Gradient ℓ1 regularization or quantization robustness [ 30 ] , and KURE [ 31 ] train model that can adaptively switch to different bitwidth configuration during inference . Most methods [ 28 , 29 , 30 ] only address the case of fixed precision quantization where all weights share the same bitwidth and activations share another . Our quantization policy search space offers much more flexibility in terms of layerwise mixed-precision quantization . 2.2 Joint Mixed-Precision Quantization and Neural Architecture Search . There exist many exciting works in the intersection of mixed precision quantization and NAS . JASQ [ 32 ] employs population-based training and evolutionary search to produce a quantized architecture under a combined accuracy and model size objective . DNAS [ 19 ] is a differentiable NAS method that optimizes for a weighted combination of accuracy and resource constraints . SPOS [ 12 ] trains a quantized one-shot supernet to search for bit-width and network channels for heavy ResNet search space . APQ [ 20 ] builds upon a full precision one-shot NAS supernet and trains 5000 quantized subnets to build a proxy accuracy predictor . The best quantized architecture proposed by the predictor is then retrained . Most existing methods fall into two categories . The first optimizes for a single weighted objective of accuracy and complexity and produces only a single or a few quantized architectures [ 32 , 19 ] , which is difficult to scale to multiple deployment scenarios . The second category is capable of estimating the performance for many quantized architectures by using proxies such as weight-sharing supernet and partially trained models [ 12 , 20 ] . Due to the accuracy degradation in weight-sharing supernet in SPOS and the use of proxy in APQ , both methods require retraining when the best quantized architecture is discovered . On the contrary , our approach does not require retraining . Training the weight-sharing supernet with BQ reduces accuracy degradation , allowing quantized subnets to reach competitive performance without retraining . As a result , QFA enjoys no marginal cost for deploying quantized subnet to new scenarios . Table 1 details the difference of our approach from other architecture search approaches . Table 2 compares the algorithmic complexity of our method with other mixed-precision quantized architecture search methods . 3 Stabilizing Mixed Precision Supernet Training with BatchQuant . We adopt the single-path weight-sharing MobileNetV3 [ 33 ] supernet formulated in [ 22 ] that supports adaptive input resolution , network kernel size , depth , and width . In addition , we assign simulated quantization operations ( quantizers ) [ 34 ] , one for each bitwidth , to each weight and activation tensor within the supernet to support quantization-aware training ( QAT ) . We allow bitwidths { 2 , 3 , 4 } for both weight quantizers and activation quantizers . The compact weight-sharing supernet formulation allows us to share a single set of quantizers across all the subnets . However , such compactness also leads to instability . Specifically , supernet training becomes highly unstable when activations from different subnets are quantized by a shared quantizer . In the following sections , we will review the fundamentals of quantization and address why shared activation quantizer can cause unstable supernet training . Then , we introduce BatchQuant as a solution to stabilize supernet training . 3.1 Quantization Preliminaries . To help address the difficulty in training the mixed-precision supernet , we start by introducing common notations for quantization . WLOG , we consider the case of uniform affine ( asymmetric ) quantization . Let x = { x1 , · · · , xN } be a floating point vector/tensor with range ( xmin , xmax ) that needs to be quantized to b-bitwidth precision . The integer coding xq will have range [ n , p ] = [ 0 , 2b−1 ] . Then we derive two parameters : Scale ( ∆ ) and Zero-point ( z ) which map the floating point values to integers ( See [ 1 ] ) . The scale specifies the step size of the quantizer and floating point zero maps to zero-point [ 34 ] , an integer which ensures that zero is quantized with no error . The procedure is as follows : ∆ = xmax − xmin p− n , z = clamp ( − ⌊xmin ∆ ⌉ , n , p ) xq = ⌊ clamp ( x ∆ + z , n , p ) ⌉ x̂ = ( xq − z ) ∆ ( 1 ) where ⌊·⌉ indicates the round function and the clamp ( · ) function clamps all values to fall between n and p. The quantized tensor x̂ is then used for efficient computation by matrix multiplication libraries that handle ∆ efficiently ( See [ 35 ] ) . Note that during training , the range ( xmin , xmax ) of activations are not known beforehand . Therefore , standard practice is to keep track of an exponential moving average ( EMA ) of past extreme values ( min & max ) to generate the scale estimator ∆̂ . As both the calculation of z and xq relies on a good ∆ estimate , we will later provide insights into why using EMA estimator ∆̂ for activation quantization in the supernet could lead to unstable training . 3.2 Weight Quantization with LSQ . Since the value distribution of a weight tensor is relatively stable across updates and is empirically observed to be symmetrical around zero , many methods treat ∆ as a learnable parameter and discard the zero point z [ 36 , 4 ] . We leverage the symmetric LSQ quantizer [ 36 ] to quantize weight tensors in our supernet as follows : xq = ⌊ clamp ( x ∆ , n , p ) ⌉ x̂ = xq∆ ( 2 ) where [ n , p ] = [ −2b−1 , 2b−1−1 ] . Note the rounding operation ⌊·⌉ has 0 derivative almost everywhere , QAT applies the straight through estimator ( STE ) [ 37 ] to allow gradients to backpropagate through . 3.3 Challenges of Activation Quantization in One-shot Supernet . While weight quantization for weight-sharing supernet is similar to that of normal networks , activation quantization requires careful treatment to enable stable supernet training . According to the insights in [ 31 ] , as well as shown in Equation ( 1 ) , an unstable ∆ estimate could impact the quantized tensor x̂ and negatively impact QAT performance . Given the significant delay in the EMA estimator ∆̂ when the ranges shift rapidly , practical approaches such as [ 34 ] completely disables activation quantization for the first 50 thousand to 2 million steps . However , we now show that one-shot supernet will always have rapidly shifting activation ranges due to subnet sampling . EMA Estimator ∆̂ is problematic . Let N = B × C ×H ×W denote the number of elements within an activation tensor with batch size B , channel number C , height H , and width W . Due to adaptive input resolution and network width ( number of channels ) , the activation map at a given layer will vary in size ( H , W ) and channel number C depending on the activated subnet . For example , at a batch size of 64 the size of incoming activations to the second mobile inverted residual block in our search space can vary from [ 64 , 72 , 64 , 64 ] to [ 64 , 144 , 112 , 112 ] . We now consider the case where elements within x are i.i.d random variables with cumulative distribution F , and let MN = max { x1 , · · · , xN } denote the maximum . Then MN follows the maximum extreme value distribution with the following cumulative distribution function : P ( MN ≤ t ) = P ( x1 ≤ t , · · · , xN ≤ t ) = P ( x1 ≤ t ) · · ·P ( xN ≤ t ) = F ( t ) N ( 3 ) Following the assumptions in [ 38 ] , we model activations within neural networks as tensor-valued Laplace random variables . WLOG , when considering the maximum value distribution , we can model each element as an exponential random variable , because the Laplace distribution can be thought of as two exponential distributions spliced together back-to-back . With xi ∼ Exp ( λ ) , i ∈ [ 1 , · · · , N ] , E [ MN ] = ∫ ∞ 0 ( 1− F ( t ) N ) dt = ∫ ∞ 0 ( 1− ( 1− e−λt ) N ) dt , ( let u = 1− e−λt ) = ∫ 1 0 1 λ 1− uN 1− u du = 1 λ N∑ k=1 1 k ∼ 1 λ log ( N ) ( 4 ) 5/19/2021 ExtremeDistStudy - Jupyter Notebook localhost:11230/notebooks/ExtremeDistStudy.ipynb 5/5 In [ 11 ] : In [ 12 ] : In [ 13 ] : In [ 14 ] : In [ ] : bgcolor='rgba ( 0,0,0,0 ) ' , ) , # legend_title= '' Activation Size '' , # legend_orientation= ' h ' , font=dict ( family= '' Courier New , monospace '' , size=16 , color= '' RebeccaPurple '' ) , ) fig.show ( ) # np.save ( 'amax_min_arch ' , amax ) # oamax = np.load ( 'amax.npy ' , allow_pickle=True ) # oamax = oamax.item ( ) # ks = sorted ( oamax.keys ( ) , key=lambda x : np.ptp ( oamax [ x ] ) , reverse=True ) # for k in ks : # print ( k ) # # plt.plot ( oamax [ k ] ) # plt.plot ( bmax [ k ] ) # plt.show ( ) We can see that E [ MN ] slowly diverges to ∞ as N increases . For a symmetric activation distribution where the minimum and maximum value distributions are symmetric about 0 , we can consider E [ ∆̂ ] = 2E [ MN ] ∝ log ( N ) for the EMA estimator . As a result , the EMA estimator ∆̂ will never become stable because it is sensitive to the activation tensor size and could change significantly for each update depending on the sampled subnet . For the example mobile inverted residual block given above , the largest difference in E [ MN ] is proportional to log ( 1511424 ∗ batch_size ) , which can lead to substantial instability in the EMA estima- tion . Furthermore , due to adaptive layer skipping , layers that are only present in deep subnets will activate less frequently and have substantially more lag in EMA update than layers shared among shallow and deep subnets , exacerbating the instability in EMA and the quantized activations . Figure 1 shows the maximum extreme value distribution for the respective ctivations collected from 1000 batches of training data . Note such a shifting extreme value distribution problem can also happen to weight tensors in weightsharing supernets because different subnets can have different weight tensor sizes . Conveniently , our supernet construction ensures that at each layer , the weight tensor for smaller subnets is nested within the weight tensor for the larger subnet . Therefore we simply perform quantization on the entire weight tensor before indexing out the weights for the given subnet to stabilize weight quantization . Learnable scale does not help . Methods that learn a fixed ∆ is also suboptimal in the case of supernet training . According to [ 38 ] the optimal quantization level depends on the variance in activation distribution ( e.g . Laplace parameter b for Laplace distribution or the variance σ2 of normal distribution ) . However , as observed in [ 29 ] , quantizing weights and input to varying precision leads to activations with different mean and variance . While we can easily resolve this issue in fixed-precision quantization by learning a different set of ∆ for each bitwidth setting , it is impossible to assign a set of ∆ parameters for each mixed-precision quantization policy on each subnet . Empirically , during the training of our compact MobileNet search space , EMA estimator ∆̂ leads to gradient explosion , and a fixed set of learnable ∆ leads to diverging training loss as well . 5

Summary: The paper presents the use of batch statics to help stabilize the NAS+Quantization process. It is based on the one-shot NAS idea, but applies the batch idea to assist the quantization policy search during the Neural Architecture Search. The results are compared to some alternative methods and show a faster convergence.