[论文解读] Trained Quantization Thresholds for Accurate and Efficient Fixed-Point Inference of Deep Neural Networks

Sambhav R. Jain, Albert Gural+2 more|arXiv (Cornell University)|Mar 19, 2019

Advanced Neural Network Applications参考文献 32被引用 55

一句话总结

本文提出 Trained Quantization Thresholds (TQT)，一种在重新训练时通过反向传播在幂为 2 的、逐张量、对称定点约束下学习量化阈值的方法，在 8 位网络上接近浮点精度，在 4 位网络上接近 FP 精度。它还提供 Graffitist，一个基于 TensorFlow 的自动量化与重新训练框架。

ABSTRACT

We propose a method of training quantization thresholds (TQT) for uniform symmetric quantizers using standard backpropagation and gradient descent. Contrary to prior work, we show that a careful analysis of the straight-through estimator for threshold gradients allows for a natural range-precision trade-off leading to better optima. Our quantizers are constrained to use power-of-2 scale-factors and per-tensor scaling of weights and activations to make it amenable for hardware implementations. We present analytical support for the general robustness of our methods and empirically validate them on various CNNs for ImageNet classification. We are able to achieve near-floating-point accuracy on traditionally difficult networks such as MobileNets with less than 5 epochs of quantized (8-bit) retraining. Finally, we present Graffitist, a framework that enables automatic quantization of TensorFlow graphs for TQT (available at https://github.com/Xilinx/graffitist ).

研究动机与目标

• 通过与权重共同学习量化阈值来推动鲁棒的低精度 DNN 推理。
• 开发一种基于梯度的在对数域参数化下训练裁剪阈值的方法。
• 确保硬件友好的量化，具备对称、逐张量、幂为 2 的尺度因子。
• 提供一个端到端框架（Graffitist）用于 TensorFlow 的自动量化与重新训练。
• 在具有挑战性的模型（如 MobileNets）上，通过有限的重新训练轮次展示接近浮点精度的性能。

提出的方法

• 定义一个对称仿射量化器，尺度为 s，阈值为 t，其中 s = 2^(-f)，f 为整数。
• 在对数域对阈值进行训练，以通过反向传播在量域范围和精度之间取得平衡。
• 推导基于 STE 的输入 x 和尺度 s 的梯度，梯度通过量化器进行反向传播，并包含对 log2(t) 的梯度。
• 对权重和激活进行逐张量量化，采用 8 位和 4 位方案，且尺度约束为幂为 2。
• 引入 Graffitist，在 TensorFlow 中自动化图变换、量化和重新训练，包括融合内核以减少内存开销。

实验结果

研究问题

• RQ1 阈值是否可以与模型权重联合学习，以优化端到端损失？
• RQ2 在对数域中训练阈值参数是否能提高量化重新训练的稳定性与收敛性？
• RQ3 幂为 2、逐张量、对称量化约束如何影响常见 CNN 的精度？
• RQ4 统一框架（Graffitist）是否对跨多种架构的自动量化与重新训练有效？
• RQ5 对于困难网络（如 MobileNets），与仅对权重重新训练相比，TQT 的精度提升有哪些？

主要发现

• 在多种 CNN 的重新训练后，8 位量化接近浮点精度。
• 在逐张量、对称、幂为 2 的尺度下，TQT 能在困难网络如 MobileNets 上达到 FP32 精度。
• 对于 INT4，需要对权重+阈值进行重新训练才能恢复精度，在受限量化下部分网络接近 FP32 表现。
• 在具有挑战性的模型上，重新训练阈值通常比仅对权重重新训练带来约 4% 的绝对 Top-1 精度增益。
• 静态量化通常不及基于重新训练的方法，因为与全局损失优化的对齐度较差。
• 在对数域训练阈值相较直接对尺度因子训练，能提升稳定性和收敛性。