Continuous Approximations for Improving Quantization Aware Training of LLMs

TL;DR

This paper introduces continuous approximation techniques for quantization aware training of large language models, significantly reducing performance loss and improving accuracy while maintaining energy efficiency.

Contribution

It proposes novel continuous approximations for the rounding and clamping functions in QAT, leading to better model performance and more accurate step size and weight learning.

Findings

Perplexity on WikiText-v2 improved to 9.0815 from 9.9621

2.76% improvement on BoolQ

5.47% improvement on MMLU

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.