HAQ: Hardware-Aware Automated Quantization with Mixed Precision

Introduction

In many real-time machine learning applications (such as robotics, autonomous driving, and mobile VR/AR), deep neural networks is strictly constrained by the latency, energy, and model size. In order to improve the hardware efficiency, many researchers have proposed to quantize the weights and activations to low precision .

Conventional quantization methods use the same number of bits for all layers , but as different layers have different redundancy and behave differently on the hardware (computation bounded or memory bounded), it is necessary to use mixed precision for different layers (as shown in Figure 1). This flexibility was originally not supported by chip vendors until recently the hardware manufacturers started to implement this feature: Apple released the A12 Bionic chip that supports mixed precision for the neural network inference ; NVIDIA recently introduced the Turing GPU architecture that supports 1-bit, 4-bit, 8-bit and 16-bit arithmetic operations ; Imagination launched a flexible neural network IP that supports per-layer bitwidth adjustment for both weights and activations . Besides industry, recently academia also works on the bit-level flexible hardware design: BISMO proposed the bit-serial multiplier to support multiplications of 1 to 8 bits; BitFusion supports multiplications of 2, 4, 8 and 16 bits in a spatial manner.

However, a very missing part is how to determine the bitwidth of both weights and activations for each layer on different hardware accelerators. This is a vast design space: with $M$ different neural network models, each with $N$ layers, on $H$ different hardware platforms, there are in total $\mathcal{O}(H\times M\times 8^{2N})$Assuming the bitwidth is 1 to 8 for both weights and activations. possible solutions. For a widely used ResNet-50 model, the size of the search space is about $8^{100}$, which is even larger than the number of particles in the universe. Conventional methods require domain experts (with knowledge of both machine learning and hardware architecture) to explore the huge design space smartly with rule-based heuristics, such as: we should retain more bits in the first layer which extracts low level features and in the last layer which computes the final outputs; also, we should use more bits in the convolution layers than in the fully-connected layers because empirically, the convolution layers are more sensitive. As the neural network becomes deeper, the search space increases exponentially, which makes it infeasible to rely on hand-crafted strategies. Therefore, these rule-based quantization policies are usually sub-optimal, and they cannot generalize from one model to another. In this paper, we would like to automate this exploration process by a learning-based framework.

Another challenge is how to optimize the latency and the energy consumption of a given model on the hardware. A widely adopted approach is to rely on some proxy signals (e.g., FLOPs, number of memory references) . However, as different hardware behaves very differently, the performance of a model on the hardware cannot always be accurately reflected by these proxy signals. Therefore, it is important to directly involve the hardware architecture’s performance feedback into the design loop. Also, as demonstrated in Table 1, the quantization solution optimized on one hardware might not be optimal on the other, which raises the demand for specialized policies for different hardware architectures.

To this end, we propose the Hardware-Aware Automated Quantization (HAQ) framework that leverages reinforcement learning to automatically predict the quantization policy given the hardware’s feedback. The RL agent decides the bitwidth of a given neural network in a layer-wise manner. For each layer, the agent receives the layer configuration and statistics as observation, and it then outputs the action which is the bitwidth of weights and activations. We then leverage the hardware accelerator as the environment to obtain the direct feedback from hardware to guide the RL agent to satisfy the resource constraints. After all layers are quantized, we finetune the quantized model for one more epoch, and feed the validation accuracy after short-term retraining as the reward signal to our RL agent. During the exploration, we leverage the deep deterministic policy gradient (DDPG) to supervise our RL agent. We also studied the quantization policy on multiple hardware architectures: both cloud and edge neural network accelerators, with spatial or temporal multi-precision design.

The contribution of this paper has four aspects:

Automation: We propose an automated framework for quantization, which does not require domain experts and rule-based heuristics. It frees the human labor from exploring the vast search space of choosing bitwidths.

Hardware-Aware: Our framework involves the hardware architecture into the loop so that it can directly reduce the latency, energy and storage on the target hardware instead of relying on proxy signals.

Specialization: For different hardware architectures, our framework can offer a specialized quantization policy that’s exactly tailored for the target hardware architecture to optimize latency and energy.

Design Insights: We interpreted the different quantization polices learned for different hardware architectures. Taking both computation and memory access into account, the interpretation offers insights on both neural network architecture and hardware architecture design.

