[论文解读] Optimal Clipping and Magnitude-aware Differentiation for Improved Quantization-aware Training

Charbel Sakr, Steve Dai+4 more|arXiv (Cornell University)|Jun 13, 2022

Advanced Neural Network Applications被引用 11

一句话总结

该论文提出OCTAV，一种基于递归牛顿-拉夫森法的算法，在量化感知训练（QAT）过程中实时计算MSE最优的裁剪标量，从而在每次迭代中最小化每层张量的量化噪声。该方法进一步引入了幅度感知的反向传播，以改进梯度估计，在ResNets、MobileNets在ImageNet上的低精度（4–6比特）训练以及BERT微调任务中实现了最先进（SOTA）的精度表现，且无需对训练方案进行任何修改，仅需在适当层插入量化操作。

ABSTRACT

Data clipping is crucial in reducing noise in quantization operations and improving the achievable accuracy of quantization-aware training (QAT). Current practices rely on heuristics to set clipping threshold scalars and cannot be shown to be optimal. We propose Optimally Clipped Tensors And Vectors (OCTAV), a recursive algorithm to determine MSE-optimal clipping scalars. Derived from the fast Newton-Raphson method, OCTAV finds optimal clipping scalars on the fly, for every tensor, at every iteration of the QAT routine. Thus, the QAT algorithm is formulated with provably minimum quantization noise at each step. In addition, we reveal limitations in common gradient estimation techniques in QAT and propose magnitude-aware differentiation as a remedy to further improve accuracy. Experimentally, OCTAV-enabled QAT achieves state-of-the-art accuracy on multiple tasks. These include training-from-scratch and retraining ResNets and MobileNets on ImageNet, and Squad fine-tuning using BERT models, where OCTAV-enabled QAT consistently preserves accuracy at low precision (4-to-6-bits). Our results require no modifications to the baseline training recipe, except for the insertion of quantization operations where appropriate.

研究动机与目标

• 解决现有QAT方法中因启发式裁剪阈值导致的次优量化噪声问题。
• 通过在训练过程中实时计算MSE最优的裁剪标量，消除对超参数调优的依赖。
• 通过解决标准裁剪梯度方法的局限性，改进QAT中的梯度估计。
• 在不修改基线训练方案的前提下，实现低精度深度神经网络训练的SOTA精度。
• 提供一种通用、高效且可复现的低精度DNN训练方法，适用于动态与静态量化设置。

提出的方法

• OCTAV使用快速牛顿-拉夫森方法，递归计算每轮QAT中各张量的MSE最优裁剪标量。
• 该算法动态确定裁剪阈值，以最小化量化噪声，确保每一步均有可证明的最优性能。
• 引入幅度感知反向传播，通过根据激活值大小加权梯度，纠正标准梯度估计器中的偏差。
• 通过在适当层插入量化操作，将该方法无缝集成到标准QAT流水线中，无需修改训练方案。
• 在静态量化中，OCTAV用于校准阶段，训练前一次性计算最优裁剪标量。
• 该方法在多个模型（ResNets、MobileNets、BERT）和任务（从头训练、微调、重训练）中，以4–8比特精度进行了评估。

实验结果

研究问题

• RQ1 我们能否在不依赖启发式方法或超参数调优的情况下，计算出最小化QAT中量化噪声的裁剪标量？
• RQ2 与基于百分位数或MSE扫描的现有校准策略相比，OCTAV实现的最优裁剪在准确率和效率方面表现如何？
• RQ3 标准裁剪QAT中的梯度估计存在哪些局限性？能否通过幅度感知反向传播方案加以缓解？
• RQ4 在张量统计随时间变化的场景（如微调）中，动态OCTAV是否优于静态量化策略？
• RQ5 OCTAV能否在动态与静态量化设置中均有效应用，同时保持或提升准确率？

主要发现

• OCTAV-enabled QAT在ImageNet上对ResNet-50和MobileNet-V2实现了SOTA精度，4比特动态量化模型的top-1准确率达到87.09%。
• 在BERT-base微调任务中，动态OCTAV在6比特时仅损失约1%准确率，4比特时损失约4%，优于所有其他方法。
• 静态OCTAV校准在BERT-base上4比特精度下比MSE扫描方法高出约1.5%，因其能更好地处理非凸量化MSE损失曲面。
• OCTAV校准过程在CPU上比MSE扫描快约10倍，使其在大规模模型中更具实用性。
• 幅度感知反向传播显著改善了梯度估计，尤其在低精度场景下，有助于提升收敛性和准确率。
• OCTAV在不修改训练方案的前提下实现最优性能，仅需在适当层插入量化操作即可。