[Paper Review] 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 ApplicationsComputer Science32 references55 citations

TL;DR

The paper introduces Trained Quantization Thresholds (TQT), a method to learn quantization thresholds during retraining with backpropagation under power-of-2, per-tensor, symmetric fixed-point constraints, achieving near-floating-point accuracy for 8-bit and close to FP accuracy for 4-bit networks. It also provides Graffitist, a TensorFlow-based framework for automatic quantization and retraining.

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 ).

Motivation & Objective

• Motivate robust low-precision DNN inference by learning quantization thresholds jointly with weights.
• Develop a gradient-based method to train clipping thresholds under a log-domain parameterization.
• Ensure hardware-friendly quantization with symmetric, per-tensor, power-of-2 scale factors.
• Provide an end-to-end framework (Graffitist) for automatic quantization and retraining in TensorFlow.
• Demonstrate near-floating-point accuracy on challenging models (e.g., MobileNets) with limited retraining epochs.

Proposed method

• Define a symmetric affine quantizer with scale s and threshold t, where s = 2^(-f) and f is integer.
• Train thresholds in the log domain to balance range and precision via backpropagation.
• Derive STE-based gradients for input x and scale s, with gradients propagated through the quantizer, and include gradients with respect to log2(t).
• Quantize weights and activations per-tensor with 8-bit and 4-bit schemes, constrained to power-of-2 scales.
• Introduce Graffitist to automate graph transforms, quantization, and retraining in TensorFlow, including fused kernels to reduce memory overhead.

Experimental results

Research questions

• RQ1 Can thresholds for quantization be learned jointly with model weights to optimize end-to-end loss?
• RQ2 Does training threshold parameters in the log domain improve stability and convergence of quantized retraining?
• RQ3 How do power-of-2, per-tensor, symmetric quantization constraints affect accuracy across common CNNs?
• RQ4 Is a unified framework (Graffitist) effective for automatic quantization and retraining across diverse architectures?
• RQ5 What accuracy gains are achievable for difficult networks (e.g., MobileNets) with TQT relative to weight-only retraining?

Key findings

• TQT enables near-floating-point accuracy for 8-bit quantization on several CNNs after retraining.
• TQT can match FP32 accuracy on difficult networks like MobileNets under per-tensor, symmetric, power-of-2 scaling.
• For INT4, weight+threshold retraining is necessary to recover accuracy, with some networks approaching FP32 performance under constrained quantization.
• Retraining thresholds often yield up to ~4% absolute gains in top-1 accuracy over weight-only retraining on challenging models.
• Static quantization generally underperforms retraining-based approaches due to misalignment with global loss optimization.
• Thresholds trained in log-domain improve stability and convergence compared to direct scale-factor training.