Source-linked AI summary
APQ: Joint Search for Network Architecture, Pruning and Quantization Policy
Tianzhe Wang, Kuan Wang, Han Cai, Ji Lin, Zhijian Liu, Song Han
TL;DR
Efficient inference requires jointly selecting architecture, pruning, and quantization, but quantization-aware predictor training is expensive. APQ transfers knowledge from an fp32 predictor built from a once-for-all network, enabling low-cost joint search and improved ImageNet efficiency results.
Problem
Separately searching architecture, pruning, and quantization can be sub-optimal, while collecting quantized model-accuracy data requires expensive fine-tuning.
Method
APQ jointly searches architecture, pruning, and quantization using a quantization-aware predictor transferred from an fp32 predictor trained with once-for-all network evaluations.
Results
APQ matches MobileNetV2+HAQ accuracy with 2× latency and 1.3× energy savings, and gains 2.3% ImageNet accuracy over ProxylessNAS+AMC+HAQ at equal latency.
Takeaways & Limitations
APQ enables low-cost mixed-precision architecture search without extra training or evaluation interaction with the target dataset during search.
Abstract
from arXiv · showhide
We present APQ for efficient deep learning inference on resource-constrained hardware. Unlike previous methods that separately search the neural architecture, pruning policy, and quantization policy, we optimize them in a joint manner. To deal with the larger design space it brings, a promising approach is to train a quantization-aware accuracy predictor to quickly get the accuracy of the quantized model and feed it to the search engine to select the best fit. However, training this quantization-aware accuracy predictor requires collecting a large number of quantized <model, accuracy> pairs, which involves quantization-aware finetuning and thus is highly time-consuming. To tackle this challenge, we propose to transfer the knowledge from a full-precision (i.e., fp32) accuracy predictor to the quantization-aware (i.e., int8) accuracy predictor, which greatly improves the sample efficiency. Besides, collecting the dataset for the fp32 accuracy predictor only requires to evaluate neural networks without any training cost by sampling from a pretrained once-for-all network, which is highly efficient. Extensive experiments on ImageNet demonstrate the benefits of our joint optimization approach. With the same accuracy, APQ reduces the latency/energy by 2x/1.3x over MobileNetV2+HAQ. Compared to the separate optimization approach (ProxylessNAS+AMC+HAQ), APQ achieves 2.3% higher ImageNet accuracy while reducing orders of magnitude GPU hours and CO2 emission, pushing the frontier for green AI that is environmental-friendly. The code and video are publicly available.
1. Introduction
APQ jointly searches architecture, pruning, and quantization to optimize accuracy, latency, and energy under resource constraints. It addresses the enlarged search cost through predictor transfer and efficient candidate evaluation.
- Motivation and approach: The combined search space requires automation because its hyper-parameters grow exponentially and can exceed acceptable human labor bandwidth.Separate optimization can select a full-precision architecture that is not necessarily best after pruning and quantization.
- Predictor-based search: APQ uses a quantization-aware accuracy predictor to estimate candidate accuracy instead of fine-tuning pruned and quantized networks during search.The predictor takes architecture and quantization information as input and enables ultra-fast search.
- Predictor-based search: Predictor transfer pretrains on fp32 accuracy data from a once-for-all network, then adapts to quantized models using a much smaller dataset and short-term fine-tuning.Direct low-bit quantization without fine-tuning usually produces near-zero accuracy, making quantized data collection expensive.
- Results: 600× fewer GPU hours and CO2 emission accompany APQ’s 2.3% higher ImageNet accuracy than ProxylessNAS+AMC+HAQ at the same latency.Against MobileNetV2+HAQ, APQ achieves the same accuracy with 2× latency and 1.3× energy savings.
- Motivation and approach: Jointly optimizing architecture, pruning, and quantization avoids the sub-optimality of separate stages while targeting accuracy, latency, and energy end-to-end.The joint search space is multiplicative and larger than stage-wise spaces, making direct optimization difficult.
2. Background and Outline
Efficient deployment combines architecture design, pruning, and quantization, but optimizing these stages separately can be costly and sub-optimal. APQ unifies them through a once-for-all network, accuracy prediction, and resource-constrained evolutionary search.
- Background: Architecture search, pruning, and quantization are established approaches for accelerating model inference on hardware platforms.Pruning removes connections or channels, while quantization supports deployment on platforms such as FPGAs and mobile phones.
- Limitations of Separate Optimization: 300 GPU hours can be required by a typical sequential design pipeline, making automatic design difficult for researchers with limited computation.Frequent target-dataset evaluations contribute to the time cost.
- Limitations of Separate Optimization: Separately optimizing architecture, pruning, and quantization can produce sub-optimal solutions because the best floating-point architecture may not remain optimal after quantization.The stages also expose a large search space because layer sensitivities differ for pruning.
- Joint Optimization: APQ jointly balances neural architecture, pruning, and quantization configurations instead of optimizing them independently.Its objective formalizes the search for an optimal strategy across these configurations.
- APQ Outline: APQ trains a once-for-all network, builds a quantization-aware accuracy predictor, and performs resource-constrained evolutionary search.The once-for-all network covers a large search space and allows sub-networks to be extracted without retraining; the search uses latency or energy lookup tables.
3. Joint Design Methodoloy
APQ jointly searches architecture, pruning, and quantization using a once-for-all network, a quantization-aware accuracy predictor, and resource-constrained evolution search. Predictor transfer reduces the data and fine-tuning burden of quantized-model accuracy prediction.
- APQ combines a flexible once-for-all network, an accuracy predictor, and evolution search to jointly optimize architecture, pruning, and quantization.
- Once-For-All Network: The once-for-all network supports directly evaluated subnetworks and a fine-grained channel search space, trained with progressive shrinking to manage accuracy variance.
- Quantization-Aware Accuracy Predictor: The predictor estimates mixed-precision accuracy from architecture, pruning, and quantization encodings, replacing repeated candidate evaluation during search.
- Quantization-Aware Accuracy Predictor: APQ encodes block-level kernel sizes, channel numbers, and weight/activation bitwidths as concatenated one-hot vectors for prediction and search.
- Quantization-Aware Accuracy Predictor: Predictor transfer pretrains on cheaply measured full-precision subnetworks, then fine-tunes with a smaller quantized dataset initialized from the full-precision predictor.
- Hardware-Aware Evolutionary Search: Resource-constrained evolution search uses predicted accuracy and lookup-table latency or energy estimates to eliminate candidates exceeding hardware budgets.
4. Implementation Details
APQ’s implementation samples architecture and quantization data, applies evolutionary search, and evaluates quantization and deployment costs with hardware-oriented procedures. Experiments compare APQ against fixed- and mixed-precision baselines under efficiency constraints.
- Data Preparation: The quantization-aware predictor uses 2,500 samples of each preparation strategy, while the full-precision predictor uses 80,000 samples.
- Evolutionary Architecture Search: Evolutionary search uses population size 100, selects Top-25 candidates, and generates 50 mutations plus 50 crossovers per generation.
- Quantization: Weights and activations are quantized with layer-specific ranges, choosing weight ranges to minimize KL-divergence from the original weights.
- Comparisons: Table 2 compares fixed- and mixed-precision efficient models under different latency constraints, with marginal search time reduced by two orders of magnitude.
- Comparisons: Under strict constraints, APQ can outperform fixed-precision models by more than 10% accuracy and HAQ by 5%.
5. Experiments
Experiments on ImageNet evaluate APQ under latency, energy, and BitOps constraints against fixed- and mixed-precision baselines, showing consistent accuracy and efficiency gains. Predictor transfer improves accuracy-prediction quality and reduces computational cost.
- Experimental setup and overall comparison: APQ consistently outperforms fixed- and mixed-precision state-of-the-art models across the evaluated efficiency constraints.Experiments use ImageNet and measure resource consumption on the BitFusion accelerator.
- Comparison with SOTA efficient models: 2.2% accuracy boost over mixed-precision MobileNetV2 searched by HAQ improves accuracy from 71.9% to 74.1%.
- Comparison with SOTA efficient models: 50% lower BitOps accompanies accuracy increasing from 74.6% to 75.1% for the large APQ model.
- Efficiency of predictor transfer: Transfer technology improves model performance from 72.1% to 74.1% while reducing marginal cloud-computing cost and CO2 emission by two orders of magnitude.
- Comparison with MobileNetV2+HAQ: 10.5% and 11.3% ImageNet top-1 accuracy gains over the MobileNetV2 baseline occur under similar latency and energy constraints, respectively.Accuracy rises from 61.4% to 71.9% for latency and 72.7% for energy constraints.
- Effectiveness of joint design: 74.1% versus 71.8% accuracy under the same latency or energy constraint favors joint optimization over multi-stage optimization.
- Effectiveness of joint design: Over 2% accuracy improvement, from 71.5% to 73.9%, is achieved under a tight BitOps constraint, while 75.1% matches 8-bit ResNet-34 at 8× lower BitOps.
- Effectiveness of predictor transfer: Predictor transfer improves pairwise accuracy from 64.6% to 75.6% with limited data and yields faster convergence.The evaluation uses 2000 held-out architecture-and-accuracy pairs.
6. Conclusion
APQ jointly searches architecture, pruning, and quantization for mixed-precision models instead of optimizing them in separate stages. Its predictor-based search and predictor transfer reduce evaluation and data-collection costs while supporting efficient deployment.
- APQ directly searches for mixed-precision architectures jointly across architecture, pruning, and quantization policies.
- The predictor-based method avoids extra target-dataset evaluation during search, reducing GPU hours, CO2 emission, and cloud-compute cost.
- Predictor transfer addresses the high cost of collecting quantized model-accuracy data and supports the joint-search approach.