Mirror Descent View for Neural Network Quantization

Abstract

Quantizing large Neural Networks (NN) while maintaining the performance is highly desirable for resource-limited devices due to reduced memory and time complexity. It is usually formulated as a constrained optimization problem and optimized via a modified version of gradient descent. In this work, by interpreting the continuous parameters (unconstrained) as the dual of the quantized ones, we introduce a Mirror Descent (MD) framework for NN quantization. Specifically, we provide conditions on the projections (i.e., mapping from continuous to quantized ones) which would enable us to derive valid mirror maps and in turn the respective MD updates. Furthermore, we present a numerically stable implementation of MD that requires storing an additional set of auxiliary variables (unconstrained), and show that it is strikingly analogous to the Straight Through Estimator (STE) based method which is typically viewed as a "trick" to avoid vanishing gradients issue. Our experiments on CIFAR-10/100, TinyImageNet, and ImageNet classification datasets with VGG-16, ResNet-18, and MobileNetV2 architectures show that our MD variants obtain quantized networks with state-of-the-art performance. Code is available at https://github.com/kartikgupta-at-anu/md-bnn.

arXiv commentThis paper was accepted at AISTATS 2021

References in corpus (19)

• Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour
• DoReFa-Net: Training Low Bitwidth Convolutional Neural Networks with Low Bitwidth Gradients
• Quantized Neural Networks: Training Neural Networks with Low Precision Weights and Activations
• Binarized Neural Networks
• SNIP: Single-shot Network Pruning based on Connection Sensitivity
• Bayesian Compression for Deep Learning
• A Survey on Methods and Theories of Quantized Neural Networks
• Loss-aware Binarization of Deep Networks
• Snapshot Ensembles: Train 1, get M for free
• Extremely Low Bit Neural Network: Squeeze the Last Bit Out with ADMM
• ProxQuant: Quantized Neural Networks via Proximal Operators
• Training Binary Neural Networks with Real-to-Binary Convolutions
• Latent Weights Do Not Exist: Rethinking Binarized Neural Network Optimization
• On the Convergence Rate of Stochastic Mirror Descent for Nonsmooth Nonconvex Optimization
• Training Binary Neural Networks through Learning with Noisy Supervision
• Progressive DNN Compression: A Key to Achieve Ultra-High Weight Pruning and Quantization Rates using ADMM
• Mirror Descent View for Neural Network Quantization
• GBCNs: Genetic Binary Convolutional Networks for Enhancing the Performance of 1-bit DCNNs

Cited by in corpus (5)

• Rotated Binary Neural Network
• Training Binary Neural Networks using the Bayesian Learning Rule
• Mirror Descent View for Neural Network Quantization
• AskewSGD : An Annealed interval-constrained Optimisation method to train Quantized Neural Networks
• ReCU: Reviving the Dead Weights in Binary Neural Networks