Nonuniform-to-Uniform Quantization: Towards Accurate Quantization via Generalized Straight-Through Estimation

Zechun Liu; Kwang-Ting Cheng; Dong Huang; Eric Xing; Zhiqiang Shen

Source-linked AI summary

Nonuniform-to-Uniform Quantization: Towards Accurate Quantization via Generalized Straight-Through Estimation

Zechun Liu, Kwang-Ting Cheng, Dong Huang, Eric Xing, Zhiqiang Shen

arXiv:2111.14826v2Nov 29, 2021cs.CVcs.AIcs.LG

TL;DR

Low-bit quantization faces an accuracy gap because fixed quantization levels poorly represent real-valued weights and activations, while nonuniform methods can incur hardware overhead. N2UQ learns flexible input thresholds with uniform outputs, using G-STE and entropy-preserving regularization, and outperforms state-of-the-art nonuniform methods by 0.5∼1.7% accuracy on ImageNet.

Problem
Low-bit quantization can lose accuracy because fixed quantization levels inadequately represent real-valued weights and activations, while nonuniform outputs require additional hardware processing.
Method
N2UQ learns input thresholds while producing uniformly quantized outputs, with G-STE for threshold gradients and entropy-preserving regularization for weights.
Results
N2UQ outperforms state-of-the-art nonuniform quantization methods by 0.5∼1.7% accuracy on ImageNet.
Takeaways & Limitations
N2UQ combines the representational flexibility of threshold learning with uniform outputs that support efficient bitwise operations.

Abstract

from arXiv · show

The nonuniform quantization strategy for compressing neural networks usually achieves better performance than its counterpart, i.e., uniform strategy, due to its superior representational capacity. However, many nonuniform quantization methods overlook the complicated projection process in implementing the nonuniformly quantized weights/activations, which incurs non-negligible time and space overhead in hardware deployment. In this study, we propose Nonuniform-to-Uniform Quantization (N2UQ), a method that can maintain the strong representation ability of nonuniform methods while being hardware-friendly and efficient as the uniform quantization for model inference. We achieve this through learning the flexible in-equidistant input thresholds to better fit the underlying distribution while quantizing these real-valued inputs into equidistant output levels. To train the quantized network with learnable input thresholds, we introduce a generalized straight-through estimator (G-STE) for intractable backward derivative calculation w.r.t. threshold parameters. Additionally, we consider entropy preserving regularization to further reduce information loss in weight quantization. Even under this adverse constraint of imposing uniformly quantized weights and activations, our N2UQ outperforms state-of-the-art nonuniform quantization methods by 0.5~1.7 on ImageNet, demonstrating the contribution of N2UQ design. Code and models are available at: https://github.com/liuzechun/Nonuniform-to-Uniform-Quantization.

1. Introduction

N2UQ addresses the accuracy–hardware trade-off in low-bit quantization by learning flexible input thresholds while retaining uniformly quantized outputs. It introduces G-STE for threshold optimization and entropy-based weight regularization, achieving higher ImageNet accuracy than prior methods.

Motivation: Quantization reduces model size and accelerates inference, but extremely low-bit models can remain substantially less accurate than full-precision networks.A 2-bit uniformly quantized ResNet-50 reaches 67.1% top-1 accuracy on ImageNet, 9.9% below its real-valued counterpart.
Motivation: Nonuniform quantization better fits value distributions, but floating-point outputs require extra mapping or lookup-table operations before binary computation.These post-processing steps increase hardware area and energy consumption.
N2UQ: N2UQ learns in-equidistant input thresholds while producing equidistant output levels, combining distributional flexibility with hardware-friendly bitwise operations.Uniform outputs avoid additional post-processing between quantization and matrix multiplication.
N2UQ: G-STE approximates the backward pass for threshold learning by encoding stochastic-quantization expectations into the deterministic quantizer’s gradients.Unlike standard STE, it adapts to nonuniform input thresholds and provides a finer-grained approximation.
N2UQ: Entropy-based weight regularization further reduces information loss from weight quantization by considering the overall weight distribution.The paper evaluates threshold learning, G-STE, and weight regularization across ImageNet architectures and bitwidth constraints.
Results: 0.5∼1.7% higher accuracy on ImageNet is reported for N2UQ than state-of-the-art nonuniform quantization methods.A 2-bit ResNet-50 reaches 76.4% top-1 accuracy, leaving a 0.6% gap to its real-valued counterpart.

