Q-BERT: Hessian Based Ultra Low Precision Quantization of BERT

TL;DR

This paper introduces a Hessian-based ultra low precision quantization method for BERT, significantly reducing model size with minimal performance loss across multiple NLP tasks.

Contribution

It proposes a novel group-wise quantization scheme combined with Hessian-based mix-precision, enabling efficient compression of BERT models for resource-constrained deployment.

Findings

Achieves up to 13x compression with only 2.3% performance degradation.

Effective on multiple NLP tasks including SST-2, MNLI, CoNLL-03, and SQuAD.

Highest performance loss observed on SQuAD due to training convergence issues.

Abstract

Transformer based architectures have become de-facto models used for a range of Natural Language Processing tasks. In particular, the BERT based models achieved significant accuracy gain for GLUE tasks, CoNLL-03 and SQuAD. However, BERT based models have a prohibitive memory footprint and latency. As a result, deploying BERT based models in resource constrained environments has become a challenging task. In this work, we perform an extensive analysis of fine-tuned BERT models using second order Hessian information, and we use our results to propose a novel method for quantizing BERT models to ultra low precision. In particular, we propose a new group-wise quantization scheme, and we use a Hessian based mix-precision method to compress the model further. We extensively test our proposed method on BERT downstream tasks of SST-2, MNLI, CoNLL-03, and SQuAD. We can achieve comparable performance to baseline with at most $2.3\%$ performance degradation, even with ultra-low precision quantization down to 2 bits, corresponding up to $13\times$ compression of the model parameters, and up to $4\times$ compression of the embedding table as well as activations. Among all tasks, we observed the highest performance loss for BERT fine-tuned on SQuAD. By probing into the Hessian based analysis as well as visualization, we show that this is related to the fact that current training/fine-tuning strategy of BERT does not converge for SQuAD.