Source-linked AI summary
TernaryBERT: Distillation-aware Ultra-low Bit BERT
Wei Zhang, Lu Hou, Yichun Yin, Lifeng Shang, Xiao Chen, Xin Jiang, Qun Liu
TL;DR
BERT's computation and memory demands hinder deployment on resource-constrained devices. TernaryBERT combines ultra-low-bit weight ternarization, granularity choices, and knowledge distillation, achieving performance comparable to full-precision BERT and other reported quantization results while using substantially less storage.
Problem
BERT models are computation- and memory-intensive because they contain hundreds of millions of parameters, limiting deployment on resource-constrained devices.
Method
TernaryBERT combines approximation-based and loss-aware weight ternarization at different BERT granularities with distillation losses during ternary-model training.
Results
TernaryBERT outperforms other BERT quantization methods and is comparable to full-precision BERT, including on GLUE and SQuAD evaluations.
Takeaways & Limitations
TernaryBERT achieves comparable full-precision performance while reducing model size, including a 14.9× smaller model on the GLUE benchmark.
Abstract
from arXiv · showhide
Transformer-based pre-training models like BERT have achieved remarkable performance in many natural language processing tasks.However, these models are both computation and memory expensive, hindering their deployment to resource-constrained devices. In this work, we propose TernaryBERT, which ternarizes the weights in a fine-tuned BERT model. Specifically, we use both approximation-based and loss-aware ternarization methods and empirically investigate the ternarization granularity of different parts of BERT. Moreover, to reduce the accuracy degradation caused by the lower capacity of low bits, we leverage the knowledge distillation technique in the training process. Experiments on the GLUE benchmark and SQuAD show that our proposed TernaryBERT outperforms the other BERT quantization methods, and even achieves comparable performance as the full-precision model while being 14.9x smaller.
1 Introduction
Transformer-based models such as BERT deliver strong NLP performance but are costly in computation and memory, motivating compression for resource-constrained devices. TernaryBERT combines ternarization and knowledge distillation to reduce model size while maintaining accuracy.
- BERT-base has 109M parameters and exceeds 400MB in 32-bit floating-point format, making inference costly for resource-constrained devices.
- Quantization reduces model size by representing weights with fewer bits without changing the model architecture.
- TernaryBERT restricts weights to {−1, 0, +1} and combines quantization with distillation during training.
- The method investigates ternarization granularity across BERT components and applies multiple distillation losses to improve the ternarized student.
- TernaryBERT outperforms other BERT compression methods on MNLI and other evaluations, with performance comparable to full-precision BERT while using much less model size.
2 Related Work
Prior work compresses BERT through distillation, quantization, and other structural methods, while ultra-low-bit quantization has mainly been developed for CNNs. TernaryBERT extends ternarization and distillation to BERT's richer Transformer structure.
- Knowledge Distillation: Knowledge distillation transfers information from teacher logits, intermediate representations, and attentions to guide smaller BERT models.
- Knowledge Distillation: Distillation has also been combined with pruning, low-rank approximation, and dynamic networks to leverage teacher-model knowledge.
- Knowledge Distillation: Prior work had not studied knowledge distillation for training quantized BERT, despite exploring quantization-distillation combinations in CNNs.
- Quantization: Approximation-based quantization keeps quantized weights close to full-precision weights, whereas loss-aware quantization optimizes weights to reduce training loss.
- Quantization: Transformer models have used 8-bit and mixed-precision quantization, with mixed precision using at least 3 bits to avoid severe accuracy loss.
- Quantization: TernaryBERT applies approximation-based and loss-aware ternarization at different granularities and uses distillation losses to address reduced ternary-model capacity.
3 Approach
TernaryBERT combines ternary weights, selective 8-bit activation quantization, and knowledge distillation to train a compact BERT with reduced quantization error and recovered teacher behavior.
- Distillation-aware ternarization: TernaryBERT trains a student whose weights take only −1, 0, or +1, while retaining full-precision weights for parameter updates.At each iteration, the student weights are ternarized for the forward pass, gradients are computed on quantized weights, and the full-precision weights are updated.
- Quantization scope: The method quantizes Transformer and word-embedding weights, linear-layer and matrix-multiplication inputs, while leaving selected embeddings, biases, softmax, normalization, and task-specific parameters unquantized.Segment and position embeddings are excluded, as are biases and operations whose parameters are negligible or whose quantization is not applied.
- Weight ternarization: TWN approximates ternarization by minimizing distance to full-precision weights, whereas LAT directly searches for ternary weights that minimize training loss.TWN uses thresholding and an approximate threshold based on the L1 norm; LAT uses a Hessian-related diagonal approximation available from adaptive optimizers.
- Weight ternarization: Layer-wise ternarization uses one scaling parameter per weight matrix, while row-wise ternarization uses one per row and therefore has finer granularity and smaller quantization error.The paper extends prior layer-level scaling schemes with these two granularities.
- Activation quantization: Activations are quantized to 8 bits with min-max quantization because Transformer hidden representations are negatively skewed, especially in early layers.The authors report that min-max quantization empirically outperforms symmetric quantization and uses a scaling parameter s = (xmax − xmin)/255.
- Knowledge distillation: Knowledge distillation trains the quantized student to recover the full-precision teacher’s behavior across Transformer layers and the prediction layer using representation and attention-based losses.The Transformer-layer objective includes MSE distillation from the teacher’s embedding and layer outputs, plus a loss based on attention scores from all heads.
4 Experiments
Experiments evaluate TernaryBERT on GLUE and SQuAD, compare quantization and compression methods, and ablate ternarization granularity, activation quantization, distillation, initialization, and augmentation. TernaryBERT generally approaches full-precision performance while using substantially smaller models.
- GLUE benchmark: TernaryBERT significantly outperforms Q-BERT and Q2BERT on GLUE, achieves comparable full-precision performance, and is 14.9× smaller.Its proposed method also outperforms Q-BERT on MNLI and SST-2 and Q8BERT on 7 of 8 tasks.
- GLUE benchmark: TernaryTinyBERT achieves comparable full-precision GLUE performance with only a 1.6 point accuracy drop while being 23.2× smaller.
- SQuAD: On SQuAD v1.1 and v2.0, TernaryBERT significantly outperforms Q-BERT and Q2BERT and is comparable to the full-precision baseline.LAT performs slightly better than TWN for this task.
- Ablation Study: Row-wise ternarization benefits word embeddings, whereas layer-wise ternarization performs slightly better for Transformer-layer weights.The resulting empirical setting uses row-wise ternarization for word embeddings and layer-wise ternarization for Transformer-layer weights.
- Ablation Study: Min-max 8-bit activation quantization outperforms symmetric quantization on SQuAD v1.1.The passage attributes this result possibly to non-symmetric hidden-representation distributions.
- Ablation Study: Removing Transformer-layer distillation reduces performance by at least 3% on CoLA and RTE, while removing logit distillation further decreases accuracy across tasks.Initialization from fine-tuned BERT and data augmentation also contribute positively, with larger improvements on CoLA and RTE.
5 Conclusion
The paper ternarizes BERT weights using approximation-based and loss-aware methods with different granularities for embeddings and Transformer-layer weights. Distillation reduces accuracy loss from quantization, and TernaryBERT performs comparably to full-precision BERT while outperforming prior quantization methods.
- TernaryBERT combines approximation-based and loss-aware weight ternarization with different granularities for word embeddings and Transformer-layer weights.
- Knowledge distillation is used to reduce the accuracy drop caused by the lower capacity of quantized weights.
- TernaryBERT outperforms state-of-the-art BERT quantization methods and performs comparably to full-precision BERT.
A Distributions of Hidden Representations on SQuAD v1.1
On SQuAD v1.1, hidden representations are distributed differently across BERT layers: early layers are biased toward negative values, while later layers are not.
- Figure 4 shows hidden-representation distributions from the embedding layer and all Transformer layers on SQuAD v1.1.
- Early layers, including the embedding layer and Transformer layers 1–8, have hidden representations biased toward negative values.
- The remaining Transformer layers do not show the same negative-value bias.
B More Comparison between TernaryBERT and Q-BERT
Additional comparisons on SST-2 and SQuAD v1.1 show that TernaryBERT performs better than mixed-precision Q-BERT.
- TernaryBERT achieves better performance than mixed-precision Q-BERT on SST-2 and SQuAD v1.1.
- Table 9 compares TernaryBERT with mixed-precision Q-BERT.
C Training Curve on MNLI
On MNLI-m, 8-bit BERT has a lower training loss and higher validation accuracy than TernaryBERT, while TernaryBERT’s TWN and LAT learning curves are not significantly different.
- 8-bit BERT has smaller training loss and higher validation accuracy than TernaryBERT on MNLI-m.
- TernaryBERT using TWN and LAT shows no significant difference in its learning curves.
D 3-bit BERT and TinyBERT
Allowing 3-bit weights improves TernaryBERT’s GLUE performance and narrows its accuracy gap with full-precision BERT.
- 3-bit BERT performs slightly better than TernaryBERTLAT on 7 of 8 GLUE tasks.The 3-bit variant replaces LAT with 3-bit Loss-aware Quantization (LAQ).
- The 3-bit model has a smaller accuracy gap with the full-precision baseline than TernaryBERTLAT.
E Attention Pattern of BERT and TernaryBERT
The attention patterns of ternarized TernaryBERT closely resemble those of fine-tuned full-precision BERT on grammatical-acceptability and sentiment-classification examples.
- TernaryBERT’s attention patterns resemble those of full-precision BERT on CoLA and SST-2.CoLA predicts grammatical acceptability, while SST-2 classifies the polarity of movie reviews.
- The comparisons use fine-tuned full-precision BERT-base and ternarized TernaryBERTTWN models.
- Table 10 reports development-set results for 3-bit quantized BERT and TinyBERT on GLUE.
- Figures 6 and 7 show attention patterns for CoLA sentences.The examples concern grammatical-acceptability prediction.
- Figures 8 and 9 show attention patterns for SST-2 sentences.The examples are movie-review phrases used for sentiment classification.