Related Work

There have been extensive explorations on compressing and accelerating deep neural networks using quantization. Han et al. quantized the network weights to reduce the model size by rule-based strategies: e.g., they used human heuristics to determine the bitwidths for convolution and fully-connected layers. Courbariaux et al. binarized the network weights into $\{-1,+1\}$; Rastegari et al. and Zhou et al. binarized each convolution filter into $\{-w,+w\}$; Zhu et al. mapped the network weights into $\{-w_{\text{N}},0,+w_{\text{P}}\}$ using two bits; Zhou et al. used one bit for network weights and two bits for activations; Jacob et al. made use of 8-bit integers for both weights and activations. We refer the reader to the survey paper by Krishnamoorthi et al. for a more detailed overview. These conventional quantization methods either simply assign the same number of bits to all layers or require domain experts to determine the bitwidths for different layers, while our framework automates this design process, and our learning-based policy outperforms rule-based strategies.

Many researchers aimed to improve the performance of deep neural networks by searching the network architectures: Zoph et al. proposed the Neural Architecture Search (NAS) to explore and design the transformable network building blocks, and their network architecture outperforms several human designed networks; Liu et al. introduced the Progressive NAS to accelerate the architecture search by 5$\times$ using sequential model-based optimization; Pham et al. introduced the Efficient NAS to speed up the exploration by 1000$\times$ using parameter sharing; Cai et al. introduced the path-level network transformation to effectively search the tree-structured architecture space. Motivated by these AutoML frameworks, He et al. leveraged the reinforcement learning to automatically prune the convolution channels. Our framework further explores the automated quantization for network weights and activations, and it takes the hardware architectures into consideration.

To facilitate the efficient deployment, researchers designed hardware-friendly approaches to slim neural network models. For instance, the coarse-grained channel pruning methods prune away the entire channel of convolution kernels to achieve speedup. Recently, researchers have explicitly optimized for various aspects of hardware properties, including the inference latency and energy: Yang et al. proposed the energy-aware pruning to directly optimize the energy consumption of neural networks; Yang et al. reduced the inference time of neural networks on the mobile devices through a lookup table. Nevertheless, these methods are still rule-based and mostly focus on pruning. Our framework automates the quantization process by taking hardware-specific metric as direct rewards using a learning based method.

Approach

We model the quantization task as a reinforcement learning problem (Figure 2). We use the actor-critic model with DDPG agent to give the action: bits for each layer. We collect hardware counters as constraints, together with accuracy as rewards to search the optimal quantization policy. We have three hardware environments that covers edge and cloud, spatial and temporal architectures for mixed-precision accelerator. Below describes the details of the RL formulation.

Our agent processes the neural network in a layer-wise manner. For each layer, our agent takes two steps: one for weights, and one for activations. In this paper, we introduce a ten-dimensional feature vector ${O}_{k}$ as our observation:

If the $k$th layer is a convolution layer, the state ${O}_{k}$ is

where $k$ is the layer index, $c_{\text{in}}$ is #input channels, $c_{\text{out}}$ is #output channels, $s_{\text{kernel}}$ is kernel size, $s_{\text{stride}}$ is the stride, $s_{\text{feat}}$ is the input feature map size, $n_{\text{params}}$ is #parameters, $i_{\text{dw}}$ is a binary indicator for depthwise convolution, $i_{\text{w/a}}$ is a binary indicator for weight/activation, and $a_{k-1}$ is the action from the last time step.

If the $k$th layer is a fully-connected layer, the state ${O}_{k}$ is

where $k$ is the layer index, $h_{\text{in}}$ is #input hidden units, $h_{\text{out}}$ is #output hidden units, $s_{\text{feat}}$ is the size of input feature vector, $n_{\text{params}}$ is #parameters, $i_{\text{w/a}}$ is a binary indicator for weight/ activation, and $a_{k-1}$ is the action from the last step.

For each dimension in the observation vector ${O}_{k}$, we normalize it into $$ to make them in the same scale.

2 Action Space

We use a continuous action space to determine the bitwidth. The reason that we do not use a discrete action space is because it loses the relative order: e.g., 2-bit quantization is more aggressive than 4-bit and even more than 8-bit. At the $k$th time step, we take the continuous action $a_{k}$ (which is in the range of $$), and round it into the discrete bitwidth value$b_{k}$:

