Understanding Straight-Through Estimator in Training Activation Quantized Neural Nets

Abstract

Training activation quantized neural networks involves minimizing a piecewise constant function whose gradient vanishes almost everywhere, which is undesirable for the standard back-propagation or chain rule. An empirical way around this issue is to use a straight-through estimator (STE) (Bengio et al., 2013) in the backward pass only, so that the "gradient" through the modified chain rule becomes non-trivial. Since this unusual "gradient" is certainly not the gradient of loss function, the following question arises: why searching in its negative direction minimizes the training loss? In this paper, we provide the theoretical justification of the concept of STE by answering this question. We consider the problem of learning a two-linear-layer network with binarized ReLU activation and Gaussian input data. We shall refer to the unusual "gradient" given by the STE-modifed chain rule as coarse gradient. The choice of STE is not unique. We prove that if the STE is properly chosen, the expected coarse gradient correlates positively with the population gradient (not available for the training), and its negation is a descent direction for minimizing the population loss. We further show the associated coarse gradient descent algorithm converges to a critical point of the population loss minimization problem. Moreover, we show that a poor choice of STE leads to instability of the training algorithm near certain local minima, which is verified with CIFAR-10 experiments.

in International Conference on Learning Representations (ICLR) 2019

References in corpus (20)

• Very Deep Convolutional Networks for Large-Scale Image Recognition
• Faster R-CNN: Towards Real-Time Object Detection with Region Proposal Networks
• Binarized Neural Networks: Training Deep Neural Networks with Weights and Activations Constrained to +1 or -1
• Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation
• DoReFa-Net: Training Low Bitwidth Convolutional Neural Networks with Low Bitwidth Gradients
• Quantized Neural Networks: Training Neural Networks with Low Precision Weights and Activations
• Obfuscated Gradients Give a False Sense of Security: Circumventing Defenses to Adversarial Examples
• Ternary Weight Networks
• Trained Ternary Quantization
• PACT: Parameterized Clipping Activation for Quantized Neural Networks
• Recovery Guarantees for One-hidden-layer Neural Networks
• Loss-aware Weight Quantization of Deep Networks
• An Analytical Formula of Population Gradient for two-layered ReLU network and its Applications in Convergence and Critical Point Analysis
• An Optimal Control Approach to Deep Learning and Applications to Discrete-Weight Neural Networks
• Learning One Convolutional Layer with Overlapping Patches
• Deep Neural Nets with Interpolating Function as Output Activation
• Deep Learning as a Mixed Convex-Combinatorial Optimization Problem
• Quantization and Training of Low Bit-Width Convolutional Neural Networks for Object Detection
• Blended Coarse Gradient Descent for Full Quantization of Deep Neural Networks
• BinaryRelax: A Relaxation Approach For Training Deep Neural Networks With Quantized Weights

Cited by in corpus (28)

• Neural Network Approximation: Three Hidden Layers Are Enough
• Chasing Sparsity in Vision Transformers: An End-to-End Exploration
• Learning N:M Fine-grained Structured Sparse Neural Networks From Scratch
• Deep Network with Approximation Error Being Reciprocal of Width to Power of Square Root of Depth
• Improving Inference for Neural Image Compression
• Mixed Precision DNNs: All you need is a good parametrization
• Explaining Question Answering Models through Text Generation
• Iteratively Training Look-Up Tables for Network Quantization
• Training Binary Neural Networks using the Bayesian Learning Rule
• Training Deep Neural Networks with Joint Quantization and Pruning of Weights and Activations
• The Bayesian Learning Rule
• On Robust Cross-View Consistency in Self-Supervised Monocular Depth Estimation
• Taming Binarized Neural Networks and Mixed-Integer Programs
• A Bop and Beyond: A Second Order Optimizer for Binarized Neural Networks
• Mixed-Precision Quantized Neural Network with Progressively Decreasing Bitwidth For Image Classification and Object Detection
• A Mean Field Theory of Quantized Deep Networks: The Quantization-Depth Trade-Off
• Alternating Direction Method of Multipliers for Quantization
• Reinforcement Learning with Feedback-modulated TD-STDP
• Path Sample-Analytic Gradient Estimators for Stochastic Binary Networks
• Differentiable Learning-to-Group Channels via Groupable Convolutional Neural Networks
• GDRQ: Group-based Distribution Reshaping for Quantization
• Integer-Only Neural Network Quantization Scheme Based on Shift-Batch-Normalization
• Scalable Approximate Inference and Some Applications
• End to End Binarized Neural Networks for Text Classification
• VC dimension of partially quantized neural networks in the overparametrized regime
• Maximin Optimization for Binary Regression
• QuantNet: Learning to Quantize by Learning within Fully Differentiable Framework
• Learning Pruned Structure and Weights Simultaneously from Scratch: an Attention based Approach