PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Sheng Shen, Zhen Dong, Jiayu Ye, Linjian Ma, Zhewei Yao, Amir Gholami, Michael W. Mahoney, Kurt Keutzer

arXiv:1909.05840v2cs.CLcs.LG

TL;DR

BERT’s memory footprint and latency hinder deployment in resource-constrained environments, while ultra-low precision can harm accuracy. Q-BERT uses Hessian-based mixed precision and group-wise quantization, achieving up to 13× weight compression with at most 2.3% accuracy loss across four tasks. SQuAD shows the largest degradation, associated with non-converged fine-tuning and unfavorable Hessian behavior.

  • Problem

    Ultra-low-precision quantization must reduce BERT’s memory and latency without unacceptable performance degradation for resource-constrained deployment.

  • Method

    Q-BERT combines layer-wise Hessian analysis using eigenvalue means and variances with group-wise quantization using separate ranges and lookup tables.

  • Results

    Across four downstream tasks, Q-BERT achieves up to 13× weight compression with at most 2.3% accuracy loss, while reducing embedding and activation sizes by up to 4×.

  • Takeaways & Limitations

    Hessian-guided mixed precision enables Q-BERT to retain performance at ultra-low precision, including 2/3-bit settings, while maintaining substantial compression.

  • Takeaways & Limitations

    SQuAD is harder to quantize because its fine-tuned BERT model has large Hessian variance and negative eigenvalues, indicating non-convergence to a local minimum.

Abstract

from arXiv · show

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.

1. Introduction.

Ultra-low-precision quantization targets BERT’s prohibitive memory and latency while limiting accuracy degradation. Q-BERT combines Hessian-guided mixed precision with group-wise quantization to improve hardware-efficient compression.

  • Motivation: Ultra-low-bit quantization can significantly degrade generalization performance, despite reducing memory footprint and accelerating inference.The resulting efficiency can support edge deployment on FPGAs and domain-specific accelerators.
  • Proposed approach: Q-BERT applies Hessian-guided mixed-precision quantization and introduces group-wise quantization with separate ranges and lookup tables for matrix groups.Its Hessian sensitivity measure uses both the mean and variance of top eigenvalues.
  • Hessian analysis: Flatter-curvature layers can be assigned lower bit precision during quantization.The loss landscape is visualized by perturbing parameters along the first two dominant Hessian eigenvectors.

2. Related Work.

Related work frames model compression as a broad set of architecture, training, pruning, hardware, and quantization approaches. The paper focuses on quantization because ultra-low precision creates accuracy and mixed-precision search challenges, especially for large Transformer models.

  • Model compression: Model compression includes compact architectures, neural architecture search, pruning, knowledge distillation, hardware co-design, and inference quantization.These categories target reduced model size or latency through different design and training strategies.
  • Quantization: Ultra-low-precision quantization can significantly degrade accuracy, while mixed-precision search has an exponentially large search space.With three precision options, a fine-tuned BERT model has approximately 5.3×10^5 possible assignments.
  • Compressed NLP models: Transformer models have grown substantially because they incorporate very large fully connected layers and attention matrices, making compression important for resource-constrained deployment.Prior NLP compression work includes LSTM- and GRU-based models and early Transformer-focused efforts.

3. Methodology.

Q-BERT quantizes BERT embeddings and encoder layers using uniform, group-wise quantization and Hessian-guided mixed precision. Its analysis identifies layer sensitivity variation and non-convergence in SQuAD as factors shaping quantization choices and degradation.

  • Model and quantization setup: BERTBASE comprises embedding, Transformer encoder, and output layers; Q-BERT quantizes the embedding and encoder while leaving the negligible-size output layer unquantized.The embedding layer is more sensitive to quantization than encoder layers, so the paper applies different schemes to them.
  • Quantization process: Uniform quantization maps real-valued weights or activations into discrete unsigned integers using equally spaced intervals and k-bit precision.The method uses a Straight-through Estimator to backpropagate through the non-differentiable quantization operator.
  • Hessian-guided mixed precision: Different encoder layers exhibit different Hessian sensitivities, making uniform bit assignments sub-optimal; Q-BERT assigns more bits to more sensitive layers.Top Hessian eigenvalues are estimated with matrix-free power iteration, avoiding explicit construction of each 7M × 7M layer Hessian.
  • Hessian-guided mixed precision: Q-BERT ranks layers using the distribution of top Hessian eigenvalues computed from 10% of the training data, then performs quantization-aware fine-tuning with the selected precisions.Using the eigenvalue distribution addresses the high variance observed across data portions, rather than relying only on average eigenvalues.
  • Convergence condition: SQuAD is expected to suffer greater quantization degradation because its fine-tuned BERT model has large negative Hessian eigenvalues and has not reached a local minimum.The method assumes the pretrained and fine-tuned model has converged to a local minimum with zero gradient and positive curvature.
  • Group-wise quantization: Group-wise quantization partitions attention weights into groups or output-neuron sub-groups, allowing each group to use its own quantization range instead of one range for an entire tensor.For multi-head self-attention, the method groups matrices by head and can further bucket continuous output neurons within each group.

