Continuous Approximations for Improving Quantization Aware Training of LLMs

He Li,Jianhang Hong,Yuanzhuo Wu,Snehal Adbol,Zonglin Li
2024-10-06
Abstract:Model compression methods are used to reduce the computation and energy requirements for Large Language Models (LLMs). Quantization Aware Training (QAT), an effective model compression method, is proposed to reduce performance degradation after quantization. To further minimize this degradation, we introduce two continuous approximations to the QAT process on the rounding function, traditionally approximated by the Straight-Through Estimator (STE), and the clamping function. By applying both methods, the perplexity (PPL) on the WikiText-v2 dataset of the quantized model reaches 9.0815, outperforming 9.9621 by the baseline. Also, we achieve a 2.76% improvement on BoolQ, and a 5.47% improvement on MMLU, proving that the step sizes and weights can be learned more accurately with our approach. Our method achieves better performance with the same precision, model size, and training setup, contributing to the development of more energy-efficient LLMs technology that aligns with global sustainability goals.
Machine Learning,Artificial Intelligence,Computation and Language
What problem does this paper attempt to address?
The paper attempts to address the issue of performance degradation during Quantization-Aware Training (QAT) of large language models (LLMs). Specifically, the paper introduces two continuous approximation methods—Sigmoid STE and SoftClamp—to improve the approximation of rounding and truncation functions in traditional QAT methods. These improvements aim to more effectively learn the step size and quantize model weights, thereby enhancing the performance of the model after quantization. The paper mentions that although quantization can significantly reduce computational complexity and energy consumption, it usually leads to a decline in model performance. Through the proposed methods, the authors hope to improve the performance of quantized models while maintaining the same accuracy, model size, and training settings, making them closer to the performance of unquantized models. This advancement would promote the development of more energy-efficient LLM technology, aligning with global sustainability goals.