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

Abstract:
As the applications of deep learning models on edge devices increase at an accelerating pace, fast adaptation to various scenarios with varying resource constraints has become a crucial aspect of model deployment. As a result, model optimization strategies with adaptive configuration are becoming increasingly popular. While single-shot quantized neural architecture search enjoys flexibility in both model architecture and quantization policy, the combined search space comes with many challenges, including instability when training the weight-sharing supernet and difficulty in navigating the exponentially growing search space. Existing methods tend to either limit the architecture search space to a small set of options or limit the quantization policy search space to fixed precision policies. To this end, we propose BatchQuant, a robust quantizer formulation that allows fast and stable training of a compact, single-shot, mixed-precision, weight-sharing supernet. We employ BatchQuant to train a compact supernet (offering over 10 quantized subnets) within substantially fewer GPU hours than previous methods. Our approach, Quantized-for-all (QFA), is the first to seamlessly extend one-shot weight-sharing NAS supernet to support subnets with arbitrary ultra-low bitwidth mixed-precision quantization policies without retraining. QFA opens up new possibilities in joint hardware-aware neural architecture search and quantization. We demonstrate the effectiveness of our method on ImageNet and achieve SOTA Top-1 accuracy under a low complexity constraint (< 20 MFLOPs). The code and models will be made publicly available at https://github.com/bhpfelix/QFA.

heading: 1 Introduction
text: 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.

heading: 2 Related Work

heading: 2.1 Mixed Precision Quantization
text: 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.

heading: 2.2 Joint Mixed-Precision Quantization and Neural Architecture Search
text: 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.

heading: 3 Stabilizing Mixed Precision Supernet Training with BatchQuant
text: 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.

heading: 3.1 Quantization Preliminaries
text: 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.

heading: 3.2 Weight Quantization with LSQ
text: 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.

heading: 3.3 Challenges of Activation Quantization in One-shot Supernet
text: 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

heading: 3.4 Batch Quantization for Robust Activation Quantization
text: To alleviate the unstable scale estimation issue, we propose the following general definition of batch quantizer (BatchQuant)
∆̂ = x̂max − x̂min
p− n , ẑ = clamp
( − ⌊ x̂min
∆̂
⌉ , n, p ) xq = ⌊ clamp ( x
∆̂γ + ẑ + β, n, p )⌉ x̂ = (xq − ẑ − β)∆̂γ.
(5)
Analogous to batch normalization [24], we leverage batch statistics to help standardize the activation distribution across different sampled subnets. Instead of using EMA, we estimate the extreme values x̂min and x̂max only from the current batch to negate the effect of changing extreme value distribution due to subnet sampling. In addition, we learn a multiplicative residual γ on ∆̂ and an additive residual β on ẑ to facilitate learning of the optimal clipping range.
Note, similar to LSQ+ [4], zero in activation may not be exactly quantizable based on such a formulation. However, we share emprical observation as LSQ+, that the learned β is often small and aids in empirical performance by reducing the quantization error for H-swish activation [33]. Because we use symmetric quantization for weights, asymmetric quantization of activations has no additional cost during inference as compared to symmetric quantization since the bias term can be precomputed and fused with the bias in the succeeding layer:
ŵx̂ = (wq∆w)(xq − ẑ − β)∆̂xγ = wqxq∆w∆̂xγ − bias︷ ︸︸ ︷ wq∆w(ẑ + β)∆̂xγ . (6)
BatchQuant (BQ) is generally applicable. Note our formulation of BatchQuant is general. We do not explicitly define the batch extreme value estimators x̂min and x̂max, because it is straightforward to plug in different definitions as the usage sees fit. We will provide three simple example formulations of x̂min and x̂max below. BatchQuant could potentially be a good drop-in replacement for conventional QAT with EMA scale estimator when adjusting training batch size or performing tasks that have varying activation size such as semantic segmentation. Due to the adaptive formulation, BatchQuant offers a much simpler initialization strategy than the optimization-based strategy in LSQ+ [4]. With a proper choice of the batch extreme value estimator (x̂min, x̂max), we can simply initialize the residual terms with γ = 1 and β = 0. As training is stabilized from the beginning, BQ can also avoid the delayed activation quantization strategy introduced in [34], which requires choosing a starting schedule for activation quantization.
Extreme Value Estimators Due to presence of learnable residuals γ and β, the extreme value estimators (x̂min, x̂max) does not need to capture the exact value of the scale as long as the calculated scale ∆̂ captures the range variation in activation. Consider 4D activations with vector index i = (iB , iC , iH , iW ) that indexes batch, channel, height, and width dimension respectively, we present three example definitions to use in our experiment:
x̂min = min i xi, x̂max = max i xi (7)
x̂min = 1
C ∑ iC min iB ,iH ,iW xi, x̂max = 1 C ∑ iC max iB ,iH ,iW xi (8)
x̂min = µ− 3σ, x̂max = µ+ 3σ (9)
where µ is the tensor-wise mean and σ is the tensor-wise variance. Note that Equation 7 provides unbiased estimation of extreme values and is analogous to EMA ∆ estimator but only based on batch statistics. Equation 9 is a biased estimation with the least variance and is equivalent to the scale parameter initialization strategy of LSQ+ [4]. Plugging Equation 9 into Equation 5 is thus analogous of using a learned scale quantizer. Equation 8 is a biased estimation with a bias-variance trade-off between Equation 7 and 9. We show empirically that Equation 8 leads to stable supernet training.
BatchQuant Calibration during test time Similar to the treatment of batchnorm in [22], to obtain a set of extreme value estimations (x̂min, x̂max) specific to a quantized subnet for inference, we simply perform BatchQuant calibration by forwarding a few batches of data and accumulate a running mean of the estimations.

