[论文解读] Channel-wise Hessian Aware trace-Weighted Quantization of Neural Networks

Xu Qian, Victor C. Li+1 more|arXiv (Cornell University)|Aug 19, 2020

Adversarial Robustness in Machine Learning参考文献 24被引用 9

一句话总结

该论文提出了一种基于深度强化学习的通道级Hessian感知迹加权量化方法（Channel-wise Hessian Aware trace-Weighted Quantization, CW-HAWQ），利用Hessian迹对通道敏感度进行排序，并采用DDPG智能体在权重和激活通道上优化混合精度位宽分配。在与AutoQ和HAWQ-V2相比压缩比相当的情况下，该方法在ResNet-50、ResNet-18和MobileNetV2上实现了最先进的模型精度，且top-1精度下降最小。

ABSTRACT

Second-order information has proven to be very effective in determining the redundancy of neural network weights and activations. Recent paper proposes to use Hessian traces of weights and activations for mixed-precision quantization and achieves state-of-the-art results. However, prior works only focus on selecting bits for each layer while the redundancy of different channels within a layer also differ a lot. This is mainly because the complexity of determining bits for each channel is too high for original methods. Here, we introduce Channel-wise Hessian Aware trace-Weighted Quantization (CW-HAWQ). CW-HAWQ uses Hessian trace to determine the relative sensitivity order of different channels of activations and weights. What's more, CW-HAWQ proposes to use deep Reinforcement learning (DRL) Deep Deterministic Policy Gradient (DDPG)-based agent to find the optimal ratios of different quantization bits and assign bits to channels according to the Hessian trace order. The number of states in CW-HAWQ is much smaller compared with traditional AutoML based mix-precision methods since we only need to search ratios for the quantization bits. Compare CW-HAWQ with state-of-the-art shows that we can achieve better results for multiple networks.

研究动机与目标

• 为解决现有混合精度量化方法仅优化每层的位宽，而非每通道的局限性，尽管层内通道间存在显著的冗余差异。
• 降低通道级混合精度量化搜索空间的复杂度，该复杂度对于穷举或标准强化学习方法而言不切实际。
• 通过利用二阶敏感度信息，联合优化权重和激活通道的量化位宽，以提升模型精度与压缩效率。
• 构建一个可扩展的强化学习框架，将状态空间简化为仅位宽比例配置，从而实现高效且有效的策略学习。

提出的方法

• 该方法计算权重和激活每个通道的Hessian迹，以确定其对量化误差的相对敏感度。
• 根据Hessian迹对通道进行排序，建立敏感度顺序，指导位宽分配决策。
• 训练一个深度确定性策略梯度（Deep Deterministic Policy Gradient, DDPG）智能体，以搜索各通道最优的量化位宽比例（例如2、3、4、5、6、7、8），而非搜索每个通道的独立位宽分配。
• 智能体在仅包含位宽比例配置的简化状态空间上运行，与传统基于AutoML的搜索相比，显著降低了计算复杂度。
• 该框架采用两阶段量化：首先使用DDPG智能体优化激活位宽，然后在所得模型上进行微调并优化权重位宽。
• 采用PACT进行激活量化，其裁剪阈值可学习；对权重使用对称三值或2-bit SAWB风格量化以减少量化误差。

实验结果

研究问题

• RQ1 Hessian迹能否有效用于在权重和激活中对通道级敏感度进行排序，以支持混合精度量化？
• RQ2 当状态空间被简化为仅位宽比例配置时，基于DDPG的强化学习智能体能否高效探索通道级混合精度搜索空间？
• RQ3 与层级或均匀量化相比，通道级混合精度量化在精度和压缩效率方面是否表现更优？
• RQ4 所提方法能否在保持极低top-1精度下降的前提下，实现优于SOTA方法（如AutoQ和HAWQ-V2）的精度？

主要发现

• 在ResNet-50上，CW-HAWQ实现了76.65%的top-1精度，平均权重位宽为2.61，激活位宽为4，相比全精度基线仅下降0.2%。
• 在14.48倍权重压缩比下，CW-HAWQ实现75.57%的top-1精度，平均权重位宽为2.21，激活位宽为3.07，精度比AutoQ高出1.2%，且top-1精度下降减少0.88%。
• 在ResNet-18上，CW-HAWQ实现69.02%的top-1精度，平均权重位宽为2.19，激活位宽为3.02，精度比AutoQ高出1.6%，top-1精度下降减少0.85%。
• 在MobileNetV2上，CW-HAWQ实现69.85%的top-1精度，平均权重位宽为2.26，激活位宽为3.13，精度比AutoQ高出1.2%，top-1精度下降减少0.47%。
• 该方法在所有评估模型上均实现了SOTA结果，表明通过DDPG实现的通道级敏感度感知位宽分配可显著提升精度与效率。