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: This paper summarizes a technique to allow quantizing a weight-sharing superset to with flexible quantization policies to allow choosing different network architecture and quantization policies for desired efficiency-accuracy trade-off. The key observation of this method is that quantizers needs to estimate the maximum and minimum values of a activations. However, the extreme values of activations are dependent on the activation tensor's size. During training, the activation size changes as we sample different architectures, so the extreme value do not converge when using EMA estimator, and the quantizer become unstable during training.  To fix this, this paper proposes a simple method to estimate the extreme values (therefore quantizer parameters) based on current batch statistics. This leads to better stability and performance.