Learned Step Size Quantization

Steven K. Esser; Jeffrey L. McKinstry; Deepika Bablani; Rathinakumar Appuswamy; Dharmendra S. Modha

Source-linked AI summary

Learned Step Size Quantization

Steven K. Esser, Jeffrey L. McKinstry, Deepika Bablani, Rathinakumar Appuswamy, Dharmendra S. Modha

arXiv:1902.08153v3Feb 21, 2019cs.LGstat.ML

TL;DR

Low-precision inference can reduce computation and memory, but maintaining accuracy requires learning effective quantizer mappings. LSQ learns layer-specific step sizes with transition-sensitive, rescaled gradients, achieving leading ImageNet accuracy across 2-, 3-, and 4-bit models and full-precision accuracy for 3-bit networks.

Problem
Low-precision networks need quantizer mappings that preserve task accuracy, but optimally learning those mappings is unresolved because quantizers are discontinuous.
Method
LSQ makes each weight and activation layer’s quantizer step size learnable using a transition-sensitive loss gradient and a layer- and precision-dependent gradient scale.
Results
LSQ achieved higher ImageNet accuracy than prior approaches across 2-, 3-, and 4-bit networks, with 3-bit networks reaching full-precision accuracy.
Takeaways & Limitations
LSQ provides a simple training modification that supports multiple precisions and requires one additional parameter per weight or activation layer.
Takeaways & Limitations
2-bit networks remained several percentage points below their full-precision counterparts, and it is unclear whether full-precision accuracy is achievable at 2 bits.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Low-precision networks reduce computation and memory demands, but configuring quantizers to preserve task accuracy remains unresolved. LSQ learns quantizer step sizes using a transition-sensitive gradient and achieves strong ImageNet results, including 3-bit models matching full-precision accuracy.

Motivation: Low-precision weights and activations can reduce computation and memory while preserving the accuracy needed by deployed deep networks.The broader system objective also includes throughput, energy efficiency, and compactness.
Open problem: Quantization maps real-valued data to discrete representation levels, but the optimal task-performance mapping remains an open question.Existing approaches use fixed, data-fitted, or task-loss-driven mappings.
Open problem: Fixed mappings lack performance guarantees, while minimizing quantization error may produce a mapping that is suboptimal for task error.Learning the mapping through task loss is appealing but requires approximating the gradient of a discontinuous quantizer.
Contribution: LSQ learns each layer’s quantizer step size with a transition-sensitive gradient and balances step-size updates against weight updates.The method applies to both activations and weights and works with existing backpropagation and stochastic-gradient-descent training.

2 METHODS

LSQ quantizes layer inputs using learnable step sizes and differentiates through the quantizer with a straight-through approximation. It additionally rescales step-size gradients to balance updates across layer sizes and precisions during training.

Quantizer: The quantizer clips and rounds scaled values into integer levels, then rescales the integer representation by the learnable step size.Unsigned activations use QN = 0, while signed weights use separate positive and negative level counts.
Step-size gradient: LSQ introduces a step-size gradient whose interior value depends on the distance between the scaled input and its rounded quantization state.At clipping boundaries, the gradient uses the corresponding positive or negative quantization limit.
Step-size gradient: The gradient uses a straight-through estimator for rounding while differentiating the remaining clipping and scaling operations normally.This keeps rounding in the forward computation while approximating its derivative during backpropagation.
Step-size gradient: Unlike related approximations, LSQ makes the gradient sensitive to proximity to transitions between quantized states.Inputs nearer a transition require smaller step-size changes to switch bins and can produce larger output jumps.
Gradient scale: Step-size gradients are scaled as g = 1/√(NW QP) for weights and g = 1/√(NF QP) for activations.The scaling compensates for differences in layer size and precision and improves trained accuracy.
Training: Training stores and updates full-precision weights while using quantized weights and activations in forward and backward passes.The experiments use 2-, 3-, 4-, or 8-bit inputs for most matrix-multiplication layers, with first and last layers kept at 8 bits.

3 RESULTS

LSQ improves low-precision network accuracy across architectures and precision levels while revealing important trade-offs in regularization, model size, gradient scaling, and quantization objectives. Knowledge distillation further improves LSQ and enables 3-bit networks to reach full-precision accuracy.