heading: 4 Mixed-Precision Quantized Architecture Search
text: Equipped with BatchQuant for stable mixed-precision supernet training, we now introduce the search procedure of our mixed-precision quantized architecture search method, QFA.
Supernet Search Space Definition Following the convention [3, 2, 19, 12, 23] of only searching and quantizing the residual blocks in the Mobilenet V3 search space, we keep the input and weights of the first convolutional layer at full precision and quantize the incoming activation and weights of the last fully-connected layer to 8-bit precision. Specifically, our Mobilenet V3 search space contains 5 stages each with 4 mobile inverted residual blocks. Each stage can optionally skip 1 or 2 of its last blocks. Each block has kernel size options {3, 5, 7} and expansion ratio choices {3, 4, 6}. Within each block, there are 3 convolution layers. Each convolution layer can independently choose an activation bitwidth and a weight bitwidth. Thus, each block has (3× 3× 33×2) = 6561 configurations, each stage has 65612 × (1 + 6561 × (1 + 6561)) ≈ 1.85 × 1015 configurations. Without considering the elastic input resolution choices, our complete search space contains over (1.85 × 1015)5 ≈ 2.19 × 1076 different quantized architectures.
Supernet Training and Elastic Quantization. Following the common practice of starting QAT from a trained full precision network [1], we start by pretraining the supernet without quantizers on ImageNet [39]. We follow the training protocol introduced in [22] to train the supernet. Then, we perform QAT starting from the full precision supernet. Unlike [22, 23] that adopts a progressive shrinking strategy which gradually opens up smaller subnet choices as training progresses, we directly allow all possible quantized subnet options. Our elastic quantization training consists of two stages. In the first stage, we train the supernet for 65 epochs with bitwidth choices {2, 3, 4, 32}, where a bitwidth of 32 means floating-point precision. Empirically, we found that mixing in the well-trained floating-point layers helps the low precision layers learn and leads to more stable training. Then, we proceed to the second
stage and train the supernet for 125 epochs with only low bitwidth choices {2, 3, 4}. Table 3 shows that our two-staged training strategy outperforms training a supernet with bitwidth choices {2, 3, 4} for the same number of epochs. Our elastic quantization training requires only 190 epochs of training. In comparison, the staged training of fixed precision supernet at each bitwidth in [23] takes 495 epochs in total.
Subnet Sampling. Similar to [22, 23, 16], for each supernet training update, we sample and accumulate gradients from multiple subnets. Specifically, we adopt the sandwich rule proposed in [16] by sampling 4 quantized subnets accompanied by the smallest architecture and the largest architecture within the supernet with the entire architecture set to a random bitwidth.
Multi-objective Evolutionary Search To balance multiple competing objectives, we leverage multiobjective evolutionary search to produce desirable subnets. Unlike the aging evolution in [22] that only produce one output architecture under a given set of constraints at a time, we adopt the NSGA-II algorithm [40], which outputs a Pareto population at once. As a result, we can quickly access the Pareto front of a trained supernet. At each iteration, the NSGA-II algorithm selects the best-fit population with non-dominated sorting. Then, a crowding distance is calculated to ensure the selected individuals cover the Pareto front evenly without concentrating at a single place. To speed up the search process, we follow [22] and train an accuracy predictor that predicts the accuracy of a given quantized architecture configuration. We use one-hot encoding to encode quantized architecture configurations into a binary vector.

