Quantization aware pre training cuts noise and boosts large model accuracy
Quantization-Aware Pre-Training with Constrained Empirical Weight Distribution
Machine Learning
Summary
Quantization-aware pre-training helps make large neural networks faster but can cause a problem where weights keep jumping back and forth, slowing down learning. The authors propose a way to keep the weights shaped like a bell curve (a normal distribution) during training to stop this jumping. This method doesn’t add extra settings or memory use but still improves training stability and final model quality. They tested it on big language models with very low precision and saw consistent benefits without extra costs.
What this means in practice
- •For machine learning engineers: Improve training stability and accuracy when preparing large language models for efficient low-precision inference.
- •For hardware acceleration teams: Deploy neural networks that maintain higher accuracy despite aggressive quantization, enabling faster inference on edge or specialized AI chips.
Authors
Ningfeng Yang, Tor M. Aamodt
Abstract
Quantization-Aware Pre-Training (QAPT) can increase the inference efficiency of DNNs, but a problematic behaviour known as rounding boundary weight oscillation can introduce detrimental noise into the training process and significantly reduce convergence speed. While existing methods can reduce this detrimental noise, they either introduce additional hyperparameters or memory overhead, or cannot consistently improve model accuracy. In this work, we propose optimization with $\textbf{C}$onstrained $\textbf{E}$mpirical $\textbf{W}$eight dis$\textbf{T}$ribution (CEWT), the first hyperparameter-free memory-overhead-free oscillation suppression method that consistently improves QAPT performance: an optimizer post-update step that projects weights to the nearest point in weight space whose empirical distribution (histogram) matches a zero-mean Gaussian. Our key insight is many quantizers are designed with the implicit assumption that the to-be-quantized data are permutations of samples from a zero-mean Gaussian, and this assumption is not true during QAPT. By enforcing the zero-mean Gaussian prior as a hard constraint, CEWT can suppress this detrimental noise. Empirical results on various combinations of SOTA quantizers and hypersphere optimizers suggest, that with a geomean increase of 4% in training time, CEWT can consistently reduce the pre-training perplexity (by an average of 2.5 and up to 21 points) of low-precision (down to 1-bit activations and weights and up to 610M parameters) LLaMA/GPT models without introducing any hyperparameters or storage overhead. Code is available at https://github.com/1733116199/cewt