3.1 WEIGHT DECAY: Lower precision required less weight decay: halving it for 3-bit networks and quartering it for 2-bit networks improved ResNet-18 performance.These values were used in subsequent experiments.
3.2 COMPARISON WITH OTHER APPROACHES: LSQ achieved higher top-1 accuracy than previous approaches for 2-, 3-, and 4-bit networks across the considered architectures.It also generally achieved the best top-5 and 8-bit accuracy, while 4-bit models usually offered no advantage over 8-bit models.
3.2 COMPARISON WITH OTHER APPROACHES: 2-bit accuracy dropped 2.9 points for ResNet-18 but 14.0 points for SqueezeNext-23-2x relative to full precision.The authors interpret SqueezeNext’s larger drop as possible sensitivity associated with its parameter-efficient design.
3.3 ACCURACY VS. MODEL SIZE: 2-bit ResNet-34 and ResNet-50 offered better accuracy on the model-size frontier than smaller higher-precision networks.The figure compares accuracy with model size, and the caption notes that some 2-bit networks provide the highest accuracy at a given size.
3.4 STEP SIZE GRADIENT SCALE IMPACT: Gradient scaling by 1/√N_W QP largely removed precision-dependent imbalance between step-size and weight updates.Without scaling, step-size updates were 2 to 3 orders of magnitude larger than weight updates; scaling by layer size alone removed much of the imbalance but not its precision dependence.
3.5 COSINE LEARNING RATE DECAY IMPACT: 67.2 top-1 accuracy was achieved by 2-bit ResNet-18 with step-based decay, 0.4 below cosine decay but 1.5 above the next-best method.Cosine decay was selected because it avoids schedule hyperparameter selection without increasing training time.
3.6 QUANTIZATION ERROR: LSQ step sizes differed substantially from quantization-error-minimizing values, with absolute differences reaching 64% across tested metrics and layers.For 2-bit ResNet-18, LSQ learned activation and weight step sizes of 0.949 ± 0.206 and 0.025 ± 0.019, respectively, suggesting task-error optimization differs from distribution fitting.
3.7 IMPROVEMENT WITH KNOWLEDGE-DISTILLATION: Knowledge distillation improved LSQ by up to 1.1 top-1 points for 3-bit ResNet-50, with 3-bit networks reaching full-precision baseline accuracy.The teacher was a frozen full-precision model with the same architecture as the low-precision student.

4 CONCLUSIONS

LSQ exceeds prior quantization approaches on ImageNet across several architectures and supports useful low-precision models, including 2-bit models under tight size limits. However, achieving full-precision accuracy remains unresolved for 2-bit networks.

LSQ exceeds all prior approaches on ImageNet across several network architectures.
2-bit networks reached accuracy several percent below their full precision counterparts.
The work continues a trend toward good ImageNet performance with steadily fewer precision bits across network architectures.
Reducing precision while maintaining high accuracy can reduce model size and increase throughput in deployed deep networks.

A STEP SIZE GRADIENT SCALE DERIVATION

The gradient-scale derivation balances step-size updates against weight or activation updates by accounting for layer size and quantization precision. It yields scale factors that decrease with the square root of the relevant layer dimension and available positive quantization states.

The derivation assumes l2-norm growth with the square root of element count and step sizes proportional to 1/QP at larger precisions.These assumptions place quantization precision and layer size into the step-size scaling heuristic.
Weight and step-size gradients are expected to have approximately the same order after applying the chain rule and treating per-weight loss gradients as uncorrelated zero-centered variables.
The weight gradient scale factor is set to g = 1/√NW QP to counter expected update imbalance.NW denotes the number of weights, while QP is the number of positive quantization states.
The scale factor is designed to balance step-size update magnitudes with parameter updates across layers.
Activation step-size scaling uses g = 1/√NF QP because update imbalance grows with feature count NF and quantization precision QP.The argument assumes batch-normalization scaling updates primarily drive changes in pre-quantization activations.

B IMPLEMENTATION

LSQ implementation uses automatic differentiation with custom gradient behavior for gradient scaling and rounding, then quantizes weights and activations before convolutional or fully connected layers. The pseudocode favors simplicity and broad framework applicability over maximal efficiency.

The implementation relies on automatic differentiation and a detach operation whose forward value is unchanged but whose backward gradient is blocked.
The pseudocode prioritizes simplicity and broad applicability, although more compute- and memory-efficient implementations are possible.
Gradient scaling returns the original forward value while multiplying its backward gradient by a specified scale.
The rounding function uses nearest-integer rounding in the forward pass while passing the unrounded input gradient backward.
Activation quantization uses Qn = 0 and Qp = 2^p - 1, whereas weight quantization uses signed ranges Qn = -2^(p-1) and Qp = 2^(p-1) - 1.
The quantize function applies a learnable step size, clipping, rounding, and rescaling to produce quantized values for layer inputs.Weights and activations are quantized before convolutional or fully connected layers.