ProxQuant: Quantized Neural Networks via Proximal Operators

TL;DR

ProxQuant introduces a principled method for training quantized neural networks using proximal operators, outperforming traditional straight-through gradient methods in stability and effectiveness, especially for binary quantization.

Contribution

It formulates quantized network training as a regularized optimization problem and applies the prox-gradient method, providing a more stable and theoretically grounded alternative to straight-through gradients.

Findings

ProxQuant outperforms state-of-the-art on binary quantization of ResNets and LSTMs.

ProxQuant achieves comparable results to state-of-the-art on multi-bit quantization.

The method is more stable than traditional straight-through gradient approaches.

Abstract

To make deep neural networks feasible in resource-constrained environments (such as mobile devices), it is beneficial to quantize models by using low-precision weights. One common technique for quantizing neural networks is the straight-through gradient method, which enables back-propagation through the quantization mapping. Despite its empirical success, little is understood about why the straight-through gradient method works. Building upon a novel observation that the straight-through gradient method is in fact identical to the well-known Nesterov's dual-averaging algorithm on a quantization constrained optimization problem, we propose a more principled alternative approach, called ProxQuant, that formulates quantized network training as a regularized learning problem instead and optimizes it via the prox-gradient method. ProxQuant does back-propagation on the underlying full-precision vector and applies an efficient prox-operator in between stochastic gradient steps to encourage quantizedness. For quantizing ResNets and LSTMs, ProxQuant outperforms state-of-the-art results on binary quantization and is on par with state-of-the-art on multi-bit quantization. For binary quantization, our analysis shows both theoretically and experimentally that ProxQuant is more stable than the straight-through gradient method (i.e. BinaryConnect), challenging the indispensability of the straight-through gradient method and providing a powerful alternative.