4. Experiment.

Experiments across four NLP tasks show that Q-BERT substantially outperforms direct quantization, especially at ultra-low precision. Hessian-guided mixed precision, group-wise quantization, and hardware-aware bit choices preserve accuracy while achieving high compression, although SQuAD remains harder to quantize.

  • 4.1. Main Results: Q-BERT outperforms DirectQ across SST-2, MNLI, CoNLL-03, and SQuAD at every tested bit setting, with larger gaps at lower precision.At 3 bits, the performance gap ranges from 9.68% to 27.83% across tasks.
  • 4.1. Main Results: 11.5% DirectQ degradation on 4-bit SQuAD falls to 0.5% with Q-BERT.Both figures compare against BERTBASE in the same 4-bit setting.
  • 4.1. Main Results: Q-BERTMP assigns higher precision to sensitive middle encoder layers, retaining performance drops within 2.3% for MNLI and SQuAD and 1.1% for SST-2 and CoNLL-03.The mixed-precision models use only 5MB additional memory and reach up to 13× weight compression.
  • 4.1. Main Results: A mixed 2/4-bit configuration matches 3-bit SST-2 model size at 53.2MB and achieves similar accuracy, addressing limited hardware support for 3-bit execution.The paper notes that 3-bit values can alternatively be cast to higher precision during execution.
  • 4.1. Main Results: SQuAD is harder to quantize: 2-bit DirectQ reaches only 10% F1, while its Hessian exhibits larger eigenvalue variance and large negative eigenvalues.The authors relate this loss landscape to BERT not reaching a local minimum for SQuAD.
  • 4.2. Effects of group-wise quantization: Group-wise quantization reduces performance degradation to below 1% with 12 groups, while gains largely saturate at 128 groups.Increasing groups from 128 to 768 improves performance by at most 0.1% and increases LUT requirements.

5. Discussion.

Q-BERT’s module-level analysis finds embeddings are more sensitive than encoder weights, with position embeddings especially fragile, while self-attention is more robust than fully connected layers. Attention-distribution analysis further indicates Q-BERT better preserves the full-precision model’s behavior than DirectQ.

  • Quantization effects on different modules: Up to 10% performance drops for 4-bit embedding quantization across SST-2, MNLI, CoNLL-03, and more than 20% for SQuAD, despite 8-bit weights and activations.Encoder layers contain around 79% of total parameters, yet 4-bit encoder quantization causes less performance loss.
  • Quantization effects on different modules: 2% additional performance degradation generally results from 4-bit position embeddings compared with 4-bit word embeddings.Position embeddings account for less than 5% of the entire embedding, motivating mixed-precision treatment.
  • Quantization effects on different modules: 7% performance drop occurs for 1/2MP self-attention, versus 11% for 1/2MP fully connected layers, indicating greater self-attention robustness.The experiments use Q-BERTMP with mixed 2- and 3-bit encoder weights, then reduce one additional bit in the tested module.
  • Qualitative Analysis: Q-BERT’s attention distributions remain closer to the full-precision baseline than DirectQ’s across SST-2, MNLI, CoNLL-03, and SQuAD.The comparison uses KL divergence over coordinated attention heads, averaged over 10% of the training dataset, with 4-bit weights and 8-bit embeddings and activations.
  • Qualitative Analysis: Q-BERT combines Hessian-guided mixed precision with fine-grained group-wise quantization to reduce model size while preserving accuracy.The discussion identifies layer-wise second-order analysis and group-wise quantization as core components.

6. Conclusion.

Q-BERT uses Hessian-informed layer-wise precision assignment and group-wise quantization to compress BERT. Across four downstream tasks, it reaches substantial compression with at most 2.3% accuracy loss.

  • Conclusion: Q-BERT’s attention-to-baseline KL divergence is much smaller than DirectQ’s, indicating closer attention distributions.Figure 5 compares attention distributions between Q-BERT or DirectQ and the baseline.
  • Conclusion: 13× weight compression, 4× smaller activations, and 4× smaller embeddings are achieved across four downstream tasks with at most 2.3% accuracy loss.The method combines a Hessian-based layer-wise approach capturing eigenvalue mean and variance with new group-wise quantization.

