Source-linked AI summary
Mixed Precision Quantization of ConvNets via Differentiable Neural Architecture Search
Bichen Wu, Yanghan Wang, Peizhao Zhang, Yuandong Tian, Peter Vajda, Kurt Keutzer
TL;DR
Uniform precision quantization does not account for layer-specific effects on accuracy and efficiency, making mixed-precision assignment a challenging exponential search problem. The paper formulates this assignment as NAS and introduces DNAS, which uses a differentiable stochastic super net and gradient-based optimization. Its ResNet models achieve up to 21.1x smaller model size or 103.9x lower computational cost while outperforming baseline quantized or full-precision models.
Problem
Uniform precision across layers can be suboptimal, while selecting among M precisions across N layers creates an exponential search space.
Method
DNAS represents candidate precision assignments in a stochastic super net and explores them through gradient-based optimization using SGD and Gumbel Softmax.
Results
21.1x smaller model size or 103.9x lower computational cost was achieved while outperforming baseline quantized or full-precision models.
Takeaways & Limitations
Layer-wise precision search can produce highly compressed ResNet models while preserving or exceeding the accuracy of comparison models.
Takeaways & Limitations
Naive Monte Carlo gradient estimation has high variance because the architecture space is much larger than any feasible batch size.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
The paper targets layer-wise mixed-precision quantization as an alternative to uniform bit-width assignment and introduces DNAS to search this exponential design space efficiently. Experiments on ResNet show strong compression while preserving or improving accuracy.
- Uniform bit-width assignment can be suboptimal because different layers affect overall network accuracy and efficiency differently.
- For N layers and M candidate precisions, exhaustive precision assignment requires exponential time complexity O(M^N).
- DNAS represents the search space with a stochastic super net whose edges are candidate operators executed according to architecture parameters θ.
- DNAS optimizes θ by training the stochastic super net with SGD, then samples a child network from the learned architecture distribution.
- 21.1x smaller model size or 103.9x lower computational cost was achieved while outperforming baseline quantized or full-precision models.The search on ResNet18 for ImageNet took less than 5 hours on 8 V100 GPUs.
2 RELATED WORK
Prior work reduced quantization and architecture-search costs, but this paper identifies neural architecture search for model quantization as a new direction.
- Weight-only quantization reduces model size, storage requirements, and over-the-air communication costs while retaining 32-bit activations.
- Neural architecture search methods aim to discover accurate and efficient architectures, but early reinforcement-learning approaches require huge computing resources.
- The authors state that they are the first to apply neural architecture search to model quantization.
3 MIXED PRECISION QUANTIZATION
Mixed-precision quantization assigns different bit-widths to different layers, turning precision selection into a cost-constrained design problem.
- Quantization maps 32-bit floating-point weights and activations to lower-bit fixed-point representations such as 8, 4, and 1 bit.
- Mixed-precision quantization allows different layers to use different precisions, supported by CPUs, FPGAs, and dedicated accelerators.
- The design objective is to choose each layer’s precision while maintaining network accuracy and minimizing model-size or computation cost.
4 DIFFERENTIABLE NEURAL ARCHITECTURE SEARCH
DNAS relaxes discrete architecture selection into differentiable optimization over a stochastic super net. Gumbel Softmax and temperature annealing manage the bias–variance trade-off during search, after which sampled architectures are trained and tested.
- NAS formulation: The NAS objective is difficult because loss is differentiable with respect to network weights but not architecture choices, making repeated convergence training expensive.
- Super net representation: A super net represents the architecture space as a computational DAG whose nodes are data tensors and whose edges are candidate operators.
- Super net representation: Candidate architectures are encoded by edge-mask vectors that select one candidate edge between each connected node pair.
- Stochastic optimization: DNAS samples architectures from a parameterized distribution and minimizes expected loss instead of directly solving the combinatorial architecture problem.
- Gumbel Softmax: Gumbel Softmax makes edge selection continuous and differentiable with respect to architecture parameters, avoiding direct backpropagation through discrete sampling.
- Gumbel Softmax: Temperature annealing favors low-variance biased gradients early and higher-variance unbiased gradients later in training.
- Relation to prior NAS: DNAS occupies a middle ground between ENAS and DARTS, behaving more like DARTS early and ENAS late as temperature decreases.
- DNAS pipeline: The pipeline alternates weight and architecture-parameter training, samples architectures, trains them on Xtrain, and evaluates them on Xtest.
5 DNAS FOR MIXED PRECISION QUANTIZATION
DNAS models mixed-precision quantization by assigning candidate bit-widths to parallel convolution operators within a fixed macro-architecture. Its cost-aware objective supports optimizing either model size or computational cost while balancing efficiency against accuracy.
- Super-net construction: Each super-net layer preserves the target ConvNet’s macro-structure while offering parallel convolution operators with different weight and activation precisions.Nodes correspond to layer output feature maps, and candidate edges represent alternative quantized convolution operators.
- Cost-aware objective: Cost(a) represents the cost of a candidate architecture, while C(·) weights the cost term relative to cross-entropy.The coefficient γ controls this trade-off: larger values favor efficiency over accuracy.
- Model-size objective: Model-size compression uses a cost based on convolution-operator parameter counts and weight bit-widths.The weighting includes the selected edge masks and the operator’s weight precision.
- Computational objective: Joint weight-and-activation compression uses floating-point operation counts multiplied by weight and activation bit-widths.This alternative cost targets computational savings rather than model size alone.
- Differentiable search: Continuous edge-selection masks in the super-net enable computation of expected cost before candidate architectures use binary selections.The relaxation makes the cost evaluable during optimization while retaining one selected edge per connected node pair in a candidate architecture.
6 EXPERIMENTS
Experiments evaluate layer-wise mixed-precision search on CIFAR10 and ImageNet, targeting model-size and computational-cost compression while preserving accuracy. DNAS finds assignments that outperform prior quantized models and remain competitive with full-precision baselines.
- 6.1 CIFAR10 EXPERIMENTS: CIFAR10 experiments search block-level weight precisions for ResNet20, ResNet56, and ResNet110 while retaining full-precision activations.The first and last layers are not quantized, and all layers within a block share one precision.
- 6.1 CIFAR10 EXPERIMENTS: 11.6–12.5X model-size reduction accompanies the most accurate CIFAR10 models, which outperform full-precision counterparts by up to 0.37%.The most efficient models achieve 16.6–20.3X compression with accuracy drops below 0.39%.
- 6.1 CIFAR10 EXPERIMENTS: Searched precision assignments generally provide higher compression and accuracy than random assignments for ResNet110 on CIFAR10.Figure 3 plots accuracy against compression rate for all searched architectures.
- 6.2 IMAGENET EXPERIMENTS: ImageNet model-size experiments search weight precisions from 1 to 32 bits and report the most accurate and most efficient architectures.The study uses ResNet18 and ResNet34, with mixed-precision search conducted at the block level.
- 6.2 IMAGENET EXPERIMENTS: 10.6–11.2X model-size reduction accompanies ImageNet models that outperform full precision by up to 0.5%, while the most efficient models reach 19.0–21.1X reduction with competitive accuracy.Compared with previous works, a less accurate searched model has almost the same accuracy as a full-precision model at 21.1X smaller size.
- 6.2 IMAGENET EXPERIMENTS: 33.2x and 40.8X computational-cost compression achieve nearly unchanged accuracy, while another searched architecture reaches 103.5x compression and exceeds GroupNet’s 102.4x rate.At similar compression to 4-bit baselines, arch-2 is 0.71–1.91% more accurate.
7 CONCLUSION
The paper formulates mixed-precision quantization as neural architecture search and uses gradient-based DNAS to explore its exponential design space. On CIFAR10 and ImageNet, searched models achieve substantial model-size or computational-cost reductions while outperforming baseline quantized or full-precision models.
- 7 CONCLUSION: DNAS searches layer-wise precision assignments for ResNet models on CIFAR10 and ImageNet through gradient-based optimization.The framework explores the exponential NAS search space efficiently.
- 7 CONCLUSION: 21.1x smaller model size or 103.9x smaller computational cost can still outperform baseline quantized or full-precision models.The reported search on ImageNet ResNet18 takes less than 5 hours.
- 7 CONCLUSION: DNAS is presented as a general architecture-search framework beyond mixed-precision quantization.Other applications are left for future publications.
APPENDIX A WEIGHT AND ACTIVATION QUANTIZATION
The appendix specifies quantization functions for weights and activations, using DoReFa-Net for weights and bounded PACT activations followed by k-bit quantization.
- Weight quantization: Weights are quantized following DoReFa-Net.The appendix introduces the weight-quantization procedure before describing activation quantization.
- Weight quantization: Qk(·) maps continuous values to their nearest k-bit quantization level.The quantization function operates on values in [0, 1] and selects among 2^k discrete levels.
- Activation quantization: Activations use a bounded PACT function followed by a quantization function.PACT bounds outputs to [0, α], where α is a learnable upper bound; x denotes the full-precision activation and yk the quantized activation.
APPENDIX B EXPERIMENT DETAILS
The experiments evaluate CIFAR10 and ImageNet using separate super-net training procedures, with dataset-specific splits, optimization settings, and evaluation protocols.
- CIFAR10 experiments: CIFAR10 experiments use 50,000 training images, 10,000 test images, 10 categories, and 32 × 32 inputs.Accuracy is reported on the test set, with 80% of the training data used for weights and 20% for architecture parameters.
- CIFAR10 experiments: CIFAR10 super-net training runs for 90 epochs with batch size 512 and separate SGD and Adam optimization settings.SGD trains weights, while Adam trains architecture parameters; the passage also specifies momentum, weight decay, learning rates, and cosine decay.
- ImageNet experiments: ImageNet experiments use 1,000 classes, roughly 1.3M training images, 50K validation images, and 224 × 224 crops.Validation accuracy is reported, and 40 randomly sampled categories are used to train the super net because full ImageNet training is computationally expensive.