Source-linked AI summary
Network Quantization with Element-wise Gradient Scaling
Junghyup Lee, Dohyung Kim, Bumsub Ham
TL;DR
Network quantization reduces resource costs, but STE-based training ignores discretization errors when propagating gradients through round functions. The paper introduces EWGS, which adaptively scales gradient elements using error and Hessian information, and reports improved stability, accuracy, and ImageNet performance across architectures and bit-widths.
Problem
STE avoids zero gradients but propagates the same gradient without considering discretization errors between discretizer inputs and outputs.
Method
EWGS scales each gradient element according to its sign and discretization error, while adapting scaling factors from Hessian information.
Results
EWGS improves quantized-network stability and accuracy, achieves state-of-the-art results across ImageNet architectures and bit-widths, and boosts other STE-based methods.
Takeaways & Limitations
EWGS is presented as an effective alternative to STE that generalizes across architectures, bit-widths, and other quantization methods.
Takeaways & Limitations
The Hessian approximation assumes the main diagonal dominates and that values from the same layer quantizer influence the loss similarly.
Abstract
from arXiv · showhide
Network quantization aims at reducing bit-widths of weights and/or activations, particularly important for implementing deep neural networks with limited hardware resources. Most methods use the straight-through estimator (STE) to train quantized networks, which avoids a zero-gradient problem by replacing a derivative of a discretizer (i.e., a round function) with that of an identity function. Although quantized networks exploiting the STE have shown decent performance, the STE is sub-optimal in that it simply propagates the same gradient without considering discretization errors between inputs and outputs of the discretizer. In this paper, we propose an element-wise gradient scaling (EWGS), a simple yet effective alternative to the STE, training a quantized network better than the STE in terms of stability and accuracy. Given a gradient of the discretizer output, EWGS adaptively scales up or down each gradient element, and uses the scaled gradient as the one for the discretizer input to train quantized networks via backpropagation. The scaling is performed depending on both the sign of each gradient element and an error between the continuous input and discrete output of the discretizer. We adjust a scaling factor adaptively using Hessian information of a network. We show extensive experimental results on the image classification datasets, including CIFAR-10 and ImageNet, with diverse network architectures under a wide range of bit-width settings, demonstrating the effectiveness of our method.
1. Introduction
Network quantization reduces memory and computation but makes training difficult because discretization disrupts gradients. EWGS addresses this by scaling gradient elements according to their signs and discretization errors, with Hessian-based adaptive factors, and achieves strong ImageNet results.
- 1. Introduction: Network quantization lowers weight or activation precision to reduce memory and computational costs for CNNs deployed on limited hardware.Binarized networks can use 32× less memory than full-precision counterparts and replace multiplication and addition with XNOR and bitcount operations.
- 1. Introduction: Discretization makes quantized-network training difficult because round-function derivatives are zero or infinite, motivating STE-based approximations.STE propagates the output gradient unchanged, without accounting for discretization error.
- 1. Introduction: Hessian information adaptively sets EWGS scaling factors across quantizers without extensive hyperparameter searches, training schedules, or additional modules.The Hessian trace is estimated efficiently with Hutchinson’s method.
- 1. Introduction: EWGS adaptively scales each gradient element using its sign and the discretization error between latent and discrete values.The scaled gradient updates the latent value during backpropagation.
- 1. Introduction: EWGS demonstrates effectiveness across CNN architectures and bit-widths, outperforming the state of the art on ImageNet and improving other STE-based quantization methods.The paper reports gains for methods including DoReFa-Net and PROFIT.
2. Related work
Prior quantization work targets quantizer design, training schedules, soft discretizers, and Hessian-guided mixed precision. These approaches often retain STE or introduce additional scheduling, temperature, regularization, or sensitivity-based mechanisms.
- 2. Related work: Quantization methods learn clipping ranges or non-uniform intervals, while activation-distribution modeling reduces quantization errors.Other formulations restrict weight and activation bit-widths to binary, ternary, or arbitrary levels.
- 2. Related work: Incremental, progressive, auxiliary-module, and PROFIT strategies improve low-precision training but rely on iterative procedures or heuristic scheduling.These methods include gradually quantizing weights, decreasing bit-widths, or freezing learned weights.
- 2. Related work: Soft discretizer methods address STE gradient mismatch with sigmoid or tanh approximations, but large temperatures can cause vanishing or exploding gradients.Their hyperparameters therefore require careful tuning during training.
- 2. Related work: Hessian-based quantization methods use eigenvalues or traces to measure layer sensitivity and allocate mixed precision, whereas EWGS uses Hessian traces to adjust gradient-scaling factors.The distinction is sensitivity-based bit allocation versus adaptive scaling during backpropagation.
3. Approach
EWGS trains quantized networks by scaling each discrete-value gradient according to its sign and discretization error, while adaptively setting the scale with Hessian information.
- 3.1. Quantization with EWGS: The quantizer normalizes and clips full-precision weights or activations, rounds them into b-bit discrete values, and applies separate interval parameters per quantized layer.Quantized activations are restricted to non-negative values, with an additional learned output-scale parameter α for each layer.
- 3.1. Quantization with EWGS: EWGS computes latent-value gradients by scaling discrete-value gradients using their sign and the discretization error xn − xq.The method reduces or increases each gradient element through the factor 1 + δsign(gxq)(xn − xq).
- 3.1. Quantization with EWGS: δ = 0 makes EWGS identical to the straight-through estimator, while positive δ enables element-wise correction of gradient mismatch.EWGS therefore contains STE as a special case and adjusts gradients only when the scaling factor is nonzero.
- 3.1. Quantization with EWGS: Figure 2 shows EWGS scaling gradients up when latent values require larger updates than discrete values and down in the opposite case.When latent and discrete values coincide, EWGS propagates the same gradient as STE.
- 3.2. Scaling factor for EWGS: The scaling-factor derivation approximates second-order loss effects because exact Hessian computation is computationally demanding and assumes diagonal Hessian terms dominate.The method also assumes discretization errors are sufficiently small for the local approximation.
- 3.2. Scaling factor for EWGS: EWGS estimates Hessian traces with Hutchinson’s method and sets layer-specific scaling factors using the trace, Hessian size, and a gradient representative.The representative is chosen as 3σ(Gxq), and scaling factors are updated periodically for each weight and activation quantizer.
4. Experiments
Experiments on CIFAR-10 and ImageNet evaluate EWGS across architectures and bit-widths, showing strong quantization accuracy and stable, layer-adaptive scaling behavior.
- 4.2. Results: EWGS outperforms prior methods across ImageNet bit-width settings, matching full-precision performance with 3-bit representations in the ResNet-18 experiments.The 4-bit EWGS result is slightly below LSQ+, while its gain relative to full precision is comparable.
- 4.2. Results: EWGS outperforms the state of the art across ResNet-34 bit-width settings, with 3-bit weights and activations causing no performance degradation relative to full precision.The comparison includes relative performance drops or gains despite LSQ using a different network structure.
- 4.2. Results: On MobileNet-V2, EWGS outperforms PACT and DSQ and remains comparable to PROFIT while using simple gradient scaling instead of PROFIT’s extensive training heuristics.The method also applies across various architectures and bit-widths, whereas PROFIT is described as effective for lightweight networks only.
- 4.3. Discussion: Scaling factors remain bounded during training, while weight factors tend to decrease and activation factors increase in deeper layers, supporting layer-specific adaptation.The exceptions are convolutional layers with 1 × 1 filters and stride 2, which reduce residual sizes in residual blocks.
- 4.3. Discussion: With fixed scaling factors on binarized ResNet-20, EWGS reaches 85.3% versus 84.7% for STE, but poor factor choices reduce performance or make it similar to STE.The best reported configuration reaches 85.6%, compared with 85.3% for the fixed-factor result.
- 4.3. Discussion: Across architectures, EWGS yields about 1% accuracy gains over STE consistently and produces lower losses and higher accuracies during binarized ResNet-18 training.The improvement is especially significant for weight-only quantization.
5. Conclusion
EWGS adaptively adjusts gradients and scaling factors for each layer, and quantized CNNs trained with it achieve state-of-the-art results across bit-widths.
- EWGS-adjusted CNNs achieve state-of-the-art results across a wide range of bit-widths.The method also improves other STE-based quantization methods without additional modules or training embellishments.