[论文解读] BatchQuant: Quantized-for-all Architecture Search with Robust Quantizer

Haoping Bai, Meng Cao+2 more|arXiv (Cornell University)|May 19, 2021

Machine Learning and Algorithms参考文献 41被引用 14

一句话总结

本文提出Quantized-for-All (QFA)，一种新颖的一次性混合精度神经架构搜索框架，可在无需微调的情况下联合搜索架构与任意超低比特宽度量化策略。它引入了BatchQuant，一种基于批量统计的鲁棒量化器，可稳定紧凑的、权重重叠的超网络训练——在仅190个GPU训练周期内实现20 MFLOPs以下的SOTA ImageNet Top-1准确率，并支持从单一超网络部署超过10^76种量化子网络。

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^{76}$ 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.

研究动机与目标

• 解决在资源受限的边缘部署场景中，联合搜索庞大且混合精度的量化与架构空间的挑战。
• 克服混合精度权重重叠超网络训练中的不稳定性，该问题阻碍了高效的一次性NAS与量化结合。
• 在部署搜索得到的量化子网络时，消除对微调或再训练的需求，实现对新场景的快速适应。
• 通过在各层灵活支持任意比特宽度组合，实现在低FLOP约束下的SOTA准确率。
• 通过最小化每个新部署场景的边际GPU小时数与碳足迹，降低搜索成本与碳排放。

提出的方法

• 提出BatchQuant (BQ)，一种即插即用的量化器，利用批量统计自适应估算缩放因子，提升在混合精度子网络选择过程中对激活分布偏移的鲁棒性。
• 引入一种新型极值估计器（公式8），通过使用可微的、基于批量的分位数估计方法替代原始的运行最小/最大值，从而稳定训练。
• 训练一个单一、紧凑的、混合精度权重重叠超网络，通过在架构与比特宽度选择间共享参数，支持超过10^76种不同的量化子网络。
• 利用NSGA-II实现Pareto最优架构发现，实现对不同量化策略下准确率与FLOPs的同步优化。
• 使用在小规模验证集上训练的轻量级准确率预测器，避免在搜索过程中进行昂贵的评估，从而实现快速、可扩展的探索。
• 将BQ集成到超网络训练流程中，保持全精度权重的同时，按子网络动态融合量化器，实现零微调推理。

实验结果

研究问题

• RQ1 是否可以稳定训练一个单一的统一超网络，在各层支持任意比特宽度的混合精度量化，而无需为每种配置重新训练？
• RQ2 如何通过量化器设计缓解在超网络训练期间因动态量化策略选择导致的激活分布偏移？
• RQ3 与运行最小/最大值或可学习量化器基线相比，所提出的BatchQuant公式在训练稳定性和收敛性方面提升程度如何？
• RQ4 通过单一超网络实现联合架构与量化策略搜索，是否能在严格FLOP约束（如<20 MFLOPs）下实现SOTA性能，且无需微调？
• RQ5 在使用该方法的动态边缘场景中，部署新型量化模型的边际成本与碳足迹是多少？

主要发现

• BatchQuant实现了对混合精度、权重重叠超网络的稳定训练，支持超过10^76种可能的量化子网络，仅用190个GPU训练周期即实现收敛，远少于OQA的495个周期。
• 该方法在20 MFLOPs以下实现了SOTA ImageNet Top-1准确率，在准确率与FLOPs的权衡中显著优于APQ与SPOS。
• 在新部署场景中搜索的边际CO2排放可忽略不计（0磅），仅需在初始超网络训练后额外40小时的数据收集与训练。
• 若移除BQ中的残差项γ与β，训练损失将发散，证明其在稳定量化过程中的关键作用。
• 仅公式8（基于批量的极值估计器）能实现稳定训练；公式7与9导致发散，证明了统计公式的必要性。
• QFA框架在搜索后无需微调或再训练，因为任何有效架构-量化组合的权重均可直接从超网络继承。