where $b_{\text{min}}$ and $b_{\text{max}}$ denote the min and max bitwidth (in our experiments, we set $b_{\text{min}}$ to $2$ and $b_{\text{max}}$ to $8$).

In real-world applications, we have limited computation budgets (i.e., latency, energy, and model size). We would like to find the quantization policy with the best performance given the constraint.

We encourage our agent to meet the computation budget by limiting the action space. After our RL agent gives actions $\{a_{k}\}$ to all layers, we measure the amount of resources that will be used by the quantized model. The feedback is directly obtained from the hardware accelerator, which we will discuss in Section 3.3. If the current policy exceeds our resource budget (on latency, energy or model size), we will sequentially decrease the bitwidth of each layer until the constraint is finally satisfied.

3 Direct Feedback from Hardware Accelerators

An intuitive feedback to our RL agent can be FLOPs or the model size. However, as these proxy signals are indirect, they are not equal to the performance (i.e., latency, energy consumption) on the hardware. Cache locality, number of kernel calls, memory bandwidth all matters. Proxy feedback can not model these hardware functionality to find the specialized strategies (see Table 1).

Instead, we use direct latency and energy feedback from the hardware accelerator as resource constraints, which enables our RL agent to determine the bitwidth allocation policy from the subtle differences between different layers: e.g., vanilla convolution has more data reuse and better locality, while depthwise convolution has less reuse and worse locality, which makes it memory bounded. Such difference impacts the optimal quantization policy.

4 Quantization

We linearly quantize the weights and activations of each layer using the action $a_{k}$ given by our agent, as linearly quantized model only needs fixed point arithmetic unit which is more efficient to implement on the hardware.

Specifically, for each weight value $w$ in the $k$th layer, we first truncate it into the range of $[-c,c]$, and we then quantize it linearly into $a_{k}$ bits:

where $\text{clamp}(\cdot,x)$ is to truncate the values into $[-x,x]$, and the scaling factor $s$ is defined as $s=c/(2^{a_{k}-1}-1)$. In this paper, we choose the value of $c$ by finding the optimal value $x$ that minimizes the KL-divergence between the original weight distribution ${W}_{k}$ and the quantized weight distribution $\text{quantize}({W}_{k},a_{k},x)$:

where $\mathcal{D}_{\text{KL}}(\cdot\mid\mid\cdot)$ is the KL-divergence that characterizes the distance between two distributions. As for activations, we quantize the values similarly except that we truncate them into the range of $[0,c]$, not $[-c,c]$ since the activation values (which are the outputs of the ReLU layers) are non-negative.

5 Reward Function

After quantization, we retrain the quantized model for one more epoch to recover the performance. As we have already imposed the resource constraints (latency, energy) by limiting the action space (Section 3.2), we define our reward function $\mathcal{R}$ to be only related to the accuracy:

where $\text{acc}_{\text{origin}}$ is the top-1 classification accuracy of the full-precision model on the training set, $\text{acc}_{\text{quant}}$ is the accuracy of the quantized model after finetuning, and $\lambda$ is a scaling factor which is set to $0.1$ in our experiments.

6 Agent

For the RL agent, we leverage the deep deterministic policy gradient (DDPG) , which is an off-policy actor-critic algorithm for continuous control problem. In our environment, one step means that our agent makes an action to decide the number of bits assigned to the weights or activations of a specific layer, while one episode is composed of multiple steps, where our RL agent makes actions to all layers. We apply a variant form of the Bellman’s Equation, where each transition in an episode is defined as ${T}_{k}=({O}_{k},a_{k},\mathcal{R},{O}_{k+1})$. During exploration, the ${Q}$-function is computed as

and the loss function can be approximated by

where $N_{\text{s}}$ denotes the number of steps in this episode, and the baseline $\mathcal{B}$ is defined as an exponential moving average of all previous rewards in order to reduce the variance of the gradient estimation. The discount factor $\gamma$ is set to 1 since we assume that the action made for each layer should contribute equally to the final result. Moreover, as the number of steps is always finite (bounded by the number of layers), the sum of the rewards will not explode.

7 Implementation Details

In this section, we present the implementation details about RL exploration and finetuning quantized models.

