Recent work in network quantization has substantially reduced the time and space complexity of neural network inference, enabling their deployment on embedded and mobile devices with limited computational and memory resources. However, existing quantization methods often represent all weights and activations with the same precision (bit-width). In this paper, we explore a new dimension of the design space: quantizing different layers with different bit-widths. We formulate this problem as a neural architecture search problem and propose a novel differentiable neural architecture search (DNAS) framework to efficiently explore its exponential search space with gradient-based optimization. Experiments show we surpass the state-of-the-art compression of ResNet on CIFAR-10 and ImageNet. Our quantized models with 21.1x smaller model size or 103.9x lower computational cost can still outperform baseline quantized or even full precision models. Recently, ConvNets have become the de-facto method in a wide range of computer vision tasks, achieving state-of-the-art performance. However, due to high computation complexity, it is nontrivial to deploy ConvNets to embedded and mobile devices with limited computational and storage budgets. In recent years, research efforts in both software and hardware have focused on lowprecision inference of ConvNets. Most of the existing quantization methods use the same precision for all (or most of) the layers of a ConvNet. However, such uniform bit-width assignment can be suboptimal since quantizing different layers can have different impact on the accuracy and efficiency of the overall network. Although mixed precision computation is widely supported in a wide range of hardware platforms such as CPUs, FPGAs, and dedicated accelerators, prior efforts have not thoroughly explored the mixed precision quantization of ConvNets.For a ConvNet with N layers and M candidate precisions in each layer, we want to find an optimal assignment of precisions to minimize the cost in terms of model size, memory footprint or computation, while keeping the accuracy. An exhaustive combinatorial search has exponential time complexity (O(M N )). Therefore, we need a more efficient approach to explore the design space.In this work, we propose a novel, effective, and efficient differentiable neural architecture search (DNAS) framework to solve this problem. The idea is illustrated in FIG0 . The problem of neural architecture search (NAS) aims to find the optimal neural net architecture in a given search space.In the DNAS framework, we represent the architecture search space with a stochastic super net where nodes represent intermediate data tensors of the super net (e.g., feature maps of a ConvNet) and edges represent operators (e.g., convolution layers in a ConvNet). Any candidate architecture can be seen as a child network (sub-graph) of the super net. When executing the super net, edges are executed stochastically and the probability of execution is parameterized by some architecture parameters θ. Under this formulation, we can relax the NAS problem and focus on finding the optimal θ that gives the optimal expected performance of the stochastic super net. The child network can then be sampled from the optimal architecture distribution.We solve for the optimal architecture parameter θ by training the stochastic super net with SGD with respect to both the network's weights and the architecture parameter θ. To compute the gradient of θ, we need to back propagate gradients through discrete random variables that control the stochastic edge execution. To address this, we use the Gumbel SoftMax function BID9 ) to "soft-control" the edges. This allows us to directly compute the gradient estimation of θ with a controllable trade-off between bias and variance. Using this technique, the stochastic super net becomes fully differentiable and can be effectively and efficiently solved by SGD. We apply the DNAS framework to solve the mixed precision quantization problem, by constructing a super net whose macro architecture (number of layers, filter size of each layer, etc.) is the same as the target network. Each layer of the super net contains several parallel edges representing convolution operators with quantized weights and activations with different precisions. We show that using DNAS to search for layer-wise precision assignments for ResNet models on CIFAR10 and ImageNet, we surpass the state-of-the-art compression. Our quantized models with 21.1x smaller model size or 103.9x smaller computational cost can still outperform baseline quantized or even full precision models. The DNAS pipeline is very fast, taking less than 5 hours on 8 V100 GPUs to complete a search on ResNet18 for ImageNet, while previous NAS algorithms (such as Zoph & Le (2016)) typically take a few hundred GPUs for several days. Last, but not least, DNAS is a general architecture search framework that can be applied to other problems such as efficient ConvNet-structure discovery. Due to the page limit, we will leave the discussion to future publications. In this work we focus on the problem of mixed precision quantization of a ConvNet to determine its layer-wise bit-widths. We formulate this problem as a neural architecture search (NAS) problem and propose a novel, efficient, and effective differentiable neural architecture search (DNAS) framework to solve it. Under the DNAS framework, we efficiently explore the exponential search space of the NAS problem through gradient based optimization (SGD). We use DNAS to search for layer-wise precision assignment for ResNet on CIFAR10 and ImageNet. Our quantized models with 21.1x smaller model size or 103.9x smaller computational cost can still outperform baseline quantized or even full precision models. DNAS is very efficient, taking less than 5 hours to finish a search on ResNet18 for ImageNet. It is also a general architecture search framework that is not limited to the mixed precision quantization problem. Its other applications will be discussed in future publications. DISPLAYFORM0 w denotes the latent full-precision weight of a network. Q k (·) denotes a k-bit quantization function that quantizes a continuous value w ∈ [0, 1] to its nearest neighbor in { DISPLAYFORM1 To quantize activations, we follow to use a bounded activation function followed by a quantization function as DISPLAYFORM2 Here, x is the full precision activation, y k is the quantized activation. P ACT (·) is a function that bounds the output between [0, α]. α is a learnable upper bound of the activation function.
A novel differentiable neural architecture search framework for mixed quantization of ConvNets. The authors introduce a new method for neural architecture search which selects the precision quantization of weights at each neural network layer, and use it in the context of network compression. The paper presents a new approach in network quantization by quantizing different layers with different bit-widths and introduces a new differentiable neural architecture search framework.