[论文解读] SDQ: Stochastic Differentiable Quantization with Mixed Precision

Xijie Huang, Zhiqiang Shen+6 more|arXiv (Cornell University)|Jun 9, 2022

Image Enhancement Techniques被引用 8

一句话总结

SDQ 提出了一种随机可微量化方法，通过可微分位宽参数（DBPs）和 Gumbel-softmax 重参数化，自动学习网络各层的最优混合精度位宽分配，实现平滑、全局优化的训练。该方法在 ResNet 和 MobileNet 上实现了最先进精度，平均位宽低于 4 位，且在 GPU 和 FPGA 上展现出优异的硬件效率。

ABSTRACT

In order to deploy deep models in a computationally efficient manner, model quantization approaches have been frequently used. In addition, as new hardware that supports mixed bitwidth arithmetic operations, recent research on mixed precision quantization (MPQ) begins to fully leverage the capacity of representation by searching optimized bitwidths for different layers and modules in a network. However, previous studies mainly search the MPQ strategy in a costly scheme using reinforcement learning, neural architecture search, etc., or simply utilize partial prior knowledge for bitwidth assignment, which might be biased and sub-optimal. In this work, we present a novel Stochastic Differentiable Quantization (SDQ) method that can automatically learn the MPQ strategy in a more flexible and globally-optimized space with smoother gradient approximation. Particularly, Differentiable Bitwidth Parameters (DBPs) are employed as the probability factors in stochastic quantization between adjacent bitwidth choices. After the optimal MPQ strategy is acquired, we further train our network with entropy-aware bin regularization and knowledge distillation. We extensively evaluate our method for several networks on different hardware (GPUs and FPGA) and datasets. SDQ outperforms all state-of-the-art mixed or single precision quantization with a lower bitwidth and is even better than the full-precision counterparts across various ResNet and MobileNet families, demonstrating the effectiveness and superiority of our method.

研究动机与目标

• 解决现有混合精度量化（MPQ）方法在位宽分配上依赖昂贵搜索或次优启发式策略的局限性。
• 开发一种可微分、全局优化的框架，实现网络各层最优位宽的端到端学习。
• 通过在反向传播过程中用随机量化替代线性插值，提升训练稳定性和收敛性。
• 通过熵感知正则化和知识蒸馏，最小化量化误差并保留信息承载组件。
• 在多种硬件平台上实现高精度、低比特宽模型，超越全精度模型和先前量化模型的性能。

提出的方法

• 引入可学习的概率因子——可微分位宽参数（DBPs），用于表示相邻候选位宽之间的随机选择。
• 在前向传播中采用随机量化，其中每个权重或激活值被量化为从由 DBPs 参数化的类别分布中采样得到的位宽。
• 使用 Gumbel-softmax 重参数化，实现对离散位宽选择的可微分反向传播，从而支持平滑的梯度更新。
• 应用熵感知二值化正则化（EBR），通过考虑各层特定的精度需求，保留信息丰富的权重组件并减少量化误差。
• 从全精度教师模型中引入知识蒸馏，以增强量化学生网络的表征能力。
• 使用标准反向传播端到端优化整个 MPQ 策略，实现所有层最优位宽分配的一次性学习。

实验结果

研究问题

• RQ1 可微分、随机的混合精度量化方法是否能在精度和优化稳定性方面超越强化学习和启发式位宽搜索方法？
• RQ2 与线性插值相比，基于 Gumbel-softmax 重参数化的随机量化在 MPQ 中如何提升训练稳定性和损失曲面平滑性？
• RQ3 熵感知正则化和知识蒸馏在低比特宽设置下，能在多大程度上减少量化误差并保持模型精度？
• RQ4 所提出的 SDQ 框架是否能在 ResNet 和 MobileNet 等多样化架构上实现平均位宽低于 4 位的最先进精度？
• RQ5 当部署在 FPGA 等真实硬件加速器上时，SDQ 在延迟和能效方面表现如何？

主要发现

• 在 ImageNet-1K 上使用 ResNet18 时，SDQ 实现了 71.7% 的 top-1 准确率，权重/激活的平均位宽为 3.85/2 位，超越了全精度模型的 70.5% 准确率。
• 在 ImageNet-1K 上使用 MobileNetV2 时，SDQ 实现了 71.8% 的 top-1 准确率，平均位宽为 4.02/3.85 位，优于所有先前的单精度和混合精度量化方法。
• 在 COCO 检测任务中使用 YOLOv4-tiny 时，SDQ 在平均位宽 3.88/4.00 的条件下实现了 15.9% 的 mAP，优于 4 位均匀量化，接近 8 位基线，同时将能效降低了 41.5%。
• 在 FPGA 上部署的 SDQ 模型实现了 21.28ms 的延迟和 167.1mJ 的能耗，其硬件效率与 4 位均匀量化模型相当。
• SDQ 的损失曲面显著比基于线性插值的方法更平滑，从而实现了更稳定、更有效的优化。
• 熵感知二值化正则化与知识蒸馏的结合，能更有效地保留信息承载组件并减少量化误差，尤其在低精度场景下表现更优。