Deep networks run with low precision operations at inference time offer power and space advantages over high precision alternatives, but need to overcome the challenge of maintaining high accuracy as precision decreases. Here, we present a method for training such networks, Learned Step Size Quantization, that achieves the highest accuracy to date on the ImageNet dataset when using models, from a variety of architectures, with weights and activations quantized to 2-, 3-or 4-bits of precision, and that can train 3-bit models that reach full precision baseline accuracy. Our approach builds upon existing methods for learning weights in quantized networks by improving how the quantizer itself is configured. Specifically, we introduce a novel means to estimate and scale the task loss gradient at each weight and activation layer's quantizer step size, such that it can be learned in conjunction with other network parameters. This approach works using different levels of precision as needed for a given system and requires only a simple modification of existing training code.
The motivation of the paper is to be able to train low precision networks to a high-accuracy. Quantization is a useful tool in model compression, and doing it well for very low-precision models (2-3 bit precision specifically), is challenging.
The main contribution of the paper comes from:
a) Step Size Gradient: They propose a gradient which is sensitive to the distance between the value and the transition point. This is different from other methods which have gradients dependent only on the clip point.
b) Step Size Gradient Scale: This is an interesting contribution, where they try to match the ratio of average update of the step size ‘s’ and average magnitude of ‘s’, with that of the network weights. This leads them to scale the gradient according to the precision and number of parameters. They demonstrate that this scaling actually helps improve the accuracy.
The results for 8-bit precision are not new. Several results (Quantization and Training of Neural Networks for Efficient
Integer-Arithmetic-Only Inference, Jacob et al., Quantizing deep convolutional networks for
efficient inference: A whitepaper, Krishnamoorthi et al.), show 8-bit quantization results where the accuracy matches floating point accuracy, and in some case exceeds it (low precision quantization acting as a regularizer). However, the results for lower precision are impressive.
There are a few questions:
1. In sec 2.1, you mention that ‘each layer of weights and activations has a distinct step size, represented as an fp32 value, initialized to …’. Can you explain the intuition behind the initial value of the step size, and how is it a function of v?
2. ‘Model Compression via Distillation and Quantization’ (Polino et al.) shows distillation actually helps significantly improve accuracy. I wonder if the authors have tried different weight combinations for the distillation loss, and using bigger models as teacher models.
3. I would like to get more details of the inference setup, specifically the size and inference latency improvements over full-precision networks. The practical applicability of low-precision networks, specifically 2-bit and 3-bit networks, equally depends on the inference infrastructure, as it does on the training improvements.
4. Have you evaluated your method for a non-Vision usecase?
Overall this is a good work, I would tend towards accepting this.