2. Related Work

Related work frames quantization as one model-compression strategy and distinguishes uniform from nonuniform approaches. Nonuniform methods can improve distribution fitting, while STE provides a widely used but limited solution for training discontinuous quantizers.

Model compression: Model compression supports deployment on devices with limited storage and computational power, using approaches such as quantization, pruning, distillation, and compact network design.This work focuses primarily on quantization.
Quantization methods: Nonuniform quantization can improve accuracy by allocating more quantization levels to important regions of the underlying value distribution.Examples include powers-of-two and additive powers-of-two quantization.
Training quantizers: STE addresses the near-zero gradients of discontinuous quantization functions but cannot incorporate quantizer input and output nonlinearities needed to fit varied distributions.The related-work discussion motivates more flexible gradient approximations.

3. Method

N2UQ learns input thresholds while keeping quantized outputs uniformly spaced, combining distribution flexibility with hardware-friendly bitwise inference. G-STE enables threshold learning by deriving backward approximations from stochastic quantization, while weight quantization uses equidistant thresholds.

3.1. Preliminaries: Nonuniform quantizers can require additional operations or lookup tables because their floating-point outputs cannot be directly mapped to fixed-bit binary representations.These extra encoding steps increase hardware area and energy consumption.
3.2.1 Forward Pass: Threshold Learning Quantization: N2UQ learns flexible input thresholds but produces uniformly quantized weights and activations for efficient bitwise operations without post-processing.The thresholds divide real-valued inputs into nonuniform ranges mapped to fixed output levels.
3.2.2 Backward Pass: Generalized Straight-Through Estimator (G-STE): G-STE addresses the zero-input-gradient and intractable-threshold-gradient problems that make N2UQ difficult to optimize.It replaces the unavailable threshold derivatives with a backward approximation derived from stochastic quantization.
3.2.2 Backward Pass: Generalized Straight-Through Estimator (G-STE): In the forward pass, deterministic quantization is obtained by setting the stochastic probability threshold to p = 0.5, avoiding random-seed generation.The remaining quantization segments are combined into the n-bit quantizer.
3.2.2 Backward Pass: Generalized Straight-Through Estimator (G-STE): G-STE represents quantization as a sum of binarization segments and encodes stochastic-quantization expectations in the backward pass.For equal-sized input intervals, it degenerates to the conventional STE; with nonuniform thresholds, it adapts the slopes.
3.2.3 Entropy Preserving Weight Regularization: For weights, N2UQ adopts equidistant thresholds because static weight parameters can be adjusted so quantized weights are approximately uniformly distributed across levels.This differs from activations, whose distributions vary across batches.

4. Experiments

Experiments on ImageNet evaluate N2UQ across architectures, bitwidths, state-of-the-art comparisons, ablations, and distribution visualizations. N2UQ improves accuracy through threshold learning with G-STE and entropy-based weight regularization while retaining hardware-friendly uniform outputs.

Experimental Setup: Experiments evaluate N2UQ on ImageNet with different architectures and bitwidth constraints, using real-valued pretrained models as initialization.Models are trained for 128 epochs with Adam and linear learning-rate decay.
State-of-the-Art Comparison: N2UQ surpasses uniform quantization methods on ResNet and achieves up to 1.7% higher accuracy than the nonuniform method LCQ.The method integrates nonlinearity learning inside the quantizer, avoiding discrepancies between nonlinear turning points and quantizer thresholds.
State-of-the-Art Comparison: N2UQ produces uniformly quantized weights and activations, avoiding look-up-table construction and floating-point-to-binary mapping before bitwise operations.This preserves hardware-friendly matrix multiplication while fixing output levels rather than using nonuniform floating-point levels.
State-of-the-Art Comparison: For MobileNetV2, N2UQ reaches a state-of-the-art result and can compress the real-valued network to 4-bit without compromising or improving accuracy.The reported behavior indicates that the learned thresholds fit underlying distributions while the quantized network remains compact.
Ablation Study: On 2-bit ImageNet ResNet-18, weight regularization and threshold learning with G-STE improve baseline accuracy by 1.9% and 3.0%, respectively, while their combination reaches 69.7%.The combined model narrows the accuracy gap to the corresponding real-valued ResNet-18 to 2.1%.
Ablation Study: Entropy-based regularization better reduces information loss than gradient-learned factors or prior weight norm because it uses overall weight-distribution information.Gradient-learned factors introduce instability from noisy gradients, while prior weight norm is not designed for quantized weights.
Visualization: When distributions are dense near zero and sparse in the tails, learned thresholds use smaller dense-region intervals and larger tail intervals to mitigate quantization error.The visualization also contrasts this behavior with a baseline biased by extrema, which maps few filter entries to some 2-bit levels.