heading: 5 Experimental Analysis and Results
text: Experiment Settings and Implementation Details We base our codebase on the open-source implementation of [22] under the MIT License, and we follow the exact training procedure on
5/26/2021 Visualization - Jupyter Notebook
localhost:8888/notebooks/Downloads/NAC_ARTIFACTS/Visualization.ipynb 4/4
In [ ]:
), row=1, col=2) fig.add_trace(go.Scatter(x=f_flops, y=f_top1, mode='lines+text', name='Ours', text=['Ours'], textposition='middle right', marker=dict(color=defaults[0]), textfont=dict(color=defaults[0]), line=dict(color=defaults[0]), showlegend=False, ), row=1, col=2) fig.add_trace(go.Scatter(x=f_flops, y=f_top1, mode='lines', name='Ours', marker=dict(color=defaults[0]), line=dict(color=defaults[0]), ), row=1, col=2) fig.update_layout(showlegend=False) fig.update_xaxes(title_text="MFLOPs", range=[0, 350], row=1, col=1) fig.update_xaxes(title_text="MFLOPs(Log Scale)", type="log", tickformat=".1r", row=1, col=2) fig.update_layout( height=380, width=900, margin=dict(l=0, b=0), yaxis_title="Top 1 Accuracy(%)", legend_title="Model Family", font=dict( family="Courier New, monospace", size=16, color="RebeccaPurple" ), ) fig.show()
ImageNet [39] to obtain the full precision supernet. For both stages of the elastic quantization procedure, we follow the common hyperparameter choice of [22] and use an initial learning rate of 0.08. For all experiments, we clip the global norm of the gradient at 500. We train with a batch size of 2048 across 32 V100 GPUs on our internal cluster. Unless otherwise mentioned, we keep all other settings the same as [22]. After training is complete, we randomly sample 16k quantized subnets and evaluate on 10k validation images sampled from the training set to train the accuracy predictor. During the evolutionary search, we keep a population size of 500 for 1000 generations. For each generation, once we identify the Pareto population based on nondominated sorting and crowding distance, we breed new genotypes through crossover and mutation with a crossover probability of 0.007 and a mutation probability of 0.02. We determine both crossover and mutation probability through grid search. To demonstrate the effectiveness of QFA, we conduct experiments that cover the accuracy and complexity trade-off of searched subnets. For a fair comparison with previous methods, we adopt FLOPs for full precision layers and BitOPs for quantized layers as our complexity measure. While there is no direct conversion from BitOPs to FLOPs, we follow the convention in [21, 23], where given a full precision layer with FLOPs a, its quantized counterpart with m bit weight and n bit activation will have a BitOPs of mn× a and a FLOPs of (mn× a)/64. To be concise, we report our complexity as FLOPs of FP layers + BitOPs of quantized layers / 64 = Total FLOPs.
Stable Mixed-precision Supernet Training with BQ To test the effectiveness of BQ in stabilizing supernet training, we attempt to train a baseline supernet with LSQ quantized weights and LSQ+ quantized activations. To test the robustness of low bitwidth training, we only allow bitwidth choices {2, 3, 4}. We follow the practice of LSQ+ and initialize the scale and offset parameter layer-bylayer through an optimization procedure that minimizes the MSE between quantized output and full precision output at the given layer. However, the training loss diverged at the start of the second learning rate warm-up epoch. Replacing LSQ+ activation quantizer with BQ, we obviate the optimization-based initialization procedure, and the divergent training loss problem no longer occurs.
Effect of Residual Terms γ and β To test the effect of γ and β in BQ, we removed γ and β from all BQ operations when training the supernet. The training loss diverged within a few update steps.
Comparing Extreme Value Estimator Choices To investigate the stability of different extreme value estimators, we train BQ supernet with bitwidth choices {2, 3, 4} and plug-in Equations 7, 8, 9 as extreme value estimator. Training loss for both Equations 7 and 9 diverged after the first few updates, and only Equation 8 led to stable training.
Comparison with SOTA Mixed-Precision Quantized Architecture Search We compare with existing quantization-aware NAS methods including APQ [20] and SPOS [12], as shown in the left subplot of Figure 2. Our results outperform both APQ and SPOS by a large margin. Thus, in terms of existing mixed-precision quantized architecture search methods, we are able to achieve the best accuracy v.s. FLOPs trade-off.
Comparison with SOTA fixed-precision quantized models We further compare with several strong fixed-precision quantization methods including OQA [23], BATS [21], BMobi [42], LSQ [36], LSQ+ [4], APOT [43], and QKD[44], as shown in Figure 2. Our results demonstrate the competitive performance of quantized subnets discovered by QFA against quantized models produced by state-ofthe-art quantization and quantization-aware NAS methods. For models > 20 MFLOPs, we are able to achieve comparable performance as OQA with much less training cost—the OQA supernets require a total of 495 epochs to train, while our mixed-precision supernet took only 190 epochs to train. We achieve SOTA ImageNet top-1 accuracy on models under 20 MFLOPs. Our advantage over OQA at < 20 MFLOPs could be mainly due to the flexibility of our search space to smoothly interpolate between the Pareto frontiers of fixed-precision architectures by mixing layers with varying precision.
Design Cost Comparison with Quantized Architecture Search Methods As shown in Table 4, our approach demonstrates better efficiency in terms of GPU hours and marginal carbon emission. We report the total GPU hours for our approach, including 1200 GPU hours of full-precision supernet training, 565 GPU hours of mixed-precision supernet training, and 40 hours of accuracy predictor data collection and training. Running NSGA-II with the trained accuracy predictor does not make use of GPU and finishes within 1 hour. For a fair comparison, we added the GPU hours to train the full precision supernet to the quantized architecture search time reported by the authors of OQA.

