Source-linked AI summary
Overcoming Oscillations in Quantization-Aware Training
Markus Nagel, Marios Fournarakis, Yelysei Bondarenko, Tijmen Blankevoort
TL;DR
The paper studies weight oscillations during quantization-aware training, an under-investigated phenomenon that can degrade inference statistics and optimization, especially in low-bit efficient networks. It analyzes their effects and proposes oscillation dampening and iterative weight freezing, achieving state-of-the-art low-bit accuracy on ImageNet.
Problem
Weight oscillations during QAT are under-investigated and can corrupt batch-normalization statistics and add optimization noise, particularly in low-bit efficient networks.
Method
The paper analyzes oscillations in existing QAT methods and proposes oscillation dampening and iterative weight freezing to reduce them during training.
Results
The two proposed methods achieve state-of-the-art accuracy for 3- and 4-bit quantization of efficient networks on ImageNet.
Takeaways & Limitations
Addressing oscillations at their source improves accuracy beyond batch-normalization re-estimation in low-bit efficient-network quantization.
Takeaways & Limitations
Oscillation dampening can increase training time by approximately 33% compared with the LSQ baseline, whereas iterative weight freezing has negligible computational overhead.
Abstract
from arXiv · showhide
When training neural networks with simulated quantization, we observe that quantized weights can, rather unexpectedly, oscillate between two grid-points. The importance of this effect and its impact on quantization-aware training (QAT) are not well-understood or investigated in literature. In this paper, we delve deeper into the phenomenon of weight oscillations and show that it can lead to a significant accuracy degradation due to wrongly estimated batch-normalization statistics during inference and increased noise during training. These effects are particularly pronounced in low-bit ($\leq$ 4-bits) quantization of efficient networks with depth-wise separable layers, such as MobileNets and EfficientNets. In our analysis we investigate several previously proposed QAT algorithms and show that most of these are unable to overcome oscillations. Finally, we propose two novel QAT algorithms to overcome oscillations during training: oscillation dampening and iterative weight freezing. We demonstrate that our algorithms achieve state-of-the-art accuracy for low-bit (3 & 4 bits) weight and activation quantization of efficient architectures, such as MobileNetV2, MobileNetV3, and EfficentNet-lite on ImageNet. Our source code is available at {https://github.com/qualcomm-ai-research/oscillations-qat}.
1. Introduction
The paper identifies quantized-weight oscillations as an under-investigated QAT problem that can harm inference statistics and accuracy, especially in low-bit efficient networks. It proposes dampening and iterative weight freezing to reduce these oscillations and improve low-bit results.
- Background: Quantization compresses weights and activations into low-bit fixed-point representations, reducing power consumption and accelerating inference on edge devices.Quantization can nevertheless add noise because of reduced precision.
- Problem: Quantized weights can oscillate between adjacent quantization levels during QAT with the straight-through estimator, adding detrimental optimization noise.The phenomenon is described as little-known and under-investigated.
- Impact: Corrupted batch-normalization statistics from weight oscillations can lead to poor validation accuracy, particularly in low-bit efficient networks with depth-wise separable layers.Reestimating batch-normalization statistics after training can address this symptom.
- Limitations of existing mitigation: Batch-normalization re-estimation addresses a significant symptom of oscillations but does not address their root cause.
- Contributions: The proposed oscillation dampening and iterative weight freezing algorithms reduce oscillations and improve accuracy beyond batch-normalization re-estimation.Both methods achieve state-of-the-art results for 3- and 4-bit quantization of efficient networks on ImageNet.
2. Oscillations in QAT
Quantization-aware training can make latent weights oscillate between adjacent quantization levels rather than converge, creating optimization noise and unstable inference statistics. The effect is especially consequential in low-bit efficient networks, while dampening and iterative freezing reduce oscillations and improve accuracy.
- Quantization-aware training: Simulated quantization uses quantized weights for inference while floating-point latent weights serve as optimization proxies.The quantizer rounds and clips latent weights using a scaling factor and quantization thresholds.
- Quantization-aware training: STE approximates the unavailable rounding gradient as 1 inside the representable quantization range, enabling gradient-based QAT.The gradient is zero outside the quantization region.
- Oscillation mechanism: Latent weights oscillate across adjacent quantization levels because piecewise-constant gradients push them toward opposite sides of the decision threshold.The oscillations occur irrespective of learning rate; reducing the learning rate lowers amplitude but not frequency, whose behavior depends on distance to the nearest quantization level.
- Oscillations in practice: Oscillations occur in MobileNetV2 depth-wise separable layers, where many 3-bit weights switch between adjacent levels and latent weights accumulate at decision boundaries near convergence.The phenomenon is therefore observed beyond the toy regression example in a larger ImageNet-trained network.
- Effect on batch normalization: Lower bit-widths enlarge quantization-level spacing, so oscillating weights cause larger output-distribution shifts; depth-wise separable layers show especially large batch-normalization KL discrepancies.These shifts can corrupt exponential-moving-average statistics and degrade inference accuracy.
- Effect on training: Batch-normalization re-estimation improves MobileNetV2 accuracy and reduces seed variance, but oscillation freezing achieves higher validation accuracy than end-of-training binary optimization.The results indicate that oscillations can both prevent convergence to the best local minimum and steer optimization toward suboptimal directions earlier in training.
3. Related work
Prior QAT work modifies STE gradients, adds regularization, or freezes layers to improve low-bit training. However, the reviewed methods generally do not directly eliminate weight oscillations.
- Earlier QAT studies mainly examined STE alternatives, but this paper instead targets oscillations as a distinct optimization problem.
- Alternative STE methods use multiplicative or additive adaptations, respectively scaling data-gradients or adding input-independent regularization terms.
- Multiplicative: EWGS and related multiplicative methods scale gradients by latent-weight distance from bin centers, while DSQ uses nonlinear gradient scaling.
- Multiplicative: Neither EWGS nor DSQ overcomes oscillations, although these methods can change their magnitude.
- Additive: Additive approaches apply regularization to make weight distributions more quantization-friendly, including penalties on quantized–floating-point differences or bin-centering terms.
- Other approaches: Progressively freezing whole layers addresses activation instability in low-bit MobileNets, while another gradient-learning approach is limited to binary weights.
4. Overcoming oscillations in QAT
The paper introduces an EMA-based oscillation metric and two targeted QAT methods: dampening latent weights toward bin centers and freezing frequently oscillating weights. These methods address oscillations during training while accounting for quantization scale and optimization behavior.
- The proposed approach first measures oscillations and then prevents them at their source during quantization-aware training.
- Oscillation measurement: An oscillation occurs when a quantized integer value changes and its change direction reverses relative to the previous integer-domain change.
- Oscillation measurement: Oscillation frequency is tracked with an exponential moving average, and a minimum frequency threshold identifies oscillating weights.
- Oscillation dampening: Oscillation dampening regularizes latent weights toward quantization-bin centers, using L = Ltask + λLdampen while clipping weights to the quantization grid.
- Oscillation dampening: Applying bin regularization in the latent-weight domain makes its gradient independent of quantization scale and indirectly independent of bit-width.
- Iterative weight freezing: Iterative weight freezing freezes any weight whose oscillation frequency exceeds fth until training ends, using the integer domain to avoid scale-dependent rounding changes.
- Iterative weight freezing: Freezing can select either oscillating quantized state because unequal residence times make the temporal expectation correspond to the optimal value.
- Iterative weight freezing: The freezing algorithm works with any gradient-based optimizer and is not limited to a particular quantization formulation or gradient estimator.
5. Experiments
Experiments on low-bit quantization of efficient networks show that both proposed methods reduce weight oscillations and improve accuracy, with iterative freezing achieving similar performance at lower computational cost.
- Experimental setup: The experiments target 3- and 4-bit quantization of efficient networks with depth-wise separable convolutions on ImageNet.The study compares competing QAT methods and evaluates both weight-only and weight-plus-activation quantization.
- Oscillation dampening: Increasing dampening strength reduces oscillating weights and narrows the pre-BN versus post-BN accuracy gap, but excessive regularization harms final accuracy.The trade-off indicates that strong dampening can inhibit beneficial weight movement between quantization levels.
- Oscillation dampening: A cosine-annealed dampening schedule improves accuracy almost 1% over post-BN re-estimation and more than 5% over pre-BN re-estimation.The schedule allows freer latent-weight movement early in training and stronger regularization near convergence.
- Oscillation dampening: Dampening clusters latent weights around quantization-bin centres with hardly any weights at the decision boundary.This distribution is shown for MobileNetV2 layer conv.3.1 after training.
- Iterative weight freezing: Iterative freezing reduces remaining oscillations as the threshold decreases, but thresholds that freeze too early reduce final accuracy.Annealing the threshold enables stronger freezing near the end of training; the best configuration leaves 0.04% oscillations versus 1.11% for dampening and improves almost 1% over post-BN re-estimation and more than 5% over pre-BN re-estimation.
- Comparison to other QAT methods: Both proposed methods outperform competing QAT techniques for 3- and 4-bit MobileNetV2 and achieve state-of-the-art results on MobileNetV3-Small and EfficientNet-lite.Across efficient networks, the methods improve more than 1% over the commonly used LSQ baseline.
- Comparison to other QAT methods: Oscillation dampening increases training time by approximately 33% versus LSQ, whereas iterative weight freezing has negligible computational overhead with similar performance.
6. Conclusion
The conclusion identifies weight oscillations as a source of degraded QAT performance, especially in low-bit lightweight networks, and presents two methods that achieve state-of-the-art ImageNet accuracy.
- Weight oscillations adversely affect quantized-network performance by corrupting inference batch-normalization statistics and harming optimization.
- The proposed oscillation dampening and iterative weight freezing methods achieve state-of-the-art ImageNet accuracy across efficient models quantized to low bit-widths.
A.1. Gradients for 1D regression problem
The appendix derives piecewise gradient-descent updates for several QAT methods and explains why multiplicative STE variants cannot stop oscillations while additive dampening can change gradient direction.
- The analysis derives gradient-descent updates for vanilla STE, PSG, EWGS, and oscillation dampening around the boundary between adjacent quantization levels.The boundary is defined relative to the optimal weight and its two closest quantization levels.
- PSG and EWGS apply positive scaling to the STE gradient, changing update magnitude but not direction, so they cannot prevent oscillations.
- Oscillation dampening adds a term with opposite sign to the STE gradient, allowing the update direction to change and preventing oscillations.
A.2. Relation of the distance and frequency
The oscillation frequency depends linearly on the distance between the optimal weight and its closest quantization level, scaled by the quantization-bin size.
- The paper defines oscillation frequency from the distance between the optimal weight and its closest quantization level.The quantization scaling factor s represents the size of the quantization bin, and the distance is d = |q(w*) − w*|.
- Experiments vary the distance d and observe the resulting oscillation frequency to test this linear dependency.
A.3. Different learning rates and multiplicative regularization methods
Learning-rate changes affect oscillation amplitude but not frequency, while multiplicative gradient scaling is therefore expected to reduce amplitude rather than eliminate oscillations. This analysis is most applicable when oscillation amplitudes are small, typically late in training.
- Different learning rates: Changing the learning rate influences oscillation amplitude but not oscillation frequency in the toy regression example.The frequency is independent of learning rate, whereas the amplitude changes across learning rates.
- Different learning rates: The distance between the nearest quantization point and the optimal value directly corresponds to oscillation frequency.
- Multiplicative regularization methods: Multiplicative STE variations such as EWGS and PSG cannot prevent oscillations because their frequency remains independent of learning rate.With a sufficiently small learning rate, weights still converge to a decision threshold and begin oscillating.
- Multiplicative regularization methods: EWGS and PSG reduce oscillation amplitude by applying combined positive gradient scaling below one, effectively reducing the learning rate.This analysis assumes a small oscillation amplitude, which is typical toward the end of training.
- Multiplicative regularization methods: Earlier in training, oscillation dynamics may be more complicated because gradient or velocity depends on distance from the closest quantization bin.
A.4. The bias of the STE
The paper distinguishes the probabilistic bias of the STE from its meaning in conventional QAT. In QAT, floating-point shadow weights are not used in the forward pass, making conventional gradient comparisons potentially misleading.
- Probabilistic bias: The STE is described as biased in probabilistic optimization because expected gradients generally differ from the gradient of the expectation.This formulation concerns discrete latent variables and their expectations.
- Bias in QAT: In conventional QAT, quantized networks lack an inherent probabilistic property unless stochastic rounding or another probabilistic formulation is explicitly used.
- Bias in QAT: Floating-point shadow weights are not used in the forward pass and can instead be viewed as a reservoir accumulating small gradient updates.Therefore, comparing their gradients with floating-point gradients may be misleading and nonconstructive for QAT.