5. Conclusions

The conclusion presents N2UQ as a quantizer that combines learnable input thresholds with uniform output levels, supported by G-STE and entropy-preserving weight regularization. It reports higher accuracy than prior nonuniform methods despite uniform outputs.

5. Conclusions: N2UQ learns input thresholds within a uniform quantizer, combining nonuniform methods’ representational capability with hardware-friendly uniform outputs.Uniform outputs avoid additional post-processing before bitwise operations.
5. Conclusions: G-STE addresses intractable threshold gradients by deriving a backward approximation from stochastic quantization.The method converts threshold-gradient calculation into slope-related optimization.
5. Conclusions: Entropy-preserving weight regularization is introduced to further reduce quantization error.The conclusion identifies it as a complementary component of N2UQ.
5. Conclusions: 0.8∼1.7% higher accuracy is achieved than previous state-of-the-art nonuniform methods, while 2-bit N2UQ ResNet-50 is 0.6% shy of its real-valued counterpart.These results are reported under the constraint of uniform output quantized levels.

Appendix

The appendix supplies additional material on entropy-preserving weight regularization, learned-parameter visualization, and experiments retaining real-valued down-sampling layers in ResNet.

Section A: Appendix Section A provides illustrations and further details of entropy-preserving weight regularization.The section concerns the regularization proposed in the main paper.
Section B: Appendix Section B visualizes learned parameters in the quantized networks.It supplements the main text’s analysis of learned distributions and thresholds.
Section C: Appendix Section C reports results for keeping down-sampling layers real-valued in ResNet structures.This isolates a deployment variant involving selected real-valued layers.

A. Entropy Preserving Weight Regularization

The entropy-preserving regularization rescales real-valued weights before quantization so the resulting weights occupy the available levels more evenly. Its factor can be computed offline after training and absorbed into BatchNorm layers when present.

Regularization Mechanism: The regularization rescales real-valued weights so their absolute mean equals 1 before quantization.The rescaling uses |W_r|/||W_r||_l1 as described in the appendix.
Regularization Mechanism: For initially uniform and symmetric weights, the regularization spreads the distribution so quantization approximately uses 2^n levels uniformly.The corresponding quantizer clips, rounds, and maps weights to uniformly spaced output levels.
Training and Deployment: During training, regularization helps distribute nonuniform real-valued weights more evenly, and its optimized factor can later be computed offline.When BatchNorm is used, the factor can be absorbed by the BatchNorm layers after quantization training.

B. Learned Parameters Visualization

The visualization shows that weight regularization reduces repeated 3 × 3 weight patterns, while learned thresholds adapt to activation magnitudes across network layers.

Weight patterns: 4.37% of quantized 3 × 3 weight matrices share identical values with the baseline, versus 1.69% with the proposed regularization.The reduction indicates fewer repeated weight patterns under the proposed method.
Learned thresholds: Learned threshold intervals are often larger and first scaling factors β1 smaller in residual-block layers with larger-magnitude activations.These layers sum residual connections with previous-layer outputs, motivating the learned parameter pattern.

C. Results without Quantizing Downsampling Layers

Keeping down-sampling layers at full precision improves low-bit ResNet accuracy, but its benefit becomes marginal at higher bitwidths.

Accuracy comparison: ∼0.3% accuracy improvement occurs when 1×1 down-sampling layers remain real-valued under lower-bit quantization.The remaining convolutional layers are quantized under this evaluation setting.
Accuracy comparison: The benefit of full-precision down-sampling layers becomes marginal at higher bitwidths.The comparison covers 2-bit, 3-bit, and 4-bit quantization of weights and activations.