[Paper Review] Q-BERT: Hessian Based Ultra Low Precision Quantization of BERT

Sheng Shen, Zhen Dong+6 more|arXiv (Cornell University)|Sep 12, 2019

Topic ModelingComputer Science42 references52 citations

TL;DR

The paper presents Q-BERT, a Hessian-informed mixed-precision and group-wise quantization scheme for BERT that achieves up to 13x weight compression with at most 2.3% accuracy loss across SST-2, MNLI, CoNLL-03, and SQuAD. Embedding and encoder layers are quantized differently, and group-wise quantization reduces degradation further, with SQuAD being the most challenging task.

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 imes$ compression of the model parameters, and up to $4 imes$ 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.

Motivation & Objective

• Motivate reducing BERT memory and latency for edge deployment without unacceptable accuracy loss.
• Analyze fine-tuned BERT using second-order Hessian information to guide quantization decisions.
• Propose a Hessian-based mixed-precision scheme for encoder layers and a group-wise quantization scheme for self-attention modules.
• Demonstrate ultra-low-precision quantization achieving substantial compression with minimal performance degradation across multiple NLP tasks.

Proposed method

• Compute Hessian-based sensitivity per encoder layer using the top eigenvalues of each layer's Hessian spectrum.
• Define a sensitivity metric Omega_i = mean(lambda_i) + std(lambda_i) from Hessian eigenvalue distributions computed on 10% of data.
• Apply mixed-precision quantization by assigning higher bits to more sensitive layers based on Omega_i.
• Introduce group-wise quantization by partitioning each dense matrix (e.g., within MHSA heads) into groups with separate quantization ranges.
• Quantize embeddings and encoder parameters with different schemes and perform quantization-aware fine-tuning.
• Use a uniform 8-bit activation scheme and compare against DirectQ baselines to measure accuracy retention.

Experimental results

Research questions

• RQ1 How does Hessian information (top eigenvalues and their distribution) correlate with quantization sensitivity in BERT layers?
• RQ2 Can mixed-precision, guided by Hessian analysis, preserve accuracy under ultra-low bit quantization (2-4 bits) for BERT?
• RQ3 Does group-wise quantization improve performance when quantizing self-attention and feed-forward components of BERT?
• RQ4 Which BERT modules (embedding vs encoder layers) are most sensitive to quantization, and how should they be quantized?
• RQ5 Why is SQuAD more challenging to quantize compared to other NLP tasks?

Key findings

• Q-BERT achieves up to 13× weight compression and 4× reductions in embedding and activation sizes with at most 2.3% accuracy loss across SST-2, MNLI, CoNLL-03, and SQuAD.
• Hessian-based mixed-precision (2/3 or 2/4 bits) outperforms uniform 2-bit quantization, especially for deeper layers; middle encoder layers are most sensitive while last layers are more robust.
• Group-wise quantization (with 128 groups) significantly reduces accuracy loss versus layer-wise quantization, with diminishing returns beyond a certain group count.
• Embedding quantization is more sensitive than encoder weights, and position embeddings are particularly critical for preserving performance.
• SQuAD exhibits larger Hessian eigenvalue variance and negative curvature at convergence, which correlates with greater accuracy loss under ultra-low precision.