The DDPG agent consists of an actor network and a critic network. Both using the same network architecture: they take the state vector and the action from the last time step as inputs and feed them into two separate fully-connected layers with hidden sizes of $400$. After that, we add the two hidden vectors together and go through another two fully-connected layers with hidden sizes of $\{300,1\}$. As for the actor network, we use an additional sigmoid function to project the output into the range of $$.

Optimization of the DDPG agent is carried out using ADAM with $\beta_{1}=0.9$ and $\beta_{2}=0.999$. We use a fixed learning rate of $10^{-4}$ for the actor network and $10^{-3}$ for the critic network. During exploration, we employ the following stochastic process of the noise:

where $\mathcal{N}_{\text{trunc}}(\mu,\sigma,a,b)$ is the truncated normal distribution, and $w$ is the model weights. The noise $\sigma$ is initialized as $0.5$, and after each episode, the noise is decayed exponentially with a decay rate of $0.99$.

During exploration, we finetune the quantized model for one epoch to help recover the performance (using SGD with a fixed learning rate of $10^{-3}$ and momentum of $0.9$). We randomly select 100 categories from ImageNet to accelerate the model finetuning during exploration. After exploration, we quantize the model with our best policy and finetune it on the full dataset.

Experiments

We conduct extensive experiments to demonstrate the consistent effectiveness of our framework for multiple objectives: latency, energy, and model size.

Our experiments are performed on the ImageNet dataset. As our focus is on more efficient models, we extensively study the quantization of MobileNet-V1 and MobileNet-V2 . Both MobileNets are inspired from the depthwise separable convolutions and replace the regular convolutions with the pointwise and depthwise convolutions: MobileNet-V1 stacks multiple “depthwise – pointwise” blocks repeatedly; while MobileNet-V2 uses the “pointwise – depthwise – pointwise” blocks as its basic building primitives.

1 Latency-Constrained Quantization

We first evaluate our framework under latency constraints on two representative hardware architectures: spatial and temporal architectures for multi-precision CNN. We show that it’s beneficial to have specialized quantization policies for different hardware architectures. We systematically interpret the policy given by AI to guide future human designs.

Bit-Serial Matrix Multiplication Overlay (BISMO) proposed by Yaman et al. is a classic temporal design of neural network accelerator on FPGA. It introduces bit-serial multipliers which are fed with one-bit digits from 256 weights and corresponding activations in parallel at one time and accumulates their partial products by shifting over time.

BitFusion architecture proposed by Hardik et al. is a state-of-the-art spatial ASIC design for neural network accelerator. It employs a 2D systolic array of Fusion Units which spatially sum the shifted partial products of two-bit elements from weights and activations.

1.1 Quantization policy for BISMO Architecture

Inferencing neural networks on edge devices and cloud severs can be quite different: batch size, memory bandwidth, peak FLOPs, etc. We use Xilinx Zynq-7020 FPGA as our edge device and Xilinx VU9P as our cloud device. Table 2 shows our experiment configurations on these two platforms along with their available resources.

As for comparison, we adopt the PACT as our baseline, which uses the same number of bits for all layers except for the first layer which extracts the low level features, they use 8 bits for both weights and activations as it has fewer parameters and is very sensitive to errors. We follow a similar setup for the first layer (8 bits), and explore the bitwidth allocation policy for all the other layers. Under the same latency, HAQ consistently achieved better accuracy than the baseline on both the cloud and the edge (Table 3). With similar accuracy, HAQ can reduce the latency by 1.4$\times$ to 1.95$\times$ compared with the baseline.

Our agent gave quite different quantization policy for edge and cloud accelerators (Figure 3). For the activations, the depthwise convolution layers are assigned less bitwidth than the pointwise layers on the edge; while on the cloud device, the bitwidth of these two types of layers are similar. For weights, the bitwidth of these types of layers are nearly the same on the edge; while on the cloud, the depthwise convolution layers got more bitwidth than the pointwise convolution layers.

We explain the difference of quantization policy between edge and cloud by the roofline model . Many previous works use FLOPs or BitOPs as metrics to measure computation complexity. However, they are not able to directly reflect the latency, since there are many other factors influencing the hardware performance, such as memory access cost and degree of parallelism . Taking computation and memory access into account, the roofline model assumes that applications are either computation-bound or memory bandwidth-bound, if not fitting in on-chip caches, depending on their operation intensity. Operation intensity is measured as operations (MACs in neural networks) per byte accessed. A lower operation intensity indicates suffering more from the memory access.