heading: 6 Discussion
text: While we demonstrate the superior performance of BatchQuant in combination with mixed-precision weight-sharing supernet, there remain many other directions for future investigation. QFA shares some limitations common to many NAS methods with weight-sharing supernets. Since we only visit a small fraction of subnets during training, some subnets may receive insufficient training. Employing techniques [45, 46] targeting this limitation may prove helpful for our especially complex search space. Finally, a thorough theoretical explanation of the effectiveness of BatchQuant on the training dynamics of weight-sharing mixed-precision supernet is still an open problem, and we leave the investigation for future exploration.

heading: 7 Conclusion
text: In this paper, we present Quantized-for-All (QFA), a novel mixed-precision quantized architecture search method that can jointly search for architecture and mixed-precision quantization policy combination and deploy quantized subnets at SOTA efficiency without retraining or finetuning. We provide analysis on the challenge of activation quantization in mixed-precision supernet. To allow stable training of QFA supernet, we proposed BatchQuant (BQ), a general, plug-and-play quantizer formulation that stabilizes the mixed-precision supernet training substantially. We demonstrate the robust performance of QFA in combination with BQ under ultra-low bitwidth settings (2/3/4). By leveraging NSGA-II, we produce a family of Pareto architectures at once. Our discovered model family achieves competitive accuracy and computational complexity trade-off in comparison to existing state-of-the-art quantization methods.

Review summaries:
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. 

Authors propose a robust quantizer formulation BatchQuant, which utilizes the batch statistics to provide a more stable quantizer for different quantization levels. With benefit of BatchQuant, they train a compact, single-shot, mix-precision, weight-sharing supernet for NAS. They evaluate BatchQuant on ImageNet.

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.