Algorithm 1: Power Iteration for Eigenvalue Computation

The power-iteration procedure estimates eigenvalue information for a BERT block’s Hessian using backpropagated gradients and a normalized random vector. The quantization pipeline then maps tensor values to integer indices within a selected range.

  • Power Iteration for Eigenvalue Computation: The algorithm takes a block parameter Wi, computes its backpropagated gradient gi, and initializes a normalized random vector v for power iteration.The supplied algorithm passage introduces the Hessian-eigenvalue computation inputs and initialization.
  • Detailed quantization process: During the forward pass, each element of a weight or activation tensor X is quantized using the paper’s quantization rule.The procedure applies to both weights and activations.
  • Detailed quantization process: The quantization rule rounds scaled tensor values to integer indices using adjacent-point spacing Δ and clamps values to the selected range [q0, q2^k−1].The range may be a subinterval of [min, max] to reduce outlier effects and better represent most tensor values.

Appendix B. Dataset.

Q-BERT is evaluated on four BERT downstream task categories using SST-2, CoNLL-2003, MNLI, and SQuAD. These datasets cover sentiment, natural-language inference, named-entity recognition, and machine reading comprehension.

  • Dataset: The evaluation covers sentiment classification, natural-language inference, named-entity recognition, and machine reading comprehension.The corresponding datasets are SST-2, MNLI, CoNLL-03, and SQuAD.
  • Dataset: SST-2 contains movie reviews with binary positive or negative labels, while MNLI predicts entailment, contradiction, or neutrality for premise-hypothesis pairs.MNLI includes matched in-domain and mismatched cross-domain splits.

Appendix C. Extra results. Here we describe several additional results.

The ablation reverses Q-BERTMP’s bit assignments, giving lower precision to sensitive layers and higher precision to less sensitive layers while preserving model size.

  • Ablation setup: Q-BERTMP-rev assigns lower bits to relatively sensitive layers and higher bits to relatively insensitive layers.The reversed assignment swaps the 2-bit and 3-bit settings used by the corresponding layers in Q-BERTMP.
  • Ablation setup: The ablation keeps model size unchanged while reversing the precision assignment.
  • Ablation setup: The comparison tests whether Hessian-guided mixed-precision assignments outperform their reversed counterpart.

C.1. Ablation Study of Hessian based Mixed Precision Assignment.

At the same model size, reversing Q-BERTMP’s Hessian-based bit assignment substantially reduces the performance gap relative to 2-bit Q-BERT, supporting Hessian-guided allocation.

  • Ablation results: Within 2% performance difference is observed for Q-BERTMP-rev versus 2-bit Q-BERT on MNLI, CoNLL-03, and SQuAD.
  • Ablation results: Within 4% performance difference is observed for Q-BERTMP-rev versus 2-bit Q-BERT on SST-2.
  • Ablation results: Beyond 5% performance differences occur for Q-BERTMP versus 2-bit Q-BERT on MNLI, CoNLL-03, and SQuAD, and beyond 8% on SST-2.
  • Ablation setup: The 2/3-bit Q-BERTMP bit settings are included in Table 6.
  • Ablation results: Table 4 reports quantization results for the reversed Hessian-based mixed-precision setting, Q-BERTMP-rev.

C.2. Mixed Precision Quantization for Embedding.

The embedding experiments further reduce model size by assigning mixed precisions to word and position embeddings, with results reported across four tasks.

  • Embedding mixed precision: The embedding table is assigned 4-bit word embeddings and 8-bit position embeddings to reduce total model size.This targets the embedding table because it bounds model size under 2/3-bit weight quantization, while word embeddings are less sensitive.
  • Embedding mixed precision: Table 5 reports embedding mixed-precision results for Q-BERT on four tasks using 128 groups in each encoder and embedding layer.
  • Embedding mixed precision: Except for the baseline, the embedding experiments use 8-bit activations and report model size with and without the embedding layer.
  • Additional analysis: Figure 6 visualizes the SST-2 loss landscape by perturbing parameters along the first two dominant Hessian eigenvectors.The silver sphere marks the point where the BERT model converged.
  • Additional analysis: Table 6 lists the 2/3-bit Q-BERTMP bit settings for all four tasks.
  • Additional analysis: Table 7 lists the 2/4-bit Q-BERTMP bit settings for all four tasks.
Loading 1909.05840v2…