[Paper Review] Searching for Low-Bit Weights in Quantized Neural Networks

Zhaohui Yang, Wang, Yunhe+5 more|arXiv (Cornell University)|Sep 18, 2020

Advanced Neural Network Applications52 references35 citations

TL;DR

This paper proposes a differentiable weight-search method for training quantized neural networks by representing low-bit weights as a probability distribution over discrete values, enabling end-to-end optimization without gradient estimation of non-differentiable quantization.

ABSTRACT

Quantized neural networks with low-bit weights and activations are attractive for developing AI accelerators. However, the quantization functions used in most conventional quantization methods are non-differentiable, which increases the optimization difficulty of quantized networks. Compared with full-precision parameters (i.e., 32-bit floating numbers), low-bit values are selected from a much smaller set. For example, there are only 16 possibilities in 4-bit space. Thus, we present to regard the discrete weights in an arbitrary quantized neural network as searchable variables, and utilize a differential method to search them accurately. In particular, each weight is represented as a probability distribution over the discrete value set. The probabilities are optimized during training and the values with the highest probability are selected to establish the desired quantized network. Experimental results on benchmarks demonstrate that the proposed method is able to produce quantized neural networks with higher performance over the state-of-the-art methods on both image classification and super-resolution tasks.

Motivation & Objective

• Motivate quantized nets with low-bit weights for memory and compute efficiency.
• Address optimization challenges from non-differentiable quantization functions.
• Propose a differentiable weight-search framework to learn quantized weights end-to-end.
• Reduce quantization Gap between training (continuous) and inference (discrete).
• Demonstrate improved accuracy on image classification and super-resolution benchmarks.

Proposed method

• Represent each weight as a probability distribution over a set of low-bit values.
• Use a continuous auxiliary tensor A to learn the distribution with a softmax over values (Eq. 5).
• Compute an expected continuous weight W_c as W_c = sum_i P_i v_i (Eq. 6) for training, and select the max-prob value at inference (Eq. 7).
• Gradients are computed w.r.t. A, avoiding gradient estimation for quantization functions.
• Apply a gradually decreasing temperature to sharpen distributions and reduce quantization gap (Temperature Limit Theorem).
• Introduce State Batch Normalization (SBN) to align statistics between continuous training outputs and discrete inference outputs.

Experimental results

Research questions

• RQ1 Can a differentiable relaxation over discrete low-bit weight values enable end-to-end optimization of quantized networks without STE-based gradient estimation?
• RQ2 Does gradually lowering temperature and using state batch normalization reduce the quantization gap and improve accuracy on classification and super-resolution tasks?
• RQ3 How does the proposed SLB method perform across different bit-widths and architectures compared to state-of-the-art quantization methods?
• RQ4 Is the learned weight distribution effective in both CNNs for image classification and super-resolution models?

Key findings

• Achieves higher accuracy than several state-of-the-art quantization methods across multiple bit-width settings on CIFAR-10 with ResNet and VGG-Small backbones.
• On ImageNet (ILSVRC2012) with ResNet18, SLB with 1-bit weights and activations reaches Top-1 61.3% and Top-5 83.1%, outperforming several baselines.
• SLB consistently improves results for 1/1, 1/2, 1/4, 1/8, 1/32 bit-width configurations and scales to 2/2, 2/4, 2/8, 2/32, 4/4, 4/8, 4/32.
• Super-resolution experiments show SLB achieving higher PSNR than DoReFa at scale factors 2 and 3, with results closer to full-precision.
• Ablation shows State Batch Normalization contributes to preserving performance when using discrete weights.