[论文解读] Mixed-Precision Neural Network Quantization via Learned Layer-wise Importance

Chen Tang, Kai Ouyang+5 more|arXiv (Cornell University)|Mar 16, 2022

Advanced Image and Video Retrieval Techniques被引用 5

一句话总结

该论文提出一种基于量化器缩放因子的可学习逐层重要性机制，通过整数线性规划（ILP）实现快速、一次性混合精度量化（MPQ）。通过在量化感知训练过程中端到端训练这些缩放因子，该方法捕捉了各层的量化敏感度，使ResNet18上的MPQ搜索时间缩短至0.06秒——比迭代方法快指数级——同时在多种约束条件下实现了ImageNet上的最先进精度。

ABSTRACT

The exponentially large discrete search space in mixed-precision quantization (MPQ) makes it hard to determine the optimal bit-width for each layer. Previous works usually resort to iterative search methods on the training set, which consume hundreds or even thousands of GPU-hours. In this study, we reveal that some unique learnable parameters in quantization, namely the scale factors in the quantizer, can serve as importance indicators of a layer, reflecting the contribution of that layer to the final accuracy at certain bit-widths. These importance indicators naturally perceive the numerical transformation during quantization-aware training, which can precisely provide quantization sensitivity metrics of layers. However, a deep network always contains hundreds of such indicators, and training them one by one would lead to an excessive time cost. To overcome this issue, we propose a joint training scheme that can obtain all indicators at once. It considerably speeds up the indicators training process by parallelizing the original sequential training processes. With these learned importance indicators, we formulate the MPQ search problem as a one-time integer linear programming (ILP) problem. That avoids the iterative search and significantly reduces search time without limiting the bit-width search space. For example, MPQ search on ResNet18 with our indicators takes only 0.06 s, which improves time efficiency exponentially compared to iterative search methods. Also, extensive experiments show our approach can achieve SOTA accuracy on ImageNet for far-ranging models with various constraints (e.g., BitOps, compress rate). Code is available on https://github.com/1hunters/LIMPQ.

研究动机与目标

• 为解决混合精度量化（MPQ）中指数级庞大的搜索空间问题，传统迭代搜索方法因此变得难以承受。
• 克服现有基于准则的MPQ方法中依赖全精度Hessian近似所带来的偏差，这些方法无法捕捉量化特有的数值行为。
• 消除MPQ中对人工指定位宽或受限搜索空间的需求，实现对权重和激活位宽选择的完全灵活性。
• 开发一种快速、一次性MPQ策略搜索方法，避免在训练数据上重复评估，从而大幅降低时间成本。
• 建立一种可学习的、端到端的重要程度指示器，反映量化感知训练过程中各层的敏感度。

提出的方法

• 利用量化器中的可学习缩放因子作为端到端的重要程度指示器，反映特定位宽下各层对量化敏感的程度。
• 通过并行化的联合训练方案联合训练所有逐层重要程度指示器，避免顺序训练，降低时间成本。
• 将MPQ搜索问题建模为一次性整数线性规划（ILP）问题，利用学习到的重要程度指示器为每层分配最优位宽。
• 使用ILP解为权重和激活分配混合精度位宽，实现细粒度、灵活的精度-效率权衡。
• 将学习到的指示器集成到完整的量化流程中：重要程度训练、基于ILP的策略搜索和模型微调。
• 确保方法支持对权重和激活的完整搜索空间探索（例如，4位、8位等），无需人工约束。

实验结果

研究问题

• RQ1 量化器中的可学习缩放因子能否作为可靠、端到端的层量化敏感度指示器？
• RQ2 与顺序训练相比，重要程度指示器的联合训练方案是否显著降低训练时间？
• RQ3 基于ILP的MPQ策略搜索能否在接近零秒的搜索时间内实现SOTA精度？
• RQ4 在精度和效率方面，该方法与基于迭代搜索和基于Hessian的准则方法相比表现如何？
• RQ5 基于缩放因子大小反转位宽分配是否会降低性能，从而验证所提出重要程度度量的有效性？

主要发现

• 该方法在ResNet18上实现MPQ策略搜索仅需0.06秒，相比AutoQ等迭代方法（需超过1000 GPU小时）实现了指数级加速。
• 对于ResNet50，该方法在50分钟内完成重要程度训练（3.3 GPU小时），ILP搜索仅需0.35秒，使多个部署设备的搜索速度相比AutoQ提升330倍。
• 与HAWQv2相比，该方法在ResNet50上将top-1精度提升1.1%，同时保持对权重和激活的完整搜索空间灵活性。
• 消融实验表明，基于缩放因子大小反转位宽分配会导致top-1精度下降6.59%，证实了所提基于重要程度的分配策略的合理性。
• 在9.68 G BitOps约束下，该方法实现71.84%的top-1精度（权重和激活均为4位），而反转分配方案降至65.25%，证明了所学习指示器的鲁棒性。
• 该方法在多种模型和约束条件（包括BitOps和压缩率）下均实现了ImageNet上的最先进精度，且未限制搜索空间。