Cluster-Promoting Quantization with Bit-Drop for Minimizing Network Quantization Loss

TL;DR

This paper introduces Cluster-Promoting Quantization (CPQ), a novel neural network quantization method that optimizes quantization grids and encourages weights to cluster around them, reducing quantization errors and improving performance.

Contribution

The work proposes a differentiable quantization approach with multi-class STE and DropBits, enabling learning of optimal and heterogeneous quantization levels during training.

Findings

CPQ outperforms fixed-level quantization on benchmarks.

Learning heterogeneous levels yields better accuracy.

DropBits effectively regularizes bit allocation across layers.

Abstract

Network quantization, which aims to reduce the bit-lengths of the network weights and activations, has emerged for their deployments to resource-limited devices. Although recent studies have successfully discretized a full-precision network, they still incur large quantization errors after training, thus giving rise to a significant performance gap between a full-precision network and its quantized counterpart. In this work, we propose a novel quantization method for neural networks, Cluster-Promoting Quantization (CPQ) that finds the optimal quantization grids while naturally encouraging the underlying full-precision weights to gather around those quantization grids cohesively during training. This property of CPQ is thanks to our two main ingredients that enable differentiable quantization: i) the use of the categorical distribution designed by a specific probabilistic parametrization in the forward pass and ii) our proposed multi-class straight-through estimator (STE) in the backward pass. Since our second component, multi-class STE, is intrinsically biased, we additionally propose a new bit-drop technique, DropBits, that revises the standard dropout regularization to randomly drop bits instead of neurons. As a natural extension of DropBits, we further introduce the way of learning heterogeneous quantization levels to find proper bit-length for each layer by imposing an additional regularization on DropBits. We experimentally validate our method on various benchmark datasets and network architectures, and also support a new hypothesis for quantization: learning heterogeneous quantization levels outperforms the case using the same but fixed quantization levels from scratch.