The bottom of Figure 3 shows the operation intensities (OPs per Byte) of convolution layers in the MobileNet-V1. Depthwise convolution is memory bounded, and the pointwise convolution is computation bounded. Our experiments show that when running MobileNet-V1 on the edge devices with small batch size, its latency is dominated by the depthwise convolution layers. Since the feature maps take a major proportion in the memory of depthwise convolution layers, our agent gives the activations less bits. In contrast, when running MobileNet-V1 on the cloud with large batch size, our agent increases the bitwidth of depthwise convolution to preserve the accuracy at low memory overhead since depthwise convolution only takes a small proportion of the total weights. A similar phenomenon can be observed in Figure 4 on MobileNet-V2. Moreover, as the activation size in deeper layers gets smaller, they get assigned more bits.

1.2 Quantization policy for BitFusion Architecture

In order to demonstrate the effectiveness of our framework on different hardware architectures, we further compare our framework with PACT under the latency constraints on the BitFusion architecture (Table 4). Our framework performs much better than the hand-craft policy with the same latency. It can achieve almost no degradation of accuracy with only half of the latency used by the original MobileNet-V1 model (from 20.08 to 11.09 ms). Therefore, our framework is flexible to provide specialized quantization policy for different hardware platforms.

2 Energy-Constrained Quantization

We then evaluate our framework under the energy constraints. Similar to the latency-constrained experiments, we compare our framework with PACT that uses fixed number of bits without hardware feedback. From Table 5, we can clearly see that our framework outperforms the rule-based baseline: it achieves much better performance while consuming similar amount of energy. In particular, our framework is able to achieve almost no loss of accuracy with nearly half of the energy consumption of the original MobileNet-V1 model (from 31.03 to 16.57 mJ), which suggests that mixed precision with hardware-aware, specialized quantization policy can indeed help reduce the energy consumption.

3 Model Size-Constrained Quantization

Finally, we evaluate our framework under the model size constraints. Following Han et al. , we employ the $k$-means algorithm to quantize the values into $k$ different centroids instead of using the linear quantization for compression, since $k$-means quantization can be more effective reducing the model size.

We compare our framework with Deep Compression on MobileNets and ResNet-50. From Table 6, we can see that our framework performs much better than Deep Compression: it achieves higher accuracy with the same model size. For compact models like MobileNets, Deep Compression significantly degrades the performance especially under aggressive quantization, while our framework can preserve the accuracy much better. For instance, when Deep Compression quantizes the weights of MobileNet-V1 to 2 bits, the accuracy drops significantly from 70.90 to 37.62; while our framework can still achieve 57.14 of accuracy with the same model size. The reason is our framework makes full use of the mixed precision by systematically searching the optimal quantization policy.

In Figure 5, we visualize the bitwidth allocation strategy for MobileNet-V2. From this figure, we can observe that our framework assigns more bitwidths to the weights in depthwise convolution layers than pointwise convolution layers. Intuitively, this is because the number of parameters in the former is much smaller than the latter. Comparing Figure 4 and Figure 5, the policies are drastically different under different optimization objectives (fewer bitwiths for depthwise convolutions under latency optimization, more bitwidths for depthwise convolutions under model size optimization). Our framework succeeds in learning to adjust its bitwidth policy under different constraints.

Conclusion

In this paper, we propose Hardware-Aware Automated Quantization (HAQ), an automated framework for quantization which does not require any domain experts and rule-based heuristics. We provide a learning based method that can search the quantization policy with hardware feedback. Compared with indirect proxy signals, our framework can offer a specialized quantization solution for different hardware platforms. Extensive experiments demonstrate that our framework performs better than conventional rule-based approaches for multiple objectives: latency, energy and model size. Our framework reveals that the optimal policies on different hardware architectures are drastically different, and we interpreted the implication of those policies. We believe the insights will inspire the future software and hardware co-design for efficient deployment of deep neural networks.

We thank MIT Quest for Intelligence, MIT-IBM Watson AI Lab, Xilinx, Samsung, Intel, ARM, Qualcomm, and SONY for supporting this research. We thank Google Cloud and AWS Machine Learning Research Awards for providing the computation resource.

References